From 05c333acdcddab2c8d81923046c3323d36e25df2 Mon Sep 17 00:00:00 2001 From: Anthony Ivan Date: Wed, 2 Sep 2026 21:49:51 +0800 Subject: [PATCH 1/3] Handle MCP transport timeouts and disconnects Signed-off-by: Anthony Ivan --- src/ucode/mcp_proxy.py | 28 +++++++++++++++++--- tests/test_mcp_proxy.py | 57 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 4 deletions(-) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index ee622884..5e84f79a 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -117,22 +117,41 @@ async def _pump( source: MemoryObjectReceiveStream, dest: MemoryObjectSendStream, ) -> None: - """Forward every message (or transport exception) from ``source`` to ``dest``. + """Forward every message from ``source`` to ``dest``. The proxy is transport-level: it never inspects or rewrites MCP method payloads, so new methods and capabilities pass through untouched.""" async with source, dest: async for message in source: + if isinstance(message, Exception): + raise message await dest.send(message) +async def _pump_upstream( + source: MemoryObjectReceiveStream, + dest: MemoryObjectSendStream, +) -> None: + """Forward upstream messages, failing if the transport closes first.""" + await _pump(source, dest) + raise RuntimeError("upstream MCP transport closed unexpectedly") + + async def _run(url: str, workspace: str, profile: str | None) -> None: httpx = _httpx() auth = _build_token_auth(workspace, profile) # 2.x-native shape: hand the transport a pre-built AsyncClient carrying our # per-request auth. Works on mcp 1.28+ and 2.x; `streamable_http_client` # yields a (read, write) pair in both. - async with httpx.AsyncClient(auth=auth) as http_client: + async with httpx.AsyncClient( + auth=auth, + timeout=httpx.Timeout( + connect=30.0, + read=300.0, + write=30.0, + pool=30.0, + ), + ) as http_client: async with streamable_http_client(url, http_client=http_client) as streams: # mcp 1.x yields (read, write, get_session_id); mcp 2.x drops the # trailing callback and yields (read, write). Take the first two @@ -141,8 +160,9 @@ async def _run(url: str, workspace: str, profile: str | None) -> None: async with stdio_server() as (stdio_read, stdio_write): # Bidirectional bridge: client stdin -> Databricks, Databricks -> client stdout. async with anyio.create_task_group() as tg: - tg.start_soon(_pump, stdio_read, http_write) - tg.start_soon(_pump, http_read, stdio_write) + tg.start_soon(_pump_upstream, http_read, stdio_write) + await _pump(stdio_read, http_write) + tg.cancel_scope.cancel() def _preflight_token(workspace: str, profile: str | None) -> None: diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index a10f8a67..6c8fbda2 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -3,6 +3,7 @@ from __future__ import annotations import tomllib +from contextlib import asynccontextmanager from pathlib import Path import anyio @@ -165,6 +166,62 @@ async def scenario() -> bool: assert anyio.run(scenario) is True + def test_raises_transport_errors(self): + async def scenario() -> None: + src_send, src_recv = anyio.create_memory_object_stream(1) + dst_send, _ = anyio.create_memory_object_stream(1) + await src_send.send(httpx.ReadTimeout("upstream timed out")) + await src_send.aclose() + + with pytest.raises(httpx.ReadTimeout, match="upstream timed out"): + await mcp_proxy._pump(src_recv, dst_send) + + anyio.run(scenario) + + def test_upstream_eof_is_an_error(self): + async def scenario() -> None: + src_send, src_recv = anyio.create_memory_object_stream(1) + dst_send, _ = anyio.create_memory_object_stream(1) + await src_send.aclose() + + with pytest.raises(RuntimeError, match="upstream MCP transport closed"): + await mcp_proxy._pump_upstream(src_recv, dst_send) + + anyio.run(scenario) + + +def test_run_uses_mcp_http_defaults(monkeypatch): + httpx_module = mcp_proxy._httpx() + captured: dict = {} + + class CapturingClient: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + class StopBridge(Exception): + pass + + @asynccontextmanager + async def stop_bridge(*args, **kwargs): + raise StopBridge + yield + + monkeypatch.setattr(httpx_module, "AsyncClient", CapturingClient) + monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args: object()) + monkeypatch.setattr(mcp_proxy, "streamable_http_client", stop_bridge) + + with pytest.raises(StopBridge): + anyio.run(mcp_proxy._run, URL, WS, None) + + timeout = captured["timeout"] + assert (timeout.connect, timeout.read, timeout.write, timeout.pool) == (30.0, 300.0, 30.0, 30.0) + class TestServe: def test_runs_the_bridge_with_parsed_args(self, monkeypatch): From 1aeeaad63ebf116a35884c399151a7e480a0452a Mon Sep 17 00:00:00 2001 From: Anthony Ivan Date: Thu, 3 Sep 2026 00:59:11 +0800 Subject: [PATCH 2/3] Handle MCP transport failures cleanly Signed-off-by: Anthony Ivan --- src/ucode/mcp_proxy.py | 48 +++++++++++++++++++++++------------------ tests/test_mcp_proxy.py | 43 +++++++++++++++++++++++++++++++----- 2 files changed, 65 insertions(+), 26 deletions(-) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 5e84f79a..2146cf4b 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -43,9 +43,8 @@ from ucode.databricks import ensure_pat_bearer, get_databricks_token -# Exit code used when the proxy cannot authenticate. MCP clients surface a -# non-zero exit far more usefully than a startup timeout, so bail out with this -# instead of letting the process hang until the client's timeout fires. +# Exit code used when the proxy cannot continue. MCP clients surface a non-zero +# exit far more usefully than a timeout, so bail out instead of hanging. AUTH_FAILURE_EXIT_CODE = 2 @@ -76,8 +75,12 @@ class ProxyAuthError(RuntimeError): stderr and exits rather than retrying.""" +class ProxyTransportError(RuntimeError): + """The upstream MCP transport failed or closed unexpectedly.""" + + def _fail_fast(message: str) -> None: - """Report a terminal auth failure on stderr and exit non-zero. + """Report a terminal proxy failure on stderr and exit non-zero. stdout is the MCP wire, so diagnostics must go to stderr — MCP clients surface a child's stderr when it fails to start.""" @@ -123,8 +126,6 @@ async def _pump( payloads, so new methods and capabilities pass through untouched.""" async with source, dest: async for message in source: - if isinstance(message, Exception): - raise message await dest.send(message) @@ -133,8 +134,13 @@ async def _pump_upstream( dest: MemoryObjectSendStream, ) -> None: """Forward upstream messages, failing if the transport closes first.""" - await _pump(source, dest) - raise RuntimeError("upstream MCP transport closed unexpectedly") + async with source, dest: + async for message in source: + if isinstance(message, Exception): + detail = " ".join(str(message).split()) or type(message).__name__ + raise ProxyTransportError(f"upstream MCP transport failed: {detail}") from message + await dest.send(message) + raise ProxyTransportError("upstream MCP transport closed unexpectedly") async def _run(url: str, workspace: str, profile: str | None) -> None: @@ -176,15 +182,15 @@ def _preflight_token(workspace: str, profile: str | None) -> None: get_databricks_token(workspace, profile) -def _unwrap_auth_error(exc: BaseException) -> ProxyAuthError | None: - """Find a ProxyAuthError anywhere in an exception (or ExceptionGroup) tree. +def _unwrap_proxy_error(exc: BaseException) -> ProxyAuthError | ProxyTransportError | None: + """Find a known proxy error in an exception (or ExceptionGroup) tree. - anyio task groups wrap failures in ExceptionGroups, so a token failure - raised inside the transport arrives nested rather than as itself.""" - if isinstance(exc, ProxyAuthError): + anyio task groups wrap failures in ExceptionGroups, so failures raised + inside the transport arrive nested rather than as themselves.""" + if isinstance(exc, (ProxyAuthError, ProxyTransportError)): return exc for nested in getattr(exc, "exceptions", ()) or (): - found = _unwrap_auth_error(nested) + found = _unwrap_proxy_error(nested) if found is not None: return found return None @@ -218,13 +224,13 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool try: anyio.run(_run, url, workspace, profile) - except BaseException as exc: # noqa: BLE001 - re-raised unless it's an auth failure - # The token can still expire mid-session; report that the same way - # rather than letting the ExceptionGroup surface as a hang or traceback. - auth_error = _unwrap_auth_error(exc) - if auth_error is None: + except BaseException as exc: # noqa: BLE001 - re-raised unless it's a known proxy failure + # Errors raised inside the transport arrive wrapped by its task group. + # Report expected auth/transport failures without hiding programming bugs. + proxy_error = _unwrap_proxy_error(exc) + if proxy_error is None: raise - _fail_fast(str(auth_error)) + _fail_fast(str(proxy_error)) -__all__ = ["AUTH_FAILURE_EXIT_CODE", "ProxyAuthError", "serve"] +__all__ = ["AUTH_FAILURE_EXIT_CODE", "ProxyAuthError", "ProxyTransportError", "serve"] diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 6c8fbda2..835aeda3 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -166,15 +166,30 @@ async def scenario() -> bool: assert anyio.run(scenario) is True - def test_raises_transport_errors(self): + def test_client_errors_are_forwarded(self): + async def scenario() -> Exception: + src_send, src_recv = anyio.create_memory_object_stream(1) + dst_send, dst_recv = anyio.create_memory_object_stream(1) + error = ValueError("malformed client message") + await src_send.send(error) + await src_send.aclose() + + await mcp_proxy._pump(src_recv, dst_send) + return await dst_recv.receive() + + error = anyio.run(scenario) + assert isinstance(error, ValueError) + assert str(error) == "malformed client message" + + def test_upstream_errors_are_raised(self): async def scenario() -> None: src_send, src_recv = anyio.create_memory_object_stream(1) dst_send, _ = anyio.create_memory_object_stream(1) await src_send.send(httpx.ReadTimeout("upstream timed out")) await src_send.aclose() - with pytest.raises(httpx.ReadTimeout, match="upstream timed out"): - await mcp_proxy._pump(src_recv, dst_send) + with pytest.raises(mcp_proxy.ProxyTransportError, match="upstream timed out"): + await mcp_proxy._pump_upstream(src_recv, dst_send) anyio.run(scenario) @@ -184,7 +199,7 @@ async def scenario() -> None: dst_send, _ = anyio.create_memory_object_stream(1) await src_send.aclose() - with pytest.raises(RuntimeError, match="upstream MCP transport closed"): + with pytest.raises(mcp_proxy.ProxyTransportError, match="upstream MCP transport closed"): await mcp_proxy._pump_upstream(src_recv, dst_send) anyio.run(scenario) @@ -345,8 +360,26 @@ def raise_group(func, *args): assert excinfo.value.code == mcp_proxy.AUTH_FAILURE_EXIT_CODE assert "token expired" in capsys.readouterr().err + def test_transport_failure_exits_with_a_one_line_message(self, monkeypatch, capsys): + def raise_group(func, *args): + raise BaseExceptionGroup( + "transport", + [mcp_proxy.ProxyTransportError("upstream MCP transport closed unexpectedly")], + ) + + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr(mcp_proxy.anyio, "run", raise_group) + + with pytest.raises(SystemExit) as excinfo: + mcp_proxy.serve(URL, WS, "p") + + assert excinfo.value.code == mcp_proxy.AUTH_FAILURE_EXIT_CODE + captured = capsys.readouterr() + assert captured.err == "ucode mcp-proxy: upstream MCP transport closed unexpectedly\n" + assert captured.out == "" + def test_non_auth_failures_still_propagate(self, monkeypatch): - # Only auth failures are converted to a clean exit; genuine transport + # Only expected proxy failures are converted to a clean exit; programming # bugs must keep their traceback so they stay debuggable. def raise_other(func, *args): raise ValueError("some transport bug") From b80d1bd23b277196d70fae2930ee28c6dd5a8683 Mon Sep 17 00:00:00 2001 From: Anthony Ivan Date: Fri, 4 Sep 2026 00:08:39 +0800 Subject: [PATCH 3/3] Format MCP proxy test Signed-off-by: Anthony Ivan --- tests/test_mcp_proxy.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 835aeda3..bf7a3536 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -199,7 +199,9 @@ async def scenario() -> None: dst_send, _ = anyio.create_memory_object_stream(1) await src_send.aclose() - with pytest.raises(mcp_proxy.ProxyTransportError, match="upstream MCP transport closed"): + with pytest.raises( + mcp_proxy.ProxyTransportError, match="upstream MCP transport closed" + ): await mcp_proxy._pump_upstream(src_recv, dst_send) anyio.run(scenario)