fix type hints for serversession

Author: asheshvidyut

examples/snippets/servers/elicitation.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field
 
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerSession
+from mcp.server.session import ServerTransportSession
 
 mcp = FastMCP(name="Elicitation Example")
 
@@ -17,7 +17,7 @@ class BookingPreferences(BaseModel):
 
 
 @mcp.tool()
-async def book_table(date: str, time: str, party_size: int, ctx: Context[ServerSession, None]) -> str:
+async def book_table(date: str, time: str, party_size: int, ctx: Context[ServerTransportSession, None]) -> str:
     """Book a table with date availability check."""
     # Check if date is available
     if date == "2024-12-25":

examples/snippets/servers/lifespan_example.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerSession
+from mcp.server.session import ServerTransportSession
 
 
 # Mock database class for example
@@ -51,7 +51,7 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
 
 # Access type-safe lifespan context in tools
 @mcp.tool()
-def query_db(ctx: Context[ServerSession, AppContext]) -> str:
+def query_db(ctx: Context[ServerTransportSession, AppContext]) -> str:
     """Tool that uses initialized resources."""
     db = ctx.request_context.lifespan_context.db
     return db.query()

examples/snippets/servers/notifications.py

@@ -1,11 +1,11 @@
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerSession
+from mcp.server.session import ServerTransportSession
 
 mcp = FastMCP(name="Notifications Example")
 
 
 @mcp.tool()
-async def process_data(data: str, ctx: Context[ServerSession, None]) -> str:
+async def process_data(data: str, ctx: Context[ServerTransportSession, None]) -> str:
     """Process data with logging."""
     # Different log levels
     await ctx.debug(f"Debug: Processing '{data}'")

examples/snippets/servers/tool_progress.py

@@ -1,11 +1,11 @@
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerSession
+from mcp.server.session import ServerTransportSession
 
 mcp = FastMCP(name="Progress Example")
 
 
 @mcp.tool()
-async def long_running_task(task_name: str, ctx: Context[ServerSession, None], steps: int = 5) -> str:
+async def long_running_task(task_name: str, ctx: Context[ServerTransportSession, None], steps: int = 5) -> str:
     """Execute a task with progress updates."""
     await ctx.info(f"Starting: {task_name}")
 

src/mcp/server/elicitation.py

@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
 
-from mcp.server.session import ServerSession
+from mcp.server.session import ServerTransportSession
 from mcp.types import RequestId
 
 ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
@@ -74,7 +74,7 @@ def _is_primitive_field(field_info: FieldInfo) -> bool:
 
 
 async def elicit_with_validation(
-    session: ServerSession,
+    session: ServerTransportSession,
     message: str,
     schema: type[ElicitSchemaModelT],
     related_request_id: RequestId | None = None,

src/mcp/server/fastmcp/server.py

@@ -54,7 +54,7 @@
 from mcp.server.lowlevel.server import LifespanResultT
 from mcp.server.lowlevel.server import Server as MCPServer
 from mcp.server.lowlevel.server import lifespan as default_lifespan
-from mcp.server.session import ServerSession, ServerSessionT
+from mcp.server.session import ServerSessionT, ServerTransportSession
 from mcp.server.sse import SseServerTransport
 from mcp.server.stdio import stdio_server
 from mcp.server.streamable_http import EventStore
@@ -315,7 +315,7 @@ async def list_tools(self) -> list[MCPTool]:
             for info in tools
         ]
 
-    def get_context(self) -> Context[ServerSession, LifespanResultT, Request]:
+    def get_context(self) -> Context[ServerTransportSession, LifespanResultT, Request]:
         """
         Returns a Context object. Note that the context will only be valid
         during a request; outside a request, most methods will error.

src/mcp/server/lowlevel/server.py

@@ -85,7 +85,7 @@ async def main():
 from mcp.server.lowlevel.func_inspection import create_call_wrapper
 from mcp.server.lowlevel.helper_types import ReadResourceContents
 from mcp.server.models import InitializationOptions
-from mcp.server.session import ServerSession
+from mcp.server.session import ServerSession, ServerTransportSession
 from mcp.shared.context import RequestContext
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
@@ -102,7 +102,7 @@ async def main():
 CombinationContent: TypeAlias = tuple[UnstructuredContent, StructuredContent]
 
 # This will be properly typed in each Server instance's context
-request_ctx: contextvars.ContextVar[RequestContext[ServerSession, Any, Any]] = contextvars.ContextVar("request_ctx")
+request_ctx: contextvars.ContextVar[RequestContext[ServerTransportSession, Any, Any]] = contextvars.ContextVar("request_ctx")
 
 
 class NotificationOptions:
@@ -231,7 +231,7 @@ def get_capabilities(
     @property
     def request_context(
         self,
-    ) -> RequestContext[ServerSession, LifespanResultT, RequestT]:
+    ) -> RequestContext[ServerTransportSession, LifespanResultT, RequestT]:
         """If called outside of a request context, this will raise a LookupError."""
         return request_ctx.get()
 

src/mcp/server/session.py

@@ -62,7 +62,7 @@ class InitializationState(Enum):
     Initialized = 3
 
 
-ServerSessionT = TypeVar("ServerSessionT", bound="ServerSession")
+ServerSessionT = TypeVar("ServerSessionT", bound="ServerTransportSession")
 
 ServerRequestResponder = (
     RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception

src/mcp/shared/context.py

@@ -4,10 +4,14 @@
 from typing_extensions import TypeVar
 
 from mcp.client.transport_session import ClientTransportSession
+from mcp.server.transport_session import ServerTransportSession
 from mcp.shared.session import BaseSession
 from mcp.types import RequestId, RequestParams
 
-SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any] | ClientTransportSession)
+SessionT = TypeVar("SessionT", 
+                   bound=BaseSession[Any, Any, Any, Any, Any] | 
+                   ClientTransportSession | 
+                   ServerTransportSession)
 LifespanContextT = TypeVar("LifespanContextT")
 RequestT = TypeVar("RequestT", default=Any)
 

tests/client/test_sampling_callback.py

@@ -1,6 +1,8 @@
 import pytest
 
 from mcp.client.session import ClientTransportSession
+from mcp.server.session import ServerSession
+from typing import cast
 from mcp.shared.context import RequestContext
 from mcp.shared.memory import (
     create_connected_server_and_client_session as create_session,
@@ -34,7 +36,8 @@ async def sampling_callback(
 
     @server.tool("test_sampling")
     async def test_sampling_tool(message: str):
-        value = await server.get_context().session.create_message(
+        session = cast(ServerSession, server.get_context().session)
+        value = await session.create_message(
             messages=[SamplingMessage(role="user", content=TextContent(type="text", text=message))],
             max_tokens=100,
         )