Add sampling and roots functionality and tests

examples/sampling.py

@@ -0,0 +1,52 @@
+"""
+Example of using sampling to request an LLM completion via Marvin
+"""
+
+import asyncio
+
+import marvin
+from mcp.types import TextContent
+
+from fastmcp import Client, Context, FastMCP
+from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+
+# -- Create a server that sends a sampling request to the LLM
+
+mcp = FastMCP("Sampling Example")
+
+
+@mcp.tool()
+async def example_tool(prompt: str, context: Context) -> str:
+    """Sample a completion from the LLM."""
+    response = await context.sample(
+        "What is your favorite programming language?",
+        system_prompt="You love languages named after snakes.",
+    )
+    assert isinstance(response, TextContent)
+    return response.text
+
+
+# -- Create a client that can handle the sampling request
+
+
+async def sampling_fn(
+    messages: list[SamplingMessage],
+    params: SamplingParams,
+    ctx: RequestContext,
+) -> str:
+    return await marvin.say_async(
+        message=[m.content.text for m in messages],
+        instructions=params.systemPrompt,
+    )
+
+
+async def run():
+    async with Client(mcp, sampling_handler=sampling_fn) as client:
+        result = await client.call_tool(
+            "example_tool", {"prompt": "What is the best programming language?"}
+        )
+        print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(run())

src/fastmcp/client/client.py

@@ -6,26 +6,22 @@ from typing import Any
 import mcp.types
 from mcp import ClientSession
 from mcp.client.session import (
-    ListRootsFnT,
     LoggingFnT,
     MessageHandlerFnT,
-    SamplingFnT,
 )
-from mcp.shared.context import LifespanContextT, RequestContext
 from pydantic import AnyUrl
 
+from fastmcp.client.roots import (
+    RootsHandler,
+    RootsList,
+    create_roots_callback,
+)
+from fastmcp.client.sampling import SamplingHandler, create_sampling_callback
 from fastmcp.server import FastMCP
 
 from .transports import ClientTransport, SessionKwargs, infer_transport
 
-
-def _get_roots_callback(roots: list[mcp.types.Root]) -> ListRootsFnT | None:
-    async def _roots_callback(
-        context: RequestContext[ClientSession, LifespanContextT],
-    ) -> mcp.types.ListRootsResult:
-        return mcp.types.ListRootsResult(roots=roots)
-
-    return _roots_callback
+__all__ = ["Client", "RootsHandler", "RootsList"]
 
 
 class Client:
@@ -40,10 +36,9 @@ class Client:
         self,
         transport: ClientTransport | FastMCP | AnyUrl | Path | str,
         # Common args
-        roots: list[mcp.types.Root] | None = None,
-        sampling_callback: SamplingFnT | None = None,
-        list_roots_callback: ListRootsFnT | None = None,
-        logging_callback: LoggingFnT | None = None,
+        roots: RootsList | RootsHandler | None = None,
+        sampling_handler: SamplingHandler | None = None,
+        log_handler: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         read_timeout_seconds: datetime.timedelta | None = None,
     ):
@@ -51,21 +46,20 @@ class Client:
         self._session: ClientSession | None = None
         self._session_cm: AbstractAsyncContextManager[ClientSession] | None = None
 
-        # Store common kwargs to pass to transport.connect_session
-        if roots is not None and list_roots_callback is not None:
-            raise ValueError("Cannot provide both `roots` and `list_roots_callback`.")
-        resolved_list_roots_callback = list_roots_callback or (
-            _get_roots_callback(roots) if roots else None
-        )
-
         self._session_kwargs: SessionKwargs = {
-            "sampling_callback": sampling_callback,
-            "list_roots_callback": resolved_list_roots_callback,
-            "logging_callback": logging_callback,
+            "sampling_callback": None,
+            "list_roots_callback": None,
+            "logging_callback": log_handler,
             "message_handler": message_handler,
             "read_timeout_seconds": read_timeout_seconds,
         }
 
+        if roots is not None:
+            self.set_roots(roots)
+
+        if sampling_handler is not None:
+            self.set_sampling_callback(sampling_handler)
+
     @property
     def session(self) -> ClientSession:
         """Get the current active session. Raises RuntimeError if not connected."""
@@ -75,6 +69,16 @@ class Client:
             )
         return self._session
 
+    def set_roots(self, roots: RootsList | RootsHandler) -> None:
+        """Set the roots for the client. This does not automatically call `send_roots_list_changed`."""
+        self._session_kwargs["list_roots_callback"] = create_roots_callback(roots)
+
+    def set_sampling_callback(self, sampling_callback: SamplingHandler) -> None:
+        """Set the sampling callback for the client."""
+        self._session_kwargs["sampling_callback"] = create_sampling_callback(
+            sampling_callback
+        )
+
     def is_connected(self) -> bool:
         """Check if the client is currently connected."""
         return self._session is not None

src/fastmcp/client/roots.py

@@ -0,0 +1,75 @@
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import TypeAlias
+
+import mcp.types
+import pydantic
+from mcp import ClientSession
+from mcp.client.session import ListRootsFnT
+from mcp.shared.context import LifespanContextT, RequestContext
+
+RootsList: TypeAlias = list[str] | list[mcp.types.Root] | list[str | mcp.types.Root]
+
+RootsHandler: TypeAlias = (
+    Callable[[RequestContext[ClientSession, LifespanContextT]], RootsList]
+    | Callable[[RequestContext[ClientSession, LifespanContextT]], Awaitable[RootsList]]
+)
+
+
+def convert_roots_list(roots: RootsList) -> list[mcp.types.Root]:
+    roots_list = []
+    for r in roots:
+        if isinstance(r, mcp.types.Root):
+            roots_list.append(r)
+        elif isinstance(r, pydantic.FileUrl):
+            roots_list.append(mcp.types.Root(uri=r))
+        elif isinstance(r, str):
+            roots_list.append(mcp.types.Root(uri=pydantic.FileUrl(r)))
+        else:
+            raise ValueError(f"Invalid root: {r}")
+    return roots_list
+
+
+def create_roots_callback(
+    handler: RootsList | RootsHandler,
+) -> ListRootsFnT:
+    if isinstance(handler, list):
+        return _create_roots_callback_from_roots(handler)
+    elif inspect.isfunction(handler):
+        return _create_roots_callback_from_fn(handler)
+    else:
+        raise ValueError(f"Invalid roots handler: {handler}")
+
+
+def _create_roots_callback_from_roots(
+    roots: RootsList,
+) -> ListRootsFnT:
+    roots = convert_roots_list(roots)
+
+    async def _roots_callback(
+        context: RequestContext[ClientSession, LifespanContextT],
+    ) -> mcp.types.ListRootsResult:
+        return mcp.types.ListRootsResult(roots=roots)
+
+    return _roots_callback
+
+
+def _create_roots_callback_from_fn(
+    fn: Callable[[RequestContext[ClientSession, LifespanContextT]], RootsList]
+    | Callable[[RequestContext[ClientSession, LifespanContextT]], Awaitable[RootsList]],
+) -> ListRootsFnT:
+    async def _roots_callback(
+        context: RequestContext[ClientSession, LifespanContextT],
+    ) -> mcp.types.ListRootsResult | mcp.types.ErrorData:
+        try:
+            roots = fn(context)
+            if inspect.isawaitable(roots):
+                roots = await roots
+            return mcp.types.ListRootsResult(roots=convert_roots_list(roots))
+        except Exception as e:
+            return mcp.types.ErrorData(
+                code=mcp.types.INTERNAL_ERROR,
+                message=str(e),
+            )
+
+    return _roots_callback

src/fastmcp/client/sampling.py

@@ -0,0 +1,50 @@
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import TypeAlias
+
+import mcp.types
+from mcp import ClientSession, CreateMessageResult
+from mcp.client.session import SamplingFnT
+from mcp.shared.context import LifespanContextT, RequestContext
+from mcp.types import CreateMessageRequestParams as SamplingParams
+from mcp.types import SamplingMessage
+
+
+class MessageResult(CreateMessageResult):
+    role: mcp.types.Role = "assistant"
+    content: mcp.types.TextContent | mcp.types.ImageContent
+    model: str = "client-model"
+
+
+SamplingHandler: TypeAlias = Callable[
+    [
+        list[SamplingMessage],
+        SamplingParams,
+        RequestContext[ClientSession, LifespanContextT],
+    ],
+    str | CreateMessageResult | Awaitable[str | CreateMessageResult],
+]
+
+
+def create_sampling_callback(sampling_handler: SamplingHandler) -> SamplingFnT:
+    async def _sampling_handler(
+        context: RequestContext[ClientSession, LifespanContextT],
+        params: SamplingParams,
+    ) -> CreateMessageResult | mcp.types.ErrorData:
+        try:
+            result = sampling_handler(params.messages, params, context)
+            if inspect.isawaitable(result):
+                result = await result
+
+            if isinstance(result, str):
+                result = MessageResult(
+                    content=mcp.types.TextContent(type="text", text=result)
+                )
+            return result
+        except Exception as e:
+            return mcp.types.ErrorData(
+                code=mcp.types.INTERNAL_ERROR,
+                message=str(e),
+            )
+
+    return _sampling_handler

src/fastmcp/server/context.py

@@ -1,6 +1,5 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import Iterable
 from typing import Any, Generic, Literal
 
 from mcp.server.lowlevel.helper_types import ReadResourceContents
@@ -9,6 +8,7 @@ from mcp.shared.context import LifespanContextT, RequestContext
 from mcp.types import (
     CreateMessageResult,
     ImageContent,
+    Root,
     SamplingMessage,
     TextContent,
 )
@@ -106,7 +106,7 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
             progress_token=progress_token, progress=progress, total=total
         )
 
-    async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]:
+    async def read_resource(self, uri: str | AnyUrl) -> list[ReadResourceContents]:
         """Read a resource by URI.
 
         Args:
@@ -175,9 +175,14 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
         """Send an error log message."""
         await self.log("error", message, **extra)
 
+    async def list_roots(self) -> list[Root]:
+        """List the roots available to the server, as indicated by the client."""
+        result = await self.request_context.session.list_roots()
+        return result.roots
+
     async def sample(
         self,
-        message: str,
+        messages: str | list[str | SamplingMessage],
         system_prompt: str | None = None,
         temperature: float | None = None,
         max_tokens: int | None = None,
@@ -193,16 +198,22 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
         if max_tokens is None:
             max_tokens = 512
 
-        assert self._request_context is not None
-        assert self._request_context.session is not None
-
-        sampling_message = SamplingMessage(
-            content=TextContent(text=message, type="text"),
-            role="user",
-        )
+        if isinstance(messages, str):
+            sampling_messages = [
+                SamplingMessage(
+                    content=TextContent(text=messages, type="text"), role="user"
+                )
+            ]
+        elif isinstance(messages, list):
+            sampling_messages = [
+                SamplingMessage(content=TextContent(text=m, type="text"), role="user")
+                if isinstance(m, str)
+                else m
+                for m in messages
+            ]
 
         result: CreateMessageResult = await self.request_context.session.create_message(
-            messages=[sampling_message],
+            messages=sampling_messages,
             system_prompt=system_prompt,
             temperature=temperature,
             max_tokens=max_tokens,

tests/client/test_roots.py

@@ -0,0 +1,48 @@
+import json
+
+import pytest
+from mcp.types import TextContent
+
+from fastmcp import Client, Context, FastMCP
+
+
+@pytest.fixture
+def fastmcp_server():
+    mcp = FastMCP()
+
+    @mcp.tool()
+    async def list_roots(context: Context) -> list[str]:
+        roots = await context.list_roots()
+        return [str(r.uri) for r in roots]
+
+    return mcp
+
+
+class TestClientRoots:
+    @pytest.mark.parametrize("roots", [["x"], ["x", "y"]])
+    async def test_invalid_roots(self, fastmcp_server: FastMCP, roots: list[str]):
+        """
+        Roots must be URIs
+        """
+        with pytest.raises(ValueError, match="Input should be a valid URL"):
+            async with Client(fastmcp_server, roots=roots):
+                pass
+
+    @pytest.mark.parametrize("roots", [["https://x.com"]])
+    async def test_invalid_urls(self, fastmcp_server: FastMCP, roots: list[str]):
+        """
+        At this time, root URIs must start with file://
+        """
+        with pytest.raises(ValueError, match="URL scheme should be 'file'"):
+            async with Client(fastmcp_server, roots=roots):
+                pass
+
+    @pytest.mark.parametrize("roots", [["file://x/y/z", "file://x/y/z"]])
+    async def test_valid_roots(self, fastmcp_server: FastMCP, roots: list[str]):
+        async with Client(fastmcp_server, roots=roots) as client:
+            result = await client.call_tool("list_roots", {})
+            assert isinstance(result.content[0], TextContent)
+            assert json.loads(result.content[0].text) == [
+                "file://x/y/z",
+                "file://x/y/z",
+            ]

tests/client/test_sampling.py

@@ -0,0 +1,85 @@
+from typing import cast
+
+import pytest
+from mcp.types import TextContent
+
+from fastmcp import Client, Context, FastMCP
+from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+
+
+@pytest.fixture
+def fastmcp_server():
+    mcp = FastMCP()
+
+    @mcp.tool()
+    async def simple_sample(message: str, context: Context) -> str:
+        result = await context.sample("Hello, world!")
+        return cast(TextContent, result).text
+
+    @mcp.tool()
+    async def sample_with_system_prompt(message: str, context: Context) -> str:
+        result = await context.sample("Hello, world!", system_prompt="You love FastMCP")
+        return cast(TextContent, result).text
+
+    @mcp.tool()
+    async def sample_with_messages(message: str, context: Context) -> str:
+        result = await context.sample(
+            [
+                "Hello!",
+                SamplingMessage(
+                    content=TextContent(
+                        type="text", text="How can I assist you today?"
+                    ),
+                    role="assistant",
+                ),
+            ]
+        )
+        return cast(TextContent, result).text
+
+    return mcp
+
+
+async def test_simple_sampling(fastmcp_server: FastMCP):
+    def sampling_handler(
+        messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
+    ) -> str:
+        return "This is the sample message!"
+
+    async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
+        result = await client.call_tool("simple_sample", {"message": "Hello, world!"})
+        reply = cast(TextContent, result.content[0])
+        assert reply.text == "This is the sample message!"
+
+
+async def test_sampling_with_system_prompt(fastmcp_server: FastMCP):
+    def sampling_handler(
+        messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
+    ) -> str:
+        assert params.systemPrompt is not None
+        return params.systemPrompt
+
+    async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
+        result = await client.call_tool(
+            "sample_with_system_prompt", {"message": "Hello, world!"}
+        )
+        reply = cast(TextContent, result.content[0])
+        assert reply.text == "You love FastMCP"
+
+
+async def test_sampling_with_messages(fastmcp_server: FastMCP):
+    def sampling_handler(
+        messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
+    ) -> str:
+        assert len(messages) == 2
+        assert messages[0].content.type == "text"
+        assert messages[0].content.text == "Hello!"
+        assert messages[1].content.type == "text"
+        assert messages[1].content.text == "How can I assist you today?"
+        return "I need to think."
+
+    async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
+        result = await client.call_tool(
+            "sample_with_messages", {"message": "Hello, world!"}
+        )
+        reply = cast(TextContent, result.content[0])
+        assert reply.text == "I need to think."