Skip to content

Commit b6be755

Browse files
committed
Transport: pre-session 404 maps to METHOD_NOT_FOUND; POSTs never read cached pv header. serve_one reshape + raise_exceptions
- Bare HTTP 404 before a session is established now maps to METHOD_NOT_FOUND (was INVALID_REQUEST/"Session terminated", which is meaningless pre-session); with a session_id, 404 keeps the session-terminated mapping - _prepare_headers split: _base_headers (POST) vs _prepare_headers (GET/DELETE). POSTs never read the cached MCP-Protocol-Version header — they get it via per-message metadata only. Prevents the discover probe's header from leaking onto a fallback initialize POST. - serve_one reshaped to (server, dctx, method, params, *, ..., raise_exceptions); modern_on_request drops the JSONRPCRequest round-trip and threads raise_exceptions through to to_jsonrpc_response(raise_unhandled=). Client's modern in-process branch now honors raise_exceptions (handler exceptions chain via __cause__ instead of being sanitized to INTERNAL_ERROR).
1 parent 52f200b commit b6be755

8 files changed

Lines changed: 194 additions & 39 deletions

File tree

src/mcp/client/client.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,10 @@ async def __aenter__(self) -> Client:
159159
client_disp, server_disp = create_direct_dispatcher_pair()
160160
tg = await exit_stack.enter_async_context(anyio.create_task_group())
161161
exit_stack.callback(server_disp.close)
162-
await tg.start(server_disp.run, modern_on_request(self._inproc_server, lifespan_state), _drop_notify)
162+
on_request = modern_on_request(
163+
self._inproc_server, lifespan_state, raise_exceptions=self.raise_exceptions
164+
)
165+
await tg.start(server_disp.run, on_request, _drop_notify)
163166
session = ClientSession(
164167
dispatcher=client_disp,
165168
read_timeout_seconds=self.read_timeout_seconds,

src/mcp/client/streamable_http.py

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from mcp.types import (
2424
INTERNAL_ERROR,
2525
INVALID_REQUEST,
26+
METHOD_NOT_FOUND,
2627
PARSE_ERROR,
2728
ErrorData,
2829
JSONRPCError,
@@ -84,19 +85,29 @@ def __init__(self, url: str) -> None:
8485
# GET/DELETE that don't carry per-message metadata.
8586
self._protocol_version_header: str | None = None
8687

87-
def _prepare_headers(self) -> dict[str, str]:
88-
"""Build MCP-specific request headers.
88+
def _base_headers(self) -> dict[str, str]:
89+
"""Build MCP-specific request headers (accept / content-type / session-id).
8990
9091
These headers will be merged with the httpx.AsyncClient's default headers,
91-
with these MCP-specific headers taking precedence.
92+
with these MCP-specific headers taking precedence. POSTs use this directly:
93+
their protocol-version header arrives per-message via ``metadata.headers``,
94+
so they must never read the cached value.
9295
"""
9396
headers: dict[str, str] = {
9497
"accept": "application/json, text/event-stream",
9598
"content-type": "application/json",
9699
}
97-
# Add session headers if available
98100
if self.session_id:
99101
headers[MCP_SESSION_ID] = self.session_id
102+
return headers
103+
104+
def _prepare_headers(self) -> dict[str, str]:
105+
"""Base headers plus the cached protocol-version header.
106+
107+
Used by transport-internal GET/DELETE (listen stream, resumption,
108+
reconnect, terminate) which don't carry per-message metadata.
109+
"""
110+
headers = self._base_headers()
100111
if self._protocol_version_header:
101112
headers[MCP_PROTOCOL_VERSION_HEADER] = self._protocol_version_header
102113
return headers
@@ -238,7 +249,7 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
238249

239250
async def _handle_post_request(self, ctx: RequestContext) -> None:
240251
"""Handle a POST request with response processing."""
241-
headers = self._prepare_headers()
252+
headers = self._base_headers()
242253
message = ctx.session_message.message
243254
if ctx.metadata is not None and ctx.metadata.headers is not None:
244255
headers.update(ctx.metadata.headers)
@@ -276,7 +287,13 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
276287
pass
277288
logger.debug("Non-2xx body was not a JSON-RPC error; using fallback")
278289
if response.status_code == 404:
279-
error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated")
290+
if self.session_id is None:
291+
# No session yet → 404 is the HTTP-level spelling of
292+
# METHOD_NOT_FOUND (gateway / legacy server doesn't know
293+
# this method); "Session terminated" would be a lie here.
294+
error_data = ErrorData(code=METHOD_NOT_FOUND, message="Not Found")
295+
else:
296+
error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated")
280297
else:
281298
error_data = ErrorData(code=INTERNAL_ERROR, message="Server returned an error response")
282299
session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data))

src/mcp/server/_streamable_http_modern.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,5 +193,5 @@ async def handle_modern_request(
193193
request_id=req.id,
194194
message_metadata=ServerMessageMetadata(request_context=request),
195195
)
196-
msg = await serve_one(app, req, connection=connection, dctx=dctx, lifespan_state=lifespan_state)
196+
msg = await serve_one(app, dctx, req.method, req.params, connection=connection, lifespan_state=lifespan_state)
197197
await _write(msg, scope, receive, send)

src/mcp/server/runner.py

Lines changed: 39 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,6 @@
4949
InitializeRequestParams,
5050
InitializeResult,
5151
JSONRPCError,
52-
JSONRPCRequest,
5352
JSONRPCResponse,
5453
RequestId,
5554
RequestParams,
@@ -183,20 +182,26 @@ async def aclose_shielded(connection: Connection) -> None:
183182
)
184183

185184

186-
async def to_jsonrpc_response(request_id: RequestId, coro: Awaitable[dict[str, Any]]) -> JSONRPCResponse | JSONRPCError:
185+
async def to_jsonrpc_response(
186+
request_id: RequestId, coro: Awaitable[dict[str, Any]], *, raise_unhandled: bool = False
187+
) -> JSONRPCResponse | JSONRPCError:
187188
"""Await ``coro`` and wrap its outcome as the JSON-RPC reply for ``request_id``.
188189
189190
The exception-to-wire boundary for the request-per-call drivers
190191
(`serve_one`, the modern HTTP entry). `MCPError` and `ValidationError`
191192
map via the shared `handler_exception_to_error_data` ladder; any other
192193
exception is logged and surfaced as `INTERNAL_ERROR` so handler internals
193-
never reach the wire.
194+
never reach the wire. Set ``raise_unhandled`` to let unmapped exceptions
195+
propagate instead of being sanitized — used by the in-process test path so
196+
handler tracebacks reach the caller.
194197
"""
195198
try:
196199
result = await coro
197200
except Exception as exc:
198201
error = handler_exception_to_error_data(exc)
199202
if error is None:
203+
if raise_unhandled:
204+
raise
200205
logger.exception("request handler raised")
201206
error = ErrorData(code=INTERNAL_ERROR, message="Internal server error")
202207
return JSONRPCError(jsonrpc="2.0", id=request_id, error=error)
@@ -497,34 +502,45 @@ async def serve_loop(
497502

498503
async def serve_one(
499504
server: Server[LifespanT],
500-
request: JSONRPCRequest,
505+
dctx: DispatchContext[TransportContext],
506+
method: str,
507+
params: Mapping[str, Any] | None,
501508
*,
502509
connection: Connection,
503-
dctx: DispatchContext[TransportContext],
504510
lifespan_state: LifespanT,
511+
raise_exceptions: bool = False,
505512
) -> JSONRPCResponse | JSONRPCError:
506-
"""Handle a single ``request`` and return its JSON-RPC reply.
507-
508-
The single-exchange driver: builds the kernel, runs `on_request` once for
509-
`request` under `dctx`, maps the outcome to a `JSONRPCResponse` /
510-
`JSONRPCError` via `to_jsonrpc_response`, and tears down
511-
`connection.exit_stack` (shielded) on the way out. The entry constructs
512-
the (born-ready) `Connection` and the `dctx`; this only consumes them.
513+
"""Handle a single request ``(method, params)`` and return its JSON-RPC reply.
514+
515+
The single-exchange driver: builds the kernel, runs `on_request` once under
516+
`dctx`, maps the outcome to a `JSONRPCResponse` / `JSONRPCError` via
517+
`to_jsonrpc_response`, and tears down `connection.exit_stack` (shielded) on
518+
the way out. The entry constructs the (born-ready) `Connection` and the
519+
`dctx`; this only consumes them. ``raise_exceptions`` lets unmapped handler
520+
exceptions propagate instead of being sanitized to `INTERNAL_ERROR`.
513521
"""
514522
runner = ServerRunner(server, connection, lifespan_state)
515523
try:
516-
return await to_jsonrpc_response(request.id, runner.on_request(dctx, request.method, request.params))
524+
# Single-exchange driver only handles requests; both entries populate `request_id`.
525+
# TODO(L54): drop once `DispatchContext` is split so `OnRequest` carries a non-Optional id.
526+
assert dctx.request_id is not None
527+
return await to_jsonrpc_response(
528+
dctx.request_id, runner.on_request(dctx, method, params), raise_unhandled=raise_exceptions
529+
)
517530
finally:
518531
await aclose_shielded(connection)
519532

520533

521-
def modern_on_request(server: Server[LifespanT], lifespan_state: LifespanT) -> OnRequest:
534+
def modern_on_request(
535+
server: Server[LifespanT], lifespan_state: LifespanT, *, raise_exceptions: bool = False
536+
) -> OnRequest:
522537
"""Return an `OnRequest` callback that serves each call via `serve_one` with a fresh per-request `Connection`.
523538
524539
Wire this into the server side of a `DirectDispatcher` peer-pair to drive an
525540
in-process server on the modern per-request-envelope path (each request
526541
carries protocol version, client info, and capabilities in `params._meta`;
527-
no `initialize` handshake).
542+
no `initialize` handshake). ``raise_exceptions`` lets unmapped handler
543+
exceptions propagate to the caller for debuggable in-process testing.
528544
"""
529545

530546
async def handle(
@@ -536,12 +552,15 @@ async def handle(
536552
meta.get(CLIENT_INFO_META_KEY),
537553
meta.get(CLIENT_CAPABILITIES_META_KEY),
538554
)
539-
# `OnRequest` is invoked for requests only, so `request_id` is always set.
540-
assert dctx.request_id is not None
541-
req = JSONRPCRequest(
542-
jsonrpc="2.0", id=dctx.request_id, method=method, params=dict(params) if params is not None else None
555+
msg = await serve_one(
556+
server,
557+
dctx,
558+
method,
559+
params,
560+
connection=connection,
561+
lifespan_state=lifespan_state,
562+
raise_exceptions=raise_exceptions,
543563
)
544-
msg = await serve_one(server, req, connection=connection, dctx=dctx, lifespan_state=lifespan_state)
545564
if isinstance(msg, JSONRPCError):
546565
raise MCPError(code=msg.error.code, message=msg.error.message, data=msg.error.data)
547566
return msg.result

tests/client/test_client.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,23 @@ async def handle_read_resource(
203203
assert exc_info.value.error.code == 404
204204

205205

206+
async def test_raise_exceptions_propagates_handler_error_on_modern_inproc_path():
207+
"""`raise_exceptions=True` on the modern in-process path: an unmapped handler
208+
exception reaches the client with its original type chained, instead of being
209+
sanitized to an opaque `INTERNAL_ERROR`."""
210+
211+
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
212+
raise ValueError("boom")
213+
214+
server = Server("test", on_call_tool=handle_call_tool)
215+
async with Client(server, mode="2026-07-28", raise_exceptions=True) as client:
216+
with pytest.raises(MCPError) as exc_info:
217+
await client.call_tool("explode", {})
218+
# The original exception is chained — not swallowed into a generic "Internal server error".
219+
assert isinstance(exc_info.value.__cause__, ValueError)
220+
assert str(exc_info.value.__cause__) == "boom"
221+
222+
206223
async def test_get_prompt(app: MCPServer):
207224
"""Test getting a prompt."""
208225
async with Client(app) as client:

tests/client/test_notification_response.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -204,12 +204,19 @@ async def test_invalid_json_response_sends_jsonrpc_error() -> None:
204204

205205

206206
def _create_non_2xx_json_body_app(status: int, body: bytes) -> Starlette:
207-
"""Server that returns a fixed non-2xx status + ``application/json`` body for non-init requests."""
207+
"""Server that returns a fixed non-2xx status + ``application/json`` body for non-init requests.
208+
209+
The initialize response carries an ``mcp-session-id`` so the client treats subsequent
210+
requests as part of an established session (needed for the 404 → session-terminated mapping).
211+
"""
208212

209213
async def handle_mcp_request(request: Request) -> Response:
210214
data = json.loads(await request.body())
211215
if data.get("method") == "initialize":
212-
return _init_json_response(data)
216+
return JSONResponse(
217+
{"jsonrpc": "2.0", "id": data["id"], "result": INIT_RESPONSE},
218+
headers={"mcp-session-id": "test-session"},
219+
)
213220
if "id" not in data:
214221
return Response(status_code=202)
215222
return Response(content=body, status_code=status, media_type="application/json")

tests/client/test_streamable_http.py

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
from inline_snapshot import snapshot
1616

1717
from mcp.client.streamable_http import streamable_http_client
18-
from mcp.shared.inbound import encode_header_value
18+
from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER, encode_header_value
1919
from mcp.shared.message import ClientMessageMetadata, SessionMessage
20-
from mcp.types import JSONRPCRequest
20+
from mcp.types import METHOD_NOT_FOUND, JSONRPCError, JSONRPCRequest
2121

2222

2323
@pytest.mark.parametrize(
@@ -78,3 +78,61 @@ def handler(request: httpx.Request) -> httpx.Response:
7878
assert isinstance(reply, SessionMessage)
7979
assert [r.method for r in recorded] == ["POST"]
8080
assert recorded[0].headers["x-test"] == "v"
81+
82+
83+
@pytest.mark.anyio
84+
async def test_pre_session_bare_404_maps_to_method_not_found() -> None:
85+
"""A bare HTTP 404 (no JSON-RPC body) before any session-id is held maps to METHOD_NOT_FOUND.
86+
87+
Gateways and legacy servers 404 at the HTTP layer for unknown methods; with no session yet,
88+
"Session terminated" is meaningless, and the discover→initialize fallback ladder keys on -32601.
89+
"""
90+
91+
def handler(request: httpx.Request) -> httpx.Response:
92+
return httpx.Response(404)
93+
94+
with anyio.fail_after(5):
95+
async with (
96+
httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http,
97+
streamable_http_client("http://test/mcp", http_client=http) as (read, write),
98+
):
99+
await write.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="server/discover", params={})))
100+
reply = await read.receive()
101+
assert isinstance(reply, SessionMessage)
102+
assert isinstance(reply.message, JSONRPCError)
103+
assert reply.message.error.code == METHOD_NOT_FOUND
104+
105+
106+
@pytest.mark.anyio
107+
async def test_post_does_not_read_cached_protocol_version_header() -> None:
108+
"""A POST's protocol-version header comes only from its own ``metadata.headers``.
109+
110+
The first POST carries (and caches) a pv header; the second POST sends no metadata
111+
and must therefore carry no pv header — a stale cached value would poison the
112+
fallback ``initialize`` after a failed discover probe. The cache exists for
113+
transport-internal GET/DELETE only.
114+
"""
115+
recorded: list[httpx.Request] = []
116+
117+
def handler(request: httpx.Request) -> httpx.Response:
118+
recorded.append(request)
119+
body = json.loads(request.content)
120+
return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": {}})
121+
122+
with anyio.fail_after(5):
123+
async with (
124+
httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http,
125+
streamable_http_client("http://test/mcp", http_client=http) as (read, write),
126+
):
127+
await write.send(
128+
SessionMessage(
129+
message=JSONRPCRequest(jsonrpc="2.0", id=1, method="server/discover", params={}),
130+
metadata=ClientMessageMetadata(headers={MCP_PROTOCOL_VERSION_HEADER: "2026-07-28"}),
131+
)
132+
)
133+
await read.receive()
134+
await write.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=2, method="initialize", params={})))
135+
await read.receive()
136+
assert [r.method for r in recorded] == ["POST", "POST"]
137+
assert recorded[0].headers[MCP_PROTOCOL_VERSION_HEADER] == "2026-07-28"
138+
assert MCP_PROTOCOL_VERSION_HEADER not in recorded[1].headers

0 commit comments

Comments
 (0)