ai-harness/crates/harness-providers/src/copilot/token.rs

//! Copilot API token strategy (flagged-risk area — verify against the live API).
//!
//! opencode sends the GitHub OAuth token directly as the Bearer (`expires: 0`). The classic
//! Copilot API instead wants a short-lived token from `copilot_internal/v2/token`. We try the
//! exchange first and cache it until shortly before expiry; if the endpoint rejects us, we fall
//! back to the direct-Bearer behavior. Refresh is single-flight under a `tokio::Mutex`.

use serde::Deserialize;
use tokio::sync::Mutex;

const EXCHANGE_URL: &str = "https://api.github.com/copilot_internal/v2/token";
/// Refresh this many seconds before the reported expiry.
const EXPIRY_SKEW_SECS: i64 = 120;
const USER_AGENT: &str = concat!("ai-harness/", env!("CARGO_PKG_VERSION"));

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CopilotToken {
    pub token: String,
    /// Unix seconds; `0` means "never expires" (direct-Bearer fallback).
    pub expires_at: i64,
}

impl CopilotToken {
    /// True when the token should be refreshed at `now_secs` (never, for `expires_at == 0`).
    pub fn needs_refresh(&self, now_secs: i64) -> bool {
        self.expires_at != 0 && now_secs + EXPIRY_SKEW_SECS >= self.expires_at
    }
}

#[derive(Debug, thiserror::Error)]
pub enum TokenError {
    #[error("network: {0}")]
    Network(String),
    /// The exchange endpoint is unavailable/unauthorized — the caller should fall back to the
    /// direct-Bearer strategy.
    #[error("exchange unsupported (status {0})")]
    ExchangeUnsupported(u16),
    #[error("unexpected response: {0}")]
    Unexpected(String),
}

#[derive(Deserialize)]
struct ExchangeResponse {
    token: String,
    #[serde(default)]
    expires_at: i64,
}

/// Performs the token exchange once. `ExchangeUnsupported` signals the caller to fall back.
pub async fn exchange(
    client: &reqwest::Client,
    oauth_token: &str,
) -> Result<CopilotToken, TokenError> {
    let resp = client
        .get(EXCHANGE_URL)
        .header("authorization", format!("token {oauth_token}"))
        .header("accept", "application/json")
        .header("user-agent", USER_AGENT)
        .send()
        .await
        .map_err(|e| TokenError::Network(e.to_string()))?;

    let status = resp.status();
    if !status.is_success() {
        // 401/403/404 → this deployment doesn't support the exchange; fall back to direct Bearer.
        if matches!(status.as_u16(), 401 | 403 | 404) {
            return Err(TokenError::ExchangeUnsupported(status.as_u16()));
        }
        return Err(TokenError::Network(format!("status {}", status.as_u16())));
    }

    let parsed: ExchangeResponse = resp
        .json()
        .await
        .map_err(|e| TokenError::Unexpected(e.to_string()))?;
    Ok(CopilotToken {
        token: parsed.token,
        expires_at: parsed.expires_at,
    })
}

/// Caches the exchanged Copilot token and refreshes it single-flight. Falls back to using the
/// OAuth token directly (never-expiring) when the exchange endpoint rejects the request.
pub struct TokenProvider {
    oauth_token: String,
    client: reqwest::Client,
    cached: Mutex<Option<CopilotToken>>,
}

impl TokenProvider {
    pub fn new(oauth_token: impl Into<String>, client: reqwest::Client) -> Self {
        Self {
            oauth_token: oauth_token.into(),
            client,
            cached: Mutex::new(None),
        }
    }

    /// Returns a usable Bearer token, refreshing/exchanging as needed. `now_secs` is injected
    /// for testability.
    pub async fn token(&self, now_secs: i64) -> CopilotToken {
        let mut guard = self.cached.lock().await;
        if let Some(tok) = guard.as_ref() {
            if !tok.needs_refresh(now_secs) {
                return tok.clone();
            }
        }
        let fresh = match exchange(&self.client, &self.oauth_token).await {
            Ok(tok) => tok,
            Err(_) => CopilotToken {
                // Direct-Bearer fallback: use the OAuth token itself, never expiring.
                token: self.oauth_token.clone(),
                expires_at: 0,
            },
        };
        *guard = Some(fresh.clone());
        fresh
    }
}

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

    #[test]
    fn never_expiring_token_does_not_refresh() {
        let tok = CopilotToken {
            token: "t".into(),
            expires_at: 0,
        };
        assert!(!tok.needs_refresh(i64::MAX));
    }

    #[test]
    fn refresh_triggers_within_skew_window() {
        let tok = CopilotToken {
            token: "t".into(),
            expires_at: 1_000,
        };
        assert!(!tok.needs_refresh(800)); // 800 + 120 < 1000
        assert!(tok.needs_refresh(881)); // 881 + 120 >= 1000
        assert!(tok.needs_refresh(1_000));
    }
}