fix: simplify implementation and add test

Author: dsp-ant

src/mcp/client/session.py

@@ -8,7 +8,7 @@
 from mcp.shared.session import BaseSession, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
-sampling_function_signature = Callable[
+SamplingFnT = Callable[
     [types.CreateMessageRequestParams], Awaitable[types.CreateMessageResult]
 ]
 
@@ -22,14 +22,14 @@ class ClientSession(
         types.ServerNotification,
     ]
 ):
-    sampling_callback: sampling_function_signature | 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: sampling_function_signature | None = None,
+        sampling_callback: SamplingFnT | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -253,4 +253,5 @@ async def _received_request(
             if self.sampling_callback is not None:
                 response = await self.sampling_callback(responder.request.root.params)
                 client_response = types.ClientResult(root=response)
-                await responder.respond(client_response)
+                with responder:
+                    await responder.respond(client_response)

src/mcp/shared/memory.py

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

tests/client/test_sampling_callback.py

@@ -0,0 +1,53 @@
+import pytest
+
+from mcp.shared.memory import (
+    create_connected_server_and_client_session as create_session,
+)
+from mcp.types import (
+    CreateMessageRequestParams,
+    CreateMessageResult,
+    SamplingMessage,
+    TextContent,
+)
+
+
+@pytest.mark.anyio
+async def test_sampling_callback():
+    from mcp.server.fastmcp import FastMCP
+
+    server = FastMCP("test")
+
+    callback_return = CreateMessageResult(
+        role="assistant",
+        content=TextContent(
+            type="text", text="This is a response from the sampling callback"
+        ),
+        model="test-model",
+        stopReason="endTurn",
+    )
+
+    async def sampling_callback(
+        message: CreateMessageRequestParams,
+    ) -> CreateMessageResult:
+        return callback_return
+
+    @server.tool("test_sampling")
+    async def test_sampling_tool(message: str):
+        value = await server.get_context().session.create_message(
+            messages=[
+                SamplingMessage(
+                    role="user", content=TextContent(type="text", text=message)
+                )
+            ],
+            max_tokens=100,
+        )
+        assert value == callback_return
+        return True
+
+    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(
+            "test_sampling", {"message": "Test message for sampling"}
+        )