From 1177a11910fe420adccd50c887781e6b689f6cee Mon Sep 17 00:00:00 2001 From: David Soria Parra Date: Mon, 11 Nov 2024 21:04:44 +0000 Subject: [PATCH] Return a specific server session instance of request context We currently return a generic instance of RequestContext without a specialization on the Session type. This makes it impossible for servers to typesafe call `list_roots()` and other methods. We now return a specific instance of `RequestContext[ServerSession]` --- src/mcp/client/sse.py | 2 +- src/mcp/server/__init__.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py index 8a90e588b..abafacb96 100644 --- a/src/mcp/client/sse.py +++ b/src/mcp/client/sse.py @@ -84,7 +84,7 @@ async def sse_reader( case "message": try: - message = types.JSONRPCMessage.model_validate_json( # noqa: E501 + message = types.JSONRPCMessage.model_validate_json( # noqa: E501 sse.data ) logger.debug( diff --git a/src/mcp/server/__init__.py b/src/mcp/server/__init__.py index 3ce8b7bf1..aa7adad79 100644 --- a/src/mcp/server/__init__.py +++ b/src/mcp/server/__init__.py @@ -16,8 +16,8 @@ logger = logging.getLogger(__name__) -request_ctx: contextvars.ContextVar[RequestContext] = contextvars.ContextVar( - "request_ctx" +request_ctx: contextvars.ContextVar[RequestContext[ServerSession]] = ( + contextvars.ContextVar("request_ctx") ) @@ -115,7 +115,7 @@ def get_capabilities( ) @property - def request_context(self) -> RequestContext: + def request_context(self) -> RequestContext[ServerSession]: """If called outside of a request context, this will raise a LookupError.""" return request_ctx.get()