Fix relayed-auth token refresh: lock contention, resilience, and 401 self-heal - #349
Merged
Conversation
…self-heal
Relayed (Claude Max/Team/Enterprise) sessions were intermittently lagging out
after ~1h, sometimes needing a manual `databricks auth login` and sometimes a
spurious Anthropic re-auth. Diagnostics traced this to the shared Databricks
token cache: many ucode helper processes (the gateway proxy, one mcp-proxy per
MCP server, web-search, tracing) each mint tokens against one lock-guarded
~/.databricks/token-cache.json, and refreshing the single expiring token at once
loses the CLI's cache-write lock ("cache update: exit status 45" -> "run
databricks auth login"). The 30-min force-refresh loop and swallowed errors made
it worse and invisible.
Fixes (all relayed-path only; non-relayed launch unchanged):
- databricks.py: get_databricks_token retries transient token-cache lock
contention with jittered backoff instead of treating it as a dead session.
Benefits every ucode process. Verified: 8-way concurrency goes 6/8 -> 8/8.
- gateway_proxy.py: refresh loop is now expiry-aware (JWT exp) and non-force, so
it refreshes ~once/hour near expiry rather than writing the cache every 30 min;
refreshes are single-flighted so a request burst at expiry triggers one mint.
- gateway_proxy.py: refresher thread can no longer be killed by a stray
non-RuntimeError, and refresh failures are surfaced instead of swallowed.
- gateway_proxy.py: lazy refresh on the request path survives laptop sleep
(which freezes the interval timer).
- gateway_proxy.py: retry-on-401 force-refreshes the swap token and retries once;
a 401 that survives is genuinely the Anthropic layer, so Claude Code's re-auth
becomes correct instead of a spurious prompt.
Also included (earlier proxy-robustness pass): pooled keep-alive httpx client
(no per-request TLS handshake), mid-stream upstream error handling, and explicit
per-op upstream timeouts.
scripts/diagnose_auth.py: safe-by-default diagnostic that reproduces and confirms
the token TTL, duplicate cache keys, rotation behavior, and lock-contention race.
Co-authored-by: Isaac
lilly-luo
approved these changes
Aug 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Relayed-auth (Claude Max/Team/Enterprise) sessions intermittently lag out after ~1h, sometimes needing a manual
databricks auth login, sometimes triggering a spurious Anthropic re-auth.Root cause
A single relayed session spawns many ucode helper processes (gateway proxy, one mcp-proxy per MCP server, web-search, tracing), and each refreshes the Databricks token against the one shared, lock-guarded
~/.databricks/token-cache.json. When the token expires, they collide on the CLI's cache-write lock; the losers fail withcache update: exit status 45— whose own message is "rundatabricks auth login."Confirmed with
scripts/diagnose_auth.py: tokens live 60 min, refresh tokens are reusable (so retrying is safe), and 8 concurrent refreshes fail 6/8.Fixes (relayed path only; non-relayed launch unchanged)
databricks.py) — the core fix.get_databricks_tokenretries the transient lock error instead of treating it as a dead session. Verified: 6/8 → 8/8 under concurrency.gateway_proxy.py) — refresh only near expiry, single-flighted, so processes stop writing the cache every 30 min.Testing
New tests for the lock retry, JWT expiry parsing, token-cache freshness/single-flight/thread-survival, and retry-on-401. Full suite: 1859 passed (2 pre-existing failures unrelated to this change). Lint + format clean.
scripts/diagnose_auth.pyis included as a safe-by-default diagnostic to reproduce this or run in a session overnight (--watch).This pull request and its description were written by Isaac.