From bf4da7a3e77aaa970b3430b69f3714bbdda2b689 Mon Sep 17 00:00:00 2001 From: zyycn Date: Wed, 29 Jul 2026 03:51:54 +0800 Subject: [PATCH] fix(accounts): unify credential and quota state Apply OpenAI reauthorization to the selected local account ID while atomically replacing its verified upstream identity. Converge successful quota observations across refresh paths and preserve explicitly disabled account state during connection checks, quota refreshes, and credential rotation. Co-authored-by: Codex --- .../src/model/provider_credentials.rs | 3 +- .../gateway-admin/tests/use_case/accounts.rs | 30 ++++ .../gateway-core/src/engine/credential.rs | 40 +++++ .../gateway-core/tests/engine/credential.rs | 15 +- .../src/postgres/provider_accounts.rs | 53 ++++++- .../src/redis/credential_cooldown.rs | 6 + .../tests/postgres/provider_accounts.rs | 146 +++++++++++++++++- .../tests/redis/credential_cooldown.rs | 51 +++++- backend/crates/providers/openai/src/admin.rs | 3 +- .../providers/openai/src/credential/admin.rs | 31 ++-- .../providers/openai/src/credential/oauth.rs | 27 +--- .../providers/openai/src/credential/quota.rs | 13 +- .../openai/src/credential/repository.rs | 3 + .../openai/tests/credential/oauth.rs | 35 ++++- .../openai/tests/credential/quota.rs | 76 ++++++++- .../openai/tests/provider/contract.rs | 10 +- backend/crates/providers/xai/src/admin.rs | 1 + 17 files changed, 473 insertions(+), 70 deletions(-) diff --git a/backend/crates/gateway-admin/src/model/provider_credentials.rs b/backend/crates/gateway-admin/src/model/provider_credentials.rs index 2676523e2..a04f3ff5d 100644 --- a/backend/crates/gateway-admin/src/model/provider_credentials.rs +++ b/backend/crates/gateway-admin/src/model/provider_credentials.rs @@ -5,7 +5,7 @@ use std::fmt; use async_trait::async_trait; use chrono::{DateTime, Utc}; use gateway_core::{ - engine::credential::{OpaqueProviderData, ProviderAccountId}, + engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity}, routing::{ProviderKind, UpstreamModelId}, }; @@ -496,6 +496,7 @@ pub struct PreparedCredentialRotationFacts { pub account_id: ProviderAccountId, pub provider_kind: ProviderKind, pub expected_credential_revision: Revision, + pub replacement_identity: Option, pub name: String, pub email: Option, pub plan_type: Option, diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index 15b1614b1..92c1038a7 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -155,6 +155,7 @@ impl FakeProviderAdmin { .current_credential_revision .lock() .expect("current credential revision"), + replacement_identity: None, name: account.name.clone(), email: account.email.clone(), plan_type: account.plan_type.clone(), @@ -801,6 +802,35 @@ async fn connection_test_should_probe_unavailable_account() { )); } +#[tokio::test] +async fn connection_test_should_preserve_disabled_account_status() { + let provider = FakeProviderAdmin::new("xai", events()); + let mut account = account_record("xai"); + account.enabled = false; + let store = FakeAccountStore::with_account(account, events()); + let services = + accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await; + + let events = services + .accounts() + .test_connection( + ProviderAccountId::new("acct_test").expect("account ID"), + gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"), + ) + .await + .expect("connection test stream") + .collect::>() + .await; + + assert!(matches!( + events.last(), + Some(AccountConnectionTestEvent::Failed { + account_status: gateway_admin::model::accounts::AccountStatus::Disabled, + .. + }) + )); +} + #[tokio::test] async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() { let provider = FakeProviderAdmin::new("openai", events()); diff --git a/backend/crates/gateway-core/src/engine/credential.rs b/backend/crates/gateway-core/src/engine/credential.rs index ae574afbb..a89dc9900 100644 --- a/backend/crates/gateway-core/src/engine/credential.rs +++ b/backend/crates/gateway-core/src/engine/credential.rs @@ -44,6 +44,46 @@ impl fmt::Display for ProviderAccountId { } } +/// credential 轮换时可选替换的上游账号身份。 +#[derive(Clone, PartialEq, Eq)] +pub struct ProviderAccountIdentity { + upstream_user_id: String, + upstream_account_id: Option, +} + +impl fmt::Debug for ProviderAccountIdentity { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProviderAccountIdentity") + .field("upstream_user_id", &"") + .field( + "upstream_account_id", + &self.upstream_account_id.as_ref().map(|_| ""), + ) + .finish() + } +} + +impl ProviderAccountIdentity { + #[must_use] + pub const fn new(upstream_user_id: String, upstream_account_id: Option) -> Self { + Self { + upstream_user_id, + upstream_account_id, + } + } + + #[must_use] + pub fn upstream_user_id(&self) -> &str { + &self.upstream_user_id + } + + #[must_use] + pub fn upstream_account_id(&self) -> Option<&str> { + self.upstream_account_id.as_deref() + } +} + /// `provider_accounts.credential_revision` 的正数 CAS revision。 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct CredentialRevision(NonZeroU64); diff --git a/backend/crates/gateway-core/tests/engine/credential.rs b/backend/crates/gateway-core/tests/engine/credential.rs index d2f42ba79..be44ff393 100644 --- a/backend/crates/gateway-core/tests/engine/credential.rs +++ b/backend/crates/gateway-core/tests/engine/credential.rs @@ -9,7 +9,8 @@ use gateway_core::engine::credential::{ AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker, AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate, CredentialRevision, OpaqueProviderData, PlaintextCredential, PreferredAccountSelection, - ProviderAccount, ProviderAccountId, ProviderAccountUpdate, RotationStrategy, + ProviderAccount, ProviderAccountId, ProviderAccountIdentity, ProviderAccountUpdate, + RotationStrategy, }; use gateway_core::routing::ProviderKind; @@ -73,6 +74,18 @@ fn opaque_provider_data_should_not_expose_quota_values_in_debug() { assert!(!format!("{quota:?}").contains("private-window")); } +#[test] +fn provider_account_identity_debug_should_redact_upstream_ids() { + let identity = ProviderAccountIdentity::new( + "private-user-id".to_owned(), + Some("private-account-id".to_owned()), + ); + let debug = format!("{identity:?}"); + + assert!(!debug.contains("private-user-id")); + assert!(!debug.contains("private-account-id")); +} + #[test] fn availability_should_round_trip_all_database_values() { assert_eq!( diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts.rs b/backend/crates/gateway-store/src/postgres/provider_accounts.rs index b7b828dc2..c28bd471a 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts.rs @@ -38,8 +38,8 @@ use gateway_core::engine::credential::{ CredentialRevision as CoreCredentialRevision, LoadedCredential, NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId, - ProviderAccountStore, ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation, - QuotaWriteOutcome, + ProviderAccountIdentity, ProviderAccountStore, + ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome, }; use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind}; use gateway_core::routing::ProviderKind; @@ -296,6 +296,7 @@ impl ImportProviderAccounts { pub struct RotateProviderAccount { pub scope: ProviderAccountAdminScope, pub profile: UpdateProviderAccount, + pub replacement_identity: Option, pub credential: ProviderCredentialUpdate, pub audit: AdminAuditEvent, } @@ -306,6 +307,7 @@ impl fmt::Debug for RotateProviderAccount { .debug_struct("RotateProviderAccount") .field("scope", &self.scope) .field("profile", &self.profile) + .field("replacement_identity", &self.replacement_identity) .field("credential", &self.credential) .field("audit", &self.audit) .finish() @@ -612,8 +614,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository { update.validate()?; let result = sqlx::query( "update provider_accounts - set availability = $3, availability_reason = $4, cooldown_until = $5, - availability_observed_at = $6, updated_at = greatest(now(), $6) + set availability = case when enabled then $3 else availability end, + availability_reason = case when enabled then $4 else availability_reason end, + cooldown_until = case when enabled then $5 else cooldown_until end, + availability_observed_at = case + when enabled then $6 + else availability_observed_at + end, + updated_at = case when enabled then greatest(now(), $6) else updated_at end where id = $1 and credential_revision = $2", ) .bind(update.account_id) @@ -755,6 +763,12 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository { command.scope.validate()?; require_nonempty(ENTITY, "account_id", &command.profile.id)?; require_nonempty(ENTITY, "name", &command.profile.name)?; + if let Some(identity) = &command.replacement_identity { + require_nonempty(ENTITY, "upstream_user_id", identity.upstream_user_id())?; + if let Some(account_id) = identity.upstream_account_id() { + require_nonempty(ENTITY, "upstream_account_id", account_id)?; + } + } if command.profile.id != command.credential.account_id { return Err(invalid("rotated profile and credential account IDs differ")); } @@ -770,6 +784,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository { &mut transaction, &command.scope, &command.profile, + command.replacement_identity.as_ref(), &command.credential, ) .await?; @@ -1040,6 +1055,7 @@ impl PgAdminAccountStore { email: prepared.email, plan_type: prepared.plan_type, }, + replacement_identity: prepared.replacement_identity, credential: ProviderCredentialUpdate { account_id: account_id.as_str().to_owned(), expected_revision: store_revision(prepared.expected_credential_revision)?, @@ -1975,8 +1991,13 @@ async fn rotate_provider_account_in_transaction( transaction: &mut Transaction<'_, Postgres>, scope: &ProviderAccountAdminScope, profile: &UpdateProviderAccount, + replacement_identity: Option<&ProviderAccountIdentity>, update: &ProviderCredentialUpdate, ) -> StoreResult { + let replace_identity = replacement_identity.is_some(); + let upstream_user_id = replacement_identity.map(ProviderAccountIdentity::upstream_user_id); + let upstream_account_id = + replacement_identity.and_then(ProviderAccountIdentity::upstream_account_id); let next = sqlx::query_scalar::<_, i64>( "update provider_accounts set name = $4, @@ -1987,19 +2008,25 @@ async fn rotate_provider_account_in_transaction( has_refresh_token = $8, access_token_expires_at = $9, next_refresh_at = $10, + upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end, + upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end, availability = case + when not enabled then availability when availability <> 'quota_exhausted' then 'ready' else availability end, availability_reason = case + when not enabled then availability_reason when availability <> 'quota_exhausted' then null else availability_reason end, cooldown_until = case + when not enabled then cooldown_until when availability <> 'quota_exhausted' then null else cooldown_until end, availability_observed_at = case + when not enabled then availability_observed_at when availability <> 'quota_exhausted' then now() else availability_observed_at end, @@ -2018,9 +2045,25 @@ async fn rotate_provider_account_in_transaction( .bind(update.has_refresh_token) .bind(update.access_token_expires_at) .bind(update.next_refresh_at) + .bind(replace_identity) + .bind(upstream_user_id) + .bind(upstream_account_id) .fetch_optional(&mut **transaction) .await - .map_err(|_| postgres_unavailable("rotate provider account in admin transaction"))? + .map_err(|error| { + if error + .as_database_error() + .is_some_and(sqlx::error::DatabaseError::is_unique_violation) + { + StoreError::Conflict { + entity: ENTITY, + id: update.account_id.clone(), + kind: ConflictKind::InvalidTransition, + } + } else { + postgres_unavailable("rotate provider account in admin transaction") + } + })? .ok_or_else(|| StoreError::Conflict { entity: ENTITY, id: update.account_id.clone(), diff --git a/backend/crates/gateway-store/src/redis/credential_cooldown.rs b/backend/crates/gateway-store/src/redis/credential_cooldown.rs index 512a6ad24..a1ccc0af1 100644 --- a/backend/crates/gateway-store/src/redis/credential_cooldown.rs +++ b/backend/crates/gateway-store/src/redis/credential_cooldown.rs @@ -504,6 +504,12 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore { let availability = change.availability; let cooldown_until = change.cooldown_until; self.authoritative.apply_state_change(change).await?; + let Some(account) = self.authoritative.get_account(&account_id).await? else { + return Ok(()); + }; + if !account.enabled() { + return Ok(()); + } self.mirror( &account_id, credential_revision, diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs index 9f2e00e3c..d8d0bd9ff 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs @@ -17,9 +17,10 @@ use gateway_admin::{ ports::store::AccountStore, }; use gateway_core::engine::credential::{ - AccountAvailability, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, - OpaqueProviderData, PlaintextCredential, ProviderAccountId, ProviderAccountStore, - ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome, + AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, + CredentialRevision, OpaqueProviderData, PlaintextCredential, ProviderAccountId, + ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation, + QuotaWriteOutcome, }; use gateway_core::routing::ProviderKind; use gateway_store::{ @@ -898,6 +899,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() { .rotate_provider_account(RotateProviderAccount { scope: wrong_scope, profile: profile("acct_admin_a", "wrong scope"), + replacement_identity: None, credential: credential_update("acct_admin_a", 1, "wrong-scope-secret"), audit: audit("audit_wrong_scope", "rotate", "acct_admin_a"), }) @@ -926,6 +928,10 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() { .rotate_provider_account(RotateProviderAccount { scope: scope.clone(), profile: profile("acct_admin_a", "rotated account"), + replacement_identity: Some(ProviderAccountIdentity::new( + "user-admin-rebound".to_owned(), + Some("workspace-admin-rebound".to_owned()), + )), credential: credential_update("acct_admin_a", 1, "rotated-secret"), audit: audit("audit_account_rotate", "rotate", "acct_admin_a"), }) @@ -933,15 +939,71 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() { .expect("rotate provider account"); assert_eq!(rotation.config_revision.get(), 3); assert_eq!(rotation.credential_revision.get(), 2); - let restored: (String, Option, Option>) = sqlx::query_as( - "select availability, availability_reason, cooldown_until + let restored: ( + String, + Option, + Option>, + String, + Option, + ) = sqlx::query_as( + "select availability, availability_reason, cooldown_until, + upstream_user_id, upstream_account_id from provider_accounts where id = $1", ) .bind("acct_admin_a") .fetch_one(&database.pool) .await .expect("load restored account state"); - assert_eq!(restored, ("ready".to_owned(), None, None)); + assert_eq!( + restored, + ( + "ready".to_owned(), + None, + None, + "user-admin-rebound".to_owned(), + Some("workspace-admin-rebound".to_owned()), + ) + ); + + let identity_conflict = repository + .rotate_provider_account(RotateProviderAccount { + scope: scope.clone(), + profile: profile("acct_admin_a", "must roll back"), + replacement_identity: Some(ProviderAccountIdentity::new( + "user-admin-b".to_owned(), + None, + )), + credential: credential_update("acct_admin_a", 2, "must-not-persist"), + audit: audit( + "audit_account_identity_conflict", + "reauthorize", + "acct_admin_a", + ), + }) + .await + .expect_err("upstream identity collision must roll back rotation"); + assert!(matches!( + identity_conflict, + StoreError::Conflict { + kind: ConflictKind::InvalidTransition, + .. + } + )); + assert_eq!(current_revision(&database.pool).await, 3); + let unchanged: (String, serde_json::Value, i64, String, Option) = sqlx::query_as( + "select name, provider_credentials_json, credential_revision, + upstream_user_id, upstream_account_id + from provider_accounts where id = $1", + ) + .bind("acct_admin_a") + .fetch_one(&database.pool) + .await + .expect("load account after identity collision"); + assert_eq!(unchanged.0, "rotated account"); + assert_eq!(unchanged.1["access_token"], "rotated-secret"); + assert_eq!(unchanged.2, 2); + assert_eq!(unchanged.3, "user-admin-rebound"); + assert_eq!(unchanged.4.as_deref(), Some("workspace-admin-rebound")); let revision = repository .set_provider_account_enabled_admin(SetProviderAccountEnabled { @@ -1013,6 +1075,7 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() { .rotate_provider_account(RotateProviderAccount { scope, profile: profile("acct_rotation_quota", "reauthorized account"), + replacement_identity: None, credential: credential_update("acct_rotation_quota", 1, "reauthorized-secret"), audit: audit( "audit_rotation_quota_rotate", @@ -1041,6 +1104,77 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() { database.close().await; } +#[tokio::test] +async fn disabled_account_preserves_user_state_during_refresh_writes() { + let Some(database) = TestDatabase::create("provider_account_disabled_refresh").await else { + return; + }; + let repository = PgProviderAccountRepository::new(database.pool.clone()); + let scope = ProviderAccountAdminScope { + provider_kind: "openai".to_owned(), + }; + let account_id = ProviderAccountId::new("acct_disabled_refresh").expect("account ID"); + let mut disabled = account(account_id.as_str(), "user-disabled-refresh"); + disabled.enabled = false; + disabled.availability = AccountAvailability::Expired; + disabled.availability_reason = Some("credential_expired".to_owned()); + repository + .import_provider_accounts(ImportProviderAccounts { + scope: scope.clone(), + accounts: vec![disabled], + audit: audit( + "audit_disabled_refresh_import", + "import", + account_id.as_str(), + ), + }) + .await + .expect("import disabled account"); + + repository + .apply_state_change(AccountStateChange { + account_id: account_id.clone(), + expected_revision: CredentialRevision::new(1).expect("credential revision"), + availability: AccountAvailability::Ready, + reason: None, + cooldown_until: None, + observed_at: SystemTime::now(), + }) + .await + .expect("disabled state write is a no-op"); + repository + .rotate_provider_account(RotateProviderAccount { + scope, + profile: profile(account_id.as_str(), "refreshed disabled account"), + replacement_identity: None, + credential: credential_update(account_id.as_str(), 1, "disabled-refreshed-secret"), + audit: audit( + "audit_disabled_refresh_rotate", + "refresh", + account_id.as_str(), + ), + }) + .await + .expect("refresh disabled account credential"); + + let current: (bool, String, Option, serde_json::Value, i64) = sqlx::query_as( + "select enabled, availability, availability_reason, + provider_credentials_json, credential_revision + from provider_accounts where id = $1", + ) + .bind(account_id.as_str()) + .fetch_one(&database.pool) + .await + .expect("load disabled account after refresh writes"); + assert!(!current.0); + assert_eq!(current.1, "expired"); + assert_eq!(current.2.as_deref(), Some("credential_expired")); + assert_eq!(current.3["access_token"], "disabled-refreshed-secret"); + assert_eq!(current.4, 2); + + database.close().await; +} + fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount { NewProviderAccount { id: id.to_owned(), diff --git a/backend/crates/gateway-store/tests/redis/credential_cooldown.rs b/backend/crates/gateway-store/tests/redis/credential_cooldown.rs index d66dc3867..3e2c83f17 100644 --- a/backend/crates/gateway-store/tests/redis/credential_cooldown.rs +++ b/backend/crates/gateway-store/tests/redis/credential_cooldown.rs @@ -50,9 +50,13 @@ impl ProviderAccountStore for CooldownTestAccountStore { async fn get_account( &self, - _account: &ProviderAccountId, + account: &ProviderAccountId, ) -> Result, CoreStoreError> { - unreachable!("cooldown tests do not get one account") + Ok(self + .accounts + .iter() + .find(|candidate| candidate.id() == account) + .cloned()) } async fn list_accounts(&self) -> Result, CoreStoreError> { @@ -190,10 +194,26 @@ fn composition_fixture( CooldownCachingProviderAccountStore, Arc, Arc>>, +) { + composition_fixture_with_accounts( + vec![ready_test_account("acct_cooldown_composition")], + fail_state_write, + fail_cache, + ) +} + +fn composition_fixture_with_accounts( + accounts: Vec, + fail_state_write: bool, + fail_cache: bool, +) -> ( + CooldownCachingProviderAccountStore, + Arc, + Arc>>, ) { let calls = Arc::new(Mutex::new(Vec::new())); let authoritative = Arc::new(CooldownTestAccountStore { - accounts: Vec::new(), + accounts, calls: Arc::clone(&calls), fail_state_write, }); @@ -249,6 +269,31 @@ async fn cooldown_composition_keeps_postgres_terminal_when_redis_fails() { assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]); } +#[tokio::test] +async fn disabled_account_state_write_does_not_change_runtime_cooldown() { + let disabled = test_provider_account("acct_cooldown_composition").with_runtime_state( + false, + AccountAvailability::Ready, + None, + ); + let (store, cache, calls) = composition_fixture_with_accounts(vec![disabled], false, false); + + store + .apply_state_change(cooldown_state_change(60)) + .await + .expect("disabled state write is a no-op"); + + assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]); + assert!(cache.cached.lock().expect("cached lock").is_empty()); + assert!( + cache + .invalidated + .lock() + .expect("invalidated lock") + .is_empty() + ); +} + #[tokio::test] async fn ready_state_invalidates_same_revision_cooldown_after_postgres_commit() { let (store, cache, calls) = composition_fixture(false, false); diff --git a/backend/crates/providers/openai/src/admin.rs b/backend/crates/providers/openai/src/admin.rs index b0dc51ae9..ecae3a396 100644 --- a/backend/crates/providers/openai/src/admin.rs +++ b/backend/crates/providers/openai/src/admin.rs @@ -586,7 +586,7 @@ fn prepared_rotation( prepared: crate::credential::PreparedCodexCredentialRotation, provider_kind: ProviderKind, ) -> Result { - let (profile, credential, guard) = prepared.into_parts(); + let (profile, credential, replacement_identity, guard) = prepared.into_parts(); let ( account_id, expected_revision, @@ -606,6 +606,7 @@ fn prepared_rotation( account_id, provider_kind, expected_credential_revision, + replacement_identity, name: profile.name, email: profile.email, plan_type: profile.plan_type, diff --git a/backend/crates/providers/openai/src/credential/admin.rs b/backend/crates/providers/openai/src/credential/admin.rs index fb0f815e9..d5fff00eb 100644 --- a/backend/crates/providers/openai/src/credential/admin.rs +++ b/backend/crates/providers/openai/src/credential/admin.rs @@ -11,7 +11,8 @@ use std::time::SystemTime; use chrono::{DateTime, FixedOffset, Utc}; use gateway_core::engine::credential::{ AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential, - NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountUpdate, + NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountIdentity, + ProviderAccountUpdate, }; use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, @@ -286,6 +287,7 @@ impl std::fmt::Debug for RotateManagedCodexCredential { pub struct PreparedCodexCredentialRotation { pub profile: ProviderAccountUpdate, pub credential: CredentialCasUpdate, + pub replacement_identity: Option, refresh_guards: Option, } @@ -300,6 +302,7 @@ impl std::fmt::Debug for PreparedCodexCredentialRotation { .debug_struct("PreparedCodexCredentialRotation") .field("profile", &self.profile) .field("credential", &self.credential) + .field("replacement_identity", &self.replacement_identity) .field( "refresh_guards", &self.refresh_guards.as_ref().map(|_| ""), @@ -316,11 +319,13 @@ impl PreparedCodexCredentialRotation { ) -> ( ProviderAccountUpdate, CredentialCasUpdate, + Option, PreparedCodexCredentialRotationGuard, ) { ( self.profile, self.credential, + self.replacement_identity, PreparedCodexCredentialRotationGuard(self.refresh_guards), ) } @@ -494,7 +499,7 @@ impl CodexCredentialAdmin { self.prepare_oauth_rotation(input, false) } - /// 完整 OAuth 授权允许认证主体投影更新，但不允许目标账号或用户发生变化。 + /// 完整 OAuth 授权按目标数据 ID 替换凭据与上游账号身份。 pub(crate) fn prepare_reauthorization( &self, input: RotateManagedCodexCredential, @@ -505,7 +510,7 @@ impl CodexCredentialAdmin { fn prepare_oauth_rotation( &self, input: RotateManagedCodexCredential, - replace_principal: bool, + replace_identity: bool, ) -> Result { let access_token_expires_at = required_time(input.verified_account.access_token_expires_at)?; @@ -516,16 +521,23 @@ impl CodexCredentialAdmin { .ok_or(CodexCredentialAdminError::IdentityMismatch)?; if input.current.account.provider().as_str() != PROVIDER_NAME || input.current.account.authentication_kind() != CODEX_AUTHENTICATION_KIND_OAUTH - || input.current.account.upstream_account_id() - != Some(input.verified_account.chatgpt_account_id.as_str()) - || input.current.account.upstream_user_id() != input.verified_account.chatgpt_user_id - || (!replace_principal - && (oauth.principal.oauth_subject != input.verified_account.oauth_subject + || (!replace_identity + && (input.current.account.upstream_account_id() + != Some(input.verified_account.chatgpt_account_id.as_str()) + || input.current.account.upstream_user_id() + != input.verified_account.chatgpt_user_id + || oauth.principal.oauth_subject != input.verified_account.oauth_subject || oauth.principal.poid != input.verified_account.poid)) { return Err(CodexCredentialAdminError::IdentityMismatch); } - if replace_principal { + let replacement_identity = replace_identity.then(|| { + ProviderAccountIdentity::new( + input.verified_account.chatgpt_user_id.clone(), + Some(input.verified_account.chatgpt_account_id.clone()), + ) + }); + if replace_identity { oauth .principal .oauth_subject @@ -567,6 +579,7 @@ impl CodexCredentialAdmin { Ok(PreparedCodexCredentialRotation { profile, credential, + replacement_identity, refresh_guards: None, }) } diff --git a/backend/crates/providers/openai/src/credential/oauth.rs b/backend/crates/providers/openai/src/credential/oauth.rs index 0e227d9a8..eca255ecf 100644 --- a/backend/crates/providers/openai/src/credential/oauth.rs +++ b/backend/crates/providers/openai/src/credential/oauth.rs @@ -29,7 +29,6 @@ use super::admin::{ use super::identity::{ CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError, }; -use super::security::CodexCredentialCodec; use super::token_client::{ AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant, OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI, @@ -500,11 +499,7 @@ impl CodexOAuthAdminService { } else { None }; - let expectation = current - .as_ref() - .map(oauth_identity_expectation) - .transpose()? - .unwrap_or_default(); + let expectation = CodexIdentityExpectation::default(); let profile = self .verifier .verify_authorization(&secret, &id_token, &pending.nonce, &expectation) @@ -679,26 +674,6 @@ impl CodexOAuthAdminService { } } -fn oauth_identity_expectation( - current: &gateway_core::engine::credential::LoadedCredential, -) -> Result { - let runtime = CodexCredentialCodec::decode(¤t.credential) - .map_err(|_| CodexOAuthAdminError::Credential)?; - if runtime.authentication.oauth().is_none() { - return Err(CodexOAuthAdminError::Credential); - } - let account_id = current - .account - .upstream_account_id() - .ok_or(CodexOAuthAdminError::Credential)?; - CodexIdentityExpectation::reauthorization( - account_id.to_owned(), - current.account.upstream_user_id().to_owned(), - runtime.installation_id, - ) - .map_err(|_| CodexOAuthAdminError::Credential) -} - fn authorization_url( pending: &CodexPendingAuthorization, oauth_client_id: &str, diff --git a/backend/crates/providers/openai/src/credential/quota.rs b/backend/crates/providers/openai/src/credential/quota.rs index 33d85c783..81d21270d 100644 --- a/backend/crates/providers/openai/src/credential/quota.rs +++ b/backend/crates/providers/openai/src/credential/quota.rs @@ -1392,10 +1392,16 @@ fn quota_success_availability( exhausted: bool, now: SystemTime, ) -> Option { + let observed = if exhausted { + AccountAvailability::QuotaExhausted + } else { + AccountAvailability::Ready + }; match current { AccountAvailability::Invalid | AccountAvailability::Expired - | AccountAvailability::Banned => None, + | AccountAvailability::Banned + | AccountAvailability::Unknown => Some(observed), AccountAvailability::QuotaExhausted => (!exhausted).then_some(AccountAvailability::Ready), AccountAvailability::Ready => exhausted.then_some(AccountAvailability::QuotaExhausted), AccountAvailability::Cooldown if exhausted => Some(AccountAvailability::QuotaExhausted), @@ -1405,11 +1411,6 @@ fn quota_success_availability( Some(AccountAvailability::Ready) } AccountAvailability::Cooldown => None, - AccountAvailability::Unknown => Some(if exhausted { - AccountAvailability::QuotaExhausted - } else { - AccountAvailability::Ready - }), } } diff --git a/backend/crates/providers/openai/src/credential/repository.rs b/backend/crates/providers/openai/src/credential/repository.rs index f33317622..f5e527820 100644 --- a/backend/crates/providers/openai/src/credential/repository.rs +++ b/backend/crates/providers/openai/src/credential/repository.rs @@ -237,6 +237,9 @@ impl CodexCredentialRepository { cooldown_until: Option, observed_at: SystemTime, ) -> Result<(), CredentialRepositoryError> { + if !account.enabled() { + return Ok(()); + } self.store .apply_state_change(AccountStateChange { account_id: account.id().clone(), diff --git a/backend/crates/providers/openai/tests/credential/oauth.rs b/backend/crates/providers/openai/tests/credential/oauth.rs index 670bf8a0c..2bf278249 100644 --- a/backend/crates/providers/openai/tests/credential/oauth.rs +++ b/backend/crates/providers/openai/tests/credential/oauth.rs @@ -600,10 +600,10 @@ async fn reauthorization_uses_revision_advanced_after_start_for_prepared_rotatio ); } -struct RebindingVerifier; +struct ReplacementIdentityVerifier; #[async_trait] -impl CodexAccountIdentityVerifier for RebindingVerifier { +impl CodexAccountIdentityVerifier for ReplacementIdentityVerifier { async fn verify( &self, _: &CodexOAuthSecret, @@ -617,8 +617,9 @@ impl CodexAccountIdentityVerifier for RebindingVerifier { _: &CodexOAuthSecret, _: &SecretString, _: &SecretString, - _: &CodexIdentityExpectation, + expectation: &CodexIdentityExpectation, ) -> Result { + assert_eq!(expectation, &CodexIdentityExpectation::default()); Ok(CodexIdentityVerification::Complete(profile( "chatgpt-different-owner", ))) @@ -699,7 +700,7 @@ async fn reauthorization_updates_changed_free_principal_for_same_account() { } #[tokio::test] -async fn reauthorization_rejects_identity_rebinding() { +async fn reauthorization_rebinds_target_record_to_authorized_identity() { let store = Arc::new(MemoryAccountStore::default()); store .seed_oauth_credential(ImportCodexOAuthCredential { @@ -711,7 +712,7 @@ async fn reauthorization_rejects_identity_rebinding() { enabled: true, }) .await; - let service = service_with_store(store, Arc::new(RebindingVerifier)); + let service = service_with_store(store, Arc::new(ReplacementIdentityVerifier)); let started = service .start_authorization(StartCodexOAuthAuthorization { mutation: reauthorization_mutation("request-rebind", "acct_oauth_rebind"), @@ -723,7 +724,7 @@ async fn reauthorization_rejects_identity_rebinding() { .query_pairs() .find_map(|(key, value)| (key == "state").then(|| value.into_owned())) .expect("state"); - let error = service + let completed = service .complete_authorization(CompleteCodexOAuthAuthorization { owner_ref: owner_ref(), flow_id: started.flow_id, @@ -732,6 +733,24 @@ async fn reauthorization_rejects_identity_rebinding() { )), }) .await - .expect_err("identity rebind"); - assert_eq!(error, CodexOAuthAdminError::Credential); + .expect("complete identity replacement"); + let CompletedCodexOAuthCredential::Reauthorize(prepared) = completed.credential else { + panic!("expected prepared reauthorization"); + }; + assert_eq!( + prepared.credential.account_id().as_str(), + "acct_oauth_rebind" + ); + let replacement = prepared + .replacement_identity + .as_ref() + .expect("replacement identity"); + assert_eq!( + replacement.upstream_account_id(), + Some("chatgpt-different-owner") + ); + assert_eq!( + replacement.upstream_user_id(), + "user-chatgpt-different-owner" + ); } diff --git a/backend/crates/providers/openai/tests/credential/quota.rs b/backend/crates/providers/openai/tests/credential/quota.rs index c17e01fa6..120264089 100644 --- a/backend/crates/providers/openai/tests/credential/quota.rs +++ b/backend/crates/providers/openai/tests/credential/quota.rs @@ -181,6 +181,14 @@ fn blocked_network_quota_service(store: &Arc) -> CodexCreden } async fn create_account(store: &Arc, account_id: &str) { + create_account_with_enabled(store, account_id, true).await; +} + +async fn create_account_with_enabled( + store: &Arc, + account_id: &str, + enabled: bool, +) { store .seed_oauth_credential(ImportCodexOAuthCredential { account_id: account_id.to_owned(), @@ -188,7 +196,7 @@ async fn create_account(store: &Arc, account_id: &str) { secret: secret(&format!("token-{account_id}")), verified_account: profile(&format!("chatgpt-{account_id}")), next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)), - enabled: true, + enabled, }) .await; } @@ -367,6 +375,17 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() { let account_id = "acct_manual_quota_state"; create_account(&store, account_id).await; let account = store.account(account_id).expect("created account"); + store + .apply_state_change(AccountStateChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + availability: AccountAvailability::Expired, + reason: Some("stale_expired_state".to_owned()), + cooldown_until: None, + observed_at: SystemTime::now(), + }) + .await + .expect("seed stale expired state"); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/api/codex/usage")) @@ -438,6 +457,61 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() { ); } +#[tokio::test] +async fn manual_quota_refresh_preserves_disabled_account_state() { + let store = Arc::new(MemoryAccountStore::default()); + let account_id = "acct_disabled_quota_state"; + create_account_with_enabled(&store, account_id, false).await; + let account = store.account(account_id).expect("created disabled account"); + store + .apply_state_change(AccountStateChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + availability: AccountAvailability::Expired, + reason: Some("disabled_expired_state".to_owned()), + cooldown_until: None, + observed_at: SystemTime::now(), + }) + .await + .expect("seed disabled account state"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .and(header( + "authorization", + format!("Bearer token-{account_id}"), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": {"used_percent": 20, "reset_at": 1_900_000_000} + } + }))) + .mount(&server) + .await; + let service = quota_service_with_base_url( + &store, + reqwest::Client::builder() + .no_proxy() + .build() + .expect("client"), + server.uri(), + ); + + service + .refresh_account(account.id()) + .await + .expect("refresh disabled account quota"); + + let current = store + .account(account_id) + .expect("disabled account after quota refresh"); + assert!(!current.enabled()); + assert_eq!(current.availability(), AccountAvailability::Expired); + assert!(store.has_quota(account_id)); +} + #[tokio::test] async fn manual_quota_auth_rejection_does_not_invalidate_refreshable_credential() { let store = Arc::new(MemoryAccountStore::default()); diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 94a27e05b..9767acf25 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -7,7 +7,7 @@ use base64::{Engine as _, engine::general_purpose::STANDARD}; use chrono::Utc; use futures::StreamExt; use gateway_core::engine::credential::{ - AccountAvailability, AccountFeedbackStats, ProviderAccountId, + AccountAvailability, AccountFeedbackStats, ProviderAccountId, ProviderAccountStore as _, }; use gateway_core::engine::provider::{Provider as _, ProviderRequest}; use gateway_core::engine::{ @@ -1218,8 +1218,8 @@ async fn response_failed_after_semantic_output_is_exposed_and_not_replay_safe() async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_state() { let store = Arc::new(MemoryAccountStore::default()); let account_id = "acct_disabled_diagnostic"; - create_account_with_enabled(&store, account_id, false).await; - let account = store.account(account_id).expect("disabled test account"); + create_account(&store, account_id).await; + let account = store.account(account_id).expect("test account"); store .repository() .apply_state( @@ -1231,6 +1231,10 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st ) .await .expect("seed quota-exhausted state"); + store + .set_enabled(account.id(), false) + .await + .expect("disable test account"); let affinity = Arc::new(MemorySessionAffinity::default()); let server = MockServer::start().await; diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs index 688ae4d20..ab39cc00d 100644 --- a/backend/crates/providers/xai/src/admin.rs +++ b/backend/crates/providers/xai/src/admin.rs @@ -942,6 +942,7 @@ fn prepared_rotation( provider_kind, expected_credential_revision: Revision::new(expected_revision.get()) .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?, + replacement_identity: None, name: profile.name, email: profile.email, plan_type: profile.plan_type,