diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index ee622884..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.""" @@ -117,7 +120,7 @@ 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.""" @@ -126,13 +129,35 @@ async def _pump( await dest.send(message) +async def _pump_upstream( + source: MemoryObjectReceiveStream, + dest: MemoryObjectSendStream, +) -> None: + """Forward upstream messages, failing if the transport closes first.""" + 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: 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 +166,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: @@ -156,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 @@ -198,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 a10f8a67..bf7a3536 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,79 @@ async def scenario() -> bool: assert anyio.run(scenario) is True + 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(mcp_proxy.ProxyTransportError, match="upstream timed out"): + await mcp_proxy._pump_upstream(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( + mcp_proxy.ProxyTransportError, 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): @@ -288,8 +362,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")