Fix stateful StreamableHTTP auth context rebinding

Author: raashish1601

src/mcp/server/lowlevel/server.py

@@ -38,8 +38,8 @@ async def main():
 
 import logging
 import warnings
-from collections.abc import AsyncIterator, Awaitable, Callable
-from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
 from importlib.metadata import version as importlib_version
 from typing import Any, Generic
 
@@ -52,8 +52,8 @@ async def main():
 from typing_extensions import TypeVar
 
 from mcp import types
-from mcp.server.auth.middleware.auth_context import AuthContextMiddleware
-from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend, RequireAuthMiddleware
+from mcp.server.auth.middleware.auth_context import AuthContextMiddleware, auth_context_var
+from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, BearerAuthBackend, RequireAuthMiddleware
 from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenVerifier
 from mcp.server.auth.routes import build_resource_metadata_url, create_auth_routes, create_protected_resource_routes
 from mcp.server.auth.settings import AuthSettings
@@ -74,6 +74,23 @@ async def main():
 LifespanResultT = TypeVar("LifespanResultT", default=Any)
 
 
+@contextmanager
+def _bind_request_auth_context(request_context: Any) -> Iterator[None]:
+    """Rebind auth context from the current transport request while handling a message."""
+    authenticated_user = None
+    scope = getattr(request_context, "scope", None)
+    if isinstance(scope, dict):
+        scope_user = scope.get("user")
+        if isinstance(scope_user, AuthenticatedUser):
+            authenticated_user = scope_user
+
+    token = auth_context_var.set(authenticated_user)
+    try:
+        yield
+    finally:
+        auth_context_var.reset(token)
+
+
 class NotificationOptions:
     def __init__(self, prompts_changed: bool = False, resources_changed: bool = False, tools_changed: bool = False):
         self.prompts_changed = prompts_changed
@@ -452,28 +469,32 @@ async def _handle_request(
                     close_sse_stream_cb = message.message_metadata.close_sse_stream
                     close_standalone_sse_stream_cb = message.message_metadata.close_standalone_sse_stream
 
-                client_capabilities = session.client_params.capabilities if session.client_params else None
-                task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
-                # Get task metadata from request params if present
-                task_metadata = None
-                if hasattr(req, "params") and req.params is not None:
-                    task_metadata = getattr(req.params, "task", None)
-                ctx = ServerRequestContext(
-                    request_id=message.request_id,
-                    meta=message.request_meta,
-                    session=session,
-                    lifespan_context=lifespan_context,
-                    experimental=Experimental(
-                        task_metadata=task_metadata,
-                        _client_capabilities=client_capabilities,
-                        _session=session,
-                        _task_support=task_support,
-                    ),
-                    request=request_data,
-                    close_sse_stream=close_sse_stream_cb,
-                    close_standalone_sse_stream=close_standalone_sse_stream_cb,
-                )
-                response = await handler(ctx, req.params)
+                # Stateful HTTP sessions process later requests on tasks that were
+                # created during session setup, so ContextVar snapshots can lag
+                # behind the current request unless we rebind them here.
+                with _bind_request_auth_context(request_data):
+                    client_capabilities = session.client_params.capabilities if session.client_params else None
+                    task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
+                    # Get task metadata from request params if present
+                    task_metadata = None
+                    if hasattr(req, "params") and req.params is not None:
+                        task_metadata = getattr(req.params, "task", None)
+                    ctx = ServerRequestContext(
+                        request_id=message.request_id,
+                        meta=message.request_meta,
+                        session=session,
+                        lifespan_context=lifespan_context,
+                        experimental=Experimental(
+                            task_metadata=task_metadata,
+                            _client_capabilities=client_capabilities,
+                            _session=session,
+                            _task_support=task_support,
+                        ),
+                        request=request_data,
+                        close_sse_stream=close_sse_stream_cb,
+                        close_standalone_sse_stream=close_standalone_sse_stream_cb,
+                    )
+                    response = await handler(ctx, req.params)
             except MCPError as err:
                 response = err.error
             except anyio.get_cancelled_exc_class():

tests/server/auth/middleware/test_auth_context_streamable_http.py

@@ -0,0 +1,104 @@
+"""Regression tests for auth context in StreamableHTTP servers."""
+
+import time
+from collections.abc import Generator
+
+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.client.session import ClientSession
+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.types import (
+    CallToolRequestParams,
+    CallToolResult,
+    ListToolsResult,
+    PaginatedRequestParams,
+    TextContent,
+    Tool,
+)
+from tests.test_helpers import run_uvicorn_in_thread
+
+
+class _EchoTokenVerifier:
+    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.fixture
+def stateful_auth_server() -> Generator[str, None, None]:
+    server = Server(
+        "auth-test-server",
+        on_call_tool=_handle_whoami,
+        on_list_tools=_handle_list_tools,
+    )
+    session_manager = StreamableHTTPSessionManager(app=server, stateless=False)
+    app = Starlette(
+        routes=[Mount("/mcp", app=session_manager.handle_request)],
+        middleware=[
+            Middleware(AuthenticationMiddleware, backend=BearerAuthBackend(_EchoTokenVerifier())),
+            Middleware(AuthContextMiddleware),
+        ],
+        lifespan=lambda app: session_manager.run(),
+    )
+
+    with run_uvicorn_in_thread(app, host="127.0.0.1", log_level="error") as base_url:
+        yield f"{base_url}/mcp"
+
+
+@pytest.mark.anyio
+async def test_get_access_token_reflects_current_request_in_stateful_session(stateful_auth_server: str) -> None:
+    auth = _MutableBearerAuth("token-A")
+    async with httpx.AsyncClient(
+        auth=auth,
+        timeout=httpx.Timeout(30, read=30),
+        follow_redirects=True,
+    ) as http_client:
+        async with streamable_http_client(stateful_auth_server, http_client=http_client) as (
+            read_stream,
+            write_stream,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+
+                first_response = await session.call_tool("whoami", {})
+                assert len(first_response.content) == 1
+                assert isinstance(first_response.content[0], TextContent)
+                assert first_response.content[0].text == "token-A"
+
+                auth.token = "token-B"
+
+                second_response = await session.call_tool("whoami", {})
+                assert len(second_response.content) == 1
+                assert isinstance(second_response.content[0], TextContent)
+                assert second_response.content[0].text == "token-B"