Merge branch 'main' into close-resources

Author: Kludex

examples/clients/simple-chatbot/mcp_simple_chatbot/main.py

@@ -23,7 +23,7 @@ class Configuration:
     def __init__(self) -> None:
         """Initialize configuration with environment variables."""
         self.load_env()
-        self.api_key = os.getenv("GROQ_API_KEY")
+        self.api_key = os.getenv("LLM_API_KEY")
 
     @staticmethod
     def load_env() -> None:

src/mcp/server/fastmcp/server.py

@@ -37,7 +37,7 @@
 from mcp.server.session import ServerSession
 from mcp.server.sse import SseServerTransport
 from mcp.server.stdio import stdio_server
-from mcp.shared.context import RequestContext
+from mcp.shared.context import LifespanContextT, RequestContext
 from mcp.types import (
     AnyFunction,
     EmbeddedResource,
@@ -564,7 +564,7 @@ def _convert_to_content(
     return [TextContent(type="text", text=result)]
 
 
-class Context(BaseModel):
+class Context(BaseModel, Generic[LifespanContextT]):
     """Context object providing access to MCP capabilities.
 
     This provides a cleaner interface to MCP's RequestContext functionality.
@@ -598,13 +598,13 @@ def my_tool(x: int, ctx: Context) -> str:
     The context is optional - tools that don't need it can omit the parameter.
     """
 
-    _request_context: RequestContext[ServerSession, Any] | None
+    _request_context: RequestContext[ServerSession, LifespanContextT] | None
     _fastmcp: FastMCP | None
 
     def __init__(
         self,
         *,
-        request_context: RequestContext | None = None,
+        request_context: RequestContext[ServerSession, LifespanContextT] | None = None,
         fastmcp: FastMCP | None = None,
         **kwargs: Any,
     ):
@@ -620,7 +620,7 @@ def fastmcp(self) -> FastMCP:
         return self._fastmcp
 
     @property
-    def request_context(self) -> RequestContext:
+    def request_context(self) -> RequestContext[ServerSession, LifespanContextT]:
         """Access to the underlying request context."""
         if self._request_context is None:
             raise ValueError("Context is not available outside of a request")

src/mcp/shared/context.py

@@ -1,11 +1,13 @@
 from dataclasses import dataclass
-from typing import Generic, TypeVar
+from typing import Any, Generic
+
+from typing_extensions import TypeVar
 
 from mcp.shared.session import BaseSession
 from mcp.types import RequestId, RequestParams
 
-SessionT = TypeVar("SessionT", bound=BaseSession)
-LifespanContextT = TypeVar("LifespanContextT")
+SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
+LifespanContextT = TypeVar("LifespanContextT", default=None)
 
 
 @dataclass

src/mcp/shared/session.py

@@ -8,6 +8,7 @@
 import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import BaseModel
+from typing_extensions import Self
 
 from mcp.shared.exceptions import McpError
 from mcp.types import (
@@ -60,7 +61,7 @@ def __init__(
         request_id: RequestId,
         request_meta: RequestParams.Meta | None,
         request: ReceiveRequestT,
-        session: "BaseSession",
+        session: "BaseSession[SendRequestT, SendNotificationT, SendResultT, ReceiveRequestT, ReceiveNotificationT]",
         on_complete: Callable[["RequestResponder[ReceiveRequestT, SendResultT]"], Any],
     ) -> None:
         self.request_id = request_id
@@ -134,7 +135,6 @@ def cancelled(self) -> bool:
 
 
 class BaseSession(
-    AbstractAsyncContextManager,
     Generic[
         SendRequestT,
         SendNotificationT,
@@ -183,7 +183,7 @@ def __init__(
             ]()
         )
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Self:
         self._task_group = anyio.create_task_group()
         await self._task_group.__aenter__()
         self._task_group.start_soon(self._receive_loop)