feat(ai): track login age, warn before a pool login ends, okena ai accounts

crates/okena-agent-usage/src/pool.rs

@@ -664,6 +664,8 @@ mod tests {
             seven_day: None,
             requests: 0,
             last_used_ms: None,
+            logged_in_ms: None,
+            refreshed_ms: None,
         }
     }
 

crates/okena-agent-usage/src/pool_summary.rs

@@ -190,7 +190,7 @@ fn blocked(account: &AiAccount, now_ms: i64) -> Option<Blocked> {
         })
     };
     if account.needs_login {
-        return untimed("Needs login");
+        return untimed("Login rejected");
     }
     if !account.enabled {
         return untimed("Disabled");
@@ -330,6 +330,8 @@ mod tests {
             seven_day: None,
             requests: 0,
             last_used_ms: None,
+            logged_in_ms: None,
+            refreshed_ms: None,
         }
     }
 
@@ -407,7 +409,7 @@ mod tests {
             [
                 ("a", None),
                 ("b", Some("Disabled")),
-                ("c", Some("Needs login"))
+                ("c", Some("Login rejected"))
             ]
         );
         assert_eq!(summary.accounts[0].tier.as_deref(), Some("Max 5x"));

crates/okena-ai/src/accounts.rs

@@ -6,7 +6,7 @@ use crate::oauth::{self, Pkce, RefreshError, Tokens, Upstream};
 use okena_agent_usage::history::{self as usage_history, LimitSample, Provider, UsageSample};
 use okena_core::ai_proxy::{
     AiAccount, AiAccountHistory, AiAccountRule, AiActivity, AiPendingLogin, AiProxyStatus,
-    AiRouting, AiRoutingPolicy, AiUsage, is_routable_project_id, write_private,
+    AiRouting, AiRoutingPolicy, AiUsage, is_routable_project_id, login_age_days, write_private,
 };
 use parking_lot::Mutex;
 use reqwest::header::HeaderMap;
@@ -34,6 +34,18 @@ struct StoredAccount {
     enabled: bool,
     needs_login: bool,
     tokens: Tokens,
+    /// The interactive login; a refresh keeps it, a new login resets it.
+    #[serde(default)]
+    logged_in_ms: Option<i64>,
+    #[serde(default)]
+    refreshed_ms: Option<i64>,
+}
+
+impl StoredAccount {
+    fn login_age_days(&self, now: i64) -> Option<i64> {
+        self.logged_in_ms
+            .map(|logged_in| login_age_days(logged_in, now))
+    }
 }
 
 #[derive(Default, Serialize, Deserialize)]
@@ -253,6 +265,8 @@ impl Accounts {
                     seven_day: account.runtime.seven_day.clone(),
                     requests: account.runtime.requests,
                     last_used_ms: account.runtime.last_used_ms,
+                    logged_in_ms: account.stored.logged_in_ms,
+                    refreshed_ms: account.stored.refreshed_ms,
                 })
                 .collect(),
             logins: state
@@ -346,6 +360,7 @@ impl Accounts {
             None => None,
         };
         let mut state = self.state.lock();
+        let now = now_ms();
         state.logins.retain(|login| login.id != login_id);
         // Logging the same provider account in again replaces its credentials.
         let existing = result.account_uuid.as_ref().and_then(|uuid| {
@@ -357,6 +372,12 @@ impl Accounts {
         if let Some(account) = existing {
             account.stored.tokens = result.tokens;
             account.stored.needs_login = false;
+            account.stored.logged_in_ms = Some(now);
+            account.stored.refreshed_ms = None;
+            log::info!(
+                "[ai-proxy] login account={:?} (logged in again)",
+                account.stored.id
+            );
             account.stored.email = result.email.or(account.stored.email.take());
             if let Some(plan) = plan {
                 account.stored.plan = Some(plan.kind);
@@ -365,9 +386,11 @@ impl Accounts {
             account.runtime.plan_checked = true;
             account.runtime.cooldown_until_ms = None;
         } else {
+            let id = okena_core::ai_proxy::random_token()?;
+            log::info!("[ai-proxy] login account={id:?} (new account)");
             state.accounts.push(Account {
                 stored: StoredAccount {
-                    id: okena_core::ai_proxy::random_token()?,
+                    id,
                     label,
                     email: result.email,
                     account_uuid: result.account_uuid,
@@ -376,6 +399,8 @@ impl Accounts {
                     enabled: true,
                     needs_login: false,
                     tokens: result.tokens,
+                    logged_in_ms: Some(now),
+                    refreshed_ms: None,
                 },
                 runtime: Runtime {
                     plan_checked: true,
@@ -588,8 +613,14 @@ impl Accounts {
         match oauth::refresh(&self.http, &self.upstream, tokens).await {
             Ok(refreshed) => {
                 let mut state = self.state.lock();
+                let now = now_ms();
                 let account = state.account(account_id)?;
                 account.stored.tokens = refreshed.clone();
+                account.stored.refreshed_ms = Some(now);
+                log::info!(
+                    "[ai-proxy] token_refresh ok account={account_id:?} login_age_days={:?}",
+                    account.stored.login_age_days(now)
+                );
                 // Keep serving from memory even if the write fails; the error is logged.
                 if let Err(error) = self.persist(&state) {
                     log::error!("[ai-proxy] failed to persist refreshed credentials: {error}");
@@ -602,7 +633,12 @@ impl Accounts {
             }
             Err(RefreshError::Rejected) => {
                 let mut state = self.state.lock();
-                state.account(account_id)?.stored.needs_login = true;
+                let account = state.account(account_id)?;
+                account.stored.needs_login = true;
+                log::warn!(
+                    "[ai-proxy] token_refresh rejected account={account_id:?} login_age_days={:?}; the account needs a new login",
+                    account.stored.login_age_days(now_ms())
+                );
                 if let Err(error) = self.persist(&state) {
                     log::error!("[ai-proxy] failed to persist account state: {error}");
                 }
@@ -631,6 +667,10 @@ impl Accounts {
         let mut state = self.state.lock();
         if let Ok(account) = state.account(account_id) {
             account.stored.needs_login = true;
+            log::warn!(
+                "[ai-proxy] token rejected after refresh account={account_id:?} login_age_days={:?}; the account needs a new login",
+                account.stored.login_age_days(now_ms())
+            );
             if let Err(error) = self.persist(&state) {
                 log::error!("[ai-proxy] failed to persist account state: {error}");
             }
@@ -958,6 +998,8 @@ mod tests {
                         refresh_token: format!("refresh-{id}"),
                         expires_at_ms: i64::MAX / 2,
                     },
+                    logged_in_ms: None,
+                    refreshed_ms: None,
                 })
                 .collect(),
         };

crates/okena-ai/src/oauth.rs

@@ -5,6 +5,7 @@ use base64::engine::general_purpose::URL_SAFE_NO_PAD;
 use okena_core::ai_proxy::AiUsage;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 
 pub const CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
@@ -101,6 +102,44 @@ struct TokenAccount {
     email_address: Option<String>,
 }
 
+/// The token reply states only the access token's lifetime. Once per process
+/// and grant, log which fields it carries (names, never values), to learn
+/// whether the provider ever states how long the login itself lasts.
+static CODE_REPLY_FIELDS_LOGGED: AtomicBool = AtomicBool::new(false);
+static REFRESH_REPLY_FIELDS_LOGGED: AtomicBool = AtomicBool::new(false);
+
+fn token_reply(
+    body: serde_json::Value,
+    grant: &str,
+    fields_logged: &AtomicBool,
+) -> anyhow::Result<TokenReply> {
+    if let Some(fields) = fields_to_log(&body, fields_logged) {
+        log::info!("[ai-proxy] token_reply_fields grant={grant} fields={fields}");
+    }
+    Ok(serde_json::from_value(body)?)
+}
+
+/// The reply's top-level field names, the first time only. Nested objects are
+/// skipped: one keyed by ids would put the ids into the log.
+fn fields_to_log(body: &serde_json::Value, fields_logged: &AtomicBool) -> Option<String> {
+    if fields_logged.swap(true, Ordering::Relaxed) {
+        return None;
+    }
+    // Sorted: serde_json keeps insertion order when a crate enables `preserve_order`.
+    let names: std::collections::BTreeSet<String> = body
+        .as_object()
+        .into_iter()
+        .flat_map(|object| object.keys())
+        .map(|key| {
+            key.chars()
+                .filter(|c| c.is_ascii_alphanumeric() || *c == '_')
+                .take(64)
+                .collect()
+        })
+        .collect();
+    Some(names.into_iter().collect::<Vec<_>>().join(","))
+}
+
 pub struct LoginResult {
     pub tokens: Tokens,
     pub account_uuid: Option<String>,
@@ -137,7 +176,11 @@ pub async fn exchange_code(
         status.is_success(),
         "Token exchange failed ({status}); the code may be expired or already used"
     );
-    let reply: TokenReply = response.json().await?;
+    let reply = token_reply(
+        response.json().await?,
+        "authorization_code",
+        &CODE_REPLY_FIELDS_LOGGED,
+    )?;
     let account = reply.account.as_ref();
     Ok(LoginResult {
         account_uuid: account.and_then(|account| account.uuid.clone()),
@@ -183,10 +226,12 @@ pub async fn refresh(
             "token refresh returned {status}"
         )));
     }
-    let reply: TokenReply = response
+    let body = response
         .json()
         .await
         .map_err(|error| RefreshError::Transient(error.into()))?;
+    let reply = token_reply(body, "refresh_token", &REFRESH_REPLY_FIELDS_LOGGED)
+        .map_err(RefreshError::Transient)?;
     tokens_from(
         reply.access_token,
         reply.refresh_token,
@@ -397,6 +442,21 @@ mod tests {
         assert_eq!(limits.seven_day.unwrap().utilization, 3.0);
     }
 
+    #[test]
+    fn reply_fields_are_logged_by_name_only_and_once() {
+        let logged = AtomicBool::new(false);
+        let reply = serde_json::json!({
+            "access_token": "secret-access",
+            "expires_in": 3600,
+            "account": {"uuid": "secret-uuid", "email_address": "someone@example.com"},
+            "bad\nkey": 1
+        });
+        let names = fields_to_log(&reply, &logged).unwrap();
+        assert_eq!(names, "access_token,account,badkey,expires_in");
+        assert!(!names.contains("secret") && !names.contains('@'));
+        assert_eq!(fields_to_log(&reply, &logged), None);
+    }
+
     #[test]
     fn only_a_dead_grant_rejects_the_refresh_token() {
         let rejected = |status: u16, body: &str| {

crates/okena-ai/src/server.rs

@@ -350,9 +350,10 @@ mod tests {
     use super::*;
     use axum::http::HeaderMap;
     use futures::StreamExt;
-    use okena_core::ai_proxy::AiActivityKind;
+    use okena_core::ai_proxy::{AiActivityKind, AiLoginState};
     use parking_lot::Mutex;
     use std::collections::HashMap;
+    use std::sync::atomic::AtomicBool;
 
     static LOG_MESSAGES: Mutex<Vec<String>> = Mutex::new(Vec::new());
     static TEST_LOGGER: TestLogger = TestLogger;
@@ -361,7 +362,7 @@ mod tests {
         static INIT: std::sync::Once = std::sync::Once::new();
         INIT.call_once(|| {
             log::set_logger(&TEST_LOGGER).unwrap();
-            log::set_max_level(log::LevelFilter::Warn);
+            log::set_max_level(log::LevelFilter::Info);
         });
     }
 
@@ -391,6 +392,8 @@ mod tests {
         error_bodies: Mutex<HashMap<String, axum::body::Bytes>>,
         error_headers: Mutex<HeaderMap>,
         token_delay: Mutex<Duration>,
+        /// Answer refreshes like a provider that ended the login.
+        reject_refresh: AtomicBool,
     }
 
     async fn fake_upstream(fake: Arc<Fake>) -> String {
@@ -446,6 +449,10 @@ mod tests {
             let delay = *fake.token_delay.lock();
             tokio::time::sleep(delay).await;
             fake.refreshes.fetch_add(1, Ordering::SeqCst);
+            if fake.reject_refresh.load(Ordering::SeqCst) {
+                let error = serde_json::json!({"error": "invalid_grant"});
+                return (StatusCode::BAD_REQUEST, axum::Json(error)).into_response();
+            }
             axum::Json(serde_json::json!({
                 "access_token": format!("{}-refreshed", body["refresh_token"].as_str().unwrap()),
                 "expires_in": 3600,
@@ -828,6 +835,84 @@ mod tests {
         assert_eq!(stored["accounts"].as_array().unwrap().len(), 1);
     }
 
+    #[tokio::test]
+    async fn login_time_survives_refreshes_until_a_new_login() {
+        init_test_logger();
+        let mut logged_in_long_ago = account("old", "token-old", "uuid-a");
+        logged_in_long_ago["logged_in_ms"] = 1_000.into();
+        let h = harness(vec![logged_in_long_ago]).await;
+        let status = h.client.command(&AiProxyCommand::Status).await.unwrap();
+        let account = &status.accounts[0];
+        assert_eq!(account.logged_in_ms, Some(1_000));
+        assert_eq!(account.login_state(crate::now_ms()), AiLoginState::Expiring);
+
+        h.fake.replies.lock().insert("token-old".into(), vec![401]);
+        assert_eq!(h.message("conv").await.0, 200);
+        let status = h.client.command(&AiProxyCommand::Status).await.unwrap();
+        let account = &status.accounts[0];
+        assert_eq!(account.logged_in_ms, Some(1_000), "a refresh keeps it");
+        let refreshed = account.refreshed_ms.expect("the refresh is recorded");
+        let stored: serde_json::Value = serde_json::from_slice(
+            &std::fs::read(proxy_dir(h._dir.path()).join("accounts.json")).unwrap(),
+        )
+        .unwrap();
+        assert_eq!(stored["accounts"][0]["logged_in_ms"], 1_000);
+        assert_eq!(stored["accounts"][0]["refreshed_ms"], refreshed);
+
+        h.fake.reject_refresh.store(true, Ordering::SeqCst);
+        h.fake
+            .replies
+            .lock()
+            .insert("old-rt-refreshed".into(), vec![401]);
+        h.message("conv").await;
+        let status = h.client.command(&AiProxyCommand::Status).await.unwrap();
+        assert_eq!(
+            status.accounts[0].login_state(crate::now_ms()),
+            AiLoginState::Rejected
+        );
+
+        let started = h
+            .client
+            .command(&AiProxyCommand::StartLogin {
+                label: "old".into(),
+            })
+            .await
+            .unwrap();
+        let finish = AiProxyCommand::FinishLogin {
+            login_id: started.logins[0].id.clone(),
+            code: "code".into(),
+        };
+        let status = h.client.command(&finish).await.unwrap();
+        let account = &status.accounts[0];
+        assert!(!account.needs_login);
+        assert!(
+            account.logged_in_ms.is_some_and(|at| at > 1_000),
+            "a new login resets it"
+        );
+        assert_eq!(account.refreshed_ms, None, "a new login starts over");
+
+        let logs: Vec<String> = LOG_MESSAGES
+            .lock()
+            .iter()
+            .filter(|line| line.contains("account=\"old\""))
+            .cloned()
+            .collect();
+        let logs = logs.join("\n");
+        for expected in [
+            "token_refresh ok account=\"old\" login_age_days=Some(",
+            "token_refresh rejected account=\"old\" login_age_days=Some(",
+            "login account=\"old\" (logged in again)",
+        ] {
+            assert!(logs.contains(expected), "{expected} missing in {logs}");
+        }
+        for token in ["token-old", "old-rt", "login-access", "login-rt"] {
+            assert!(
+                !logs.contains(token),
+                "token value {token} in the log: {logs}"
+            );
+        }
+    }
+
     #[tokio::test]
     async fn all_accounts_limited_returns_the_provider_429() {
         let h = harness(vec![account("a", "token-a", "uuid-a")]).await;