Fix stdio client shutdown bugs and rebuild the stdio test suite (#2773)

tests/server/mcpserver/test_elicitation.py

@@ -1,4 +1,4 @@
-"""Test the elicitation feature using stdio transport."""
+"""Test the elicitation feature over the in-memory client transport."""
 
 from typing import Any
 
@@ -58,9 +58,9 @@ async def call_tool_and_assert(
 
 
 @pytest.mark.anyio
-async def test_stdio_elicitation():
-    """Test the elicitation feature using stdio transport."""
-    mcp = MCPServer(name="StdioElicitationServer")
+async def test_elicitation_accept_returns_the_users_answer_to_the_tool():
+    """An accepted elicitation delivers the user's content back to the requesting tool."""
+    mcp = MCPServer(name="ElicitationServer")
     create_ask_user_tool(mcp)
 
     # Create a custom handler for elicitation requests
@@ -76,9 +76,9 @@ async def test_stdio_elicitation():
 
 
 @pytest.mark.anyio
-async def test_stdio_elicitation_decline():
-    """Test elicitation with user declining."""
-    mcp = MCPServer(name="StdioElicitationDeclineServer")
+async def test_elicitation_decline_reaches_the_tool_without_content():
+    """A declined elicitation reports the decline to the tool, with no content attached."""
+    mcp = MCPServer(name="ElicitationDeclineServer")
     create_ask_user_tool(mcp)
 
     async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):

tests/server/test_stdio.py

@@ -1,17 +1,26 @@
 import io
 import sys
+import threading
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from io import TextIOWrapper
 
 import anyio
 import pytest
 
+from mcp.server.mcpserver import MCPServer
 from mcp.server.stdio import stdio_server
 from mcp.shared.message import SessionMessage
 from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse, jsonrpc_message_adapter
 
 
 @pytest.mark.anyio
-async def test_stdio_server():
+async def test_stdio_server_round_trips_messages_over_injected_streams() -> None:
+    """stdio_server frames JSON-RPC messages as one line each in both directions.
+
+    Parses one message per stdin line and writes each outgoing message as exactly one
+    line, driven over injected in-process streams.
+    """
     stdin = io.StringIO()
     stdout = io.StringIO()
 
@@ -24,52 +33,45 @@ async def test_stdio_server():
         stdin.write(message.model_dump_json(by_alias=True, exclude_none=True) + "\n")
     stdin.seek(0)
 
-    async with stdio_server(stdin=anyio.AsyncFile(stdin), stdout=anyio.AsyncFile(stdout)) as (
-        read_stream,
-        write_stream,
-    ):
-        received_messages: list[JSONRPCMessage] = []
-        async with read_stream:
-            async for message in read_stream:
-                if isinstance(message, Exception):  # pragma: no cover
-                    raise message
-                received_messages.append(message.message)
-                if len(received_messages) == 2:
-                    break
+    with anyio.fail_after(5):
+        async with stdio_server(stdin=anyio.AsyncFile(stdin), stdout=anyio.AsyncFile(stdout)) as (
+            read_stream,
+            write_stream,
+        ):
+            async with read_stream:
+                received_messages: list[JSONRPCMessage] = []
+                for _ in range(2):
+                    received = await read_stream.receive()
+                    assert not isinstance(received, Exception)
+                    received_messages.append(received.message)
 
-        # Verify received messages
-        assert len(received_messages) == 2
-        assert received_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
-        assert received_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
+            assert received_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+            assert received_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
 
-        # Test sending responses from the server
-        responses = [
-            JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"),
-            JSONRPCResponse(jsonrpc="2.0", id=4, result={}),
-        ]
+            responses = [
+                JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"),
+                JSONRPCResponse(jsonrpc="2.0", id=4, result={}),
+            ]
 
-        async with write_stream:
             for response in responses:
-                session_message = SessionMessage(response)
-                await write_stream.send(session_message)
+                await write_stream.send(SessionMessage(response))
+            await write_stream.aclose()
 
     stdout.seek(0)
     output_lines = stdout.readlines()
     assert len(output_lines) == 2
 
     received_responses = [jsonrpc_message_adapter.validate_json(line.strip()) for line in output_lines]
-    assert len(received_responses) == 2
     assert received_responses[0] == JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
     assert received_responses[1] == JSONRPCResponse(jsonrpc="2.0", id=4, result={})
 
 
 @pytest.mark.anyio
-async def test_stdio_server_invalid_utf8(monkeypatch: pytest.MonkeyPatch):
-    """Non-UTF-8 bytes on stdin must not crash the server.
+async def test_stdio_server_invalid_utf8(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Non-UTF-8 stdin bytes surface as an in-stream exception without killing the stream.
 
-    Invalid bytes are replaced with U+FFFD, which then fails JSON parsing and
-    is delivered as an in-stream exception. Subsequent valid messages must
-    still be processed.
+    Invalid bytes are replaced with U+FFFD, fail JSON parsing, and arrive as an in-stream
+    exception; subsequent valid messages are still processed.
     """
     # \xff\xfe are invalid UTF-8 start bytes.
     valid = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
@@ -92,3 +94,78 @@ async def test_stdio_server_invalid_utf8(monkeypatch: pytest.MonkeyPatch):
                 second = await read_stream.receive()
                 assert isinstance(second, SessionMessage)
                 assert second.message == valid
+
+
+class _KeepOpenBytesIO(io.BytesIO):
+    """A BytesIO that survives its TextIOWrapper being closed.
+
+    Lets the test read what was written after `run()` has torn the wrapper down.
+    """
+
+    def close(self) -> None:
+        pass
+
+
+def _run_stdio_bounded(server: MCPServer) -> None:
+    """Run the blocking `server.run("stdio")` in a daemon thread joined with a 5s bound.
+
+    `run()` creates its own event loop, so a sync test cannot arm `anyio.fail_after`;
+    the join timeout turns a run loop that never returns on stdin EOF into a red test
+    instead of a silent CI hang. An exception escaping `run()` still fails the test:
+    pytest's unhandled-thread warning is escalated by `filterwarnings = ["error"]`.
+    """
+
+    def target() -> None:
+        server.run("stdio")
+
+    thread = threading.Thread(target=target, daemon=True)
+    thread.start()
+    thread.join(5)
+    assert not thread.is_alive(), 'run("stdio") did not return after stdin EOF'
+
+
+def test_mcpserver_run_stdio_serves_until_stdin_closes(monkeypatch: pytest.MonkeyPatch) -> None:
+    """`MCPServer.run("stdio")` serves over process stdio and returns at stdin EOF.
+
+    Answers a request over the process's stdio and returns when stdin reaches EOF,
+    rather than serving forever.
+    """
+    ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+    stdin_bytes = io.BytesIO(ping.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n")
+    captured = _KeepOpenBytesIO()
+    monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin_bytes, encoding="utf-8"))
+    monkeypatch.setattr(sys, "stdout", TextIOWrapper(captured, encoding="utf-8"))
+
+    _run_stdio_bounded(MCPServer(name="RunStdioServer"))
+
+    response = jsonrpc_message_adapter.validate_json(captured.getvalue().decode().strip())
+    assert response == JSONRPCResponse(jsonrpc="2.0", id=1, result={})
+
+
+def test_mcpserver_run_stdio_runs_lifespan_cleanup_after_stdin_closes(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Code after `yield` in a lifespan runs when stdin EOF ends `run("stdio")`.
+
+    Regression lock for the issue #1027 shutdown chain: the run loop must end on
+    stdin EOF and unwind the lifespan rather than be killed before returning.
+    """
+    events: list[str] = []
+
+    @asynccontextmanager
+    async def lifespan(server: MCPServer) -> AsyncIterator[None]:
+        events.append("setup")
+        try:
+            yield
+        finally:
+            events.append("cleanup")
+
+    ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+    stdin_bytes = io.BytesIO(ping.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n")
+    captured = _KeepOpenBytesIO()
+    monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin_bytes, encoding="utf-8"))
+    monkeypatch.setattr(sys, "stdout", TextIOWrapper(captured, encoding="utf-8"))
+
+    _run_stdio_bounded(MCPServer(name="LifespanStdioServer", lifespan=lifespan))
+
+    assert events == ["setup", "cleanup"]
+    response = jsonrpc_message_adapter.validate_json(captured.getvalue().decode().strip())
+    assert response == JSONRPCResponse(jsonrpc="2.0", id=1, result={})