fix auth context reset for streamable HTTP

Author: Jianke LIN

src/mcp/server/auth/middleware/auth_context.py

@@ -38,9 +38,7 @@ def push_auth_context_from_request(request: Request | None) -> Token[Authenticat
             user = getattr(request, "user", None)
         except AssertionError:
             user = None
-    if isinstance(user, AuthenticatedUser):
-        return auth_context_var.set(user)
-    return None
+    return auth_context_var.set(user if isinstance(user, AuthenticatedUser) else None)
 
 
 def pop_auth_context(token: Token[AuthenticatedUser | None] | None) -> None:

tests/server/auth/test_get_access_token_streamable_http.py

@@ -2,6 +2,14 @@
 
 import httpx
 import pytest
+from mcp_types import (
+    CallToolRequestParams,
+    CallToolResult,
+    ListToolsResult,
+    PaginatedRequestParams,
+    TextContent,
+    Tool,
+)
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
@@ -14,15 +22,6 @@
 from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend
 from mcp.server.auth.provider import AccessToken
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
-from mcp.server.transport_security import TransportSecuritySettings
-from mcp.types import (
-    CallToolRequestParams,
-    CallToolResult,
-    ListToolsResult,
-    PaginatedRequestParams,
-    TextContent,
-    Tool,
-)
 
 
 class _EchoTokenVerifier:
@@ -43,14 +42,34 @@ async def _handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequest
 
 
 class _MutableBearerAuth(httpx.Auth):
-    def __init__(self, token: str) -> None:
+    def __init__(self, token: str | None) -> None:
         self.token = token
 
     def auth_flow(self, request: httpx.Request):
-        request.headers["Authorization"] = f"Bearer {self.token}"
+        if self.token is not None:
+            request.headers["Authorization"] = f"Bearer {self.token}"
         yield request
 
 
+async def _call_whoami(asgi_app: Starlette, host: str, token: str | None) -> str:
+    auth = _MutableBearerAuth(token)
+    async with (
+        httpx.ASGITransport(asgi_app) as transport,
+        httpx.AsyncClient(
+            transport=transport,
+            base_url=f"http://{host}",
+            auth=auth,
+            timeout=httpx.Timeout(30, read=30),
+            follow_redirects=True,
+        ) as http_client,
+    ):
+        transport_ctx = streamable_http_client(f"http://{host}/mcp", http_client=http_client)
+        async with Client(transport_ctx) as client:  # pragma: no branch
+            result = await client.call_tool("whoami", {})
+            assert isinstance(result.content[0], TextContent)
+            return result.content[0].text
+
+
 @pytest.mark.anyio
 async def test_get_access_token_reflects_current_request_in_stateful_session() -> None:
     host = "testserver"
@@ -61,11 +80,7 @@ async def test_get_access_token_reflects_current_request_in_stateful_session() -
         on_list_tools=_handle_list_tools,
     )
 
-    security = TransportSecuritySettings(
-        allowed_hosts=[host, f"{host}:*"],
-        allowed_origins=[f"http://{host}:*"],
-    )
-    session_manager = StreamableHTTPSessionManager(app=server, security_settings=security, stateless=False)
+    session_manager = StreamableHTTPSessionManager(app=server, stateless=False)
 
     asgi_app = Starlette(
         routes=[Mount("/mcp", app=session_manager.handle_request)],
@@ -76,25 +91,7 @@ async def test_get_access_token_reflects_current_request_in_stateful_session() -
         lifespan=lambda app: session_manager.run(),
     )
 
-    auth = _MutableBearerAuth("token-A")
     async with asgi_app.router.lifespan_context(asgi_app):
-        async with (
-            httpx.ASGITransport(asgi_app) as transport,
-            httpx.AsyncClient(
-                transport=transport,
-                base_url=f"http://{host}",
-                auth=auth,
-                timeout=httpx.Timeout(30, read=30),
-                follow_redirects=True,
-            ) as http_client,
-        ):
-            transport_ctx = streamable_http_client(f"http://{host}/mcp", http_client=http_client)
-            async with Client(transport_ctx) as client:
-                r1 = await client.call_tool("whoami", {})
-                assert isinstance(r1.content[0], TextContent)
-                assert r1.content[0].text == "token-A"
-
-                auth.token = "token-B"
-                r2 = await client.call_tool("whoami", {})
-                assert isinstance(r2.content[0], TextContent)
-                assert r2.content[0].text == "token-B"
+        assert await _call_whoami(asgi_app, host, "token-A") == "token-A"
+        assert await _call_whoami(asgi_app, host, "token-B") == "token-B"
+        assert await _call_whoami(asgi_app, host, None) == "<none>"