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,