From 144db31f492677e966582a33a0cbfd1dcd50a407 Mon Sep 17 00:00:00 2001 From: Abdullah Elkady Date: Tue, 18 Feb 2025 19:45:26 -0500 Subject: [PATCH 1/3] enable http headers in sse in the context --- src/mcp/server/lowlevel/server.py | 6 ++++++ src/mcp/server/sse.py | 3 ++- src/mcp/shared/context.py | 3 ++- src/mcp/types.py | 1 + 4 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index c0008b329..1a5c666b9 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -526,6 +526,11 @@ async def _handle_request( logger.debug(f"Dispatching request of type {type(req).__name__}") token = None + headers = {} + try: + headers = message.request.root.headers # type: ignore + except Exception: + pass try: # Set our global state that can be retrieved via # app.get_request_context() @@ -535,6 +540,7 @@ async def _handle_request( message.request_meta, session, lifespan_context, + headers, ) ) response = await handler(req) diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py index 0127753d0..2ce4be779 100644 --- a/src/mcp/server/sse.py +++ b/src/mcp/server/sse.py @@ -160,7 +160,8 @@ async def handle_post_message( logger.debug(f"Received JSON: {json}") try: - message = types.JSONRPCMessage.model_validate(json) + message_with_headers = {**json, "headers": dict(request.headers)} + message = types.JSONRPCMessage.model_validate(message_with_headers) logger.debug(f"Validated client message: {message}") except ValidationError as err: logger.error(f"Failed to parse message: {err}") diff --git a/src/mcp/shared/context.py b/src/mcp/shared/context.py index a45fdacd4..6140d9eca 100644 --- a/src/mcp/shared/context.py +++ b/src/mcp/shared/context.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Generic, TypeVar from mcp.shared.session import BaseSession @@ -14,3 +14,4 @@ class RequestContext(Generic[SessionT, LifespanContextT]): meta: RequestParams.Meta | None session: SessionT lifespan_context: LifespanContextT + headers: dict[str, str] = field(default_factory=dict) diff --git a/src/mcp/types.py b/src/mcp/types.py index 7d867bd3b..362779e58 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -118,6 +118,7 @@ class JSONRPCRequest(Request): jsonrpc: Literal["2.0"] id: RequestId params: dict[str, Any] | None = None + headers: dict[str, str] | None = None class JSONRPCNotification(Notification): From e402103483fcfa403b413e8c4d319408d400fa26 Mon Sep 17 00:00:00 2001 From: Abdullah Elkady Date: Tue, 18 Feb 2025 19:55:29 -0500 Subject: [PATCH 2/3] add todo comment --- src/mcp/server/lowlevel/server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 1a5c666b9..193c204c3 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -528,6 +528,7 @@ async def _handle_request( token = None headers = {} try: + # TODO: This try/catch and ignoring the type is wrong. headers = message.request.root.headers # type: ignore except Exception: pass From b6e45768f2a9764a4f91e4c1d19d7cb0272daeca Mon Sep 17 00:00:00 2001 From: Abdullah Elkady Date: Wed, 19 Feb 2025 10:25:23 -0500 Subject: [PATCH 3/3] type instead of try/except --- src/mcp/server/lowlevel/server.py | 10 ++-------- src/mcp/types.py | 2 +- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 193c204c3..e02c9444d 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -514,7 +514,7 @@ async def _handle_message( async def _handle_request( self, - message: RequestResponder, + message: RequestResponder[types.ClientRequest, types.ServerResult], req: Any, session: ServerSession, lifespan_context: LifespanResultT, @@ -526,12 +526,6 @@ async def _handle_request( logger.debug(f"Dispatching request of type {type(req).__name__}") token = None - headers = {} - try: - # TODO: This try/catch and ignoring the type is wrong. - headers = message.request.root.headers # type: ignore - except Exception: - pass try: # Set our global state that can be retrieved via # app.get_request_context() @@ -541,7 +535,7 @@ async def _handle_request( message.request_meta, session, lifespan_context, - headers, + message.request.root.headers or {}, ) ) response = await handler(req) diff --git a/src/mcp/types.py b/src/mcp/types.py index 362779e58..41f737f63 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -74,6 +74,7 @@ class Request(BaseModel, Generic[RequestParamsT, MethodT]): method: MethodT params: RequestParamsT + headers: dict[str, str] | None = None model_config = ConfigDict(extra="allow") @@ -118,7 +119,6 @@ class JSONRPCRequest(Request): jsonrpc: Literal["2.0"] id: RequestId params: dict[str, Any] | None = None - headers: dict[str, str] | None = None class JSONRPCNotification(Notification):