buzz/crates/buzz-agent/src/auth.rs

//! Token sources for the LLM transport layer.
//!
//! [`TokenSource`] decouples request auth from `Config::api_key`: providers
//! can supply a static string ([`StaticTokenSource`]) or a refreshable OAuth
//! 2.0 PKCE engine ([`PkceOAuthTokenSource`]). Engines own their own cache
//! and refresh logic; the [`Llm`] just asks for a bearer per request.
//!
//! The PKCE engine implements RFC 6749 + RFC 7636 with on-disk token
//! caching keyed by `sha256(discovery_url|client_id|scopes)`. It's the
//! same shape goose uses for Databricks, but we own the wire format and
//! cache directory so the two are independently upgradable.
//!
//! First-use (cache empty) requires a browser: the engine opens
//! `authorization_endpoint` in `webbrowser`, listens on `127.0.0.1:0`,
//! captures the redirect, and exchanges the code for a token. Subsequent
//! calls hit the cache and silently refresh when expired.

use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use base64::Engine;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::Digest;
use tokio::sync::Mutex;

use crate::types::AgentError;

/// Buffer before `expires_at` to consider a cached token "still good".
/// Keeps us off the cliff if the clock or the server's clock drifts.
const TOKEN_REFRESH_LEEWAY: Duration = Duration::from_secs(60);

/// Wall-clock budget for the interactive browser dance. Goose uses 60s.
/// We match: any longer and the user has gone to lunch.
const BROWSER_AUTH_TIMEOUT: Duration = Duration::from_secs(60);

/// 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]
pub trait TokenSource: Send + Sync {
    async fn bearer(&self) -> Result<String, AgentError>;

    /// Return a bearer token from cache or refresh, **never** opening a browser.
    ///
    /// The default delegates to [`bearer`](Self::bearer) — correct for token
    /// sources (e.g. static API keys) that can never trigger a browser flow.
    /// [`PkceOAuthTokenSource`] overrides this to stop before the browser step.
    async fn bearer_no_browser(&self) -> Result<String, AgentError> {
        self.bearer().await
    }

    /// Force a fresh bearer after the server rejected the current one (401).
    ///
    /// `rejected` is the exact access token that just got the 401. Unlike
    /// [`bearer`](Self::bearer), which trusts the local expiry clock, this is
    /// driven by the server's verdict: the cached token looked valid to us
    /// (well within its local expiry) but the provider rejected it — clock
    /// skew, server-side revocation, or a node that never saw it. The clock
    /// therefore can't decide whether to refresh; the caller passes the
    /// rejected token so the impl can refresh unless a concurrent caller has
    /// *already* replaced it. Implementations must obtain a new token without
    /// any interactive step, so a headless harness never hangs. The default
    /// returns the existing bearer — correct for sources that can't refresh
    /// (a static key); the caller's retry then fails terminally rather than
    /// looping.
    async fn refresh_now(&self, _rejected: &str) -> Result<String, AgentError> {
        self.bearer().await
    }
}

/// A token that never changes for the life of the process.
pub struct StaticTokenSource(String);

impl StaticTokenSource {
    pub fn new(token: impl Into<String>) -> Self {
        Self(token.into())
    }
}

#[async_trait]
impl TokenSource for StaticTokenSource {
    async fn bearer(&self) -> Result<String, AgentError> {
        Ok(self.0.clone())
    }
}

/// Static config for an OAuth 2.0 Authorization Code + PKCE provider.
///
/// The `discovery_url` must return a JSON document with at least
/// `authorization_endpoint` and `token_endpoint` (RFC 8414). The
/// `cache_namespace` is the directory under `~/.config/buzz-agent/oauth/`
/// the token JSON lives in — separates providers' caches cleanly.
#[derive(Debug, Clone)]
pub struct PkceOAuthConfig {
    pub discovery_url: String,
    pub client_id: String,
    pub scopes: Vec<String>,
    pub cache_namespace: String,
    /// When `Some`, the engine writes tokens here instead of
    /// `~/.config/buzz-agent/oauth/<cache_namespace>/`. Production code
    /// leaves this `None`. Integration tests use it to avoid stomping on
    /// a shared `$HOME` when running in parallel.
    pub cache_dir_override: Option<PathBuf>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
struct CachedToken {
    access_token: String,
    refresh_token: Option<String>,
    /// Unix seconds. `None` means the server didn't advertise an expiry;
    /// we use it without checking and rely on refresh on 401.
    expires_at: Option<u64>,
}

#[derive(Debug, Clone)]
struct OidcEndpoints {
    authorization_endpoint: String,
    token_endpoint: String,
}

/// PKCE OAuth token source with on-disk refresh cache.
///
/// First call:
///   1. Loads from cache if present and unexpired.
///   2. Otherwise tries `refresh_token` if cached.
///   3. Otherwise runs the full browser flow.
///
/// Subsequent calls hit an in-memory copy of the cached token and only
/// touch disk/network if the access token is past `expires_at`.
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.
    state: Mutex<Option<CachedToken>>,
}

impl PkceOAuthTokenSource {
    pub fn new(cfg: PkceOAuthConfig) -> Result<Arc<Self>, 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}")))?;
        }
        let initial = read_cache(&cache_path);
        Ok(Arc::new(Self {
            cfg,
            http: Client::new(),
            cache_path,
            state: Mutex::new(initial),
        }))
    }

    /// Discover authorization + token endpoints from the well-known URL.
    async fn endpoints(&self) -> Result<OidcEndpoints, AgentError> {
        let v: Value = self
            .http
            .get(&self.cfg.discovery_url)
            .send()
            .await
            .map_err(|e| AgentError::Llm(format!("oauth discovery: {e}")))?
            .error_for_status()
            .map_err(|e| AgentError::Llm(format!("oauth discovery status: {e}")))?
            .json()
            .await
            .map_err(|e| AgentError::Llm(format!("oauth discovery json: {e}")))?;
        let auth = v
            .get("authorization_endpoint")
            .and_then(Value::as_str)
            .ok_or_else(|| {
                AgentError::Llm("oauth discovery: authorization_endpoint missing".into())
            })?
            .to_string();
        let token = v
            .get("token_endpoint")
            .and_then(Value::as_str)
            .ok_or_else(|| AgentError::Llm("oauth discovery: token_endpoint missing".into()))?
            .to_string();
        Ok(OidcEndpoints {
            authorization_endpoint: auth,
            token_endpoint: token,
        })
    }

    /// Persist a token to disk and the in-memory cell.
    fn save(&self, state: &mut Option<CachedToken>, token: CachedToken) -> Result<(), AgentError> {
        let body = serde_json::to_vec_pretty(&token)
            .map_err(|e| AgentError::Llm(format!("oauth cache serialize: {e}")))?;
        // Atomic rename so a concurrent reader never sees a partial write.
        let tmp = self.cache_path.with_extension("json.tmp");
        fs::write(&tmp, &body)
            .map_err(|e| AgentError::Llm(format!("oauth cache write {tmp:?}: {e}")))?;
        fs::rename(&tmp, &self.cache_path)
            .map_err(|e| AgentError::Llm(format!("oauth cache rename: {e}")))?;
        *state = Some(token);
        Ok(())
    }

    /// Exchange a refresh token for a fresh access token.
    async fn refresh(
        &self,
        endpoints: &OidcEndpoints,
        refresh_token: &str,
    ) -> Result<CachedToken, AgentError> {
        let params = [
            ("grant_type", "refresh_token"),
            ("refresh_token", refresh_token),
            ("client_id", &self.cfg.client_id),
        ];
        let resp = self
            .http
            .post(&endpoints.token_endpoint)
            .form(&params)
            .send()
            .await
            .map_err(|e| AgentError::Llm(format!("oauth refresh: {e}")))?;
        if !resp.status().is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(AgentError::Llm(format!("oauth refresh failed: {body}")));
        }
        let v: Value = resp
            .json()
            .await
            .map_err(|e| AgentError::Llm(format!("oauth refresh json: {e}")))?;
        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.
    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)?;
        Ok(())
    }
}

#[async_trait]
impl TokenSource for PkceOAuthTokenSource {
    async fn bearer(&self) -> Result<String, AgentError> {
        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)
    }

    async fn bearer_no_browser(&self) -> Result<String, AgentError> {
        self.try_bearer_no_browser().await
    }

    /// Force-refresh after a 401, never touching the browser flow.
    ///
    /// `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.
    async fn refresh_now(&self, rejected: &str) -> Result<String, AgentError> {
        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<String, AgentError> {
        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(),
        ))
    }
}

// ---- helpers -------------------------------------------------------------

/// Aborts a spawned task when dropped. Used to guarantee the localhost
/// callback server doesn't outlive a failed/abandoned PKCE attempt.
struct AbortOnDrop(tokio::task::JoinHandle<()>);

impl Drop for AbortOnDrop {
    fn drop(&mut self) {
        self.0.abort();
    }
}

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
}

fn cache_path_for(cfg: &PkceOAuthConfig) -> Result<PathBuf, AgentError> {
    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 dir = match &cfg.cache_dir_override {
        Some(p) => p.join(&cfg.cache_namespace),
        None => dirs::home_dir()
            .ok_or_else(|| AgentError::Llm("oauth cache: home directory not found".into()))?
            .join(".config")
            .join("buzz-agent")
            .join("oauth")
            .join(&cfg.cache_namespace),
    };
    Ok(dir.join(format!("{hash}.json")))
}

fn read_cache(path: &PathBuf) -> Option<CachedToken> {
    let body = fs::read(path).ok()?;
    serde_json::from_slice(&body).ok()
}

/// Parse a token-endpoint JSON response. Fails loudly when `access_token`
/// is missing or empty — without this, a malformed server response would
/// be cached and `bearer()` would silently return `""` until the entry
/// expires or is deleted by hand.
fn token_from_response(
    v: &Value,
    fallback_refresh: Option<&str>,
) -> Result<CachedToken, AgentError> {
    let access_token = v
        .get("access_token")
        .and_then(Value::as_str)
        .filter(|s| !s.is_empty())
        .ok_or_else(|| AgentError::Llm("oauth: token response missing/empty access_token".into()))?
        .to_string();
    let refresh_token = v
        .get("refresh_token")
        .and_then(Value::as_str)
        .map(str::to_string)
        .or_else(|| fallback_refresh.map(str::to_string));
    let expires_at = v.get("expires_in").and_then(Value::as_u64).map(|secs| {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0)
            + secs
    });
    Ok(CachedToken {
        access_token,
        refresh_token,
        expires_at,
    })
}

/// PKCE pieces: URL-safe random verifier (~64 chars) and its SHA-256
/// challenge (RFC 7636 §4.2).
fn pkce_pair() -> Result<(String, String), AgentError> {
    let mut bytes = [0u8; 48];
    getrandom::fill(&mut bytes).map_err(|e| AgentError::Llm(format!("pkce rng: {e}")))?;
    let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
    let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .encode(sha2::Sha256::digest(verifier.as_bytes()));
    Ok((verifier, challenge))
}

fn random_state() -> Result<String, AgentError> {
    let mut bytes = [0u8; 16];
    getrandom::fill(&mut bytes).map_err(|e| AgentError::Llm(format!("state rng: {e}")))?;
    Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes))
}

/// 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.
async fn browser_pkce_flow(
    http: &Client,
    cfg: &PkceOAuthConfig,
    endpoints: &OidcEndpoints,
) -> Result<CachedToken, AgentError> {
    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 (tx, rx) = oneshot::channel::<Result<String, String>>();
    let tx = Arc::new(Mutex::new(Some(tx)));

    let expected_state = state.clone();
    let app = Router::new().route(
        "/",
        get(move |Query(params): Query<HashMap<String, String>>| {
            let tx = Arc::clone(&tx);
            let expected = expected_state.clone();
            async move {
                let result = match (params.get("code"), params.get("state")) {
                    (Some(code), Some(st)) if st == &expected => Ok(code.clone()),
                    (Some(_), Some(_)) => Err("state mismatch".to_string()),
                    _ => Err(params
                        .get("error")
                        .cloned()
                        .unwrap_or_else(|| "missing code".into())),
                };
                if let Some(sender) = tx.lock().await.take() {
                    let _ = sender.send(result.clone());
                }
                match result {
                    Ok(_) => Html(
                        "<h2>Buzz: signed in</h2><p>You can close this window.</p>".to_string(),
                    ),
                    Err(e) => Html(format!("<h2>Buzz auth failed</h2><pre>{e}</pre>")),
                }
            }
        }),
    );

    let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
        .await
        .map_err(|e| AgentError::Llm(format!("oauth callback bind: {e}")))?;
    let port = listener
        .local_addr()
        .map_err(|e| AgentError::Llm(format!("oauth callback addr: {e}")))?
        .port();
    let redirect_uri = format!("http://localhost:{port}");

    // `_server` is held until this function returns; the drop guard aborts
    // the axum task on every exit path (timeout, callback error, token
    // exchange failure, or success), so we never leak a listener bound to
    // 127.0.0.1 past the auth attempt.
    let _server = AbortOnDrop(tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    }));

    let auth_url = format!(
        "{}?response_type=code&client_id={}&redirect_uri={}&scope={}&state={}&code_challenge={}&code_challenge_method=S256",
        endpoints.authorization_endpoint,
        urlencoding::encode(&cfg.client_id),
        urlencoding::encode(&redirect_uri),
        urlencoding::encode(&cfg.scopes.join(" ")),
        urlencoding::encode(&state),
        urlencoding::encode(&challenge),
    );

    eprintln!("Opening browser for authentication. If it doesn't open, visit:\n  {auth_url}");
    let _ = webbrowser::open(&auth_url);

    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}")))?;

    // Exchange code for token.
    let params = [
        ("grant_type", "authorization_code"),
        ("code", &code),
        ("redirect_uri", &redirect_uri),
        ("code_verifier", &verifier),
        ("client_id", &cfg.client_id),
    ];
    let resp = http
        .post(&endpoints.token_endpoint)
        .form(&params)
        .send()
        .await
        .map_err(|e| AgentError::Llm(format!("oauth exchange: {e}")))?;
    if !resp.status().is_success() {
        let body = resp.text().await.unwrap_or_default();
        return Err(AgentError::Llm(format!("oauth exchange failed: {body}")));
    }
    let v: Value = resp
        .json()
        .await
        .map_err(|e| AgentError::Llm(format!("oauth exchange json: {e}")))?;
    token_from_response(&v, None)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pkce_pair_produces_valid_challenge() {
        let (verifier, challenge) = pkce_pair().unwrap();
        assert!(verifier.len() >= 43);
        let expected = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(sha2::Sha256::digest(verifier.as_bytes()));
        assert_eq!(expected, challenge);
    }

    #[test]
    fn cached_token_no_expiry_is_not_expired() {
        let t = CachedToken {
            access_token: "x".into(),
            refresh_token: None,
            expires_at: None,
        };
        assert!(!is_expired(&t));
    }

    #[test]
    fn cached_token_far_future_is_not_expired() {
        let future = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + 3600;
        let t = CachedToken {
            access_token: "x".into(),
            refresh_token: None,
            expires_at: Some(future),
        };
        assert!(!is_expired(&t));
    }

    #[test]
    fn cached_token_within_leeway_is_expired() {
        let near = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + 10; // 10s away, leeway is 60s → counts as expired
        let t = CachedToken {
            access_token: "x".into(),
            refresh_token: None,
            expires_at: Some(near),
        };
        assert!(is_expired(&t));
    }

    #[test]
    fn cache_path_uses_platform_home_directory() {
        let cfg = PkceOAuthConfig {
            discovery_url: "https://example.com/.well-known".into(),
            client_id: "abc".into(),
            scopes: vec!["a".into(), "b".into()],
            cache_namespace: "demo".into(),
            cache_dir_override: None,
        };
        let p = cache_path_for(&cfg).unwrap();
        let expected_dir = dirs::home_dir()
            .unwrap()
            .join(".config")
            .join("buzz-agent")
            .join("oauth")
            .join("demo");
        assert_eq!(p.parent(), Some(expected_dir.as_path()));
        assert_eq!(p.extension().and_then(|s| s.to_str()), Some("json"));
    }

    #[test]
    fn token_from_response_uses_fallback_refresh() {
        let v: Value = serde_json::from_str(r#"{"access_token":"abc","expires_in":3600}"#).unwrap();
        let t = token_from_response(&v, Some("old-refresh")).unwrap();
        assert_eq!(t.access_token, "abc");
        assert_eq!(t.refresh_token.as_deref(), Some("old-refresh"));
        assert!(t.expires_at.is_some());
    }

    #[test]
    fn token_from_response_rejects_missing_access_token() {
        let v: Value = serde_json::from_str(r#"{"expires_in":3600}"#).unwrap();
        assert!(token_from_response(&v, None).is_err());
    }

    #[test]
    fn token_from_response_rejects_empty_access_token() {
        let v: Value = serde_json::from_str(r#"{"access_token":""}"#).unwrap();
        assert!(token_from_response(&v, None).is_err());
    }

    #[tokio::test]
    async fn test_bearer_reuses_disk_token_after_expiry() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = PkceOAuthConfig {
            discovery_url: "https://example.com/.well-known".into(),
            client_id: "test-client".into(),
            scopes: vec!["offline_access".into()],
            cache_namespace: "test".into(),
            cache_dir_override: Some(dir.path().to_path_buf()),
        };
        let source = PkceOAuthTokenSource::new(cfg).unwrap();

        // Expire the in-memory state.
        {
            let mut state = source.state.lock().await;
            *state = Some(CachedToken {
                access_token: "stale".into(),
                refresh_token: None,
                expires_at: Some(0), // long expired
            });
        }

        // Write a valid token to disk (simulating another process refreshing).
        let future_exp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + 7200;
        let fresh_token = CachedToken {
            access_token: "fresh-from-disk".into(),
            refresh_token: Some("rt".into()),
            expires_at: Some(future_exp),
        };
        let body = serde_json::to_vec_pretty(&fresh_token).unwrap();
        fs::write(&source.cache_path, &body).unwrap();

        // bearer() should pick up the disk token without any network call.
        let result = source.bearer().await.unwrap();
        assert_eq!(result, "fresh-from-disk");
    }

    #[tokio::test]
    async fn test_bearer_falls_through_to_browser_when_disk_also_expired() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = PkceOAuthConfig {
            discovery_url: "https://example.com/.well-known".into(),
            client_id: "test-client".into(),
            scopes: vec!["offline_access".into()],
            cache_namespace: "test".into(),
            cache_dir_override: Some(dir.path().to_path_buf()),
        };
        let source = PkceOAuthTokenSource::new(cfg).unwrap();

        // Expire the in-memory state.
        {
            let mut state = source.state.lock().await;
            *state = Some(CachedToken {
                access_token: "stale".into(),
                refresh_token: None,
                expires_at: Some(0),
            });
        }

        // Write an expired token to disk too.
        let expired_token = CachedToken {
            access_token: "also-stale".into(),
            refresh_token: None,
            expires_at: Some(0),
        };
        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}"
        );
    }

    /// `try_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.
    #[tokio::test]
    async fn test_try_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.
        let cfg = PkceOAuthConfig {
            discovery_url: "https://invalid.example.test/.well-known/oauth-authorization-server"
                .into(),
            client_id: "test-client".into(),
            scopes: vec!["offline_access".into()],
            cache_namespace: "test".into(),
            cache_dir_override: Some(dir.path().to_path_buf()),
        };
        let source = PkceOAuthTokenSource::new(cfg).unwrap();

        // Empty in-memory state (no token, no refresh token).
        {
            let mut state = source.state.lock().await;
            *state = None;
        }

        // No disk cache file either — dir is empty.

        let result = source.try_bearer_no_browser().await;
        assert!(result.is_err(), "expected Err, got Ok");
        match result.unwrap_err() {
            AgentError::LlmAuth(_) => {} // correct: graceful fallback
            other => panic!(
                "expected LlmAuth (no discovery attempted), got: {other:?}\n\
                 This means endpoints() was called before the refresh-token check."
            ),
        }
    }
}