From 1824e51dcd00c95baba1ad33069769392d32a594 Mon Sep 17 00:00:00 2001 From: Duncan Date: Tue, 11 Aug 2026 00:57:47 -0400 Subject: [PATCH] feat(buzz-agent): add DatabricksAuthCoordinator single-flight OAuth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Databricks OAuth acquisition was scattered across the TokenSource methods with no coordination: concurrent callers (Desktop discovery, the model picker, managed-runtime inference) could each pop their own browser, and a just-denied attempt would immediately re-prompt. Consolidate every acquisition path behind one coordinator on PkceOAuthTokenSource. - Intent policy (Auto/UserInitiated/Headless) decides browser and cooldown behavior; all four TokenSource methods plus interactive_login route through one acquire()/acquire_locked() core. - Single-flight per cache key via a std File advisory lock. try_lock is per open-file-description, so distinct handles contend in-process AND across processes — one primitive covers both with no in-memory registry. RAII drop releases, so a crashed holder never wedges a successor. - Typed AuthError outcomes with stable code()/from_code() for the Phase 2 Tauri boundary, replacing display-text matching. - Durable cooldown sidecar: every failed browser attempt is recorded; Auto honors an unexpired record instead of re-launching, UserInitiated bypasses and clears it. Success clears it. - Browser opener injected and invoked while the localhost callback listener is live, so a launch failure never returns a URL pointing at a torn-down listener. The crate now uses std File::try_lock/unlock (stable in 1.89), so its rust-version is pinned above the 1.88 workspace floor. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-agent/Cargo.toml | 4 +- crates/buzz-agent/src/auth.rs | 938 +++++++++++++----- .../tests/databricks_auth_coordinator.rs | 589 +++++++++++ 3 files changed, 1307 insertions(+), 224 deletions(-) create mode 100644 crates/buzz-agent/tests/databricks_auth_coordinator.rs diff --git a/crates/buzz-agent/Cargo.toml b/crates/buzz-agent/Cargo.toml index fabf75754..1ba9c55fb 100644 --- a/crates/buzz-agent/Cargo.toml +++ b/crates/buzz-agent/Cargo.toml @@ -2,7 +2,9 @@ name = "buzz-agent" version.workspace = true edition.workspace = true -rust-version.workspace = true +# Above the 1.88 workspace floor: the auth coordinator's cross-process +# single-flight uses `std::fs::File::try_lock`/`unlock`, stable since 1.89. +rust-version = "1.89.0" license.workspace = true repository.workspace = true description = "Minimal, unbreakable ACP-compliant agent. Non-streaming. Tool-calls-as-output." diff --git a/crates/buzz-agent/src/auth.rs b/crates/buzz-agent/src/auth.rs index a78a499bd..5588415b2 100644 --- a/crates/buzz-agent/src/auth.rs +++ b/crates/buzz-agent/src/auth.rs @@ -39,6 +39,202 @@ const TOKEN_REFRESH_LEEWAY: Duration = Duration::from_secs(60); /// We match: any longer and the user has gone to lunch. const BROWSER_AUTH_TIMEOUT: Duration = Duration::from_secs(60); +/// Per-request network timeout for every OAuth HTTP call (discovery, refresh +/// grant, code exchange). Without this, a hung provider connection would stall +/// the caller — and, worse, stall every same-key caller waiting on the +/// cross-process lock this holder owns. +const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); + +/// Longest an in-flight auth attempt can legitimately run: cold discovery +/// (`30s`) + browser wait (`60s`) + code exchange (`30s`), plus a failed +/// refresh (`30s`) ahead of the browser. Rounded to `150s`. A waiter derives +/// its lock-wait bound from this so it never times out ahead of a healthy +/// holder. +const AUTH_ATTEMPT_DEADLINE: Duration = Duration::from_secs(150); + +/// How long a same-key caller waits to acquire the cross-process lock before +/// giving up with [`AuthError::LockTimeout`]. Deliberately longer than +/// [`AUTH_ATTEMPT_DEADLINE`] so a waiter outlasts any legitimate holder rather +/// than timing out mid-flow. +const LOCK_WAIT_TIMEOUT: Duration = Duration::from_secs(165); + +/// Poll interval for deadline-aware lock acquisition. `try_lock` is +/// non-blocking, so we sleep between attempts rather than blocking a worker. +const LOCK_POLL_INTERVAL: Duration = Duration::from_millis(100); + +/// How long a failed interactive (browser) attempt suppresses automatic +/// re-launch for the same key. Long enough that a spurned dropdown does not +/// re-pop a browser on the next debounced refresh, short enough that a user +/// who fixes the problem is not locked out. +const COOLDOWN_DURATION: Duration = Duration::from_secs(300); + +/// Why an auth acquisition wants a token, which decides whether it may open a +/// browser and whether it honors a cooldown. +/// +/// - [`Auto`](Self::Auto): passive Desktop discovery (create/edit/defaults/ +/// onboarding). May open a browser, but honors an unexpired cooldown and +/// returns its recorded outcome instead of re-launching. +/// - [`UserInitiated`](Self::UserInitiated): an explicit human action — the +/// saved-agent model picker or `buzz-agent auth databricks`. May open a +/// browser and *bypasses* the cooldown (the user asked for it now). +/// - [`Headless`](Self::Headless): managed-runtime inference and provider +/// preflight. Never opens a browser; may consume another attempt's cached +/// success but never becomes the initiator. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthIntent { + Auto, + UserInitiated, + Headless, +} + +impl AuthIntent { + /// `true` for the intents permitted to open a browser. + fn may_open_browser(self) -> bool { + matches!(self, Self::Auto | Self::UserInitiated) + } + + /// `true` for the one intent that honors a recorded cooldown on read. + fn honors_cooldown(self) -> bool { + matches!(self, Self::Auto) + } +} + +/// Typed result of an auth acquisition. `Ok` carries the bearer; the error +/// arm classifies *why* no token was produced so callers (and, in Phase 2, the +/// Tauri boundary) can branch on a stable code instead of matching display +/// text. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AuthError { + /// No cached token, no refresh grant, and the caller may not open a + /// browser (`Headless`). + NoCredential, + /// The user (or provider) rejected the browser authorization. + Denied, + /// The browser flow was not completed within [`BROWSER_AUTH_TIMEOUT`]. + TimedOut, + /// Every browser-launch strategy failed, so the flow never started. + BrowserOpenFailed, + /// An OAuth network call (discovery/refresh/exchange) could not reach the + /// provider or timed out. + NetworkUnavailable, + /// A refresh-token grant was rejected (dead/rotated refresh token) and the + /// caller may not fall back to a browser. + RefreshRejected, + /// The authorization-code exchange itself was rejected by the token + /// endpoint (distinct from a refresh rejection). + ExchangeFailed, + /// Could not acquire the cross-process auth lock within + /// [`LOCK_WAIT_TIMEOUT`]. + LockTimeout, +} + +impl AuthError { + /// Stable machine-readable code. Phase 2 serializes this across the Tauri + /// boundary (the `project_git_merge_error` `{code, message}` precedent) so + /// the Desktop formatter switches on the code, never on display text. + pub fn code(&self) -> &'static str { + match self { + Self::NoCredential => "no_credential", + Self::Denied => "denied", + Self::TimedOut => "timed_out", + Self::BrowserOpenFailed => "browser_open_failed", + Self::NetworkUnavailable => "network_unavailable", + Self::RefreshRejected => "refresh_rejected", + Self::ExchangeFailed => "exchange_failed", + Self::LockTimeout => "lock_timeout", + } + } + + /// `true` for the browser-attempt outcomes worth recording in the cooldown + /// sidecar — the failures that would otherwise re-pop a browser on the + /// next automatic attempt. Non-browser failures (no credential, refresh + /// rejection, lock timeout, network) are not recorded. + fn is_cooldown_worthy(&self) -> bool { + matches!( + self, + Self::Denied | Self::TimedOut | Self::BrowserOpenFailed | Self::ExchangeFailed + ) + } + + /// Reconstruct a recorded outcome from its [`code`](Self::code). Only the + /// cooldown-worthy variants round-trip; any other code (a forward-compat + /// sidecar written by a newer buzz-agent) yields `None`, so a stale or + /// unrecognized record is treated as "no cooldown" rather than a hard + /// failure. + fn from_code(code: &str) -> Option { + match code { + "denied" => Some(Self::Denied), + "timed_out" => Some(Self::TimedOut), + "browser_open_failed" => Some(Self::BrowserOpenFailed), + "exchange_failed" => Some(Self::ExchangeFailed), + _ => None, + } + } + + fn message(&self) -> String { + match self { + Self::NoCredential => { + "no cached Databricks token; run `buzz-agent auth databricks` first".into() + } + Self::Denied => "Databricks authorization was denied".into(), + Self::TimedOut => "Databricks browser authorization timed out".into(), + Self::BrowserOpenFailed => "could not open a browser for Databricks sign-in".into(), + Self::NetworkUnavailable => "could not reach Databricks to authenticate".into(), + Self::RefreshRejected => "Databricks rejected the refresh token; sign in again".into(), + Self::ExchangeFailed => "Databricks rejected the authorization code".into(), + Self::LockTimeout => "timed out waiting for a concurrent Databricks sign-in".into(), + } + } +} + +impl From for AgentError { + /// Map a typed auth failure onto the crate error the [`TokenSource`] trait + /// returns. Auth-decision failures become [`AgentError::LlmAuth`] so the + /// caller's retry loop stops instead of hammering a rejected credential; + /// purely infrastructural failures (network, lock contention) become + /// [`AgentError::Llm`], matching the pre-coordinator classification of a + /// discovery/network error. + fn from(e: AuthError) -> Self { + match e { + AuthError::NetworkUnavailable | AuthError::LockTimeout => AgentError::Llm(e.message()), + AuthError::NoCredential + | AuthError::Denied + | AuthError::TimedOut + | AuthError::BrowserOpenFailed + | AuthError::RefreshRejected + | AuthError::ExchangeFailed => AgentError::LlmAuth(e.message()), + } + } +} + +/// Opens a URL for the interactive browser step. Injected so the PKCE +/// continuation (callback listener, verifier, timeout) stays alive across the +/// launch: the coordinator calls this *while* the localhost listener is +/// bound, so a launch failure never leaves a returned URL pointing at a torn +/// down listener. Desktop (Phase 2) supplies the Tauri opener; the CLI uses +/// [`DefaultBrowserOpener`], which prints the URL and opens the system +/// browser. +pub trait BrowserOpener: Send + Sync { + /// Attempt to present `url` to the user. Returning `Err` means every + /// launch strategy for this opener failed; the coordinator then reports + /// [`AuthError::BrowserOpenFailed`] without waiting on a listener nobody + /// will reach. + fn open(&self, url: &str) -> Result<(), String>; +} + +/// Default opener: print the URL (so a user on a headless box can copy it) +/// and open the system browser. Printing is itself a launch strategy, so this +/// never reports failure — the URL is always visible to the waiting user. +pub struct DefaultBrowserOpener; + +impl BrowserOpener for DefaultBrowserOpener { + fn open(&self, url: &str) -> Result<(), String> { + eprintln!("Opening browser for authentication. If it doesn't open, visit:\n {url}"); + let _ = webbrowser::open(url); + Ok(()) + } +} + /// Asynchronous source of a bearer token. The [`Llm`] calls this per /// request, so impls are expected to be cheap on the cache-hit path. #[async_trait] @@ -136,27 +332,66 @@ pub struct PkceOAuthTokenSource { cfg: PkceOAuthConfig, http: Client, cache_path: PathBuf, - /// Single-flight guard: only one refresh/browser flow at a time, even - /// if many tool calls land concurrently. + /// Injected browser launcher, called inside [`browser_pkce_flow`] while the + /// localhost listener is live. Production uses [`DefaultBrowserOpener`]; + /// Phase 2 supplies the Tauri opener. + opener: Arc, + /// In-memory single-flight *and* fast-path cache. The cross-process file + /// lock serializes slow-path work; this cell keeps the fast path off disk + /// during a turn and off the lock entirely. state: Mutex>, } impl PkceOAuthTokenSource { + /// Construct with the default browser opener (prints the URL and opens the + /// system browser). This is the signature every production call site uses. pub fn new(cfg: PkceOAuthConfig) -> Result, AgentError> { + Self::new_with(cfg, Arc::new(DefaultBrowserOpener)) + } + + /// Construct with an injected [`BrowserOpener`]. Tests substitute a + /// recording/failing opener to exercise the browser branch without a real + /// window; Phase 2 Desktop injects the Tauri opener. + pub fn new_with( + cfg: PkceOAuthConfig, + opener: Arc, + ) -> Result, AgentError> { let cache_path = cache_path_for(&cfg)?; if let Some(parent) = cache_path.parent() { fs::create_dir_all(parent) .map_err(|e| AgentError::Llm(format!("oauth cache dir {parent:?}: {e}")))?; } + // Every OAuth HTTP call inherits this timeout so a hung provider can + // never stall the caller — nor the same-key callers waiting on the + // cross-process lock this holder owns. A build failure falls back to + // the untimed default rather than making construction fallible. + let http = Client::builder() + .timeout(HTTP_REQUEST_TIMEOUT) + .build() + .unwrap_or_else(|_| Client::new()); let initial = read_cache(&cache_path); Ok(Arc::new(Self { cfg, - http: Client::new(), + http, cache_path, + opener, state: Mutex::new(initial), })) } + /// Path of the cross-process advisory lock file guarding slow-path auth + /// for this cache key. Co-located with the cache so it shares the + /// per-key directory and `$HOME` override. + fn lock_path(&self) -> PathBuf { + append_ext(&self.cache_path, "lock") + } + + /// Path of the cooldown sidecar recording the last browser-attempt + /// failure for this cache key. + fn cooldown_path(&self) -> PathBuf { + append_ext(&self.cache_path, "cooldown") + } + /// Discover authorization + token endpoints from the well-known URL. async fn endpoints(&self) -> Result { let v: Value = self @@ -232,192 +467,244 @@ impl PkceOAuthTokenSource { token_from_response(&v, Some(refresh_token)) } - /// Run the full browser-mediated Authorization Code + PKCE flow. - /// Caller must hold a TTY/browser: this opens a window and blocks. + /// Run the full browser-mediated Authorization Code + PKCE flow and cache + /// the result. Routes through the coordinator as a [`UserInitiated`] + /// acquisition: it may open a browser, bypasses (and clears) any cooldown, + /// and single-flights with concurrent callers on the cross-process lock. A + /// still-valid cached token short-circuits to success without re-prompting. + /// + /// [`UserInitiated`]: AuthIntent::UserInitiated pub async fn interactive_login(&self) -> Result<(), AgentError> { - let endpoints = self.endpoints().await?; - let token = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?; - let mut state = self.state.lock().await; - self.save(&mut state, token)?; + self.acquire(AuthIntent::UserInitiated, None).await?; Ok(()) } + + /// Public entry for passive Desktop discovery (Phase 2): acquire a bearer + /// under an explicit [`AuthIntent`], returning the typed [`AuthError`] so + /// the caller can branch on a stable `code` rather than display text. The + /// [`TokenSource`] trait methods wrap this and flatten the error into + /// [`AgentError`]. + pub async fn acquire_with_intent(&self, intent: AuthIntent) -> Result { + self.acquire(intent, None).await + } + + /// Return a usable cached bearer, applying the identity rule for a + /// 401-driven acquisition. + /// + /// `rejected = None` (normal): a not-yet-expired cached token is a hit. + /// `rejected = Some(t)`: the expiry clock is untrustworthy — the rejected + /// token looked locally fresh — so a hit requires the cached token to + /// *differ* from `t`, meaning a sibling already replaced it. Checks the + /// in-memory cell first, then re-reads disk (a sibling process may have + /// written a newer token) and adopts it into the cell on a hit. + fn cached_hit( + &self, + state: &mut Option, + rejected: Option<&str>, + ) -> Option { + let usable = |tok: &CachedToken| match rejected { + Some(r) => tok.access_token != r, + None => !is_expired(tok), + }; + if let Some(tok) = state.as_ref() { + if usable(tok) { + return Some(tok.access_token.clone()); + } + } + if let Some(disk) = read_cache(&self.cache_path) { + if usable(&disk) { + let bearer = disk.access_token.clone(); + *state = Some(disk); + return Some(bearer); + } + } + None + } + + /// Discover OIDC endpoints once per flow, memoizing into `slot` so the + /// refresh and browser branches share a single discovery call. A discovery + /// failure (unreachable URL or malformed document) maps to + /// [`AuthError::NetworkUnavailable`] — the infrastructural bucket, so the + /// caller's retry loop treats it as transient rather than as an auth + /// decision. + async fn discover<'a>( + &self, + slot: &'a mut Option, + ) -> Result<&'a OidcEndpoints, AuthError> { + if slot.is_none() { + let eps = self + .endpoints() + .await + .map_err(|_| AuthError::NetworkUnavailable)?; + *slot = Some(eps); + } + Ok(slot.as_ref().expect("endpoints just populated")) + } + + /// The single acquisition entry point behind every [`TokenSource`] method. + /// + /// `intent` decides browser and cooldown policy; `rejected` (`Some` only on + /// a 401-driven refresh) switches cache checks from clock-based to + /// identity-based. The fast path returns a usable cached token without + /// touching the lock or the network. Otherwise the slow path serializes + /// every same-key caller — in this process *and* across processes — on the + /// cross-process advisory lock, so concurrent dialogs coalesce onto one + /// refresh/browser flow instead of racing browsers. + async fn acquire( + &self, + intent: AuthIntent, + rejected: Option<&str>, + ) -> Result { + // Fast path: no lock, no network. + { + let mut state = self.state.lock().await; + if let Some(hit) = self.cached_hit(&mut state, rejected) { + return Ok(hit); + } + } + + // Slow path: one flow at a time per cache key. The waiter's deadline + // exceeds a healthy holder's attempt deadline, so it never gives up on + // a live holder. + let deadline = std::time::Instant::now() + LOCK_WAIT_TIMEOUT; + let _guard = acquire_auth_lock(&self.lock_path(), deadline).await?; + + // Bound the whole locked attempt so a wedged flow can't hold the lock + // past the waiters' patience. On expiry the lock releases (guard drop) + // and the attempt reports TimedOut. + match tokio::time::timeout(AUTH_ATTEMPT_DEADLINE, self.acquire_locked(intent, rejected)) + .await + { + Ok(result) => result, + Err(_) => Err(AuthError::TimedOut), + } + } + + /// Slow-path body, run while holding the cross-process auth lock. + async fn acquire_locked( + &self, + intent: AuthIntent, + rejected: Option<&str>, + ) -> Result { + let mut state = self.state.lock().await; + + // Re-check under the lock: a holder we queued behind may have already + // produced a token (this process or a sibling wrote the cache). + if let Some(hit) = self.cached_hit(&mut state, rejected) { + return Ok(hit); + } + + // Refresh-token grant, if we have one. Endpoints are discovered lazily + // here (and reused by the browser branch) so a no-refresh headless + // failure never depends on reaching the discovery URL. + let mut endpoints: Option = None; + let mut refresh_failed = false; + if let Some(rt) = state.as_ref().and_then(|t| t.refresh_token.clone()) { + let eps = self.discover(&mut endpoints).await?; + match self.refresh(eps, &rt).await { + Ok(fresh) => return self.finish(&mut state, fresh), + Err(e) => { + tracing::warn!(error = %e, "oauth refresh failed; falling back"); + // A sibling may still have won the race while we ran. + if let Some(hit) = self.cached_hit(&mut state, rejected) { + return Ok(hit); + } + refresh_failed = true; + } + } + } + + // No token from cache or refresh. Browser or terminal failure. + if !intent.may_open_browser() { + return Err(if refresh_failed { + AuthError::RefreshRejected + } else { + AuthError::NoCredential + }); + } + + let cooldown_path = self.cooldown_path(); + if intent.honors_cooldown() { + // A recent browser attempt failed; surface its recorded outcome + // instead of re-popping a browser on this automatic attempt. + if let Some(recorded) = read_cooldown(&cooldown_path) { + return Err(recorded); + } + } else { + // An explicit user retry clears any prior suppression. + clear_cooldown(&cooldown_path); + } + + let eps = self.discover(&mut endpoints).await?; + match browser_pkce_flow(&self.http, &self.cfg, eps, self.opener.as_ref()).await { + // `finish` clears the cooldown on success. + Ok(fresh) => self.finish(&mut state, fresh), + Err(e) => { + if e.is_cooldown_worthy() { + write_cooldown(&cooldown_path, &e); + } + Err(e) + } + } + } + + /// Persist a freshly-obtained token, clear any cooldown, and return its + /// bearer. A cache-write failure maps to [`AuthError::NetworkUnavailable`] + /// (the infrastructural bucket) — the token was valid but couldn't be + /// persisted, which the caller should treat as transient, not as a + /// credential rejection. + fn finish( + &self, + state: &mut Option, + token: CachedToken, + ) -> Result { + let bearer = token.access_token.clone(); + self.save(state, token) + .map_err(|_| AuthError::NetworkUnavailable)?; + clear_cooldown(&self.cooldown_path()); + Ok(bearer) + } } #[async_trait] impl TokenSource for PkceOAuthTokenSource { + /// Acquire a bearer for a request. Routes through the coordinator as a + /// [`Headless`](AuthIntent::Headless) acquisition: it serves a cached or + /// refreshed token but never opens a browser, so a managed runtime with no + /// interactive display can never hang on inference. First-use auth is the + /// job of `buzz-agent auth databricks` ([`interactive_login`]). + /// + /// [`interactive_login`]: PkceOAuthTokenSource::interactive_login async fn bearer(&self) -> Result { - let mut state = self.state.lock().await; - - // 1. In-memory cache hit, still fresh. - if let Some(tok) = state.as_ref() { - if !is_expired(tok) { - return Ok(tok.access_token.clone()); - } - } - - // 2. Re-read disk — another process may have refreshed already. - if let Some(disk_tok) = read_cache(&self.cache_path) { - if !is_expired(&disk_tok) { - let bearer = disk_tok.access_token.clone(); - *state = Some(disk_tok); - return Ok(bearer); - } - } - - // 3. Try refresh if we have a refresh token. Discover endpoints once - // here — deliberately hoisted above the refresh-token check so the - // browser flow at step 5 (which also needs them) reuses this call. - let endpoints = self.endpoints().await?; - let refresh = state.as_ref().and_then(|t| t.refresh_token.clone()); - if let Some(rt) = refresh { - match self.refresh(&endpoints, &rt).await { - Ok(fresh) => { - let bearer = fresh.access_token.clone(); - self.save(&mut state, fresh)?; - return Ok(bearer); - } - Err(e) => { - tracing::warn!(error = %e, "oauth refresh failed; falling back to browser flow"); - } - } - - // 4. Re-read disk after refresh failure — another process may have won the race. - if let Some(disk_tok) = read_cache(&self.cache_path) { - if !is_expired(&disk_tok) { - let bearer = disk_tok.access_token.clone(); - *state = Some(disk_tok); - return Ok(bearer); - } - } - } - - // 5. No usable cache: full browser dance. - let fresh = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?; - let bearer = fresh.access_token.clone(); - self.save(&mut state, fresh)?; - Ok(bearer) + self.acquire(AuthIntent::Headless, None) + .await + .map_err(Into::into) } + /// Identical to [`bearer`](Self::bearer) for this source — both are + /// headless. Retained as a distinct method so callers can state the + /// no-browser requirement at the call site (and so other [`TokenSource`] + /// impls that *would* browse in `bearer` can still expose a safe path). async fn bearer_no_browser(&self) -> Result { - self.try_bearer_no_browser().await + self.acquire(AuthIntent::Headless, None) + .await + .map_err(Into::into) } - /// Force-refresh after a 401, never touching the browser flow. + /// Force a fresh bearer after the server rejected `rejected` with a 401. /// - /// `rejected` is the access token the server just 401'd. Coalescing keys - /// off token *identity*, not the expiry clock: a 401 means the token was - /// rejected while it still looked locally fresh, so `is_expired()` would - /// say "keep it" and no grant would ever run. Instead, under the lock we - /// compare the current cached token to `rejected` — if they differ, a - /// concurrent caller (this process or a sibling) already refreshed, so we - /// return the new token without burning a second grant. If they still - /// match, this is the rejected token and we run the refresh-token grant - /// unconditionally. The whole check→refresh→save runs under one lock hold - /// so concurrent callers serialize. On any failure the refresh token is - /// preserved (never nulled) and the error is terminal `LlmAuth` — no - /// browser, no hang. + /// A [`Headless`](AuthIntent::Headless) acquisition keyed by token + /// *identity* rather than the expiry clock: a 401 means the cached token + /// was rejected while still locally fresh, so [`is_expired`] would wrongly + /// keep it. Passing `rejected` makes the coordinator run the refresh-token + /// grant unless a concurrent caller already replaced the token, in which + /// case that newer token is returned without a second grant. Never opens a + /// browser; a dead refresh token surfaces terminally so the retry loop + /// stops instead of hanging. async fn refresh_now(&self, rejected: &str) -> Result { - let mut state = self.state.lock().await; - - // 1. Coalesce by identity: if the cached token (in-memory, then disk) - // is no longer the one the server rejected, someone already - // refreshed it. Return that instead of grabbing another grant. - if let Some(tok) = state.as_ref() { - if tok.access_token != rejected { - return Ok(tok.access_token.clone()); - } - } - if let Some(disk_tok) = read_cache(&self.cache_path) { - if disk_tok.access_token != rejected { - let bearer = disk_tok.access_token.clone(); - *state = Some(disk_tok); - return Ok(bearer); - } - } - - // 2. The cached token is still the rejected one. Run the refresh-token - // grant unconditionally — the expiry clock can't be trusted here, a - // locally-fresh token is exactly what got 401'd. - let refresh = state.as_ref().and_then(|t| t.refresh_token.clone()); - let Some(rt) = refresh else { - return Err(AgentError::LlmAuth( - "token rejected and no refresh token available".into(), - )); - }; - let endpoints = self.endpoints().await?; - match self.refresh(&endpoints, &rt).await { - Ok(fresh) => { - let bearer = fresh.access_token.clone(); - self.save(&mut state, fresh)?; - Ok(bearer) - } - // 3. Refresh token is itself dead. Terminal — surfacing LlmAuth - // stops the retry loop instead of falling to the browser flow, - // which would hang a headless harness. - Err(e) => Err(AgentError::LlmAuth(format!("token refresh failed: {e}"))), - } - } -} - -impl PkceOAuthTokenSource { - /// Return a bearer token from cache or refresh, **never** opening a browser. - /// - /// Follows the same steps as [`bearer`](TokenSource::bearer) but stops at - /// step 4 — if no usable token is available after cache + refresh attempts, - /// returns `Err(LlmAuth(...))` instead of launching the browser PKCE flow. - /// Used by model-discovery paths that must not block on user interaction. - pub(crate) async fn try_bearer_no_browser(&self) -> Result { - let mut state = self.state.lock().await; - - // 1. In-memory cache hit, still fresh. - if let Some(tok) = state.as_ref() { - if !is_expired(tok) { - return Ok(tok.access_token.clone()); - } - } - - // 2. Re-read disk — another process may have refreshed already. - if let Some(disk_tok) = read_cache(&self.cache_path) { - if !is_expired(&disk_tok) { - let bearer = disk_tok.access_token.clone(); - *state = Some(disk_tok); - return Ok(bearer); - } - } - - // 3. Try refresh if we have a refresh token. Endpoints are discovered - // lazily here — only when a refresh token is actually present — so - // that an unreachable OIDC discovery URL cannot prevent the - // no-token/no-cache path from returning `LlmAuth` (graceful - // fallback) instead of `Llm` (hard error). - let refresh = state.as_ref().and_then(|t| t.refresh_token.clone()); - if let Some(rt) = refresh { - let endpoints = self.endpoints().await?; - match self.refresh(&endpoints, &rt).await { - Ok(fresh) => { - let bearer = fresh.access_token.clone(); - self.save(&mut state, fresh)?; - return Ok(bearer); - } - Err(e) => { - tracing::warn!(error = %e, "oauth refresh failed during model discovery"); - } - } - - // 4. Re-read disk after refresh failure. - if let Some(disk_tok) = read_cache(&self.cache_path) { - if !is_expired(&disk_tok) { - let bearer = disk_tok.access_token.clone(); - *state = Some(disk_tok); - return Ok(bearer); - } - } - } - - // No usable token — return error instead of opening a browser. - Err(AgentError::LlmAuth( - "no cached Databricks token; run `buzz-agent auth databricks` first".into(), - )) + self.acquire(AuthIntent::Headless, Some(rejected)) + .await + .map_err(Into::into) } } @@ -437,11 +724,7 @@ fn is_expired(t: &CachedToken) -> bool { let Some(exp) = t.expires_at else { return false; }; - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); - now + TOKEN_REFRESH_LEEWAY.as_secs() >= exp + now_secs() + TOKEN_REFRESH_LEEWAY.as_secs() >= exp } fn cache_path_for(cfg: &PkceOAuthConfig) -> Result { @@ -465,6 +748,123 @@ fn cache_path_for(cfg: &PkceOAuthConfig) -> Result { Ok(dir.join(format!("{hash}.json"))) } +/// Append `ext` as an extra extension onto `base` (e.g. `.json` → +/// `.json.lock`). Keeps the lock and cooldown sidecars in the same +/// per-key directory as the cache, so they inherit its `$HOME` override and +/// owner-only parent without a second key derivation. +fn append_ext(base: &Path, ext: &str) -> PathBuf { + let mut name = base.as_os_str().to_owned(); + name.push("."); + name.push(ext); + PathBuf::from(name) +} + +/// Durable record of the last browser-attempt failure for a cache key. Written +/// while holding the auth lock so concurrent writers can't interleave, read by +/// `Auto` callers to decide whether to suppress an automatic browser re-launch. +#[derive(Debug, Serialize, Deserialize)] +struct CooldownRecord { + /// [`AuthError::code`] of the failure being cooled down. + code: String, + /// Unix seconds after which the cooldown lapses and an `Auto` caller may + /// launch a browser again. + until: u64, +} + +fn now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// Return the still-active cooldown outcome for `path`, if any. +/// +/// `None` when the sidecar is absent, unparseable, expired, or records a code +/// this build doesn't recognize — every one of those means "no active +/// cooldown", so the caller proceeds to a normal attempt. An expired record is +/// removed opportunistically so the directory doesn't accumulate stale files. +fn read_cooldown(path: &Path) -> Option { + let body = fs::read(path).ok()?; + let record: CooldownRecord = serde_json::from_slice(&body).ok()?; + if record.until > now_secs() { + AuthError::from_code(&record.code) + } else { + let _ = fs::remove_file(path); + None + } +} + +/// Record `err` as a fresh cooldown at `path`, expiring [`COOLDOWN_DURATION`] +/// from now. Best-effort: a write failure only means the next automatic +/// attempt may re-pop a browser, never a hard auth failure, so errors are +/// swallowed. Called while holding the auth lock. +fn write_cooldown(path: &Path, err: &AuthError) { + let record = CooldownRecord { + code: err.code().to_string(), + until: now_secs() + COOLDOWN_DURATION.as_secs(), + }; + if let Ok(body) = serde_json::to_vec(&record) { + let _ = write_private_cache(path, &body); + } +} + +/// Remove any cooldown sidecar at `path`. Called on a successful acquisition +/// (the problem is resolved) and by `UserInitiated` callers that bypass the +/// cooldown (an explicit retry clears the suppression). Best-effort. +fn clear_cooldown(path: &Path) { + let _ = fs::remove_file(path); +} + +/// Hold on the cross-process auth lock. Dropping it (or the owning process +/// dying) releases the OS advisory lock — no PID files, no manual break. +#[derive(Debug)] +struct AuthLockGuard(fs::File); + +impl Drop for AuthLockGuard { + fn drop(&mut self) { + // Explicit for intent; closing the fd would release it regardless. + let _ = self.0.unlock(); + } +} + +/// Acquire the cross-process auth lock at `path`, polling until `deadline`. +/// +/// `File::try_lock` is per–open-file-description, so a lock taken on one +/// handle blocks every other handle — same process or not — which is exactly +/// the single-flight guarantee we want without a separate in-memory registry. +/// The lock is non-blocking, so we poll on [`LOCK_POLL_INTERVAL`] rather than +/// parking a worker thread in a blocking `lock()`. A waiter whose `deadline` +/// lapses returns [`AuthError::LockTimeout`]; because the caller sets that +/// deadline longer than [`AUTH_ATTEMPT_DEADLINE`], a healthy holder always +/// finishes first. +async fn acquire_auth_lock( + path: &Path, + deadline: std::time::Instant, +) -> Result { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|_| AuthError::LockTimeout)?; + } + let file = fs::OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .open(path) + .map_err(|_| AuthError::LockTimeout)?; + loop { + match file.try_lock() { + Ok(()) => return Ok(AuthLockGuard(file)), + Err(fs::TryLockError::WouldBlock) => { + if std::time::Instant::now() >= deadline { + return Err(AuthError::LockTimeout); + } + tokio::time::sleep(LOCK_POLL_INTERVAL).await; + } + Err(fs::TryLockError::Error(_)) => return Err(AuthError::LockTimeout), + } + } +} + /// Load a cached token, enforcing the owner-only invariant on load. /// /// Owner-only permissions are a cache *lifecycle* invariant, not just a @@ -712,21 +1112,37 @@ fn sanitize_callback_detail(raw: &str) -> String { .collect() } -/// Spin up a localhost callback server, open the authorize URL in a -/// browser, wait up to [`BROWSER_AUTH_TIMEOUT`] for the redirect, then -/// exchange the code for a token. +/// Spin up a localhost callback server, hand the authorize URL to `opener`, +/// wait up to [`BROWSER_AUTH_TIMEOUT`] for the redirect, then exchange the +/// code for a token. +/// +/// `opener` is invoked *after* the listener is bound and the abort guard is +/// armed, so a launch failure never returns a URL pointing at a torn-down +/// listener. Every failure is a typed [`AuthError`] so the coordinator can +/// record a cooldown (or not) by category: an open failure is +/// [`BrowserOpenFailed`], a redirect that never arrives is [`TimedOut`], a +/// provider-reported denial is [`Denied`], and a rejected code exchange is +/// [`ExchangeFailed`]; infrastructure faults (bind/exchange transport) are +/// [`NetworkUnavailable`]. +/// +/// [`BrowserOpenFailed`]: AuthError::BrowserOpenFailed +/// [`TimedOut`]: AuthError::TimedOut +/// [`Denied`]: AuthError::Denied +/// [`ExchangeFailed`]: AuthError::ExchangeFailed +/// [`NetworkUnavailable`]: AuthError::NetworkUnavailable async fn browser_pkce_flow( http: &Client, cfg: &PkceOAuthConfig, endpoints: &OidcEndpoints, -) -> Result { + opener: &dyn BrowserOpener, +) -> Result { use axum::{extract::Query, response::Html, routing::get, Router}; use std::collections::HashMap; use std::net::SocketAddr; use tokio::sync::oneshot; - let (verifier, challenge) = pkce_pair()?; - let state = random_state()?; + let (verifier, challenge) = pkce_pair().map_err(|_| AuthError::NetworkUnavailable)?; + let state = random_state().map_err(|_| AuthError::NetworkUnavailable)?; let (tx, rx) = oneshot::channel::>(); let tx = Arc::new(Mutex::new(Some(tx))); @@ -749,10 +1165,10 @@ async fn browser_pkce_flow( let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))) .await - .map_err(|e| AgentError::Llm(format!("oauth callback bind: {e}")))?; + .map_err(|_| AuthError::NetworkUnavailable)?; let port = listener .local_addr() - .map_err(|e| AgentError::Llm(format!("oauth callback addr: {e}")))? + .map_err(|_| AuthError::NetworkUnavailable)? .port(); let redirect_uri = format!("http://localhost:{port}"); @@ -774,14 +1190,25 @@ async fn browser_pkce_flow( urlencoding::encode(&challenge), ); - eprintln!("Opening browser for authentication. If it doesn't open, visit:\n {auth_url}"); - let _ = webbrowser::open(&auth_url); + // Launch the browser while the listener is live. A launch failure aborts + // before we wait on a redirect nobody can send. + opener.open(&auth_url).map_err(|e| { + tracing::warn!(error = %e, "oauth browser launch failed"); + AuthError::BrowserOpenFailed + })?; - let code = tokio::time::timeout(BROWSER_AUTH_TIMEOUT, rx) - .await - .map_err(|_| AgentError::Llm("oauth: browser auth timed out".into()))? - .map_err(|_| AgentError::Llm("oauth: callback sender dropped".into()))? - .map_err(|e| AgentError::Llm(format!("oauth callback: {e}")))?; + let code = match tokio::time::timeout(BROWSER_AUTH_TIMEOUT, rx).await { + // Timed out waiting for the redirect. + Err(_) => return Err(AuthError::TimedOut), + // Callback task dropped the sender without sending — treat as timeout. + Ok(Err(_)) => return Err(AuthError::TimedOut), + // Provider/user reported an error (denial, state mismatch, missing code). + Ok(Ok(Err(detail))) => { + tracing::warn!(detail = %detail, "oauth callback reported failure"); + return Err(AuthError::Denied); + } + Ok(Ok(Ok(code))) => code, + }; // Exchange code for token. let params = [ @@ -796,21 +1223,20 @@ async fn browser_pkce_flow( .form(¶ms) .send() .await - .map_err(|e| AgentError::Llm(format!("oauth exchange: {e}")))?; + .map_err(|_| AuthError::NetworkUnavailable)?; if !resp.status().is_success() { let body = resp.text().await.unwrap_or_default(); - return Err(AgentError::Llm(format!("oauth exchange failed: {body}"))); + tracing::warn!(status = "error", body = %body, "oauth code exchange rejected"); + return Err(AuthError::ExchangeFailed); } - let v: Value = resp - .json() - .await - .map_err(|e| AgentError::Llm(format!("oauth exchange json: {e}")))?; - token_from_response(&v, None) + let v: Value = resp.json().await.map_err(|_| AuthError::ExchangeFailed)?; + token_from_response(&v, None).map_err(|_| AuthError::ExchangeFailed) } #[cfg(test)] mod tests { use super::*; + use std::time::Instant; #[test] fn pkce_pair_produces_valid_challenge() { @@ -944,10 +1370,13 @@ mod tests { } #[tokio::test] - async fn test_bearer_falls_through_to_browser_when_disk_also_expired() { + async fn test_bearer_headless_no_credential_is_terminal_without_browser() { let dir = tempfile::tempdir().unwrap(); let cfg = PkceOAuthConfig { - discovery_url: "https://example.com/.well-known".into(), + // Unreachable discovery URL: if bearer() ever attempts discovery or + // a browser flow, this test would hang or error differently. The + // headless path must not touch either. + discovery_url: "https://invalid.example.test/.well-known".into(), client_id: "test-client".into(), scopes: vec!["offline_access".into()], cache_namespace: "test".into(), @@ -955,7 +1384,7 @@ mod tests { }; let source = PkceOAuthTokenSource::new(cfg).unwrap(); - // Expire the in-memory state. + // Expire the in-memory state with no refresh token. { let mut state = source.state.lock().await; *state = Some(CachedToken { @@ -965,7 +1394,7 @@ mod tests { }); } - // Write an expired token to disk too. + // Write an expired, refresh-less token to disk too. let expired_token = CachedToken { access_token: "also-stale".into(), refresh_token: None, @@ -974,27 +1403,25 @@ mod tests { let body = serde_json::to_vec_pretty(&expired_token).unwrap(); fs::write(&source.cache_path, &body).unwrap(); - // bearer() should fall through past the disk check. - // It will fail at the endpoints() discovery call since there's no server, - // which proves it didn't short-circuit on the expired disk token. - let result = source.bearer().await; - assert!(result.is_err()); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("oauth discovery"), - "expected discovery error, got: {err_msg}" - ); + // bearer() is a Headless acquisition: past the cache checks with no + // refresh token, it returns terminally instead of opening a browser. + // With no refresh token it never even discovers endpoints, so the + // unreachable URL is never contacted — the error is a graceful + // LlmAuth, not a hard Llm/discovery error. + match source.bearer().await.unwrap_err() { + AgentError::LlmAuth(_) => {} // correct: terminal, no browser + other => panic!("expected terminal LlmAuth, got: {other:?}"), + } } - /// `try_bearer_no_browser` with an empty cache and no refresh token must + /// `bearer_no_browser` with an empty cache and no refresh token must /// return `LlmAuth` immediately — it must NOT attempt OIDC discovery even - /// when the `discovery_url` is unreachable/invalid. This guards the - /// regression where `endpoints()` was called unconditionally before the - /// refresh-token check, causing an `Llm` error (hard failure) instead of - /// the intended graceful `LlmAuth` fallback. + /// when the `discovery_url` is unreachable/invalid, and must never browse. + /// This guards the regression where `endpoints()` was called + /// unconditionally before the refresh-token check, causing an `Llm` error + /// (hard failure) instead of the intended graceful `LlmAuth` fallback. #[tokio::test] - async fn test_try_bearer_no_browser_empty_cache_no_refresh_returns_llm_auth_without_discovery() - { + async fn test_bearer_no_browser_empty_cache_no_refresh_returns_llm_auth_without_discovery() { let dir = tempfile::tempdir().unwrap(); // Intentionally invalid/unreachable discovery URL — if endpoints() is // called, the test will get an `Llm` error and the assertion below fails. @@ -1016,7 +1443,7 @@ mod tests { // No disk cache file either — dir is empty. - let result = source.try_bearer_no_browser().await; + let result = source.bearer_no_browser().await; assert!(result.is_err(), "expected Err, got Ok"); match result.unwrap_err() { AgentError::LlmAuth(_) => {} // correct: graceful fallback @@ -1342,4 +1769,69 @@ mod tests { "read_cache followed a symlinked cache path" ); } + + // ---- cross-process advisory lock primitive -------------------------- + // + // The full 165s waiter bound (`LOCK_WAIT_TIMEOUT`) is not exercisable in a + // unit test, so these drive `acquire_auth_lock` with explicit deadlines to + // pin the three properties the coordinator relies on: a contended waiter + // times out (never blocks forever), a timeout leaves the *holder* + // untouched (never cancels the in-flight attempt), and releasing the + // holder — the RAII stand-in for a crashed process — lets a successor + // proceed with no wedge and no lock-breaking. + + #[tokio::test] + async fn test_lock_wait_times_out_and_leaves_holder_untouched() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("cache.json.lock"); + + // Holder takes the lock with a generous deadline. + let holder = acquire_auth_lock(&path, Instant::now() + Duration::from_secs(30)) + .await + .expect("holder should acquire the free lock"); + + // A waiter with an already-lapsed deadline must give up with + // LockTimeout rather than block — this is the deadline-aware polling + // that replaces a blocking `lock()`. + let waiter = acquire_auth_lock(&path, Instant::now()).await; + assert!( + matches!(waiter, Err(AuthError::LockTimeout)), + "contended waiter past its deadline must return LockTimeout, got {waiter:?}" + ); + + // The timeout did not cancel or steal the holder: a second immediate + // waiter still cannot acquire, proving the holder is intact. + let still_held = acquire_auth_lock(&path, Instant::now()).await; + assert!( + matches!(still_held, Err(AuthError::LockTimeout)), + "holder must remain intact after a waiter times out, got {still_held:?}" + ); + + drop(holder); + } + + #[tokio::test] + async fn test_lock_released_on_holder_drop_lets_successor_proceed() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("cache.json.lock"); + + let holder = acquire_auth_lock(&path, Instant::now() + Duration::from_secs(30)) + .await + .expect("holder should acquire the free lock"); + // Confirm contention while held. + assert!(matches!( + acquire_auth_lock(&path, Instant::now()).await, + Err(AuthError::LockTimeout) + )); + + // Dropping the guard is the RAII stand-in for the holder process + // dying: the kernel releases the advisory lock, so a successor + // acquires without any PID inspection or lock breaking. + drop(holder); + let successor = acquire_auth_lock(&path, Instant::now() + Duration::from_secs(30)).await; + assert!( + successor.is_ok(), + "successor must acquire after the holder releases, got {successor:?}" + ); + } } diff --git a/crates/buzz-agent/tests/databricks_auth_coordinator.rs b/crates/buzz-agent/tests/databricks_auth_coordinator.rs new file mode 100644 index 000000000..a5b0798f2 --- /dev/null +++ b/crates/buzz-agent/tests/databricks_auth_coordinator.rs @@ -0,0 +1,589 @@ +//! Concurrency-matrix tests for the Databricks auth coordinator. +//! +//! The coordinator single-flights OAuth acquisition per cache key using an OS +//! advisory lock, so one browser dance is shared and failures are coalesced +//! through a durable cooldown sidecar. These tests drive the public API +//! (`acquire_with_intent`, `interactive_login`) with an injected +//! [`BrowserOpener`] that scripts the localhost callback instead of popping a +//! real window — the browser step becomes deterministic and countable. +//! +//! Two `PkceOAuthTokenSource` instances sharing one cache path model two +//! processes: `File::try_lock` is per open-file-description, so distinct +//! handles contend whether or not they live in the same process. The +//! lock-primitive, crash-release, and lock-timeout edges live in the in-crate +//! `auth::tests` module where the private helpers are reachable. + +use std::io::Write; +use std::net::{SocketAddr, TcpStream}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use axum::extract::Form; +use axum::{routing::get, routing::post, Json, Router}; +use buzz_agent::auth::{ + AuthError, AuthIntent, BrowserOpener, PkceOAuthConfig, PkceOAuthTokenSource, +}; +use serde::Deserialize; +use serde_json::json; +use tempfile::TempDir; + +// ---- scripted browser opener -------------------------------------------- + +/// What the scripted "user" does when the coordinator opens a browser. +#[derive(Clone, Copy)] +enum Script { + /// Redirect with a valid `code`+`state` → the flow exchanges it for a + /// token and succeeds. + Approve, + /// Redirect with `error=access_denied` → the flow returns `Denied`. + Deny, + /// Every launch strategy fails → the flow returns `BrowserOpenFailed` + /// without waiting on a listener nobody will reach. + FailToOpen, +} + +/// A [`BrowserOpener`] that counts launches and drives the localhost callback +/// on a background thread, so the caller's callback wait observes the redirect +/// exactly as a real browser would deliver it. +#[derive(Clone)] +struct ScriptedOpener { + script: Script, + calls: Arc, +} + +impl ScriptedOpener { + fn new(script: Script) -> Self { + Self { + script, + calls: Arc::new(AtomicU64::new(0)), + } + } + + fn call_count(&self) -> u64 { + self.calls.load(Ordering::SeqCst) + } +} + +impl BrowserOpener for ScriptedOpener { + fn open(&self, url: &str) -> Result<(), String> { + self.calls.fetch_add(1, Ordering::SeqCst); + let query = match self.script { + Script::FailToOpen => return Err("no browser available".into()), + Script::Approve => "code=scripted-code", + Script::Deny => "error=access_denied", + }; + // Pull the loopback redirect target and the anti-CSRF state out of the + // authorize URL, then fire the callback from a separate thread so this + // synchronous `open()` returns and the flow proceeds to await it. + let parsed = url::Url::parse(url).expect("authorize URL must parse"); + let redirect = parsed + .query_pairs() + .find(|(k, _)| k == "redirect_uri") + .map(|(_, v)| v.into_owned()) + .expect("authorize URL carries redirect_uri"); + let state = parsed + .query_pairs() + .find(|(k, _)| k == "state") + .map(|(_, v)| v.into_owned()) + .expect("authorize URL carries state"); + let redirect = url::Url::parse(&redirect).expect("redirect_uri must parse"); + // The coordinator's listener binds 127.0.0.1; connect there directly so + // the callback can't land on an IPv6 `localhost` (::1) with no listener. + let port = redirect.port().expect("loopback redirect carries a port"); + // `state` is base64url (no reserved characters), safe to inline. + let request = format!( + "GET /?{query}&state={state} HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nConnection: close\r\n\r\n" + ); + std::thread::spawn(move || { + // A real browser holds the connection open until the callback page + // responds; do the same so hyper dispatches the request before the + // socket closes (a bare write+drop races the server and is lost). + if let Ok(mut sock) = TcpStream::connect(("127.0.0.1", port)) { + use std::io::Read; + let _ = sock.write_all(request.as_bytes()); + let _ = sock.flush(); + let mut discard = Vec::new(); + let _ = sock.read_to_end(&mut discard); + } + }); + Ok(()) + } +} + +// ---- stub OIDC provider -------------------------------------------------- + +#[derive(Deserialize)] +struct TokenForm { + grant_type: String, +} + +struct Stub { + base: String, + /// authorization-code exchanges served (browser flows completed). + code_grants: Arc, + /// refresh-token grants served. + refresh_grants: Arc, +} + +/// Boot a stub provider. `reject_refresh` makes the token endpoint 401 every +/// refresh-token grant (a dead refresh token); authorization-code grants +/// always succeed with a fresh token. +async fn spawn_stub(reject_refresh: bool) -> Stub { + let code_grants = Arc::new(AtomicU64::new(0)); + let refresh_grants = Arc::new(AtomicU64::new(0)); + + let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))) + .await + .unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let disco_base = base.clone(); + + let discovery = move || { + let base = disco_base.clone(); + async move { + Json(json!({ + "authorization_endpoint": format!("{base}/authorize"), + "token_endpoint": format!("{base}/token"), + })) + } + }; + + let code_for_token = code_grants.clone(); + let refresh_for_token = refresh_grants.clone(); + let app = Router::new() + // Two discovery paths so distinct-host tests derive distinct cache + // keys (the key hashes the discovery URL) from one stub. + .route("/disco/a", get(discovery.clone())) + .route("/disco/b", get(discovery)) + .route( + "/token", + post(move |Form(form): Form| { + let code_grants = code_for_token.clone(); + let refresh_grants = refresh_for_token.clone(); + let reject_refresh = reject_refresh; + async move { + if form.grant_type == "refresh_token" { + let n = refresh_grants.fetch_add(1, Ordering::SeqCst) + 1; + if reject_refresh { + return ( + axum::http::StatusCode::UNAUTHORIZED, + Json(json!({ "error": "invalid_grant" })), + ); + } + return ( + axum::http::StatusCode::OK, + Json(json!({ + "access_token": format!("refreshed-token-{n}"), + "refresh_token": "rotated-refresh", + "expires_in": 3600, + })), + ); + } + let n = code_grants.fetch_add(1, Ordering::SeqCst) + 1; + ( + axum::http::StatusCode::OK, + Json(json!({ + "access_token": format!("browser-token-{n}"), + "refresh_token": "browser-refresh", + "expires_in": 3600, + })), + ) + } + }), + ); + + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + Stub { + base, + code_grants, + refresh_grants, + } +} + +fn config(stub: &Stub, disco_path: &str, cache_dir: &std::path::Path) -> PkceOAuthConfig { + PkceOAuthConfig { + discovery_url: format!("{}{disco_path}", stub.base), + client_id: "test-client".into(), + scopes: vec!["offline_access".into()], + cache_namespace: "databricks".into(), + cache_dir_override: Some(cache_dir.to_path_buf()), + } +} + +fn future_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600 +} + +fn seed_cache(cfg: &PkceOAuthConfig, cache_dir: &std::path::Path, body: serde_json::Value) { + use sha2::Digest; + let mut h = sha2::Sha256::new(); + h.update(cfg.discovery_url.as_bytes()); + h.update(b"|"); + h.update(cfg.client_id.as_bytes()); + h.update(b"|"); + h.update(cfg.scopes.join(",").as_bytes()); + let hash = hex::encode(h.finalize()); + let path = cache_dir + .join(&cfg.cache_namespace) + .join(format!("{hash}.json")); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, serde_json::to_vec(&body).unwrap()).unwrap(); +} + +// ---- acceptance matrix --------------------------------------------------- + +#[tokio::test] +async fn test_same_key_concurrent_callers_share_one_browser_attempt() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Approve); + + // Two independent sources on the SAME cache key = two processes racing. + let a = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(opener.clone()), + ) + .unwrap(); + let b = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(opener.clone()), + ) + .unwrap(); + + let (ra, rb) = tokio::join!( + a.acquire_with_intent(AuthIntent::Auto), + b.acquire_with_intent(AuthIntent::Auto), + ); + let ta = ra.expect("first caller authenticates"); + let tb = rb.expect("second caller authenticates"); + + // One browser launch, one code exchange, one shared token. + assert_eq!( + opener.call_count(), + 1, + "only one browser attempt for one key" + ); + assert_eq!( + stub.code_grants.load(Ordering::SeqCst), + 1, + "exactly one authorization-code exchange" + ); + assert_eq!(ta, tb, "both callers observe the same token"); + assert_eq!(ta, "browser-token-1"); +} + +#[tokio::test] +async fn test_denied_then_auto_reads_cooldown_without_second_launch() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Deny); + + let src = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(opener.clone()), + ) + .unwrap(); + + let first = src.acquire_with_intent(AuthIntent::Auto).await; + assert_eq!( + first, + Err(AuthError::Denied), + "first Auto attempt is denied" + ); + assert_eq!(opener.call_count(), 1); + + // The denial wrote a cooldown; a subsequent Auto caller reads it and + // returns the recorded outcome instead of popping a second browser. + let second = src.acquire_with_intent(AuthIntent::Auto).await; + assert_eq!( + second, + Err(AuthError::Denied), + "queued Auto caller honors the cooldown" + ); + assert_eq!( + opener.call_count(), + 1, + "cooldown suppresses the second browser launch" + ); +} + +#[tokio::test] +async fn test_userinitiated_denial_is_visible_to_crossprocess_auto() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Deny); + + // Process A: an explicit UserInitiated attempt is denied. + let proc_a = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(opener.clone()), + ) + .unwrap(); + let denied = proc_a.acquire_with_intent(AuthIntent::UserInitiated).await; + assert_eq!(denied, Err(AuthError::Denied)); + assert_eq!(opener.call_count(), 1); + + // Process B: a passive Auto caller (distinct instance = distinct process) + // reads the durable sidecar A wrote and does not launch a second browser. + // This is the cross-policy edge: the sidecar is written for ANY failed + // interactive attempt, only the reader policy differs. + let proc_b = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(opener.clone()), + ) + .unwrap(); + let auto = proc_b.acquire_with_intent(AuthIntent::Auto).await; + assert_eq!( + auto, + Err(AuthError::Denied), + "cross-process Auto reads the UserInitiated failure sidecar" + ); + assert_eq!( + opener.call_count(), + 1, + "no second browser across the policy/process boundary" + ); +} + +#[tokio::test] +async fn test_userinitiated_retry_bypasses_cooldown_and_reopens() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + + // First attempt: denied, writes a cooldown. + let deny_opener = ScriptedOpener::new(Script::Deny); + let denier = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(deny_opener.clone()), + ) + .unwrap(); + assert_eq!( + denier.acquire_with_intent(AuthIntent::UserInitiated).await, + Err(AuthError::Denied) + ); + + // The user explicitly retries: UserInitiated bypasses (and clears) the + // cooldown and opens a fresh browser, which now succeeds. + let approve_opener = ScriptedOpener::new(Script::Approve); + let retrier = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(approve_opener.clone()), + ) + .unwrap(); + let token = retrier + .acquire_with_intent(AuthIntent::UserInitiated) + .await + .expect("explicit retry re-launches the browser and succeeds"); + assert_eq!(token, "browser-token-1"); + assert_eq!( + approve_opener.call_count(), + 1, + "UserInitiated retry launches despite the prior cooldown" + ); + + // Cooldown cleared on success: a follow-up Auto now sees a valid token, + // never the stale denial. + let auto = retrier.acquire_with_intent(AuthIntent::Auto).await; + assert_eq!(auto, Ok("browser-token-1".to_string())); +} + +#[tokio::test] +async fn test_distinct_hosts_do_not_inherit_cooldown() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + + // Host A is denied and records a cooldown under key A. + let deny_opener = ScriptedOpener::new(Script::Deny); + let host_a = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(deny_opener.clone()), + ) + .unwrap(); + assert_eq!( + host_a.acquire_with_intent(AuthIntent::Auto).await, + Err(AuthError::Denied) + ); + + // Host B is a different key (different discovery URL). It must NOT inherit + // A's cooldown: an Auto caller launches its own browser and succeeds. + let approve_opener = ScriptedOpener::new(Script::Approve); + let host_b = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/b", cache.path()), + Arc::new(approve_opener.clone()), + ) + .unwrap(); + let token = host_b + .acquire_with_intent(AuthIntent::Auto) + .await + .expect("distinct host is unaffected by another key's cooldown"); + assert_eq!(token, "browser-token-1"); + assert_eq!(approve_opener.call_count(), 1); +} + +#[tokio::test] +async fn test_browser_open_failure_is_typed_and_retryable_by_user() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + + // Every launch strategy fails: the flow reports the typed BrowserOpenFailed + // without waiting on a listener nobody will reach. + let fail_opener = ScriptedOpener::new(Script::FailToOpen); + let failing = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(fail_opener.clone()), + ) + .unwrap(); + let result = failing.acquire_with_intent(AuthIntent::UserInitiated).await; + assert_eq!( + result, + Err(AuthError::BrowserOpenFailed), + "a failed launch surfaces as the typed BrowserOpenFailed" + ); + assert_eq!(fail_opener.call_count(), 1); + + // A failed launch writes a cooldown, but a UserInitiated retry bypasses it + // and reopens — a transient "no browser" (e.g. race with a display coming + // up) must never wedge an explicit user sign-in. + let approve_opener = ScriptedOpener::new(Script::Approve); + let retrier = PkceOAuthTokenSource::new_with( + config(&stub, "/disco/a", cache.path()), + Arc::new(approve_opener.clone()), + ) + .unwrap(); + let token = retrier + .acquire_with_intent(AuthIntent::UserInitiated) + .await + .expect("explicit retry reopens despite the prior launch failure"); + assert_eq!(token, "browser-token-1"); + assert_eq!(approve_opener.call_count(), 1); +} + +#[tokio::test] +async fn test_headless_dead_refresh_returns_refresh_rejected_without_browser() { + let stub = spawn_stub(true).await; // refresh grants 401 + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Approve); + let cfg = config(&stub, "/disco/a", cache.path()); + + // Expired token WITH a refresh token, but the server rejects the refresh + // grant (dead/rotated). A Headless caller must classify this terminally as + // RefreshRejected and never open a browser. + seed_cache( + &cfg, + cache.path(), + json!({ + "access_token": "stale", + "refresh_token": "dead-refresh", + "expires_at": 1u64, + }), + ); + + let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap(); + let result = src.acquire_with_intent(AuthIntent::Headless).await; + assert_eq!( + result, + Err(AuthError::RefreshRejected), + "Headless dead-refresh is terminal RefreshRejected" + ); + assert_eq!(opener.call_count(), 0, "Headless never opens a browser"); + assert_eq!( + stub.refresh_grants.load(Ordering::SeqCst), + 1, + "the refresh grant was attempted exactly once" + ); +} + +#[tokio::test] +async fn test_interactive_dead_refresh_converts_to_browser() { + let stub = spawn_stub(true).await; // refresh grants 401 + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Approve); + let cfg = config(&stub, "/disco/a", cache.path()); + + // Same dead-refresh seed, but an interactive intent must fall through to a + // browser flow instead of failing terminally. + seed_cache( + &cfg, + cache.path(), + json!({ + "access_token": "stale", + "refresh_token": "dead-refresh", + "expires_at": 1u64, + }), + ); + + let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap(); + let token = src + .acquire_with_intent(AuthIntent::UserInitiated) + .await + .expect("interactive intent recovers via the browser"); + assert_eq!(token, "browser-token-1"); + assert_eq!(opener.call_count(), 1, "interactive intent opens a browser"); + assert_eq!(stub.refresh_grants.load(Ordering::SeqCst), 1); + assert_eq!(stub.code_grants.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_headless_expired_token_live_refresh_recovers_silently() { + let stub = spawn_stub(false).await; // refresh succeeds + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Approve); + let cfg = config(&stub, "/disco/a", cache.path()); + + seed_cache( + &cfg, + cache.path(), + json!({ + "access_token": "stale", + "refresh_token": "live-refresh", + "expires_at": 1u64, + }), + ); + + let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap(); + let token = src + .acquire_with_intent(AuthIntent::Headless) + .await + .expect("live refresh recovers a Headless caller silently"); + assert_eq!(token, "refreshed-token-1"); + assert_eq!(opener.call_count(), 0, "no browser on a live refresh"); + assert_eq!(stub.refresh_grants.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_interactive_login_reuses_valid_cache_without_browser() { + let stub = spawn_stub(false).await; + let cache = TempDir::new().unwrap(); + let opener = ScriptedOpener::new(Script::Approve); + let cfg = config(&stub, "/disco/a", cache.path()); + + // A still-valid cached token short-circuits interactive_login: an explicit + // sign-in should not re-prompt when a good token is already present. + seed_cache( + &cfg, + cache.path(), + json!({ + "access_token": "already-valid", + "refresh_token": "rt", + "expires_at": future_secs(), + }), + ); + + let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener.clone())).unwrap(); + src.interactive_login() + .await + .expect("interactive_login succeeds off the valid cache"); + assert_eq!( + opener.call_count(), + 0, + "a valid cached token means no browser prompt" + ); +}