Added roots callback also

Author: jerome3o-anthropic

src/mcp/client/session.py

@@ -1,16 +1,27 @@
 from datetime import timedelta
-from typing import Awaitable, Callable
+from typing import Protocol, Any
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import AnyUrl
 
+from mcp.shared.context import RequestContext
 import mcp.types as types
 from mcp.shared.session import BaseSession, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
-SamplingFnT = Callable[
-    [types.CreateMessageRequestParams], Awaitable[types.CreateMessageResult]
-]
+
+class SamplingFnT(Protocol):
+    async def __call__(
+        self, context: RequestContext["ClientSession", Any], params: types.CreateMessageRequestParams
+    ) -> types.CreateMessageResult:
+        ...
+
+
+class ListRootsFnT(Protocol):
+    async def __call__(
+        self, context: RequestContext["ClientSession", Any]
+    ) -> types.ListRootsResult:
+        ...
 
 
 class ClientSession(
@@ -22,14 +33,15 @@ class ClientSession(
         types.ServerNotification,
     ]
 ):
-    sampling_callback: SamplingFnT | None = None
+    _sampling_callback: SamplingFnT | None = None
 
     def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
         write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
         read_timeout_seconds: timedelta | None = None,
         sampling_callback: SamplingFnT | None = None,
+        list_roots_callback: ListRootsFnT | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -38,11 +50,22 @@ def __init__(
             types.ServerNotification,
             read_timeout_seconds=read_timeout_seconds,
         )
-        self.sampling_callback = sampling_callback
+        self._sampling_callback = sampling_callback
+        self._list_roots_callback = list_roots_callback
 
     async def initialize(self) -> types.InitializeResult:
         sampling = (
-            types.SamplingCapability() if self.sampling_callback is not None else None
+            types.SamplingCapability() if self._sampling_callback is not None else None
+        )
+        roots = (
+            types.RootsCapability(
+                # TODO: Should this be based on whether we
+                # _will_ send notifications, or only whether
+                # they're supported?
+                listChanged=True,
+            )
+            if self._list_roots_callback is not None
+            else None
         )
 
         result = await self.send_request(
@@ -54,12 +77,7 @@ async def initialize(self) -> types.InitializeResult:
                         capabilities=types.ClientCapabilities(
                             sampling=sampling,
                             experimental=None,
-                            roots=types.RootsCapability(
-                                # TODO: Should this be based on whether we
-                                # _will_ send notifications, or only whether
-                                # they're supported?
-                                listChanged=True
-                            ),
+                            roots=roots,
                         ),
                         clientInfo=types.Implementation(name="mcp", version="0.1.0"),
                     ),
@@ -258,11 +276,29 @@ async def send_roots_list_changed(self) -> None:
         )
 
     async def _received_request(
-        self, responder: RequestResponder["types.ServerRequest", "types.ClientResult"]
+        self, responder: RequestResponder[types.ServerRequest, types.ClientResult]
     ) -> None:
-        if isinstance(responder.request.root, types.CreateMessageRequest):
-            if self.sampling_callback is not None:
-                response = await self.sampling_callback(responder.request.root.params)
-                client_response = types.ClientResult(root=response)
+
+        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:
+                if self._sampling_callback is not None:
+                    response = await self._sampling_callback(ctx, responder.request.root.params)
+                    client_response = types.ClientResult(root=response)
+                    with responder:
+                        await responder.respond(client_response)
+            case types.ListRootsRequest:
+                if self._list_roots_callback is not None:
+                    response = await self._list_roots_callback(ctx)
+                    client_response = types.ClientResult(root=response)
+                    with responder:
+                        await responder.respond(client_response)
+            case types.PingRequest:
                 with responder:
-                    await responder.respond(client_response)
+                    await responder.respond(types.ClientResult(root=types.EmptyResult()))

src/mcp/shared/memory.py

@@ -9,7 +9,7 @@
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
-from mcp.client.session import ClientSession, SamplingFnT
+from mcp.client.session import ClientSession, ListRootsFnT, SamplingFnT
 from mcp.server import Server
 from mcp.types import JSONRPCMessage
 
@@ -55,6 +55,7 @@ async def create_connected_server_and_client_session(
     server: Server,
     read_timeout_seconds: timedelta | None = None,
     sampling_callback: SamplingFnT | None = None,
+    list_roots_callback: ListRootsFnT | None = None,
     raise_exceptions: bool = False,
 ) -> AsyncGenerator[ClientSession, None]:
     """Creates a ClientSession that is connected to a running MCP server."""
@@ -82,6 +83,7 @@ async def create_connected_server_and_client_session(
                     write_stream=client_write,
                     read_timeout_seconds=read_timeout_seconds,
                     sampling_callback=sampling_callback,
+                    list_roots_callback=list_roots_callback,
                 ) as client_session:
                     await client_session.initialize()
                     yield client_session

tests/client/test_list_roots_callback.py

@@ -0,0 +1,50 @@
+from pydantic import FileUrl
+import pytest
+
+from mcp.client.session import ClientSession
+from mcp.server.fastmcp.server import Context
+from mcp.shared.context import RequestContext
+from mcp.shared.memory import (
+    create_connected_server_and_client_session as create_session,
+)
+from mcp.types import (
+    ListRootsResult,
+    Root,
+)
+
+
+@pytest.mark.anyio
+async def test_list_roots_callback():
+    from mcp.server.fastmcp import FastMCP
+
+    server = FastMCP("test")
+
+    callback_return = ListRootsResult(roots=[
+        Root(
+            uri=FileUrl("test://users/fake/test"),
+            name="Test Root 1",
+        ),
+        Root(
+            uri=FileUrl("test://users/fake/test/2"),
+            name="Test Root 2",
+        )
+    ])
+
+    async def list_roots_callback(
+        context: RequestContext[ClientSession, None]
+    ) -> ListRootsResult:
+        return callback_return
+
+    @server.tool("test_list_roots")
+    async def test_list_roots(context: Context, message: str):
+        roots = context.session.list_roots()
+        assert roots == callback_return
+        return True
+
+    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(
+            "test_list_roots", {"message": "test message"}
+        )

tests/client/test_sampling_callback.py

@@ -1,5 +1,7 @@
 import pytest
 
+from mcp.client.session import ClientSession
+from mcp.shared.context import RequestContext
 from mcp.shared.memory import (
     create_connected_server_and_client_session as create_session,
 )
@@ -27,7 +29,8 @@ async def test_sampling_callback():
     )
 
     async def sampling_callback(
-        message: CreateMessageRequestParams,
+        context: RequestContext[ClientSession, None],
+        params: CreateMessageRequestParams,
     ) -> CreateMessageResult:
         return callback_return