modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[v1.x] Bind transport sessions to the authenticated principal (#2719)

Browse source
This commit is contained in:
Max 2026-05-29 17:46:37 +01:00 • committed by GitHub
commit ce267b6fc5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 404 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

59
src/mcp/server/sse.py
View file

@ -52,6 +52,7 @@ from starlette.responses import Response
from starlette.types import Receive, Scope, Send
import mcp.types as types
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context
from mcp.server.transport_security import (
TransportSecurityMiddleware,
TransportSecuritySettings,
@ -75,6 +76,9 @@ class SseServerTransport:
_endpoint: str
_read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
# Identity of the credential that created each session; requests for a
# session must present the same credential.
_session_owners: dict[UUID, AuthorizationContext]
_security: TransportSecurityMiddleware
def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | None = None) -> None:
@ -115,6 +119,7 @@ class SseServerTransport:
self._endpoint = endpoint
self._read_stream_writers = {}
self._session_owners = {}
self._security = TransportSecurityMiddleware(security_settings)
logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
@ -142,6 +147,9 @@ class SseServerTransport:
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
session_id = uuid4()
user = scope.get("user")
if isinstance(user, AuthenticatedUser):
self._session_owners[session_id] = authorization_context(user)
self._read_stream_writers[session_id] = read_stream_writer
logger.debug(f"Created new session with ID: {session_id}")
@ -177,26 +185,34 @@ class SseServerTransport:
}
)
async with anyio.create_task_group() as tg:
try:
async with anyio.create_task_group() as tg:
async def response_wrapper(scope: Scope, receive: Receive, send: Send):
"""
The EventSourceResponse returning signals a client close / disconnect.
In this case we close our side of the streams to signal the client that
the connection has been closed.
"""
await EventSourceResponse(content=sse_stream_reader, data_sender_callable=sse_writer)(
scope, receive, send
)
await read_stream_writer.aclose()
await write_stream_reader.aclose()
logging.debug(f"Client session disconnected {session_id}")
async def response_wrapper(scope: Scope, receive: Receive, send: Send):
"""
The EventSourceResponse returning signals a client close / disconnect.
In this case we close our side of the streams to signal the client that
the connection has been closed.
"""
await EventSourceResponse(content=sse_stream_reader, data_sender_callable=sse_writer)(
scope, receive, send
)
await read_stream_writer.aclose()
await write_stream_reader.aclose()
await sse_stream_reader.aclose()
logging.debug(f"Client session disconnected {session_id}")
logger.debug("Starting SSE response task")
tg.start_soon(response_wrapper, scope, receive, send)
logger.debug("Starting SSE response task")
tg.start_soon(response_wrapper, scope, receive, send)
logger.debug("Yielding read and write streams")
yield (read_stream, write_stream)
logger.debug("Yielding read and write streams")
yield (read_stream, write_stream)
finally:
# The connection is gone: stop routing messages to this session
# and drop its entries so they do not accumulate for the lifetime
# of the transport.
self._read_stream_writers.pop(session_id, None)
self._session_owners.pop(session_id, None)
async def handle_post_message(self, scope: Scope, receive: Receive, send: Send) -> None: # pragma: no cover
logger.debug("Handling POST message")
@ -227,6 +243,15 @@ class SseServerTransport:
response = Response("Could not find session", status_code=404)
return await response(scope, receive, send)
user = scope.get("user")
requestor = authorization_context(user) if isinstance(user, AuthenticatedUser) else None
if requestor != self._session_owners.get(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 message for session %s: credential does not match", session_id)
response = Response("Could not find session", status_code=404)
return await response(scope, receive, send)
body = await request.body()
logger.debug(f"Received JSON: {body}")