Agent/gateway_server.py

"""
A2A IM Gateway Server

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

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

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)

    # 创建 Router
    router = GatewayRouter(registry)

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

    # 启动和关闭事件
    @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()