|
1 | 1 | import gc |
2 | 2 | import io |
| 3 | +import json |
3 | 4 | import os |
4 | 5 | import sys |
5 | 6 | import threading |
|
13 | 14 | from mcp_types import ( |
14 | 15 | CLIENT_CAPABILITIES_META_KEY, |
15 | 16 | CLIENT_INFO_META_KEY, |
| 17 | + INVALID_REQUEST, |
| 18 | + PARSE_ERROR, |
16 | 19 | PROTOCOL_VERSION_META_KEY, |
17 | 20 | SERVER_INFO_META_KEY, |
| 21 | + ErrorData, |
| 22 | + JSONRPCError, |
18 | 23 | JSONRPCMessage, |
19 | 24 | JSONRPCRequest, |
20 | 25 | JSONRPCResponse, |
@@ -75,26 +80,130 @@ async def test_stdio_server_round_trips_messages_over_injected_streams() -> None |
75 | 80 | assert received_responses[1] == JSONRPCResponse(jsonrpc="2.0", id=4, result={}) |
76 | 81 |
|
77 | 82 |
|
| 83 | +async def _serve_raw_lines(lines: list[str]) -> tuple[list[SessionMessage | Exception], list[JSONRPCMessage]]: |
| 84 | + """Feed raw stdin lines through `stdio_server()`; return (read-stream items, parsed stdout frames). |
| 85 | +
|
| 86 | + Raw lines are the only way to exercise malformed input: the typed API cannot produce it. |
| 87 | + """ |
| 88 | + stdin = io.StringIO("".join(line + "\n" for line in lines)) |
| 89 | + stdout = io.StringIO() |
| 90 | + |
| 91 | + received: list[SessionMessage | Exception] = [] |
| 92 | + with anyio.fail_after(5): |
| 93 | + async with stdio_server(stdin=anyio.AsyncFile(stdin), stdout=anyio.AsyncFile(stdout)) as ( |
| 94 | + read_stream, |
| 95 | + write_stream, |
| 96 | + ): |
| 97 | + async with read_stream: |
| 98 | + try: |
| 99 | + while True: |
| 100 | + received.append(await read_stream.receive()) |
| 101 | + except anyio.EndOfStream: |
| 102 | + pass |
| 103 | + await write_stream.aclose() |
| 104 | + |
| 105 | + stdout.seek(0) |
| 106 | + sent = [jsonrpc_message_adapter.validate_json(line.strip()) for line in stdout.readlines()] |
| 107 | + return received, sent |
| 108 | + |
| 109 | + |
78 | 110 | @pytest.mark.anyio |
79 | 111 | async def test_stdio_server_invalid_utf8(monkeypatch: pytest.MonkeyPatch) -> None: |
80 | | - """Non-UTF-8 stdin bytes surface as an in-stream exception without killing the stream.""" |
| 112 | + """Non-UTF-8 stdin bytes are answered with a parse error frame without killing the stream. |
| 113 | +
|
| 114 | + Invalid bytes are replaced with U+FFFD, fail JSON parsing, and complete as a |
| 115 | + JSON-RPC parse error with `id: null` (no id can be recovered); subsequent |
| 116 | + valid messages are still processed. |
| 117 | + """ |
81 | 118 | valid = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping") |
82 | 119 | raw_stdin = io.BytesIO(b"\xff\xfe\n" + valid.model_dump_json(by_alias=True, exclude_none=True).encode() + b"\n") |
| 120 | + stdout = io.StringIO() |
83 | 121 |
|
84 | 122 | # stdio_server()'s default path wraps sys.stdin.buffer with errors='replace'. |
85 | 123 | monkeypatch.setattr(sys, "stdin", TextIOWrapper(raw_stdin, encoding="utf-8")) |
86 | | - monkeypatch.setattr(sys, "stdout", TextIOWrapper(io.BytesIO(), encoding="utf-8")) |
87 | 124 |
|
88 | 125 | with anyio.fail_after(5): |
89 | | - async with stdio_server() as (read_stream, write_stream): |
90 | | - await write_stream.aclose() |
| 126 | + async with stdio_server(stdout=anyio.AsyncFile(stdout)) as (read_stream, write_stream): |
91 | 127 | async with read_stream: # pragma: no branch |
92 | | - first = await read_stream.receive() |
93 | | - assert isinstance(first, Exception) |
| 128 | + received = await read_stream.receive() |
| 129 | + assert isinstance(received, SessionMessage) |
| 130 | + assert received.message == valid |
| 131 | + await write_stream.aclose() |
94 | 132 |
|
95 | | - second = await read_stream.receive() |
96 | | - assert isinstance(second, SessionMessage) |
97 | | - assert second.message == valid |
| 133 | + stdout.seek(0) |
| 134 | + sent = [jsonrpc_message_adapter.validate_json(line.strip()) for line in stdout.readlines()] |
| 135 | + assert sent == [JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=PARSE_ERROR, message="Parse error"))] |
| 136 | + |
| 137 | + |
| 138 | +@pytest.mark.anyio |
| 139 | +async def test_stdio_server_deeply_nested_request_completes_with_a_parse_error() -> None: |
| 140 | + """A request too deeply nested to validate completes with a parse error carrying its id. |
| 141 | +
|
| 142 | + Regression for #2751: the failure used to surface only as a bare exception on |
| 143 | + the read stream, which the server dispatcher drops, leaving the request id |
| 144 | + pending until the client timed out. |
| 145 | + """ |
| 146 | + params: object = {"leaf": True} |
| 147 | + for index in reversed(range(256)): |
| 148 | + params = {f"p{index}": params} |
| 149 | + deep_line = json.dumps({"jsonrpc": "2.0", "id": 900256, "method": "ping", "params": params}) |
| 150 | + shallow = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping") |
| 151 | + |
| 152 | + received, sent = await _serve_raw_lines([deep_line, shallow.model_dump_json(by_alias=True, exclude_none=True)]) |
| 153 | + |
| 154 | + # The shallow request after the failed one still arrives; the failed one does not. |
| 155 | + assert len(received) == 1 |
| 156 | + assert isinstance(received[0], SessionMessage) |
| 157 | + assert received[0].message == shallow |
| 158 | + assert sent == [JSONRPCError(jsonrpc="2.0", id=900256, error=ErrorData(code=PARSE_ERROR, message="Parse error"))] |
| 159 | + |
| 160 | + |
| 161 | +@pytest.mark.anyio |
| 162 | +@pytest.mark.parametrize( |
| 163 | + ("line", "expected"), |
| 164 | + [ |
| 165 | + # Truncated JSON: the id is recovered from the raw line, ahead of the break. |
| 166 | + ( |
| 167 | + '{"jsonrpc":"2.0","id":"bad-1","method":"ping","params":[', |
| 168 | + JSONRPCError(jsonrpc="2.0", id="bad-1", error=ErrorData(code=PARSE_ERROR, message="Parse error")), |
| 169 | + ), |
| 170 | + ( |
| 171 | + '{"jsonrpc":"2.0","id":42,"method":"ping","params":[', |
| 172 | + JSONRPCError(jsonrpc="2.0", id=42, error=ErrorData(code=PARSE_ERROR, message="Parse error")), |
| 173 | + ), |
| 174 | + # No id anywhere in the line: the parse error falls back to `id: null`. |
| 175 | + ( |
| 176 | + "not json", |
| 177 | + JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=PARSE_ERROR, message="Parse error")), |
| 178 | + ), |
| 179 | + # An id whose string token is itself invalid JSON is not recoverable either. |
| 180 | + ( |
| 181 | + '{"jsonrpc":"2.0","id":"a\\q","method":"ping"}', |
| 182 | + JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=PARSE_ERROR, message="Parse error")), |
| 183 | + ), |
| 184 | + # Valid JSON that is not a valid message: invalid request, id preserved where present. |
| 185 | + ( |
| 186 | + '{"jsonrpc":"2.0","id":"bad-method","method":1}', |
| 187 | + JSONRPCError( |
| 188 | + jsonrpc="2.0", id="bad-method", error=ErrorData(code=INVALID_REQUEST, message="Invalid request") |
| 189 | + ), |
| 190 | + ), |
| 191 | + ( |
| 192 | + '{"jsonrpc":"2.0","id":null,"method":1}', |
| 193 | + JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Invalid request")), |
| 194 | + ), |
| 195 | + ( |
| 196 | + "[]", |
| 197 | + JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Invalid request")), |
| 198 | + ), |
| 199 | + ], |
| 200 | +) |
| 201 | +async def test_stdio_server_unparseable_line_completes_with_an_error_frame(line: str, expected: JSONRPCError) -> None: |
| 202 | + """A line that fails validation is answered on stdout with a single JSON-RPC error frame.""" |
| 203 | + received, sent = await _serve_raw_lines([line]) |
| 204 | + |
| 205 | + assert received == [] |
| 206 | + assert sent == [expected] |
98 | 207 |
|
99 | 208 |
|
100 | 209 | @contextmanager |
@@ -608,14 +717,19 @@ def close(self) -> None: |
608 | 717 |
|
609 | 718 |
|
610 | 719 | def _serve_stdio_and_collect( |
611 | | - monkeypatch: pytest.MonkeyPatch, server: MCPServer, frames: list[JSONRPCRequest], responses: int |
| 720 | + monkeypatch: pytest.MonkeyPatch, |
| 721 | + server: MCPServer, |
| 722 | + frames: list[JSONRPCRequest], |
| 723 | + responses: int, |
| 724 | + raw_prefix: bytes = b"", |
612 | 725 | ) -> list[JSONRPCMessage]: |
613 | 726 | """Serve `frames` over process stdio and return the parsed response lines. |
614 | 727 |
|
615 | 728 | Runs the blocking server.run("stdio") in a daemon thread and signals stdin EOF only after |
616 | 729 | `responses` lines arrive - as a real client would - so handlers never race EOF cancellation. |
| 730 | + `raw_prefix` prepends untyped wire bytes (e.g. a malformed line) the typed frames cannot express. |
617 | 731 | """ |
618 | | - payload = "".join(f.model_dump_json(by_alias=True, exclude_none=True) + "\n" for f in frames).encode() |
| 732 | + payload = raw_prefix + "".join(f.model_dump_json(by_alias=True, exclude_none=True) + "\n" for f in frames).encode() |
619 | 733 | stdin = _GatedStdin(payload) |
620 | 734 | stdout = _NotifyingStdout() |
621 | 735 | monkeypatch.setattr(sys, "stdin", TextIOWrapper(stdin, encoding="utf-8")) |
@@ -643,6 +757,31 @@ def test_mcpserver_run_stdio_serves_until_stdin_closes(monkeypatch: pytest.Monke |
643 | 757 | assert responses == [JSONRPCResponse(jsonrpc="2.0", id=1, result={})] |
644 | 758 |
|
645 | 759 |
|
| 760 | +def test_mcpserver_run_stdio_answers_a_deeply_nested_request_with_a_parse_error( |
| 761 | + monkeypatch: pytest.MonkeyPatch, |
| 762 | +) -> None: |
| 763 | + """A deeply nested request over real stdio gets a parse error response instead of hanging. |
| 764 | +
|
| 765 | + End-to-end regression for #2751 through `MCPServer.run("stdio")`: the client |
| 766 | + waiting on the nested request's id receives a response, and a later shallow |
| 767 | + request on the same connection is still served. |
| 768 | + """ |
| 769 | + params: object = {"leaf": True} |
| 770 | + for index in reversed(range(256)): |
| 771 | + params = {f"p{index}": params} |
| 772 | + deep_line = (json.dumps({"jsonrpc": "2.0", "id": 900256, "method": "ping", "params": params}) + "\n").encode() |
| 773 | + ping = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping") |
| 774 | + |
| 775 | + responses = _serve_stdio_and_collect( |
| 776 | + monkeypatch, MCPServer(name="DeepRequestStdioServer"), [ping], 2, raw_prefix=deep_line |
| 777 | + ) |
| 778 | + |
| 779 | + assert responses == [ |
| 780 | + JSONRPCError(jsonrpc="2.0", id=900256, error=ErrorData(code=PARSE_ERROR, message="Parse error")), |
| 781 | + JSONRPCResponse(jsonrpc="2.0", id=1, result={}), |
| 782 | + ] |
| 783 | + |
| 784 | + |
646 | 785 | def test_mcpserver_run_stdio_runs_lifespan_cleanup_after_stdin_closes(monkeypatch: pytest.MonkeyPatch) -> None: |
647 | 786 | """Code after `yield` in a lifespan runs when stdin EOF ends `run("stdio")`. |
648 | 787 |
|
|
0 commit comments