Memoize Databricks CLI tokens per process and skip impossible `--no-browser` re-auth by yaoharry · Pull Request #1017 · databricks/unity-gateway

yaoharry · 2026-10-06T22:58:18Z

What did you change, and why?

Why: On a host whose first-matching ~/.databrickscfg profile has an invalid refresh token, every get_databricks_token call paid auth token (fail) -> auth login --no-browser -> auth token (fail). Current CLIs (v1.17.0 and v1.19.0) reject --no-browser with Error: unknown flag: --no-browser, so the re-auth never did anything.

Change:

get_databricks_token memoizes per process, keyed by (workspace, resolved profile). A minted token is reused until 5 minutes before its expires_in/expiry, taking the earlier of the two. A token whose JSON has no expiry is reused for 60s. A definitive failure is memoized for 60s and re-raised with the same message, including the stale-profile logout/login hint. Timeouts and token-cache lock contention are not memoized. force_refresh=True skips the memo read and never stores a failure, so a failed forced retry cannot hide a token other callers still use. Nothing persists across processes. The DATABRICKS_BEARER / DATABRICKS_BEARER_COMMAND short-circuits are unchanged.
The short TTLs for no-expiry tokens and failures let long-lived processes (MCP proxy, web-search MCP, relay proxy) pick up a re-login done in another terminal within a minute.
has_valid_databricks_auth shares memoized successes, not failures, so the decision to offer an interactive login always uses a real check after a cached failure. Regular and custom-OAuth CLI logins clear the memo.
--no-browser re-auth runs only when a successful databricks auth login --help probe lists the flag. The probe runs lazily, once per CLI binary per process, only after a token fetch fails. If help or re-auth fails, the unnecessary login or second token fetch is skipped.
Concurrent callers coordinate token fetches per workspace/profile, and cache invalidation prevents an old in-flight result from restoring a cleared memo. Both token-fetch and re-auth timeouts remain transient.
The port adds coverage for failed help probes, refresh/expiry edge cases, and concurrent requests, and updates the test coverage notes without changing credential routing.

Why the stale generic profile is used at all (not changed here)

This is not a fallback order. The two profiles serve different callers:

configure_shared_state always sets state["profile"] with find_profile_name_for_host, which returns the first CLI profile whose host matches. That happens even in custom-OAuth CLI mode (ENABLE_CUSTOM_OAUTH_FROM_CLI=1), and in the measured setup it picks the stale profile.
The custom-client profile (ug-oauth-<host>-<client-id>) is stored separately in state["custom_oauth"]["profile"]. Only these paths use it: Claude's apiKeyHelper (what the gateway actually sees), the preflight probe in configure_shared_state, smart routing's _launch_token, and the web-search MCP (agents/claude.py already prefers it for the web-search profile).
On this launch, both stale-profile calls come from refresh_managed_config (managed_config.py), which uses state.get("profile"). MCP discovery, skills download and similar callers use the same source. When that token fails, refresh_managed_config quietly falls back to the persisted config, or to none. So a workspace's admin-published managed config is currently not applied for a user whose generic profile is stale, even though their custom-client login is healthy.

Changing which credential these callers use has auth and governance implications, so this PR leaves it alone.

Open question for reviewers: which of these, if any, should a follow-up do? I don't have enough context on the intended identity and scope model for custom OAuth clients to choose.

A. In custom-OAuth CLI mode, use the custom-client profile for state["profile"] callers (managed config, MCP discovery, skills). This matches the identity the gateway already sees, removes the remaining stale call (~115ms), and makes managed config apply. Risks: those workspace APIs would be called with a token from a different OAuth app, with whatever --scopes the user chose. Narrower scopes would turn into 403s on MCP/UC/skills. If the two profiles belong to different principals, the caller's identity changes. A narrower variant applies this only to refresh_managed_config.
B. Use the first source that works: try the generic profile, and fall back to the custom-client profile on failure. Behavior stays the same when the generic profile is healthy. But identity and scopes then switch silently per call depending on token health, and the stale call is still paid once.
C. Keep the current routing (this PR). The stale profile is tried once per process (~115ms), and managed config is skipped for these users as it is today.

How do you know it works?

Verification on current main (4d8ff8d):

442 focused tests pass across tests/test_databricks.py,
tests/test_mcp_web_search.py, tests/test_custom_oauth.py, and
tests/test_integration_contract.py, including strengthened threaded regressions.
The full clean-environment uv run --frozen pytest -q result is 3,396 passed,
42 skipped, 4 failed. The same four installed-agent capture tests fail on
unmodified main (3,375 passed, 42 skipped, 4 failed): both the User-Agent
and managed-header cases for Claude and Codex in test_e2e_user_agent.py.
This host's Claude rejects its machine-wide model selection, and its Codex
uses the configured gateway rather than the local capture server. No new
full-suite failures were introduced.
Representative new tests fail behaviorally against unmodified main:
three stale-profile calls produce six token fetches instead of one, and a
repeated successful fetch returns a new token instead of the memoized token.
Ruff check, Ruff format check, ty check src/, and git diff --check pass.
Live workspace tests and launch-time measurements have not been run for this port.

The original author measured the original patch on the earlier base 3108c95 with
a stale generic profile before a healthy custom-client profile (Databricks CLI
v1.19.0). Their handoff reports 9 -> 3 databricks auth subprocesses, 593ms -> 152ms
CLI wall time, and 1.89s -> 1.59s median time to the final claude exec. These
measurements have not been repeated here.

This pull request and its description were written by Isaac.

…owser re-auth A ug claude launch calls get_databricks_token from many places, and a host whose first-matching CLI profile has a stale refresh token paid fail -> auth login --no-browser -> fail on every call. Current CLIs reject --no-browser, so the re-auth never did anything. Memoize tokens per workspace and resolved profile until five minutes before their earliest stated expiry, or for 60 seconds without expiry. Cache definitive failures briefly, but never cache token-fetch or re-auth timeouts and lock contention. Forced refresh bypasses the memo and does not replace a success with failure. Auth validation shares successful tokens and logins clear the memo. Only attempt non-interactive re-auth when the CLI advertises --no-browser, and skip the second fetch after a failed re-auth. Ported from databricks-eng/universe#2742026 onto Unity Gateway main at 4d8ff8d. Coordinate concurrent fetches per key, reject in-flight memo writes after login invalidation, and cover failed help probes, expiry, refresh, and concurrent callers without changing credential/profile routing. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Harry Yao <harry.yao@databricks.com>

david-siqi-liu

Thanks for porting this, the latency win is worth having. Two things I'd like addressed before this gets pinned in Isaac, inline below: transient CLI failures getting memoized as definitive, and the monotonic-clock memo deadline across laptop suspend. Both are small fixes.

One more that I couldn't anchor inline because the file isn't in the diff: run_connection_login in src/ucode/mcp_connection_login.py (~L119) runs databricks auth login --resource ... --profile P but doesn't clear the token memo, unlike the other login paths. mcp_proxy's post-login _mint retry is non-forced, so it can get the pre-login token back (e.g. if the browser login switched accounts), and the once-per-session login guard prevents another attempt. Suggest calling clear_databricks_token_cache() in a finally there. Non-blocking.

Require explicit invalid or missing credential diagnostics before negative caching, expire memos when either wall or monotonic time reaches its deadline, and clear the shared token memo after every attempted MCP connection login. Add regressions for transient failure recovery, suspend and clock rollback, login failure paths, and fresh-token use on sync and async proxy retries. Update the component coverage notes. Signed-off-by: Harry Yao <harry.yao@databricks.com>

yaoharry · 2026-10-08T06:07:20Z

Also addressed the non-inline connection-login note in ef1d269: run_connection_login clears the shared token/failure memo in finally after every attempted login. Regression coverage includes success, nonzero exit, timeout, missing-binary errors, and fresh-token use on both sync and async proxy retries. Unsupported-CLI prechecks leave the existing memo alone.

Validation:

Focused auth/MCP component suite: 467 passed.
Ruff check, Ruff format check, and ty: passed.
CI-equivalent unit command (uv run --frozen pytest --ignore=tests/test_e2e.py): 3,427 passed, 4 failed, 9 skipped. The four failures are the existing Claude/Codex header-capture tests in tests/test_e2e_user_agent.py; all four also fail on a clean, unchanged checkout of the previous PR head, 397722f, because no request reaches their local capture server in this environment.

Live Isaac/workspace E2E was not rerun for this feedback update.

Signed-off-by: Harry Yao <harry.yao@databricks.com> # Conflicts: # src/ucode/databricks.py # tests/test_databricks.py

github-actions · 2026-10-09T21:01:07Z

UG review

The change adds process-local credential caching, login invalidation, and capability-gated re-authentication. Long-lived processes still have recovery gaps around rejected cached tokens and transient capability-probe failures.

Major

Refresh cached tokens after rejection — src/ucode/databricks.py:1577: can we refresh cached tokens on non-connection MCP 401s so revocation followed by external re-login does not leave subsequent requests failing until the memo expires?
Retry inconclusive capability probes — src/ucode/databricks.py:1557: can we avoid caching probe timeouts as unsupported so one slow --help does not disable working --no-browser re-auth for the rest of the process?

Some patches were unavailable or truncated to fit the review context. Treat this as a partial review.

Automated advisory review of 7e5f8682d41c using Lilly's UG review rubric. It does not approve or block this PR.

After this PR, no workflow uses `pull_request_target`. Fork PRs get credentialed CI when an org member (`MEMBER` or `OWNER`) comments a command on the PR. Every comment-triggered run uses the default branch's workflow definitions. Design doc (option B): https://docs.google.com/document/d/1v-6znVBCb8YfGxr12SBd80w3i9n4U0fMEYohyz9Aoyw/edit ## Commands Org members comment these on a PR: - `/integration-test` runs the full Integration workflow with secrets on the PR's current head commit. - `/integration-test <sha>` does the same, but only if the head is still `<sha>` (7 to 40 hex characters). If the fork pushed since, it refuses and replies with the new head. - `/ug-review` runs the UG rubric review. - `/user-journey-check` runs the user-journey gate. For in-repo PRs it also runs automatically, so the command is mostly for forks. Each command posts a commit status and a comment with the run link on the PR. The statuses are `Fork integration` and `User Journey Test Required`. Each new push needs a new command. Anyone else who comments `/integration-test` gets a reply saying only org members can run integration tests. ## Before and after Live integration on fork PRs: - Before: every credentialed job skips, because fork PRs get no secrets on `pull_request`. That covers the E2E CUJs, workspace validation, and the smoke, full, and managed journeys. - After: a member reviews, then comments `/integration-test` or `/integration-test <sha>`. It runs `integration.yml` on the pinned commit with secrets. Live integration on in-repo PRs: - Before: runs automatically. - After: unchanged. `/integration-test` also works as an on-demand rerun, and it never cancels the automatic run. UG review: - Before: adding the `ug-review` label, via `pull_request_target`. - After: a member comments `/ug-review`. It still runs only default-branch code and never executes the PR. User Journey Test Required: - Before: `pull_request_target` on every PR. - After: `pull_request` for in-repo PRs, and `/user-journey-check` for fork PRs. It still reads the PR only through the API. The gate checks the head is unchanged across its diff reads and reports the fork-PR status only on that evaluated commit, so a push during the check never inherits the old verdict. `pull_request_target`: - Before: used by `ug-review.yml` and `user-journey-test-required.yml`. - After: not used anywhere, and a contract test fails if it comes back. No admin event-policy exception is needed before GitHub blocks it on public repos on 2026-11-02. ## Fork safety in `integration.yml` - It takes an optional `ref` input. When the input is empty, existing callers behave exactly as before. - Every checkout pins `ref` and sets `persist-credentials: false`. - setup-uv skips its cache for ref-pinned runs, so fork builds cannot poison it. - Concurrency groups key on the PR number, so runs for different PRs never cancel each other. Concurrency is set per job, so ordinary PR comments never cancel a run. `AGENTS.md` now says workflows must not use `pull_request_target`. The new `fork-integration.yml` call to `integration.yml` carries the same inline `zizmor: ignore[secrets-inherit]` as `ci.yml` (databricks#1089). ## Testing - Contract tests cover the triggers, the member-only gate, and the authorize script, run against a fake `gh`. They also check the non-member reply, the pinned credential-free checkouts, the disabled caches, and the status reports. - `issue_comment` workflows only run from the default branch, so the first live run happens after merge. The plan is to comment `/integration-test` on a fork PR such as databricks#1017 or databricks#1018, and `/ug-review` on any PR. OIDC note: under `issue_comment`, the OIDC subject is the default branch, not a PR ref. If the JFrog trust policy for the Windows installation lane keys on PR refs, that one lane fails on fork runs. The other lanes use secrets, not OIDC. This pull request and its description were written by Isaac. --------- Co-authored-by: Isaac <no-reply@databricks.com>

david-siqi-liu · 2026-10-12T17:08:14Z

/integration-test

…icks#1101) After this PR, `/integration-test` works on fork PRs. Its three small jobs (authorize, non-member reply, result report) run on `databricks-protected-runner-group`, the same runners `ug-review.yml` and the user-journey gate already use. ## Before and after `/integration-test` comment: - Before: the authorize job fails right away with HTTP 403 ("the `databricks` organization has an IP allow list enabled"). GitHub-hosted `ubuntu-latest` runners are not on the org allow list, so their GitHub API calls are rejected. Nothing replies on the PR and no status is posted. First seen on databricks#1017. - After: the jobs run on allowlisted runners, so the command posts the `Fork integration` status, runs Integration on the pinned commit, and reports back. Reading the PR head in the authorize step: - Before: `gh api` piped into a standalone `jq`. - After: one `gh api --jq` call. The protected runner uses a custom image (`ubuntu-22-04-crowdstrike-image`) that does not guarantee a `jq` binary, and gh has jq built in. ## Testing - `tests/test_integration_contract.py` (72 passed). The fake `gh` in the authorize-script test answers only the exact `--jq` query the step sends, so changing that expression fails 6 cases. A new test fails if any fork-integration job loses the protected runner group or its `linux-ubuntu-latest` label (a misspelled label fails it), goes back to `ubuntu-latest`, or calls a bare `jq`. - Read-only probe on `databricks-protected-runner-group` (scratch branch, since deleted): `gh` 2.102.0 is installed, and the exact `gh api --jq` line read databricks#1017's head (`f2db89b`, `yaoharry/unity-gateway`) through the IP allow list with the job token. - The rest of the chain is unchanged and already proven. The failed databricks#1017 run shows the member gate admits a `MEMBER` comment. The `integration` call grants the same permissions as `ci.yml`'s. A fork head commit is fetchable by SHA from this repo. Windows JFrog OIDC accepts the `refs/heads/main` subject that `issue_comment` runs get, as every main push run shows. - `issue_comment` workflows only run from the default branch, so the end-to-end check is commenting `/integration-test` on databricks#1017 after this merges. This pull request and its description were written by Isaac. --------- Co-authored-by: Isaac <no-reply@databricks.com>

david-siqi-liu · 2026-10-12T19:11:57Z

/integration-test

github-actions · 2026-10-12T19:45:29Z

Fork integration failure on f2db89b: https://github.com/databricks/unity-gateway/actions/runs/38296038771

github-actions · 2026-10-12T19:48:50Z

Fork integration success on f2db89b: https://github.com/databricks/unity-gateway/actions/runs/38296038771

yaoharry marked this pull request as ready for review October 6, 2026 23:57

lilly-luo added the ug-review Run the automated UG review label Oct 7, 2026

lilly-luo requested a review from david-siqi-liu October 7, 2026 05:57

david-siqi-liu reviewed Oct 7, 2026

View reviewed changes

Comment thread src/ucode/databricks.py

Comment thread src/ucode/databricks.py Outdated

This was referenced Oct 8, 2026

[test, do not merge] Fork integration check on #1017 #1045

Closed

Move fork-PR credentialed CI off pull_request_target #1044

Merged

david-siqi-liu approved these changes Oct 9, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into codex/pr1017

7e5f868

Signed-off-by: Harry Yao <harry.yao@databricks.com> # Conflicts: # src/ucode/databricks.py # tests/test_databricks.py

Merge remote-tracking branch 'ugh/main' into pr-1017-update

f2db89b

david-siqi-liu mentioned this pull request Oct 12, 2026

Run fork integration's GitHub API jobs on allowlisted runners #1101

Merged

david-siqi-liu added this pull request to the merge queue Oct 12, 2026

Merged via the queue into databricks:main with commit 0d52d5b Oct 12, 2026
28 checks passed

Navigation Menu

Repository navigation

Memoize Databricks CLI tokens per process and skip impossible `--no-browser` re-auth - #1017

Memoize Databricks CLI tokens per process and skip impossible `--no-browser` re-auth#1017
david-siqi-liu merged 4 commits into
databricks:mainfrom
yaoharry:harryyao/port-token-memoization

yaoharry commented Oct 6, 2026 •

edited

Loading

Uh oh!

david-siqi-liu left a comment

Uh oh!

Uh oh!

Uh oh!

yaoharry commented Oct 8, 2026

Uh oh!

github-actions Bot commented Oct 9, 2026

Uh oh!

david-siqi-liu commented Oct 12, 2026

Uh oh!

david-siqi-liu commented Oct 12, 2026

Uh oh!

github-actions Bot commented Oct 12, 2026

Uh oh!

github-actions Bot commented Oct 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

yaoharry commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What did you change, and why?

Why the stale generic profile is used at all (not changed here)

How do you know it works?

Uh oh!

david-siqi-liu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yaoharry commented Oct 8, 2026

Uh oh!

github-actions Bot commented Oct 9, 2026

UG review

Major

Uh oh!

david-siqi-liu commented Oct 12, 2026

Uh oh!

david-siqi-liu commented Oct 12, 2026

Uh oh!

github-actions Bot commented Oct 12, 2026

Uh oh!

github-actions Bot commented Oct 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yaoharry commented Oct 6, 2026 •

edited

Loading