mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-09 16:06:27 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7542ed365e | ||
|
|
bf4da7a3e7 | ||
|
|
0fe354b083 |
@@ -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<ProviderAccountIdentity>,
|
||||
pub name: String,
|
||||
pub email: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
|
||||
@@ -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::<Vec<_>>()
|
||||
.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());
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ProviderAccountIdentity {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ProviderAccountIdentity")
|
||||
.field("upstream_user_id", &"<redacted>")
|
||||
.field(
|
||||
"upstream_account_id",
|
||||
&self.upstream_account_id.as_ref().map(|_| "<redacted>"),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderAccountIdentity {
|
||||
#[must_use]
|
||||
pub const fn new(upstream_user_id: String, upstream_account_id: Option<String>) -> 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);
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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<ProviderAccountIdentity>,
|
||||
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<Revision> {
|
||||
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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<String>, Option<chrono::DateTime<Utc>>) = sqlx::query_as(
|
||||
"select availability, availability_reason, cooldown_until
|
||||
let restored: (
|
||||
String,
|
||||
Option<String>,
|
||||
Option<chrono::DateTime<Utc>>,
|
||||
String,
|
||||
Option<String>,
|
||||
) = 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<String>) = 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<String>, 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(),
|
||||
|
||||
@@ -50,9 +50,13 @@ impl ProviderAccountStore for CooldownTestAccountStore {
|
||||
|
||||
async fn get_account(
|
||||
&self,
|
||||
_account: &ProviderAccountId,
|
||||
account: &ProviderAccountId,
|
||||
) -> Result<Option<ProviderAccount>, 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<Vec<ProviderAccount>, CoreStoreError> {
|
||||
@@ -190,10 +194,26 @@ fn composition_fixture(
|
||||
CooldownCachingProviderAccountStore,
|
||||
Arc<CooldownTestCache>,
|
||||
Arc<Mutex<Vec<&'static str>>>,
|
||||
) {
|
||||
composition_fixture_with_accounts(
|
||||
vec![ready_test_account("acct_cooldown_composition")],
|
||||
fail_state_write,
|
||||
fail_cache,
|
||||
)
|
||||
}
|
||||
|
||||
fn composition_fixture_with_accounts(
|
||||
accounts: Vec<ProviderAccount>,
|
||||
fail_state_write: bool,
|
||||
fail_cache: bool,
|
||||
) -> (
|
||||
CooldownCachingProviderAccountStore,
|
||||
Arc<CooldownTestCache>,
|
||||
Arc<Mutex<Vec<&'static str>>>,
|
||||
) {
|
||||
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);
|
||||
|
||||
@@ -586,7 +586,7 @@ fn prepared_rotation(
|
||||
prepared: crate::credential::PreparedCodexCredentialRotation,
|
||||
provider_kind: ProviderKind,
|
||||
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
|
||||
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,
|
||||
|
||||
@@ -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<ProviderAccountIdentity>,
|
||||
refresh_guards: Option<ProviderRefreshGuards>,
|
||||
}
|
||||
|
||||
@@ -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(|_| "<held>"),
|
||||
@@ -316,11 +319,13 @@ impl PreparedCodexCredentialRotation {
|
||||
) -> (
|
||||
ProviderAccountUpdate,
|
||||
CredentialCasUpdate,
|
||||
Option<ProviderAccountIdentity>,
|
||||
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<PreparedCodexCredentialRotation, CodexCredentialAdminError> {
|
||||
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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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<CodexIdentityExpectation, CodexOAuthAdminError> {
|
||||
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,
|
||||
|
||||
@@ -1392,10 +1392,16 @@ fn quota_success_availability(
|
||||
exhausted: bool,
|
||||
now: SystemTime,
|
||||
) -> Option<AccountAvailability> {
|
||||
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
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -237,6 +237,9 @@ impl CodexCredentialRepository {
|
||||
cooldown_until: Option<SystemTime>,
|
||||
observed_at: SystemTime,
|
||||
) -> Result<(), CredentialRepositoryError> {
|
||||
if !account.enabled() {
|
||||
return Ok(());
|
||||
}
|
||||
self.store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: account.id().clone(),
|
||||
|
||||
@@ -68,6 +68,7 @@ use crate::transport::request::{
|
||||
scope_request_to_account,
|
||||
};
|
||||
use crate::transport::session::CodexSessionIdentity;
|
||||
use crate::transport::usage::normalize_service_tier;
|
||||
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
|
||||
use crate::transport::{
|
||||
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
|
||||
@@ -822,8 +823,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
|
||||
// WS 帧由 reducer 以 encode_sse_event(&event, raw) 逐字节内嵌上游原始 JSON
|
||||
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
|
||||
let mut decoder =
|
||||
CodexCanonicalDecoder::new(upstream_model.as_str()).with_raw_sse_passthrough();
|
||||
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str())
|
||||
.with_requested_service_tier(request.service_tier())
|
||||
.with_raw_sse_passthrough();
|
||||
let mut pre_commit_events = PreCommitClientEvents::new(stream_commit_policy);
|
||||
loop {
|
||||
let Some(stream_deadline) = remaining(context.deadline()) else {
|
||||
@@ -895,7 +897,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
};
|
||||
pre_commit_events.observe_chunk(chunk_len);
|
||||
let service_tier_changed = observation_state
|
||||
.observe_service_tier(decoder.response_service_tier());
|
||||
.observe_upstream_service_tier(decoder.response_service_tier());
|
||||
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
|
||||
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
|
||||
(
|
||||
@@ -990,8 +992,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
)
|
||||
});
|
||||
let timing_signals = decoder.take_timing_signals();
|
||||
let service_tier_changed =
|
||||
observation_state.observe_service_tier(decoder.response_service_tier());
|
||||
let service_tier_changed = observation_state
|
||||
.observe_upstream_service_tier(decoder.response_service_tier());
|
||||
let timing_changed = observation_state
|
||||
.observe_timing_signals(timing_signals, context.timing_started_at());
|
||||
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
|
||||
@@ -1070,7 +1072,8 @@ struct OpenAiResponseObservationState {
|
||||
requested_model: String,
|
||||
stream: bool,
|
||||
compact: bool,
|
||||
service_tier: Option<String>,
|
||||
requested_service_tier: Option<String>,
|
||||
upstream_service_tier: Option<String>,
|
||||
rate_limit_headers: Vec<(String, String)>,
|
||||
account_id: String,
|
||||
attempt_index: u32,
|
||||
@@ -1102,8 +1105,8 @@ impl OpenAiResponseObservationState {
|
||||
requested_model: request.model().to_owned(),
|
||||
stream: request.stream(),
|
||||
compact: semantics.compact,
|
||||
// 请求档位保留在 requestSummary;一等观测只接受响应确认的实际档位。
|
||||
service_tier: None,
|
||||
requested_service_tier: normalize_service_tier(request.service_tier()),
|
||||
upstream_service_tier: None,
|
||||
rate_limit_headers: selected_observation_headers(&response.rate_limit_headers),
|
||||
account_id: account_id.to_owned(),
|
||||
attempt_index,
|
||||
@@ -1127,7 +1130,7 @@ impl OpenAiResponseObservationState {
|
||||
if let Some(metadata) = self.provider_metadata() {
|
||||
observation = observation.with_provider_metadata(metadata);
|
||||
}
|
||||
if let Some(service_tier) = self.service_tier.as_deref() {
|
||||
if let Some(service_tier) = self.effective_service_tier() {
|
||||
observation = observation.with_service_tier_if_valid(service_tier.to_owned());
|
||||
}
|
||||
Some(observation)
|
||||
@@ -1158,17 +1161,23 @@ impl OpenAiResponseObservationState {
|
||||
changed
|
||||
}
|
||||
|
||||
fn observe_service_tier(&mut self, service_tier: Option<&str>) -> bool {
|
||||
let Some(service_tier) = service_tier else {
|
||||
fn observe_upstream_service_tier(&mut self, service_tier: Option<&str>) -> bool {
|
||||
let Some(service_tier) = normalize_service_tier(service_tier) else {
|
||||
return false;
|
||||
};
|
||||
if self.service_tier.as_deref() == Some(service_tier) {
|
||||
if self.upstream_service_tier.as_deref() == Some(service_tier.as_str()) {
|
||||
return false;
|
||||
}
|
||||
self.service_tier = Some(service_tier.to_owned());
|
||||
self.upstream_service_tier = Some(service_tier);
|
||||
true
|
||||
}
|
||||
|
||||
fn effective_service_tier(&self) -> Option<&str> {
|
||||
self.requested_service_tier
|
||||
.as_deref()
|
||||
.or(self.upstream_service_tier.as_deref())
|
||||
}
|
||||
|
||||
fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool {
|
||||
let mut changed = false;
|
||||
for (name, value) in selected_observation_headers(updates) {
|
||||
@@ -1245,9 +1254,21 @@ impl OpenAiResponseObservationState {
|
||||
)),
|
||||
);
|
||||
metadata.insert("requestSummary".to_owned(), self.request_summary.clone());
|
||||
if let Some(service_tier) = &self.service_tier {
|
||||
if let Some(service_tier) = self.effective_service_tier() {
|
||||
metadata.insert(
|
||||
"serviceTier".to_owned(),
|
||||
Value::String(service_tier.to_owned()),
|
||||
);
|
||||
}
|
||||
if let Some(service_tier) = &self.requested_service_tier {
|
||||
metadata.insert(
|
||||
"requestedServiceTier".to_owned(),
|
||||
Value::String(service_tier.clone()),
|
||||
);
|
||||
}
|
||||
if let Some(service_tier) = &self.upstream_service_tier {
|
||||
metadata.insert(
|
||||
"upstreamServiceTier".to_owned(),
|
||||
Value::String(service_tier.clone()),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -21,10 +21,9 @@ use thiserror::Error;
|
||||
use super::protocol::responses::{
|
||||
ResponseEventSignals, ResponsesSseFailure, response_event_signals,
|
||||
};
|
||||
use super::usage::openai_billing_breakdown;
|
||||
use super::usage::{normalize_service_tier, openai_billing_breakdown};
|
||||
|
||||
const CONTENTS_PER_OUTPUT: u32 = 1_024;
|
||||
const MAX_SERVICE_TIER_BYTES: usize = 64;
|
||||
|
||||
/// 单 attempt 的增量 Responses decoder。
|
||||
///
|
||||
@@ -42,6 +41,7 @@ pub struct CodexCanonicalDecoder {
|
||||
reasoning_output_seen: BTreeSet<u32>,
|
||||
usage_emitted: bool,
|
||||
semantic_output_seen: bool,
|
||||
requested_service_tier: Option<String>,
|
||||
response_service_tier: Option<String>,
|
||||
timing_signals: ResponseEventSignals,
|
||||
raw_sse_passthrough: bool,
|
||||
@@ -137,6 +137,7 @@ impl CodexCanonicalDecoder {
|
||||
reasoning_output_seen: BTreeSet::new(),
|
||||
usage_emitted: false,
|
||||
semantic_output_seen: false,
|
||||
requested_service_tier: None,
|
||||
response_service_tier: None,
|
||||
timing_signals: ResponseEventSignals::default(),
|
||||
raw_sse_passthrough: false,
|
||||
@@ -152,6 +153,13 @@ impl CodexCanonicalDecoder {
|
||||
self
|
||||
}
|
||||
|
||||
/// 请求显式指定的档位优先用于计费;响应档位只在请求未指定时兜底。
|
||||
#[must_use]
|
||||
pub fn with_requested_service_tier(mut self, service_tier: Option<&str>) -> Self {
|
||||
self.requested_service_tier = normalize_service_tier(service_tier);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
|
||||
self.timing_signals = ResponseEventSignals::default();
|
||||
if self.raw_sse_passthrough {
|
||||
@@ -308,19 +316,14 @@ impl CodexCanonicalDecoder {
|
||||
}
|
||||
|
||||
fn observe_response_service_tier(&mut self, value: &Value) {
|
||||
let Some(service_tier) = response_object(value)
|
||||
.and_then(|response| response.get("service_tier"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| {
|
||||
!value.is_empty()
|
||||
&& value.len() <= MAX_SERVICE_TIER_BYTES
|
||||
&& !value.chars().any(char::is_control)
|
||||
})
|
||||
else {
|
||||
let Some(service_tier) = normalize_service_tier(
|
||||
response_object(value)
|
||||
.and_then(|response| response.get("service_tier"))
|
||||
.and_then(Value::as_str),
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
self.response_service_tier = Some(service_tier.to_ascii_lowercase());
|
||||
self.response_service_tier = Some(service_tier);
|
||||
}
|
||||
|
||||
/// 构造下发的 wire event。
|
||||
@@ -836,9 +839,9 @@ impl CodexCanonicalDecoder {
|
||||
.filter(|model| !model.is_empty())
|
||||
.unwrap_or(&self.fallback_model)
|
||||
.to_owned();
|
||||
let service_tier = response
|
||||
.get("service_tier")
|
||||
.and_then(Value::as_str)
|
||||
let service_tier = self
|
||||
.requested_service_tier
|
||||
.as_deref()
|
||||
.or(self.response_service_tier.as_deref());
|
||||
if let Some(breakdown) = usage
|
||||
.filter(|usage| billable_usage_is_complete(response, *usage))
|
||||
|
||||
@@ -298,10 +298,13 @@ fn model_matches_rule(model: &str, rule: &str) -> bool {
|
||||
.is_some_and(|suffix| matches!(suffix.as_bytes().first(), Some(b'-' | b'.' | b':')))
|
||||
}
|
||||
|
||||
fn normalize_service_tier(service_tier: Option<&str>) -> Option<String> {
|
||||
/// 规范化请求或响应携带的服务档位,供观测与计费共用。
|
||||
pub(crate) fn normalize_service_tier(service_tier: Option<&str>) -> Option<String> {
|
||||
service_tier
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| {
|
||||
!value.is_empty() && value.len() <= 64 && !value.chars().any(char::is_control)
|
||||
})
|
||||
.map(str::to_ascii_lowercase)
|
||||
}
|
||||
|
||||
|
||||
@@ -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<CodexIdentityVerification, CodexIdentityVerificationError> {
|
||||
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"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -181,6 +181,14 @@ fn blocked_network_quota_service(store: &Arc<MemoryAccountStore>) -> CodexCreden
|
||||
}
|
||||
|
||||
async fn create_account(store: &Arc<MemoryAccountStore>, account_id: &str) {
|
||||
create_account_with_enabled(store, account_id, true).await;
|
||||
}
|
||||
|
||||
async fn create_account_with_enabled(
|
||||
store: &Arc<MemoryAccountStore>,
|
||||
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<MemoryAccountStore>, 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());
|
||||
|
||||
@@ -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::{
|
||||
@@ -29,7 +29,7 @@ use provider_openai::credential::{
|
||||
use provider_openai::transport::CodexWebSocketPool;
|
||||
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
|
||||
use provider_openai::{CodexProvider, OFFICIAL_CODEX_BASE_URL};
|
||||
use serde_json::{Map, json};
|
||||
use serde_json::{Map, Value, json};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::{TcpListener, TcpStream},
|
||||
@@ -827,13 +827,23 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
|
||||
.expect("prepare provider stream");
|
||||
|
||||
let mut observed_service_tier = None;
|
||||
let mut upstream_service_tier = None;
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.expect("provider event");
|
||||
if let Some(service_tier) = event
|
||||
.response_observation()
|
||||
.and_then(|observation| observation.service_tier())
|
||||
{
|
||||
observed_service_tier = Some(service_tier.to_owned());
|
||||
if let Some(observation) = event.response_observation() {
|
||||
if let Some(service_tier) = observation.service_tier() {
|
||||
observed_service_tier = Some(service_tier.to_owned());
|
||||
}
|
||||
upstream_service_tier = observation
|
||||
.provider_metadata()
|
||||
.and_then(|metadata| serde_json::from_str::<Value>(metadata.as_json()).ok())
|
||||
.and_then(|metadata| {
|
||||
metadata
|
||||
.get("upstreamServiceTier")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned)
|
||||
})
|
||||
.or(upstream_service_tier);
|
||||
}
|
||||
if event
|
||||
.canonical_facts()
|
||||
@@ -846,7 +856,8 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
|
||||
drop(stream);
|
||||
|
||||
assert_eq!(affinity.binding_count(), 1);
|
||||
assert_eq!(observed_service_tier.as_deref(), Some("default"));
|
||||
assert_eq!(observed_service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(upstream_service_tier.as_deref(), Some("default"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1207,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(
|
||||
@@ -1220,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;
|
||||
|
||||
@@ -203,18 +203,40 @@ fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
|
||||
fn decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"default\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
|
||||
);
|
||||
let mut decoder = CodexCanonicalDecoder::new("fallback");
|
||||
let mut decoder =
|
||||
CodexCanonicalDecoder::new("fallback").with_requested_service_tier(Some("priority"));
|
||||
let events = decoder
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical priority response");
|
||||
|
||||
assert_eq!(decoder.response_service_tier(), Some("default"));
|
||||
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
|
||||
event,
|
||||
GatewayEvent::CalculatedCost(cost)
|
||||
if cost.total().amount().scaled() == 6_875_000
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_use_response_service_tier_when_request_omits_it() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
|
||||
);
|
||||
let mut decoder = CodexCanonicalDecoder::new("fallback");
|
||||
let events = decoder
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical response-tier priority response");
|
||||
|
||||
assert_eq!(decoder.response_service_tier(), Some("priority"));
|
||||
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
|
||||
event,
|
||||
@@ -223,6 +245,31 @@ fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
|
||||
let created = websocket_event_to_sse_frame(
|
||||
r#"{"type":"response.created","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","service_tier":"default"}}"#,
|
||||
)
|
||||
.expect("created frame");
|
||||
let completed = websocket_event_to_sse_frame(
|
||||
r#"{"type":"response.completed","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","status":"completed","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25,"cache_write_tokens":0},"total_tokens":110}}}"#,
|
||||
)
|
||||
.expect("completed frame");
|
||||
let mut decoder = CodexCanonicalDecoder::new("fallback")
|
||||
.with_requested_service_tier(Some("priority"))
|
||||
.with_raw_sse_passthrough();
|
||||
let events = decoder
|
||||
.push(format!("{created}{completed}").as_bytes())
|
||||
.expect("canonical priority WebSocket response");
|
||||
|
||||
assert_eq!(decoder.response_service_tier(), Some("default"));
|
||||
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
|
||||
event,
|
||||
GatewayEvent::CalculatedCost(cost)
|
||||
if cost.total().amount().scaled() == 6_875_000
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_normalize_text_usage_and_completion() {
|
||||
let body = concat!(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -1 +1 @@
|
||||
version: "3.0.22"
|
||||
version: "3.0.23"
|
||||
|
||||
Reference in New Issue
Block a user