Agent/gateway_server.py

"""
A2A IM Gateway Server

启动 Gateway 服务器，提供 Agent 注册和消息路由
"""

import logging

import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from gateway.core.channels.loader import load_enabled_channels
from gateway.core.executor import TaskManager, build_executor_router
from gateway.core.lifecycle import TraceManager, WorkspaceManager
from gateway.core.registry import AgentRegistry
from gateway.core.router import GatewayRouter

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)

logger = logging.getLogger(__name__)


def create_gateway_app() -> FastAPI:
    """创建 Gateway FastAPI 应用"""

    app = FastAPI(
        title="A2A IM Gateway",
        description="Agent 即时通讯网关",
        version="1.0.0",
    )

    # 添加 CORS 中间件
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # 创建 Registry
    registry = AgentRegistry(heartbeat_timeout=60)

    # 创建 Gateway Router
    gateway_router = GatewayRouter(registry)

    # 注册 Gateway 路由
    app.include_router(gateway_router.router)

    for router in load_enabled_channels():
        app.include_router(router)

    _wm = WorkspaceManager.from_env()
    _tm = TraceManager.from_env(_wm)
    _task_mgr = TaskManager.from_env(_wm, _tm)
    app.include_router(build_executor_router(_task_mgr))
    logger.info("Gateway Executor mounted at /gateway/executor")

    # 启动和关闭事件
    @app.on_event("startup")
    async def startup():
        await registry.start()
        logger.info("Gateway started")

    @app.on_event("shutdown")
    async def shutdown():
        await registry.stop()
        logger.info("Gateway stopped")

    # 健康检查
    @app.get("/health")
    async def health():
        return {"status": "ok"}

    return app


def main():
    """启动 Gateway 服务器"""

    app = create_gateway_app()

    # 启动服务器
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info",
    )


if __name__ == "__main__":
    main()