refactor args for clearer meaning, use error vs returning none on timeout

Author: davemssavage

src/mcp/client/session.py

@@ -307,14 +307,14 @@ async def join_call_tool(
         request_id: types.RequestId,
         progress_callback: ProgressFnT | None = None,
         request_read_timeout_seconds: timedelta | None = None,
-        fail_on_timeout: bool = True,
+        done_on_timeout: bool = True,
     ) -> types.CallToolResult | None:
         return await self.join_request(
             request_id,
             types.CallToolResult,
             request_read_timeout_seconds=request_read_timeout_seconds,
             progress_callback=progress_callback,
-            fail_on_timeout=fail_on_timeout,
+            done_on_timeout=done_on_timeout,
         )
 
     async def cancel_call_tool(

src/mcp/shared/session.py

@@ -176,8 +176,10 @@ async def send_progress(
     ): ...
 
     async def receive_response(
-        self, request_id: RequestId, timeout: float | None = None, fail_on_timeout: bool = True
-    ) -> JSONRPCResponse | JSONRPCError | None: ...
+        self,
+        request_id: RequestId,
+        timeout: float | None = None,
+    ) -> JSONRPCResponse | JSONRPCError: ...
 
     async def handle_response(self, message: JSONRPCResponse | JSONRPCError) -> bool: ...
 
@@ -242,8 +244,10 @@ async def send_progress(
                 )
 
     async def receive_response(
-        self, request_id: RequestId, timeout: float | None = None, fail_on_timeout: bool = True
-    ) -> JSONRPCResponse | JSONRPCError | None:
+        self,
+        request_id: RequestId,
+        timeout: float | None = None,
+    ) -> JSONRPCResponse | JSONRPCError:
         request, _, response_stream_reader = self._response_streams.get(request_id, [None, None, None])
 
         if response_stream_reader is None:
@@ -254,24 +258,20 @@ async def receive_response(
                 )
             )
 
-        if fail_on_timeout:
-            try:
-                with anyio.fail_after(timeout):
-                    return await response_stream_reader.receive()
-            except TimeoutError:
-                raise McpError(
-                    ErrorData(
-                        code=httpx.codes.REQUEST_TIMEOUT,
-                        message=(
-                            f"Timed out while waiting for response to "
-                            f"{request.__class__.__name__}. Waited "
-                            f"{timeout} seconds."
-                        ),
-                    )
-                )
-        else:
-            with anyio.move_on_after(timeout):
+        try:
+            with anyio.fail_after(timeout):
                 return await response_stream_reader.receive()
+        except TimeoutError:
+            raise McpError(
+                ErrorData(
+                    code=httpx.codes.REQUEST_TIMEOUT,
+                    message=(
+                        f"Timed out while waiting for response to "
+                        f"{request.__class__.__name__}. Waited "
+                        f"{timeout} seconds."
+                    ),
+                )
+            )
 
     async def handle_response(self, message: JSONRPCResponse | JSONRPCError) -> bool:
         _, stream, _ = self._response_streams.get(message.id, [None, None, None])
@@ -405,8 +405,8 @@ async def join_request(
         result_type: type[ReceiveResultT],
         request_read_timeout_seconds: timedelta | None = None,
         progress_callback: ProgressFnT | None = None,
-        fail_on_timeout: bool = True,
-    ) -> ReceiveResultT | None:
+        done_on_timeout: bool = True,
+    ) -> ReceiveResultT:
         """
         Joins a request previously started via start_request
         """
@@ -420,16 +420,18 @@ async def join_request(
         elif self._session_read_timeout_seconds is not None:
             timeout = self._session_read_timeout_seconds.total_seconds()
 
-        response_or_error = await self._request_state_manager.receive_response(request_id, timeout, fail_on_timeout)
+        response_or_error = await self._request_state_manager.receive_response(request_id, timeout)
 
-        if response_or_error is None:
-            return None
+        if isinstance(response_or_error, JSONRPCError):
+            if response_or_error.error.code == httpx.codes.REQUEST_TIMEOUT.value:
+                if done_on_timeout:
+                    await self._request_state_manager.close_request(request_id)
+            else:
+                await self._request_state_manager.close_request(request_id)
+            raise McpError(response_or_error.error)
         else:
             await self._request_state_manager.close_request(request_id)
-            if isinstance(response_or_error, JSONRPCError):
-                raise McpError(response_or_error.error)
-            else:
-                return result_type.model_validate(response_or_error.result)
+            return result_type.model_validate(response_or_error.result)
 
     async def cancel_request(self, request_id: RequestId) -> bool:
         """
@@ -464,10 +466,10 @@ async def send_request(
         instead.
         """
         request_id = await self.start_request(request, metadata, progress_callback)
-        result = await self.join_request(request_id, result_type, request_read_timeout_seconds)
-        if result is None:
-            raise RuntimeError("Should not be possible")
-        return result
+        try:
+            return await self.join_request(request_id, result_type, request_read_timeout_seconds)
+        finally:
+            await self._request_state_manager.close_request(request_id)
 
     async def send_notification(
         self,

tests/client/test_session.py

@@ -2,11 +2,13 @@
 from typing import Any
 
 import anyio
+import httpx
 import pytest
 
 import mcp.types as types
 from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
 from mcp.shared.context import RequestContext
+from mcp.shared.exceptions import McpError
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import ImMemoryRequestStateManager, RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
@@ -636,13 +638,17 @@ async def message_handler(
         request_id = await session.request_call_tool("hello", {"name": "world"})
 
         with anyio.fail_after(1):
-            result = await session.join_call_tool(
-                request_id, request_read_timeout_seconds=timedelta(microseconds=1), fail_on_timeout=False
-            )
-            assert result is None
+            try:
+                result = await session.join_call_tool(
+                    request_id, request_read_timeout_seconds=timedelta(microseconds=1), done_on_timeout=False
+                )
+            except McpError as e:
+                if not e.error.code == httpx.codes.REQUEST_TIMEOUT:
+                    raise e
+
             send_result.set()
             result = await session.join_call_tool(
-                request_id, request_read_timeout_seconds=timedelta(seconds=1), fail_on_timeout=False
+                request_id, request_read_timeout_seconds=timedelta(seconds=1), done_on_timeout=False
             )
 
             # Assert the result