Refactored default behaviour, updated tests

Author: jerome3o-anthropic

src/mcp/client/session.py

@@ -1,27 +1,49 @@
 from datetime import timedelta
-from typing import Protocol, Any
+from typing import Any, Protocol
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from pydantic import AnyUrl
+from pydantic import AnyUrl, TypeAdapter
 
-from mcp.shared.context import RequestContext
 import mcp.types as types
+from mcp.shared.context import RequestContext
 from mcp.shared.session import BaseSession, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
 
 class SamplingFnT(Protocol):
     async def __call__(
-        self, context: RequestContext["ClientSession", Any], params: types.CreateMessageRequestParams
-    ) -> types.CreateMessageResult:
-        ...
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.CreateMessageRequestParams,
+    ) -> types.CreateMessageResult | types.ErrorData: ...
 
 
 class ListRootsFnT(Protocol):
     async def __call__(
         self, context: RequestContext["ClientSession", Any]
-    ) -> types.ListRootsResult:
-        ...
+    ) -> types.ListRootsResult |  types.ErrorData: ...    
+
+
+async def _default_sampling_callback(
+    context: RequestContext["ClientSession", Any],
+    params: types.CreateMessageRequestParams,
+) -> types.CreateMessageResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.INVALID_REQUEST,
+        message="Sampling not supported",
+    )
+
+
+async def _default_list_roots_callback(
+    context: RequestContext["ClientSession", Any],
+) -> types.ListRootsResult |  types.ErrorData:
+    return types.ErrorData(
+        code=types.INVALID_REQUEST,
+        message="List roots not supported",
+    )
+
+
+ClientResponse = TypeAdapter(types.ClientResult | types.ErrorData)
 
 
 class ClientSession(
@@ -33,8 +55,6 @@ class ClientSession(
         types.ServerNotification,
     ]
 ):
-    _sampling_callback: SamplingFnT | None = None
-
     def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
@@ -50,8 +70,8 @@ def __init__(
             types.ServerNotification,
             read_timeout_seconds=read_timeout_seconds,
         )
-        self._sampling_callback = sampling_callback
-        self._list_roots_callback = list_roots_callback
+        self._sampling_callback = sampling_callback or _default_sampling_callback
+        self._list_roots_callback = list_roots_callback or _default_list_roots_callback
 
     async def initialize(self) -> types.InitializeResult:
         sampling = (
@@ -278,27 +298,28 @@ async def send_roots_list_changed(self) -> None:
     async def _received_request(
         self, responder: RequestResponder[types.ServerRequest, types.ClientResult]
     ) -> None:
-
         ctx = RequestContext[ClientSession, Any](
             request_id=responder.request_id,
             meta=responder.request_meta,
             session=self,
             lifespan_context=None,
         )
-        
+
         match responder.request.root:
             case types.CreateMessageRequest(params=params):
-                if self._sampling_callback is not None:
+                with responder:
                     response = await self._sampling_callback(ctx, params)
-                    client_response = types.ClientResult(root=response)
-                    with responder:
-                        await responder.respond(client_response)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
             case types.ListRootsRequest():
-                if self._list_roots_callback is not None:
+                with responder:
                     response = await self._list_roots_callback(ctx)
-                    client_response = types.ClientResult(root=response)
-                    with responder:
-                        await responder.respond(client_response)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+                    
             case types.PingRequest():
                 with responder:
-                    await responder.respond(types.ClientResult(root=types.EmptyResult()))
+                    return await responder.respond(
+                        types.ClientResult(root=types.EmptyResult())
+                    )

tests/client/test_list_roots_callback.py

@@ -1,5 +1,5 @@
-from pydantic import FileUrl
 import pytest
+from pydantic import FileUrl
 
 from mcp.client.session import ClientSession
 from mcp.server.fastmcp.server import Context
@@ -10,6 +10,7 @@
 from mcp.types import (
     ListRootsResult,
     Root,
+    TextContent,
 )
 
 
@@ -21,11 +22,11 @@ async def test_list_roots_callback():
 
     callback_return = ListRootsResult(roots=[
         Root(
-            uri=FileUrl("test://users/fake/test"),
+            uri=FileUrl("file://users/fake/test"),
             name="Test Root 1",
         ),
         Root(
-            uri=FileUrl("test://users/fake/test/2"),
+            uri=FileUrl("file://users/fake/test/2"),
             name="Test Root 2",
         )
     ])
@@ -37,14 +38,29 @@ async def list_roots_callback(
 
     @server.tool("test_list_roots")
     async def test_list_roots(context: Context, message: str):
-        roots = context.session.list_roots()
+        roots = await context.session.list_roots()
         assert roots == callback_return
         return True
 
+    # Test with list_roots callback
     async with create_session(
         server._mcp_server, list_roots_callback=list_roots_callback
     ) as client_session:
         # Make a request to trigger sampling callback
-        assert await client_session.call_tool(
+        result = await client_session.call_tool(
             "test_list_roots", {"message": "test message"}
         )
+        assert result.isError is False
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == 'true'
+
+    # Test without list_roots callback
+    async with create_session(server._mcp_server) as client_session:
+        # Make a request to trigger sampling callback
+        result = await client_session.call_tool(
+            "test_list_roots", {"message": "test message"}
+        )
+        assert result.isError is True
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == 'Error executing tool test_list_roots: List roots not supported'
+

tests/client/test_sampling_callback.py

@@ -47,10 +47,24 @@ async def test_sampling_tool(message: str):
         assert value == callback_return
         return True
 
+    # Test with sampling callback
     async with create_session(
         server._mcp_server, sampling_callback=sampling_callback
     ) as client_session:
         # Make a request to trigger sampling callback
-        assert await client_session.call_tool(
+        result = await client_session.call_tool(
             "test_sampling", {"message": "Test message for sampling"}
         )
+        assert result.isError is False
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == 'true'
+
+    # Test without sampling callback
+    async with create_session(server._mcp_server) as client_session:
+        # Make a request to trigger sampling callback
+        result = await client_session.call_tool(
+            "test_sampling", {"message": "Test message for sampling"}
+        )
+        assert result.isError is True
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == 'Error executing tool test_sampling: Sampling not supported'