starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): sync account plans with quota updates (#112)

Browse Source 
* fix(openai): sync account plans with quota updates

Co-authored-by: Codex <noreply@openai.com>

* test(store): preserve settings in account plan refresh fixture

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-16 18:22:01 +08:00
committed by GitHub
co-authored by Codex Zyy
parent 47ab48b4cf
commit 866f71e3fb
28 changed files with 522 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-admin/src/model/provider_credentials.rs
+2
View File
@@ -614,6 +614,8 @@ pub struct PreparedCredentialRotationFacts {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
/// Token 刷新保留提交时的资料，不以准备阶段的副本覆盖新套餐。
pub preserve_profile: bool,
pub provider_material: ProviderDocument,
pub has_refresh_token: bool,
pub access_token_expires_at: Option<DateTime<Utc>>,
backend/crates/gateway-admin/tests/use_case/accounts.rs
+1
View File
@@ -235,6 +235,7 @@ impl FakeProviderAdmin {
name: account.name.clone(),
email: account.email.clone(),
plan_type: account.plan_type.clone(),
preserve_profile: false,
provider_material: document(),
has_refresh_token: account.has_refresh_token,
access_token_expires_at: account
backend/crates/gateway-core/src/account/model.rs
+15
View File
@@ -1065,6 +1065,7 @@ pub struct CredentialCasUpdateParts {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub profile: ProviderAccountUpdate,
pub preserve_profile: bool,
pub credential: PlaintextCredential,
pub has_refresh_token: bool,
pub access_token_expires_at: Option<SystemTime>,
@@ -1078,6 +1079,7 @@ pub struct CredentialCasUpdate {
account_id: ProviderAccountId,
expected_revision: CredentialRevision,
profile: ProviderAccountUpdate,
preserve_profile: bool,
credential: PlaintextCredential,
has_refresh_token: bool,
access_token_expires_at: Option<SystemTime>,
@@ -1092,6 +1094,7 @@ impl fmt::Debug for CredentialCasUpdate {
.field("account_id", &self.account_id)
.field("expected_revision", &self.expected_revision)
.field("profile", &self.profile)
.field("preserve_profile", &self.preserve_profile)
.field("credential", &self.credential)
.field("has_refresh_token", &self.has_refresh_token)
.field("access_token_expires_at", &self.access_token_expires_at)
@@ -1126,6 +1129,7 @@ impl CredentialCasUpdate {
account_id,
expected_revision,
profile,
preserve_profile: false,
credential,
has_refresh_token,
access_token_expires_at,
@@ -1134,6 +1138,13 @@ impl CredentialCasUpdate {
})
}
/// 仅轮换凭据，保留提交时的账号资料，避免覆盖并发额度观测更新的套餐。
#[must_use]
pub const fn preserving_profile(mut self) -> Self {
self.preserve_profile = true;
self
}
/// 将刷新调度与账号错误事实放入同一个 revision CAS。
#[must_use]
pub fn with_account_state(
@@ -1199,6 +1210,7 @@ impl CredentialCasUpdate {
account_id: self.account_id,
expected_revision: self.expected_revision,
profile: self.profile,
preserve_profile: self.preserve_profile,
credential: self.credential,
has_refresh_token: self.has_refresh_token,
access_token_expires_at: self.access_token_expires_at,
@@ -1221,6 +1233,8 @@ pub struct QuotaObservation {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub quota: OpaqueProviderData,
/// Provider 确认的账号套餐，与额度原子写入；`None` 保留已有套餐。
pub plan_type: Option<String>,
/// Provider 原始 quota document 的观察时间；不代表访问结论发生变化。
pub observed_at: SystemTime,
/// Provider 已从私有 JSON 归一化出的额度访问事实。
@@ -1234,6 +1248,7 @@ impl fmt::Debug for QuotaObservation {
.field("account_id", &self.account_id)
.field("expected_revision", &self.expected_revision)
.field("quota", &self.quota)
.field("plan_type", &self.plan_type)
.field("observed_at", &self.observed_at)
.field("state", &self.state)
.finish()
backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs
+1
View File
@@ -323,6 +323,7 @@ impl PgAdminAccountStore {
has_refresh_token: prepared.has_refresh_token,
access_token_expires_at: prepared.access_token_expires_at,
next_refresh_at: prepared.next_refresh_at,
preserve_profile: prepared.preserve_profile,
},
audit: mutation_audit(
context,
backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs
+10 -2
View File
@@ -148,6 +148,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
account_id,
expected_revision,
profile,
preserve_profile,
credential,
has_refresh_token,
access_token_expires_at,
@@ -178,7 +179,9 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.and_then(|state| state.message.as_deref());
let next = sqlx::query_scalar::<_, i64>(
"update provider_accounts
set name = $3, email = $4, plan_type = $5,
set name = case when $14 then name else $3 end,
email = case when $14 then email else $4 end,
plan_type = case when $14 then plan_type else $5 end,
provider_credentials_json = $6,
credential_revision = credential_revision + 1,
has_refresh_token = $7, access_token_expires_at = $8,
@@ -212,6 +215,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.bind(credential_observed_at)
.bind(error_reason)
.bind(message)
.bind(preserve_profile)
.fetch_optional(&self.pool)
.await
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::Unavailable))?;
@@ -236,7 +240,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(|account| account.as_str().to_owned())
.collect::<Vec<_>>();
let rows = sqlx::query(
"select id, credential_revision, provider_quota_json, quota_observed_at, \
"select id, credential_revision, plan_type, provider_quota_json, quota_observed_at, \
quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at \
from provider_accounts \
where id = any($1) and quota_observed_at is not null",
@@ -294,6 +298,9 @@ impl ProviderAccountStore for PgProviderAccountRepository {
)
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
Ok(QuotaObservation {
plan_type: row
.try_get("plan_type")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
account_id: CoreProviderAccountId::new(account_id)
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
expected_revision: revision,
@@ -322,6 +329,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
quota,
DateTime::<Utc>::from(observation.observed_at),
observation.state,
observation.plan_type.as_deref(),
)
.await
.map_err(core_store_error)?;
backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs
+8 -3
View File
@@ -28,6 +28,7 @@ pub trait ProviderAccountRepository: Send + Sync {
quota: JsonObject,
observed_at: DateTime<Utc>,
state: QuotaState,
plan_type: Option<&str>,
) -> StoreResult<bool>;
async fn touch_provider_quota_observation(
&self,
@@ -302,6 +303,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
quota: JsonObject,
observed_at: DateTime<Utc>,
state: QuotaState,
plan_type: Option<&str>,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "account_id", account_id)?;
validate_object_size("provider_quota_json", &quota, QUOTA_MAX_BYTES)?;
@@ -309,6 +311,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
let result = sqlx::query(
"update provider_accounts
set provider_quota_json = $3, quota_observed_at = $4,
plan_type = coalesce($9, plan_type),
quota_access_state = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
@@ -337,6 +340,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(state.evidence().map(QuotaEvidence::as_str))
.bind(access_observed_at)
.bind(state.reset_at().map(DateTime::<Utc>::from))
.bind(plan_type)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("compare and swap provider quota"))?;
@@ -904,9 +908,9 @@ pub(crate) async fn rotate_provider_account_in_transaction(
// 事务时间可能早于应用写入的额度观测，保留既有时间下界，避免轮换破坏时间约束。
let next = sqlx::query_scalar::<_, i64>(
"update provider_accounts
set name = $4,
email = $5,
plan_type = $6,
set name = case when $14 then name else $4 end,
email = case when $14 then email else $5 end,
plan_type = case when $14 then plan_type else $6 end,
provider_credentials_json = $7,
credential_revision = credential_revision + 1,
has_refresh_token = $8,
@@ -943,6 +947,7 @@ pub(crate) async fn rotate_provider_account_in_transaction(
.bind(replace_identity)
.bind(upstream_user_id)
.bind(upstream_account_id)
.bind(update.preserve_profile)
.fetch_optional(&mut **transaction)
.await
.map_err(|error| {
backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs
+1
View File
@@ -357,6 +357,7 @@ pub struct ProviderCredentialUpdate {
pub has_refresh_token: bool,
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub preserve_profile: bool,
}
impl fmt::Debug for ProviderCredentialUpdate {
backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs
+128
View File
@@ -204,6 +204,7 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
let observed_at = SystemTime::now();
let outcome = repository
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: ProviderAccountId::new(id).expect("account id"),
expected_revision: revision,
quota: OpaqueProviderData::new(
@@ -279,6 +280,7 @@ async fn quota_observation_touch_preserves_quota_state_and_advances_account_upda
assert_eq!(
repository
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account_id.clone(),
expected_revision: revision,
quota: OpaqueProviderData::new(
@@ -377,6 +379,7 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
.expect("older observation time");
let quota = |marker: &str, observed_at| QuotaObservation {
plan_type: Some(if marker == "newer" { "pro" } else { "plus" }.to_owned()),
account_id: account_id.clone(),
expected_revision: revision,
quota: OpaqueProviderData::new(
@@ -435,6 +438,7 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
.expect("load fenced account")
.expect("fenced account");
assert_eq!(current.credential_state(), CredentialState::Invalid);
assert_eq!(current.plan_type(), Some("pro"));
let observed = repository
.get_quotas(std::slice::from_ref(&account_id))
.await
@@ -446,6 +450,128 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
database.close().await;
}
#[tokio::test]
async fn quota_plan_changes_survive_inflight_background_and_manual_token_refresh() {
let Some(database) = TestDatabase::create("quota_plan_refresh").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let id = ProviderAccountId::new("acct_plan_refresh").unwrap();
let mut seed = account(id.as_str(), "user-plan-refresh");
seed.plan_type = Some("plus".to_owned());
repository.insert_provider_account(seed).await.unwrap();
let original = repository.get_account(&id).await.unwrap().unwrap();
let config_revision = current_revision(&database.pool).await;
let observed_at = SystemTime::now();
let observation = QuotaObservation {
account_id: id.clone(),
expected_revision: original.revision(),
plan_type: Some("pro".to_owned()),
quota: OpaqueProviderData::new(json!({"plan_type": "pro"}).as_object().unwrap().clone()),
observed_at,
state: QuotaState::allowed(observed_at),
};
assert_eq!(
repository
.compare_and_swap_quota(observation.clone())
.await
.unwrap(),
QuotaWriteOutcome::Updated
);
let upgraded = repository.get_account(&id).await.unwrap().unwrap();
assert_eq!(upgraded.plan_type(), Some("pro"));
assert_eq!(upgraded.revision(), original.revision());
assert_eq!(current_revision(&database.pool).await, config_revision);
// 刷新在额度更新前已准备好旧资料，CAS 提交仍须保留数据库中的新套餐。
let refresh = CredentialCasUpdate::new(
id.clone(),
original.revision(),
ProviderAccountUpdate {
account_id: id.clone(),
name: original.name().to_owned(),
email: original.email().map(str::to_owned),
plan_type: Some("plus".to_owned()),
},
PlaintextCredential::new(
json!({"access_token": "background-refreshed"})
.as_object()
.unwrap()
.clone(),
),
false,
original.access_token_expires_at(),
None,
)
.unwrap()
.preserving_profile();
assert!(matches!(
repository
.compare_and_swap_credential(refresh)
.await
.unwrap(),
CredentialCasOutcome::Updated(_)
));
let current = repository.get_account(&id).await.unwrap().unwrap();
assert_eq!(current.plan_type(), Some("pro"));
let stale = QuotaObservation {
plan_type: Some("free".to_owned()),
..observation
};
assert_eq!(
repository.compare_and_swap_quota(stale).await.unwrap(),
QuotaWriteOutcome::Conflict
);
let mut credential =
credential_update(id.as_str(), current.revision().get(), "manual-refreshed");
credential.preserve_profile = true;
repository
.rotate_provider_account(RotateProviderAccount {
settings: None,
scope: ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
},
profile: profile(id.as_str(), "stale profile"),
replacement_identity: None,
credential,
audit: audit("audit_plan_refresh", "refresh", id.as_str()),
})
.await
.unwrap();
let refreshed = repository.get_account(&id).await.unwrap().unwrap();
assert_eq!(refreshed.plan_type(), Some("pro"));
assert_eq!(refreshed.name(), original.name());
assert_eq!(refreshed.email(), original.email());
let observed_at = SystemTime::now();
repository
.compare_and_swap_quota(QuotaObservation {
account_id: id.clone(),
expected_revision: refreshed.revision(),
plan_type: None,
quota: OpaqueProviderData::new(
json!({"rate_limit": {"allowed": true}})
.as_object()
.unwrap()
.clone(),
),
observed_at,
state: QuotaState::allowed(observed_at),
})
.await
.unwrap();
assert_eq!(
repository
.get_account(&id)
.await
.unwrap()
.unwrap()
.plan_type(),
Some("pro")
);
database.close().await;
}
#[tokio::test]
async fn credential_error_message_is_persisted_and_cleared_on_recovery() {
let Some(database) = TestDatabase::create("provider_account_state_message").await else {
@@ -1411,6 +1537,7 @@ async fn account_recovery_resets_status_facts_without_changing_credentials_or_sc
let observed_at = SystemTime::now();
repository
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: ProviderAccountId::new("acct_recovery").expect("account ID"),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
quota: OpaqueProviderData::new(
@@ -2831,6 +2958,7 @@ pub(super) fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
fn credential_update(account_id: &str, revision: u64, marker: &str) -> ProviderCredentialUpdate {
ProviderCredentialUpdate {
preserve_profile: false,
account_id: account_id.to_owned(),
expected_revision: Revision::new(revision).expect("credential revision"),
provider_credentials_json: credential_json(marker),
backend/crates/providers/openai/src/admin.rs
+2
View File
@@ -784,6 +784,7 @@ fn prepared_rotation(
account_id,
expected_revision,
profile: credential_profile,
preserve_profile,
credential,
has_refresh_token,
access_token_expires_at,
@@ -804,6 +805,7 @@ fn prepared_rotation(
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
preserve_profile,
provider_material: ProviderDocument::new(OpaqueProviderData::new(
credential.into_inner(),
)),
backend/crates/providers/openai/src/credential/admin.rs
+2 -1
View File
@@ -737,7 +737,8 @@ impl CodexCredentialAdmin {
optional_time(access_token_expires_at),
optional_time(next_refresh_at),
)
.map_err(|_| CodexCredentialAdminError::InvalidCredential)?;
.map_err(|_| CodexCredentialAdminError::InvalidCredential)?
.preserving_profile();
Ok(PreparedCodexCredentialRotation {
profile,
credential,
backend/crates/providers/openai/src/credential/quota/mod.rs
+36
View File
@@ -827,6 +827,7 @@ impl CodexCredentialQuotaService {
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()),
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(object),
@@ -890,6 +891,17 @@ impl CodexCredentialQuotaService {
let existing = existing
.map(|observation| observation.quota.into_inner())
.unwrap_or_default();
// 只用本次响应明确携带的套餐更新账号，不能把合并前的旧快照重新当作新证据。
let observed_plan = rate_limits
.iter()
.rev()
.filter_map(|observation| {
observation.plan_type.as_deref().filter(|plan| {
!plan.trim().is_empty() && !plan.trim().eq_ignore_ascii_case("unknown")
})
})
.next();
let plan_type = observed_account_plan(account.plan_type(), observed_plan);
// 套餐、credits 等元数据可以更新，但没有额度窗口事实时必须保留旧观察时刻，
// 也不能借旧快照重新推导 quota state。
if !has_quota_facts {
@@ -899,6 +911,7 @@ impl CodexCredentialQuotaService {
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
plan_type,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)),
@@ -921,6 +934,7 @@ impl CodexCredentialQuotaService {
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
plan_type,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(quota),
@@ -1152,6 +1166,7 @@ impl CodexCredentialQuotaService {
if self
.store
.compare_and_swap_quota(QuotaObservation {
plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()),
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(object),
@@ -1386,6 +1401,27 @@ async fn fetch_usage_with_5xx_retry(
}
}
/// 上游额度可确认套餐变更；同族泛化值不能丢弃 JWT 已给出的具体 SKU。
fn observed_account_plan(current: Option<&str>, observed: Option<&str>) -> Option<String> {
let plan = observed?.trim().to_ascii_lowercase();
if plan.is_empty() || plan == "unknown" {
return None;
}
// 套餐族沿用官方 codex_protocol::account::PlanType 的分类。
let current = current.unwrap_or_default().trim().to_ascii_lowercase();
let generalized = matches!(
(plan.as_str(), current.as_str()),
(
"team",
"self_serve_business_prolite" | "self_serve_business_usage_based"
) | (
"business",
"ent26" | "enterprise_cbp_automation" | "enterprise_cbp_usage_based"
) | ("edu" | "education", "edu_plus" | "edu_pro")
);
(!generalized).then_some(plan)
}
fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool {
account.enabled() && access_token_is_current(account, now)
}
backend/crates/providers/openai/src/credential/repository.rs
+7 -8
View File
@@ -48,9 +48,8 @@ impl CodexCredentialRepository {
.store
.load_credential(account.id(), account.revision())
.await?;
if current.account != *account {
return Err(CredentialRepositoryError::RevisionConflict);
}
// load_credential 已校验凭据版本；套餐与额度可在 RT exchange 期间更新，
// 这些运行时事实变化不能使已成功轮换的 token 丢失。
let mut data = CodexCredentialCodec::decode_complete(&current.credential)?;
let oauth = data
.oauth_mut()
@@ -75,6 +74,7 @@ impl CodexCredentialRepository {
next_refresh_at,
)
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?
.preserving_profile()
.with_account_state(
oauth_account_state(account, CredentialState::Ready),
SystemTime::now(),
@@ -96,9 +96,6 @@ impl CodexCredentialRepository {
.store
.load_credential(account.id(), account.revision())
.await?;
if current.account != *account {
return Err(CredentialRepositoryError::RevisionConflict);
}
let data = CodexCredentialCodec::decode_complete(&current.credential)?;
let has_refresh_token = data.has_refresh_token();
let credential = CodexCredentialCodec::encode_complete(data)?;
@@ -111,7 +108,8 @@ impl CodexCredentialRepository {
account.access_token_expires_at(),
Some(next_refresh_at),
)
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?
.preserving_profile();
if let Some(error_reason) = error_reason {
update = update.with_account_state(
account.credential_state(),
@@ -207,7 +205,8 @@ impl CodexCredentialRepository {
account.access_token_expires_at(),
account.next_refresh_at(),
)
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?
.preserving_profile();
cas_revision(self.store.compare_and_swap_credential(update).await?)
}
backend/crates/providers/openai/tests/admin.rs
+74
View File
@@ -721,6 +721,54 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export
assert!(exported_account.get("token").is_none());
}
#[tokio::test]
async fn openai_admin_quota_refresh_updates_the_account_plan() {
let store = Arc::new(MemoryAccountStore::default());
let mut verified_account = profile("chatgpt-upgraded-plan");
verified_account.plan_type = Some("plus".to_owned());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_upgraded_plan".to_owned(),
name: "upgraded plan".to_owned(),
secret: secret("upgraded-plan-test-token"),
verified_account,
next_refresh_at: None,
enabled: true,
})
.await;
let account = store.account("acct_upgraded_plan").unwrap();
let server = MockServer::start().await;
Mock::given(method("GET")).and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"plan_type": "pro", "rate_limit": {"allowed": true, "primary_window": {"used_percent": 1}}
})))
.expect(1).mount(&server).await;
let mut config = valid_config();
config.config.api.base_url = server.uri();
let bundle = provider_openai::initialize(
config.config,
provider_ports_with(store.clone(), Arc::new(TestOAuthPending::default())),
)
.await
.unwrap();
for refresh in [true, false] {
let quota = bundle
.admin_provider()
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh,
rolling_usage: None,
})
.await
.unwrap();
assert_eq!(quota.plan_type.as_deref(), Some("pro"));
assert_eq!(
store.account("acct_upgraded_plan").unwrap().plan_type(),
Some("pro")
);
}
}
#[tokio::test]
async fn openai_admin_projects_free_plan_from_cached_quota_when_account_claims_omit_it() {
let store = Arc::new(MemoryAccountStore::default());
@@ -740,6 +788,7 @@ async fn openai_admin_projects_free_plan_from_cached_quota_when_account_claims_o
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
@@ -853,6 +902,7 @@ async fn openai_admin_provider_projects_official_codex_quota_and_independent_buc
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
@@ -953,6 +1003,7 @@ async fn openai_admin_keeps_confirmed_exhaustion_separate_from_raw_usage_display
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
@@ -1911,6 +1962,29 @@ mod errors {
assert_eq!(store.account("acct_refresh_error").unwrap(), before);
}
#[tokio::test]
async fn manual_token_refresh_prepares_to_preserve_concurrent_profile_changes() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token": "new-access-token", "expires_in": 3600
})))
.expect(1)
.mount(&server)
.await;
let (bundle, store, _config) = refresh_fixture(&server, false, true).await;
let before = store.account("acct_refresh_error").unwrap();
let prepared = bundle
.admin_provider()
.prepare_refresh(PrepareCredentialRefresh {
account: account_record(&before),
})
.await
.unwrap();
assert!(prepared.facts().preserve_profile);
}
async fn refresh_fixture(
server: &MockServer,
busy: bool,
backend/crates/providers/openai/tests/credential/contract.rs
+1
View File
@@ -1235,6 +1235,7 @@ fn native_continuation_surfaces_the_original_accounts_quota_signal_to_the_coordi
});
let observed_at = SystemTime::now();
let outcome = block_on(store.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: original.id().clone(),
expected_revision: original.revision(),
quota: OpaqueProviderData::new(quota.as_object().expect("quota object").clone()),
backend/crates/providers/openai/tests/credential/quota/mod.rs
+182
View File
@@ -387,6 +387,188 @@ async fn websocket_additional_rate_limit_resolves_by_name_without_touching_core(
assert_eq!(by_limit_id.len(), 2);
}
#[tokio::test]
async fn quota_refresh_synchronizes_plan_changes_without_losing_subtypes() {
for (current_plan, observed_plan, expected) in [
("plus", Some("pro"), "pro"),
("pro", Some("plus"), "plus"),
("plus", Some("free"), "free"),
("unknown", Some("pro"), "pro"),
("plus", Some(" PRO "), "pro"),
("pro", None, "pro"),
("pro", Some("unknown"), "pro"),
("pro", Some(" "), "pro"),
(
"self_serve_business_prolite",
Some("team"),
"self_serve_business_prolite",
),
(
"self_serve_business_usage_based",
Some("team"),
"self_serve_business_usage_based",
),
(
"enterprise_cbp_usage_based",
Some("business"),
"enterprise_cbp_usage_based",
),
("edu_plus", Some("edu"), "edu_plus"),
("self_serve_business_prolite", Some("pro"), "pro"),
(
"team",
Some("self_serve_business_prolite"),
"self_serve_business_prolite",
),
("plus", Some("future_plan"), "future_plan"),
] {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_plan_refresh";
let mut verified_account = profile("chatgpt-plan-refresh");
verified_account.plan_type = Some(current_plan.to_owned());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: account_id.to_owned(),
secret: secret("plan-refresh-token"),
verified_account,
next_refresh_at: None,
enabled: true,
})
.await;
let before = store.account(account_id).unwrap();
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"plan_type": observed_plan,
"rate_limit": { "allowed": true, "primary_window": {"used_percent": 12} }
})))
.expect(1)
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
);
service
.refresh_account(before.id())
.await
.expect("refresh plan and quota");
let after = store.account(account_id).unwrap();
assert_eq!(
after.plan_type(),
Some(expected),
"{current_plan} -> {observed_plan:?}"
);
assert_eq!(after.revision(), before.revision());
assert_eq!(after.email(), before.email());
assert_eq!(after.upstream_user_id(), before.upstream_user_id());
assert_eq!(after.upstream_account_id(), before.upstream_account_id());
}
}
#[tokio::test]
async fn passive_plan_observations_update_account_without_replaying_stale_plan() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_passive_plan").await;
let original = store.account("acct_passive_plan").unwrap();
let service = quota_service(&store);
service
.synchronize_passive_headers(
&original,
&[
("x-codex-plan-type".to_owned(), "plus".to_owned()),
("x-codex-primary-used-percent".to_owned(), "5".to_owned()),
],
)
.await
.unwrap();
assert_eq!(
store.account("acct_passive_plan").unwrap().plan_type(),
Some("plus")
);
let upgrade = parse_rate_limits_event(&json!({
"type": "codex.rate_limits",
"plan_type": "pro",
"rate_limits": {"primary": {"used_percent": 5, "window_minutes": 300, "reset_at": 1900000000}}
})).unwrap();
service
.synchronize_passive_rate_limits(&original, &[upgrade])
.await
.unwrap();
assert_eq!(
store.account("acct_passive_plan").unwrap().plan_type(),
Some("pro")
);
service
.synchronize_passive_headers(
&original,
&[("x-codex-primary-used-percent".to_owned(), "6".to_owned())],
)
.await
.unwrap();
assert_eq!(
store.account("acct_passive_plan").unwrap().plan_type(),
Some("pro")
);
assert_eq!(
store.account("acct_passive_plan").unwrap().revision(),
original.revision()
);
let state = exhausted_quota(None);
persist_quota_state(&store, &original, state).await;
// 仅有套餐的新观察不能清除既有额度耗尽结论。
service
.synchronize_passive_headers(
&original,
&[("x-codex-plan-type".to_owned(), "plus".to_owned())],
)
.await
.unwrap();
let changed = store.account("acct_passive_plan").unwrap();
assert_eq!(changed.plan_type(), Some("plus"));
assert_eq!(changed.quota(), state);
// 后台 RT 刷新拿着旧账号快照返回时，仍要提交新 token 并保留刚观测的套餐。
store
.repository()
.rotate_refreshed_oauth_secret(
&original,
secret("rotated-after-plan-update"),
original.access_token_expires_at(),
None,
)
.await
.expect("quota changes must not invalidate credential rotation");
assert_eq!(
store.account("acct_passive_plan").unwrap().plan_type(),
Some("plus")
);
}
#[tokio::test]
async fn initial_quota_worker_synchronizes_account_plan() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_worker_plan").await;
let server = MockServer::start().await;
Mock::given(method("GET")).and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"plan_type": "plus", "rate_limit": {"allowed": true, "primary_window": {"used_percent": 10}}
})))
.expect(1).mount(&server).await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
);
assert_eq!(service.synchronize().await.unwrap().updated, 1);
assert_eq!(
store.account("acct_worker_plan").unwrap().plan_type(),
Some("plus")
);
}
#[tokio::test]
async fn quota_refresh_persists_explicit_provider_denial() {
let store = Arc::new(MemoryAccountStore::default());
backend/crates/providers/openai/tests/credential/quota/scheduling.rs
+1
View File
@@ -25,6 +25,7 @@ async fn persist_snapshot(
let account = store.account(id).expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
backend/crates/providers/openai/tests/credential/quota/slots.rs
+1
View File
@@ -495,6 +495,7 @@ async fn stored_weekly_exhaustion_ignores_the_old_primary_only_account_reset() {
old["rate_limit"]["allowed"] = json!(false);
store
.compare_and_swap_quota(gateway_core::account::QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: gateway_core::account::OpaqueProviderData::new(
backend/crates/providers/openai/tests/credential/quota/snapshot.rs
+6
View File
@@ -187,6 +187,7 @@ async fn persisted_provider_quota_projects_dynamic_windows_without_network_io()
});
let outcome = store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
@@ -232,6 +233,7 @@ async fn expired_window_remains_the_last_observation_until_refresh() {
});
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
@@ -270,6 +272,7 @@ async fn persisted_access_fact_survives_limit_without_percent_or_reset() {
});
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
@@ -319,6 +322,7 @@ async fn persisted_codex_alias_keeps_the_top_level_rate_limit_canonical() {
});
let outcome = store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
@@ -376,6 +380,7 @@ async fn persisted_quota_orders_core_window_before_additional_limit() {
});
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
@@ -440,6 +445,7 @@ async fn code_review_limit_projects_as_one_snapshot_per_limit_id() {
});
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(raw.as_object().expect("quota object").clone()),
backend/crates/providers/openai/tests/provider/contract.rs
+3
View File
@@ -4507,6 +4507,7 @@ async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_st
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: first_account.id().clone(),
expected_revision: first_account.revision(),
quota: OpaqueProviderData::new(Map::new()),
@@ -5273,6 +5274,7 @@ async fn official_usage_limit_failure_persists_fact_without_fabricating_usage()
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
@@ -6472,6 +6474,7 @@ async fn quota_limited_account_diagnostic_uses_upstream() {
let observed_at = SystemTime::now();
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
backend/crates/providers/openai/tests/support.rs
+15 -5
View File
@@ -278,6 +278,7 @@ impl ProviderAccountStore for MemoryAccountStore {
account_id,
expected_revision,
profile,
preserve_profile,
credential,
has_refresh_token,
access_token_expires_at,
@@ -320,7 +321,11 @@ impl ProviderAccountStore for MemoryAccountStore {
access_token_expires_at,
has_refresh_token,
next_refresh_at,
profile: Some((profile.name, profile.email, profile.plan_type)),
profile: (!preserve_profile).then_some((
profile.name,
profile.email,
profile.plan_type,
)),
},
);
stored.credential = credential;
@@ -362,11 +367,16 @@ impl ProviderAccountStore for MemoryAccountStore {
return Ok(QuotaWriteOutcome::Conflict);
}
let quota = observation.state;
let mut replacement = AccountRebuild::preserving(&stored.account).with_quota(quota);
if let Some(plan_type) = &observation.plan_type {
replacement.profile = Some((
stored.account.name().to_owned(),
stored.account.email().map(str::to_owned),
Some(plan_type.clone()),
));
}
stored.quota = Some(observation);
stored.account = rebuild_account(
&stored.account,
AccountRebuild::preserving(&stored.account).with_quota(quota),
);
stored.account = rebuild_account(&stored.account, replacement);
Ok(QuotaWriteOutcome::Updated)
}
backend/crates/providers/xai/src/admin.rs
+2
View File
@@ -957,6 +957,7 @@ fn prepared_rotation(
account_id,
expected_revision,
profile: credential_profile,
preserve_profile,
credential,
has_refresh_token,
access_token_expires_at,
@@ -976,6 +977,7 @@ fn prepared_rotation(
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
preserve_profile,
provider_material: ProviderDocument::new(OpaqueProviderData::new(
credential.into_inner(),
)),
backend/crates/providers/xai/src/credential/repository.rs
+1
View File
@@ -545,6 +545,7 @@ impl GrokCredentialRepository {
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id,
expected_revision,
quota: OpaqueProviderData::new(document),
backend/crates/providers/xai/tests/admin.rs
+1
View File
@@ -433,6 +433,7 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export()
assert!(quota.windows.is_empty());
store
.compare_and_swap_quota(gateway_core::account::QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
backend/crates/providers/xai/tests/credential/catalog.rs
+1
View File
@@ -1025,6 +1025,7 @@ async fn quota_read_rejects_corrupt_provider_document() {
document.insert("config".to_owned(), serde_json::json!([]));
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account_id("corrupt"),
expected_revision: CredentialRevision::new(1).expect("revision"),
quota: OpaqueProviderData::new(document),
backend/crates/providers/xai/tests/credential/selector.rs
+1
View File
@@ -267,6 +267,7 @@ impl SelectorFixture {
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: id.clone(),
expected_revision: CredentialRevision::new(1).expect("revision"),
quota: OpaqueProviderData::new(document.as_object().expect("quota object").clone()),
backend/crates/providers/xai/tests/support.rs
+16 -3
View File
@@ -407,6 +407,7 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
account_id,
expected_revision,
profile,
preserve_profile,
credential,
has_refresh_token,
access_token_expires_at,
@@ -444,9 +445,21 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
enabled: stored.account.enabled(),
has_refresh_token,
next_refresh_at,
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
name: if preserve_profile {
stored.account.name().to_owned()
} else {
profile.name
},
email: if preserve_profile {
stored.account.email().map(str::to_owned)
} else {
profile.email
},
plan_type: if preserve_profile {
stored.account.plan_type().map(str::to_owned)
} else {
profile.plan_type
},
},
);
stored.credential = credential;
docs/api.md
+2
View File
@@ -352,6 +352,8 @@ Token 明细、费用明细、用时/首字与状态。Token 和费用复用现
账号视图和 Dashboard 账号概览中的 `planType` 保留原始套餐值；`planTypeDisplay` 由后端先按 Provider 解析名称，
再统一为大驼峰格式，前端直接展示该字段，例如 `Free`、`SuperGrokPro`、`EduPlus`。
OpenAI 的 `self_serve_business_prolite` 等 Team 套餐显示为 `Business`；新套餐也使用相同格式。
OpenAI 主动额度刷新和正常响应携带的明确套餐会同步到账号，支持升级与降级；
空值或 `unknown` 不覆盖已有套餐，同族泛化值（如 `team`）保留已知的具体套餐子类型。
账号套餐为空或 `unknown` 时，后端优先用已保存的上游额度响应
中的明确套餐值补全 `planType` 和 `planTypeDisplay`；两处均无套餐信息时才显示“未知套餐”。
docs/architecture.md
+2
View File
@@ -423,6 +423,8 @@ credential 与 quota 是两组独立事实：credential refresh 不等于 quota
手工和后台 credential refresh 仍不隐式刷新 quota。
- quota refresh、正常推理返回的 rate-limit headers 和后台健康任务汇入同一额度事实；套餐只用于展示与
目录 cache 隔离，不创建套餐专属状态机。
OpenAI Provider 将其中明确的套餐变更与额度原子提交，共用凭据版本和观察时间保护；
空值及 `unknown` 不覆盖套餐，同族泛化值保留具体子类型。Token 刷新保留提交时的账号资料。
OpenAI 订阅周期属于按需个人信息，不是额度事实。Admin 账号用例通过现有 Provider 管理端口并发读取
个人资料统计与订阅，汇聚为一次只读响应；Provider 继续拥有各自的认证、出站代理与上游协议处理。