Add client handling for sampling, list roots, ping (#218)

Adds sampling and list roots callbacks to the ClientSession, allowing the client to handle requests from the server.

Co-authored-by: TerminalMan <84923604+SecretiveShell@users.noreply.github.com>
Co-authored-by: David Soria Parra <davidsp@anthropic.com>
This commit is contained in:
Jerome
2025-02-20 10:49:43 +00:00
committed by GitHub
parent 106619967b
commit ff22f48365
6 changed files with 256 additions and 12 deletions

View File

@@ -9,7 +9,7 @@ from typing import AsyncGenerator
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