test: cover EOF drain branches

Author: Epochex

src/mcp/shared/jsonrpc_dispatcher.py

@@ -445,7 +445,9 @@ async def run(
                         # On normal EOF, let already-received handlers drain
                         # their responses before the task group exits.
                         tg.cancel_scope.cancel()
-                    elif self._read_eof_drain_timeout_seconds is not None:
+                    elif self._read_eof_drain_timeout_seconds is None:
+                        pass
+                    else:
                         tg.cancel_scope.deadline = anyio.current_time() + self._read_eof_drain_timeout_seconds
         finally:
             # Covers the cancel/crash paths where the inline fan-out above is

tests/shared/test_jsonrpc_dispatcher.py

@@ -249,6 +249,62 @@ async def drive() -> None:
     s2c_recv.close()
 
 
+@pytest.mark.anyio
+async def test_run_closes_write_stream_after_clean_eof_without_drain_timeout():
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
+        c2s_recv,
+        s2c_send,
+        close_write_stream_on_read_close=False,
+        read_eof_drain_timeout_seconds=None,
+    )
+    on_request, on_notify = echo_handlers(Recorder())
+
+    with anyio.fail_after(5):
+        async with anyio.create_task_group() as tg, c2s_send, c2s_recv, s2c_send, s2c_recv:
+            await tg.start(server.run, on_request, on_notify)
+            c2s_send.close()
+            with pytest.raises(anyio.EndOfStream):
+                await s2c_recv.receive()
+
+
+@pytest.mark.anyio
+async def test_run_drains_in_flight_handlers_on_clean_eof_without_timeout():
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
+        c2s_recv,
+        s2c_send,
+        close_write_stream_on_read_close=False,
+        read_eof_drain_timeout_seconds=None,
+    )
+    handler_started = anyio.Event()
+    handler_allowed_to_finish = anyio.Event()
+
+    async def handle_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        handler_started.set()
+        await handler_allowed_to_finish.wait()
+        return {"drained": True}
+
+    async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
+        raise NotImplementedError
+
+    with anyio.fail_after(5):
+        async with anyio.create_task_group() as tg, c2s_send, c2s_recv, s2c_send, s2c_recv:
+            await tg.start(server.run, handle_request, on_notify)
+            await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="x", params=None)))
+            await handler_started.wait()
+            c2s_send.close()
+            handler_allowed_to_finish.set()
+
+            response = await s2c_recv.receive()
+            assert isinstance(response, SessionMessage)
+            assert isinstance(response.message, JSONRPCResponse)
+            assert response.message.id == 1
+            assert response.message.result == {"drained": True}
+
+
 @pytest.mark.anyio
 async def test_run_closes_write_stream_on_exit():
     """run() enters both streams; the write end is released on EOF."""

tests/shared/test_session.py

@@ -1,13 +1,16 @@
+from typing import Any
+
 import anyio
 import pytest
+from pydantic import TypeAdapter
 
 from mcp import Client, types
 from mcp.client.session import ClientSession
 from mcp.server import Server, ServerRequestContext
 from mcp.shared.exceptions import MCPError
 from mcp.shared.memory import create_client_server_memory_streams
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
+from mcp.shared.session import BaseSession, RequestId, RequestResponder
 from mcp.types import (
     PARSE_ERROR,
     CancelledNotification,
@@ -16,6 +19,7 @@
     EmptyResult,
     ErrorData,
     JSONRPCError,
+    JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
     ServerNotification,
@@ -291,6 +295,40 @@ async def mock_server():
                 await ev_response.wait()
 
 
+@pytest.mark.anyio
+async def test_receive_loop_can_leave_write_stream_open_on_read_eof():
+    class TestSession(BaseSession[Any, Any, Any, Any, Any]):
+        async def _send_response(self, request_id: RequestId, response: Any | ErrorData) -> None:
+            raise NotImplementedError  # pragma: no cover
+
+        @property
+        def _receive_request_adapter(self) -> TypeAdapter[Any]:
+            return TypeAdapter(object)  # pragma: no cover
+
+        @property
+        def _receive_notification_adapter(self) -> TypeAdapter[Any]:
+            return TypeAdapter(object)  # pragma: no cover
+
+    read_send, read_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+    write_send, write_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    session = TestSession(read_receive, write_send, close_write_stream_on_read_close=False)
+    receive_loop_returned = anyio.Event()
+
+    async def drive_receive_loop() -> None:
+        await session._receive_loop()  # pyright: ignore[reportPrivateUsage]
+        receive_loop_returned.set()
+
+    with anyio.fail_after(5):
+        async with anyio.create_task_group() as tg, read_send, read_receive, write_send, write_receive:
+            tg.start_soon(drive_receive_loop)
+            await read_send.aclose()
+            await receive_loop_returned.wait()
+
+            marker = SessionMessage(message=JSONRPCNotification(jsonrpc="2.0", method="still-open"))
+            await write_send.send(marker)
+            assert await write_receive.receive() is marker
+
+
 @pytest.mark.anyio
 async def test_null_id_error_surfaced_via_message_handler():
     """Test that a JSONRPCError with id=None is surfaced to the message handler.