From 1ecb4dad55ddb6e209b0e78512bf9a5f144eb818 Mon Sep 17 00:00:00 2001 From: dave Date: Mon, 27 Apr 2026 17:36:00 +0000 Subject: [PATCH] huskies: merge 724_story_per_account_oauth_credential_storage_with_login_pool --- server/src/http/oauth.rs | 15 ++-- server/src/llm/oauth.rs | 88 ++++++++++++++++++++++ server/src/service/oauth/flow.rs | 50 ++++++++++++- server/src/service/oauth/io.rs | 123 ++++++++++++++++++++++++++----- server/src/service/oauth/mod.rs | 22 ++++-- 5 files changed, 265 insertions(+), 33 deletions(-) diff --git a/server/src/http/oauth.rs b/server/src/http/oauth.rs index 90c45ad7..f1869af5 100644 --- a/server/src/http/oauth.rs +++ b/server/src/http/oauth.rs @@ -93,16 +93,15 @@ pub async fn oauth_callback( } } -/// `GET /oauth/status` — Check whether valid (non-expired) OAuth credentials exist. +/// `GET /oauth/status` — Return status for all stored OAuth accounts in the login pool. +/// +/// Each entry includes the account email, status (`active`, `expired`, or +/// `rate-limited`), expiry timestamp, and refresh-token presence flag. +/// Returns an empty `accounts` array when no accounts are stored. #[handler] pub async fn oauth_status() -> poem::Response { - let status = svc::check_status(); - let body = serde_json::json!({ - "authenticated": status.authenticated, - "expired": status.expired, - "expires_at": status.expires_at, - "has_refresh_token": status.has_refresh_token, - }); + let accounts = svc::check_all_accounts(); + let body = serde_json::json!({ "accounts": accounts }); poem::Response::builder() .status(StatusCode::OK) .header("Content-Type", "application/json") diff --git a/server/src/llm/oauth.rs b/server/src/llm/oauth.rs index 261ff760..38616ee7 100644 --- a/server/src/llm/oauth.rs +++ b/server/src/llm/oauth.rs @@ -1,6 +1,7 @@ //! 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. @@ -85,6 +86,93 @@ pub fn write_credentials(creds: &CredentialsFile) -> Result<(), String> { 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, + /// Claude subscription tier (e.g. "max"), if known. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subscription_type: Option, + /// Rate-limit tier string, if known. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rate_limit_tier: Option, + /// Whether the server has observed a rate-limit response for this account. + #[serde(default)] + pub rate_limited: bool, +} + +/// 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, +} + +/// Returns the path to `~/.claude/oauth_pool.json`. +fn pool_path() -> Result { + 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 { + 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) +} + /// Refresh the OAuth access token using the stored refresh token. /// /// On success, updates `~/.claude/.credentials.json` with the new access diff --git a/server/src/service/oauth/flow.rs b/server/src/service/oauth/flow.rs index be127570..95f06927 100644 --- a/server/src/service/oauth/flow.rs +++ b/server/src/service/oauth/flow.rs @@ -3,7 +3,53 @@ //! All functions here are pure — no I/O, no network, no clocks. //! Side-effectful operations live exclusively in `io.rs`. -use crate::llm::oauth::CredentialsFile; +use crate::llm::oauth::{CredentialsFile, PoolAccount}; + +// ── Account pool status types ───────────────────────────────────────────────── + +/// Status of a single account in the login pool as returned by `GET /oauth/status`. +#[derive(Debug, Clone, serde::Serialize)] +pub struct AccountInfo { + /// Email address that identifies the account. + pub email: String, + /// Human-readable account status. + pub status: AccountStatus, + /// Unix-epoch millisecond expiry timestamp. + pub expires_at: u64, + /// Whether a non-empty refresh token is present. + pub has_refresh_token: bool, +} + +/// The status of a stored OAuth account. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum AccountStatus { + /// Token is valid and not rate-limited. + Active, + /// Token expiry is in the past. + Expired, + /// A rate-limit response was observed for this account. + RateLimited, +} + +/// Build an [`AccountInfo`] from a [`PoolAccount`] and the current time. +pub fn build_account_info(account: &PoolAccount, now_ms: u64) -> AccountInfo { + let status = if account.rate_limited { + AccountStatus::RateLimited + } else if is_token_expired(account.expires_at, now_ms) { + AccountStatus::Expired + } else { + AccountStatus::Active + }; + AccountInfo { + email: account.email.clone(), + status, + expires_at: account.expires_at, + has_refresh_token: !account.refresh_token.is_empty(), + } +} + +// ── PKCE flow ───────────────────────────────────────────────────────────────── /// A pending PKCE flow waiting for an OAuth callback. pub struct PendingFlow { @@ -17,6 +63,7 @@ pub struct PendingFlow { #[derive(Debug, Clone)] pub struct FlowStatus { /// Whether valid credentials were found on disk. + #[allow(dead_code)] pub authenticated: bool, /// Whether the access token is past its expiry timestamp. pub expired: bool, @@ -45,6 +92,7 @@ pub fn build_flow_status(creds: &CredentialsFile, now_ms: u64) -> FlowStatus { } /// Return the unauthenticated `FlowStatus` (no credentials on disk). +#[allow(dead_code)] pub fn unauthenticated_status() -> FlowStatus { FlowStatus { authenticated: false, diff --git a/server/src/service/oauth/io.rs b/server/src/service/oauth/io.rs index 2e9bb085..9e44b651 100644 --- a/server/src/service/oauth/io.rs +++ b/server/src/service/oauth/io.rs @@ -5,9 +5,9 @@ //! All business logic and branching belong in `mod.rs`, `pkce.rs`, or `flow.rs`. use super::Error; -use super::flow::FlowStatus; +use super::flow::{AccountInfo, build_account_info}; use super::pkce::SCOPES; -use crate::llm::oauth::{self, CredentialsFile}; +use crate::llm::oauth::{self, CredentialsFile, PoolAccount}; use crate::slog; /// Raw token exchange result returned by the Anthropic OAuth endpoint. @@ -80,33 +80,122 @@ pub(super) async fn exchange_code_for_tokens( .map_err(|e| Error::Parse(format!("Unexpected response from Anthropic: {e}"))) } -/// Persist a token exchange result to `~/.claude/.credentials.json`. +/// Persist a token exchange result to `~/.claude/.credentials.json` and to the +/// multi-account pool (`~/.claude/oauth_pool.json`). /// -/// Builds a [`CredentialsFile`] from the token response and `now_ms`, then -/// delegates to [`oauth::write_credentials`]. -pub(super) fn save_credentials(token: &TokenExchangeResult, now_ms: u64) -> Result<(), Error> { +/// Writes the credentials file for Claude Code CLI compatibility, then upserts +/// the account into the pool keyed by `email`. If `email` is empty a fallback +/// key is derived from the first 16 characters of the access token. +pub(super) fn save_credentials( + token: &TokenExchangeResult, + now_ms: u64, + email: &str, +) -> Result<(), Error> { + let expires_at = now_ms + (token.expires_in * 1000); + let scopes: Vec = SCOPES.split(' ').map(|s| s.to_string()).collect(); + + // Write to .credentials.json for Claude Code CLI compatibility. let creds = CredentialsFile { claude_ai_oauth: oauth::OAuthCredentials { access_token: token.access_token.clone(), refresh_token: token.refresh_token.clone().unwrap_or_default(), - expires_at: now_ms + (token.expires_in * 1000), - scopes: SCOPES.split(' ').map(|s| s.to_string()).collect(), + expires_at, + scopes: scopes.clone(), subscription_type: None, rate_limit_tier: None, }, }; - oauth::write_credentials(&creds).map_err(Error::TokenStorage) + oauth::write_credentials(&creds).map_err(Error::TokenStorage)?; + + // Build the pool key: prefer email, fall back to token prefix. + let key = if email.is_empty() { + format!( + "account-{}", + &token.access_token[..token.access_token.len().min(16)] + ) + } else { + email.to_string() + }; + + let account = PoolAccount { + email: key, + access_token: token.access_token.clone(), + refresh_token: token.refresh_token.clone().unwrap_or_default(), + expires_at, + scopes, + subscription_type: None, + rate_limit_tier: None, + rate_limited: false, + }; + oauth::upsert_pool_account(account).map_err(Error::TokenStorage)?; + + Ok(()) } -/// Load OAuth credentials from disk and compute a [`FlowStatus`]. +/// Attempt to fetch the authenticated user's email from Anthropic's userinfo endpoint. /// -/// Returns `Ok(None)` when no credentials file exists yet (user not logged in). -pub(super) fn load_status() -> FlowStatus { - match oauth::read_credentials() { - Ok(creds) => { - let now_ms = current_time_ms(); - super::flow::build_flow_status(&creds, now_ms) +/// Uses the access token as a Bearer credential. Returns `None` when the +/// request fails or the response does not contain an email field. +pub(super) async fn fetch_user_email(access_token: &str) -> Option { + const USERINFO_URL: &str = "https://claude.ai/api/me"; + + #[derive(serde::Deserialize)] + struct UserInfo { + email: Option, + } + + let client = reqwest::Client::new(); + let resp = client + .get(USERINFO_URL) + .header("Authorization", format!("Bearer {access_token}")) + .send() + .await + .ok()?; + + if !resp.status().is_success() { + return None; + } + + let info: UserInfo = resp.json().await.ok()?; + info.email.filter(|e| !e.is_empty()) +} + +/// Load all accounts from the pool and compute their [`AccountInfo`]. +/// +/// Falls back to reading the legacy `.credentials.json` when the pool is empty +/// or missing, synthesising a single unauthenticated-or-authenticated entry so +/// existing deployments that have not yet run a new login continue to work. +pub(super) fn load_all_accounts() -> Vec { + let now_ms = current_time_ms(); + + match oauth::read_pool() { + Ok(pool) if !pool.accounts.is_empty() => { + let mut accounts: Vec = pool + .accounts + .values() + .map(|a| build_account_info(a, now_ms)) + .collect(); + accounts.sort_by(|a, b| a.email.cmp(&b.email)); + accounts + } + _ => { + // Pool empty or unreadable — fall back to legacy single-account file. + match oauth::read_credentials() { + Ok(creds) => { + let status = super::flow::build_flow_status(&creds, now_ms); + vec![AccountInfo { + email: "default".to_string(), + status: if status.expired { + super::flow::AccountStatus::Expired + } else { + super::flow::AccountStatus::Active + }, + expires_at: status.expires_at, + has_refresh_token: status.has_refresh_token, + }] + } + Err(_) => vec![], + } } - Err(_) => super::flow::unauthenticated_status(), } } diff --git a/server/src/service/oauth/mod.rs b/server/src/service/oauth/mod.rs index 0ec1b2c4..dba52c01 100644 --- a/server/src/service/oauth/mod.rs +++ b/server/src/service/oauth/mod.rs @@ -11,7 +11,7 @@ pub mod flow; pub(super) mod io; pub mod pkce; -pub use flow::FlowStatus; +pub use flow::AccountInfo; use flow::PendingFlow; use std::collections::HashMap; @@ -154,18 +154,26 @@ pub async fn exchange_code(state: &OAuthState, code: &str, csrf_state: &str) -> io::exchange_code_for_tokens(code, &flow.redirect_uri, &flow.code_verifier, csrf_state) .await?; let now_ms = io::current_time_ms(); - io::save_credentials(&token, now_ms)?; + + // Attempt to resolve the email for this account; silently fall back to an + // empty string so that credential storage always succeeds. + let email = io::fetch_user_email(&token.access_token) + .await + .unwrap_or_default(); + + io::save_credentials(&token, now_ms, &email)?; crate::slog!("[oauth] Successfully authenticated and saved credentials"); Ok(()) } -/// Return the current OAuth credential status without performing any I/O beyond -/// reading the credentials file. +/// Return status information for every account in the login pool. /// -/// Returns an unauthenticated [`FlowStatus`] when no credentials file exists. -pub fn check_status() -> FlowStatus { - io::load_status() +/// If no pool exists yet, falls back to the legacy single-account credentials +/// file so that existing deployments continue to work. Returns an empty `Vec` +/// when neither the pool nor the legacy file is present. +pub fn check_all_accounts() -> Vec { + io::load_all_accounts() } // ── Tests ─────────────────────────────────────────────────────────────────────