Skip to content

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

Merged
sunishsheth2009 merged 1 commit into
databricks:mainfrom
sunishsheth2009:sunish-sheth_data/mcp-proxy-fail-fast
Aug 4, 2026
Merged

mcp-proxy: fail fast on dead Databricks auth instead of hanging#259
sunishsheth2009 merged 1 commit into
databricks:mainfrom
sunishsheth2009:sunish-sheth_data/mcp-proxy-fail-fast

Conversation

@sunishsheth2009

Copy link
Copy Markdown
Collaborator

Summary

When the Databricks CLI can't mint a token (expired refresh token, logged-out profile), ucode mcp-proxy hung until the MCP client's startup timeout instead of reporting the problem. Users saw this, with no hint that the real cause was expired auth:

MCP client for `foo` timed out after 30 seconds.
  Add or adjust `startup_timeout_sec` in your config.toml

Because every registered MCP server shares the same Databricks profile, all of them fail at once in that state — and raising startup_timeout_sec never helps, because the proxy never answers at all.

Root cause

get_databricks_token raises RuntimeError when auth is dead. The proxy called it inside httpx.Auth.auth_flow, i.e. inside the streamable-HTTP transport's anyio task group. The exception tore through the task group instead of surfacing as a clean failure, and the process stalled.

Worse, the failure path is slow by construction: get_databricks_token does a 15s token fetch → 30s non-interactive auth login retry → 15s refetch. That's up to 60s, which structurally cannot fit inside a client's 30s startup budget.

Measured before/after

Reproduced against a workspace whose OAuth refresh token had expired:

Before After
Time to fail >60s (hang) ~1s
Exit code 0 2 (AUTH_FAILURE_EXIT_CODE)
What the user sees generic client timeout the CLI's own message + databricks auth login commands

Actual stderr now:

ucode mcp-proxy: Databricks CLI returned no access token for https://…
Run `databricks auth login` to re-authenticate. The saved Databricks CLI profile
may be stale or invalid. Try:
  databricks auth logout --profile <p>
  databricks auth login --host https://… --profile <p>

Changes

  • Pre-flight the token in serve() before opening the bridge, so a dead profile is diagnosed up front rather than 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 is reported the same way.
  • Report on stderr and exit AUTH_FAILURE_EXIT_CODE (2) — stdout is the MCP wire, and MCP clients surface a child's stderr + non-zero exit far more usefully than a hang.
  • Non-auth failures still propagate with their traceback; KeyboardInterrupt still works (both covered by tests).

How do you know it works?

  • Verified against the real failure (expired ml-inference OAuth profile): previously >60s hang, now exits in ~1s with code 2, stdout empty, actionable stderr.
  • Happy path unaffected: preflight passes in 0s when a token is available.
  • Tests added in tests/test_mcp_proxy.py (16 pass, up from 8) covering: dead auth → ProxyAuthError; preflight runs before the bridge; dead auth exits fast without opening the bridge and keeps stdout clean; mid-session expiry inside an ExceptionGroup still reports; non-auth errors still propagate; preflight validates the same workspace/profile the bridge uses.
  • Note: two pre-existing TestServe tests only passed incidentally (their anyio.run stub no-op'd the token call); they now stub _preflight_token explicitly.
  • Full suite: 1116 passed; the only 2 failures are the pre-existing test_e2e_user_agent live-gateway network tests, which fail identically on main.
  • ruff format, ruff check, and ty check all clean.

This is a follow-up fix to #201 (the uniform mcp-proxy bridge).

This pull request and its description were written by Isaac.

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
@sunishsheth2009
sunishsheth2009 marked this pull request as ready for review August 3, 2026 22:13
@sunishsheth2009
sunishsheth2009 merged commit e87e2d4 into databricks:main Aug 4, 2026
2 checks passed
@sunishsheth2009
sunishsheth2009 deleted the sunish-sheth_data/mcp-proxy-fail-fast branch August 4, 2026 00:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants