huskies/server/src/llm/oauth.rs

//! Anthropic OAuth — token refresh and credential management for Claude API access.
use crate::slog;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

/// The client ID used by Claude Code for OAuth.
const CLAUDE_CODE_CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
const TOKEN_ENDPOINT: &str = "https://platform.claude.com/v1/oauth/token";

/// OAuth credentials as stored in `~/.claude/.credentials.json`.
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OAuthCredentials {
    pub access_token: String,
    pub refresh_token: String,
    pub expires_at: u64,
    #[serde(default)]
    pub scopes: Vec<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub subscription_type: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub rate_limit_tier: Option<String>,
}

/// Top-level structure of `~/.claude/.credentials.json`.
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CredentialsFile {
    pub claude_ai_oauth: OAuthCredentials,
}

/// Response from the Anthropic OAuth token refresh endpoint.
#[derive(Debug, Deserialize)]
struct TokenRefreshResponse {
    access_token: String,
    expires_in: u64,
    #[allow(dead_code)]
    token_type: Option<String>,
}

/// Error from the Anthropic OAuth token refresh endpoint.
#[derive(Debug, Deserialize)]
struct TokenRefreshError {
    #[allow(dead_code)]
    error: String,
    error_description: Option<String>,
}

/// Returns the path to `~/.claude/.credentials.json`.
fn credentials_path() -> Result<PathBuf, String> {
    let home = std::env::var("HOME").map_err(|_| "HOME not set".to_string())?;
    Ok(PathBuf::from(home)
        .join(".claude")
        .join(".credentials.json"))
}

/// Read OAuth credentials from disk.
pub fn read_credentials() -> Result<CredentialsFile, String> {
    let path = credentials_path()?;
    let data = std::fs::read_to_string(&path).map_err(|e| {
        format!(
            "Cannot read {}: {e}. Run `claude login` to authenticate.",
            path.display()
        )
    })?;
    serde_json::from_str(&data).map_err(|e| format!("Failed to parse {}: {e}", path.display()))
}

/// Write updated credentials back to disk with 0600 permissions.
pub fn write_credentials(creds: &CredentialsFile) -> Result<(), String> {
    let path = credentials_path()?;
    let data = serde_json::to_string_pretty(creds)
        .map_err(|e| format!("Failed to serialize credentials: {e}"))?;
    std::fs::write(&path, &data).map_err(|e| format!("Failed to write {}: {e}", path.display()))?;

    // Restore 0600 permissions
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let perms = std::fs::Permissions::from_mode(0o600);
        std::fs::set_permissions(&path, perms)
            .map_err(|e| format!("Failed to set permissions on {}: {e}", path.display()))?;
    }

    Ok(())
}

// ── Multi-account pool ────────────────────────────────────────────────────────

/// A single account entry in the multi-account login pool.
///
/// Stored in the pool file, keyed by email address.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PoolAccount {
    /// Email address that identifies this account.
    pub email: String,
    /// OAuth access token.
    pub access_token: String,
    /// OAuth refresh token (empty string when not present).
    pub refresh_token: String,
    /// Unix-epoch milliseconds when the access token expires.
    pub expires_at: u64,
    /// Scopes granted by this token.
    #[serde(default)]
    pub scopes: Vec<String>,
    /// Claude subscription tier (e.g. "max"), if known.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub subscription_type: Option<String>,
    /// Rate-limit tier string, if known.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub rate_limit_tier: Option<String>,
    /// Whether the server has observed a rate-limit response for this account.
    #[serde(default)]
    pub rate_limited: bool,
    /// Unix-epoch milliseconds when the rate limit resets, if known.
    ///
    /// Used to compute "time until earliest reset" when all accounts are exhausted.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub rate_limit_reset_at: Option<u64>,
}

/// The multi-account login pool, stored in `~/.claude/oauth_pool.json`.
///
/// Accounts are keyed by email address so that repeated logins with the same
/// account update the existing entry rather than creating a duplicate.
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct AccountPool {
    /// Email → account credentials.
    pub accounts: HashMap<String, PoolAccount>,
}

/// Returns the path to `~/.claude/oauth_pool.json`.
fn pool_path() -> Result<PathBuf, String> {
    let home = std::env::var("HOME").map_err(|_| "HOME not set".to_string())?;
    Ok(PathBuf::from(home).join(".claude").join("oauth_pool.json"))
}

/// Read the account pool from disk.
///
/// Returns an empty pool when no file exists yet.
pub fn read_pool() -> Result<AccountPool, String> {
    let path = pool_path()?;
    if !path.exists() {
        return Ok(AccountPool::default());
    }
    let data = std::fs::read_to_string(&path)
        .map_err(|e| format!("Cannot read {}: {e}", path.display()))?;
    serde_json::from_str(&data).map_err(|e| format!("Failed to parse {}: {e}", path.display()))
}

/// Write the account pool to disk with 0600 permissions.
pub fn write_pool(pool: &AccountPool) -> Result<(), String> {
    let path = pool_path()?;
    let data =
        serde_json::to_string_pretty(pool).map_err(|e| format!("Failed to serialize pool: {e}"))?;
    std::fs::write(&path, &data).map_err(|e| format!("Failed to write {}: {e}", path.display()))?;

    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let perms = std::fs::Permissions::from_mode(0o600);
        std::fs::set_permissions(&path, perms)
            .map_err(|e| format!("Failed to set permissions on {}: {e}", path.display()))?;
    }

    Ok(())
}

/// Insert or update an account in the pool, then persist to disk.
///
/// If an account with the same email already exists its credentials are
/// overwritten; all other accounts are left unchanged.
pub fn upsert_pool_account(account: PoolAccount) -> Result<(), String> {
    let mut pool = read_pool()?;
    pool.accounts.insert(account.email.clone(), account);
    write_pool(&pool)
}

/// Perform a token refresh HTTP call using `refresh_token`.
///
/// Returns `(new_access_token, new_expires_at_ms)` on success.
async fn do_token_refresh(refresh_token: &str) -> Result<(String, u64), String> {
    let client = reqwest::Client::new();
    let resp = client
        .post(TOKEN_ENDPOINT)
        .json(&serde_json::json!({
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
            "client_id": CLAUDE_CODE_CLIENT_ID,
        }))
        .send()
        .await
        .map_err(|e| format!("OAuth refresh request failed: {e}"))?;

    let status = resp.status();
    let body = resp
        .text()
        .await
        .map_err(|e| format!("Failed to read refresh response: {e}"))?;

    if !status.is_success() {
        if let Ok(err) = serde_json::from_str::<TokenRefreshError>(&body) {
            let desc = err
                .error_description
                .unwrap_or_else(|| "unknown error".to_string());
            slog!("[oauth] Refresh failed: {desc} (full body: {body})");
            return Err(format!(
                "OAuth session expired. Please run `claude login` to re-authenticate. ({desc})"
            ));
        }
        return Err(format!(
            "OAuth session expired. Please run `claude login` to re-authenticate. (HTTP {status})"
        ));
    }

    let token_resp: TokenRefreshResponse = serde_json::from_str(&body)
        .map_err(|e| format!("Failed to parse refresh response: {e}"))?;

    let now_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0);

    let expires_at = now_ms + (token_resp.expires_in * 1000);
    Ok((token_resp.access_token, expires_at))
}

/// Refresh the OAuth access token using the stored refresh token.
///
/// On success, updates `~/.claude/.credentials.json` with the new access
/// token and expiry, then returns `Ok(())`.
///
/// On failure (e.g. refresh token expired), returns an error string.
pub async fn refresh_access_token() -> Result<(), String> {
    slog!("[oauth] Attempting to refresh OAuth access token");

    let mut creds = read_credentials()?;
    let refresh_token = creds.claude_ai_oauth.refresh_token.clone();

    if refresh_token.is_empty() {
        return Err("No refresh token found. Run `claude login` to authenticate.".to_string());
    }

    let (new_access_token, new_expires_at) = do_token_refresh(&refresh_token).await?;

    creds.claude_ai_oauth.access_token = new_access_token;
    creds.claude_ai_oauth.expires_at = new_expires_at;

    write_credentials(&creds)?;

    slog!(
        "[oauth] Successfully refreshed access token, expires at {}",
        creds.claude_ai_oauth.expires_at
    );

    Ok(())
}

// ── Account pool swap ─────────────────────────────────────────────────────────

/// Select the next non-rate-limited account for swapping, given the current
/// credentials' refresh token.
///
/// Marks the current account (identified by `current_refresh_token`) as
/// rate-limited with `reset_at_ms`, then returns a clone of the first
/// non-rate-limited account in the pool.
///
/// Returns:
/// - `Ok(account)` — the next available pool account to switch to
/// - `Err(msg)` — all accounts are rate-limited (includes earliest reset time
///   in the message), or the pool has fewer than 2 accounts
///
/// This is a pure function exposed for unit testing.  The async
/// [`swap_to_next_available_account`] calls this and then performs the HTTP
/// token refresh.
pub(crate) fn select_next_account_for_swap(
    pool: &mut AccountPool,
    current_refresh_token: &str,
    reset_at_ms: u64,
) -> Result<PoolAccount, String> {
    if pool.accounts.len() < 2 {
        return Err(
            "No account pool available for automatic account rotation (fewer than 2 accounts)."
                .to_string(),
        );
    }

    // Mark the current account as rate-limited.
    for account in pool.accounts.values_mut() {
        if account.refresh_token == current_refresh_token {
            account.rate_limited = true;
            account.rate_limit_reset_at = Some(reset_at_ms);
            break;
        }
    }

    // Find the next non-rate-limited account.
    if let Some(next) = pool.accounts.values().find(|a| !a.rate_limited).cloned() {
        return Ok(next);
    }

    // All accounts are rate-limited — compute the earliest reset time.
    let earliest_reset_ms = pool
        .accounts
        .values()
        .filter_map(|a| a.rate_limit_reset_at)
        .min();

    let msg = match earliest_reset_ms {
        Some(ms) => {
            let reset_secs = (ms / 1000) as i64;
            let reset_dt =
                chrono::DateTime::from_timestamp(reset_secs, 0).unwrap_or_else(chrono::Utc::now);
            let until = reset_dt.signed_duration_since(chrono::Utc::now());
            let mins = until.num_minutes().max(0);
            format!(
                "All OAuth accounts are rate-limited. \
                 Earliest reset in {} minute{} (at {} UTC).",
                mins,
                if mins == 1 { "" } else { "s" },
                reset_dt.format("%H:%M")
            )
        }
        None => "All OAuth accounts are rate-limited. No reset time available.".to_string(),
    };

    Err(msg)
}

/// Swap to the next available OAuth account in the pool on rate-limit.
///
/// Identifies the currently active account (via `~/.claude/.credentials.json`),
/// marks it as rate-limited, then refreshes and activates the next available
/// account by writing its credentials to `~/.claude/.credentials.json`.
///
/// Returns the email of the newly activated account on success, or an error
/// message when no account pool exists or all accounts are exhausted.
pub async fn swap_to_next_available_account(
    reset_at: chrono::DateTime<chrono::Utc>,
) -> Result<String, String> {
    slog!("[oauth] Rate limit hit — attempting account pool swap");

    let mut pool = read_pool()?;

    let reset_at_ms = reset_at.timestamp_millis() as u64;

    let creds = read_credentials().unwrap_or_else(|_| CredentialsFile {
        claude_ai_oauth: OAuthCredentials {
            access_token: String::new(),
            refresh_token: String::new(),
            expires_at: 0,
            scopes: vec![],
            subscription_type: None,
            rate_limit_tier: None,
        },
    });
    let current_refresh_token = creds.claude_ai_oauth.refresh_token.clone();

    let next_account =
        select_next_account_for_swap(&mut pool, &current_refresh_token, reset_at_ms)?;

    slog!(
        "[oauth] Swapping to account '{}' — refreshing token",
        next_account.email
    );

    // Refresh the next account's token via HTTP.
    let (new_access_token, new_expires_at) = do_token_refresh(&next_account.refresh_token).await?;

    // Update the pool account with the fresh token and clear rate-limited flag.
    if let Some(pool_account) = pool.accounts.get_mut(&next_account.email) {
        pool_account.access_token = new_access_token.clone();
        pool_account.expires_at = new_expires_at;
        pool_account.rate_limited = false;
        pool_account.rate_limit_reset_at = None;
    }

    // Persist the updated pool (rate-limited flag on old account, fresh token on new).
    write_pool(&pool)?;

    // Write the new account's credentials to .credentials.json so the next
    // Claude Code CLI invocation uses the freshly activated account.
    let new_creds = CredentialsFile {
        claude_ai_oauth: OAuthCredentials {
            access_token: new_access_token,
            refresh_token: next_account.refresh_token.clone(),
            expires_at: new_expires_at,
            scopes: next_account.scopes.clone(),
            subscription_type: next_account.subscription_type.clone(),
            rate_limit_tier: next_account.rate_limit_tier.clone(),
        },
    };
    write_credentials(&new_creds)?;

    slog!(
        "[oauth] Account swap complete — now using '{}'",
        next_account.email
    );

    Ok(next_account.email)
}

/// Extract the OAuth login URL from an error message produced by the Claude Code provider.
///
/// The provider returns errors like:
/// `"OAuth session expired or credentials missing. Please log in: http://localhost:3001/oauth/authorize"`
///
/// Returns the URL portion when the error indicates missing or expired credentials,
/// `None` otherwise.
pub fn extract_login_url_from_error(err: &str) -> Option<&str> {
    let marker = "Please log in: ";
    let start = err.find(marker)?;
    Some(err[start + marker.len()..].trim())
}

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

    #[test]
    fn extract_login_url_from_oauth_error() {
        let err = "OAuth session expired or credentials missing. Please log in: http://localhost:3001/oauth/authorize";
        let url = extract_login_url_from_error(err);
        assert_eq!(url, Some("http://localhost:3001/oauth/authorize"));
    }

    #[test]
    fn extract_login_url_returns_none_for_unrelated_error() {
        let err = "Some other error occurred";
        assert!(extract_login_url_from_error(err).is_none());
    }

    #[test]
    fn extract_login_url_with_different_port() {
        let err = "OAuth session expired or credentials missing. Please log in: http://localhost:3002/oauth/authorize";
        let url = extract_login_url_from_error(err);
        assert_eq!(url, Some("http://localhost:3002/oauth/authorize"));
    }

    #[test]
    fn parse_credentials_file() {
        let json = r#"{
            "claudeAiOauth": {
                "accessToken": "sk-ant-oat01-test",
                "refreshToken": "sk-ant-ort01-test",
                "expiresAt": 1774466144677,
                "scopes": ["user:inference"],
                "subscriptionType": "max",
                "rateLimitTier": "default_claude_max_20x"
            }
        }"#;
        let creds: CredentialsFile = serde_json::from_str(json).unwrap();
        assert_eq!(creds.claude_ai_oauth.access_token, "sk-ant-oat01-test");
        assert_eq!(creds.claude_ai_oauth.refresh_token, "sk-ant-ort01-test");
        assert_eq!(creds.claude_ai_oauth.expires_at, 1774466144677);
        assert_eq!(
            creds.claude_ai_oauth.subscription_type.as_deref(),
            Some("max")
        );
    }

    #[test]
    fn serialize_credentials_roundtrip() {
        let creds = CredentialsFile {
            claude_ai_oauth: OAuthCredentials {
                access_token: "access".to_string(),
                refresh_token: "refresh".to_string(),
                expires_at: 12345,
                scopes: vec!["user:inference".to_string()],
                subscription_type: Some("max".to_string()),
                rate_limit_tier: None,
            },
        };
        let json = serde_json::to_string(&creds).unwrap();
        let parsed: CredentialsFile = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.claude_ai_oauth.access_token, "access");
        assert_eq!(parsed.claude_ai_oauth.refresh_token, "refresh");
        // rate_limit_tier should be omitted from JSON (skip_serializing_if)
        assert!(!json.contains("rateLimitTier"));
    }

    #[test]
    fn credentials_path_uses_home() {
        // Just verify it doesn't panic and returns a path ending in .credentials.json
        if std::env::var("HOME").is_ok() {
            let path = credentials_path().unwrap();
            assert!(path.ends_with(".claude/.credentials.json"));
        }
    }

    // ── select_next_account_for_swap ─────────────────────────────────────────

    fn make_account(email: &str, refresh_token: &str, rate_limited: bool) -> PoolAccount {
        PoolAccount {
            email: email.to_string(),
            access_token: format!("access-{email}"),
            refresh_token: refresh_token.to_string(),
            expires_at: 9_999_999_999_999,
            scopes: vec![],
            subscription_type: None,
            rate_limit_tier: None,
            rate_limited,
            rate_limit_reset_at: None,
        }
    }

    fn make_pool(accounts: Vec<PoolAccount>) -> AccountPool {
        let mut pool = AccountPool::default();
        for a in accounts {
            pool.accounts.insert(a.email.clone(), a);
        }
        pool
    }

    /// AC3a: single rate-limit — current account is marked limited, swap to next.
    #[test]
    fn swap_selects_next_account_when_current_is_rate_limited() {
        let account_a = make_account("a@example.com", "refresh-a", false);
        let account_b = make_account("b@example.com", "refresh-b", false);
        let mut pool = make_pool(vec![account_a, account_b]);

        let reset_ms = 1_700_000_000_000u64;
        let result = select_next_account_for_swap(&mut pool, "refresh-a", reset_ms);

        assert!(result.is_ok(), "should find account B: {result:?}");
        let next = result.unwrap();
        assert_eq!(next.email, "b@example.com");

        // Account A should now be marked rate-limited.
        let a = pool.accounts.get("a@example.com").unwrap();
        assert!(a.rate_limited, "account A should be marked rate-limited");
        assert_eq!(a.rate_limit_reset_at, Some(reset_ms));
    }

    /// AC3b: cascade swap — next account is already rate-limited, picks the one after.
    #[test]
    fn swap_skips_already_rate_limited_account_and_picks_next_available() {
        let account_a = make_account("a@example.com", "refresh-a", false);
        let mut account_b = make_account("b@example.com", "refresh-b", false);
        account_b.rate_limited = true; // B is already rate-limited
        account_b.rate_limit_reset_at = Some(1_700_000_000_000);
        let account_c = make_account("c@example.com", "refresh-c", false);

        let mut pool = make_pool(vec![account_a, account_b, account_c]);

        let reset_ms = 1_700_000_001_000u64;
        let result = select_next_account_for_swap(&mut pool, "refresh-a", reset_ms);

        assert!(result.is_ok(), "should find account C: {result:?}");
        let next = result.unwrap();
        assert_eq!(next.email, "c@example.com");
    }

    /// AC3c: all-exhausted error path — error message includes reset time.
    #[test]
    fn swap_returns_error_with_reset_time_when_all_accounts_exhausted() {
        let account_a = make_account("a@example.com", "refresh-a", false);
        let mut account_b = make_account("b@example.com", "refresh-b", false);
        // B is already rate-limited with a known reset time.
        account_b.rate_limited = true;
        account_b.rate_limit_reset_at =
            Some((chrono::Utc::now() + chrono::Duration::hours(2)).timestamp_millis() as u64);

        let mut pool = make_pool(vec![account_a, account_b]);

        // A also hits rate limit now → all exhausted.
        let reset_ms = (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp_millis() as u64;
        let result = select_next_account_for_swap(&mut pool, "refresh-a", reset_ms);

        assert!(result.is_err(), "should return exhausted error");
        let msg = result.unwrap_err();
        assert!(
            msg.contains("All OAuth accounts are rate-limited"),
            "error should say all exhausted: {msg}"
        );
        assert!(
            msg.contains("minute"),
            "error should include time until reset: {msg}"
        );
    }

    /// Pool with fewer than 2 accounts returns no-pool error.
    #[test]
    fn swap_returns_error_when_pool_has_fewer_than_two_accounts() {
        let mut pool = make_pool(vec![make_account("a@example.com", "refresh-a", false)]);
        let result = select_next_account_for_swap(&mut pool, "refresh-a", 0);
        assert!(result.is_err());
        let msg = result.unwrap_err();
        assert!(
            msg.contains("fewer than 2 accounts"),
            "should explain pool too small: {msg}"
        );
    }

    /// Empty pool also returns no-pool error.
    #[test]
    fn swap_returns_error_for_empty_pool() {
        let mut pool = AccountPool::default();
        let result = select_next_account_for_swap(&mut pool, "any-token", 0);
        assert!(result.is_err());
    }
}