# main.py
import contextlib
from starlette.applications import Starlette
from starlette.routing import Mount

from test_tools import create_test_mcp
from system_tools import create_system_mcp  # 如果暂时不用可先不挂

# 先创建 MCP 实例
test_mcp = create_test_mcp()
# system_mcp = create_system_mcp()

# 关键：在 Starlette 的 lifespan 中启动 MCP 的 session manager
@contextlib.asynccontextmanager
async def lifespan(app: Starlette):
    async with contextlib.AsyncExitStack() as stack:
        await stack.enter_async_context(test_mcp.session_manager.run())
        # await stack.enter_async_context(system_mcp.session_manager.run())
        yield  # 服务器运行期间
        # 退出时自动清理

app = Starlette(
    lifespan=lifespan,
    routes=[
        Mount("/test", app=test_mcp.streamable_http_app()),
        # Mount("/system", app=system_mcp.streamable_http_app()),
    ],
)