Handle MCP transport timeouts and disconnects

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"]

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")