Skip to content

Add force-refresh path for cached U2M tokens - #1552

Merged
mihaimitrea-db merged 8 commits into
mainfrom
mihaimitrea-db/token-force-refresh
Mar 19, 2026
Merged

mihaimitrea-db merged 8 commits into
mainfrom
mihaimitrea-db/token-force-refresh

Conversation

@mihaimitrea-db

@mihaimitrea-db mihaimitrea-db commented Mar 17, 2026 •

Copy link
Copy Markdown
Contributor

Summary

This PR adds an explicit PersistentAuth.ForceRefreshToken() path for cached U2M OAuth tokens so the CLI can implement databricks auth token --force-refresh on top of the 5-minute proactive refresh buffer introduced in #1535.

For maintainability reasons it also splits token loading from refresh decisions and documents a cross-process refresh race for a follow-up fix.

Why

#1535 changed the default U2M token path to proactively refresh tokens that are within 5 minutes of expiry so callers of databricks auth token are less likely to receive a token that expires before it can be used. That addresses the "give me a useful token" case, which came out of the auth token discussion in cli issue #4564.

That proactive refresh is intentionally best-effort, though. PersistentAuth.Token() still returns the existing access token when it is valid and a proactive refresh fails, because callers did not explicitly ask to block on a fresh token. That behavior is appropriate for the default path, but it is not strong enough for integrations that want to treat the CLI as a token minter or that want to manage their own cache/TTL policy. In those cases, "return a still-usable token" is different from "refresh now and give me a newly minted token or fail."

This PR adds that stricter SDK primitive. It gives the CLI a clean way to implement --force-refresh without changing the semantics of the existing Token() flow.

The raw oauth2 error text for forced refresh failures can be misleading. In particular, when a forced refresh is requested for a cached access token that is still valid but does not include a refresh token, the default oauth2 error suggests that the token is expired even though the real problem is that the cached token is not refresh-capable.

This PR surfaces explicit refresh-token errors instead of relying on the default oauth2 wording, so callers can distinguish "the current access token is expired" from "a new token could not be minted because the refresh token is missing or was rejected by the IdP".

What changed

Interface changes

  • Added (*PersistentAuth).ForceRefreshToken() (*oauth2.Token, error).
    This always attempts to refresh the cached U2M token, regardless of its current expiry, and redacts the refresh token from the returned value just like Token() does.

Behavioral changes

  • For Token(), the only behavioral change is that expired tokens with no refresh token now return ErrMissingRefreshToken instead of a less specific oauth2 library error. Callers can use errors.Is(err, ErrMissingRefreshToken) to detect this case.
  • PersistentAuth.ForceRefreshToken() always tries to refresh, even when the cached access token is still valid.
  • If forced refresh fails, ForceRefreshToken() returns an error instead of falling back to the existing token.
  • If the cached token does not contain a refresh token, ForceRefreshToken() fails explicitly.

Internal changes

  • Extracted cache lookup into loadToken() so token loading and refresh policy are separated.
  • Reused the existing refresh path so forced refresh still goes through the same endpoint resolution, token exchange, cache update, and refresh-token redaction logic as the normal path.
  • Added dedicated tests for force-refresh success, missing refresh token, transient refresh failure, invalid refresh token handling, and the profile-to-host fallback path.
  • Documented the remaining cross-process refresh race in refresh() with a follow-up TODO rather than expanding the scope of this PR to add inter-process coordination.
  • Added ErrMissingRefreshToken for the "cached token has no refresh token" case.
  • Added missing-refresh-token check into the shared refresh() path so both Token() and ForceRefreshToken() wrap the same semantic error.

How is this tested?

  • go test ./credentials/u2m
  • Added unit tests covering:
    • refreshing a still-valid cached token via ForceRefreshToken()
    • failing ForceRefreshToken() when the cached token has no refresh token
    • surfacing the exact wrapped error on refresh failure
    • preserving InvalidRefreshTokenError behavior
    • exercising forced refresh through the existing profile-key to host-key fallback path

Split token loading from refresh decisions, add dedicated force-refresh coverage, and document the remaining cross-process refresh race for a follow-up fix.

Signed-off-by: Mihai Mitrea <[email protected]>
Document the new PersistentAuth.ForceRefreshToken method in NEXT_CHANGELOG so the release notes cover the new explicit cached-token refresh behavior.

Signed-off-by: Mihai Mitrea <[email protected]>
Move the missing refresh token check into the shared refresh path, expose it as ErrMissingRefreshToken, and cover the wrapped error behavior from both Token and ForceRefreshToken.

Signed-off-by: Mihai Mitrea <[email protected]>

@simonfaltum simonfaltum left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review (automated, 2 agents)

Verdict: Approved

0 Critical | 0 Major | 1 Gap | 2 Nit | 1 Suggestion

See inline comments for details.

Comment thread credentials/u2m/persistent_auth.go Resolved
Comment thread credentials/u2m/persistent_auth.go Resolved
mihaimitrea-db

This comment was marked as outdated.

mihaimitrea-db

This comment was marked as outdated.

@mihaimitrea-db mihaimitrea-db left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated Review by 2 Agents

Verdict: Approve with comments

  • 🔴 0 Critical
  • 🟠 0 Major
  • 🟡 3 Minor
  • 🔵 3 Nit

The core logic changes look sound; the main concerns are a misclassified changelog entry, a missing test for the new ErrMissingRefreshToken graceful-degradation path, and a misleading doc comment on ForceRefreshToken().

See inline comments for details.

Comment thread NEXT_CHANGELOG.md Resolved
Comment thread credentials/u2m/persistent_auth_test.go Resolved
Comment thread credentials/u2m/error.go Resolved
Comment thread credentials/u2m/persistent_auth.go Resolved
Comment thread credentials/u2m/persistent_auth.go Resolved
Comment thread NEXT_CHANGELOG.md Resolved
Comment thread credentials/u2m/persistent_auth.go Resolved
- NEXT_CHANGELOG: Move ErrMissingRefreshToken to API Changes, remove
  proactive-refresh and ParentPath (rebase artifacts / unrelated)
- error.go: Fix InvalidRefreshTokenError doc (Load -> Token/ForceRefreshToken)
- persistent_auth.go: Clarify ForceRefreshToken dual-write vs Token();
  unify error prefix to 'token refresh:'; simplify loadToken doc
- persistent_auth_test.go: Add TestToken_ReturnsExistingTokenWhenNearExpiryAndNoRefreshToken;
  add defensive store t.Fatalf to ForceRefreshToken tests

Signed-off-by: Mihai Mitrea <[email protected]>

@simonfaltum simonfaltum left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review (automated, 3 agents: Claude Code + Isaac + Cursor, 2-round deep swarm)

Verdict: Approved (with suggestions)

0 Critical | 1 Major | 0 Gap | 2 Nit | 0 Suggestion

See inline comments for details. Overall the PR is well-designed: the loadToken() extraction is clean, ForceRefreshToken() has the right strict-error semantics, ErrMissingRefreshToken is idiomatic Go, and test coverage is thorough across all key paths.

Comment thread credentials/u2m/persistent_auth.go Resolved
Comment thread credentials/u2m/error.go Resolved
@simonfaltum

Copy link
Copy Markdown
Member

[Nit] PR description claims "no default behavior changed" for Token(), but it did

The PR description states "No default behavior changed for existing callers of PersistentAuth.Token()." The new ErrMissingRefreshToken check in refresh() changes the error returned when Token() is called with an expired token that has no refresh token: from a misleading oauth2 library error to "token refresh: cached token has no refresh token". Callers can now use errors.Is(err, ErrMissingRefreshToken). This is a net improvement, but should be acknowledged as a behavioral change.

Suggestion: Update the PR description, e.g.: "For Token(), the only behavioral change is that expired tokens with no refresh token now return ErrMissingRefreshToken instead of a misleading oauth2 library error."

Found by: Isaac (R1) + Claude Code (R1), confirmed by Cursor (R2)

…ration TODO

- Move ForceRefreshToken and ErrMissingRefreshToken changelog entries to
  "New Features and Improvements" (hand-written SDK features, not
  auto-generated API surface changes).
- Use "i.e." instead of "e.g." for exhaustive method list in changelog.
- Reword InvalidRefreshTokenError doc comment to remove overstated expiry
  precondition.
- Correct ForceRefreshToken doc comment to note Token() has the same
  migration behavior on refresh, and add TODO to thread provenance out of
  loadToken() in a follow-up.

Signed-off-by: Mihai Mitrea <[email protected]>
…force-refresh

# Conflicts:
#	NEXT_CHANGELOG.md
@github-actions

Copy link
Copy Markdown

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-go

Inputs:

  • PR number: 1552
  • Commit SHA: b3d26c1ce7aa447534ca760df009cfd7985babc4

Checks will be approved automatically on success.

@mihaimitrea-db
mihaimitrea-db added this pull request to the merge queue Mar 19, 2026
Merged via the queue into main with commit b805cf1 Mar 19, 2026
15 checks passed
@mihaimitrea-db
mihaimitrea-db deleted the mihaimitrea-db/token-force-refresh branch March 19, 2026 15:41
github-merge-queue Bot pushed a commit to databricks/cli that referenced this pull request Mar 27, 2026
Wire a new `--force-refresh` CLI flag that delegates to the SDK's
`ForceRefreshToken()` method, bypassing the cached token validity check.
The default path through `Token()` is unchanged.

> **Note:** this will not compile until the SDK ships
`ForceRefreshToken()`
([databricks/databricks-sdk-go#1552](databricks/databricks-sdk-go#1552))
and the CLI''s Go SDK dependency is bumped to v0.126.

## Changes

- Add a `--force-refresh` boolean flag to `databricks auth token`. When
set, the command calls `PersistentAuth.ForceRefreshToken()` instead of
`PersistentAuth.Token()`, which always performs a token refresh against
the IdP regardless of the cached token's remaining TTL.
- The error handling is shared between both paths: invalid refresh
tokens produce the same actionable "To reauthenticate, run..." message,
and the backward-compat `cache.ErrNotFound` rewrite is preserved.
- Add three acceptance tests covering the force-refresh flag end-to-end:
- **`force-refresh-success`**: pre-populated cache with a valid token,
`--force-refresh` returns the server's new token (not the cached one).
- **`force-refresh-invalid-refresh-token`**: server returns 401 with
invalid refresh token; asserts the actionable re-login error message.
- **`force-refresh-no-cache`**: no cached token exists; asserts the
backward-compat error message is preserved for `--force-refresh`.
- Add unit tests verifying the default path still returns a cached valid
token without refreshing, and that `--force-refresh` correctly delegates
to `ForceRefreshToken()` for both success and failure cases.

## Why

[#4564](#4564) reports that
external consumers using `databricks auth token` as a credential helper
(e.g. Claude Code via `apiKeyHelper`) can receive near-expired tokens
that expire before they can be used.


[databricks/databricks-sdk-go#1535](databricks/databricks-sdk-go#1535)
addressed the common case by adding a 5-minute proactive refresh buffer
to `Token()`. However, that proactive refresh is intentionally
best-effort: `Token()` still returns the existing access token when it
is valid and a proactive refresh fails, because callers did not
explicitly ask for a fresh token.

For integrations that treat the CLI as a token minter or want to manage
their own cache/TTL policy, "return a still-usable token" is different
from "refresh now and give me a newly minted token or fail." The
`--force-refresh` flag gives those integrations an explicit way to
guarantee a fresh token.

## Tests

- Unit tests in `cmd/auth/token_test.go`: table-driven cases covering
default-path cache reuse, force-refresh success, and force-refresh error
preservation.
- Acceptance tests in
`acceptance/cmd/auth/token/force-refresh-{success,invalid-refresh-token,no-cache}/`:
end-to-end tests against the mock OIDC server, run for both `terraform`
and `direct` engine variants.

---------

Co-authored-by: simon <[email protected]>
denik pushed a commit to databricks/cli that referenced this pull request May 20, 2026
Wire a new `--force-refresh` CLI flag that delegates to the SDK's
`ForceRefreshToken()` method, bypassing the cached token validity check.
The default path through `Token()` is unchanged.

> **Note:** this will not compile until the SDK ships
`ForceRefreshToken()`
([databricks/databricks-sdk-go#1552](databricks/databricks-sdk-go#1552))
and the CLI''s Go SDK dependency is bumped to v0.126.

## Changes

- Add a `--force-refresh` boolean flag to `databricks auth token`. When
set, the command calls `PersistentAuth.ForceRefreshToken()` instead of
`PersistentAuth.Token()`, which always performs a token refresh against
the IdP regardless of the cached token's remaining TTL.
- The error handling is shared between both paths: invalid refresh
tokens produce the same actionable "To reauthenticate, run..." message,
and the backward-compat `cache.ErrNotFound` rewrite is preserved.
- Add three acceptance tests covering the force-refresh flag end-to-end:
- **`force-refresh-success`**: pre-populated cache with a valid token,
`--force-refresh` returns the server's new token (not the cached one).
- **`force-refresh-invalid-refresh-token`**: server returns 401 with
invalid refresh token; asserts the actionable re-login error message.
- **`force-refresh-no-cache`**: no cached token exists; asserts the
backward-compat error message is preserved for `--force-refresh`.
- Add unit tests verifying the default path still returns a cached valid
token without refreshing, and that `--force-refresh` correctly delegates
to `ForceRefreshToken()` for both success and failure cases.

## Why

[#4564](#4564) reports that
external consumers using `databricks auth token` as a credential helper
(e.g. Claude Code via `apiKeyHelper`) can receive near-expired tokens
that expire before they can be used.


[databricks/databricks-sdk-go#1535](databricks/databricks-sdk-go#1535)
addressed the common case by adding a 5-minute proactive refresh buffer
to `Token()`. However, that proactive refresh is intentionally
best-effort: `Token()` still returns the existing access token when it
is valid and a proactive refresh fails, because callers did not
explicitly ask for a fresh token.

For integrations that treat the CLI as a token minter or want to manage
their own cache/TTL policy, "return a still-usable token" is different
from "refresh now and give me a newly minted token or fail." The
`--force-refresh` flag gives those integrations an explicit way to
guarantee a fresh token.

## Tests

- Unit tests in `cmd/auth/token_test.go`: table-driven cases covering
default-path cache reuse, force-refresh success, and force-refresh error
preservation.
- Acceptance tests in
`acceptance/cmd/auth/token/force-refresh-{success,invalid-refresh-token,no-cache}/`:
end-to-end tests against the mock OIDC server, run for both `terraform`
and `direct` engine variants.

---------

Co-authored-by: simon <[email protected]>

This branch was previously deployed

1 inactive deployment
test-trigger-is — b3d26c1c Deployed Mar 19, 2026 by mihaimitrea-db via Trigger Tests #1520
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