fix type hints without cast

Author: asheshvidyut

src/mcp/client/session_group.py

@@ -96,10 +96,10 @@ class _ComponentNames(BaseModel):
     _tools: dict[str, types.Tool]
 
     # Client-server connection management.
-    _sessions: dict["mcp.ClientTransportSession", _ComponentNames]
-    _tool_to_session: dict[str, "mcp.ClientTransportSession"]
+    _sessions: dict[mcp.ClientTransportSession, _ComponentNames]
+    _tool_to_session: dict[str, mcp.ClientTransportSession]
     _exit_stack: contextlib.AsyncExitStack
-    _session_exit_stacks: dict["mcp.ClientTransportSession", contextlib.AsyncExitStack]
+    _session_exit_stacks: dict[mcp.ClientTransportSession, contextlib.AsyncExitStack]
 
     # Optional fn consuming (component_name, serverInfo) for custom names.
     # This is provide a means to mitigate naming conflicts across servers.
@@ -153,7 +153,7 @@ async def __aexit__(
                 tg.start_soon(exit_stack.aclose)
 
     @property
-    def sessions(self) -> list["mcp.ClientTransportSession"]:
+    def sessions(self) -> list[mcp.ClientTransportSession]:
         """Returns the list of sessions being managed."""
         return list(self._sessions.keys())  # pragma: no cover
 
@@ -178,7 +178,7 @@ async def call_tool(self, name: str, args: dict[str, Any]) -> types.CallToolResu
         session_tool_name = self.tools[name].name
         return await session.call_tool(session_tool_name, args)
 
-    async def disconnect_from_server(self, session: "mcp.ClientTransportSession") -> None:
+    async def disconnect_from_server(self, session: mcp.ClientTransportSession) -> None:
         """Disconnects from a single MCP server."""
 
         session_known_for_components = session in self._sessions
@@ -216,23 +216,23 @@ async def disconnect_from_server(self, session: "mcp.ClientTransportSession") ->
             await session_stack_to_close.aclose()  # pragma: no cover
 
     async def connect_with_session(
-        self, server_info: types.Implementation, session: "mcp.ClientTransportSession"
-    ) -> "mcp.ClientTransportSession":
+        self, server_info: types.Implementation, session: mcp.ClientTransportSession
+    ) -> mcp.ClientTransportSession:
         """Connects to a single MCP server."""
         await self._aggregate_components(server_info, session)
         return session
 
     async def connect_to_server(
         self,
         server_params: ServerParameters,
-    ) -> "mcp.ClientTransportSession":
+    ) -> mcp.ClientTransportSession:
         """Connects to a single MCP server."""
         server_info, session = await self._establish_session(server_params)
         return await self.connect_with_session(server_info, session)
 
     async def _establish_session(
         self, server_params: ServerParameters
-    ) -> tuple[types.Implementation, "mcp.ClientTransportSession"]:
+    ) -> tuple[types.Implementation, mcp.ClientTransportSession]:
         """Establish a client session to an MCP server."""
 
         session_stack = contextlib.AsyncExitStack()
@@ -277,7 +277,7 @@ async def _establish_session(
             raise
 
     async def _aggregate_components(
-        self, server_info: types.Implementation, session: "mcp.ClientTransportSession"
+        self, server_info: types.Implementation, session: mcp.ClientTransportSession
     ) -> None:
         """Aggregates prompts, resources, and tools from a given session."""
 

src/mcp/shared/memory.py

@@ -14,6 +14,7 @@
 
 import mcp.types as types
 from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+from mcp.client.session import ClientTransportSession
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
 from mcp.shared.message import SessionMessage
@@ -57,7 +58,7 @@ async def create_connected_server_and_client_session(
     client_info: types.Implementation | None = None,
     raise_exceptions: bool = False,
     elicitation_callback: ElicitationFnT | None = None,
-) -> AsyncGenerator[ClientSession, None]:
+) -> AsyncGenerator[ClientTransportSession, None]:
     """Creates a ClientSession that is connected to a running MCP server."""
 
     # TODO(Marcelo): we should have a proper `Client` that can use this "in-memory transport",

tests/client/test_sampling_callback.py

@@ -1,5 +1,3 @@
-from typing import cast
-
 import pytest
 
 from mcp.client.session import ClientTransportSession
@@ -37,7 +35,8 @@ async def sampling_callback(
 
     @server.tool("test_sampling")
     async def test_sampling_tool(message: str):
-        session = cast(ServerSession, server.get_context().session)
+        session = server.get_context().session
+        assert isinstance(session, ServerSession)
         value = await session.create_message(
             messages=[SamplingMessage(role="user", content=TextContent(type="text", text=message))],
             max_tokens=100,

tests/server/test_cancel_handling.py

@@ -9,6 +9,7 @@
 from mcp.server.lowlevel.server import Server
 from mcp.shared.exceptions import McpError
 from mcp.shared.memory import create_connected_server_and_client_session
+from mcp.client.session import ClientSession
 from mcp.types import (
     CallToolRequest,
     CallToolRequestParams,
@@ -56,6 +57,7 @@ async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[
 
     async with create_connected_server_and_client_session(server) as client:
         # First request (will be cancelled)
+        assert isinstance(client, ClientSession)
         async def first_request():
             try:
                 await client.send_request(

tests/shared/test_memory.py

@@ -2,7 +2,7 @@
 from pydantic import AnyUrl
 from typing_extensions import AsyncGenerator
 
-from mcp.client.session import ClientSession
+from mcp.client.session import ClientSession, ClientTransportSession
 from mcp.server import Server
 from mcp.shared.memory import create_connected_server_and_client_session
 from mcp.types import EmptyResult, Resource
@@ -28,7 +28,7 @@ async def handle_list_resources():  # pragma: no cover
 @pytest.fixture
 async def client_connected_to_server(
     mcp_server: Server,
-) -> AsyncGenerator[ClientSession, None]:
+) -> AsyncGenerator[ClientTransportSession, None]:
     async with create_connected_server_and_client_session(mcp_server) as client_session:
         yield client_session
 

tests/shared/test_progress_notifications.py

@@ -370,6 +370,7 @@ async def handle_list_tools() -> list[types.Tool]:
     with patch("mcp.shared.session.logging.error", side_effect=mock_log_error):
         async with create_connected_server_and_client_session(server) as client_session:
             # Send a request with a failing progress callback
+            assert isinstance(client_session, ClientSession)
             result = await client_session.send_request(
                 types.ClientRequest(
                     types.CallToolRequest(

tests/shared/test_session.py

@@ -5,7 +5,7 @@
 import pytest
 
 import mcp.types as types
-from mcp.client.session import ClientSession
+from mcp.client.session import ClientSession, ClientTransportSession
 from mcp.server.lowlevel.server import Server
 from mcp.shared.exceptions import McpError
 from mcp.shared.memory import create_client_server_memory_streams, create_connected_server_and_client_session
@@ -27,19 +27,20 @@ def mcp_server() -> Server:
 @pytest.fixture
 async def client_connected_to_server(
     mcp_server: Server,
-) -> AsyncGenerator[ClientSession, None]:
+) -> AsyncGenerator[ClientTransportSession, None]:
     async with create_connected_server_and_client_session(mcp_server) as client_session:
         yield client_session
 
 
 @pytest.mark.anyio
 async def test_in_flight_requests_cleared_after_completion(
-    client_connected_to_server: ClientSession,
+    client_connected_to_server: ClientTransportSession,
 ):
     """Verify that _in_flight is empty after all requests complete."""
     # Send a request and wait for response
     response = await client_connected_to_server.send_ping()
     assert isinstance(response, EmptyResult)
+    assert isinstance(client_connected_to_server, ClientSession)
 
     # Verify _in_flight is empty
     assert len(client_connected_to_server._in_flight) == 0
@@ -101,6 +102,7 @@ async def make_request(client_session: ClientSession):
 
     async with create_connected_server_and_client_session(make_server()) as client_session:
         async with anyio.create_task_group() as tg:
+            assert isinstance(client_session, ClientSession)
             tg.start_soon(make_request, client_session)
 
             # Wait for the request to be in-flight

tests/shared/test_sse.py

@@ -17,7 +17,7 @@
 from starlette.routing import Mount, Route
 
 import mcp.types as types
-from mcp.client.session import ClientSession
+from mcp.client.session import ClientSession, ClientTransportSession
 from mcp.client.sse import sse_client
 from mcp.server import Server
 from mcp.server.sse import SseServerTransport
@@ -185,7 +185,7 @@ async def test_sse_client_basic_connection(server: None, server_url: str) -> Non
 
 
 @pytest.fixture
-async def initialized_sse_client_session(server: None, server_url: str) -> AsyncGenerator[ClientSession, None]:
+async def initialized_sse_client_session(server: None, server_url: str) -> AsyncGenerator[ClientTransportSession, None]:
     async with sse_client(server_url + "/sse", sse_read_timeout=0.5) as streams:
         async with ClientSession(*streams) as session:
             await session.initialize()

tests/shared/test_streamable_http.py

@@ -8,7 +8,7 @@
 import multiprocessing
 import socket
 from collections.abc import Generator
-from typing import Any, cast
+from typing import Any
 
 import anyio
 import httpx
@@ -199,7 +199,8 @@ async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]
 
             elif name == "test_sampling_tool":
                 # Test sampling by requesting the client to sample a message
-                session = cast(ServerSession, ctx.session)
+                session = ctx.session
+                assert isinstance(session, ServerSession)
                 sampling_result = await session.create_message(
                     messages=[
                         types.SamplingMessage(

tests/shared/test_ws.py

@@ -12,7 +12,7 @@
 from starlette.routing import WebSocketRoute
 from starlette.websockets import WebSocket
 
-from mcp.client.session import ClientSession
+from mcp.client.session import ClientSession, ClientTransportSession
 from mcp.client.websocket import websocket_client
 from mcp.server import Server
 from mcp.server.websocket import websocket_server
@@ -125,7 +125,7 @@ def server(server_port: int) -> Generator[None, None, None]:
 
 
 @pytest.fixture()
-async def initialized_ws_client_session(server: None, server_url: str) -> AsyncGenerator[ClientSession, None]:
+async def initialized_ws_client_session(server: None, server_url: str) -> AsyncGenerator[ClientTransportSession, None]:
     """Create and initialize a WebSocket client session"""
     async with websocket_client(server_url + "/ws") as streams:
         async with ClientSession(*streams) as session: