huskies: merge 724_story_per_account_oauth_credential_storage_with_login_pool

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<String> = 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<String> {
+    const USERINFO_URL: &str = "https://claude.ai/api/me";
+
+    #[derive(serde::Deserialize)]
+    struct UserInfo {
+        email: Option<String>,
+    }
+
+    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<AccountInfo> {
+    let now_ms = current_time_ms();
+
+    match oauth::read_pool() {
+        Ok(pool) if !pool.accounts.is_empty() => {
+            let mut accounts: Vec<AccountInfo> = 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(),
     }
 }