suggested changes

Author: ihrpr

src/mcp/client/session.py

@@ -267,7 +267,29 @@ async def call_tool(
         on_resumption_token_update: ResumptionTokenUpdateCallback | None = None,
         resumption_token: ResumptionToken | None = None,
     ) -> types.CallToolResult:
-        """Send a tools/call request."""
+        """Send a tools/call request to invoke or resume a tool operation.
+
+        This method supports two modes of operation:
+
+        Initial Invocation:
+          - Tool is executed from the beginning
+          - If on_resumption_token_update is provided, it will receive resumption tokens
+            if transport supports resumable operations
+
+        Resumption:
+          - Requires a resumption_token from a previous execution
+          - Continues execution from the specified resumption token
+
+        Args:
+            name: Tool name to call
+            arguments: Optional arguments dictionary for the tool
+            read_timeout_seconds: request read timeout
+            on_resumption_token_update: Optional callback that receives
+                resumption tokens during execution, which can be
+                used later to resume an interrupted operation
+            resumption_token: Optional token to resume from a specific token
+                (supported by some transports)
+        """
         metadata = None
         if on_resumption_token_update or resumption_token:
             metadata = ClientMessageMetadata(

src/mcp/client/streamable_http.py

@@ -16,7 +16,7 @@
 import anyio
 import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from httpx_sse import EventSource, aconnect_sse
+from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
 
 from mcp.shared.message import ClientMessageMetadata, SessionMessage
 from mcp.types import (
@@ -26,13 +26,14 @@
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    RequestId,
 )
 
 logger = logging.getLogger(__name__)
 
 
-MessageOrError = SessionMessage | Exception
-StreamWriter = MemoryObjectSendStream[MessageOrError]
+SessionMessageOrError = SessionMessage | Exception
+StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
 StreamReader = MemoryObjectReceiveStream[SessionMessage]
 
 
@@ -123,23 +124,21 @@ def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
             and message.root.method == "notifications/initialized"
         )
 
-    def _extract_session_id_from_response(
+    def _maybe_extract_session_id_from_response(
         self,
         response: httpx.Response,
-        is_initialization: bool,
     ) -> None:
         """Extract and store session ID from response headers."""
-        if is_initialization:
-            new_session_id = response.headers.get(MCP_SESSION_ID)
-            if new_session_id:
-                self.session_id = new_session_id
-                logger.info(f"Received session ID: {self.session_id}")
+        new_session_id = response.headers.get(MCP_SESSION_ID)
+        if new_session_id:
+            self.session_id = new_session_id
+            logger.info(f"Received session ID: {self.session_id}")
 
     async def _handle_sse_event(
         self,
-        sse: Any,
+        sse: ServerSentEvent,
         read_stream_writer: StreamWriter,
-        original_request_id: Any | None = None,
+        original_request_id: RequestId | None = None,
         resumption_callback: Callable[[str], Awaitable[None]] | None = None,
     ) -> bool:
         """Handle an SSE event, returning True if the response is complete."""
@@ -161,7 +160,8 @@ async def _handle_sse_event(
                 if sse.id and resumption_callback:
                     await resumption_callback(sse.id)
 
-                # If this is a response or error, we're done
+                # If this is a response or error return True indicating completion
+                # Otherwise, return False to continue listening
                 return isinstance(message.root, JSONRPCResponse | JSONRPCError)
 
             except Exception as exc:
@@ -262,7 +262,8 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                 return
 
             response.raise_for_status()
-            self._extract_session_id_from_response(response, is_initialization)
+            if is_initialization:
+                self._maybe_extract_session_id_from_response(response)
 
             content_type = response.headers.get(CONTENT_TYPE, "").lower()
 
@@ -324,7 +325,7 @@ async def _handle_unexpected_content_type(
     async def _send_session_terminated_error(
         self,
         read_stream_writer: StreamWriter,
-        request_id: Any,
+        request_id: RequestId,
     ) -> None:
         """Send a session terminated error response."""
         jsonrpc_error = JSONRPCError(