cleanup from review comments

Author: nick-merrill

mcp_python/client/session.py

@@ -1,3 +1,5 @@
+from datetime import timedelta
+
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import AnyUrl
 
@@ -36,7 +38,7 @@ def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
         write_stream: MemoryObjectSendStream[JSONRPCMessage],
-        read_timeout_seconds: int | float | None = None,
+        read_timeout_seconds: timedelta | None = None,
     ) -> None:
         super().__init__(
             read_stream,

mcp_python/server/memory.py

@@ -3,6 +3,7 @@
 """
 
 from contextlib import asynccontextmanager
+from datetime import timedelta
 from typing import AsyncGenerator
 
 import anyio
@@ -52,10 +53,10 @@ async def create_client_server_memory_streams() -> AsyncGenerator[
 @asynccontextmanager
 async def create_connected_server_and_client_session(
     server: Server,
-    read_timeout_seconds: int | float | None = None,
+    read_timeout_seconds: timedelta | None = None,
     raise_exceptions: bool = False,
 ) -> AsyncGenerator[ClientSession, None]:
-    """Creates a ServerSession that is connected to the `server`."""
+    """Creates a ClientSession that is connected to a running MCP server."""
     async with create_client_server_memory_streams() as (
         client_streams,
         server_streams,
@@ -75,9 +76,6 @@ async def create_connected_server_and_client_session(
             )
 
             try:
-                # Client session could be created here using client_read and
-                # client_write This would allow testing the server with a client
-                # in the same process
                 async with ClientSession(
                     read_stream=client_read, write_stream=client_write,
                     read_timeout_seconds=read_timeout_seconds,

mcp_python/shared/session.py

@@ -1,5 +1,7 @@
 from contextlib import AbstractAsyncContextManager
+from datetime import timedelta
 from typing import Generic, TypeVar
+import httpx
 
 import anyio
 import anyio.lowlevel
@@ -88,7 +90,7 @@ def __init__(
         receive_request_type: type[ReceiveRequestT],
         receive_notification_type: type[ReceiveNotificationT],
         # If none, reading will never time out
-        read_timeout_seconds: int | float | None = None,
+        read_timeout_seconds: timedelta | None = None,
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
@@ -151,12 +153,15 @@ async def send_request(
         await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
 
         try:
-            with anyio.fail_after(self._read_timeout_seconds):
+            with anyio.fail_after(
+                None if self._read_timeout_seconds is None
+                else self._read_timeout_seconds.total_seconds()
+            ):
                 response_or_error = await response_stream_reader.receive()
         except TimeoutError:
             raise McpError(
                 ErrorData(
-                    code=408,
+                    code=httpx.codes.REQUEST_TIMEOUT,
                     message=(
                         f"Timed out while waiting for response to "
                         f"{request.__class__.__name__}. Waited "

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mcp-python"
-version = "0.4.0.dev"
+version = "0.3.1.dev"
 description = "Model Context Protocol implementation for Python"
 readme = "README.md"
 requires-python = ">=3.10"