From 421394bc5c708b4b12e965fcd48a301d17cfefba Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Mon, 3 Aug 2026 18:39:26 +0000 Subject: [PATCH] mcp-proxy: fail fast on dead Databricks auth instead of hanging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the Databricks CLI can't mint a token (expired refresh token, logged-out profile), `get_databricks_token` raised a RuntimeError from inside `httpx.Auth.auth_flow` — i.e. inside the streamable-HTTP transport's anyio task group. Rather than surfacing as an error, that stalled the process until the MCP client's startup timeout fired (~30s in Codex/Claude), so the user saw: MCP client for `foo` timed out after 30 seconds. Add or adjust `startup_timeout_sec` in your config.toml with no hint that the real problem was expired auth. Because every registered server shares the profile, all of them failed at once — and raising the timeout never helps, since the proxy never answers. Measured on a workspace whose refresh token had expired: the proxy hung >60s (15s token fetch + 30s non-interactive re-auth + 15s refetch, all inside the client's 30s budget). It now exits in ~1s with the CLI's own message. - Pre-flight the token in `serve()` before opening the bridge, so a dead profile is diagnosed up front instead of from inside the transport. - Translate token failures in `auth_flow` into a terminal `ProxyAuthError`, and unwrap it from anyio ExceptionGroups so a mid-session expiry reports too. - Report on stderr (stdout is the MCP wire) and exit AUTH_FAILURE_EXIT_CODE=2; MCP clients surface a child's stderr and non-zero exit far better than a hang. - Non-auth failures still propagate with their traceback, and KeyboardInterrupt still works (verified). Co-authored-by: Isaac --- src/ucode/mcp_proxy.py | 94 ++++++++++++++++++++++++++++++++-- tests/test_mcp_proxy.py | 109 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 5 deletions(-) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 1538accf..1f9e5055 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -13,10 +13,20 @@ ${OAUTH_TOKEN}``, Claude ``headersHelper``, Cursor literal-token rewrites): one uniform mechanism, token refresh in a single place, and the proxy is an invisible implementation detail baked into each client's config. + +Auth failures are terminal and are reported *fast*. When the Databricks CLI +can't mint a token (expired refresh token, logged-out profile), the proxy prints +the CLI's own message to stderr and exits ``AUTH_FAILURE_EXIT_CODE`` rather than +letting the client wait out its MCP startup timeout. Every server registered +against the same profile fails at once in that state, so a silent hang is +especially confusing -- the user needs to be told to re-run +``databricks auth login``. """ from __future__ import annotations +import sys + import anyio import httpx from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream @@ -25,6 +35,28 @@ from ucode.databricks import 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. +AUTH_FAILURE_EXIT_CODE = 2 + + +class ProxyAuthError(RuntimeError): + """The proxy could not mint a Databricks token, so it cannot serve requests. + + Kept distinct from a transport error: this one is terminal and actionable + (the user must re-run `databricks auth login`), so `serve` reports it on + stderr and exits rather than retrying.""" + + +def _fail_fast(message: str) -> None: + """Report a terminal auth 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.""" + print(f"ucode mcp-proxy: {message}", file=sys.stderr, flush=True) + raise SystemExit(AUTH_FAILURE_EXIT_CODE) + class _DatabricksTokenAuth(httpx.Auth): """Injects a fresh Databricks OAuth bearer on every request. @@ -41,7 +73,14 @@ def __init__(self, workspace: str, profile: str | None, *, use_pat: bool) -> Non def auth_flow(self, request: httpx.Request): # get_databricks_token honors the DATABRICKS_BEARER short-circuit and PAT # profiles internally; --use-pat is surfaced via the env ucode already set. - token = get_databricks_token(self._workspace, self._profile) + # A RuntimeError here means auth is dead (expired refresh token, logged-out + # profile). Raising it from inside httpx's auth_flow would tear through the + # transport's task group and stall the process until the client times out, + # so translate it into a terminal ProxyAuthError the caller reports cleanly. + try: + token = get_databricks_token(self._workspace, self._profile) + except RuntimeError as exc: + raise ProxyAuthError(str(exc)) from exc request.headers["Authorization"] = f"Bearer {token}" yield request @@ -69,9 +108,54 @@ async def _run(url: str, workspace: str, profile: str | None, use_pat: bool) -> tg.start_soon(_pump, http_read, stdio_write) -def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool = False) -> None: - """Run the stdio<->streamable-HTTP MCP proxy until the client closes stdin.""" - anyio.run(_run, url, workspace, profile, use_pat) +def _preflight_token(workspace: str, profile: str | None) -> None: + """Verify a Databricks token can be minted before opening the bridge. + + Raises ``RuntimeError`` (with the CLI's own message) when auth is dead. This + is a plain synchronous call: ``get_databricks_token`` already bounds itself + with subprocess timeouts, so it returns or fails on its own — the point here + is only to *locate* the failure before the transport starts, where it can be + reported instead of stalling the session.""" + get_databricks_token(workspace, profile) + +def _unwrap_auth_error(exc: BaseException) -> ProxyAuthError | None: + """Find a ProxyAuthError anywhere in an exception (or ExceptionGroup) tree. -__all__ = ["serve"] + 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): + return exc + for nested in getattr(exc, "exceptions", ()) or (): + found = _unwrap_auth_error(nested) + if found is not None: + return found + return None + + +def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool = False) -> None: + """Run the stdio<->streamable-HTTP MCP proxy until the client closes stdin. + + Authentication is checked up front: a dead profile is a terminal condition, + and failing here (fast, with the CLI's own message) is far better than + letting the client wait out its MCP startup timeout with no explanation.""" + # Pre-flight the token before opening the bridge. Without this, the first + # token failure surfaces from inside the transport's task group, where it can + # stall the process instead of erroring out. + try: + _preflight_token(workspace, profile) + except RuntimeError as exc: + _fail_fast(str(exc)) + + try: + anyio.run(_run, url, workspace, profile, use_pat) + 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: + raise + _fail_fast(str(auth_error)) + + +__all__ = ["AUTH_FAILURE_EXIT_CODE", "ProxyAuthError", "serve"] diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 0ec17dd3..b7233be4 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -60,6 +60,19 @@ def test_auth_flow_yields_the_same_request(self, monkeypatch): assert yielded == [request] + def test_dead_auth_becomes_a_terminal_proxy_auth_error(self, monkeypatch): + # A raw RuntimeError escaping auth_flow tears through httpx's transport + # task group and stalls the proxy until the client's startup timeout. + # Translating it keeps the failure reportable by `serve`. + def boom(ws, profile): + raise RuntimeError("no access token; run `databricks auth login`") + + monkeypatch.setattr(mcp_proxy, "get_databricks_token", boom) + auth = mcp_proxy._DatabricksTokenAuth(WS, "p", use_pat=False) + + with pytest.raises(mcp_proxy.ProxyAuthError, match="databricks auth login"): + list(auth.auth_flow(httpx.Request("POST", URL))) + class TestPump: def test_forwards_all_messages_in_order(self): @@ -108,6 +121,7 @@ def fake_run(func, *args): captured["func"] = func captured["args"] = args + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) monkeypatch.setattr(mcp_proxy.anyio, "run", fake_run) mcp_proxy.serve(URL, WS, "uc-dogfood", use_pat=True) @@ -117,8 +131,103 @@ def fake_run(func, *args): def test_defaults_profile_none_and_use_pat_false(self, monkeypatch): captured: dict = {} + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: captured.update(args=args)) mcp_proxy.serve(URL, WS) assert captured["args"] == (URL, WS, None, False) + + def test_preflights_auth_before_opening_the_bridge(self, monkeypatch): + # Order matters: a dead profile must be caught before the stdio bridge + # starts, so the failure is a fast exit rather than a stalled session. + order: list[str] = [] + monkeypatch.setattr( + mcp_proxy, "_preflight_token", lambda ws, profile: order.append("preflight") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: order.append("bridge")) + + mcp_proxy.serve(URL, WS, "p") + + assert order == ["preflight", "bridge"] + + def test_dead_auth_exits_fast_without_starting_the_bridge(self, monkeypatch, capsys): + # The regression this fix targets: previously the token failure surfaced + # from inside the transport and the proxy hung until the MCP client's + # startup timeout (~30s) with no explanation. + started: list[str] = [] + + def dead_auth(ws, profile): + raise RuntimeError("no access token for " + ws + "; run `databricks auth login`") + + monkeypatch.setattr(mcp_proxy, "_preflight_token", dead_auth) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: started.append("bridge")) + + with pytest.raises(SystemExit) as excinfo: + mcp_proxy.serve(URL, WS, "p") + + assert excinfo.value.code == mcp_proxy.AUTH_FAILURE_EXIT_CODE + assert started == [] # the bridge never opened + # Diagnostics go to stderr; stdout is the MCP wire and must stay clean. + captured = capsys.readouterr() + assert "databricks auth login" in captured.err + assert captured.out == "" + + def test_auth_expiring_mid_session_exits_with_the_actionable_message(self, monkeypatch, capsys): + # A ProxyAuthError raised once the bridge is running arrives wrapped in an + # anyio ExceptionGroup; it must still be reported, not surface as a crash. + def raise_group(func, *args): + raise BaseExceptionGroup( + "transport", + [mcp_proxy.ProxyAuthError("token expired; run `databricks auth login`")], + ) + + 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 + assert "token expired" in capsys.readouterr().err + + def test_non_auth_failures_still_propagate(self, monkeypatch): + # Only auth failures are converted to a clean exit; genuine transport + # bugs must keep their traceback so they stay debuggable. + def raise_other(func, *args): + raise ValueError("some transport bug") + + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr(mcp_proxy.anyio, "run", raise_other) + + with pytest.raises(ValueError, match="some transport bug"): + mcp_proxy.serve(URL, WS, "p") + + +class TestPreflightToken: + def test_passes_through_when_a_token_is_available(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + mcp_proxy._preflight_token(WS, "p") # no exception + + def test_surfaces_the_cli_error_message(self, monkeypatch): + def boom(ws, profile): + raise RuntimeError("profile is stale; run `databricks auth logout`") + + monkeypatch.setattr(mcp_proxy, "get_databricks_token", boom) + + with pytest.raises(RuntimeError, match="databricks auth logout"): + mcp_proxy._preflight_token(WS, "p") + + def test_checks_the_same_workspace_and_profile_the_bridge_will_use(self, monkeypatch): + # The preflight must validate the exact credentials the request-time auth + # hook uses, or it could pass while the bridge still fails. + calls: list[tuple[str, str | None]] = [] + monkeypatch.setattr( + mcp_proxy, + "get_databricks_token", + lambda ws, profile: calls.append((ws, profile)) or "tok", + ) + + mcp_proxy._preflight_token(WS, "myprofile") + + assert calls == [(WS, "myprofile")]