moved core logic for progress call back to BaseSession

Author: davemssavage

src/mcp/client/session.py

@@ -8,7 +8,7 @@
 import mcp.types as types
 from mcp.shared.context import RequestContext
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import BaseSession, RequestResponder
+from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
 DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
@@ -35,13 +35,6 @@ async def __call__(
     ) -> None: ...
 
 
-class ProgressFnT(Protocol):
-    async def __call__(
-        self,
-        params: types.ProgressNotificationParams,
-    ) -> None: ...
-
-
 class MessageHandlerFnT(Protocol):
     async def __call__(
         self,
@@ -98,9 +91,6 @@ class ClientSession(
         types.ServerNotification,
     ]
 ):
-    _progress_id: int
-    _in_progress: dict[types.ProgressToken, ProgressFnT]
-
     def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -124,8 +114,6 @@ def __init__(
         self._list_roots_callback = list_roots_callback or _default_list_roots_callback
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
-        self._progress_id = 0
-        self._in_progress = {}
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability()
@@ -274,34 +262,17 @@ async def call_tool(
         progress_callback: ProgressFnT | None = None,
     ) -> types.CallToolResult:
         """Send a tools/call request."""
-
-        if progress_callback is None:
-            progress_id = None
-            call_params = types.CallToolRequestParams(name=name, arguments=arguments)
-        else:
-            progress_id = self._progress_id
-            self._progress_id = progress_id + 1
-
-            call_meta = types.RequestParams.Meta(progressToken=progress_id)
-            call_params = types.CallToolRequestParams(
-                name=name, arguments=arguments, _meta=call_meta
-            )
-            self._in_progress[progress_id] = progress_callback
-
-        try:
-            return await self.send_request(
-                types.ClientRequest(
-                    types.CallToolRequest(
-                        method="tools/call",
-                        params=call_params,
-                    )
-                ),
-                types.CallToolResult,
-                request_read_timeout_seconds=read_timeout_seconds,
-            )
-        finally:
-            if progress_id is not None:
-                self._in_progress.pop(progress_id, None)
+        return await self.send_request(
+            types.ClientRequest(
+                types.CallToolRequest(
+                    method="tools/call",
+                    params=types.CallToolRequestParams(name=name, arguments=arguments),
+                )
+            ),
+            types.CallToolResult,
+            request_read_timeout_seconds=read_timeout_seconds,
+            progress_callback=progress_callback,
+        )
 
     async def list_prompts(self) -> types.ListPromptsResult:
         """Send a prompts/list request."""
@@ -414,9 +385,5 @@ async def _received_notification(
         match notification.root:
             case types.LoggingMessageNotification(params=params):
                 await self._logging_callback(params)
-            case types.ProgressNotification(params=params):
-                if params.progressToken in self._in_progress:
-                    progress_callback = self._in_progress[params.progressToken]
-                    await progress_callback(params)
             case _:
                 pass

src/mcp/shared/session.py

@@ -3,7 +3,7 @@
 from contextlib import AsyncExitStack
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Protocol, TypeVar
 
 import anyio
 import httpx
@@ -24,6 +24,9 @@
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    ProgressNotification,
+    ProgressNotificationParams,
+    ProgressToken,
     RequestParams,
     ServerNotification,
     ServerRequest,
@@ -39,6 +42,14 @@
     "ReceiveNotificationT", ClientNotification, ServerNotification
 )
 
+
+class ProgressFnT(Protocol):
+    async def __call__(
+        self,
+        params: ProgressNotificationParams,
+    ) -> None: ...
+
+
 RequestId = str | int
 
 
@@ -168,7 +179,9 @@ class BaseSession(
         RequestId, MemoryObjectSendStream[JSONRPCResponse | JSONRPCError]
     ]
     _request_id: int
+    _progress_id: int
     _in_flight: dict[RequestId, RequestResponder[ReceiveRequestT, SendResultT]]
+    _in_progress: dict[ProgressToken, ProgressFnT]
 
     def __init__(
         self,
@@ -187,6 +200,8 @@ def __init__(
         self._receive_notification_type = receive_notification_type
         self._session_read_timeout_seconds = read_timeout_seconds
         self._in_flight = {}
+        self._progress_id = 0
+        self._in_progress = {}
         self._exit_stack = AsyncExitStack()
 
     async def __aenter__(self) -> Self:
@@ -214,19 +229,44 @@ async def send_request(
         result_type: type[ReceiveResultT],
         request_read_timeout_seconds: timedelta | None = None,
         metadata: MessageMetadata = None,
+        progress_callback: ProgressFnT | None = None,
     ) -> ReceiveResultT:
         """
         Sends a request and wait for a response. Raises an McpError if the
         response contains an error. If a request read timeout is provided, it
         will take precedence over the session read timeout.
 
+        If progress_callback is provided any progress notifications sent from the
+        receiver will be passed back to the sender
+
         Do not use this method to emit notifications! Use send_notification()
         instead.
         """
 
         request_id = self._request_id
         self._request_id = request_id + 1
 
+        progress_id = None
+        send_request = None
+
+        if progress_callback is not None:
+            if request.root.params is not None:
+                progress_id = self._progress_id
+                self._progress_id = progress_id + 1
+                new_params = request.root.params.model_copy(
+                    update={"meta": RequestParams.Meta(progressToken=progress_id)}
+                )
+                new_root = request.root.model_copy(update={"params": new_params})
+                send_request = request.model_copy(update={"root": new_root})
+                self._in_progress[progress_id] = progress_callback
+            else:
+                raise ValueError(
+                    f"{type(request.root).__name__} does not support progress"
+                )
+
+        if send_request is None:
+            send_request = request
+
         response_stream, response_stream_reader = anyio.create_memory_object_stream[
             JSONRPCResponse | JSONRPCError
         ](1)
@@ -236,11 +276,11 @@ async def send_request(
             jsonrpc_request = JSONRPCRequest(
                 jsonrpc="2.0",
                 id=request_id,
-                **request.model_dump(by_alias=True, mode="json", exclude_none=True),
+                **send_request.model_dump(
+                    by_alias=True, mode="json", exclude_none=True
+                ),
             )
 
-            # TODO: Support progress callbacks
-
             await self._write_stream.send(
                 SessionMessage(
                     message=JSONRPCMessage(jsonrpc_request), metadata=metadata
@@ -276,6 +316,8 @@ async def send_request(
 
         finally:
             self._response_streams.pop(request_id, None)
+            if progress_id is not None:
+                self._in_progress.pop(progress_id, None)
             await response_stream.aclose()
             await response_stream_reader.aclose()
 
@@ -364,6 +406,20 @@ async def _receive_loop(self) -> None:
                             if cancelled_id in self._in_flight:
                                 await self._in_flight[cancelled_id].cancel()
                         else:
+                            match notification.root:
+                                case ProgressNotification(params=params):
+                                    if params.progressToken in self._in_progress:
+                                        progress_callback = self._in_progress[
+                                            params.progressToken
+                                        ]
+                                        await progress_callback(params)
+                                    else:
+                                        logging.warning(
+                                            "Unknown progress token %s",
+                                            params.progressToken,
+                                        )
+                                case _:
+                                    pass
                             await self._received_notification(notification)
                             await self._handle_incoming(notification)
                     except Exception as e: