add Generic types to be used with the new BaseClientSession class

Author: sreenithi

src/mcp/client/base_client_session.py

@@ -1,11 +1,12 @@
 from abc import abstractmethod
-from typing import Any
+from typing import Any, TypeVar
 
 from mcp import types
 from mcp.shared.session import CommonBaseSession, ProgressFnT
 from mcp.types._types import RequestParamsMeta
 
-# from mcp.shared.session import CommonBaseSession
+ClientSessionT_contra = TypeVar("ClientSessionT_contra", bound="BaseClientSession", contravariant=True)
+
 
 class BaseClientSession(
     CommonBaseSession[

src/mcp/client/client.py

@@ -76,10 +76,10 @@ async def main():
     read_timeout_seconds: float | None = None
     """Timeout for read operations."""
 
-    sampling_callback: SamplingFnT | None = None
+    sampling_callback: SamplingFnT[ClientSession] | None = None
     """Callback for handling sampling requests."""
 
-    list_roots_callback: ListRootsFnT | None = None
+    list_roots_callback: ListRootsFnT[ClientSession] | None = None
     """Callback for handling list roots requests."""
 
     logging_callback: LoggingFnT | None = None
@@ -92,7 +92,7 @@ async def main():
     client_info: Implementation | None = None
     """Client implementation info to send to server."""
 
-    elicitation_callback: ElicitationFnT | None = None
+    elicitation_callback: ElicitationFnT[ClientSession] | None = None
     """Callback for handling elicitation requests."""
 
     _session: ClientSession | None = field(init=False, default=None)

src/mcp/client/context.py

@@ -1,9 +1,11 @@
 """Request context for MCP client handlers."""
 
-from mcp.client.session import ClientSession
+from mcp.client import BaseClientSession
+
+# from mcp.client.session import ClientSession
 from mcp.shared._context import RequestContext
 
-ClientRequestContext = RequestContext[ClientSession]
+ClientRequestContext = RequestContext[BaseClientSession]
 """Context for handling incoming requests in a client session.
 
 This context is passed to client-side callbacks (sampling, elicitation, list_roots) when the server sends requests

src/mcp/client/experimental/task_handlers.py

@@ -14,71 +14,69 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Protocol
+from typing import Generic, Protocol
 
 from pydantic import TypeAdapter
 
 from mcp import types
+from mcp.client.base_client_session import BaseClientSession, ClientSessionT_contra
 from mcp.shared._context import RequestContext
 from mcp.shared.session import RequestResponder
 
-if TYPE_CHECKING:
-    from mcp.client.session import ClientSession
 
-
-class GetTaskHandlerFnT(Protocol):
+class GetTaskHandlerFnT(Protocol[ClientSessionT_contra]):
     """Handler for tasks/get requests from server.
 
     WARNING: This is experimental and may change without notice.
     """
 
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.GetTaskRequestParams,
     ) -> types.GetTaskResult | types.ErrorData: ...  # pragma: no branch
 
 
-class GetTaskResultHandlerFnT(Protocol):
+class GetTaskResultHandlerFnT(Protocol[ClientSessionT_contra]):
     """Handler for tasks/result requests from server.
 
     WARNING: This is experimental and may change without notice.
     """
 
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.GetTaskPayloadRequestParams,
     ) -> types.GetTaskPayloadResult | types.ErrorData: ...  # pragma: no branch
 
 
-class ListTasksHandlerFnT(Protocol):
+class ListTasksHandlerFnT(Protocol[ClientSessionT_contra]):
     """Handler for tasks/list requests from server.
 
     WARNING: This is experimental and may change without notice.
     """
 
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.PaginatedRequestParams | None,
     ) -> types.ListTasksResult | types.ErrorData: ...  # pragma: no branch
 
 
-class CancelTaskHandlerFnT(Protocol):
+class CancelTaskHandlerFnT(Protocol[ClientSessionT_contra]):
     """Handler for tasks/cancel requests from server.
 
     WARNING: This is experimental and may change without notice.
     """
 
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.CancelTaskRequestParams,
     ) -> types.CancelTaskResult | types.ErrorData: ...  # pragma: no branch
 
 
-class TaskAugmentedSamplingFnT(Protocol):
+class TaskAugmentedSamplingFnT(Protocol[ClientSessionT_contra]):
     """Handler for task-augmented sampling/createMessage requests from server.
 
     When server sends a CreateMessageRequest with task field, this callback
@@ -90,13 +88,13 @@ class TaskAugmentedSamplingFnT(Protocol):
 
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.CreateMessageRequestParams,
         task_metadata: types.TaskMetadata,
     ) -> types.CreateTaskResult | types.ErrorData: ...  # pragma: no branch
 
 
-class TaskAugmentedElicitationFnT(Protocol):
+class TaskAugmentedElicitationFnT(Protocol[ClientSessionT_contra]):
     """Handler for task-augmented elicitation/create requests from server.
 
     When server sends an ElicitRequest with task field, this callback
@@ -108,14 +106,14 @@ class TaskAugmentedElicitationFnT(Protocol):
 
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.ElicitRequestParams,
         task_metadata: types.TaskMetadata,
     ) -> types.CreateTaskResult | types.ErrorData: ...  # pragma: no branch
 
 
 async def default_get_task_handler(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.GetTaskRequestParams,
 ) -> types.GetTaskResult | types.ErrorData:
     return types.ErrorData(
@@ -125,7 +123,7 @@ async def default_get_task_handler(
 
 
 async def default_get_task_result_handler(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.GetTaskPayloadRequestParams,
 ) -> types.GetTaskPayloadResult | types.ErrorData:
     return types.ErrorData(
@@ -135,7 +133,7 @@ async def default_get_task_result_handler(
 
 
 async def default_list_tasks_handler(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.PaginatedRequestParams | None,
 ) -> types.ListTasksResult | types.ErrorData:
     return types.ErrorData(
@@ -145,7 +143,7 @@ async def default_list_tasks_handler(
 
 
 async def default_cancel_task_handler(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.CancelTaskRequestParams,
 ) -> types.CancelTaskResult | types.ErrorData:
     return types.ErrorData(
@@ -155,7 +153,7 @@ async def default_cancel_task_handler(
 
 
 async def default_task_augmented_sampling(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.CreateMessageRequestParams,
     task_metadata: types.TaskMetadata,
 ) -> types.CreateTaskResult | types.ErrorData:
@@ -166,7 +164,7 @@ async def default_task_augmented_sampling(
 
 
 async def default_task_augmented_elicitation(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.ElicitRequestParams,
     task_metadata: types.TaskMetadata,
 ) -> types.CreateTaskResult | types.ErrorData:
@@ -177,7 +175,7 @@ async def default_task_augmented_elicitation(
 
 
 @dataclass
-class ExperimentalTaskHandlers:
+class ExperimentalTaskHandlers(Generic[ClientSessionT_contra]):
     """Container for experimental task handlers.
 
     Groups all task-related handlers that handle server -> client requests.
@@ -195,14 +193,16 @@ class ExperimentalTaskHandlers:
     """
 
     # Pure task request handlers
-    get_task: GetTaskHandlerFnT = field(default=default_get_task_handler)
-    get_task_result: GetTaskResultHandlerFnT = field(default=default_get_task_result_handler)
-    list_tasks: ListTasksHandlerFnT = field(default=default_list_tasks_handler)
-    cancel_task: CancelTaskHandlerFnT = field(default=default_cancel_task_handler)
+    get_task: GetTaskHandlerFnT[ClientSessionT_contra] = field(default=default_get_task_handler)
+    get_task_result: GetTaskResultHandlerFnT[ClientSessionT_contra] = field(default=default_get_task_result_handler)
+    list_tasks: ListTasksHandlerFnT[ClientSessionT_contra] = field(default=default_list_tasks_handler)
+    cancel_task: CancelTaskHandlerFnT[ClientSessionT_contra] = field(default=default_cancel_task_handler)
 
     # Task-augmented request handlers
-    augmented_sampling: TaskAugmentedSamplingFnT = field(default=default_task_augmented_sampling)
-    augmented_elicitation: TaskAugmentedElicitationFnT = field(default=default_task_augmented_elicitation)
+    augmented_sampling: TaskAugmentedSamplingFnT[ClientSessionT_contra] = field(default=default_task_augmented_sampling)
+    augmented_elicitation: TaskAugmentedElicitationFnT[ClientSessionT_contra] = field(
+        default=default_task_augmented_elicitation
+    )
 
     def build_capability(self) -> types.ClientTasksCapability | None:
         """Build ClientTasksCapability from the configured handlers.
@@ -250,7 +250,7 @@ def handles_request(request: types.ServerRequest) -> bool:
 
     async def handle_request(
         self,
-        ctx: RequestContext[ClientSession],
+        ctx: RequestContext[ClientSessionT_contra],
         responder: RequestResponder[types.ServerRequest, types.ClientResult],
     ) -> None:
         """Handle a task-related request from the server.

src/mcp/client/session.py

@@ -8,7 +8,7 @@
 from pydantic import TypeAdapter
 
 from mcp import types
-from mcp.client.base_client_session import BaseClientSession
+from mcp.client.base_client_session import BaseClientSession, ClientSessionT_contra
 from mcp.client.experimental import ExperimentalClientFeatures
 from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
 from mcp.shared._context import RequestContext
@@ -22,25 +22,25 @@
 logger = logging.getLogger("client")
 
 
-class SamplingFnT(Protocol):
+class SamplingFnT(Protocol[ClientSessionT_contra]):
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.CreateMessageRequestParams,
     ) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: ...  # pragma: no branch
 
 
-class ElicitationFnT(Protocol):
+class ElicitationFnT(Protocol[ClientSessionT_contra]):
     async def __call__(
         self,
-        context: RequestContext[ClientSession],
+        context: RequestContext[ClientSessionT_contra],
         params: types.ElicitRequestParams,
     ) -> types.ElicitResult | types.ErrorData: ...  # pragma: no branch
 
 
-class ListRootsFnT(Protocol):
+class ListRootsFnT(Protocol[ClientSessionT_contra]):
     async def __call__(
-        self, context: RequestContext[ClientSession]
+        self, context: RequestContext[ClientSessionT_contra]
     ) -> types.ListRootsResult | types.ErrorData: ...  # pragma: no branch
 
 
@@ -62,7 +62,7 @@ async def _default_message_handler(
 
 
 async def _default_sampling_callback(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.CreateMessageRequestParams,
 ) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData:
     return types.ErrorData(
@@ -72,7 +72,7 @@ async def _default_sampling_callback(
 
 
 async def _default_elicitation_callback(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
     params: types.ElicitRequestParams,
 ) -> types.ElicitResult | types.ErrorData:
     return types.ErrorData(  # pragma: no cover
@@ -82,7 +82,7 @@ async def _default_elicitation_callback(
 
 
 async def _default_list_roots_callback(
-    context: RequestContext[ClientSession],
+    context: RequestContext[BaseClientSession],
 ) -> types.ListRootsResult | types.ErrorData:
     return types.ErrorData(
         code=types.INVALID_REQUEST,
@@ -114,15 +114,15 @@ def __init__(
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
         write_stream: MemoryObjectSendStream[SessionMessage],
         read_timeout_seconds: float | None = None,
-        sampling_callback: SamplingFnT | None = None,
-        elicitation_callback: ElicitationFnT | None = None,
-        list_roots_callback: ListRootsFnT | None = None,
+        sampling_callback: SamplingFnT[ClientSession] | None = None,
+        elicitation_callback: ElicitationFnT[ClientSession] | None = None,
+        list_roots_callback: ListRootsFnT[ClientSession] | None = None,
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
         *,
         sampling_capabilities: types.SamplingCapability | None = None,
-        experimental_task_handlers: ExperimentalTaskHandlers | None = None,
+        experimental_task_handlers: ExperimentalTaskHandlers[ClientSession] | None = None,
     ) -> None:
         super().__init__(read_stream, write_stream, read_timeout_seconds=read_timeout_seconds)
         self._client_info = client_info or DEFAULT_CLIENT_INFO

src/mcp/client/session_group.py

@@ -11,7 +11,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, TypeAlias
+from typing import Any, Generic, TypeAlias
 
 import anyio
 import httpx
@@ -20,7 +20,8 @@
 
 import mcp
 from mcp import types
-from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+from mcp.client.base_client_session import ClientSessionT_contra
+from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters
 from mcp.client.streamable_http import streamable_http_client
@@ -70,13 +71,13 @@ class StreamableHttpParameters(BaseModel):
 # Use dataclass instead of pydantic BaseModel
 # because pydantic BaseModel cannot handle Protocol fields.
 @dataclass
-class ClientSessionParameters:
+class ClientSessionParameters(Generic[ClientSessionT_contra]):
     """Parameters for establishing a client session to an MCP server."""
 
     read_timeout_seconds: float | None = None
-    sampling_callback: SamplingFnT | None = None
-    elicitation_callback: ElicitationFnT | None = None
-    list_roots_callback: ListRootsFnT | None = None
+    sampling_callback: SamplingFnT[ClientSessionT_contra] | None = None
+    elicitation_callback: ElicitationFnT[ClientSessionT_contra] | None = None
+    list_roots_callback: ListRootsFnT[ClientSessionT_contra] | None = None
     logging_callback: LoggingFnT | None = None
     message_handler: MessageHandlerFnT | None = None
     client_info: types.Implementation | None = None
@@ -254,7 +255,7 @@ async def connect_with_session(
     async def connect_to_server(
         self,
         server_params: ServerParameters,
-        session_params: ClientSessionParameters | None = None,
+        session_params: ClientSessionParameters[ClientSession] | None = None,
     ) -> mcp.ClientSession:
         """Connects to a single MCP server."""
         server_info, session = await self._establish_session(server_params, session_params or ClientSessionParameters())
@@ -263,7 +264,7 @@ async def connect_to_server(
     async def _establish_session(
         self,
         server_params: ServerParameters,
-        session_params: ClientSessionParameters,
+        session_params: ClientSessionParameters[ClientSession],
     ) -> tuple[types.Implementation, mcp.ClientSession]:
         """Establish a client session to an MCP server."""