Bind transport sessions to the authenticated principal (#2718)

src/mcp/server/streamable_http_manager.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import contextlib
 import logging
 from collections.abc import AsyncIterator
-from http import HTTPStatus
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
@@ -15,6 +14,7 @@ from starlette.requests import Request
 from starlette.responses import Response
 from starlette.types import Receive, Scope, Send
 
+from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context
 from mcp.server.streamable_http import (
     MCP_SESSION_ID_HEADER,
     EventStore,
@@ -89,6 +89,9 @@ class StreamableHTTPSessionManager:
         # Session tracking (only used if not stateless)
         self._session_creation_lock = anyio.Lock()
         self._server_instances: dict[str, StreamableHTTPServerTransport] = {}
+        # Identity of the credential that created each session; requests for a
+        # session must present the same credential.
+        self._session_owners: dict[str, AuthorizationContext] = {}
 
         # The task group will be set during lifespan
         self._task_group = None
@@ -135,6 +138,7 @@ class StreamableHTTPSessionManager:
                 self._task_group = None
                 # Clear any remaining server instances
                 self._server_instances.clear()
+                self._session_owners.clear()
 
     async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Process ASGI request with proper session handling and transport setup.
@@ -192,9 +196,29 @@ class StreamableHTTPSessionManager:
         request = Request(scope, receive)
         request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
 
+        user = scope.get("user")
+        requestor = authorization_context(user) if isinstance(user, AuthenticatedUser) else None
+
         # Existing session case
         if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances:
             transport = self._server_instances[request_mcp_session_id]
+            if requestor != self._session_owners.get(request_mcp_session_id):
+                # A session can only be used with the credential that created
+                # it. Respond exactly as if the session did not exist.
+                logger.warning(
+                    "Rejecting request for session %s: credential does not match the one that created the session",
+                    request_mcp_session_id[:64],
+                )
+                body = JSONRPCError(
+                    jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
+                )
+                response = Response(
+                    body.model_dump_json(by_alias=True, exclude_unset=True),
+                    status_code=404,
+                    media_type="application/json",
+                )
+                await response(scope, receive, send)
+                return
             logger.debug("Session already exists, handling request directly")
             # Push back idle deadline on activity
             if transport.idle_scope is not None and self.session_idle_timeout is not None:
@@ -216,6 +240,8 @@ class StreamableHTTPSessionManager:
                 )
 
                 assert http_transport.mcp_session_id is not None
+                if requestor is not None:
+                    self._session_owners[http_transport.mcp_session_id] = requestor
                 self._server_instances[http_transport.mcp_session_id] = http_transport
                 logger.info(f"Created new transport with session ID: {new_session_id}")
 
@@ -246,6 +272,7 @@ class StreamableHTTPSessionManager:
                                 assert http_transport.mcp_session_id is not None
                                 logger.info(f"Session {http_transport.mcp_session_id} idle timeout")
                                 self._server_instances.pop(http_transport.mcp_session_id, None)
+                                self._session_owners.pop(http_transport.mcp_session_id, None)
                                 await http_transport.terminate()
                         except Exception:
                             logger.exception(f"Session {http_transport.mcp_session_id} crashed")
@@ -260,6 +287,7 @@ class StreamableHTTPSessionManager:
                                     f"{http_transport.mcp_session_id} from active instances."
                                 )
                                 del self._server_instances[http_transport.mcp_session_id]
+                                self._session_owners.pop(http_transport.mcp_session_id, None)
 
                 # Assert task group is not None for type checking
                 assert self._task_group is not None
@@ -273,15 +301,11 @@ class StreamableHTTPSessionManager:
             # TODO: Align error code once spec clarifies
             # See: https://github.com/modelcontextprotocol/python-sdk/issues/1821
             logger.info(f"Rejected request with unknown or expired session ID: {request_mcp_session_id[:64]}")
-            error_response = JSONRPCError(
-                jsonrpc="2.0",
-                id=None,
-                error=ErrorData(code=INVALID_REQUEST, message="Session not found"),
+            body = JSONRPCError(
+                jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
             )
             response = Response(
-                content=error_response.model_dump_json(by_alias=True, exclude_unset=True),
-                status_code=HTTPStatus.NOT_FOUND,
-                media_type="application/json",
+                body.model_dump_json(by_alias=True, exclude_unset=True), status_code=404, media_type="application/json"
             )
             await response(scope, receive, send)