Merge kziemski transport abstractions review suggestions with test fixes

Author: sreenithi

docs/migration.md

@@ -471,6 +471,193 @@ await client.read_resource("test://resource")
 await client.read_resource(str(my_any_url))
 ```
 
+### Transport Abstractions Refactored
+
+The session hierarchy has been refactored to support pluggable transport implementations. This introduces several breaking changes:
+
+#### `ClientRequestContext` type changed
+
+`ClientRequestContext` is now `RequestContext[BaseClientSession]` instead of `RequestContext[ClientSession]`. This means callbacks receive the more general `BaseClientSession` type, which may not have all methods available on `ClientSession`.
+
+**Before:**
+
+```python
+from mcp.client.context import ClientRequestContext
+from mcp.client.session import ClientSession
+
+async def my_callback(context: ClientRequestContext) -> None:
+    # Could access ClientSession-specific methods
+    caps = context.session.get_server_capabilities()
+```
+
+**After:**
+
+```python
+from mcp.client.context import ClientRequestContext
+from mcp.client.session import ClientSession
+
+async def my_callback(context: ClientRequestContext) -> None:
+    # context.session is BaseClientSession - narrow the type if needed
+    if isinstance(context.session, ClientSession):
+        caps = context.session.get_server_capabilities()
+```
+
+#### Callback protocols are now generic
+
+`sampling_callback`, `elicitation_callback`, and `list_roots_callback` protocols now require explicit type parameters.
+
+**Before:**
+
+```python
+from mcp.client.session import SamplingFnT
+
+async def my_sampling(context, params) -> CreateMessageResult:
+    ...
+
+# Type inferred as SamplingFnT
+session = ClientSession(..., sampling_callback=my_sampling)
+```
+
+**After:**
+
+```python
+from mcp.client.session import SamplingFnT, ClientSession
+
+async def my_sampling(
+    context: RequestContext[ClientSession],
+    params: CreateMessageRequestParams
+) -> CreateMessageResult:
+    ...
+
+# Explicit type annotation recommended
+my_sampling_typed: SamplingFnT[ClientSession] = my_sampling
+session = ClientSession(..., sampling_callback=my_sampling_typed)
+```
+
+#### `SessionT` renamed to `SessionT_co`
+
+In `mcp.shared._context` and `mcp.shared.progress`, the `SessionT` TypeVar has been renamed to `SessionT_co` to follow naming conventions for covariant type variables.
+
+**Before:**
+
+```python
+from mcp.shared._context import SessionT
+```
+
+**After:**
+
+```python
+from mcp.shared._context import SessionT_co
+```
+
+#### New `AbstractBaseSession` structural interface
+
+The session hierarchy now uses a new **runtime-checkable Protocol** called `AbstractBaseSession` to define the shared contract for all MCP sessions. This protocol enables structural subtyping, allowing different transport implementations to be used interchangeably without requiring rigid inheritance.
+
+Key characteristics of `AbstractBaseSession`:
+1.  **Pure Interface**: It is a structural protocol with no implementation state or `__init__` method.
+2.  **Simplified Type Parameters**: It takes two parameters: `AbstractBaseSession[SendRequestT, SendNotificationT]`. Contravariant variance is used for these parameters to ensure that sessions can be used safely in generic contexts (like `RequestContext`).
+3.  **BaseSession Implementation**: The concrete implementation logic (state management, response routing) is provided by the `BaseSession` class, which satisfies the protocol.
+
+**Before:**
+
+```python
+from mcp.shared.session import AbstractBaseSession
+
+class MySession(AbstractBaseSession[MyMessage, ...]):
+    def __init__(self):
+        super().__init__()  # Would set up _response_streams, _task_group
+```
+
+**After:**
+
+```python
+from mcp.shared.session import AbstractBaseSession
+
+class MySession(AbstractBaseSession[...]):
+    def __init__(self):
+        # Manage your own state - no super().__init__() to call
+        self._my_state = {}
+```
+
+#### `SendRequestT` changed to contravariant
+
+The `SendRequestT` TypeVar is now defined as **contravariant** to support its use in the `AbstractBaseSession` Protocol.
+
+**Before:**
+
+```python
+SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
+```
+
+**After:**
+
+```python
+SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest, contravariant=True)
+```
+
+#### `SendNotificationT` changed to contravariant
+
+The `SendNotificationT` TypeVar is now defined as **contravariant** to support its use in the `AbstractBaseSession` Protocol.
+
+**Before:**
+
+```python
+SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
+```
+
+**After:**
+
+```python
+SendNotificationT = TypeVar(
+    "SendNotificationT", ClientNotification, ServerNotification, contravariant=True
+)
+```
+
+#### `ReceiveResultT` changed to covariant
+
+The `ReceiveResultT` TypeVar is now defined as **covariant** to support its use in the `AbstractBaseSession` Protocol.
+
+**Before:**
+
+```python
+ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
+```
+
+**After:**
+
+```python
+ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel, covariant=True)
+```
+
+#### `BaseClientSession` is now a Protocol
+
+`BaseClientSession` is now a `typing.Protocol` (structural subtyping) instead of an abstract base class. It no longer inherits from `AbstractBaseSession` and requires no inheritance to satisfy.
+
+**Before:**
+
+```python
+from mcp.client.base_client_session import BaseClientSession
+
+class MyClientSession(BaseClientSession):
+    async def initialize(self) -> InitializeResult:
+        ...
+```
+
+**After:**
+
+```python
+from mcp.client.base_client_session import BaseClientSession
+
+class MyClientSession:
+    # Just implement the methods - no inheritance needed
+    async def initialize(self) -> InitializeResult:
+        ...
+
+# Verify protocol satisfaction at runtime
+assert isinstance(MyClientSession(), BaseClientSession)
+```
+
 ## Deprecations
 
 <!-- Add deprecations below -->

src/mcp/client/base_client_session.py

@@ -1,75 +1,68 @@
-from abc import abstractmethod
 from typing import Any, TypeVar
 
+from typing_extensions import Protocol, runtime_checkable
+
 from mcp import types
-from mcp.shared.session import AbstractBaseSession, ProgressFnT
+from mcp.shared.session import ProgressFnT
 from mcp.types._types import RequestParamsMeta
 
 ClientSessionT_contra = TypeVar("ClientSessionT_contra", bound="BaseClientSession", contravariant=True)
 
 
-class BaseClientSession(
-    AbstractBaseSession[
-        Any,
-        types.ClientRequest,
-        types.ClientNotification,
-        types.ClientResult,
-        types.ServerRequest,
-        types.ServerNotification,
-    ]
-):
-    """Base class for client transport sessions.
-
-    The class provides all the methods that a client session should implement,
-    irrespective of the transport used.
+@runtime_checkable
+class BaseClientSession(Protocol):
+    """Protocol defining the interface for MCP client sessions.
+
+    This protocol specifies all methods that a client session must implement,
+    irrespective of the transport used. Implementations satisfy this protocol
+    through structural subtyping — no inheritance required.
     """
 
-    @abstractmethod
-    async def initialize(self) -> types.InitializeResult:
-        """Initialize the client session."""
-        raise NotImplementedError
+    # Methods from AbstractBaseSession (must be explicitly declared in Protocol)
+    async def send_request(
+        self,
+        request: types.ClientRequest,
+        result_type: type,
+        request_read_timeout_seconds: float | None = None,
+        metadata: Any = None,
+        progress_callback: ProgressFnT | None = None,
+    ) -> Any: ...
+
+    async def send_notification(
+        self,
+        notification: types.ClientNotification,
+        related_request_id: Any = None,
+    ) -> None: ...
+
+    async def send_progress_notification(
+        self,
+        progress_token: types.ProgressToken,
+        progress: float,
+        total: float | None = None,
+        message: str | None = None,
+        *,
+        meta: RequestParamsMeta | None = None,
+    ) -> None: ...
 
-    @abstractmethod
-    async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult:
-        """Send a ping request."""
-        raise NotImplementedError
+    # Client-specific methods
+    async def initialize(self) -> types.InitializeResult: ...
 
-    @abstractmethod
-    async def list_resources(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListResourcesResult:
-        """Send a resources/list request.
+    async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: ...
 
-        Args:
-            params: Full pagination parameters including cursor and any future fields
-        """
-        raise NotImplementedError
+    async def list_resources(
+        self, *, params: types.PaginatedRequestParams | None = None
+    ) -> types.ListResourcesResult: ...
 
-    @abstractmethod
     async def list_resource_templates(
         self, *, params: types.PaginatedRequestParams | None = None
-    ) -> types.ListResourceTemplatesResult:
-        """Send a resources/templates/list request.
-
-        Args:
-            params: Full pagination parameters including cursor and any future fields
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def read_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.ReadResourceResult:
-        """Send a resources/read request."""
-        raise NotImplementedError
-
-    @abstractmethod
-    async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult:
-        """Send a resources/subscribe request."""
-        raise NotImplementedError
-
-    @abstractmethod
-    async def unsubscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult:
-        """Send a resources/unsubscribe request."""
-        raise NotImplementedError
-
-    @abstractmethod
+    ) -> types.ListResourceTemplatesResult: ...
+
+    async def read_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.ReadResourceResult: ...
+
+    async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: ...
+
+    async def unsubscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult: ...
+
     async def call_tool(
         self,
         name: str,
@@ -78,40 +71,33 @@ async def call_tool(
         progress_callback: ProgressFnT | None = None,
         *,
         meta: RequestParamsMeta | None = None,
-    ) -> types.CallToolResult:
-        """Send a tools/call request with optional progress callback support."""
-        raise NotImplementedError
+    ) -> types.CallToolResult: ...
 
-    @abstractmethod
-    async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult:
-        """Send a prompts/list request.
+    async def list_prompts(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListPromptsResult: ...
 
-        Args:
-            params: Full pagination parameters including cursor and any future fields
-        """
-        raise NotImplementedError
-
-    @abstractmethod
     async def get_prompt(
         self,
         name: str,
         arguments: dict[str, str] | None = None,
         *,
         meta: RequestParamsMeta | None = None,
-    ) -> types.GetPromptResult:
-        """Send a prompts/get request."""
-        raise NotImplementedError
-
-    @abstractmethod
-    async def list_tools(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult:
-        """Send a tools/list request.
-
-        Args:
-            params: Full pagination parameters including cursor and any future fields
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def send_roots_list_changed(self) -> None:  # pragma: no cover
-        """Send a roots/list_changed notification."""
-        raise NotImplementedError
+    ) -> types.GetPromptResult: ...
+
+    async def list_tools(self, *, params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult: ...
+
+    # Missing methods added per review
+    async def complete(
+        self,
+        ref: types.ResourceTemplateReference | types.PromptReference,
+        argument: dict[str, str],
+        context_arguments: dict[str, str] | None = None,
+    ) -> types.CompleteResult: ...
+
+    async def set_logging_level(
+        self,
+        level: types.LoggingLevel,
+        *,
+        meta: RequestParamsMeta | None = None,
+    ) -> types.EmptyResult: ...
+
+    async def send_roots_list_changed(self) -> None: ...

src/mcp/client/context.py

@@ -1,8 +1,6 @@
 """Request context for MCP client handlers."""
 
 from mcp.client import BaseClientSession
-
-# from mcp.client.session import ClientSession
 from mcp.shared._context import RequestContext
 
 ClientRequestContext = RequestContext[BaseClientSession]

src/mcp/client/session.py

@@ -107,7 +107,6 @@ class ClientSession(
         types.ServerRequest,
         types.ServerNotification,
     ],
-    BaseClientSession,
 ):
     def __init__(
         self,