diff --git a/nocobase_crm_mcp/server.py b/nocobase_crm_mcp/server.py index 52f51e2..0e67f1c 100644 --- a/nocobase_crm_mcp/server.py +++ b/nocobase_crm_mcp/server.py @@ -402,11 +402,13 @@ def _read_token_from_env(env_name: str) -> str: def create_http_app(server: Any, token: str, upload_endpoint: UploadEndpoint | None = None) -> Any: if upload_endpoint is None: return BearerTokenAuthMiddleware(server.streamable_http_app(), token) + mcp_app = server.streamable_http_app() app = Starlette( routes=[ Route("/uploads", endpoint=upload_endpoint, methods=["POST"]), - Mount("/", app=server.streamable_http_app()), - ] + Mount("/", app=mcp_app), + ], + lifespan=mcp_app.router.lifespan_context, ) return BearerTokenAuthMiddleware(app, token) diff --git a/tests/test_server.py b/tests/test_server.py index 0fc2cac..f8ac7bb 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -1,9 +1,11 @@ import asyncio import unittest +from contextlib import asynccontextmanager from unittest.mock import patch import httpx from starlette.responses import JSONResponse +from starlette.applications import Starlette from nocobase_crm_mcp import server as server_module from nocobase_crm_mcp.nb_client import NbAuthenticationError, NbCommandError @@ -102,6 +104,21 @@ class FakeUploadEndpoint: return JSONResponse({"uploaded": True}) +async def run_asgi_lifespan(app, on_started): + messages = asyncio.Queue() + await messages.put({"type": "lifespan.startup"}) + + async def receive(): + return await messages.get() + + async def send(message): + if message["type"] == "lifespan.startup.complete": + on_started() + await messages.put({"type": "lifespan.shutdown"}) + + await app({"type": "lifespan"}, receive, send) + + class ServerTests(unittest.TestCase): def test_bearer_token_middleware_rejects_missing_or_wrong_token(self) -> None: app = BearerTokenAuthMiddleware(_ok_app, "secret") @@ -248,6 +265,26 @@ class ServerTests(unittest.TestCase): self.assertEqual(response.status_code, 200) self.assertEqual(response.json(), {"uploaded": True}) + def test_upload_route_preserves_mcp_application_lifespan(self) -> None: + lifecycle = {"started": False, "stopped": False} + + @asynccontextmanager + async def lifespan(_app): + lifecycle["started"] = True + try: + yield + finally: + lifecycle["stopped"] = True + + class FakeServer: + def streamable_http_app(self): + return Starlette(lifespan=lifespan) + + app = create_http_app(FakeServer(), "secret", upload_endpoint=FakeUploadEndpoint()) + + asyncio.run(run_asgi_lifespan(app, lambda: self.assertTrue(lifecycle["started"]))) + self.assertTrue(lifecycle["stopped"]) + def test_nocobase_environment_prefers_canonical_shared_credentials(self) -> None: with patch.dict( "os.environ",