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::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(())
}

/// 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 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() {
        // Try to parse a structured error
        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);

    creds.claude_ai_oauth.access_token = token_resp.access_token;
    creds.claude_ai_oauth.expires_at = now_ms + (token_resp.expires_in * 1000);

    write_credentials(&creds)?;

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

    Ok(())
}

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