fix(auth): make get_access_token per-request in stateful sessions

Author: Epochex

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

@@ -1,5 +1,8 @@
 import contextvars
 
+from contextvars import Token
+
+from starlette.requests import Request
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
@@ -20,6 +23,26 @@ def get_access_token() -> AccessToken | None:
     return auth_user.access_token if auth_user else None
 
 
+def _push_auth_context_from_request(request: Request | None) -> Token[AuthenticatedUser | None] | None:
+    """Set auth context for the current task from an incoming request.
+
+    This is primarily used by server transports where request handlers may run
+    in background tasks that are not part of the original ASGI request task.
+    """
+    if request is None:
+        return None
+    user = getattr(request, "user", None)
+    if isinstance(user, AuthenticatedUser):
+        return auth_context_var.set(user)
+    return None
+
+
+def _pop_auth_context(token: Token[AuthenticatedUser | None] | None) -> None:
+    if token is None:
+        return
+    auth_context_var.reset(token)
+
+
 class AuthContextMiddleware:
     """Middleware that extracts the authenticated user from the request
     and sets it in a contextvar for easy access throughout the request lifecycle.

src/mcp/server/runner.py

@@ -26,6 +26,7 @@
 from pydantic import BaseModel, ValidationError
 from typing_extensions import TypeVar
 
+from mcp.server.auth.middleware.auth_context import _pop_auth_context, _push_auth_context_from_request
 from mcp.server.connection import Connection
 from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
 from mcp.server.models import InitializationOptions
@@ -259,7 +260,11 @@ async def _inner() -> HandlerResult:
             return result
 
         call = self._compose_server_middleware(ctx, method, params, _inner)
-        result = _dump_result(await call())
+        auth_token = _push_auth_context_from_request(ctx.request)
+        try:
+            result = _dump_result(await call())
+        finally:
+            _pop_auth_context(auth_token)
         if method == "initialize":
             # Commit only on chain success, so a middleware veto leaves no state.
             # Race-free: the read loop is parked until this call returns.

tests/server/auth/test_get_access_token_streamable_http.py

@@ -0,0 +1,100 @@
+import time
+
+import httpx
+import pytest
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.routing import Mount
+
+from mcp import Client
+from mcp.client.streamable_http import streamable_http_client
+from mcp.server import Server, ServerRequestContext
+from mcp.server.auth.middleware.auth_context import AuthContextMiddleware, get_access_token
+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:
+    """Accepts any bearer token and echoes it back as the verified AccessToken."""
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        return AccessToken(token=token, client_id=token, scopes=[], expires_at=int(time.time()) + 3600)
+
+
+async def _handle_whoami(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
+    access = get_access_token()
+    text = access.token if access else "<none>"
+    return CallToolResult(content=[TextContent(type="text", text=text)])
+
+
+async def _handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
+    return ListToolsResult(tools=[Tool(name="whoami", input_schema={"type": "object", "properties": {}})])
+
+
+class _MutableBearerAuth(httpx.Auth):
+    def __init__(self, token: str) -> None:
+        self.token = token
+
+    def auth_flow(self, request: httpx.Request):
+        request.headers["Authorization"] = f"Bearer {self.token}"
+        yield request
+
+
+@pytest.mark.anyio
+async def test_get_access_token_reflects_current_request_in_stateful_session() -> None:
+    host = "testserver"
+
+    server = Server(
+        "auth-test-server",
+        on_call_tool=_handle_whoami,
+        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)
+
+    asgi_app = Starlette(
+        routes=[Mount("/mcp", app=session_manager.handle_request)],
+        middleware=[
+            Middleware(AuthenticationMiddleware, backend=BearerAuthBackend(_EchoTokenVerifier())),
+            Middleware(AuthContextMiddleware),
+        ],
+        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"