stacklok / toolhive / 24988918773

67%

Build Type

push

github

Committed by web-flow

Commit Message

pkg/llm: fix withPreemptiveRefresh composition and stale _AT on rotation (#5052)

* pkg/llm: fix withPreemptiveRefresh composition and stale _AT on rotation

Two bugs in tokensource.go flagged in PR review (#5051):

1. withPreemptiveRefresh composition: both callers were passing a source
   already wrapped in oauth2.ReuseTokenSource (from CreateTokenSourceFromCached
   and flow.TokenSource()), creating a double-caching chain. In the preemptive
   window the inner ReuseTokenSource returned the same cached real token, causing
   preemptiveTokenSource to shift an already-expired expiry, so every Token() call
   re-entered the inner chain instead of one refresh per window.

   Fix: build the raw refresher directly (NewResourceTokenSource for the resource
   path, config.TokenSource for the non-resource path) without an inner
   ReuseTokenSource, then stack: ReuseTokenSource(nil, preemptive{Persisting{raw}}).
   Extract oauth2ConfigFrom helper to reuse the already-discovered flowCfg in
   performBrowserFlow without a second OIDC round-trip.

2. Stale _AT cache after refresh token rotation: makeTokenPersister wrote the new
   refresh token but left the _AT access-token cache entry unchanged. A concurrent
   process could read the stale _AT between the write and the next Token() call.

   Fix: after SetSecret for the refresh token succeeds, overwrite the _AT entry
   with an empty string so the next Token() call falls through to tier 2.

Adds a regression test verifying that multiple Token() calls within the preemptive
window produce exactly two inner calls: one for the initial short-lived token and
one for the preemptive refresh that yields a fresh cached token.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* changes from review

* fixes from review

---------

Co-authored-by: taskbot <taskbot@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Coverage Stats

86 of 106 new or added lines in 3 files covered. (81.13%)

2 existing lines in 1 file now uncovered.

59382 of 89598 relevant lines covered (66.28%)

60.42 hits per line