Merge branch 'main' into main

Author: abdullahkady

README.md

@@ -476,9 +476,21 @@ server_params = StdioServerParameters(
     env=None # Optional environment variables
 )
 
+# Optional: create a sampling callback
+async def handle_sampling_message(message: types.CreateMessageRequestParams) -> types.CreateMessageResult:
+    return types.CreateMessageResult(
+        role="assistant",
+        content=types.TextContent(
+            type="text",
+            text="Hello, world! from model",
+        ),
+        model="gpt-3.5-turbo",
+        stopReason="endTurn",
+    )
+
 async def run():
     async with stdio_client(server_params) as (read, write):
-        async with ClientSession(read, write) as session:
+        async with ClientSession(read, write, sampling_callback=handle_sampling_message) as session:
             # Initialize the connection
             await session.initialize()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcp"
-version = "1.3.0.dev0"
+version = "1.4.0.dev0"
 description = "Model Context Protocol SDK"
 readme = "README.md"
 requires-python = ">=3.10"

src/mcp/client/session.py

@@ -1,13 +1,51 @@
 from datetime import timedelta
+from typing import Any, Protocol
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from pydantic import AnyUrl
+from pydantic import AnyUrl, TypeAdapter
 
 import mcp.types as types
-from mcp.shared.session import BaseSession
+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 | types.ErrorData: ...
+
+
+class ListRootsFnT(Protocol):
+    async def __call__(
+        self, context: RequestContext["ClientSession", Any]
+    ) -> 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(
     BaseSession[
         types.ClientRequest,
@@ -22,6 +60,8 @@ def __init__(
         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,
@@ -30,23 +70,34 @@ def __init__(
             types.ServerNotification,
             read_timeout_seconds=read_timeout_seconds,
         )
+        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 = (
+            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(
             types.ClientRequest(
                 types.InitializeRequest(
                     method="initialize",
                     params=types.InitializeRequestParams(
                         protocolVersion=types.LATEST_PROTOCOL_VERSION,
                         capabilities=types.ClientCapabilities(
-                            sampling=None,
+                            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"),
                     ),
@@ -243,3 +294,32 @@ 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):
+                with responder:
+                    response = await self._sampling_callback(ctx, params)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
+            case types.ListRootsRequest():
+                with responder:
+                    response = await self._list_roots_callback(ctx)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
+            case types.PingRequest():
+                with responder:
+                    return await responder.respond(
+                        types.ClientResult(root=types.EmptyResult())
+                    )

src/mcp/server/fastmcp/server.py

@@ -3,7 +3,7 @@
 import inspect
 import json
 import re
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Iterable
 from contextlib import (
     AbstractAsyncContextManager,
     asynccontextmanager,
@@ -34,6 +34,7 @@
 from mcp.server.lowlevel.server import (
     lifespan as default_lifespan,
 )
+from mcp.server.session import ServerSession
 from mcp.server.sse import SseServerTransport
 from mcp.server.stdio import stdio_server
 from mcp.shared.context import RequestContext
@@ -235,7 +236,7 @@ async def list_resource_templates(self) -> list[MCPResourceTemplate]:
             for template in templates
         ]
 
-    async def read_resource(self, uri: AnyUrl | str) -> ReadResourceContents:
+    async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContents]:
         """Read a resource by URI."""
 
         resource = await self._resource_manager.get_resource(uri)
@@ -244,7 +245,7 @@ async def read_resource(self, uri: AnyUrl | str) -> ReadResourceContents:
 
         try:
             content = await resource.read()
-            return ReadResourceContents(content=content, mime_type=resource.mime_type)
+            return [ReadResourceContents(content=content, mime_type=resource.mime_type)]
         except Exception as e:
             logger.error(f"Error reading resource {uri}: {e}")
             raise ResourceError(str(e))
@@ -597,7 +598,7 @@ def my_tool(x: int, ctx: Context) -> str:
     The context is optional - tools that don't need it can omit the parameter.
     """
 
-    _request_context: RequestContext | None
+    _request_context: RequestContext[ServerSession, Any] | None
     _fastmcp: FastMCP | None
 
     def __init__(
@@ -648,7 +649,7 @@ async def report_progress(
             progress_token=progress_token, progress=progress, total=total
         )
 
-    async def read_resource(self, uri: str | AnyUrl) -> ReadResourceContents:
+    async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]:
         """Read a resource by URI.
 
         Args:

src/mcp/server/lowlevel/server.py

@@ -67,9 +67,9 @@ async def main():
 import contextvars
 import logging
 import warnings
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Iterable
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
-from typing import Any, AsyncIterator, Generic, Sequence, TypeVar
+from typing import Any, AsyncIterator, Generic, TypeVar
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@@ -279,7 +279,9 @@ async def handler(_: Any):
 
     def read_resource(self):
         def decorator(
-            func: Callable[[AnyUrl], Awaitable[str | bytes | ReadResourceContents]],
+            func: Callable[
+                [AnyUrl], Awaitable[str | bytes | Iterable[ReadResourceContents]]
+            ],
         ):
             logger.debug("Registering handler for ReadResourceRequest")
 
@@ -307,13 +309,22 @@ def create_content(data: str | bytes, mime_type: str | None):
                     case str() | bytes() as data:
                         warnings.warn(
                             "Returning str or bytes from read_resource is deprecated. "
-                            "Use ReadResourceContents instead.",
+                            "Use Iterable[ReadResourceContents] instead.",
                             DeprecationWarning,
                             stacklevel=2,
                         )
                         content = create_content(data, None)
-                    case ReadResourceContents() as contents:
-                        content = create_content(contents.content, contents.mime_type)
+                    case Iterable() as contents:
+                        contents_list = [
+                            create_content(content_item.content, content_item.mime_type)
+                            for content_item in contents
+                            if isinstance(content_item, ReadResourceContents)
+                        ]
+                        return types.ServerResult(
+                            types.ReadResourceResult(
+                                contents=contents_list,
+                            )
+                        )
                     case _:
                         raise ValueError(
                             f"Unexpected return type from read_resource: {type(result)}"
@@ -387,7 +398,7 @@ def decorator(
             func: Callable[
                 ...,
                 Awaitable[
-                    Sequence[
+                    Iterable[
                         types.TextContent | types.ImageContent | types.EmbeddedResource
                     ]
                 ],

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, ListRootsFnT, SamplingFnT
 from mcp.server import Server
 from mcp.types import JSONRPCMessage
 
@@ -54,6 +54,8 @@ 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,
+    list_roots_callback: ListRootsFnT | None = None,
     raise_exceptions: bool = False,
 ) -> AsyncGenerator[ClientSession, None]:
     """Creates a ClientSession that is connected to a running MCP server."""
@@ -80,6 +82,8 @@ 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,
+                    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,70 @@
+import pytest
+from pydantic import FileUrl
+
+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,
+    TextContent,
+)
+
+
+@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("file://users/fake/test"),
+                name="Test Root 1",
+            ),
+            Root(
+                uri=FileUrl("file://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 = 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
+        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"
+        )