From c71138d6b822d78cac072bd94a04bdaa5021567a Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 20 May 2025 09:27:47 -0400 Subject: [PATCH 1/2] Add progress handler to client --- docs/clients/client.mdx | 79 +++++++++++++++++++++++++++++++- src/fastmcp/client/client.py | 30 ++++++++++-- src/fastmcp/client/logging.py | 8 ++++ src/fastmcp/client/progress.py | 38 +++++++++++++++ src/fastmcp/client/transports.py | 2 + tests/client/test_progress.py | 68 +++++++++++++++++++++++++++ 6 files changed, 221 insertions(+), 4 deletions(-) create mode 100644 src/fastmcp/client/progress.py create mode 100644 tests/client/test_progress.py diff --git a/docs/clients/client.mdx b/docs/clients/client.mdx index 5c9607800..3b086ab5c 100644 --- a/docs/clients/client.mdx +++ b/docs/clients/client.mdx @@ -18,6 +18,17 @@ The FastMCP Client architecture separates the protocol logic (`Client`) from the - **`Client`**: Handles sending MCP requests (like `tools/call`, `resources/read`), receiving responses, and managing callbacks. - **`Transport`**: Responsible for establishing and maintaining the connection to the server (e.g., via WebSockets, SSE, Stdio, or in-memory). +```python +from fastmcp import Client, FastMCP +from fastmcp.client import ( + RootsHandler, + RootsList, + LogHandler, + MessageHandler, + SamplingHandler, + ProgressHandler # For handling progress notifications +) +``` ### Transports @@ -114,7 +125,7 @@ The standard client methods return user-friendly representations that may change tools = await client.list_tools() # tools -> list[mcp.types.Tool] ``` -* **`call_tool(name: str, arguments: dict[str, Any] | None = None, timeout: float | None = None)`**: Executes a tool on the server. +* **`call_tool(name: str, arguments: dict[str, Any] | None = None, timeout: float | None = None, progress_handler: ProgressHandler | None = None)`**: Executes a tool on the server. ```python result = await client.call_tool("add", {"a": 5, "b": 3}) # result -> list[mcp.types.TextContent | mcp.types.ImageContent | ...] @@ -122,10 +133,18 @@ The standard client methods return user-friendly representations that may change # With timeout (aborts if execution takes longer than 2 seconds) result = await client.call_tool("long_running_task", {"param": "value"}, timeout=2.0) + + # With progress handler (to track execution progress) + result = await client.call_tool( + "long_running_task", + {"param": "value"}, + progress_handler=my_progress_handler + ) ``` * Arguments are passed as a dictionary. FastMCP servers automatically handle JSON string parsing for complex types if needed. * Returns a list of content objects (usually `TextContent` or `ImageContent`). * The optional `timeout` parameter limits the maximum execution time (in seconds) for this specific call, overriding any client-level timeout. + * The optional `progress_handler` parameter receives progress updates during execution, overriding any client-level progress handler. #### Resource Operations @@ -234,6 +253,64 @@ Timeout behavior varies between transport types: For consistent behavior across all transports, we recommend explicitly setting timeouts at the individual tool call level when needed, rather than relying on client-level timeouts. +#### Progress Tracking + + + +MCP servers can report progress during long-running operations. The client can set a progress handler to receive and process these updates. + +```python +from fastmcp import Client +from fastmcp.client.progress import ProgressHandler + +# A simple progress handler that prints progress updates +async def my_progress_handler( + progress: float, + total: float | None, + message: str | None +) -> None: + """Handle progress updates from the server.""" + if total is not None: + percent = (progress / total) * 100 + print(f"Progress: {percent:.1f}% ({progress}/{total})") + else: + print(f"Progress: {progress}") + + if message: + print(f"Message: {message}") + +# Set the progress handler at client level +client = Client( + my_mcp_server, + progress_handler=my_progress_handler +) +``` + +By default, FastMCP uses a handler that logs progress updates at the debug level. This default handler properly handles cases where `total` or `message` might be None. + +You can override the progress handler for specific tool calls: + +```python +# Client uses the default debug logger for progress +client = Client(my_mcp_server) + +async with client: + # Use default progress handler (debug logging) + result1 = await client.call_tool("long_task", {"param": "value"}) + + # Override with custom progress handler just for this call + result2 = await client.call_tool( + "another_task", + {"param": "value"}, + progress_handler=my_progress_handler + ) +``` + +A typical progress update includes: +- Current progress value (e.g., 2 of 5 steps completed) +- Total expected value (may be None) +- Status message (may be None) + #### LLM Sampling MCP Servers can request LLM completions from clients. The client can provide a `sampling_handler` to handle these requests. The sampling handler receives a list of messages and other parameters from the server, and should return a string completion. diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index 320addb8c..803b5a205 100644 --- a/src/fastmcp/client/client.py +++ b/src/fastmcp/client/client.py @@ -8,7 +8,8 @@ from exceptiongroup import catch from mcp import ClientSession from pydantic import AnyUrl -from fastmcp.client.logging import LogHandler, MessageHandler +from fastmcp.client.logging import LogHandler, MessageHandler, default_log_handler +from fastmcp.client.progress import ProgressHandler, default_progress_handler from fastmcp.client.roots import ( RootsHandler, RootsList, @@ -28,6 +29,7 @@ __all__ = [ "LogHandler", "MessageHandler", "SamplingHandler", + "ProgressHandler", ] @@ -50,6 +52,7 @@ class Client: sampling_handler: Optional handler for sampling requests log_handler: Optional handler for log messages message_handler: Optional handler for protocol messages + progress_handler: Optional handler for progress notifications timeout: Optional timeout for requests (seconds or timedelta) Examples: @@ -74,6 +77,7 @@ class Client: sampling_handler: SamplingHandler | None = None, log_handler: LogHandler | None = None, message_handler: MessageHandler | None = None, + progress_handler: ProgressHandler | None = None, timeout: datetime.timedelta | float | int | None = None, ): self.transport = infer_transport(transport) @@ -81,6 +85,14 @@ class Client: self._exit_stack: AsyncExitStack | None = None self._nesting_counter: int = 0 + if log_handler is None: + log_handler = default_log_handler + + if progress_handler is None: + progress_handler = default_progress_handler + + self._progress_handler = progress_handler + if isinstance(timeout, int | float): timeout = datetime.timedelta(seconds=timeout) @@ -96,7 +108,9 @@ class Client: self.set_roots(roots) if sampling_handler is not None: - self.set_sampling_callback(sampling_handler) + self._session_kwargs["sampling_callback"] = create_sampling_callback( + sampling_handler + ) @property def session(self) -> ClientSession: @@ -433,6 +447,7 @@ class Client: self, name: str, arguments: dict[str, Any], + progress_handler: ProgressHandler | None = None, timeout: datetime.timedelta | float | int | None = None, ) -> mcp.types.CallToolResult: """Send a tools/call request and return the complete MCP protocol result. @@ -444,6 +459,8 @@ class Client: name (str): The name of the tool to call. arguments (dict[str, Any]): Arguments to pass to the tool. timeout (datetime.timedelta | float | int | None, optional): The timeout for the tool call. Defaults to None. + progress_handler (ProgressHandler | None, optional): The progress handler to use for the tool call. Defaults to None. + Returns: mcp.types.CallToolResult: The complete response object from the protocol, containing the tool result and any additional metadata. @@ -455,7 +472,10 @@ class Client: if isinstance(timeout, int | float): timeout = datetime.timedelta(seconds=timeout) result = await self.session.call_tool( - name=name, arguments=arguments, read_timeout_seconds=timeout + name=name, + arguments=arguments, + read_timeout_seconds=timeout, + progress_callback=progress_handler or self._progress_handler, ) return result @@ -464,6 +484,7 @@ class Client: name: str, arguments: dict[str, Any] | None = None, timeout: datetime.timedelta | float | int | None = None, + progress_handler: ProgressHandler | None = None, ) -> list[ mcp.types.TextContent | mcp.types.ImageContent | mcp.types.EmbeddedResource ]: @@ -474,6 +495,8 @@ class Client: Args: name (str): The name of the tool to call. arguments (dict[str, Any] | None, optional): Arguments to pass to the tool. Defaults to None. + timeout (datetime.timedelta | float | int | None, optional): The timeout for the tool call. Defaults to None. + progress_handler (ProgressHandler | None, optional): The progress handler to use for the tool call. Defaults to None. Returns: list[mcp.types.TextContent | mcp.types.ImageContent | mcp.types.EmbeddedResource]: @@ -487,6 +510,7 @@ class Client: name=name, arguments=arguments or {}, timeout=timeout, + progress_handler=progress_handler, ) if result.isError: msg = cast(mcp.types.TextContent, result.content[0]).text diff --git a/src/fastmcp/client/logging.py b/src/fastmcp/client/logging.py index 2047633e3..826eb0d28 100644 --- a/src/fastmcp/client/logging.py +++ b/src/fastmcp/client/logging.py @@ -6,8 +6,16 @@ from mcp.client.session import ( ) from mcp.types import LoggingMessageNotificationParams +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) + LogMessage: TypeAlias = LoggingMessageNotificationParams LogHandler: TypeAlias = LoggingFnT MessageHandler: TypeAlias = MessageHandlerFnT __all__ = ["LogMessage", "LogHandler", "MessageHandler"] + + +async def default_log_handler(params: LogMessage) -> None: + logger.debug(f"Log received: {params}") diff --git a/src/fastmcp/client/progress.py b/src/fastmcp/client/progress.py new file mode 100644 index 000000000..2470f18ce --- /dev/null +++ b/src/fastmcp/client/progress.py @@ -0,0 +1,38 @@ +from typing import TypeAlias + +from mcp.shared.session import ProgressFnT + +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) + +ProgressHandler: TypeAlias = ProgressFnT + + +async def default_progress_handler( + progress: float, total: float | None, message: str | None +) -> None: + """Default handler for progress notifications. + + Logs progress updates at debug level, properly handling missing total or message values. + + Args: + progress: Current progress value + total: Optional total expected value + message: Optional status message + """ + if total is not None: + # We have both progress and total + percent = (progress / total) * 100 + progress_str = f"{progress}/{total} ({percent:.1f}%)" + else: + # We only have progress + progress_str = f"{progress}" + + # Include message if available + if message: + log_msg = f"Progress: {progress_str} - {message}" + else: + log_msg = f"Progress: {progress_str}" + + logger.debug(log_msg) diff --git a/src/fastmcp/client/transports.py b/src/fastmcp/client/transports.py index 7faeab613..e1597bec8 100644 --- a/src/fastmcp/client/transports.py +++ b/src/fastmcp/client/transports.py @@ -22,6 +22,7 @@ from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamablehttp_client from mcp.client.websocket import websocket_client from mcp.shared.memory import create_connected_server_and_client_session +from mcp.shared.session import ProgressFnT from pydantic import AnyUrl from typing_extensions import Unpack @@ -35,6 +36,7 @@ class SessionKwargs(TypedDict, total=False): list_roots_callback: ListRootsFnT | None logging_callback: LoggingFnT | None message_handler: MessageHandlerFnT | None + progress_callback: ProgressFnT | None read_timeout_seconds: datetime.timedelta | None diff --git a/tests/client/test_progress.py b/tests/client/test_progress.py new file mode 100644 index 000000000..5f61c3487 --- /dev/null +++ b/tests/client/test_progress.py @@ -0,0 +1,68 @@ +import pytest + +from fastmcp import Client, Context, FastMCP + +PROGRESS_MESSAGES = [] + + +@pytest.fixture(autouse=True) +def clear_progress_messages(): + PROGRESS_MESSAGES.clear() + yield + PROGRESS_MESSAGES.clear() + + +@pytest.fixture +def fastmcp_server(): + mcp = FastMCP() + + @mcp.tool() + async def progress_tool(context: Context) -> int: + for i in range(3): + await context.report_progress( + progress=i + 1, + total=3, + message=f"{(i + 1) / 3 * 100:.2f}% complete", + ) + return 100 + + return mcp + + +EXPECTED_PROGRESS_MESSAGES = [ + dict(progress=1, total=3, message="33.33% complete"), + dict(progress=2, total=3, message="66.67% complete"), + dict(progress=3, total=3, message="100.00% complete"), +] + + +async def progress_handler( + progress: float, total: float | None, message: str | None +) -> None: + PROGRESS_MESSAGES.append(dict(progress=progress, total=total, message=message)) + + +async def test_progress_handler(fastmcp_server: FastMCP): + async with Client(fastmcp_server, progress_handler=progress_handler) as client: + await client.call_tool("progress_tool", {}) + + assert PROGRESS_MESSAGES == EXPECTED_PROGRESS_MESSAGES + + +async def test_progress_handler_can_be_supplied_on_tool_call(fastmcp_server: FastMCP): + async with Client(fastmcp_server) as client: + await client.call_tool("progress_tool", {}, progress_handler=progress_handler) + + assert PROGRESS_MESSAGES == EXPECTED_PROGRESS_MESSAGES + + +async def test_progress_handler_supplied_on_tool_call_overrides_default( + fastmcp_server: FastMCP, +): + async def bad_progress_handler(*args, **kwargs): + 1 / 0 + + async with Client(fastmcp_server, progress_handler=bad_progress_handler) as client: + await client.call_tool("progress_tool", {}, progress_handler=progress_handler) + + assert PROGRESS_MESSAGES == EXPECTED_PROGRESS_MESSAGES From 5a89e0cc2b6c19a734db426fde79a17f80d821cc Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 20 May 2025 09:36:00 -0400 Subject: [PATCH 2/2] Fix typing --- src/fastmcp/client/transports.py | 2 -- tests/client/test_progress.py | 6 ++++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/fastmcp/client/transports.py b/src/fastmcp/client/transports.py index e1597bec8..7faeab613 100644 --- a/src/fastmcp/client/transports.py +++ b/src/fastmcp/client/transports.py @@ -22,7 +22,6 @@ from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamablehttp_client from mcp.client.websocket import websocket_client from mcp.shared.memory import create_connected_server_and_client_session -from mcp.shared.session import ProgressFnT from pydantic import AnyUrl from typing_extensions import Unpack @@ -36,7 +35,6 @@ class SessionKwargs(TypedDict, total=False): list_roots_callback: ListRootsFnT | None logging_callback: LoggingFnT | None message_handler: MessageHandlerFnT | None - progress_callback: ProgressFnT | None read_timeout_seconds: datetime.timedelta | None diff --git a/tests/client/test_progress.py b/tests/client/test_progress.py index 5f61c3487..f67f5c54c 100644 --- a/tests/client/test_progress.py +++ b/tests/client/test_progress.py @@ -59,8 +59,10 @@ async def test_progress_handler_can_be_supplied_on_tool_call(fastmcp_server: Fas async def test_progress_handler_supplied_on_tool_call_overrides_default( fastmcp_server: FastMCP, ): - async def bad_progress_handler(*args, **kwargs): - 1 / 0 + async def bad_progress_handler( + progress: float, total: float | None, message: str | None + ) -> None: + raise Exception("This should not be called") async with Client(fastmcp_server, progress_handler=bad_progress_handler) as client: await client.call_tool("progress_tool", {}, progress_handler=progress_handler)