mcp-proxy: fail fast on dead Databricks auth instead of hanging

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

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