diff --git a/backend/src/upstream/accounts/store/queries.rs b/backend/src/upstream/accounts/store/queries.rs index 6569fb94e..1a652896e 100644 --- a/backend/src/upstream/accounts/store/queries.rs +++ b/backend/src/upstream/accounts/store/queries.rs @@ -502,10 +502,10 @@ where id = ?"; pub(super) const UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL: &str = r" update accounts set - email = ?, - chatgpt_account_id = ?, - chatgpt_user_id = ?, - plan_type = ?, + email = coalesce(?, email), + chatgpt_account_id = coalesce(?, chatgpt_account_id), + chatgpt_user_id = coalesce(?, chatgpt_user_id), + plan_type = coalesce(?, plan_type), access_token = ?, refresh_token = ?, access_token_expires_at = ?, @@ -521,10 +521,10 @@ where id = ?"; pub(super) const UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL: &str = r" update accounts set - email = ?, - chatgpt_account_id = ?, - chatgpt_user_id = ?, - plan_type = ?, + email = coalesce(?, email), + chatgpt_account_id = coalesce(?, chatgpt_account_id), + chatgpt_user_id = coalesce(?, chatgpt_user_id), + plan_type = coalesce(?, plan_type), access_token = ?, access_token_expires_at = ?, next_refresh_at = ?, diff --git a/backend/src/upstream/accounts/token_refresh/runtime.rs b/backend/src/upstream/accounts/token_refresh/runtime.rs index 7bf005bcc..4d474eeff 100644 --- a/backend/src/upstream/accounts/token_refresh/runtime.rs +++ b/backend/src/upstream/accounts/token_refresh/runtime.rs @@ -13,7 +13,7 @@ use chrono::{DateTime, Utc}; use secrecy::{ExposeSecret, SecretString}; use thiserror::Error; use tokio::{sync::Mutex, task::JoinHandle, time::sleep}; -use tracing::{debug, warn}; +use tracing::{debug, info, warn}; use uuid::Uuid; use crate::upstream::accounts::{ @@ -627,7 +627,19 @@ where updated.access_token_expires_at = jwt_expiration(&updated.access_token).or(updated.access_token_expires_at); updated.next_refresh_at = self.computed_refresh_at(&updated); + let refresh_token_rotated = updated.refresh_token != attempt.refresh_token; persist_token_update(&self.store, &updated).await?; + info!( + account_id = %updated.id, + trigger = ?trigger, + attempt = attempt_index + 1, + status = %updated.status, + access_token_expires_at = ?updated.access_token_expires_at, + next_refresh_at = ?updated.next_refresh_at, + refresh_token_present = updated.refresh_token.is_some(), + refresh_token_rotated, + "token refresh succeeded" + ); return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated))); } Err(RefreshError::RetryableTransport | RefreshError::Transport) diff --git a/backend/tests/upstream/accounts/account_repository.rs b/backend/tests/upstream/accounts/account_repository.rs index b3e209a39..4cf888d3d 100644 --- a/backend/tests/upstream/accounts/account_repository.rs +++ b/backend/tests/upstream/accounts/account_repository.rs @@ -252,6 +252,75 @@ async fn account_repository_should_not_reactivate_disabled_or_banned_accounts_fr ); } +#[tokio::test] +async fn account_repository_claims_update_should_preserve_optional_claims_when_missing() { + let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 49).await; + let repo = SqliteAccountStore::new(pool); + seed_plan_account(&repo, "acct_preserve_without_refresh").await; + + repo.update_from_claims( + "acct_preserve_without_refresh", + AccountClaimsUpdate { + email: None, + account_id: None, + user_id: None, + plan_type: None, + access_token: SecretString::new("new-access".to_string().into()), + refresh_token: None, + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + next_refresh_at: Some(Utc::now() + Duration::minutes(30)), + status: AccountStatus::Active, + }, + ) + .await + .expect("claims update should persist"); + let stored = repo + .get("acct_preserve_without_refresh") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(stored.email.as_deref(), Some("existing@example.com")); + assert_eq!(stored.account_id.as_deref(), Some("chatgpt-existing")); + assert_eq!(stored.user_id.as_deref(), Some("user-existing")); + assert_eq!(stored.plan_type.as_deref(), Some("plus")); +} + +#[tokio::test] +async fn account_repository_claims_update_should_preserve_optional_claims_when_refresh_token_rotates( +) { + let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 50).await; + let repo = SqliteAccountStore::new(pool); + seed_plan_account(&repo, "acct_preserve_with_refresh").await; + + repo.update_from_claims( + "acct_preserve_with_refresh", + AccountClaimsUpdate { + email: None, + account_id: None, + user_id: None, + plan_type: None, + access_token: SecretString::new("new-access".to_string().into()), + refresh_token: Some(SecretString::new("new-refresh".to_string().into())), + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + next_refresh_at: Some(Utc::now() + Duration::minutes(30)), + status: AccountStatus::Active, + }, + ) + .await + .expect("claims update should persist"); + let stored = repo + .get("acct_preserve_with_refresh") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(stored.email.as_deref(), Some("existing@example.com")); + assert_eq!(stored.account_id.as_deref(), Some("chatgpt-existing")); + assert_eq!(stored.user_id.as_deref(), Some("user-existing")); + assert_eq!(stored.plan_type.as_deref(), Some("plus")); +} + #[tokio::test] async fn account_repository_should_accumulate_usage_counters() { let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 7).await; @@ -740,6 +809,24 @@ async fn seed_new_account(repo: &SqliteAccountStore, id: &str) { .unwrap(); } +async fn seed_plan_account(repo: &SqliteAccountStore, id: &str) { + repo.insert(NewAccount { + id: id.to_string(), + email: Some("existing@example.com".to_string()), + account_id: Some("chatgpt-existing".to_string()), + user_id: Some("user-existing".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: Some(SecretString::new("refresh-secret".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + added_at: None, + }) + .await + .unwrap(); +} + async fn seed_repo_account(pool: &sqlx::SqlitePool, id: &str, added_at: &str) { sqlx::query( "insert into accounts (id, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?)", diff --git a/backend/tests/upstream/accounts/token_refresh/success.rs b/backend/tests/upstream/accounts/token_refresh/success.rs index 7be369f22..ed2984076 100644 --- a/backend/tests/upstream/accounts/token_refresh/success.rs +++ b/backend/tests/upstream/accounts/token_refresh/success.rs @@ -75,6 +75,71 @@ async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refre ); } +#[tokio::test] +async fn token_refresh_task_should_persist_rotated_refresh_token_when_returned() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-rotated.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool); + let now = Utc.with_ymd_and_hms(2026, 6, 18, 9, 0, 0).unwrap(); + let new_expires_at = now + Duration::hours(1); + let new_access_token = test_jwt(new_expires_at.timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-rotated".to_string(), + email: Some("rotated@example.com".to_string()), + account_id: Some("chatgpt-rotated".to_string()), + user_id: Some("user-rotated".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now + Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new("refresh-old".to_string().into())), + added_at: None, + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = StaticTokenRefresher { + response: Arc::new(Mutex::new(Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: Some("refresh-new".to_string()), + }))), + }; + let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("refresh should succeed"); + let stored = store + .get("acct-refresh-rotated") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(summary.refreshed, 1); + assert_eq!(stored.access_token.expose_secret(), new_access_token); + assert_eq!( + stored + .refresh_token + .as_ref() + .map(ExposeSecret::expose_secret), + Some("refresh-new") + ); +} + #[tokio::test] async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clearing_status() { let dir = tempfile::tempdir().expect("temp dir");