Compare commits

...
11 Commits
Author SHA1 Message Date
zyycnandCodex 3110b71512 chore: bump version to 3.0.20
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 18:46:20 +08:00
zyycnandCodex eac1b2c39e fix(xai): classify upstream failures by scope
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 18:31:06 +08:00
zyycnandCodex 337bb80230 fix(openai): restore streaming failure fallback
Restore the v2.x stream commit policies so ordinary requests forward the first event immediately while continuation and same-batch failures remain behind the retry boundary. Preserve typed upstream failure batches through Core and both streaming adapters, with a 64 KiB prefetch limit and end-to-end regression coverage.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 17:29:15 +08:00
zyycnandCodex 82f2a0da06 fix(openai): preserve quota exhaustion state
Prefer explicit quota evidence over generic 429 feedback, keep exhausted and terminal accounts from being downgraded by temporary rate limits, and reconcile historical cooldown states from current quota observations.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 14:47:21 +08:00
zyycnandCodex f0d8dad71e chore: bump version to 3.0.19
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 13:05:04 +08:00
zyycnandCodex f5c1d773d7 fix(accounts): use current credential revisions
Load the current credential revision server-side for reauthorization and manual refresh so background token refreshes do not make admin mutations fail stale client fences.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 13:04:12 +08:00
zyycnandCodex 9148e6b11e chore: bump version to 3.0.18
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 09:48:21 +08:00
zyycnandCodex a7d8c3d3df fix(accounts): make fixed diagnostics resilient
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 09:45:35 +08:00
zyycnandCodex 997eb86aa0 fix(openai): refresh expired OAuth access tokens
Force refresh scheduling for OAuth accounts whose access token has already expired, even when a persisted next_refresh_at points to the future.

Co-authored-by: Codex <noreply@openai.com>
2026-07-28 09:41:44 +08:00
zyycnandCodex 3866acc269 chore: bump version to 3.0.17
Co-authored-by: Codex <noreply@openai.com>
2026-07-28 00:06:15 +08:00
zyycnandCodex 38477f4328 fix(openai): normalize OAuth import credentials
Import OAuth accounts from access and refresh tokens only, then verify identity and observe the initial quota through official endpoints.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 23:54:59 +08:00
72 changed files with 6835 additions and 1694 deletions
+2
View File
@@ -1194,6 +1194,7 @@ dependencies = [
"subtle",
"thiserror",
"tokio",
"tracing",
"uuid",
]
@@ -1214,6 +1215,7 @@ dependencies = [
"serde_json",
"thiserror",
"tokio",
"tokio-tungstenite 0.29.0",
"tower",
"tower-http 0.7.0",
"tracing",
+1
View File
@@ -24,6 +24,7 @@ serde = { version = "1.0.228", features = ["derive"] }
secrecy = { version = "0.10.3", features = ["serde"] }
subtle = "2.6.1"
thiserror = "2.0.18"
tracing = "0.1.44"
uuid = { version = "1.23.3", features = ["v7"] }
[dev-dependencies]
@@ -2,6 +2,7 @@
use std::fmt;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
@@ -9,7 +10,7 @@ use gateway_core::{
};
use super::{
MutationActor, MutationContext, PageSize, Revision,
AdminError, MutationActor, MutationContext, PageSize, Revision,
accounts::{AccountAvailability, AccountRecord, AccountStatus, AccountSummary, AccountUsage},
};
@@ -170,13 +171,6 @@ pub struct CredentialImportCommit {
pub prepared: PreparedCredentialImport,
}
/// 可选的重新授权目标。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReauthorizationTarget {
pub account_id: ProviderAccountId,
pub credential_revision: Revision,
}
/// OAuth pending owner 的中立身份;不编码具体 Provider 的 Redis key 或 JSON。
#[derive(Clone, PartialEq, Eq)]
pub enum AuthorizationOwner {
@@ -237,15 +231,13 @@ impl fmt::Debug for AuthorizationOwnerBinding {
}
/// OAuth 完成时应创建新账号还是 CAS 更新既有 credential。
///
/// 重新授权只绑定稳定的账号身份:credential revision 由服务端在临近写入时读取,
/// 长流程期间的后台刷新不得让恢复操作失效。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthorizationMutationTarget {
Create {
name: String,
},
Reauthorize {
account_id: ProviderAccountId,
expected_credential_revision: Revision,
},
Create { name: String },
Reauthorize { account_id: ProviderAccountId },
}
/// 必须完整进入 Provider opaque pending payload、并在 complete 后原样恢复的事务信封。
@@ -291,7 +283,7 @@ impl PendingAuthorizationMutation {
pub struct StartAuthorization {
pub context: MutationContext,
pub name: String,
pub reauthorization: Option<ReauthorizationTarget>,
pub reauthorization: Option<ProviderAccountId>,
}
/// OAuth 流程启动结果。
@@ -310,6 +302,20 @@ pub trait CredentialCommitGuard: Send + 'static {
fn finish(self: Box<Self>);
}
/// OAuth pending claim 在 Store 事务结果确定后执行的显式结算动作。
///
/// OAuth provider 必须在 credential 准备阶段持有 claim:Store 提交成功后消费 flow,提交或
/// 校验失败后释放 flow。异步结算不能依赖 `Drop`,否则请求提前返回会让同一授权流无谓地
/// 进入短暂的不可重试状态。
#[async_trait]
pub trait AuthorizationCommitGuard: Send + 'static {
/// Store 已提交 OAuth credential 后消费对应的一次性 flow。
async fn commit(self: Box<Self>) -> Result<(), AdminError>;
/// Store 未提交 OAuth credential 时释放 claim,允许同一 flow 重试。
async fn abort(self: Box<Self>) -> Result<(), AdminError>;
}
/// OAuth complete 后由 Provider 返回的准备结果;Store 仍是唯一提交者。
pub enum PreparedAuthorizationCredential {
Create(PreparedCredentialCreate),
@@ -320,6 +326,7 @@ pub enum PreparedAuthorizationCredential {
pub struct PreparedAuthorizationCommit {
pub pending: PendingAuthorizationMutation,
pub credential: PreparedAuthorizationCredential,
authorization_guard: Option<Box<dyn AuthorizationCommitGuard>>,
}
impl fmt::Debug for PreparedAuthorizationCommit {
@@ -346,9 +353,37 @@ pub struct AuthorizationCommit {
pub credential: AuthorizationCredentialCommit,
}
/// OAuth 准备结果拆解后的 Store 命令与两个结算 guard。
///
/// Store 成功时先结束 credential guard,再消费 OAuth claim;失败时丢弃 credential guard 并
/// 释放 OAuth claim,使同一授权流可以重试。
pub(crate) struct AuthorizationCommitSettlement {
pub(crate) command: AuthorizationCommit,
pub(crate) credential_guard: Option<Box<dyn CredentialCommitGuard>>,
pub(crate) authorization_guard: Option<Box<dyn AuthorizationCommitGuard>>,
}
impl PreparedAuthorizationCommit {
#[must_use]
pub fn into_commit(self) -> (AuthorizationCommit, Option<Box<dyn CredentialCommitGuard>>) {
pub fn new(
pending: PendingAuthorizationMutation,
credential: PreparedAuthorizationCredential,
) -> Self {
Self {
pending,
credential,
authorization_guard: None,
}
}
/// 让 OAuth claim 覆盖准备完成到 Store 事务结算之间的全部窗口。
#[must_use]
pub fn with_authorization_guard(mut self, guard: Box<dyn AuthorizationCommitGuard>) -> Self {
self.authorization_guard = Some(guard);
self
}
pub(crate) fn into_commit(self) -> AuthorizationCommitSettlement {
let (credential, guard) = match self.credential {
PreparedAuthorizationCredential::Create(credential) => {
(AuthorizationCredentialCommit::Create(credential), None)
@@ -361,13 +396,27 @@ impl PreparedAuthorizationCommit {
)
}
};
(
AuthorizationCommit {
AuthorizationCommitSettlement {
command: AuthorizationCommit {
pending: self.pending,
credential,
},
guard,
)
credential_guard: guard,
authorization_guard: self.authorization_guard,
}
}
pub(crate) async fn abort(self) -> Result<(), AdminError> {
let AuthorizationCommitSettlement {
credential_guard,
authorization_guard,
..
} = self.into_commit();
drop(credential_guard);
if let Some(guard) = authorization_guard {
guard.abort().await?;
}
Ok(())
}
}
@@ -422,14 +471,12 @@ pub struct CredentialDeletionResult {
/// Provider-owned token 轮换命令。
pub struct RotateCredential {
pub mutation: CredentialMutation,
pub expected_credential_revision: Revision,
pub provider_material: ProviderDocument,
}
/// Provider 校验手工轮换材料时所需的非事务输入。
pub struct PrepareCredentialRotation {
pub account: AccountRecord,
pub expected_credential_revision: Revision,
pub provider_material: ProviderDocument,
}
@@ -438,10 +485,6 @@ impl fmt::Debug for PrepareCredentialRotation {
formatter
.debug_struct("PrepareCredentialRotation")
.field("account", &self.account)
.field(
"expected_credential_revision",
&self.expected_credential_revision,
)
.field("provider_material", &self.provider_material)
.finish()
}
@@ -514,10 +557,6 @@ impl fmt::Debug for RotateCredential {
formatter
.debug_struct("RotateCredential")
.field("mutation", &self.mutation)
.field(
"expected_credential_revision",
&self.expected_credential_revision,
)
.field("provider_material", &self.provider_material)
.finish()
}
@@ -90,17 +90,10 @@ async fn pending_authorization(
resource: &'static str,
) -> Result<PendingAuthorizationMutation, AdminError> {
let target = match &command.reauthorization {
Some(target) => {
let details =
required_credential(accounts, provider_kind, &target.account_id, resource).await?;
if details.credential.credential_revision != target.credential_revision {
return Err(AdminError::conflict(format!(
"{resource} reauthorization target is stale"
)));
}
Some(account_id) => {
required_credential(accounts, provider_kind, account_id, resource).await?;
AuthorizationMutationTarget::Reauthorize {
account_id: target.account_id.clone(),
expected_credential_revision: target.credential_revision,
account_id: account_id.clone(),
}
}
None => AuthorizationMutationTarget::Create {
@@ -142,7 +135,6 @@ fn validate_prepared_rotation(
let facts = prepared.facts();
if facts.account_id.as_str() != account.id.as_str()
|| facts.provider_kind != account.provider_kind
|| facts.expected_credential_revision != account.credential_revision
{
return Err(AdminError::conflict(format!(
"{resource} prepared facts do not match the current credential"
@@ -151,18 +143,27 @@ fn validate_prepared_rotation(
Ok(())
}
fn validate_authorization_commit(
async fn validate_authorization_commit(
provider_kind: &ProviderKind,
context: &MutationContext,
prepared: &PreparedAuthorizationCommit,
prepared: PreparedAuthorizationCommit,
resource: &'static str,
) -> Result<(), AdminError> {
) -> Result<PreparedAuthorizationCommit, AdminError> {
if prepared.pending.provider_kind() != provider_kind
|| !prepared.pending.owner_binding().matches_context(context)
{
return Err(AdminError::conflict(format!(
let validation_error = AdminError::conflict(format!(
"{resource} pending authorization binding is invalid"
)));
));
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
let matches_target = match (prepared.pending.target(), &prepared.credential) {
(
@@ -170,25 +171,29 @@ fn validate_authorization_commit(
PreparedAuthorizationCredential::Create(credential),
) => credential.provider_kind == *provider_kind,
(
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
},
AuthorizationMutationTarget::Reauthorize { account_id },
PreparedAuthorizationCredential::Reauthorize(credential),
) => {
let facts = credential.facts();
facts.provider_kind == *provider_kind
&& facts.account_id == *account_id
&& facts.expected_credential_revision == *expected_credential_revision
facts.provider_kind == *provider_kind && facts.account_id == *account_id
}
_ => false,
};
if !matches_target {
return Err(AdminError::conflict(format!(
let validation_error = AdminError::conflict(format!(
"{resource} prepared credential does not match its pending target"
)));
));
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
Ok(())
Ok(prepared)
}
async fn commit_authorization(
@@ -197,16 +202,34 @@ async fn commit_authorization(
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (command, guard) = prepared.into_commit();
let crate::model::provider_credentials::AuthorizationCommitSettlement {
command,
credential_guard,
authorization_guard,
} = prepared.into_commit();
match accounts.commit_authorization(command, context).await {
Ok(result) => {
if let Some(guard) = guard {
if let Some(guard) = credential_guard {
guard.finish();
}
if let Some(guard) = authorization_guard {
guard.commit().await?;
}
Ok(result)
}
Err(error) => {
drop(guard);
drop(credential_guard);
if let Some(guard) = authorization_guard
&& let Err(settlement_error) = guard.abort().await
{
tracing::warn!(
resource,
store_error = %error,
settlement_error = %settlement_error,
"OAuth authorization claim release failed after Store commit failure"
);
return Err(settlement_error);
}
Err(map_store_error(error, resource))
}
}
@@ -13,7 +13,8 @@ use crate::{
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutation,
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
PrepareCredentialRotation, ProviderQuotaRequest, RotateCredential, StartAuthorization,
PrepareCredentialRotation, PreparedAuthorizationCredential, ProviderQuotaRequest,
RotateCredential, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -83,6 +84,27 @@ impl DefaultOpenAiService {
snapshot,
}
}
async fn observe_initial_quotas(&self, account_ids: &[ProviderAccountId], request_id: &str) {
for account_id in account_ids {
if let Err(error) = self
.provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await
{
tracing::warn!(
request_id,
account_id = %account_id.as_str(),
quota_error = ?error.kind(),
"OpenAI initial quota observation failed"
);
}
}
}
}
#[async_trait]
@@ -130,17 +152,9 @@ impl OpenAiService for DefaultOpenAiService {
.await
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
// 导入已经提交,额度属于可重建观察事实;单个账号查询失败由周期任务重试。
for account_id in &result.credential_ids {
let _ = self
.provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await;
}
// 导入完成后立即做一次观察;失败不回滚已经提交的 credential。
self.observe_initial_quotas(&result.credential_ids, &context.request_id)
.await;
Ok(result)
}
@@ -171,12 +185,17 @@ impl OpenAiService for DefaultOpenAiService {
.complete_authorization(command)
.await
.map_err(|error| map_provider_error(error, "OpenAI authorization"))?;
validate_authorization_commit(
let prepared = validate_authorization_commit(
self.provider.provider_kind(),
&context,
&prepared,
prepared,
"OpenAI authorization",
)?;
)
.await?;
let creates_account = matches!(
&prepared.credential,
PreparedAuthorizationCredential::Create(_)
);
let result = commit_authorization(
self.accounts.as_ref(),
prepared,
@@ -185,6 +204,13 @@ impl OpenAiService for DefaultOpenAiService {
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
if creates_account {
self.observe_initial_quotas(
std::slice::from_ref(&result.account_id),
&context.request_id,
)
.await;
}
Ok(result)
}
@@ -201,17 +227,11 @@ impl OpenAiService for DefaultOpenAiService {
"OpenAI credential rotation",
)
.await?;
if details.credential.credential_revision != command.expected_credential_revision {
return Err(AdminError::conflict(
"OpenAI credential rotation revision is stale",
));
}
let account = details.credential;
let prepared = self
.provider
.prepare_rotation(PrepareCredentialRotation {
account: account.clone(),
expected_credential_revision: command.expected_credential_revision,
provider_material: command.provider_material,
})
.await
@@ -148,12 +148,13 @@ impl XaiService for DefaultXaiService {
.complete_authorization(command)
.await
.map_err(|error| map_provider_error(error, "xAI authorization"))?;
validate_authorization_commit(
let prepared = validate_authorization_commit(
self.provider.provider_kind(),
&context,
&prepared,
prepared,
"xAI authorization",
)?;
)
.await?;
let result = commit_authorization(
self.accounts.as_ref(),
prepared,
@@ -16,7 +16,7 @@ use gateway_core::{
use gateway_admin::{
AdminServices,
model::{
MutationContext, Revision,
AdminError, MutationContext, Revision,
accounts::{
AccountAvailability, AccountConnectionTestEvent, AccountListQuery, AccountPage,
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
@@ -24,12 +24,12 @@ use gateway_admin::{
},
observability::TimeRange,
provider_credentials::{
AuthorizationCommit, AuthorizationCredentialCommit, AuthorizationMutationTarget,
AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit,
AuthorizationCommit, AuthorizationCommitGuard, AuthorizationCredentialCommit,
AuthorizationMutationTarget, AuthorizationStarted, CompleteAuthorization,
CredentialCommitGuard, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
@@ -58,11 +58,13 @@ pub(super) struct FakeProviderAdmin {
events: EventLog,
failure: Mutex<Option<ProviderAdminErrorKind>>,
quota_failure: Mutex<Option<ProviderAdminErrorKind>>,
pending: Mutex<Option<PendingAuthorizationMutation>>,
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
retry_authorization_after_abort: Mutex<bool>,
export_inputs: Mutex<Vec<ProviderExportCredentialInput>>,
import_account_ids: Mutex<Vec<String>>,
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
quota: Mutex<ProviderQuota>,
current_credential_revision: Mutex<Revision>,
}
impl FakeProviderAdmin {
@@ -72,11 +74,13 @@ impl FakeProviderAdmin {
events,
failure: Mutex::new(None),
quota_failure: Mutex::new(None),
pending: Mutex::new(None),
pending: Arc::new(Mutex::new(None)),
retry_authorization_after_abort: Mutex::new(false),
export_inputs: Mutex::new(Vec::new()),
import_account_ids: Mutex::new(vec!["acct_prepared".to_owned()]),
quota_requests: Mutex::new(Vec::new()),
quota: Mutex::new(empty_quota()),
current_credential_revision: Mutex::new(revision(1)),
})
}
@@ -113,6 +117,20 @@ impl FakeProviderAdmin {
self.pending.lock().expect("pending authorization").clone()
}
pub(super) fn retry_authorization_after_abort(&self) {
*self
.retry_authorization_after_abort
.lock()
.expect("authorization retry") = true;
}
pub(super) fn set_current_credential_revision(&self, revision: Revision) {
*self
.current_credential_revision
.lock()
.expect("current credential revision") = revision;
}
fn export_inputs(&self) -> Vec<ProviderExportCredentialInput> {
self.export_inputs.lock().expect("export inputs").clone()
}
@@ -133,7 +151,10 @@ impl FakeProviderAdmin {
PreparedCredentialRotationFacts {
account_id: ProviderAccountId::new(account.id.clone()).expect("account ID"),
provider_kind: account.provider_kind.clone(),
expected_credential_revision: account.credential_revision,
expected_credential_revision: *self
.current_credential_revision
.lock()
.expect("current credential revision"),
name: account.name.clone(),
email: account.email.clone(),
plan_type: account.plan_type.clone(),
@@ -246,23 +267,31 @@ impl ProviderAdmin for FakeProviderAdmin {
if !pending.owner_binding().matches_context(&command.context) {
return Err(ProviderAdminError::new(ProviderAdminErrorKind::NotFound));
}
let retry_pending = (*self
.retry_authorization_after_abort
.lock()
.expect("authorization retry"))
.then(|| pending.clone());
let credential = match pending.target() {
AuthorizationMutationTarget::Create { name } => {
PreparedAuthorizationCredential::Create(prepared_create(self.kind.clone(), name))
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
AuthorizationMutationTarget::Reauthorize { account_id } => {
let mut account = account_record(self.kind.as_str());
account.id = account_id.as_str().to_owned();
account.credential_revision = *expected_credential_revision;
PreparedAuthorizationCredential::Reauthorize(self.prepared_rotation(&account))
}
};
Ok(PreparedAuthorizationCommit {
pending,
credential,
let prepared = PreparedAuthorizationCommit::new(pending, credential);
Ok(match retry_pending {
Some(pending) => {
prepared.with_authorization_guard(Box::new(RetryableAuthorizationGuard::new(
self.events.clone(),
Arc::clone(&self.pending),
pending,
)))
}
None => prepared,
})
}
@@ -664,6 +693,51 @@ impl Drop for RecordingGuard {
}
}
struct RetryableAuthorizationGuard {
events: EventLog,
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
authorization: PendingAuthorizationMutation,
}
impl RetryableAuthorizationGuard {
fn new(
events: EventLog,
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
authorization: PendingAuthorizationMutation,
) -> Self {
Self {
events,
pending,
authorization,
}
}
}
#[async_trait]
impl AuthorizationCommitGuard for RetryableAuthorizationGuard {
async fn commit(self: Box<Self>) -> Result<(), AdminError> {
self.events
.lock()
.expect("authorization guard events")
.push("authorization_guard.commit");
Ok(())
}
async fn abort(self: Box<Self>) -> Result<(), AdminError> {
let Self {
events,
pending,
authorization,
} = *self;
*pending.lock().expect("pending authorization") = Some(authorization);
events
.lock()
.expect("authorization guard events")
.push("authorization_guard.abort");
Ok(())
}
}
#[test]
fn provider_registry_should_resolve_custom_kind_without_central_match() {
let provider = FakeProviderAdmin::new("custom-provider", events());
@@ -6,7 +6,7 @@ use gateway_admin::{
AdminServices,
model::provider_credentials::{
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, CredentialMutation,
ImportCredentials, ProviderQuotaRequest, ReauthorizationTarget, StartAuthorization,
ImportCredentials, ProviderQuotaRequest, StartAuthorization,
},
ports::provider::ProviderAdminErrorKind,
};
@@ -203,6 +203,133 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
assert_eq!(provider.quota_requests().len(), 1);
}
#[tokio::test]
async fn openai_authorization_create_should_observe_initial_quota() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider.clone(), store.clone()).await;
services
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-start-openai-create"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
services
.openai()
.complete_authorization(CompleteAuthorization {
context: context("oauth-start-openai-create"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
})
.await
.expect("complete authorization");
assert_eq!(
recorded(&events),
[
"provider.start_authorization",
"provider.complete_authorization",
"store.commit_authorization",
"provider.quota",
]
);
assert_eq!(
provider.quota_requests(),
[ProviderQuotaRequest {
account_id: ProviderAccountId::new("acct_prepared").expect("account ID"),
refresh: true,
rolling_usage: None,
}]
);
assert_eq!(store.audit_requests(), ["oauth-start-openai-create"]);
}
#[tokio::test]
async fn openai_authorization_create_should_remain_successful_when_initial_quota_fails() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
provider.fail_next_quota(ProviderAdminErrorKind::Unavailable);
let store = FakeAccountStore::new("openai", events);
let services = service(provider.clone(), store).await;
services
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-start-openai-quota-failure"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
let result = services
.openai()
.complete_authorization(CompleteAuthorization {
context: context("oauth-start-openai-quota-failure"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
})
.await
.expect("quota observation must not fail authorization");
assert_eq!(result.account_id.as_str(), "acct_prepared");
assert_eq!(provider.quota_requests().len(), 1);
}
#[tokio::test]
async fn openai_authorization_store_failure_should_release_claim_for_retry() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
provider.retry_authorization_after_abort();
let store = FakeAccountStore::new("openai", events.clone());
store.fail_next_commit();
let services = service(provider.clone(), store).await;
let command = || CompleteAuthorization {
context: context("oauth-store-retry"),
flow_id: "flow-test".to_owned(),
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
};
services
.openai()
.start_authorization(StartAuthorization {
context: context("oauth-store-retry"),
name: "new OpenAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
services
.openai()
.complete_authorization(command())
.await
.expect_err("Store commit must fail");
services
.openai()
.complete_authorization(command())
.await
.expect("released claim must permit retry");
assert_eq!(
recorded(&events),
[
"provider.start_authorization",
"provider.complete_authorization",
"store.commit_authorization",
"authorization_guard.abort",
"provider.complete_authorization",
"store.commit_authorization",
"authorization_guard.commit",
"provider.quota",
]
);
assert!(provider.pending().is_none());
}
#[tokio::test]
async fn openai_import_provider_error_should_not_touch_store_transaction() {
let events = events();
@@ -224,7 +351,7 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
}
#[tokio::test]
async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_through_commit() {
async fn openai_reauthorization_should_commit_after_credential_revision_advances() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
@@ -234,10 +361,7 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
.start_authorization(StartAuthorization {
context: context("oauth-start-openai"),
name: "reauthorize".to_owned(),
reauthorization: Some(ReauthorizationTarget {
account_id: ProviderAccountId::new("acct_test").expect("account ID"),
credential_revision: revision(1),
}),
reauthorization: Some(ProviderAccountId::new("acct_test").expect("account ID")),
})
.await
.expect("start reauthorization");
@@ -250,15 +374,14 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
);
assert!(matches!(
pending.target(),
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} if account_id.as_str() == "acct_test"
&& *expected_credential_revision == revision(1)
AuthorizationMutationTarget::Reauthorize { account_id }
if account_id.as_str() == "acct_test"
));
assert!(!format!("{pending:?}").contains("admin-test"));
services
provider.set_current_credential_revision(revision(2));
let result = services
.openai()
.complete_authorization(CompleteAuthorization {
context: context("oauth-complete-openai"),
@@ -268,6 +391,8 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
.await
.expect("complete reauthorization");
assert_eq!(result.credential_revision, Some(revision(3)));
assert_eq!(
recorded(&events),
[
+3 -1
View File
@@ -31,7 +31,9 @@ uuid = { version = "1.23.3", features = ["v7"] }
[dev-dependencies]
async-trait = "0.1.89"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] }
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread"] }
tokio-tungstenite = "0.29.0"
tower = { version = "0.5.3", features = ["util"] }
[lints]
@@ -16,7 +16,7 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, FixedOffset, Utc};
use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
AdminError as AdminServiceError, AdminErrorKind, PageSize, Revision,
AdminError as AdminServiceError, AdminErrorKind, PageSize,
accounts::{
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
@@ -26,8 +26,8 @@ use gateway_admin::model::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion, CredentialDeletionResult,
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow,
ReauthorizationTarget, RotateCredential, StartAuthorization,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow, RotateCredential,
StartAuthorization,
},
};
use gateway_core::{
@@ -196,8 +196,6 @@ pub struct AccountView {
pub token_refreshing: bool,
pub availability: String,
pub enabled: bool,
pub credential_revision: u64,
pub state_revision: Option<u64>,
pub access_token_expires_at: Option<String>,
pub access_token_expires_at_display: Option<String>,
pub refresh_token_expires_at: Option<String>,
@@ -604,24 +602,16 @@ pub struct StartAccountAuthorizationRequest {
pub provider: String,
pub name: String,
pub account_id: Option<String>,
pub expected_credential_revision: Option<u64>,
}
impl StartAccountAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_text(&self.name, MAX_NAME_BYTES, "name")?;
match (
self.account_id.as_deref(),
self.expected_credential_revision,
) {
(None, None) => Ok(()),
(Some(account_id), Some(credential_revision)) => {
require_account_id(account_id, "accountId")?;
require_positive(credential_revision, "expectedCredentialRevision")
}
_ => Err(WireValidationError::new("reauthorization")),
if let Some(account_id) = self.account_id.as_deref() {
require_account_id(account_id, "accountId")?;
}
Ok(())
}
fn into_command(
@@ -630,15 +620,11 @@ impl StartAccountAuthorizationRequest {
) -> Result<(AccountProvider, StartAuthorization), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
let reauthorization = match (self.account_id, self.expected_credential_revision) {
(None, None) => None,
(Some(account_id), Some(credential_revision)) => Some(ReauthorizationTarget {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
credential_revision: revision(credential_revision, "expectedCredentialRevision")?,
}),
_ => return Err(WireValidationError::new("reauthorization")),
};
let reauthorization = self
.account_id
.map(ProviderAccountId::new)
.transpose()
.map_err(|_| WireValidationError::new("accountId"))?;
Ok((
provider,
StartAuthorization {
@@ -771,7 +757,6 @@ impl AccountDeletionRequest {
pub struct RotateAccountRequest {
pub provider: String,
pub account_id: String,
pub expected_credential_revision: u64,
pub access_token: String,
pub refresh_token: Option<String>,
}
@@ -782,10 +767,6 @@ impl RotateAccountRequest {
return Err(WireValidationError::new("provider"));
}
require_account_id(&self.account_id, "accountId")?;
require_positive(
self.expected_credential_revision,
"expectedCredentialRevision",
)?;
validate_oauth_material(&self.access_token, self.refresh_token.as_deref())
}
@@ -806,10 +787,6 @@ impl RotateAccountRequest {
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
expected_credential_revision: revision(
self.expected_credential_revision,
"expectedCredentialRevision",
)?,
provider_material: ProviderDocument::new(OpaqueProviderData::new(material)),
})
}
@@ -858,15 +835,12 @@ impl From<AuthorizationStarted> for AccountAuthorizationData {
#[serde(rename_all = "camelCase")]
pub struct AccountMutationData {
pub account_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub credential_revision: Option<u64>,
}
impl From<CredentialMutationResult> for AccountMutationData {
fn from(result: CredentialMutationResult) -> Self {
Self {
account_id: result.account_id.to_string(),
credential_revision: result.credential_revision.map(Revision::get),
}
}
}
@@ -1409,8 +1383,6 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
token_refreshing: false,
availability: account.availability.as_str().to_owned(),
enabled: account.enabled,
credential_revision: account.credential_revision.get(),
state_revision: None,
access_token_expires_at: expires_at,
access_token_expires_at_display: account
.access_token_expires_at
@@ -1661,10 +1633,6 @@ fn empty_account_usage() -> AccountUsageView {
}
}
fn revision(value: u64, field: &'static str) -> Result<Revision, WireValidationError> {
Revision::new(value).map_err(|_| WireValidationError::new(field))
}
fn provider_document(
value: Value,
field: &'static str,
@@ -1698,13 +1666,6 @@ fn require_text(
Ok(())
}
fn require_positive(value: u64, field: &'static str) -> Result<(), WireValidationError> {
if value == 0 || i64::try_from(value).is_err() {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn validate_oauth_material(
access_token: &str,
refresh_token: Option<&str>,
@@ -105,12 +105,15 @@ mod actions {
use gateway_admin::model::{
Revision,
accounts::{AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus},
provider_credentials::{CredentialDeletionResult, CredentialImportResult},
provider_credentials::{
CredentialDeletionResult, CredentialImportResult, CredentialMutationResult,
},
};
use gateway_api::admin::accounts::{
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
AccountImportData, AccountImportRequest, AccountRefreshRequest, AccountTestQuery,
AccountImportData, AccountImportRequest, AccountMutationData, AccountRefreshRequest,
AccountTestQuery, RotateAccountRequest, StartAccountAuthorizationRequest,
};
use gateway_core::engine::credential::ProviderAccountId;
use serde_json::json;
@@ -168,6 +171,58 @@ mod actions {
);
}
#[test]
fn credential_recovery_requests_should_not_accept_client_revision_fences() {
let authorization: StartAccountAuthorizationRequest = serde_json::from_value(json!({
"provider": "openai",
"name": "reauthorize",
"accountId": "acct_1"
}))
.expect("decode reauthorization");
assert!(authorization.validate().is_ok());
assert!(
serde_json::from_value::<StartAccountAuthorizationRequest>(json!({
"provider": "openai",
"name": "reauthorize",
"accountId": "acct_1",
"expectedCredentialRevision": 1
}))
.is_err()
);
let rotation: RotateAccountRequest = serde_json::from_value(json!({
"provider": "openai",
"accountId": "acct_1",
"accessToken": "header.payload.signature",
"refreshToken": "refresh-token"
}))
.expect("decode rotation");
assert!(rotation.validate().is_ok());
assert!(
serde_json::from_value::<RotateAccountRequest>(json!({
"provider": "openai",
"accountId": "acct_1",
"expectedCredentialRevision": 1,
"accessToken": "header.payload.signature"
}))
.is_err()
);
}
#[test]
fn credential_mutation_response_should_not_expose_internal_revision() {
let response = AccountMutationData::from(CredentialMutationResult {
config_revision: Revision::new(8).expect("config revision"),
account_id: ProviderAccountId::new("acct_1").expect("account ID"),
credential_revision: Some(Revision::new(9).expect("credential revision")),
});
assert_eq!(
serde_json::to_value(response).expect("serialize credential mutation"),
json!({ "accountId": "acct_1" })
);
}
#[test]
fn account_import_should_use_provider_and_opaque_data_fields() {
let valid: AccountImportRequest = serde_json::from_value(json!({
@@ -1055,6 +1055,86 @@ async fn streaming_upstream_wire_failure_should_not_be_rewritten_as_a_gateway_er
);
}
#[tokio::test]
async fn atomic_uncommitted_upstream_failure_batch_should_be_forwarded_once() {
let trace = Arc::new(Trace::default());
let raw_created = Bytes::from_static(
b"event: response.created\r\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_atomic_api\",\"model\":\"public-model\",\"status\":\"in_progress\"}}\r\n\r\n",
);
let raw_failure = Bytes::from_static(
b"event: response.failed\r\ndata: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_atomic_api\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"atomic upstream marker\"}}}\r\n\r\n",
);
let started_wire = ProviderEvent::canonical_with_wire(
vec![started()],
ProtocolWireEvent::json_with_raw_sse_metadata(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {
"id": "resp_atomic_api",
"model": "public-model",
"status": "in_progress"
}
}),
raw_created.clone(),
None,
None,
)
.expect("valid upstream started wire"),
);
let failed_wire = ProviderEvent::wire(
ProtocolWireEvent::json_with_raw_sse_metadata(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": "resp_atomic_api",
"status": "failed",
"error": {
"code": "rate_limit_exceeded",
"message": "atomic upstream marker"
}
}
}),
raw_failure.clone(),
None,
None,
)
.expect("valid upstream failure wire"),
);
let batch = CoordinatedEvent::try_batch(
vec![started_wire, failed_wire],
CommitRequirement::CommitBeforeDelivery,
)
.expect("atomic failure delivery");
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::Event(batch),
NextStep::Error(EngineError::Provider(ProviderError::new(
ProviderErrorKind::RateLimited,
UpstreamSendState::Sent,
))),
],
);
let response = stream_execution_response(Box::new(session), None).await;
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("read SSE body");
let body = String::from_utf8(body.to_vec()).expect("SSE is UTF-8");
assert_eq!(body.matches("atomic upstream marker").count(), 1);
assert_eq!(body.matches("event: response.failed").count(), 1);
assert!(body.contains(std::str::from_utf8(&raw_created).expect("created frame")));
assert!(body.contains(std::str::from_utf8(&raw_failure).expect("failure frame")));
assert!(!body.contains("upstream capacity is temporarily unavailable"));
assert!(body.ends_with("data: [DONE]\n\n"));
assert_eq!(trace.snapshot(), vec!["next_event", "commit", "next_error"]);
}
#[tokio::test]
async fn streaming_second_commit_request_should_fail_and_finalize_once() {
let trace = Arc::new(Trace::default());
@@ -1,14 +1,34 @@
use std::sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering},
};
use std::time::{Duration, SystemTime};
use axum::{
body::Body,
http::{Request, StatusCode, header::AUTHORIZATION},
};
use futures::{SinkExt, StreamExt, future::BoxFuture};
use gateway_api::openai::responses::ResponseCreateFrameError;
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
StartExecution, StartedExecution,
};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, UpstreamSendState};
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind};
use gateway_core::event::{
GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta,
};
use gateway_core::operation::Operation;
use serde_json::json;
use gateway_core::routing::PublicModelId;
use serde_json::{Value, json};
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message as ClientMessage;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tower::ServiceExt;
use super::decode_response_create;
use crate::openai::{api_router, models::ModelsExecution};
use crate::openai::{api_router, authenticated_client, models::ModelsExecution};
#[test]
fn response_create_should_default_to_the_websocket_streaming_contract() {
@@ -148,11 +168,233 @@ fn response_create_should_reject_non_boolean_stream_without_disclosing_body_valu
assert!(!rendered.contains(opaque_stream_value));
}
// `serve_responses_websocket`/`forward_execution` 只有拿到 `hyper::upgrade::OnUpgrade`
// 的真实 HTTP/1.1 升级连接才会运行;本 crate 的依赖图按 axum `ws` feature 编译
// hyper(无 `http1`),测试进程内不存在能完成升级的服务端,`hyper::upgrade` 也不
// 提供公开构造器。集成测试因此止步于 upgrade 边界:会话循环内的事件转发、错误帧
// 与断连取消仍未被覆盖,需要 dev 依赖开启 `http1` 或 socket 泛化后才能落地。
#[derive(Default)]
struct AtomicFailureTrace {
next_calls: AtomicUsize,
committed: AtomicBool,
finalized: AtomicBool,
}
struct AtomicFailureSession {
trace: Arc<AtomicFailureTrace>,
}
impl ExecutionSession for AtomicFailureSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
match self.trace.next_calls.fetch_add(1, Ordering::AcqRel) {
0 => Ok(Some(atomic_failure_batch())),
1 => {
self.trace.finalized.store(true, Ordering::Release);
Err(EngineError::Provider(ProviderError::new(
ProviderErrorKind::RateLimited,
UpstreamSendState::Sent,
)))
}
_ => Ok(None),
}
})
}
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
Box::pin(async { Err(EngineError::InvalidDeliveryState) })
}
fn response_headers(&self) -> &[ProviderResponseHeader] {
&[]
}
fn commit_downstream(&mut self, _: Option<u16>) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
self.trace.committed.store(true, Ordering::Release);
Ok(())
})
}
fn record_client_status(&mut self, _: u16) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async { Ok(()) })
}
fn is_finalized(&self) -> bool {
self.trace.finalized.load(Ordering::Acquire)
}
fn cancel(&self) {}
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move {
self.trace.finalized.store(true, Ordering::Release);
})
}
}
struct AtomicFailureExecution {
client: AuthenticatedClient,
trace: Arc<AtomicFailureTrace>,
}
impl ExecutionService for AtomicFailureExecution {
fn authenticate(
&self,
plaintext: &str,
) -> Result<AuthenticatedClient, ClientAuthenticationError> {
(plaintext == "sk_ws_atomic")
.then(|| self.client.clone())
.ok_or(ClientAuthenticationError::InvalidKey)
}
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
vec![PublicModelId::new("model-a").expect("public model")]
}
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
model.as_str() == "model-a"
}
fn start(
&self,
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
Ok(StartedExecution {
request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic")
.expect("request id"),
created_at: SystemTime::now(),
stream: request.metadata.stream,
session: Box::new(AtomicFailureSession {
trace: Arc::clone(&self.trace),
}),
})
})
}
}
fn atomic_failure_batch() -> CoordinatedEvent {
let started = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
"resp_ws_atomic",
"model-a",
))],
ProtocolWireEvent::json(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {
"id": "resp_ws_atomic",
"model": "model-a",
"status": "in_progress"
}
}),
)
.expect("created wire"),
);
let failed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": "resp_ws_atomic",
"status": "failed",
"error": {
"code": "rate_limit_exceeded",
"message": "websocket atomic upstream marker"
}
}
}),
)
.expect("failed wire"),
);
CoordinatedEvent::try_batch(
vec![started, failed],
CommitRequirement::CommitBeforeDelivery,
)
.expect("atomic WebSocket batch")
}
#[tokio::test]
async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() {
let trace = Arc::new(AtomicFailureTrace::default());
let execution = Arc::new(AtomicFailureExecution {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_atomic".parse().expect("authorization"),
);
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
socket
.send(ClientMessage::Text(
json!({
"type": "response.create",
"model": "model-a",
"input": "hello"
})
.to_string()
.into(),
))
.await
.expect("send response.create");
let mut messages = Vec::<Value>::new();
loop {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("WebSocket response timeout")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
continue;
};
let value: Value = serde_json::from_str(&text).expect("JSON WebSocket event");
let failed = value.get("type").and_then(Value::as_str) == Some("response.failed");
messages.push(value);
if failed {
break;
}
}
assert_eq!(
messages
.iter()
.filter(|message| message.get("type").and_then(Value::as_str) == Some("response.failed"))
.count(),
1
);
let failed = messages
.iter()
.find(|message| message.get("type").and_then(Value::as_str) == Some("response.failed"))
.expect("upstream failed event");
assert_eq!(
failed
.pointer("/response/error/message")
.and_then(Value::as_str),
Some("websocket atomic upstream marker")
);
assert!(trace.committed.load(Ordering::Acquire));
assert_eq!(trace.next_calls.load(Ordering::Acquire), 2);
assert!(trace.finalized.load(Ordering::Acquire));
socket.close(None).await.expect("close WebSocket");
server.abort();
}
const TEST_WEBSOCKET_KEY: &str = "AAAAAAAAAAAAAAAAAAAAAA==";
@@ -32,6 +32,21 @@ pub struct AttemptCoordinator<S: ?Sized> {
engine: Arc<GatewayEngine<S>>,
}
#[derive(Debug, Clone)]
enum AccountSelection {
Scheduled(Option<crate::engine::credential::ProviderAccountId>),
Diagnostic(crate::engine::credential::ProviderAccountId),
}
impl AccountSelection {
fn required_account(&self) -> Option<&crate::engine::credential::ProviderAccountId> {
match self {
Self::Scheduled(account) => account.as_ref(),
Self::Diagnostic(account) => Some(account),
}
}
}
impl<S: ?Sized> AttemptCoordinator<S>
where
S: ExecutionStore,
@@ -56,6 +71,49 @@ where
required_account: Option<crate::engine::credential::ProviderAccountId>,
continuation: Option<ContinuationBinding>,
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
self.start_with_account_selection(
request,
operation,
plan,
AccountSelection::Scheduled(required_account),
continuation,
cancellation,
)
.await
}
/// 对固定账号执行诊断请求。
///
/// 仅跳过该账号的本地可用性投影;账号租约、并发和请求间隔仍必须满足。
pub async fn start_diagnostic(
&self,
request: NewModelRequest,
operation: Operation,
plan: RoutingPlan,
required_account: crate::engine::credential::ProviderAccountId,
continuation: Option<ContinuationBinding>,
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
self.start_with_account_selection(
request,
operation,
plan,
AccountSelection::Diagnostic(required_account),
continuation,
cancellation,
)
.await
}
async fn start_with_account_selection(
&self,
request: NewModelRequest,
operation: Operation,
plan: RoutingPlan,
account_selection: AccountSelection,
continuation: Option<ContinuationBinding>,
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
let request_id = request.id.clone();
let client_api_key_ref = request.client_api_key_ref.clone();
@@ -84,7 +142,7 @@ where
request_persisted: false,
operation,
plan,
required_account,
account_selection,
continuation,
continuation_attempt,
account_state_owner,
@@ -107,7 +165,9 @@ where
client_response_id: None,
upstream_response_id: None,
last_account_exhaustion: None,
last_account_exhaustion_events: Vec::new(),
provider_attempt_outcomes: Vec::new(),
pending_terminal_failure: None,
};
if session.cancellation.is_cancelled() {
@@ -158,7 +218,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
request_persisted: bool,
operation: Operation,
plan: RoutingPlan,
required_account: Option<crate::engine::credential::ProviderAccountId>,
account_selection: AccountSelection,
continuation: Option<ContinuationBinding>,
continuation_attempt: ContinuationAttempt,
account_state_owner: Option<ProviderAccountStateOwner>,
@@ -185,7 +245,16 @@ pub struct ResponseExecutionSession<S: ?Sized> {
client_response_id: Option<String>,
upstream_response_id: Option<String>,
last_account_exhaustion: Option<ProviderError>,
/// 最近一次账号耗尽在重试时被丢弃的原始失败批次;只用于后续空选路终态。
last_account_exhaustion_events: Vec<ProviderEvent>,
provider_attempt_outcomes: Vec<ProviderAttemptOutcome>,
/// 原子失败批次已交给协议层、但尚待下游提交后收敛的原 Provider 错误。
pending_terminal_failure: Option<PendingTerminalFailure>,
}
struct PendingTerminalFailure {
error: ProviderError,
send_state: UpstreamSendState,
}
impl<S: ?Sized> ResponseExecutionSession<S>
@@ -201,6 +270,12 @@ where
if self.delivery_pending {
return Err(EngineError::DownstreamCommitRequired);
}
if let Some(pending) = self.pending_terminal_failure.take() {
let engine_error = provider_engine_error(&pending.error);
self.finish_provider_error_with_send_state(&pending.error, pending.send_state)
.await?;
return Err(engine_error);
}
if self.finalized {
return Ok(None);
}
@@ -220,6 +295,21 @@ where
)));
}
PullOutcome::AttemptDiscarded => {}
PullOutcome::TerminalFailure {
events,
error,
send_state,
} => {
let requirement = if self.downstream_committed_at.is_some() {
CommitRequirement::AlreadyCommitted
} else {
self.delivery_pending = true;
CommitRequirement::CommitBeforeDelivery
};
self.pending_terminal_failure =
Some(PendingTerminalFailure { error, send_state });
return CoordinatedEvent::try_batch(events, requirement).map(Some);
}
PullOutcome::End => {
if self.downstream_committed_at.is_some() {
self.finish_success().await?;
@@ -248,6 +338,14 @@ where
match self.pull().await? {
PullOutcome::Event(event) => events.push(event),
PullOutcome::AttemptDiscarded => events.clear(),
PullOutcome::TerminalFailure {
error, send_state, ..
} => {
let engine_error = provider_engine_error(&error);
self.finish_provider_error_with_send_state(&error, send_state)
.await?;
return Err(engine_error);
}
PullOutcome::End => {
if events.is_empty() {
return Err(EngineError::InvalidDeliveryState);
@@ -376,13 +474,20 @@ where
return Ok(());
}
self.delivery_pending = false;
if let Some(pending) = self.pending_terminal_failure.take() {
return self
.finish_provider_error_with_send_state(&pending.error, pending.send_state)
.await;
}
self.finish_interruption(EngineError::Cancelled).await
}
async fn pull(&mut self) -> Result<PullOutcome, EngineError> {
loop {
if self.current.is_none() {
self.prepare_attempt().await?;
if self.current.is_none()
&& let Some(outcome) = self.prepare_attempt().await?
{
return Ok(outcome);
}
let boundary = {
@@ -412,10 +517,14 @@ where
if let Some(observation) = event.take_observation()
&& let Err(error) = self.observe_response(observation)
{
if self.handle_stream_error(error).await? {
return Ok(PullOutcome::AttemptDiscarded);
match self.handle_stream_error(error).await? {
StreamErrorOutcome::AttemptDiscarded => {
return Ok(PullOutcome::AttemptDiscarded);
}
terminal @ StreamErrorOutcome::TerminalFailure { .. } => {
return Ok(terminal.into_pull_outcome());
}
}
continue;
}
let mut identity_error = None;
for fact in event.canonical_facts() {
@@ -425,10 +534,14 @@ where
}
}
if let Some(error) = identity_error {
if self.handle_stream_error(error).await? {
return Ok(PullOutcome::AttemptDiscarded);
match self.handle_stream_error(error).await? {
StreamErrorOutcome::AttemptDiscarded => {
return Ok(PullOutcome::AttemptDiscarded);
}
terminal @ StreamErrorOutcome::TerminalFailure { .. } => {
return Ok(terminal.into_pull_outcome());
}
}
continue;
}
for fact in event.canonical_facts() {
self.observe_event(fact).await?;
@@ -442,8 +555,13 @@ where
return Ok(PullOutcome::Event(event));
}
PollBoundary::Item(Some(Err(error))) => {
if self.handle_stream_error(error).await? {
return Ok(PullOutcome::AttemptDiscarded);
match self.handle_stream_error(error).await? {
StreamErrorOutcome::AttemptDiscarded => {
return Ok(PullOutcome::AttemptDiscarded);
}
terminal @ StreamErrorOutcome::TerminalFailure { .. } => {
return Ok(terminal.into_pull_outcome());
}
}
}
PollBoundary::Item(None) => {
@@ -455,7 +573,7 @@ where
}
}
async fn prepare_attempt(&mut self) -> Result<(), EngineError> {
async fn prepare_attempt(&mut self) -> Result<Option<PullOutcome>, EngineError> {
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
@@ -483,27 +601,36 @@ where
.ok_or(EngineError::EmptyRoutingPlan)?;
// recovery 钉选在此被一次性消费,只绑定本次 replay attempt;
// 外部 required_account 每次 attempt 都重新生效。
let pinned_account = self
.recovery_account
.take()
.or_else(|| self.required_account.clone());
let pinned_account = match &self.account_selection {
AccountSelection::Diagnostic(account) => Some(account.clone()),
AccountSelection::Scheduled(_) => self
.recovery_account
.take()
.or_else(|| self.account_selection.required_account().cloned()),
};
let account_context = match &self.account_selection {
AccountSelection::Diagnostic(account) => AccountAttemptContext::diagnostic(
self.excluded_accounts.clone(),
account.clone(),
self.account_state_owner.clone(),
),
AccountSelection::Scheduled(_) => AccountAttemptContext::new(
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
),
}
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
.contains(account)
}));
let context = AttemptContext::new(
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
.with_timing_started_at(self.timing_started_at),
next_attempt,
self.deadline,
self.plan.account_selection_policy(),
AccountAttemptContext::new(
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
)
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(
|account| {
self.credential_recovery_attempted_accounts
.contains(account)
},
)),
account_context,
self.continuation.clone(),
self.cancellation.clone(),
)
@@ -548,6 +675,16 @@ where
if error.kind() == ProviderErrorKind::NoEligibleAccount
&& let Some(exhaustion) = self.last_account_exhaustion.clone()
{
let send_state = self.current_send_state();
let mut events = std::mem::take(&mut self.last_account_exhaustion_events);
if !events.is_empty() {
self.observe_atomic_terminal_events(&mut events)?;
return Ok(Some(PullOutcome::TerminalFailure {
events,
error: exhaustion,
send_state,
}));
}
let engine_error = provider_engine_error(&exhaustion);
self.finish_provider_error(&exhaustion).await?;
return Err(engine_error);
@@ -672,10 +809,24 @@ where
send_observed: false,
response_observation: None,
});
Ok(())
Ok(None)
}
async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> {
self.observe_event_facts(event);
let current = self.current.as_mut().ok_or(EngineError::NoActiveAttempt)?;
if !current.send_observed {
self.engine
.store()
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
.await?;
current.send_observed = true;
self.send_state_watermark = UpstreamSendState::Sent;
}
Ok(())
}
fn observe_event_facts(&mut self, event: &GatewayEvent) {
let elapsed = elapsed_ms(self.timing_started_at);
observe_event_timing(&mut self.timings, event, elapsed);
if let GatewayEvent::Usage(observed) = event {
@@ -689,16 +840,6 @@ where
if let GatewayEvent::ProviderCost(observed) = event {
self.cost = observed.into_estimate();
}
let current = self.current.as_mut().ok_or(EngineError::NoActiveAttempt)?;
if !current.send_observed {
self.engine
.store()
.mark_send_state(&self.request_id, UpstreamSendState::Sent)
.await?;
current.send_observed = true;
self.send_state_watermark = UpstreamSendState::Sent;
}
Ok(())
}
async fn observe_wire_event(&mut self) -> Result<(), EngineError> {
@@ -788,14 +929,22 @@ where
Ok(())
}
/// 返回 `true` 表示调用方必须丢弃本 attempt 已收集的未提交事件。
async fn handle_stream_error(&mut self, error: ProviderError) -> Result<bool, EngineError> {
/// 失败可重试时要求丢弃本 attempt;预算耗尽时可返回最后一批原始失败事件。
async fn handle_stream_error(
&mut self,
mut error: ProviderError,
) -> Result<StreamErrorOutcome, EngineError> {
// 原始 wire 只能活在本次决策栈内;后续 clone、attempt 记录与终态对象均只
// 接触已剥离的稳定错误字段。
let mut atomic_client_events = error.take_atomic_client_events();
let current = self.current.take().ok_or(EngineError::NoActiveAttempt)?;
if matches!(
let account_exhausted = matches!(
error.kind(),
ProviderErrorKind::RateLimited | ProviderErrorKind::QuotaExhausted
) {
);
if account_exhausted {
self.last_account_exhaustion = Some(error.clone());
self.last_account_exhaustion_events.clear();
}
self.record_provider_failure(current.metadata.provider().clone(), error.kind());
// attempt_send_state 是本 attempt 自身的发送事实,驱动重试门;
@@ -818,7 +967,7 @@ where
attempt_send_state,
provider_proved_replay_safe,
);
let ordinary_retry = self.required_account.is_none()
let ordinary_retry = self.account_selection.required_account().is_none()
&& self.continuation_attempt == ContinuationAttempt::None
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
@@ -843,6 +992,9 @@ where
let retryable = continuation_retry || same_account_retry || ordinary_retry;
if retryable {
if account_exhausted {
self.last_account_exhaustion_events = atomic_client_events;
}
if same_account_retry {
if let Some(account) = current.metadata.provider_account_id() {
self.credential_recovery_attempted_accounts
@@ -879,16 +1031,42 @@ where
})
.await?;
self.reset_uncommitted_observations();
return Ok(true);
return Ok(StreamErrorOutcome::AttemptDiscarded);
}
let engine_error = provider_engine_error(&error);
self.current = Some(current);
if !atomic_client_events.is_empty() {
self.observe_atomic_terminal_events(&mut atomic_client_events)?;
return Ok(StreamErrorOutcome::TerminalFailure {
events: atomic_client_events,
error,
send_state,
});
}
self.finish_provider_error_with_send_state(&error, send_state)
.await?;
Err(engine_error)
}
fn observe_atomic_terminal_events(
&mut self,
events: &mut [ProviderEvent],
) -> Result<(), EngineError> {
for event in events {
if let Some(observation) = event.take_observation() {
self.observe_response(observation)
.map_err(EngineError::Provider)?;
}
for fact in event.canonical_facts() {
self.observe_response_identity(fact)
.map_err(EngineError::Provider)?;
self.observe_event_facts(fact);
}
}
Ok(())
}
fn prepare_continuation_retry(
&mut self,
current: &CurrentAttempt,
@@ -896,7 +1074,7 @@ where
send_state: UpstreamSendState,
provider_proved_replay_safe: bool,
) -> bool {
if self.required_account.is_some()
if self.account_selection.required_account().is_some()
|| self.continuation_attempt == ContinuationAttempt::None
|| self.downstream_committed_at.is_some()
|| self.delivery_pending
@@ -1238,9 +1416,40 @@ where
enum PullOutcome {
Event(ProviderEvent),
AttemptDiscarded,
TerminalFailure {
events: Vec<ProviderEvent>,
error: ProviderError,
send_state: UpstreamSendState,
},
End,
}
enum StreamErrorOutcome {
AttemptDiscarded,
TerminalFailure {
events: Vec<ProviderEvent>,
error: ProviderError,
send_state: UpstreamSendState,
},
}
impl StreamErrorOutcome {
fn into_pull_outcome(self) -> PullOutcome {
match self {
Self::AttemptDiscarded => PullOutcome::AttemptDiscarded,
Self::TerminalFailure {
events,
error,
send_state,
} => PullOutcome::TerminalFailure {
events,
error,
send_state,
},
}
}
}
enum PollBoundary {
Item(Option<Result<ProviderEvent, ProviderError>>),
Cancelled,
@@ -602,6 +602,15 @@ pub trait ProviderAccountStore: Send + Sync {
expected_revision: CredentialRevision,
) -> Result<LoadedCredential, StoreError>;
/// 读取账号当前 credential 及其 revision,不做任何版本比对。
///
/// 管理写入必须在临近 CAS 时用它取 fence,而不是携带调用方持有的旧 revision:
/// 后台刷新随时会推进 revision,用陈旧快照会把正常的恢复操作误判为冲突。
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, StoreError>;
async fn compare_and_swap_credential(
&self,
update: CredentialCasUpdate,
@@ -901,6 +910,26 @@ pub struct AccountSelectionContext {
pub excluded_accounts: BTreeSet<ProviderAccountId>,
pub preferred_account: Option<ProviderAccountId>,
pub round_robin_cursor: u64,
pub availability: AccountAvailabilityPolicy,
}
/// 选择账号时是否信任本地可用性投影。
///
/// 管理端对指定账号执行诊断时,会直接向上游确认实际状态;该模式跳过
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约、并发和
/// 请求间隔约束,且只能与固定账号约束组合使用。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AccountAvailabilityPolicy {
#[default]
Enforce,
BypassForDiagnostic,
}
impl AccountAvailabilityPolicy {
#[must_use]
pub const fn bypasses_local_availability(self) -> bool {
matches!(self, Self::BypassForDiagnostic)
}
}
/// 同一 target 内唯一的账号排序器。
@@ -1101,7 +1130,8 @@ fn smart_preference_should_escape(signals: &AccountRuntimeSignals) -> bool {
}
fn eligible(candidate: &AccountCandidate, context: &AccountSelectionContext) -> bool {
if !candidate.account.is_schedulable(context.now)
if (!context.availability.bypasses_local_availability()
&& !candidate.account.is_schedulable(context.now))
|| context.excluded_accounts.contains(candidate.account.id())
|| candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get()
{
@@ -477,11 +477,11 @@ impl DefaultExecutionService {
};
let mut session = match self
.probe_coordinator
.start(
.start_diagnostic(
new_request,
operation,
plan,
Some(account_id),
account_id,
None,
CancellationToken::new(),
)
@@ -525,7 +525,8 @@ impl DefaultExecutionService {
})
}
/// 探测失败先落库再转成客户端错误;转换会丢掉上游状态码与 provider code。
/// 探测失败先记录完整分类事实,再仅把 Provider 显式标记为客户端可见的
/// 结构化 message/code/type 带入协议错误。
async fn observe_probe_failure(
&self,
observed: &ProbeObservation,
@@ -232,6 +232,7 @@ pub struct AccountAttemptContext {
required_account: Option<ProviderAccountId>,
state_owner: Option<ProviderAccountStateOwner>,
credential_recovery_attempted: bool,
diagnostic_required_account: bool,
}
impl AccountAttemptContext {
@@ -246,6 +247,25 @@ impl AccountAttemptContext {
required_account,
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: false,
}
}
/// 为管理端对固定账号的上游诊断创建上下文。
///
/// 诊断只绕过本地的被动可用性投影,不能换号,并继续受账号租约保护。
#[must_use]
pub const fn diagnostic(
excluded_accounts: BTreeSet<ProviderAccountId>,
required_account: ProviderAccountId,
state_owner: Option<ProviderAccountStateOwner>,
) -> Self {
Self {
excluded_accounts,
required_account: Some(required_account),
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: true,
}
}
@@ -275,6 +295,12 @@ impl AccountAttemptContext {
pub const fn credential_recovery_attempted(&self) -> bool {
self.credential_recovery_attempted
}
/// 此尝试是否要由真实上游而非本地投影确认固定账号的可用性。
#[must_use]
pub const fn is_diagnostic_required_account(&self) -> bool {
self.diagnostic_required_account
}
}
/// 请求中 Provider 账号绑定状态的唯一归属。
@@ -474,6 +500,12 @@ impl AttemptContext {
self.account.credential_recovery_attempted()
}
/// 管理端固定账号诊断会跳过本地可用性投影,保留全部租约约束。
#[must_use]
pub const fn is_diagnostic_required_account(&self) -> bool {
self.account.is_diagnostic_required_account()
}
#[must_use]
pub const fn continuation(&self) -> Option<&ContinuationBinding> {
self.continuation.as_ref()
+46 -1
View File
@@ -6,6 +6,7 @@ use std::time::Duration;
use thiserror::Error;
use crate::engine::UpstreamSendState;
use crate::event::ProviderEvent;
/// 应用层标识不满足约束。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
@@ -266,9 +267,13 @@ impl fmt::Debug for ClientVisibleUpstreamError {
/// 单次 Provider 调用的稳定错误。
///
/// 该类型不接收原始响应正文,也不会在 `Debug` 或 `Display` 中打印上游
/// 稳定诊断字段不接收原始响应正文,也不会在 `Debug` 或 `Display` 中打印上游
/// code/request ID/response ID。Adapter 若捕获到可能含 secret 的上下文,只能调用
/// [`ProviderError::redact_sensitive_context`] 丢弃正文并留下脱敏标记。
///
/// 唯一例外是 [`ProviderError::with_atomic_client_events`]:它只承载尚未交付客户端、
/// 必须与本错误一起完成重试判断的协议事件。Core 会在任何持久化或日志记录前取走
/// 这些事件,不能把它们当作诊断上下文使用。
#[derive(Clone)]
pub struct ProviderError {
kind: ProviderErrorKind,
@@ -284,8 +289,12 @@ pub struct ProviderError {
retry_same_account: bool,
sensitive_context_redacted: bool,
client_visible_upstream_error: Option<Box<ClientVisibleUpstreamError>>,
atomic_client_events: Option<Box<AtomicClientEvents>>,
}
#[derive(Clone)]
struct AtomicClientEvents(Vec<ProviderEvent>);
impl ProviderError {
/// 创建 Provider 错误。
#[must_use]
@@ -304,6 +313,7 @@ impl ProviderError {
retry_same_account: false,
sensitive_context_redacted: false,
client_visible_upstream_error: None,
atomic_client_events: None,
}
}
@@ -379,6 +389,17 @@ impl ProviderError {
self
}
/// 附加必须先由 Core 判断重试、再决定丢弃或交付的客户端协议事件。
///
/// Provider 只能放入当前 attempt 尚未交付的事件;调用方应设置明确的有界缓冲,
/// 并保证失败事件本身位于这批事件的末尾。
#[must_use]
pub fn with_atomic_client_events(mut self, events: Vec<ProviderEvent>) -> Self {
self.atomic_client_events =
(!events.is_empty()).then(|| Box::new(AtomicClientEvents(events)));
self
}
/// 丢弃敏感正文,只记录“发生过脱敏”这一事实。
#[must_use]
pub fn redact_sensitive_context(mut self, _sensitive_context: impl AsRef<str>) -> Self {
@@ -462,6 +483,23 @@ impl ProviderError {
pub fn client_visible_upstream_error(&self) -> Option<&ClientVisibleUpstreamError> {
self.client_visible_upstream_error.as_deref()
}
/// 取走只供本次重试/提交决策使用的原子客户端事件。
///
/// Core 必须在克隆错误、记录中间失败或构造终态前调用本方法,避免原始 wire
/// 进入诊断或持久化对象。
#[must_use]
pub fn take_atomic_client_events(&mut self) -> Vec<ProviderEvent> {
self.atomic_client_events
.take()
.map_or_else(Vec::new, |events| events.0)
}
/// 返回是否携带尚未提交的原子客户端事件;不暴露事件内容。
#[must_use]
pub fn has_atomic_client_events(&self) -> bool {
self.atomic_client_events.is_some()
}
}
impl fmt::Debug for ProviderError {
@@ -508,6 +546,13 @@ impl fmt::Debug for ProviderError {
.as_ref()
.map(|_| "<present>"),
)
.field(
"atomic_client_events",
&self
.atomic_client_events
.as_ref()
.map_or(0, |events| events.0.len()),
)
.finish()
}
}
@@ -13,7 +13,7 @@ use crate::engine::credential::{
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
};
use crate::error::{IdentifierError, validate_text};
use crate::routing::ProviderKind;
use crate::routing::{ProviderKind, UpstreamModelId};
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
@@ -527,6 +527,80 @@ impl ProviderCooldown {
}
}
/// 可丢失 cooldown 的细粒度作用域。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum ProviderCooldownScope {
/// 只阻止同一账号调用指定上游模型。
UpstreamModel(UpstreamModelId),
}
impl ProviderCooldownScope {
#[must_use]
pub const fn upstream_model(model: UpstreamModelId) -> Self {
Self::UpstreamModel(model)
}
#[must_use]
pub const fn kind(&self) -> &'static str {
match self {
Self::UpstreamModel(_) => "model",
}
}
#[must_use]
pub fn value(&self) -> &str {
match self {
Self::UpstreamModel(model) => model.as_str(),
}
}
}
/// 不进入账号持久状态的账号+作用域 cooldown。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderScopedCooldown {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
scope: ProviderCooldownScope,
until: SystemTime,
}
impl ProviderScopedCooldown {
#[must_use]
pub const fn new(
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
scope: ProviderCooldownScope,
until: SystemTime,
) -> Self {
Self {
account_id,
credential_revision,
scope,
until,
}
}
#[must_use]
pub const fn account_id(&self) -> &ProviderAccountId {
&self.account_id
}
#[must_use]
pub const fn credential_revision(&self) -> CredentialRevision {
self.credential_revision
}
#[must_use]
pub const fn scope(&self) -> &ProviderCooldownScope {
&self.scope
}
#[must_use]
pub const fn until(&self) -> SystemTime {
self.until
}
}
pub trait ProviderCooldownPort: Send + Sync {
fn put_if_later(
&self,
@@ -543,6 +617,24 @@ pub trait ProviderCooldownPort: Send + Sync {
account_id: &'a ProviderAccountId,
through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
fn put_scoped_if_later(
&self,
cooldown: ProviderScopedCooldown,
) -> BoxFuture<'_, Result<bool, ProviderStoreError>>;
fn read_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>>;
fn clear_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -778,35 +870,75 @@ pub enum OAuthPendingPutOutcome {
AlreadyExists,
}
/// OAuth 回调处理取得临时 flow 的独占权结果。
///
/// 与一次性消费不同,Provider 在上游换取令牌失败时可以释放 claim,让同一 flow
/// 使用新的回调地址重试;只有完整校验成功后才会消费 flow。
#[derive(Clone, PartialEq)]
pub enum OAuthPendingTakeOutcome {
Taken(OpaqueProviderData),
pub enum OAuthPendingClaimOutcome {
Claimed(OpaqueProviderData),
NotFound,
OwnerMismatch,
InProgress,
}
impl fmt::Debug for OAuthPendingTakeOutcome {
impl fmt::Debug for OAuthPendingClaimOutcome {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Taken(_) => formatter.write_str("Taken([PROVIDER-OWNED])"),
Self::Claimed(_) => formatter.write_str("Claimed([PROVIDER-OWNED])"),
Self::NotFound => formatter.write_str("NotFound"),
Self::OwnerMismatch => formatter.write_str("OwnerMismatch"),
Self::InProgress => formatter.write_str("InProgress"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OAuthPendingReleaseOutcome {
Released,
NotFound,
OwnerMismatch,
ClaimMismatch,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OAuthPendingConsumeOutcome {
Consumed,
NotFound,
OwnerMismatch,
ClaimMismatch,
}
pub trait OAuthPendingFlowPort: Send + Sync {
fn put_if_absent(
&self,
flow: NewOAuthPendingFlow,
) -> BoxFuture<'_, Result<OAuthPendingPutOutcome, ProviderStoreError>>;
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>>;
claim: &'a OAuthPendingBinding,
claim_ttl: Duration,
) -> BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>>;
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>>;
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>>;
}
/// Provider 只能按能力取用端口,无法取得 Redis client 或 repository 集合。
@@ -399,6 +399,69 @@ fn complete_stream(total_tokens: Option<u64>) -> Vec<Result<GatewayEvent, Provid
events
}
fn atomic_response_failed(response_id: &str, marker: &str) -> ProviderError {
let started = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
response_id,
"gpt-5",
))],
ProtocolWireEvent::json(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {"id": response_id, "model": "gpt-5", "status": "in_progress"}
}),
)
.expect("started wire"),
);
let failed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": response_id,
"status": "failed",
"error": {"code": "rate_limit_exceeded", "message": marker}
}
}),
)
.expect("failed wire"),
);
ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_status(429)
.with_replay_safe()
.with_atomic_client_events(vec![started, failed])
}
fn bare_atomic_response_failed(response_id: &str) -> ProviderError {
let failed = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
response_id,
"fallback-model",
))],
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": response_id,
"status": "failed",
"error": {"code": "rate_limit_exceeded", "message": "bare failure"}
}
}),
)
.expect("bare failed wire"),
);
ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_status(429)
.with_replay_safe()
.with_atomic_client_events(vec![failed])
}
fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, ProviderError>> {
let mut usage = Usage::new();
usage.input_tokens = Some(12);
@@ -1202,7 +1265,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
}
#[test]
fn required_account_reaches_provider_and_matching_metadata_succeeds() {
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
@@ -1210,11 +1273,11 @@ fn required_account_reaches_provider_and_matching_metadata_succeeds() {
items: complete_stream(None),
}]);
let required = ProviderAccountId::new("acct_required").expect("account id");
let mut session = block_on(coordinator.start(
let mut session = block_on(coordinator.start_diagnostic(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
Some(required.clone()),
required.clone(),
None,
CancellationToken::new(),
))
@@ -1225,6 +1288,7 @@ fn required_account_reaches_provider_and_matching_metadata_succeeds() {
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 1);
assert_eq!(contexts[0].required_account(), Some(&required));
assert!(contexts[0].is_diagnostic_required_account());
}
#[test]
@@ -1836,6 +1900,277 @@ fn explicit_429_after_structural_event_should_retry_before_commit() {
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn atomic_response_failed_should_rotate_account_before_stream_commit() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_first",
items: vec![Err(atomic_response_failed(
"response-first-failed",
"first account failure",
))],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let first = block_on(session.next_event())
.expect("second account should produce the first delivery")
.expect("started event");
let delivered = first.into_provider_events();
assert!(delivered.iter().flat_map(ProviderEvent::canonical_facts).any(
|event| matches!(event, GatewayEvent::Started(meta) if meta.response_id() == "response-1")
));
assert!(
!delivered
.iter()
.filter_map(ProviderEvent::wire_event)
.any(|wire| {
wire.data().pointer("/response/id").and_then(Value::as_str)
== Some("response-first-failed")
})
);
block_on(session.commit_downstream(Some(200))).expect("commit second account");
while block_on(session.next_event())
.expect("successful stream")
.is_some()
{}
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.commits, 1);
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn bare_atomic_response_failed_should_rotate_before_the_first_delivery() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_first",
items: vec![Err(bare_atomic_response_failed("response-bare-failed"))],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let first = block_on(session.next_event())
.expect("second account should replace the bare failure")
.expect("second account started event");
assert!(first
.into_provider_events()
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Started(meta) if meta.response_id() == "response-1")));
block_on(session.commit_downstream(Some(200))).expect("commit second account");
while block_on(session.next_event())
.expect("successful stream")
.is_some()
{}
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn exhausted_atomic_response_failed_should_deliver_only_the_last_failure_once() {
let operation = generate_operation();
let route_plan = plan(&operation);
let scripts = (1..=route_plan.max_attempts().get())
.map(|attempt| {
let response_id = format!("response-attempt-{attempt}-failed");
let marker = if attempt == route_plan.max_attempts().get() {
"last upstream failure".to_owned()
} else {
format!("discarded failure {attempt}")
};
Script::ObservedStream {
account_id: "acct_failed",
items: vec![Err(atomic_response_failed(&response_id, &marker))],
}
})
.collect::<Vec<_>>();
let expected_attempts = route_plan.max_attempts().get();
let expected_last_response = format!("response-attempt-{expected_attempts}-failed");
let (coordinator, store, provider) = coordinator(scripts);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let terminal = block_on(session.next_event())
.expect("last failure remains deliverable")
.expect("terminal failure batch");
assert_eq!(
terminal.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
let events = terminal.into_provider_events();
let failures = events
.iter()
.filter_map(ProviderEvent::wire_event)
.filter(|wire| wire.event_type() == Some("response.failed"))
.collect::<Vec<_>>();
assert_eq!(failures.len(), 1);
assert_eq!(
failures[0]
.data()
.pointer("/response/id")
.and_then(Value::as_str),
Some(expected_last_response.as_str())
);
assert!(
failures[0]
.data()
.to_string()
.contains("last upstream failure")
);
block_on(session.commit_downstream(Some(200))).expect("commit original terminal failure");
let error = block_on(session.next_event()).expect_err("typed failure finalizes the request");
assert!(matches!(error, EngineError::Provider(_)));
assert!(session.is_finalized());
assert_eq!(
provider.contexts.lock().expect("contexts lock").len(),
usize::try_from(expected_attempts).expect("attempt count fits usize")
);
let state = store.state.lock().expect("store lock");
assert_eq!(
state.intermediate_failures,
usize::try_from(expected_attempts - 1).expect("attempt count fits usize")
);
assert_eq!(state.commits, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Incomplete);
assert!(state.finalizations[0].committed);
}
#[test]
fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_only",
items: vec![Err(bare_atomic_response_failed("response-only-failed"))],
},
Script::Error(ProviderError::new(
ProviderErrorKind::NoEligibleAccount,
UpstreamSendState::NotSent,
)),
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let terminal = block_on(session.next_event())
.expect("last account failure remains deliverable")
.expect("terminal failure batch");
let events = terminal.into_provider_events();
let failure = events
.iter()
.filter_map(ProviderEvent::wire_event)
.find(|wire| wire.event_type() == Some("response.failed"))
.expect("last upstream response.failed");
assert_eq!(
failure
.data()
.pointer("/response/id")
.and_then(Value::as_str),
Some("response-only-failed")
);
block_on(session.commit_downstream(Some(200))).expect("commit terminal failure");
assert!(matches!(
block_on(session.next_event()),
Err(EngineError::Provider(_))
));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.attempts.len(), 1);
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Incomplete);
assert_eq!(
state.finalizations[0].client_response_id.as_deref(),
Some("response-only-failed")
);
assert_eq!(
state.finalizations[0].upstream_response_id.as_deref(),
Some("response-only-failed")
);
}
#[test]
fn atomic_response_failed_keeps_collect_uncommitted_error_semantics() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![Script::ObservedStream {
account_id: "acct_required",
items: vec![Err(atomic_response_failed(
"response-required-failed",
"required account failure",
))],
}]);
let required = ProviderAccountId::new("acct_required").expect("required account");
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
Some(required),
None,
CancellationToken::new(),
))
.expect("start execution");
let error = block_on(session.collect_uncommitted())
.expect_err("non-streaming collection must keep returning the typed failure");
assert!(matches!(error, EngineError::Provider(_)));
assert!(session.is_finalized());
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
let state = store.state.lock().expect("store lock");
assert_eq!(state.commits, 0);
assert_eq!(state.intermediate_failures, 0);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn native_continuation_explicit_429_is_not_retried() {
let continuation = NativeContinuationPin::new(
@@ -5,11 +5,11 @@ use std::time::{Duration, SystemTime};
use serde_json::{Map, Value};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountCandidate, AccountFeedbackStats,
AccountQuotaSignals, AccountRuntimeSignals, AccountSelectionContext, AccountSelectionPolicy,
AccountSelector, CredentialCasUpdate, CredentialRevision, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountUpdate,
RotationStrategy,
AccountAttemptFeedback, AccountAvailability, AccountAvailabilityPolicy, AccountCandidate,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSelectionContext,
AccountSelectionPolicy, AccountSelector, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountUpdate, RotationStrategy,
};
use gateway_core::routing::ProviderKind;
@@ -51,6 +51,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
excluded_accounts: BTreeSet::new(),
preferred_account: None,
round_robin_cursor: 0,
availability: AccountAvailabilityPolicy::Enforce,
}
}
@@ -80,6 +81,55 @@ fn availability_should_round_trip_all_database_values() {
);
}
#[test]
fn diagnostic_selection_bypasses_all_local_account_availability() {
let exhausted = AccountCandidate {
account: account("acct_exhausted").with_runtime_state(
true,
AccountAvailability::QuotaExhausted,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let disabled = AccountCandidate {
account: account("acct_disabled").with_runtime_state(
false,
AccountAvailability::QuotaExhausted,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let mut context = context(RotationStrategy::Sticky);
context.availability = AccountAvailabilityPolicy::BypassForDiagnostic;
assert_eq!(
AccountSelector
.select(std::slice::from_ref(&exhausted), &context)
.map(|candidate| candidate.account.id()),
Some(exhausted.account.id())
);
assert_eq!(
AccountSelector
.select(std::slice::from_ref(&disabled), &context)
.map(|candidate| candidate.account.id()),
Some(disabled.account.id())
);
}
#[test]
fn disabled_account_should_not_be_schedulable() {
let account =
@@ -2,6 +2,8 @@ use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{
ClientVisibleUpstreamError, GatewayError, ProviderError, ProviderErrorKind, SafeUpstreamValue,
};
use gateway_core::event::{ProtocolWireEvent, ProviderEvent};
use serde_json::json;
#[test]
fn provider_error_debug_should_not_expose_sensitive_context() {
@@ -37,6 +39,26 @@ fn provider_error_replay_proof_should_be_explicit() {
assert!(error.replay_is_safe());
}
#[test]
fn provider_error_atomic_client_events_should_be_take_only_and_debug_redacted() {
let marker = "atomic-wire-must-not-enter-debug";
let event = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({"type": "response.failed", "message": marker}),
)
.expect("atomic wire"),
);
let mut error = ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
.with_atomic_client_events(vec![event]);
assert!(error.has_atomic_client_events());
assert!(!format!("{error:?}").contains(marker));
assert_eq!(error.take_atomic_client_events().len(), 1);
assert!(!error.has_atomic_client_events());
}
#[test]
fn gateway_error_should_keep_client_visible_upstream_fields_out_of_safe_diagnostics() {
let message = "Your Codex quota is exhausted";
@@ -2191,15 +2191,23 @@ impl ProviderAccountStore for PgProviderAccountRepository {
&self,
account: &CoreProviderAccountId,
expected_revision: CoreCredentialRevision,
) -> Result<LoadedCredential, CoreStoreError> {
let loaded = self.load_current_credential(account).await?;
if loaded.account.revision().get() != expected_revision.get() {
return Err(CoreStoreError::new(CoreStoreErrorKind::Conflict));
}
Ok(loaded)
}
async fn load_current_credential(
&self,
account: &CoreProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
let record = self
.load_provider_account(account.as_str())
.await
.map_err(core_store_error)?
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
if record.summary.credential_revision.get() != expected_revision.get() {
return Err(CoreStoreError::new(CoreStoreErrorKind::Conflict));
}
Ok(LoadedCredential {
account: core_account_from_summary(record.summary)?,
credential: PlaintextCredential::new(record.provider_credentials_json.fields().clone()),
@@ -13,7 +13,8 @@ use gateway_core::{
},
error::StoreError as CoreStoreError,
provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderStoreError, ProviderStoreErrorKind,
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown,
ProviderStoreError, ProviderStoreErrorKind,
},
routing::ProviderKind,
};
@@ -104,24 +105,35 @@ impl RedisCredentialCooldownRepository {
let fingerprint = resource_fingerprint("credential cooldown", provider_account_id)?;
Ok(format!("{}:account:{fingerprint}:cooldown", self.namespace))
}
}
#[async_trait]
impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
async fn cache_credential_cooldown(&self, cooldown: &CredentialCooldown) -> StoreResult<bool> {
require_nonempty(
"credential cooldown",
"provider_account_id",
&cooldown.provider_account_id,
)?;
let until_ms = cooldown.cooldown_until.timestamp_millis();
fn scoped_key(
&self,
provider_account_id: &str,
scope: &ProviderCooldownScope,
) -> StoreResult<String> {
let account_fingerprint = resource_fingerprint("credential cooldown", provider_account_id)?;
let scope_fingerprint = resource_fingerprint("credential cooldown scope", scope.value())?;
Ok(format!(
"{}:account:{account_fingerprint}:cooldown:{}:{scope_fingerprint}",
self.namespace,
scope.kind(),
))
}
async fn cache_at_key(
&self,
key: String,
credential_revision: Revision,
cooldown_until: DateTime<Utc>,
) -> StoreResult<bool> {
let until_ms = cooldown_until.timestamp_millis();
if until_ms <= 0 {
return Err(invalid("cooldown expiry must be positive"));
}
let mut connection = self.connection.clone();
let written = Script::new(WRITE_SCRIPT)
.key(self.key(&cooldown.provider_account_id)?)
.arg(cooldown.credential_revision.get())
.key(key)
.arg(credential_revision.get())
.arg(until_ms)
.invoke_async::<i64>(&mut connection)
.await
@@ -129,18 +141,10 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
Ok(written == 1)
}
async fn read_credential_cooldown(
&self,
provider_account_id: &str,
) -> StoreResult<Option<CredentialCooldown>> {
require_nonempty(
"credential cooldown",
"provider_account_id",
provider_account_id,
)?;
async fn read_at_key(&self, key: String) -> StoreResult<Option<(Revision, DateTime<Utc>)>> {
let mut connection = self.connection.clone();
let (present, revision, until_ms): (i64, String, String) = Script::new(READ_SCRIPT)
.key(self.key(provider_account_id)?)
.key(key)
.invoke_async(&mut connection)
.await
.map_err(|_| redis_unavailable("read credential cooldown"))?;
@@ -155,11 +159,59 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
.map_err(|_| invalid("cached cooldown expiry is invalid"))?;
let cooldown_until = DateTime::from_timestamp_millis(until_ms)
.ok_or_else(|| invalid("cached cooldown expiry is invalid"))?;
Ok(Some(CredentialCooldown {
provider_account_id: provider_account_id.to_owned(),
credential_revision: Revision::new(revision)?,
cooldown_until,
}))
Ok(Some((Revision::new(revision)?, cooldown_until)))
}
async fn invalidate_at_key(
&self,
key: String,
through_revision: Revision,
) -> StoreResult<bool> {
let mut connection = self.connection.clone();
let removed = Script::new(INVALIDATE_SCRIPT)
.key(key)
.arg(through_revision.get())
.invoke_async::<i64>(&mut connection)
.await
.map_err(|_| redis_unavailable("invalidate credential cooldown"))?;
Ok(removed == 1)
}
}
#[async_trait]
impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
async fn cache_credential_cooldown(&self, cooldown: &CredentialCooldown) -> StoreResult<bool> {
require_nonempty(
"credential cooldown",
"provider_account_id",
&cooldown.provider_account_id,
)?;
self.cache_at_key(
self.key(&cooldown.provider_account_id)?,
cooldown.credential_revision,
cooldown.cooldown_until,
)
.await
}
async fn read_credential_cooldown(
&self,
provider_account_id: &str,
) -> StoreResult<Option<CredentialCooldown>> {
require_nonempty(
"credential cooldown",
"provider_account_id",
provider_account_id,
)?;
self.read_at_key(self.key(provider_account_id)?)
.await
.map(|value| {
value.map(|(credential_revision, cooldown_until)| CredentialCooldown {
provider_account_id: provider_account_id.to_owned(),
credential_revision,
cooldown_until,
})
})
}
async fn invalidate_credential_cooldown(
@@ -172,14 +224,8 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
"provider_account_id",
provider_account_id,
)?;
let mut connection = self.connection.clone();
let removed = Script::new(INVALIDATE_SCRIPT)
.key(self.key(provider_account_id)?)
.arg(through_revision.get())
.invoke_async::<i64>(&mut connection)
self.invalidate_at_key(self.key(provider_account_id)?, through_revision)
.await
.map_err(|_| redis_unavailable("invalidate credential cooldown"))?;
Ok(removed == 1)
}
}
@@ -241,6 +287,69 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
.map_err(|_| provider_unavailable("clear credential cooldown"))
})
}
fn put_scoped_if_later(
&self,
cooldown: ProviderScopedCooldown,
) -> futures::future::BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let revision = Revision::new(cooldown.credential_revision().get())
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?;
self.cache_at_key(
self.scoped_key(cooldown.account_id().as_str(), cooldown.scope())
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
revision,
cooldown.until().into(),
)
.await
.map_err(|_| provider_unavailable("cache scoped credential cooldown"))
})
}
fn read_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
) -> futures::future::BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>>
{
Box::pin(async move {
self.read_at_key(
self.scoped_key(account_id.as_str(), scope)
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
)
.await
.map_err(|_| provider_unavailable("read scoped credential cooldown"))?
.map(|(revision, until)| {
Ok(ProviderScopedCooldown::new(
account_id.clone(),
CredentialRevision::new(revision.get())
.map_err(|_| provider_invalid("decode scoped credential cooldown"))?,
scope.clone(),
until.into(),
))
})
.transpose()
})
}
fn clear_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
through_revision: CredentialRevision,
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let revision = Revision::new(through_revision.get())
.map_err(|_| provider_invalid("encode scoped cooldown revision"))?;
self.invalidate_at_key(
self.scoped_key(account_id.as_str(), scope)
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
revision,
)
.await
.map_err(|_| provider_unavailable("clear scoped credential cooldown"))
})
}
}
/// 账号状态写入到可丢失 cooldown cache 的桥接 adapter。
@@ -348,6 +457,13 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore {
.await
}
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
self.authoritative.load_current_credential(account).await
}
async fn compare_and_swap_credential(
&self,
update: CredentialCasUpdate,
@@ -3,8 +3,9 @@
use std::time::Duration;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderStoreError, ProviderStoreErrorKind,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome, ProviderStoreError,
ProviderStoreErrorKind,
};
use gateway_core::routing::ProviderKind;
use redis::{Script, aio::ConnectionManager};
@@ -31,7 +32,7 @@ redis.call('PEXPIREAT', KEYS[1], expires_at_ms)
return 1
"#;
const TAKE_SCRIPT: &str = r#"
const CLAIM_SCRIPT: &str = r#"
local function equal_bytes(left, right)
local different = math.abs(string.len(left) - string.len(right))
local length = math.max(string.len(left), string.len(right))
@@ -51,14 +52,82 @@ end
if not equal_bytes(owner, ARGV[1]) then
return {-1, ''}
end
local clock = redis.call('TIME')
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
local existing_claim = redis.call('HGET', KEYS[1], 'claim_fingerprint')
if existing_claim ~= false then
local claim_expires_at = redis.call('HGET', KEYS[1], 'claim_expires_at_epoch_millis')
if claim_expires_at ~= false and tonumber(claim_expires_at) > now_ms then
return {-2, ''}
end
redis.call('HDEL', KEYS[1], 'claim_fingerprint', 'claim_expires_at_epoch_millis')
end
local payload = redis.call('HGET', KEYS[1], 'provider_payload')
if payload == false then
return redis.error_reply('OAuth pending flow payload is missing')
end
redis.call('DEL', KEYS[1])
redis.call('HSET', KEYS[1],
'claim_fingerprint', ARGV[2],
'claim_expires_at_epoch_millis', now_ms + tonumber(ARGV[3]))
return {1, payload}
"#;
const RELEASE_CLAIM_SCRIPT: &str = r#"
local function equal_bytes(left, right)
local different = math.abs(string.len(left) - string.len(right))
local length = math.max(string.len(left), string.len(right))
for index = 1, length do
local left_byte = string.byte(left, index) or 0
local right_byte = string.byte(right, index) or 0
if left_byte ~= right_byte then
different = different + 1
end
end
return different == 0
end
local owner = redis.call('HGET', KEYS[1], 'owner_fingerprint')
if owner == false then
return 0
end
if not equal_bytes(owner, ARGV[1]) then
return -1
end
local claim = redis.call('HGET', KEYS[1], 'claim_fingerprint')
if claim == false or not equal_bytes(claim, ARGV[2]) then
return -2
end
redis.call('HDEL', KEYS[1], 'claim_fingerprint', 'claim_expires_at_epoch_millis')
return 1
"#;
const CONSUME_CLAIM_SCRIPT: &str = r#"
local function equal_bytes(left, right)
local different = math.abs(string.len(left) - string.len(right))
local length = math.max(string.len(left), string.len(right))
for index = 1, length do
local left_byte = string.byte(left, index) or 0
local right_byte = string.byte(right, index) or 0
if left_byte ~= right_byte then
different = different + 1
end
end
return different == 0
end
local owner = redis.call('HGET', KEYS[1], 'owner_fingerprint')
if owner == false then
return 0
end
if not equal_bytes(owner, ARGV[1]) then
return -1
end
local claim = redis.call('HGET', KEYS[1], 'claim_fingerprint')
if claim == false or not equal_bytes(claim, ARGV[2]) then
return -2
end
redis.call('DEL', KEYS[1])
return 1
"#;
/// OAuth pending flow 的 Redis 原子持久化能力。
#[derive(Clone)]
pub struct RedisOAuthPendingFlowRepository {
@@ -118,39 +187,115 @@ impl OAuthPendingFlowPort for RedisOAuthPendingFlowRepository {
})
}
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>> {
claim: &'a OAuthPendingBinding,
claim_ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = self.key(provider_kind, flow);
let owner = provider_scoped_fingerprint(provider_kind, owner);
let claim = provider_scoped_fingerprint(provider_kind, claim);
let claim_ttl = ttl_millis(claim_ttl)
.map_err(|_| provider_invalid("encode OAuth pending claim TTL"))?;
let mut connection = self.connection.clone();
let (status, payload): (i64, Vec<u8>) = Script::new(TAKE_SCRIPT)
let (status, payload): (i64, Vec<u8>) = Script::new(CLAIM_SCRIPT)
.key(key)
.arg(owner)
.arg(claim)
.arg(claim_ttl)
.invoke_async(&mut connection)
.await
.map_err(|_| provider_unavailable("take OAuth pending flow"))?;
.map_err(|_| provider_unavailable("claim OAuth pending flow"))?;
match status {
0 => Ok(OAuthPendingTakeOutcome::NotFound),
-1 => Ok(OAuthPendingTakeOutcome::OwnerMismatch),
0 => Ok(OAuthPendingClaimOutcome::NotFound),
-1 => Ok(OAuthPendingClaimOutcome::OwnerMismatch),
-2 => Ok(OAuthPendingClaimOutcome::InProgress),
1 => {
let value: serde_json::Value = serde_json::from_slice(&payload)
.map_err(|_| provider_invalid("decode OAuth pending payload"))?;
let serde_json::Value::Object(fields) = value else {
return Err(provider_invalid("decode OAuth pending payload"));
};
Ok(OAuthPendingTakeOutcome::Taken(
Ok(OAuthPendingClaimOutcome::Claimed(
gateway_core::engine::credential::OpaqueProviderData::new(fields),
))
}
_ => Err(provider_invalid("decode OAuth pending outcome")),
_ => Err(provider_invalid("decode OAuth pending claim outcome")),
}
})
}
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>>
{
Box::pin(async move {
let key = self.key(provider_kind, flow);
let owner = provider_scoped_fingerprint(provider_kind, owner);
let claim = provider_scoped_fingerprint(provider_kind, claim);
let mut connection = self.connection.clone();
let status = Script::new(RELEASE_CLAIM_SCRIPT)
.key(key)
.arg(owner)
.arg(claim)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|_| provider_unavailable("release OAuth pending claim"))?;
Ok(match status {
0 => OAuthPendingReleaseOutcome::NotFound,
-1 => OAuthPendingReleaseOutcome::OwnerMismatch,
-2 => OAuthPendingReleaseOutcome::ClaimMismatch,
1 => OAuthPendingReleaseOutcome::Released,
_ => {
return Err(provider_invalid(
"decode OAuth pending claim release outcome",
));
}
})
})
}
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> futures::future::BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>>
{
Box::pin(async move {
let key = self.key(provider_kind, flow);
let owner = provider_scoped_fingerprint(provider_kind, owner);
let claim = provider_scoped_fingerprint(provider_kind, claim);
let mut connection = self.connection.clone();
let status = Script::new(CONSUME_CLAIM_SCRIPT)
.key(key)
.arg(owner)
.arg(claim)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|_| provider_unavailable("consume OAuth pending claim"))?;
Ok(match status {
0 => OAuthPendingConsumeOutcome::NotFound,
-1 => OAuthPendingConsumeOutcome::OwnerMismatch,
-2 => OAuthPendingConsumeOutcome::ClaimMismatch,
1 => OAuthPendingConsumeOutcome::Consumed,
_ => {
return Err(provider_invalid(
"decode OAuth pending claim consume outcome",
));
}
})
})
}
}
fn provider_scoped_fingerprint(
@@ -769,6 +769,25 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
assert_eq!(row.3["access_token"], "after-secret");
assert_eq!(row.4, 2);
assert!(
repository
.load_credential(
&account_id,
CredentialRevision::new(1).expect("stale credential revision"),
)
.await
.is_err()
);
let current = repository
.load_current_credential(&account_id)
.await
.expect("load current credential without a caller revision");
assert_eq!(current.account.revision().get(), 2);
assert_eq!(
current.credential.expose_to_provider()["access_token"],
"after-secret"
);
let stale = CredentialCasUpdate::new(
account_id.clone(),
CredentialRevision::new(1).expect("stale credential revision"),
@@ -13,6 +13,7 @@ use gateway_core::{
QuotaObservation, QuotaWriteOutcome,
},
error::{StoreError as CoreStoreError, StoreErrorKind},
provider_ports::{ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown},
routing::ProviderKind,
};
use gateway_store::{
@@ -73,6 +74,13 @@ impl ProviderAccountStore for CooldownTestAccountStore {
unreachable!("cooldown tests do not load credentials")
}
async fn load_current_credential(
&self,
_account: &ProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
unreachable!("cooldown tests do not load credentials")
}
async fn compare_and_swap_credential(
&self,
_update: CredentialCasUpdate,
@@ -523,6 +531,109 @@ async fn credential_cooldown_read_removes_expired_grace_key() {
assert!(namespace_keys(&mut connection, &namespace).await.is_empty());
}
#[tokio::test]
async fn scoped_cooldown_isolated_by_model_and_revision_fenced() {
let Some((repository, mut connection, namespace)) = repository().await else {
return;
};
let account_id = ProviderAccountId::new("acct_scoped_cooldown").expect("valid account ID");
let revision = CredentialRevision::new(2).expect("positive revision");
let model_a = ProviderCooldownScope::upstream_model(
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
);
let model_b = ProviderCooldownScope::upstream_model(
gateway_core::routing::UpstreamModelId::new("grok-4.6").expect("model"),
);
let until_a: SystemTime = millisecond_precision(Utc::now() + Duration::seconds(30)).into();
let until_b: SystemTime = millisecond_precision(Utc::now() + Duration::seconds(60)).into();
assert!(
repository
.put_scoped_if_later(ProviderScopedCooldown::new(
account_id.clone(),
revision,
model_a.clone(),
until_a,
))
.await
.expect("cache model A cooldown")
);
assert!(
!repository
.put_scoped_if_later(ProviderScopedCooldown::new(
account_id.clone(),
CredentialRevision::new(1).expect("stale revision"),
model_a.clone(),
until_b,
))
.await
.expect("reject stale model A cooldown")
);
assert!(
repository
.put_scoped_if_later(ProviderScopedCooldown::new(
account_id.clone(),
revision,
model_b.clone(),
until_b,
))
.await
.expect("cache model B cooldown")
);
assert_eq!(
repository
.read_scoped(&account_id, &model_a)
.await
.expect("read model A cooldown")
.expect("model A cooldown")
.until(),
until_a
);
assert_eq!(
repository
.read_scoped(&account_id, &model_b)
.await
.expect("read model B cooldown")
.expect("model B cooldown")
.until(),
until_b
);
assert!(
!repository
.clear_scoped(
&account_id,
&model_a,
CredentialRevision::new(1).expect("stale revision"),
)
.await
.expect("fence stale model A invalidation")
);
assert!(
repository
.clear_scoped(&account_id, &model_a, revision)
.await
.expect("clear model A cooldown")
);
assert!(
repository
.read_scoped(&account_id, &model_a)
.await
.expect("read cleared model A cooldown")
.is_none()
);
assert!(
repository
.read_scoped(&account_id, &model_b)
.await
.expect("read retained model B cooldown")
.is_some()
);
let keys = namespace_keys(&mut connection, &namespace).await;
assert_eq!(keys.len(), 1);
assert!(!keys[0].contains(account_id.as_str()));
assert!(!keys[0].contains(model_b.value()));
}
fn cooldown(account_id: &str, revision: u64, seconds: i64) -> CredentialCooldown {
CredentialCooldown {
provider_account_id: account_id.to_owned(),
@@ -3,8 +3,8 @@ use std::{sync::Arc, time::Duration};
use gateway_core::{
engine::credential::OpaqueProviderData,
provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome,
OAuthPendingConsumeOutcome, OAuthPendingFlowPort, OAuthPendingPutOutcome,
},
routing::ProviderKind,
};
@@ -23,6 +23,7 @@ async fn codex_pending_flow_wrong_owner_does_not_consume_it() {
OAuthPendingBinding::try_new(format!("flow-{}", Uuid::new_v4())).expect("flow binding");
let owner = OAuthPendingBinding::try_new("owner-a").expect("owner binding");
let other_owner = OAuthPendingBinding::try_new("owner-b").expect("other owner binding");
let claim = OAuthPendingBinding::try_new("claim-a").expect("claim binding");
let mut payload = serde_json::Map::new();
payload.insert(
"state".to_owned(),
@@ -77,24 +78,37 @@ async fn codex_pending_flow_wrong_owner_does_not_consume_it() {
assert!(ttl > 0 && ttl <= 60_000);
assert_eq!(
repository
.take_if_owner(&provider, &flow, &other_owner)
.claim_if_owner(
&provider,
&flow,
&other_owner,
&claim,
Duration::from_secs(30),
)
.await
.expect("reject wrong owner"),
OAuthPendingTakeOutcome::OwnerMismatch
OAuthPendingClaimOutcome::OwnerMismatch
);
assert!(matches!(
repository
.take_if_owner(&provider, &flow, &owner)
.claim_if_owner(&provider, &flow, &owner, &claim, Duration::from_secs(30),)
.await
.expect("consume right owner"),
OAuthPendingTakeOutcome::Taken(_)
.expect("claim right owner"),
OAuthPendingClaimOutcome::Claimed(_)
));
assert_eq!(
repository
.take_if_owner(&provider, &flow, &owner)
.consume_claim(&provider, &flow, &owner, &claim)
.await
.expect("consume right owner claim"),
OAuthPendingConsumeOutcome::Consumed
);
assert_eq!(
repository
.claim_if_owner(&provider, &flow, &owner, &claim, Duration::from_secs(30),)
.await
.expect("flow is one shot"),
OAuthPendingTakeOutcome::NotFound
OAuthPendingClaimOutcome::NotFound
);
}
+140 -96
View File
@@ -31,8 +31,9 @@ use gateway_core::engine::credential::{
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use secrecy::{ExposeSecret as _, SecretString};
@@ -47,10 +48,11 @@ use crate::credential::{
CodexAccountQuotaSnapshot, CodexCredentialAdmin, CodexCredentialAdminError,
CodexCredentialAdminService, CodexCredentialCatalogError, CodexCredentialCatalogService,
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexPendingAuthorization,
CodexQuotaWindowKind, CodexQuotaWindowRole, CompleteCodexOAuthAuthorization,
CompletedCodexOAuthCredential, ExportManagedCodexCredential, RotateManagedCodexCredential,
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization, refresh_time,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexPendingAuthorization, CodexQuotaWindowKind, CodexQuotaWindowRole,
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
RotateManagedCodexCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
refresh_time,
};
use crate::transport::CodexWebSocketPool;
use crate::transport::openai_billing_breakdown;
@@ -59,7 +61,7 @@ use crate::transport::profile::{
};
const PROVIDER_NAME: &str = "openai";
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 2;
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 3;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
@@ -144,7 +146,7 @@ impl OpenAiAdminProvider {
}
let current = self
.accounts
.load_credential(existing.id(), existing.revision())
.load_current_credential(existing.id())
.await
.map_err(map_store_error)?;
incoming.credential = CodexCredentialCodec::preserve_installation_id(
@@ -316,29 +318,37 @@ impl ProviderAdmin for OpenAiAdminProvider {
);
})
.map_err(map_oauth_error)?;
let credential = match completed.credential {
CompletedCodexOAuthCredential::Create(credential) => {
PreparedAuthorizationCredential::Create(prepared_create(credential, Utc::now())?)
let (mutation, completed_credential, authorization_guard) = completed.into_parts();
let credential = match (mutation.target(), completed_credential) {
(
AuthorizationMutationTarget::Create { .. },
CompletedCodexOAuthCredential::Create(credential),
) => {
prepared_create(credential, Utc::now()).map(PreparedAuthorizationCredential::Create)
}
CompletedCodexOAuthCredential::Reauthorize(credential) => {
let provider_kind = match completed.mutation.target() {
AuthorizationMutationTarget::Reauthorize { .. } => {
completed.mutation.provider_kind().clone()
}
AuthorizationMutationTarget::Create { .. } => {
return Err(provider_admin_error(ProviderAdminErrorKind::Internal));
}
};
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
credential,
provider_kind,
)?)
(
AuthorizationMutationTarget::Reauthorize { .. },
CompletedCodexOAuthCredential::Reauthorize(credential),
) => prepared_rotation(credential, mutation.provider_kind().clone())
.map(PreparedAuthorizationCredential::Reauthorize),
_ => Err(provider_admin_error(ProviderAdminErrorKind::Internal)),
};
let credential = match credential {
Ok(credential) => credential,
Err(error) => {
if let Err(settlement_error) = authorization_guard.abort().await {
tracing::warn!(
request_id = %request_id,
settlement_error = %settlement_error,
"OpenAI OAuth claim release failed after credential preparation"
);
return Err(provider_admin_error(ProviderAdminErrorKind::Unavailable));
}
return Err(error);
}
};
Ok(PreparedAuthorizationCommit {
pending: completed.mutation,
credential,
})
Ok(PreparedAuthorizationCommit::new(mutation, credential)
.with_authorization_guard(authorization_guard))
}
async fn prepare_rotation(
@@ -346,16 +356,11 @@ impl ProviderAdmin for OpenAiAdminProvider {
command: PrepareCredentialRotation,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
validate_account_record(&command.account, &self.provider_kind)?;
if command.account.credential_revision != command.expected_credential_revision {
return Err(provider_admin_error(ProviderAdminErrorKind::Conflict));
}
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let expected_revision = CredentialRevision::new(command.expected_credential_revision.get())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let current = self
.accounts
.load_credential(&account_id, expected_revision)
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
@@ -415,11 +420,17 @@ impl ProviderAdmin for OpenAiAdminProvider {
validate_account_record(&command.account, &self.provider_kind)?;
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let revision = CredentialRevision::new(command.account.credential_revision.get())
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
let current = self
.accounts
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
return Err(provider_admin_error(ProviderAdminErrorKind::Conflict));
}
let prepared = self
.credentials
.manual_refresh(account_id, revision)
.manual_refresh(current)
.await
.map_err(map_credential_admin_error)?;
prepared_rotation(prepared, command.account.provider_kind)
@@ -707,20 +718,9 @@ fn dashboard_desktop_release(
fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool {
account.id().as_str() == record.id
&& account.provider() == &record.provider_kind
&& account.revision().get() == record.credential_revision.get()
&& account.name() == record.name
&& account.email() == record.email.as_deref()
&& account.upstream_user_id() == record.upstream_user_id
&& account.upstream_account_id() == record.upstream_account_id.as_deref()
&& account.plan_type() == record.plan_type.as_deref()
&& account.authentication_kind() == record.authentication_kind
&& account.enabled() == record.enabled
&& account.availability() == record.availability
&& account.cooldown_until().map(DateTime::<Utc>::from) == record.cooldown_until
&& account.access_token_expires_at().map(DateTime::<Utc>::from)
== record.access_token_expires_at
&& account.next_refresh_at().map(DateTime::<Utc>::from) == record.next_refresh_at
&& account.has_refresh_token() == record.has_refresh_token
}
fn empty_quota() -> ProviderQuota {
@@ -896,21 +896,86 @@ impl CodexOAuthPendingStore for OpenAiOAuthPendingStore {
}
}
async fn take(
async fn claim(
&self,
owner_ref: &str,
flow_id: &str,
) -> Result<Option<CodexPendingAuthorization>, CodexOAuthPendingStoreError> {
claim_ref: &str,
claim_ttl: std::time::Duration,
) -> Result<CodexOAuthPendingClaimOutcome, CodexOAuthPendingStoreError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
let outcome = self
.port
.claim_if_owner(&self.provider_kind, &flow, &owner, &claim, claim_ttl)
.await
.map_err(map_pending_store_error)?;
match outcome {
OAuthPendingClaimOutcome::Claimed(payload) => match decode_pending(payload) {
Ok(pending) => Ok(CodexOAuthPendingClaimOutcome::Claimed(Box::new(pending))),
Err(error) => match self
.port
.release_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_pending_store_error)?
{
OAuthPendingReleaseOutcome::Released => Err(error),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => {
Err(CodexOAuthPendingStoreError::Unavailable)
}
},
},
OAuthPendingClaimOutcome::NotFound | OAuthPendingClaimOutcome::OwnerMismatch => {
Ok(CodexOAuthPendingClaimOutcome::NotFound)
}
OAuthPendingClaimOutcome::InProgress => Ok(CodexOAuthPendingClaimOutcome::InProgress),
}
}
async fn release_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
match self
.port
.take_if_owner(&self.provider_kind, &flow, &owner)
.release_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_pending_store_error)?
{
OAuthPendingTakeOutcome::Taken(payload) => decode_pending(payload).map(Some),
OAuthPendingTakeOutcome::NotFound | OAuthPendingTakeOutcome::OwnerMismatch => Ok(None),
OAuthPendingReleaseOutcome::Released => Ok(true),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => Ok(false),
}
}
async fn consume_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
match self
.port
.consume_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_pending_store_error)?
{
OAuthPendingConsumeOutcome::Consumed => Ok(true),
OAuthPendingConsumeOutcome::NotFound
| OAuthPendingConsumeOutcome::OwnerMismatch
| OAuthPendingConsumeOutcome::ClaimMismatch => Ok(false),
}
}
}
@@ -927,7 +992,6 @@ struct PendingDocument {
nonce: String,
code_verifier: String,
reauthorization_account_id: Option<String>,
reauthorization_credential_revision: Option<u64>,
mutation: PendingMutationDocument,
}
@@ -944,13 +1008,8 @@ struct PendingMutationDocument {
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum PendingTargetDocument {
Create {
name: String,
},
Reauthorize {
account_id: String,
expected_credential_revision: u64,
},
Create { name: String },
Reauthorize { account_id: String },
}
#[derive(Deserialize)]
@@ -994,15 +1053,9 @@ fn encode_pending(pending: &CodexPendingAuthorization) -> Map<String, Value> {
);
document.insert(
"reauthorization_account_id".to_owned(),
pending.reauthorization().map_or(Value::Null, |target| {
Value::String(target.account_id().to_string())
}),
);
document.insert(
"reauthorization_credential_revision".to_owned(),
pending.reauthorization().map_or(Value::Null, |target| {
Value::Number(Number::from(target.credential_revision().get()))
}),
pending
.reauthorization()
.map_or(Value::Null, |target| Value::String(target.to_string())),
);
document.insert(
"mutation".to_owned(),
@@ -1043,19 +1096,12 @@ fn encode_target(target: &AuthorizationMutationTarget) -> Map<String, Value> {
document.insert("kind".to_owned(), Value::String("create".to_owned()));
document.insert("name".to_owned(), Value::String(name.clone()));
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
AuthorizationMutationTarget::Reauthorize { account_id } => {
document.insert("kind".to_owned(), Value::String("reauthorize".to_owned()));
document.insert(
"account_id".to_owned(),
Value::String(account_id.to_string()),
);
document.insert(
"expected_credential_revision".to_owned(),
Value::Number(Number::from(expected_credential_revision.get())),
);
}
}
document
@@ -1096,7 +1142,6 @@ fn decode_pending(
nonce: SecretString::from(document.nonce),
code_verifier: SecretString::from(document.code_verifier),
reauthorization_account_id: document.reauthorization_account_id,
reauthorization_credential_revision: document.reauthorization_credential_revision,
mutation: decode_mutation(document.mutation)?,
})
}
@@ -1111,15 +1156,12 @@ fn decode_mutation(
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?;
let target = match document.target {
PendingTargetDocument::Create { name } => AuthorizationMutationTarget::Create { name },
PendingTargetDocument::Reauthorize {
account_id,
expected_credential_revision,
} => AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
expected_credential_revision: Revision::new(expected_credential_revision)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
},
PendingTargetDocument::Reauthorize { account_id } => {
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
}
}
};
let actor = match document.owner {
PendingOwnerDocument::AdminSession { admin_user_id } => {
@@ -1222,10 +1264,10 @@ fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdmin
| Error::AccountBanned
| Error::IdentityRejected => ProviderAdminErrorKind::Invalid,
Error::NotFound => ProviderAdminErrorKind::NotFound,
Error::RevisionConflict | Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => {
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => {
ProviderAdminErrorKind::Conflict
}
Error::StoreUnavailable | Error::RefreshUnavailable | Error::IdentityUnavailable => {
Error::RefreshUnavailable | Error::IdentityUnavailable => {
ProviderAdminErrorKind::Unavailable
}
})
@@ -1237,8 +1279,6 @@ const fn credential_admin_error_code(error: CodexCredentialAdminError) -> &'stat
CodexCredentialAdminError::IdentityMismatch => "identity_mismatch",
CodexCredentialAdminError::InvalidCredential => "invalid_credential",
CodexCredentialAdminError::NotFound => "not_found",
CodexCredentialAdminError::RevisionConflict => "revision_conflict",
CodexCredentialAdminError::StoreUnavailable => "store_unavailable",
CodexCredentialAdminError::MissingRefreshToken => "missing_refresh_token",
CodexCredentialAdminError::RefreshLeaseUnavailable => "refresh_lease_unavailable",
CodexCredentialAdminError::RefreshRejected => "refresh_rejected",
@@ -1272,9 +1312,11 @@ fn log_import_failure(stage: &'static str, error: &'static str) {
fn map_oauth_error(error: CodexOAuthAdminError) -> ProviderAdminError {
use CodexOAuthAdminError as Error;
provider_admin_error(match error {
Error::InvalidInput | Error::UpstreamRejected | Error::Credential => {
ProviderAdminErrorKind::Invalid
}
Error::InvalidInput
| Error::CallbackRejected
| Error::TokenRejected
| Error::IdentityRejected
| Error::Credential => ProviderAdminErrorKind::Invalid,
Error::NotFound | Error::FlowExpired => ProviderAdminErrorKind::NotFound,
Error::Conflict | Error::Ambiguous => ProviderAdminErrorKind::Conflict,
Error::UpstreamUnavailable | Error::StorageUnavailable => {
@@ -1289,7 +1331,9 @@ const fn oauth_error_code(error: &CodexOAuthAdminError) -> &'static str {
CodexOAuthAdminError::NotFound => "not_found",
CodexOAuthAdminError::Conflict => "conflict",
CodexOAuthAdminError::FlowExpired => "flow_expired",
CodexOAuthAdminError::UpstreamRejected => "upstream_rejected",
CodexOAuthAdminError::CallbackRejected => "callback_rejected",
CodexOAuthAdminError::TokenRejected => "token_rejected",
CodexOAuthAdminError::IdentityRejected => "identity_rejected",
CodexOAuthAdminError::UpstreamUnavailable => "upstream_unavailable",
CodexOAuthAdminError::Ambiguous => "ambiguous",
CodexOAuthAdminError::StorageUnavailable => "storage_unavailable",
@@ -13,7 +13,6 @@ use gateway_core::engine::credential::{
AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountUpdate,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy,
@@ -28,7 +27,6 @@ use thiserror::Error;
use super::identity::{
CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError,
};
use super::repository::CodexCredentialRepository;
use super::security::CodexCredentialCodec;
use super::token_client::{RefreshFailure, TokenRefresher};
use super::types::{
@@ -42,32 +40,6 @@ const CODEX_CREDENTIAL_SCHEMA_VERSION: u32 = 1;
const MAX_BATCH: usize = 200;
const MAX_IMPORT_DOCUMENT_BYTES: usize = 64 * 1024 * 1024;
const CPR_CONTAINER_KEYS: &[&str] = &["sourceFormat", "source_format", "accounts"];
const CPR_ACCOUNT_KEYS: &[&str] = &[
"id",
"email",
"accountId",
"account_id",
"userId",
"user_id",
"label",
"planType",
"plan_type",
"token",
"at",
"accessToken",
"access_token",
"refreshToken",
"refresh_token",
"accessTokenExpiresAt",
"access_token_expires_at",
"status",
"addedAt",
"added_at",
"updatedAt",
"updated_at",
];
pub struct ImportCodexOAuthCredential {
pub account_id: String,
pub name: String,
@@ -133,7 +105,6 @@ enum ParsedCodexAuthentication {
OAuth {
access_token: Option<String>,
refresh_token: Option<String>,
id_token: Option<String>,
},
AgentIdentity {
runtime_id: String,
@@ -148,7 +119,6 @@ impl fmt::Debug for ParsedCodexAuthentication {
Self::OAuth {
access_token,
refresh_token,
id_token,
} => formatter
.debug_struct("OAuth")
.field("access_token", &access_token.as_ref().map(|_| "<redacted>"))
@@ -156,7 +126,6 @@ impl fmt::Debug for ParsedCodexAuthentication {
"refresh_token",
&refresh_token.as_ref().map(|_| "<redacted>"),
)
.field("id_token", &id_token.as_ref().map(|_| "<redacted>"))
.finish(),
Self::AgentIdentity {
runtime_id: _,
@@ -379,10 +348,6 @@ pub enum CodexCredentialAdminError {
InvalidCredential,
#[error("Codex account was not found")]
NotFound,
#[error("Codex credential revision is stale")]
RevisionConflict,
#[error("Codex account store is unavailable")]
StoreUnavailable,
#[error("Codex account has no refresh token")]
MissingRefreshToken,
#[error("Codex refresh lease is unavailable")]
@@ -580,9 +545,8 @@ impl CodexCredentialAdmin {
}
}
/// 有状态的 Codex 手工刷新边界;只读取 Store 并准备 CAS,不自行持久化。
/// 有状态的 Codex 手工刷新边界;消费调用方刚读取的当前 credential 并准备 CAS。
pub struct CodexCredentialAdminService {
repository: CodexCredentialRepository,
refresher: Arc<dyn TokenRefresher>,
verifier: Arc<dyn CodexAccountIdentityVerifier>,
leases: Arc<dyn ProviderLeasePort>,
@@ -593,7 +557,6 @@ impl fmt::Debug for CodexCredentialAdminService {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexCredentialAdminService")
.field("repository", &"CodexCredentialRepository")
.field("refresher", &"TokenRefresher")
.field("verifier", &"CodexAccountIdentityVerifier")
.field("leases", &"ProviderLeasePort")
@@ -604,14 +567,12 @@ impl fmt::Debug for CodexCredentialAdminService {
impl CodexCredentialAdminService {
pub fn new(
repository: CodexCredentialRepository,
refresher: Arc<dyn TokenRefresher>,
verifier: Arc<dyn CodexAccountIdentityVerifier>,
leases: Arc<dyn ProviderLeasePort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
) -> Self {
Self {
repository,
refresher,
verifier,
leases,
@@ -622,19 +583,11 @@ impl CodexCredentialAdminService {
/// 官方 RT exchange + AT 身份验证;结果由 App 在同一 revision/audit 事务中提交。
pub async fn manual_refresh(
&self,
account_id: ProviderAccountId,
expected_revision: CredentialRevision,
current: LoadedCredential,
) -> Result<PreparedCodexCredentialRotation, CodexCredentialAdminError> {
let current = self
.repository
.store()
.load_credential(&account_id, expected_revision)
.await
.map_err(map_admin_store_error)?;
if current.account.provider().as_str() != PROVIDER_NAME
|| current.account.id() != &account_id
|| current.account.revision() != expected_revision
{
let account_id = current.account.id().clone();
let expected_revision = current.account.revision();
if current.account.provider().as_str() != PROVIDER_NAME {
return Err(CodexCredentialAdminError::NotFound);
}
let runtime = CodexCredentialCodec::decode(&current.credential)
@@ -720,7 +673,7 @@ impl CodexCredentialAdminService {
Ok(prepared)
}
/// 按文档结构严格判别账号布局,并归一到唯一 `NewProviderAccount` 写入路径。
/// 从 OAuth token 对补全账号事实,并归一到唯一 `NewProviderAccount` 写入路径。
pub async fn prepare_import_document(
&self,
payload: Value,
@@ -741,19 +694,18 @@ impl CodexCredentialAdminService {
let mut accounts = Vec::with_capacity(candidates.len());
let mut policy = None;
for candidate in candidates {
let account_id = candidate
.id
.clone()
.filter(|id| ProviderAccountId::new(id.clone()).is_ok())
.unwrap_or_else(|| format!("acct_{}", uuid::Uuid::now_v7().simple()));
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (enabled, availability) = import_runtime_state(candidate.status.as_deref())?;
let mut prepared = match &candidate.authentication {
ParsedCodexAuthentication::OAuth { .. } => {
let (secret, verified_account) =
self.verify_import_candidate(&candidate).await?;
let prepared = match &candidate.authentication {
ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
} => {
let account_id = format!("acct_{}", uuid::Uuid::now_v7().simple());
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (secret, verified_account) = self
.verify_import_tokens(access_token.clone(), refresh_token.clone())
.await?;
if !upstream_identities.insert((
verified_account.chatgpt_user_id.clone(),
verified_account.chatgpt_account_id.clone(),
@@ -780,19 +732,29 @@ impl CodexCredentialAdminService {
)?;
CodexCredentialAdmin.prepare_import(ImportCodexOAuthCredential {
account_id,
name: candidate
.name
name: verified_account
.email
.clone()
.filter(|name| !name.trim().is_empty())
.unwrap_or_else(|| "Codex OAuth".to_owned()),
secret,
verified_account,
next_refresh_at,
enabled,
enabled: true,
})?
}
ParsedCodexAuthentication::AgentIdentity { .. } => {
let account =
let account_id = candidate
.id
.clone()
.filter(|id| ProviderAccountId::new(id.clone()).is_ok())
.unwrap_or_else(|| format!("acct_{}", uuid::Uuid::now_v7().simple()));
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (enabled, availability) =
import_runtime_state(candidate.status.as_deref())?;
let mut account =
self.prepare_agent_identity_import(&candidate, account_id, enabled)?;
if !upstream_identities.insert((
account.account.upstream_user_id().to_owned(),
@@ -804,29 +766,23 @@ impl CodexCredentialAdminService {
)) {
return Err(CodexCredentialAdminError::InvalidInput);
}
account.account =
account
.account
.with_runtime_state(enabled, availability, None);
account
}
};
prepared.account = prepared
.account
.with_runtime_state(enabled, availability, None);
accounts.push(prepared);
}
Ok(PreparedCodexAccountImport { accounts })
}
async fn verify_import_candidate(
async fn verify_import_tokens(
&self,
candidate: &ParsedCodexImportAccount,
access_token: Option<String>,
refresh_token: Option<String>,
) -> Result<(CodexOAuthSecret, CodexAccountProfile), CodexCredentialAdminError> {
let ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token,
} = &candidate.authentication
else {
return Err(CodexCredentialAdminError::InvalidCredential);
};
let access_token = access_token
.as_deref()
.map(normalize_bearer)
@@ -835,26 +791,14 @@ impl CodexCredentialAdminService {
.as_deref()
.map(normalize_bearer)
.filter(|token| !token.is_empty());
let id_token = id_token
.as_deref()
.map(normalize_bearer)
.filter(|token| !token.is_empty());
let expectation = CodexIdentityExpectation::imported(
candidate.chatgpt_account_id.clone(),
candidate.chatgpt_user_id.clone(),
)
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
if let Some(access_token) = access_token {
let secret = CodexOAuthSecret {
access_token: SecretString::from(access_token),
refresh_token: refresh_token.clone().map(SecretString::from),
id_token: id_token.clone().map(SecretString::from),
id_token: None,
};
match self.verifier.verify(&secret, &expectation).await {
Ok(verification) => {
let profile = verification.into_complete().map_err(map_identity_error)?;
return Ok((secret, profile));
}
match self.verifier.verify_import(&secret).await {
Ok(profile) => return Ok((secret, profile)),
Err(CodexIdentityVerificationError::Unavailable) => {
return Err(CodexCredentialAdminError::IdentityUnavailable);
}
@@ -873,17 +817,19 @@ impl CodexCredentialAdminService {
if tokens.access_token.trim().is_empty() || tokens.expires_in.is_zero() {
return Err(CodexCredentialAdminError::InvalidCredential);
}
// 导入时已经消费来源 RT;只有上游明确轮换返回的新 RT 才能继续保存。
let secret = CodexOAuthSecret {
access_token: SecretString::from(tokens.access_token),
refresh_token: tokens.refresh_token.map(SecretString::from),
id_token: id_token.map(SecretString::from),
// RT 未轮换时仍保留导入提供的 RT,保证本次补全不会丢失后续刷新能力。
refresh_token: tokens
.refresh_token
.map(SecretString::from)
.or_else(|| Some(SecretString::from(refresh_token))),
id_token: None,
};
let profile = self
.verifier
.verify(&secret, &expectation)
.verify_import(&secret)
.await
.and_then(super::identity::CodexIdentityVerification::into_complete)
.map_err(map_identity_error)?;
Ok((secret, profile))
}
@@ -1026,17 +972,6 @@ fn china_rfc3339(value: DateTime<Utc>) -> String {
.to_rfc3339()
}
fn map_admin_store_error(error: gateway_core::error::StoreError) -> CodexCredentialAdminError {
match error.kind() {
StoreErrorKind::Conflict => CodexCredentialAdminError::RevisionConflict,
StoreErrorKind::Unavailable => CodexCredentialAdminError::StoreUnavailable,
StoreErrorKind::InvalidState | StoreErrorKind::InvalidData => {
CodexCredentialAdminError::NotFound
}
_ => CodexCredentialAdminError::StoreUnavailable,
}
}
const fn map_refresh_failure(error: RefreshFailure) -> CodexCredentialAdminError {
match error {
RefreshFailure::InvalidGrant => CodexCredentialAdminError::RefreshRejected,
@@ -1058,67 +993,29 @@ const fn map_identity_error(error: CodexIdentityVerificationError) -> CodexCrede
fn parse_import_document(
payload: &Value,
) -> Result<Vec<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let shape = if looks_like_agent_identity_document(payload) {
ImportDocumentShape::AgentIdentity
} else if looks_like_credential_bundle(payload) {
ImportDocumentShape::CredentialBundle
} else if looks_like_auth_document(payload) {
ImportDocumentShape::AuthDocument
} else {
ImportDocumentShape::Native
};
if shape == ImportDocumentShape::Native
&& payload.get("accounts").is_some()
&& payload.as_object().is_none_or(|object| {
object
.keys()
.any(|key| !CPR_CONTAINER_KEYS.contains(&key.as_str()))
})
{
return Err(CodexCredentialAdminError::InvalidInput);
}
let values = import_account_values(payload)?;
let mut accounts = Vec::new();
for value in values {
// CPR 导出允许 OAuth 与 Agent Identity 共存;按账号选择 Provider
// 解析器,避免一个 Agent Identity 条目把整份文档误判成单一格式。
let account_shape = if looks_like_agent_identity_account(value) {
ImportDocumentShape::AgentIdentity
} else if shape == ImportDocumentShape::AgentIdentity {
if value.get("credentials").is_some() || value.get("tokens").is_some() {
ImportDocumentShape::CredentialBundle
} else if value
.as_object()
.is_some_and(auth_document_provider_is_openai)
{
ImportDocumentShape::AuthDocument
} else {
ImportDocumentShape::Native
}
} else {
shape
};
let parsed = match account_shape {
ImportDocumentShape::Native => Some(parse_native_account(value)?),
ImportDocumentShape::AgentIdentity => Some(parse_agent_identity_account(value)?),
ImportDocumentShape::CredentialBundle => parse_credential_bundle_account(value)?,
ImportDocumentShape::AuthDocument => parse_auth_document_account(value)?,
};
if let Some(parsed) = parsed {
accounts.push(parsed);
for value in import_account_values(payload)? {
if looks_like_agent_identity_account(value) {
accounts.push(parse_agent_identity_account(value)?);
continue;
}
if !is_openai_oauth_candidate(value) {
continue;
}
accounts.push(ParsedCodexImportAccount {
id: None,
name: None,
email: None,
plan_type: None,
chatgpt_account_id: None,
chatgpt_user_id: None,
authentication: parse_oauth_import_tokens(value)?,
status: None,
});
}
Ok(accounts)
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum ImportDocumentShape {
Native,
AgentIdentity,
CredentialBundle,
AuthDocument,
}
fn parse_agent_identity_account(
value: &Value,
) -> Result<ParsedCodexImportAccount, CodexCredentialAdminError> {
@@ -1189,274 +1086,44 @@ fn import_account_values(payload: &Value) -> Result<Vec<&Value>, CodexCredential
Ok(vec![payload])
}
fn parse_native_account(
fn parse_oauth_import_tokens(
value: &Value,
) -> Result<ParsedCodexImportAccount, CodexCredentialAdminError> {
if value.get("accounts").is_some() {
return Err(CodexCredentialAdminError::InvalidInput);
) -> Result<ParsedCodexAuthentication, CodexCredentialAdminError> {
let mut access_token = None;
let mut refresh_token = None;
let mut pending = vec![value];
while let Some(current) = pending.pop() {
match current {
Value::Object(object) => {
for (key, value) in object {
let token = match key.as_str() {
"access_token" | "accessToken" | "at" | "token" => &mut access_token,
"refresh_token" | "refreshToken" | "rt" => &mut refresh_token,
_ => {
pending.push(value);
continue;
}
};
let value = value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or(CodexCredentialAdminError::InvalidCredential)?;
if token.replace(value.to_owned()).is_some() {
return Err(CodexCredentialAdminError::InvalidInput);
}
}
}
Value::Array(values) => pending.extend(values),
_ => {}
}
}
let account = value
.as_object()
.ok_or(CodexCredentialAdminError::InvalidInput)?;
if account
.keys()
.any(|key| !CPR_ACCOUNT_KEYS.contains(&key.as_str()))
{
return Err(CodexCredentialAdminError::InvalidInput);
}
let access_token = first_string(value, &["token", "at", "accessToken", "access_token"]);
let refresh_token = first_string(value, &["refreshToken", "refresh_token"]);
if access_token.is_none() && refresh_token.is_none() {
return Err(CodexCredentialAdminError::InvalidCredential);
}
Ok(ParsedCodexImportAccount {
id: first_string(value, &["id"]),
name: first_string(value, &["label", "email"]),
email: first_string(value, &["email"]),
plan_type: first_string(value, &["planType", "plan_type"]),
chatgpt_account_id: first_string(value, &["accountId", "account_id"]),
chatgpt_user_id: first_string(value, &["userId", "user_id"]),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: None,
},
status: first_string(value, &["status"]),
})
}
fn parse_credential_bundle_account(
value: &Value,
) -> Result<Option<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let account = value
.as_object()
.ok_or(CodexCredentialAdminError::InvalidInput)?;
if let Some(credentials) = value.get("credentials") {
if account
.get("platform")
.and_then(Value::as_str)
.is_some_and(|value| !value.eq_ignore_ascii_case("openai"))
|| account
.get("type")
.and_then(Value::as_str)
.is_some_and(|value| !value.eq_ignore_ascii_case("oauth"))
{
return Ok(None);
}
let access_token = first_path_string(
credentials,
&[
&["access_token"],
&["accessToken"],
&["at"],
&["token", "access_token"],
&["token", "accessToken"],
&["token", "at"],
],
);
let refresh_token = first_path_string(
credentials,
&[
&["refresh_token"],
&["refreshToken"],
&["rt"],
&["token", "refresh_token"],
&["token", "refreshToken"],
&["token", "rt"],
],
);
if access_token.is_none() && refresh_token.is_none() {
return Ok(None);
}
return Ok(Some(ParsedCodexImportAccount {
id: first_string(value, &["id"]),
name: first_string(value, &["label", "name"]),
email: first_path_string(credentials, &[&["email"], &["user", "email"]]),
plan_type: first_path_string(
credentials,
&[&["plan_type"], &["planType"], &["plan", "type"]],
),
chatgpt_account_id: first_path_string(
credentials,
&[
&["chatgpt_account_id"],
&["chatgptAccountId"],
&["account_id"],
&["accountId"],
&["account", "id"],
&["account", "account_id"],
&["account", "chatgpt_account_id"],
],
),
chatgpt_user_id: first_path_string(
credentials,
&[
&["chatgpt_user_id"],
&["chatgptUserId"],
&["user_id"],
&["userId"],
&["user", "id"],
&["account", "user_id"],
],
),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: first_path_string(credentials, &[&["id_token"], &["idToken"]]),
},
status: first_string(value, &["status"]),
}));
}
let access_token = first_path_string(
value,
&[
&["tokens", "access_token"],
&["tokens", "accessToken"],
&["tokens", "at"],
&["access_token"],
&["accessToken"],
&["token"],
&["at"],
],
);
let refresh_token = first_path_string(
value,
&[
&["tokens", "refresh_token"],
&["tokens", "refreshToken"],
&["tokens", "rt"],
&["refresh_token"],
&["refreshToken"],
&["rt"],
],
);
if access_token.is_none() && refresh_token.is_none() {
return Ok(None);
}
Ok(Some(ParsedCodexImportAccount {
id: first_string(value, &["id"]),
name: first_path_string(value, &[&["label"], &["name"], &["user", "name"]]),
email: first_path_string(value, &[&["email"], &["user", "email"]]),
plan_type: first_path_string(value, &[&["plan_type"], &["planType"], &["plan", "type"]]),
chatgpt_account_id: first_path_string(
value,
&[
&["chatgpt_account_id"],
&["chatgptAccountId"],
&["account_id"],
&["accountId"],
&["account", "id"],
&["account", "account_id"],
&["account", "chatgpt_account_id"],
],
),
chatgpt_user_id: first_path_string(
value,
&[
&["chatgpt_user_id"],
&["chatgptUserId"],
&["user_id"],
&["userId"],
&["user", "id"],
&["account", "user_id"],
],
),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: first_path_string(
value,
&[&["tokens", "id_token"], &["id_token"], &["idToken"]],
),
},
status: first_string(value, &["status"]),
}))
}
fn parse_auth_document_account(
value: &Value,
) -> Result<Option<ParsedCodexImportAccount>, CodexCredentialAdminError> {
let account = value
.as_object()
.ok_or(CodexCredentialAdminError::InvalidInput)?;
if !auth_document_provider_is_openai(account) {
return Ok(None);
}
let access_token = first_path_string(
value,
&[
&["access_token"],
&["accessToken"],
&["at"],
&["token", "access_token"],
&["token", "accessToken"],
&["token", "at"],
&["token"],
],
);
let refresh_token = first_path_string(
value,
&[
&["refresh_token"],
&["refreshToken"],
&["rt"],
&["token", "refresh_token"],
&["token", "refreshToken"],
&["token", "rt"],
],
);
if access_token.is_none() && refresh_token.is_none() {
return Ok(None);
}
Ok(Some(ParsedCodexImportAccount {
id: None,
name: first_string(value, &["label", "name", "email"]),
email: first_string(value, &["email"]),
plan_type: first_string(value, &["plan_type", "planType"]),
chatgpt_account_id: first_string(
value,
&[
"chatgpt_account_id",
"chatgptAccountId",
"account_id",
"accountId",
],
),
chatgpt_user_id: first_string(
value,
&["chatgpt_user_id", "chatgptUserId", "user_id", "userId"],
),
authentication: ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
id_token: first_string(value, &["id_token", "idToken"]),
},
status: account
.get("disabled")
.and_then(Value::as_bool)
.filter(|disabled| *disabled)
.map(|_| "disabled".to_owned()),
}))
}
fn looks_like_credential_bundle(value: &Value) -> bool {
value.get("exported_at").is_some()
|| value.get("proxies").is_some()
|| import_account_values(value).is_ok_and(|accounts| {
accounts.iter().any(|account| {
account.get("credentials").is_some()
|| account.get("tokens").is_some()
|| account.get("cachedQuota").is_some()
|| account.get("cached_quota").is_some()
})
})
}
fn looks_like_agent_identity_document(value: &Value) -> bool {
import_account_values(value).is_ok_and(|accounts| {
accounts
.iter()
.any(|account| looks_like_agent_identity_account(account))
Ok(ParsedCodexAuthentication::OAuth {
access_token,
refresh_token,
})
}
@@ -1471,26 +1138,24 @@ fn looks_like_agent_identity_account(value: &Value) -> bool {
|| first_string(identity, &["agent_runtime_id", "agentRuntimeId"]).is_some()
}
fn looks_like_auth_document(value: &Value) -> bool {
import_account_values(value).is_ok_and(|accounts| {
accounts.iter().any(|account| {
account
.as_object()
.is_some_and(auth_document_provider_is_openai)
})
})
}
fn auth_document_provider_is_openai(account: &serde_json::Map<String, Value>) -> bool {
account
.get("type")
fn is_openai_oauth_candidate(value: &Value) -> bool {
let Some(account) = value.as_object() else {
return false;
};
if let Some(provider) = account
.get("platform")
.or_else(|| account.get("provider"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|provider| {
provider.eq_ignore_ascii_case("openai")
// CLIProxyAPI 的 OpenAI Codex OAuth 文件使用 `type: codex`。
|| provider.eq_ignore_ascii_case("codex")
{
return provider.eq_ignore_ascii_case("openai") || provider.eq_ignore_ascii_case("codex");
}
account
.get("type")
.and_then(Value::as_str)
.is_none_or(|kind| {
kind.eq_ignore_ascii_case("openai")
|| kind.eq_ignore_ascii_case("codex")
|| kind.eq_ignore_ascii_case("oauth")
})
}
@@ -1502,23 +1167,6 @@ fn first_string(value: &Value, keys: &[&str]) -> Option<String> {
.map(str::to_owned)
}
fn first_path_string(value: &Value, paths: &[&[&str]]) -> Option<String> {
paths
.iter()
.find_map(|path| value_at_path(value, path).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_owned)
}
fn value_at_path<'a>(value: &'a Value, path: &[&str]) -> Option<&'a Value> {
let mut current = value;
for key in path {
current = current.get(*key)?;
}
Some(current)
}
fn normalize_bearer(value: &str) -> String {
let value = value.trim();
value
@@ -347,7 +347,7 @@ impl CodexAgentIdentityTaskService {
self.load_current(account_id).await
}
async fn load_current(
pub(crate) async fn load_current(
&self,
account_id: &ProviderAccountId,
) -> Result<PreparedCodexRuntimeCredential, CodexAgentIdentityError> {
@@ -116,6 +116,24 @@ impl CodexSignedIdentity {
pub const fn access_token_expires_at(&self) -> DateTime<Utc> {
self.access_token_expires_at
}
fn into_import_profile(self) -> Result<CodexAccountProfile, CodexIdentityVerificationError> {
let chatgpt_account_id = self
.claimed_account_id
.ok_or(CodexIdentityVerificationError::Rejected)?;
let chatgpt_user_id = self
.claimed_user_id
.ok_or(CodexIdentityVerificationError::Rejected)?;
Ok(CodexAccountProfile {
email: self.email,
oauth_subject: self.oauth_subject,
poid: self.poid,
chatgpt_account_id,
chatgpt_user_id,
plan_type: self.plan_type,
access_token_expires_at: Some(self.access_token_expires_at),
})
}
}
/// 导入文档或现有 credential 对认证结果施加的一致性约束。
@@ -256,6 +274,19 @@ pub trait CodexAccountIdentityVerifier: Send + Sync {
expectation: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError>;
/// 为无状态导入补全 OAuth 账号事实。
///
/// 导入不信任导出文件中的身份投影;实现必须只使用已验签的 token
/// 与官方认证接口返回的事实生成可持久化 profile。
async fn verify_import(
&self,
secret: &CodexOAuthSecret,
) -> Result<CodexAccountProfile, CodexIdentityVerificationError> {
self.verify(secret, &CodexIdentityExpectation::default())
.await?
.into_complete()
}
async fn verify_authorization(
&self,
secret: &CodexOAuthSecret,
@@ -580,6 +611,31 @@ impl CodexAccountIdentityVerifier for CodexAccountIdentityService {
self.complete(secret, signed, expectation).await
}
async fn verify_import(
&self,
secret: &CodexOAuthSecret,
) -> Result<CodexAccountProfile, CodexIdentityVerificationError> {
let signed = self.signed.verify_access(secret).await?;
let expectation = CodexIdentityExpectation::imported(
signed.claimed_account_id.clone(),
signed.claimed_user_id.clone(),
)?;
let fallback = signed.clone().into_import_profile().ok();
match self.accounts.fetch(secret, &expectation).await {
Ok(account) => match complete_profile(signed, account, &expectation) {
Ok(profile) => Ok(profile),
Err(_) => fallback.ok_or(CodexIdentityVerificationError::Rejected),
},
// `/wham/usage` 是导入时的补全来源,不是 token 有效性的裁决者。AT
// 已完成官方验签且带有完整账号声明时,接口拒绝、暂时不可用或返回默认
// workspace 的冲突投影,都不能阻止恢复该 OAuth 账号。
Err(
CodexIdentityVerificationError::Rejected
| CodexIdentityVerificationError::Unavailable,
) => fallback.ok_or(CodexIdentityVerificationError::Rejected),
}
}
async fn verify_authorization(
&self,
secret: &CodexOAuthSecret,
@@ -47,7 +47,7 @@ pub use identity::{
};
pub use oauth::{
CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthAuthorizationStarted,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthReauthorizationTarget,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthAuthorization,
CompletedCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
};
@@ -1,16 +1,19 @@
//! Codex Authorization Code + PKCE/OIDC 管理流。
use std::fmt;
use std::sync::Arc;
use std::{fmt, sync::Arc, time::Duration};
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwner, PendingAuthorizationMutation,
use gateway_admin::model::{
AdminError,
provider_credentials::{
AuthorizationCommitGuard, AuthorizationMutationTarget, AuthorizationOwner,
PendingAuthorizationMutation,
},
};
use gateway_core::engine::credential::{
CredentialRevision, NewProviderAccount, ProviderAccountId, ProviderAccountStore,
NewProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::provider_ports::ProviderRuntimePolicyPort;
use secrecy::{ExposeSecret, SecretString};
@@ -35,6 +38,7 @@ use super::token_client::{
const AUTHORIZATION_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
const AUTHORIZATION_SCOPE: &str = "openid profile email offline_access";
const AUTHORIZATION_TTL: TimeDelta = TimeDelta::minutes(10);
const AUTHORIZATION_CLAIM_TTL: Duration = Duration::from_secs(90);
const MAX_CALLBACK_BYTES: usize = 16 * 1024;
const MAX_TEXT_BYTES: usize = 512;
@@ -43,24 +47,6 @@ pub struct StartCodexOAuthAuthorization {
pub mutation: PendingAuthorizationMutation,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CodexOAuthReauthorizationTarget {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
}
impl CodexOAuthReauthorizationTarget {
#[must_use]
pub const fn account_id(&self) -> &ProviderAccountId {
&self.account_id
}
#[must_use]
pub const fn credential_revision(&self) -> CredentialRevision {
self.credential_revision
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct CodexOAuthAuthorizationStarted {
pub flow_id: String,
@@ -89,6 +75,7 @@ pub struct CompleteCodexOAuthAuthorization {
pub struct CompletedCodexOAuthAuthorization<T> {
pub mutation: PendingAuthorizationMutation,
pub credential: T,
authorization_guard: Box<dyn AuthorizationCommitGuard>,
}
/// OAuth exchange 后唯一的 credential preparation 结果。
@@ -116,6 +103,18 @@ impl<T> fmt::Debug for CompletedCodexOAuthAuthorization<T> {
}
}
impl<T> CompletedCodexOAuthAuthorization<T> {
pub(crate) fn into_parts(
self,
) -> (
PendingAuthorizationMutation,
T,
Box<dyn AuthorizationCommitGuard>,
) {
(self.mutation, self.credential, self.authorization_guard)
}
}
impl fmt::Debug for CompleteCodexOAuthAuthorization {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
@@ -136,7 +135,7 @@ pub struct CodexPendingAuthorization {
state: SecretString,
nonce: SecretString,
code_verifier: SecretString,
reauthorization: Option<CodexOAuthReauthorizationTarget>,
reauthorization: Option<ProviderAccountId>,
mutation: PendingAuthorizationMutation,
}
@@ -150,7 +149,6 @@ pub struct StoredCodexPendingAuthorization {
pub nonce: SecretString,
pub code_verifier: SecretString,
pub reauthorization_account_id: Option<String>,
pub reauthorization_credential_revision: Option<u64>,
pub mutation: PendingAuthorizationMutation,
}
@@ -158,19 +156,11 @@ impl CodexPendingAuthorization {
pub fn from_stored(
input: StoredCodexPendingAuthorization,
) -> Result<Self, CodexOAuthPendingStoreError> {
let reauthorization = match (
input.reauthorization_account_id,
input.reauthorization_credential_revision,
) {
(None, None) => None,
(Some(account_id), Some(revision)) => Some(CodexOAuthReauthorizationTarget {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
credential_revision: CredentialRevision::new(revision)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
}),
_ => return Err(CodexOAuthPendingStoreError::InvalidValue),
};
let reauthorization = input
.reauthorization_account_id
.map(ProviderAccountId::new)
.transpose()
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?;
let pending = Self {
flow_id: input.flow_id,
owner_ref: input.owner_ref,
@@ -239,7 +229,7 @@ impl CodexPendingAuthorization {
}
#[must_use]
pub const fn reauthorization(&self) -> Option<&CodexOAuthReauthorizationTarget> {
pub const fn reauthorization(&self) -> Option<&ProviderAccountId> {
self.reauthorization.as_ref()
}
@@ -284,11 +274,33 @@ pub trait CodexOAuthPendingStore: Send + Sync {
pending: &CodexPendingAuthorization,
) -> Result<(), CodexOAuthPendingStoreError>;
async fn take(
async fn claim(
&self,
owner_ref: &str,
flow_id: &str,
) -> Result<Option<CodexPendingAuthorization>, CodexOAuthPendingStoreError>;
claim_ref: &str,
claim_ttl: Duration,
) -> Result<CodexOAuthPendingClaimOutcome, CodexOAuthPendingStoreError>;
async fn release_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError>;
async fn consume_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError>;
}
pub enum CodexOAuthPendingClaimOutcome {
Claimed(Box<CodexPendingAuthorization>),
NotFound,
InProgress,
}
#[derive(Debug, PartialEq, Eq, thiserror::Error)]
@@ -301,8 +313,12 @@ pub enum CodexOAuthAdminError {
Conflict,
#[error("Codex OAuth flow expired")]
FlowExpired,
#[error("Codex OAuth upstream rejected the operation")]
UpstreamRejected,
#[error("Codex OAuth callback was rejected")]
CallbackRejected,
#[error("Codex OAuth token exchange was rejected")]
TokenRejected,
#[error("Codex OAuth identity verification was rejected")]
IdentityRejected,
#[error("Codex OAuth upstream is unavailable")]
UpstreamUnavailable,
#[error("Codex OAuth exchange send state is ambiguous")]
@@ -336,6 +352,46 @@ pub struct CodexOAuthAdminService {
oauth_client_id: String,
}
struct CodexOAuthAuthorizationCommitGuard {
pending: Arc<dyn CodexOAuthPendingStore>,
owner_ref: String,
flow_id: String,
claim_ref: String,
}
#[async_trait]
impl AuthorizationCommitGuard for CodexOAuthAuthorizationCommitGuard {
async fn commit(self: Box<Self>) -> Result<(), AdminError> {
let consumed = self
.pending
.consume_claim(&self.owner_ref, &self.flow_id, &self.claim_ref)
.await
.map_err(map_authorization_settlement_error)?;
if consumed {
Ok(())
} else {
Err(AdminError::conflict(
"OpenAI OAuth pending claim is no longer current",
))
}
}
async fn abort(self: Box<Self>) -> Result<(), AdminError> {
let released = self
.pending
.release_claim(&self.owner_ref, &self.flow_id, &self.claim_ref)
.await
.map_err(map_authorization_settlement_error)?;
if released {
Ok(())
} else {
Err(AdminError::unavailable(
"OpenAI OAuth pending claim could not be released",
))
}
}
}
impl CodexOAuthAdminService {
#[must_use]
pub fn new(
@@ -372,15 +428,10 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
let (name, reauthorization) = match command.mutation.target() {
AuthorizationMutationTarget::Create { name } => (name.clone(), None),
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
let expected_revision = CredentialRevision::new(expected_credential_revision.get())
.map_err(|_| CodexOAuthAdminError::InvalidInput)?;
AuthorizationMutationTarget::Reauthorize { account_id } => {
let current = self
.store
.load_credential(account_id, expected_revision)
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
if current.account.provider().as_str() != "openai"
@@ -388,13 +439,7 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
{
return Err(CodexOAuthAdminError::NotFound);
}
(
current.account.name().to_owned(),
Some(CodexOAuthReauthorizationTarget {
account_id: account_id.clone(),
credential_revision: expected_revision,
}),
)
(current.account.name().to_owned(), Some(account_id.clone()))
}
};
self.start_pending(command.mutation, name, reauthorization)
@@ -412,12 +457,43 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
{
return Err(CodexOAuthAdminError::InvalidInput);
}
let (pending, mut secret, id_token) = self.exchange_pending(command).await?;
let claim_ref = random_secret()?;
let pending = self.claim_pending(&command, &claim_ref).await?;
let completed = self
.complete_claimed_authorization(pending, command.callback_url.expose_secret())
.await;
match completed {
Ok((mutation, credential)) => Ok(CompletedCodexOAuthAuthorization {
mutation,
credential,
authorization_guard: Box::new(CodexOAuthAuthorizationCommitGuard {
pending: Arc::clone(&self.pending),
owner_ref: command.owner_ref,
flow_id: command.flow_id,
claim_ref,
}),
}),
Err(error) => {
self.release_claim(&command, &claim_ref).await?;
Err(error)
}
}
}
}
impl CodexOAuthAdminService {
async fn complete_claimed_authorization(
&self,
pending: CodexPendingAuthorization,
callback_url: &str,
) -> Result<(PendingAuthorizationMutation, CompletedCodexOAuthCredential), CodexOAuthAdminError>
{
let (mut secret, id_token) = self.exchange_pending(&pending, callback_url).await?;
let mutation = pending.mutation.clone();
let current = if let Some(target) = pending.reauthorization() {
let current = self
.store
.load_credential(target.account_id(), target.credential_revision())
.load_current_credential(target)
.await
.map_err(map_store_error)?;
Some(current)
@@ -483,10 +559,7 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
.map_err(map_admin_error)?,
)
};
Ok(CompletedCodexOAuthAuthorization {
mutation,
credential,
})
Ok((mutation, credential))
}
}
@@ -495,7 +568,7 @@ impl CodexOAuthAdminService {
&self,
mutation: PendingAuthorizationMutation,
name: String,
reauthorization: Option<CodexOAuthReauthorizationTarget>,
reauthorization: Option<ProviderAccountId>,
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
if !valid_text(&name) {
return Err(CodexOAuthAdminError::InvalidInput);
@@ -514,11 +587,7 @@ impl CodexOAuthAdminService {
state: SecretString::from(random_secret()?),
nonce: SecretString::from(random_secret()?),
code_verifier: SecretString::from(random_secret()?),
reauthorization_account_id: reauthorization
.as_ref()
.map(|target| target.account_id().to_string()),
reauthorization_credential_revision: reauthorization
.map(|target| target.credential_revision().get()),
reauthorization_account_id: reauthorization.as_ref().map(ToString::to_string),
mutation,
})
.map_err(map_pending_error)?;
@@ -533,38 +602,70 @@ impl CodexOAuthAdminService {
})
}
async fn exchange_pending(
async fn claim_pending(
&self,
command: CompleteCodexOAuthAuthorization,
) -> Result<
(
CodexPendingAuthorization,
super::types::CodexOAuthSecret,
SecretString,
),
CodexOAuthAdminError,
> {
command: &CompleteCodexOAuthAuthorization,
claim_ref: &str,
) -> Result<CodexPendingAuthorization, CodexOAuthAdminError> {
if !valid_text(&command.owner_ref)
|| !valid_text(&command.flow_id)
|| !valid_secret(claim_ref)
|| command.callback_url.expose_secret().len() > MAX_CALLBACK_BYTES
{
return Err(CodexOAuthAdminError::InvalidInput);
}
let pending = self
let pending = match self
.pending
.take(&command.owner_ref, &command.flow_id)
.claim(
&command.owner_ref,
&command.flow_id,
claim_ref,
AUTHORIZATION_CLAIM_TTL,
)
.await
.map_err(map_pending_error)?
.ok_or(CodexOAuthAdminError::NotFound)?;
{
CodexOAuthPendingClaimOutcome::Claimed(pending) => *pending,
CodexOAuthPendingClaimOutcome::NotFound => return Err(CodexOAuthAdminError::NotFound),
CodexOAuthPendingClaimOutcome::InProgress => {
return Err(CodexOAuthAdminError::Conflict);
}
};
if pending.expires_at <= Utc::now() {
self.release_claim(command, claim_ref).await?;
return Err(CodexOAuthAdminError::FlowExpired);
}
let (code, callback_state) = callback_parts(command.callback_url.expose_secret())?;
Ok(pending)
}
async fn release_claim(
&self,
command: &CompleteCodexOAuthAuthorization,
claim_ref: &str,
) -> Result<(), CodexOAuthAdminError> {
if self
.pending
.release_claim(&command.owner_ref, &command.flow_id, claim_ref)
.await
.map_err(map_pending_error)?
{
Ok(())
} else {
Err(CodexOAuthAdminError::StorageUnavailable)
}
}
async fn exchange_pending(
&self,
pending: &CodexPendingAuthorization,
callback_url: &str,
) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> {
let (code, callback_state) = callback_parts(callback_url)?;
if !constant_time_equal(
pending.state.expose_secret().as_bytes(),
callback_state.expose_secret().as_bytes(),
) {
return Err(CodexOAuthAdminError::UpstreamRejected);
return Err(CodexOAuthAdminError::CallbackRejected);
}
let tokens = self
.exchanger
@@ -574,7 +675,7 @@ impl CodexOAuthAdminService {
})
.await
.map_err(map_exchange_error)?;
Ok((pending, tokens.secret, tokens.id_token))
Ok((tokens.secret, tokens.id_token))
}
}
@@ -622,14 +723,14 @@ fn authorization_url(
}
fn callback_parts(value: &str) -> Result<(SecretString, SecretString), CodexOAuthAdminError> {
let url = Url::parse(value).map_err(|_| CodexOAuthAdminError::UpstreamRejected)?;
let url = Url::parse(value).map_err(|_| CodexOAuthAdminError::CallbackRejected)?;
if url.scheme() != "http"
|| url.host_str() != Some("localhost")
|| url.port() != Some(1455)
|| url.path() != "/auth/callback"
|| url.fragment().is_some()
{
return Err(CodexOAuthAdminError::UpstreamRejected);
return Err(CodexOAuthAdminError::CallbackRejected);
}
let mut code = None;
let mut state = None;
@@ -646,7 +747,7 @@ fn callback_parts(value: &str) -> Result<(SecretString, SecretString), CodexOAut
let state = state.filter(|value| valid_secret(value));
match (code, state) {
(Some(code), Some(state)) => Ok((SecretString::from(code), SecretString::from(state))),
_ => Err(CodexOAuthAdminError::UpstreamRejected),
_ => Err(CodexOAuthAdminError::CallbackRejected),
}
}
@@ -655,7 +756,7 @@ fn set_unique_callback_parameter(
value: String,
) -> Result<(), CodexOAuthAdminError> {
if target.replace(value).is_some() {
return Err(CodexOAuthAdminError::UpstreamRejected);
return Err(CodexOAuthAdminError::CallbackRejected);
}
Ok(())
}
@@ -700,15 +801,8 @@ fn pending_mutation_matches(pending: &CodexPendingAuthorization) -> bool {
}
match (pending.mutation.target(), pending.reauthorization.as_ref()) {
(AuthorizationMutationTarget::Create { name }, None) => name == &pending.name,
(
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
},
Some(target),
) => {
account_id == target.account_id()
&& expected_credential_revision.get() == target.credential_revision().get()
(AuthorizationMutationTarget::Reauthorize { account_id }, Some(target)) => {
account_id == target
}
_ => false,
}
@@ -722,9 +816,23 @@ fn map_pending_error(error: CodexOAuthPendingStoreError) -> CodexOAuthAdminError
}
}
fn map_authorization_settlement_error(error: CodexOAuthPendingStoreError) -> AdminError {
match error {
CodexOAuthPendingStoreError::InvalidValue => {
AdminError::internal("OpenAI OAuth pending claim is invalid")
}
CodexOAuthPendingStoreError::Conflict => {
AdminError::conflict("OpenAI OAuth pending claim conflicts with current state")
}
CodexOAuthPendingStoreError::Unavailable => {
AdminError::unavailable("OpenAI OAuth pending claim store is unavailable")
}
}
}
fn map_exchange_error(error: AuthorizationCodeExchangeError) -> CodexOAuthAdminError {
match error {
AuthorizationCodeExchangeError::Rejected => CodexOAuthAdminError::UpstreamRejected,
AuthorizationCodeExchangeError::Rejected => CodexOAuthAdminError::TokenRejected,
AuthorizationCodeExchangeError::Unavailable => CodexOAuthAdminError::UpstreamUnavailable,
AuthorizationCodeExchangeError::Ambiguous => CodexOAuthAdminError::Ambiguous,
}
@@ -732,7 +840,7 @@ fn map_exchange_error(error: AuthorizationCodeExchangeError) -> CodexOAuthAdminE
fn map_identity_error(error: CodexIdentityVerificationError) -> CodexOAuthAdminError {
match error {
CodexIdentityVerificationError::Rejected => CodexOAuthAdminError::UpstreamRejected,
CodexIdentityVerificationError::Rejected => CodexOAuthAdminError::IdentityRejected,
CodexIdentityVerificationError::Unavailable => CodexOAuthAdminError::UpstreamUnavailable,
}
}
@@ -30,6 +30,9 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60);
const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
// ready 账号主要依赖真实请求的被动观察;这里只提供低频、定额的状态兜底。
const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60);
const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8;
/// OpenAI Provider 主动额度刷新的调度策略。
///
@@ -175,6 +178,14 @@ impl CodexQuotaFact {
}
}
fn quota_fact_is_active_exhaustion(fact: CodexQuotaFact, now: SystemTime) -> bool {
fact.exhausted()
&& fact
.resets_at()
.map(SystemTime::from)
.is_none_or(|reset_at| reset_at > now)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CodexQuotaSyncSummary {
pub updated: u64,
@@ -235,6 +246,8 @@ struct CodexQuotaProjectionState {
next_version: u64,
entries: BTreeMap<ProviderAccountId, CodexQuotaSchedulingEntry>,
last_periodic_refresh_at: BTreeMap<ProviderAccountId, Instant>,
last_exhausted_reset_recheck: BTreeMap<ProviderAccountId, CodexQuotaResetMarker>,
ready_refresh_cursor: Option<ProviderAccountId>,
}
#[derive(Debug, Clone, Copy)]
@@ -245,6 +258,24 @@ struct CodexQuotaSchedulingEntry {
signals: Option<AccountQuotaSignals>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CodexQuotaResetMarker {
revision: CredentialRevision,
reset_at: SystemTime,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodexQuotaRefreshKind {
Exhausted { reset_at: Option<SystemTime> },
Ready,
}
#[derive(Clone)]
struct CodexQuotaRefreshCandidate {
account: ProviderAccount,
kind: CodexQuotaRefreshKind,
}
#[derive(Clone)]
struct CodexQuotaHydrationTarget {
account: ProviderAccount,
@@ -395,17 +426,18 @@ impl CodexQuotaSchedulingProjection {
fn reserve_periodic_refreshes(
&self,
accounts: Vec<ProviderAccount>,
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
now: SystemTime,
) -> Vec<ProviderAccount> {
let candidates = accounts
.into_iter()
.filter(|account| eligible_quota_sync_account(account, now))
.filter_map(|account| quota_refresh_candidate(account, observations, now))
.collect::<Vec<_>>();
let candidate_ids = candidates
.iter()
.map(|account| account.id().clone())
.map(|candidate| candidate.account.id().clone())
.collect::<BTreeSet<_>>();
let now = Instant::now();
let refreshed_at = Instant::now();
let mut state = self
.state
.write()
@@ -413,26 +445,185 @@ impl CodexQuotaSchedulingProjection {
state
.last_periodic_refresh_at
.retain(|account_id, _| candidate_ids.contains(account_id));
candidates
.into_iter()
.filter(|account| {
let due = state
.last_periodic_refresh_at
.get(account.id())
.is_none_or(|last| {
now.saturating_duration_since(*last) >= QUOTA_PERIODIC_REFRESH_MIN_INTERVAL
});
if due {
state
.last_periodic_refresh_at
.insert(account.id().clone(), now);
state
.last_exhausted_reset_recheck
.retain(|account_id, _| candidate_ids.contains(account_id));
let mut reserved = Vec::new();
let mut ready = Vec::new();
for candidate in candidates {
match candidate.kind {
CodexQuotaRefreshKind::Exhausted { reset_at } => {
if reserve_exhausted_quota_refresh(
&mut state,
&candidate.account,
reset_at,
now,
refreshed_at,
) {
reserved.push(candidate.account);
}
}
due
})
.collect()
CodexQuotaRefreshKind::Ready => ready.push(candidate.account),
}
}
ready.sort_by(|left, right| left.id().cmp(right.id()));
let ready = ready
.into_iter()
.filter(|account| periodic_quota_refresh_due(&state, account.id(), refreshed_at))
.collect::<Vec<_>>();
let ready_len = ready.len();
if ready_len == 0 {
return reserved;
}
let start = state
.ready_refresh_cursor
.as_ref()
.and_then(|cursor| ready.iter().position(|account| account.id() > cursor))
.unwrap_or(0);
let ready = ready
.into_iter()
.cycle()
.skip(start)
.take(ready_len.min(QUOTA_READY_REFRESH_BATCH_SIZE))
.collect::<Vec<_>>();
if let Some(last) = ready.last() {
state.ready_refresh_cursor = Some(last.id().clone());
}
for account in &ready {
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
}
reserved.extend(ready);
reserved
}
}
fn quota_refresh_candidate(
account: ProviderAccount,
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
now: SystemTime,
) -> Option<CodexQuotaRefreshCandidate> {
if !eligible_periodic_quota_refresh(&account, now) {
return None;
}
let observation = observations.and_then(|observations| observations.get(account.id()));
let snapshot = quota_snapshot_for_account(&account, observation);
let kind = match account.availability() {
AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
},
AccountAvailability::Cooldown
if snapshot
.as_ref()
.is_some_and(|snapshot| snapshot.fact().exhausted()) =>
{
CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
}
}
AccountAvailability::Ready
if observations.is_some() && quota_observation_is_stale(&account, observation, now) =>
{
CodexQuotaRefreshKind::Ready
}
_ => return None,
};
Some(CodexQuotaRefreshCandidate { account, kind })
}
fn quota_snapshot_for_account(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
) -> Option<CodexAccountQuotaSnapshot> {
observation
.filter(|observation| observation.expected_revision == account.revision())
.and_then(quota_snapshot_from_observation)
}
fn quota_reset_at_from_snapshot(
snapshot: Option<&CodexAccountQuotaSnapshot>,
) -> Option<SystemTime> {
snapshot
.and_then(|snapshot| snapshot.fact().resets_at())
.map(SystemTime::from)
}
fn quota_observation_confirms_active_exhaustion(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
now: SystemTime,
) -> bool {
quota_snapshot_for_account(account, observation)
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))
}
fn quota_observation_is_stale(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
now: SystemTime,
) -> bool {
let Some(observation) = observation.filter(|observation| {
observation.expected_revision == account.revision()
&& quota_snapshot_from_observation(observation).is_some()
}) else {
return true;
};
observation
.observed_at
.and_then(|observed_at| now.duration_since(observed_at).ok())
.is_none_or(|age| age >= QUOTA_READY_REFRESH_MAX_AGE)
}
fn reserve_exhausted_quota_refresh(
state: &mut CodexQuotaProjectionState,
account: &ProviderAccount,
reset_at: Option<SystemTime>,
now: SystemTime,
refreshed_at: Instant,
) -> bool {
let reset_recheck_due = reset_at.is_some_and(|reset_at| {
reset_at <= now
&& state
.last_exhausted_reset_recheck
.get(account.id())
.is_none_or(|marker| {
marker.revision != account.revision() || marker.reset_at != reset_at
})
});
if !reset_recheck_due && !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
if let Some(reset_at) = reset_at.filter(|reset_at| *reset_at <= now) {
state.last_exhausted_reset_recheck.insert(
account.id().clone(),
CodexQuotaResetMarker {
revision: account.revision(),
reset_at,
},
);
}
true
}
fn periodic_quota_refresh_due(
state: &CodexQuotaProjectionState,
account_id: &ProviderAccountId,
now: Instant,
) -> bool {
state
.last_periodic_refresh_at
.get(account_id)
.is_none_or(|last| {
now.saturating_duration_since(*last) >= QUOTA_PERIODIC_REFRESH_MIN_INTERVAL
})
}
fn insert_projection_entry(
state: &mut CodexQuotaProjectionState,
account_id: ProviderAccountId,
@@ -516,10 +707,31 @@ impl CodexCredentialQuotaService {
}
pub async fn synchronize(&self) -> Result<CodexQuotaSyncSummary, CodexCredentialQuotaError> {
let accounts = self.repository.list_for_provider().await?;
let mut accounts = self.repository.list_for_provider().await?;
let mut summary = CodexQuotaSyncSummary::default();
let now = SystemTime::now();
let accounts = self.scheduling.reserve_periodic_refreshes(accounts, now);
let account_ids = accounts
.iter()
.map(|account| account.id().clone())
.collect::<Vec<_>>();
let observations = match self.store.get_quotas(&account_ids).await {
Ok(observations) => Some(
observations
.into_iter()
.map(|observation| (observation.account_id.clone(), observation))
.collect::<BTreeMap<_, _>>(),
),
Err(error) => {
// 持久观察不可用时不盲扫 ready;保留已耗尽账号的既有恢复轮询。
tracing::warn!(error = %error, "OpenAI quota refresh skipped ready-account scan");
None
}
};
self.reconcile_persisted_exhaustion(&mut accounts, observations.as_ref(), now)
.await;
let accounts =
self.scheduling
.reserve_periodic_refreshes(accounts, observations.as_ref(), now);
if accounts.is_empty() {
return Ok(summary);
}
@@ -553,6 +765,9 @@ impl CodexCredentialQuotaService {
}
Err(CodexQuotaFetchError::Upstream { account, error }) => {
let (availability, reason, cooldown) = classify_error(&error, observed_at);
let preserve_exhaustion = account.availability()
== AccountAvailability::QuotaExhausted
&& availability == Some(AccountAvailability::Cooldown);
match availability {
Some(AccountAvailability::Invalid | AccountAvailability::Banned) => {
summary.invalid += 1;
@@ -564,6 +779,9 @@ impl CodexCredentialQuotaService {
continue;
}
}
if preserve_exhaustion {
continue;
}
let _ = self
.repository
.apply_state(
@@ -580,6 +798,63 @@ impl CodexCredentialQuotaService {
Ok(summary)
}
async fn reconcile_persisted_exhaustion(
&self,
accounts: &mut [ProviderAccount],
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
observed_at: SystemTime,
) {
let Some(observations) = observations else {
return;
};
for account in accounts {
if account.availability() != AccountAvailability::Cooldown
|| !quota_observation_confirms_active_exhaustion(
account,
observations.get(account.id()),
observed_at,
)
{
continue;
}
let Ok(Some(current)) = self.store.get_account(account.id()).await else {
continue;
};
if current.revision() != account.revision()
|| current.availability() != AccountAvailability::Cooldown
{
continue;
}
if !matches!(
self.has_active_exhaustion(&current, observed_at).await,
Ok(true)
) {
continue;
}
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
observed_at,
)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
"OpenAI persisted quota exhaustion reconciliation failed"
);
continue;
}
let enabled = current.enabled();
*account =
current.with_runtime_state(enabled, AccountAvailability::QuotaExhausted, None);
}
}
/// 解析并 revision-fenced 落库单账号的 Provider quota JSON。
async fn apply_fetched_quota(
&self,
@@ -717,6 +992,30 @@ impl CodexCredentialQuotaService {
.await?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialQuotaError::NotFound)?;
self.read_snapshot_for(&account).await
}
pub(super) async fn has_active_exhaustion(
&self,
account: &ProviderAccount,
now: SystemTime,
) -> Result<bool, CodexCredentialQuotaError> {
match self.read_snapshot_for(account).await {
Ok(snapshot) => Ok(snapshot
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))),
Err(CodexCredentialQuotaError::RevisionConflict) => Ok(false),
Err(error) => Err(error),
}
}
async fn read_snapshot_for(
&self,
account: &ProviderAccount,
) -> Result<Option<CodexAccountQuotaSnapshot>, CodexCredentialQuotaError> {
if account.provider().as_str() != "openai" {
return Err(CodexCredentialQuotaError::NotFound);
}
let account_id = account.id();
let Some(observation) = self
.store
.get_quotas(std::slice::from_ref(account_id))
@@ -918,12 +1217,11 @@ async fn fetch_usage_once(
.map_err(CodexQuotaFetchAttemptError::Upstream)
}
fn eligible_quota_sync_account(account: &ProviderAccount, now: SystemTime) -> bool {
fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool {
account.enabled()
&& account
.access_token_expires_at()
.is_none_or(|expires_at| expires_at > now)
&& account.availability() == AccountAvailability::QuotaExhausted
}
fn merge_passive_quota(
@@ -28,6 +28,16 @@ const MAX_REFRESH_BATCH: u32 = 1_000;
/// 连续失败计数窗口;每次瞬态失败刷新该 TTL,静默满窗后计数过期归零。
const REFRESH_BACKOFF_WINDOW: Duration = Duration::from_secs(30 * 60);
fn refresh_due_at(account: &ProviderAccount, now: SystemTime) -> Option<SystemTime> {
if let Some(expires_at) = account.access_token_expires_at()
&& expires_at <= now
{
// AT 已过期时不能继续等待持久退避;后续额度/请求探测都必须先用 RT 续出可用 AT。
return Some(expires_at);
}
account.next_refresh_at().filter(|next| *next <= now)
}
pub struct DueCodexCredential {
pub account: ProviderAccount,
pub secret: CodexOAuthSecret,
@@ -248,9 +258,9 @@ impl CodexCredentialRefreshService {
)
&& (account.availability() != AccountAvailability::Cooldown
|| account.cooldown_until().is_some_and(|until| until <= now))
&& account.next_refresh_at().is_some_and(|next| next <= now)
&& refresh_due_at(account, now).is_some()
});
accounts.sort_by_key(|account| (account.next_refresh_at(), account.id().clone()));
accounts.sort_by_key(|account| (refresh_due_at(account, now), account.id().clone()));
accounts.truncate(MAX_REFRESH_BATCH as usize);
let mut due = Vec::with_capacity(accounts.len());
let mut failures = Vec::new();
@@ -6,8 +6,9 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
use gateway_core::engine::credential::{
AccountAvailability, AccountCandidate, AccountFeedbackStats, AccountRuntimeSignals,
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
AccountRuntimeSignals, AccountSelectionContext, AccountSelector, ProviderAccount,
ProviderAccountId,
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt};
use gateway_core::provider_ports::{
@@ -162,22 +163,26 @@ impl CodexCredentialSelector {
request: &SelectCodexCredential<'_>,
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let diagnostic = request.attempt.is_diagnostic_required_account();
let accounts = self.repository.list_for_provider().await?;
let accounts = accounts
.into_iter()
.filter(|account| {
account.provider() == &self.provider_kind
&& !matches!(
self.catalog.observed_model_support(
account.id(),
account.revision(),
request.upstream_model,
),
Ok(Some(false))
)
&& (diagnostic
|| !matches!(
self.catalog.observed_model_support(
account.id(),
account.revision(),
request.upstream_model,
),
Ok(Some(false))
))
})
.collect::<Vec<_>>();
self.quota.prepare_scheduling(&accounts).await;
if !diagnostic {
self.quota.prepare_scheduling(&accounts).await;
}
let account_ids = accounts
.iter()
.map(|account| account.id().clone())
@@ -223,23 +228,27 @@ impl CodexCredentialSelector {
.collect::<Vec<_>>()
});
let scheduling_candidates = quota_exhausted_candidates.as_deref().unwrap_or(&candidates);
let affinity_account = match request.session_affinity_key {
Some(key) => match self.lookup_session_affinity(key).await {
SessionAffinityLookup::Bound(account_id)
if scheduling_candidates.iter().any(|candidate| {
candidate.account.id() == &account_id
&& candidate.account.is_schedulable(SystemTime::now())
}) =>
{
Some(account_id)
}
SessionAffinityLookup::Bound(_) => {
self.clear_session_affinity(key).await;
None
}
SessionAffinityLookup::Missing | SessionAffinityLookup::Unavailable => None,
},
None => None,
let affinity_account = if diagnostic {
None
} else {
match request.session_affinity_key {
Some(key) => match self.lookup_session_affinity(key).await {
SessionAffinityLookup::Bound(account_id)
if scheduling_candidates.iter().any(|candidate| {
candidate.account.id() == &account_id
&& candidate.account.is_schedulable(SystemTime::now())
}) =>
{
Some(account_id)
}
SessionAffinityLookup::Bound(_) => {
self.clear_session_affinity(key).await;
None
}
SessionAffinityLookup::Missing | SessionAffinityLookup::Unavailable => None,
},
None => None,
}
};
let cyber_policy_scope = self
.prepare_cyber_policy_scope(cyber_policy_session_key)
@@ -290,6 +299,11 @@ impl CodexCredentialSelector {
excluded_accounts: excluded.clone(),
preferred_account: preferred.clone(),
round_robin_cursor,
availability: if diagnostic {
AccountAvailabilityPolicy::BypassForDiagnostic
} else {
AccountAvailabilityPolicy::Enforce
},
};
let Some(selected) = AccountSelector.select(scheduling_candidates, &context) else {
return match shortest_retry {
@@ -304,6 +318,7 @@ impl CodexCredentialSelector {
.find(|candidate| candidate.account.id() == selected.account.id())
.map(|candidate| candidate.account.clone())
.ok_or(CredentialSelectionError::InvalidCredential)?;
let allows_account_state_mutation = !diagnostic || account.enabled();
let policy = request.attempt.account_selection_policy();
match self
.leases
@@ -327,11 +342,12 @@ impl CodexCredentialSelector {
let (account, runtime) = if account.authentication_kind()
== CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY
{
let prepared = self
.agent_identity
.prepare(&account)
.await
.map_err(|_| CredentialSelectionError::InvalidCredential)?;
let prepared = if allows_account_state_mutation {
self.agent_identity.prepare(&account).await
} else {
self.agent_identity.load_current(account.id()).await
}
.map_err(|_| CredentialSelectionError::InvalidCredential)?;
(prepared.account, prepared.credential)
} else {
let runtime = self.repository.load_runtime_credential(&account).await?;
@@ -359,6 +375,7 @@ impl CodexCredentialSelector {
authentication: runtime.authentication,
cookies,
cyber_policy_scope,
allows_account_state_mutation,
_guard: guard,
});
}
@@ -535,13 +552,9 @@ impl CodexCredentialSelector {
CookieRecovery::None,
),
CodexAccountFailure::RateLimited { retry_after } => {
let until = now.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
(
AccountAvailability::QuotaExhausted,
Some("rate_limited".to_owned()),
until,
CookieRecovery::None,
)
return self
.record_rate_limit_failure(account, retry_after, now)
.await;
}
CodexAccountFailure::CloudflareChallenge { retry_after } => {
let delay = self.cloudflare_challenge_delay(account.id(), now, retry_after);
@@ -579,11 +592,55 @@ impl CodexCredentialSelector {
Ok(())
}
async fn record_rate_limit_failure(
&self,
account: &ProviderAccount,
retry_after: Option<Duration>,
observed_at: SystemTime,
) -> Result<(), CredentialSelectionError> {
let current = self.current_account(account.id()).await?;
if current.revision() != account.revision() {
return Err(CredentialRepositoryError::RevisionConflict.into());
}
if matches!(
current.availability(),
AccountAvailability::Invalid
| AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::QuotaExhausted
) {
return Ok(());
}
let quota_exhausted = self
.quota
.has_active_exhaustion(&current, observed_at)
.await
.map_err(|_| CredentialSelectionError::Store)?;
let (availability, reason, cooldown_until) = if quota_exhausted {
(
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
)
} else {
(
AccountAvailability::Cooldown,
Some("rate_limited".to_owned()),
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)),
)
};
self.repository
.apply_state(&current, availability, reason, cooldown_until, observed_at)
.await?;
Ok(())
}
pub async fn record_success(
&self,
account: &ProviderAccount,
session_affinity_key: Option<&ProviderSessionAffinityKey>,
) {
self.restore_recoverable_account_state(account).await;
self.risk_recovery
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
@@ -626,6 +683,42 @@ impl CodexCredentialSelector {
}
}
async fn restore_recoverable_account_state(&self, account: &ProviderAccount) {
let Ok(current) = self.current_account(account.id()).await else {
return;
};
if current.provider() != &self.provider_kind
|| current.revision() != account.revision()
|| !current.enabled()
|| !matches!(
current.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
)
{
return;
}
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::Ready,
None,
None,
SystemTime::now(),
)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
"OpenAI account state recovery after successful upstream response failed"
);
}
}
pub async fn current_account(
&self,
account_id: &ProviderAccountId,
@@ -801,6 +894,7 @@ pub struct CodexCredentialLease {
cookies: Vec<RuntimeCodexCookie>,
installation_id: String,
cyber_policy_scope: Option<CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
_guard: Box<dyn ProviderLeaseGuard>,
}
@@ -834,6 +928,12 @@ impl CodexCredentialLease {
pub(crate) const fn cyber_policy_scope(&self) -> Option<&CodexCyberPolicyScope> {
self.cyber_policy_scope.as_ref()
}
/// 禁用账号的管理端诊断必须只返回真实上游结果,不能回写账号侧状态。
#[must_use]
pub(crate) const fn allows_account_state_mutation(&self) -> bool {
self.allows_account_state_mutation
}
}
impl fmt::Debug for CodexCredentialLease {
@@ -160,7 +160,6 @@ pub async fn initialize(
let refresher: Arc<dyn TokenRefresher> = token_client.clone();
let exchanger: Arc<dyn AuthorizationCodeExchanger> = token_client;
let credential_admin = Arc::new(CodexCredentialAdminService::new(
repository.clone(),
Arc::clone(&refresher),
Arc::clone(&identity),
Arc::clone(&leases),
+289 -123
View File
@@ -79,6 +79,8 @@ const PROVIDER_NAME: &str = "openai";
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const WEBSOCKET_TRANSPORT: &str = "websocket";
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
/// 与 2.x 预取器一致:提交边界前最多保留 64 KiB 原始上游 chunk。
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api";
pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api";
@@ -376,6 +378,7 @@ impl Provider for CodexProvider {
response_store,
continuation_scope: None,
});
let allows_account_state_mutation = lease.allows_account_state_mutation();
let events = cold_response_stream(ColdResponse {
client: self.client.clone(),
response_origin: self.responses_url.clone(),
@@ -391,8 +394,12 @@ impl Provider for CodexProvider {
session_affinity_key,
session_capture,
});
Ok(ProviderStream::new(metadata, events, lease)
.with_account_feedback(Arc::clone(&self.account_feedback)))
let stream = ProviderStream::new(metadata, events, lease);
Ok(if allows_account_state_mutation {
stream.with_account_feedback(Arc::clone(&self.account_feedback))
} else {
stream
})
}
}
@@ -657,6 +664,15 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
} = response;
Box::pin(async_stream::try_stream! {
let cyber_policy_scope = lease.cyber_policy_scope().cloned();
let allows_account_state_mutation = lease.allows_account_state_mutation();
let failure_context = OpenAiFailureContext {
client: &client,
selector: &selector,
quota: &quota,
response_origin: &response_origin,
cyber_policy_scope: cyber_policy_scope.as_ref(),
allows_account_state_mutation,
};
let mut active_account = lease.account().clone();
let cookie_header = build_cookie_header(lease.cookies())?;
let mut authorization = lease
@@ -670,6 +686,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
})?;
let request_id = context.request_id().as_str().to_owned();
let cancellation = context.cancellation().clone();
let stream_commit_policy =
StreamCommitPolicy::for_continuation(context.continuation_attempt());
let mut response = create_response_attempt(
&client,
&request,
@@ -686,7 +704,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
&cancellation,
)
.await;
if let Err(CodexHandshakeAttemptError::Client(error)) = &response {
if allows_account_state_mutation
&& let Err(CodexHandshakeAttemptError::Client(error)) = &response
{
match agent_identity
.recover_after_rejected_task(
active_account.id(),
@@ -735,15 +755,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if let Some(observation) = failure.observation.take() {
yield ProviderEvent::observation(observation);
}
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
return;
@@ -754,15 +766,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
ProviderErrorKind::Protocol,
UpstreamSendState::Sent,
));
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
return;
@@ -786,7 +790,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
context.attempt_index().get(),
);
yield ProviderEvent::observation(observation_state.observation()?);
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
}
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(etag)
{
@@ -795,7 +801,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
"OpenAI model ETag observation was rejected"
);
}
if !response.set_cookie_headers.is_empty()
if allows_account_state_mutation
&& !response.set_cookie_headers.is_empty()
&& let Ok(outcome) = selector
.capture_response_cookies(
&active_account,
@@ -821,6 +828,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
let mut decoder =
CodexCanonicalDecoder::new(upstream_model.as_str()).with_raw_sse_passthrough();
let mut pre_commit_events = PreCommitClientEvents::new(stream_commit_policy);
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
@@ -847,20 +855,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure) => {
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
if observation_state.merge_rate_limit_headers(&updates) {
yield ProviderEvent::observation(observation_state.observation()?);
}
}
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
return;
@@ -871,7 +873,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let rate_limits_changed = if updates.is_empty() {
false
} else {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
};
if let (Some(capture), Some(updates)) =
@@ -883,56 +887,48 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, context.timing_started_at());
let (mut events, terminal_failure) = match decoder.push(&chunk) {
let chunk_len = chunk.len();
let (mut events, canonical_failure) = match decoder.push(&chunk) {
CodexCanonicalOutcome::Events(events) => (events, None),
CodexCanonicalOutcome::Failed(failure) => {
let (events, error, semantic_output_seen) = failure.into_parts();
let failure = map_canonical_error(
(events, Some((error, semantic_output_seen)))
}
};
pre_commit_events.observe_chunk(chunk_len);
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
map_canonical_error(
error,
&failure_diagnostics,
&failure_set_cookie_headers,
&failure_rate_limit_headers,
ReplayBoundary::from_semantic_output(semantic_output_seen),
);
(events, Some(failure))
}
};
ReplayBoundary::from_semantic_output(
semantic_output_seen || pre_commit_events.is_committed(),
),
),
atomic_upstream_failure,
)
});
let timing_signals = decoder.take_timing_signals();
let timing_changed = first_event_changed
|| observation_state
.observe_timing_signals(decoder.take_timing_signals(), context.timing_started_at());
.observe_timing_signals(timing_signals, context.timing_started_at());
let completed = events
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if rate_limits_changed || timing_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
failure,
)
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
if terminal_failure.is_none() {
let failure = MappedProviderFailure::plain(error);
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
} else {
@@ -940,56 +936,78 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
return;
}
for event in events {
yield event;
}
if let Some(failure) = terminal_failure {
Err(failure.error)?;
return;
}
if completed {
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 完成事件一旦交给下游,Core 可以立刻停止轮询 Provider stream;
// 在此之前持久化亲和关系,保证成功请求不会因流被提前 drop 而丢失绑定。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if rate_limits_changed || timing_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some((mut failure, atomic_upstream_failure)) = terminal_failure {
let failure_after_commit =
timing_signals.semantic_output || pre_commit_events.is_committed();
if failure_after_commit {
for event in pre_commit_events.commit(events) {
yield event;
}
} else if atomic_upstream_failure {
failure.error = failure
.error
.with_atomic_client_events(pre_commit_events.take_for_failure(events));
}
Err(failure.error)?;
return;
}
let events = pre_commit_events.stage(events, timing_signals, completed)?;
for event in events {
yield event;
}
if completed {
return;
}
}
let (mut events, terminal_failure) = match decoder.finish() {
let (mut events, canonical_failure) = match decoder.finish() {
CodexCanonicalOutcome::Events(events) => (events, None),
CodexCanonicalOutcome::Failed(failure) => {
let (events, error, semantic_output_seen) = failure.into_parts();
let failure = map_canonical_error(
(events, Some((error, semantic_output_seen)))
}
};
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
map_canonical_error(
error,
&failure_diagnostics,
&failure_set_cookie_headers,
&failure_rate_limit_headers,
ReplayBoundary::from_semantic_output(semantic_output_seen),
);
(events, Some(failure))
}
};
ReplayBoundary::from_semantic_output(
semantic_output_seen || pre_commit_events.is_committed(),
),
),
atomic_upstream_failure,
)
});
let timing_signals = decoder.take_timing_signals();
let timing_changed = observation_state
.observe_timing_signals(decoder.take_timing_signals(), context.timing_started_at());
.observe_timing_signals(timing_signals, context.timing_started_at());
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
let rate_limits_changed = if updates.is_empty() {
false
} else {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
};
if let Some(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
failure,
)
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let (Some(capture), Some(updates)) =
@@ -1001,15 +1019,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
if terminal_failure.is_none() {
let failure = MappedProviderFailure::plain(error);
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
} else {
@@ -1023,17 +1033,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
for event in events {
yield event;
}
if let Some(failure) = terminal_failure {
Err(failure.error)?;
return;
}
if completed {
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 同上:尾部 finish() 也可能产出 completed,亲和记录必须先于任何下游 yield。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
@@ -1041,6 +1042,28 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some((mut failure, atomic_upstream_failure)) = terminal_failure {
let failure_after_commit =
timing_signals.semantic_output || pre_commit_events.is_committed();
if failure_after_commit {
for event in pre_commit_events.commit(events) {
yield event;
}
} else if atomic_upstream_failure {
failure.error = failure
.error
.with_atomic_client_events(pre_commit_events.take_for_failure(events));
}
Err(failure.error)?;
return;
}
let events = pre_commit_events.finish(events, timing_signals, completed)?;
for event in events {
yield event;
}
})
}
@@ -1761,6 +1784,128 @@ impl MappedProviderFailure {
}
}
/// 与 2.x 一致的流提交策略:普通请求优先首事件,续写/重放优先可恢复性。
#[derive(Clone, Copy)]
enum StreamCommitPolicy {
FirstForwardableEvent,
UntilOutputOrTerminal,
}
impl StreamCommitPolicy {
const fn for_continuation(attempt: ContinuationAttempt) -> Self {
match attempt {
ContinuationAttempt::None => Self::FirstForwardableEvent,
ContinuationAttempt::Native
| ContinuationAttempt::ReplayOwner
| ContinuationAttempt::ReplayAny => Self::UntilOutputOrTerminal,
}
}
fn boundary_reached(
self,
events: &[ProviderEvent],
timing_signals: ResponseEventSignals,
completed: bool,
) -> bool {
match self {
Self::FirstForwardableEvent => events.iter().any(ProviderEvent::has_client_event),
Self::UntilOutputOrTerminal => timing_signals.semantic_output || completed,
}
}
}
/// 提交边界前的上游事件预取。
///
/// 原始 chunk 计数而不是重编码后的 event 大小,保持 2.x 的 64 KiB 资源边界;一旦
/// 提交,后续事件不再具备无痕重放资格。
struct PreCommitClientEvents {
policy: StreamCommitPolicy,
pending: Vec<ProviderEvent>,
prefetched_bytes: usize,
committed: bool,
}
impl PreCommitClientEvents {
const fn new(policy: StreamCommitPolicy) -> Self {
Self {
policy,
pending: Vec::new(),
prefetched_bytes: 0,
committed: false,
}
}
fn observe_chunk(&mut self, bytes: usize) {
if !self.committed {
self.prefetched_bytes = self.prefetched_bytes.saturating_add(bytes);
}
}
fn stage(
&mut self,
incoming: Vec<ProviderEvent>,
timing_signals: ResponseEventSignals,
completed: bool,
) -> Result<Vec<ProviderEvent>, ProviderError> {
if self.committed {
return Ok(incoming);
}
let boundary_reached = self
.policy
.boundary_reached(&incoming, timing_signals, completed);
self.pending.extend(incoming);
if boundary_reached {
return Ok(self.commit_pending());
}
if self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
return Err(prefetch_protocol_error());
}
Ok(Vec::new())
}
fn finish(
&mut self,
incoming: Vec<ProviderEvent>,
timing_signals: ResponseEventSignals,
completed: bool,
) -> Result<Vec<ProviderEvent>, ProviderError> {
let events = self.stage(incoming, timing_signals, completed)?;
if self.committed {
Ok(events)
} else {
Err(prefetch_protocol_error())
}
}
fn commit(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
if self.committed {
return incoming;
}
self.pending.extend(incoming);
self.commit_pending()
}
fn take_for_failure(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
self.pending.extend(incoming);
self.prefetched_bytes = 0;
std::mem::take(&mut self.pending)
}
const fn is_committed(&self) -> bool {
self.committed
}
fn commit_pending(&mut self) -> Vec<ProviderEvent> {
self.committed = true;
self.prefetched_bytes = 0;
std::mem::take(&mut self.pending)
}
}
fn prefetch_protocol_error() -> ProviderError {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
}
#[derive(Clone, Copy)]
enum ReplayBoundary {
BeforeSemanticOutput,
@@ -1781,24 +1926,40 @@ impl ReplayBoundary {
}
}
struct OpenAiFailureContext<'a> {
client: &'a CodexBackendClient,
selector: &'a CodexCredentialSelector,
quota: &'a CodexCredentialQuotaService,
response_origin: &'a Url,
cyber_policy_scope: Option<&'a CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
}
async fn apply_failure(
client: &CodexBackendClient,
selector: &CodexCredentialSelector,
quota: &CodexCredentialQuotaService,
context: &OpenAiFailureContext<'_>,
account: &ProviderAccount,
response_origin: &Url,
cyber_policy_scope: Option<&CodexCyberPolicyScope>,
failure: &MappedProviderFailure,
) {
synchronize_passive_quota(quota, account, &failure.rate_limit_headers).await;
if !context.allows_account_state_mutation {
return;
}
synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await;
if failure.cyber_policy_failure {
selector
.record_cyber_policy_failure(cyber_policy_scope, account)
context
.selector
.record_cyber_policy_failure(context.cyber_policy_scope, account)
.await;
}
if let Some(account_failure) = failure.account_failure {
client.evict_websocket_account(account.id().as_str()).await;
if let Err(error) = selector.record_failure(account, account_failure).await {
context
.client
.evict_websocket_account(account.id().as_str())
.await;
if let Err(error) = context
.selector
.record_failure(account, account_failure)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
@@ -1808,8 +1969,13 @@ async fn apply_failure(
}
if failure.capture_response_cookies
&& !failure.set_cookie_headers.is_empty()
&& let Err(error) = selector
.capture_response_cookies(account, response_origin, &failure.set_cookie_headers)
&& let Err(error) = context
.selector
.capture_response_cookies(
account,
context.response_origin,
&failure.set_cookie_headers,
)
.await
{
tracing::warn!(
@@ -273,8 +273,20 @@ impl CodexCanonicalDecoder {
self.merge_timing_signals(signals);
if matches!(event_type, Some("response.failed" | "error")) {
let failure = ResponsesSseFailure::from_event(event_type.unwrap_or_default(), &value);
let mut canonical = Vec::new();
if !self.started {
// 失败首帧可能是唯一携带 response_id 的上游事实。身份投影只做
// 最佳努力:缺少 response/id 时仍必须保留原 typed failure。
let _ = self.start(&value, &mut canonical);
}
if let Some(wire) = Self::wire_for_event(event, value, raw_sse_frame) {
output.push(ProviderEvent::wire(wire));
output.push(if canonical.is_empty() {
ProviderEvent::wire(wire)
} else {
ProviderEvent::canonical_with_wire(canonical, wire)
});
} else {
output.extend(canonical.into_iter().map(ProviderEvent::canonical));
}
return Err(CodexCanonicalError::Upstream(failure));
}
@@ -415,15 +415,6 @@ fn classify_upstream_failure(
if status == Some(StatusCode::NOT_FOUND) && body.trim().is_empty() {
return CodexFailureCategory::CloudflarePathBlocked;
}
if [code.as_str(), error_type.as_str()]
.into_iter()
.any(is_rate_limit_signal)
|| is_rate_limit_message(&message)
|| is_rate_limit_message(&body)
|| status == Some(StatusCode::TOO_MANY_REQUESTS)
{
return CodexFailureCategory::RateLimited;
}
if [code.as_str(), error_type.as_str()]
.into_iter()
.any(is_quota_signal)
@@ -433,6 +424,15 @@ fn classify_upstream_failure(
{
return CodexFailureCategory::QuotaExhausted;
}
if [code.as_str(), error_type.as_str()]
.into_iter()
.any(is_rate_limit_signal)
|| is_rate_limit_message(&message)
|| is_rate_limit_message(&body)
|| status == Some(StatusCode::TOO_MANY_REQUESTS)
{
return CodexFailureCategory::RateLimited;
}
if is_upstream_overload(&code) || is_upstream_overload(&message) {
return CodexFailureCategory::Unavailable;
}
+204 -16
View File
@@ -33,10 +33,12 @@ use gateway_core::operation::OperationKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingTakeOutcome,
NewOAuthPendingFlow, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCredentialState, ProviderCredentialStatePort, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts,
ProviderCooldownScope, ProviderCredentialState, ProviderCredentialStatePort,
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderStoreError,
ProviderStorePorts,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
@@ -339,13 +341,29 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
.expect("start authorization");
{
let values = pending.values.lock().expect("OAuth pending");
let (_, payload, _) = values.values().next().expect("stored pending");
let (_, payload, _, _) = values.values().next().expect("stored pending");
let mutation = payload
.expose_to_provider()
.get("mutation")
.and_then(Value::as_object)
.expect("pending mutation");
assert!(
payload
.expose_to_provider()
.get("reauthorization_credential_revision")
.is_none()
);
assert_eq!(
mutation.get("schema_version").and_then(Value::as_u64),
Some(3)
);
assert!(mutation.get("expected_config_revision").is_none());
assert!(
mutation
.get("target")
.and_then(Value::as_object)
.is_some_and(|target| target.get("expected_credential_revision").is_none())
);
assert_eq!(
mutation.get("started_request_id").and_then(Value::as_str),
Some("request-start")
@@ -369,6 +387,81 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
assert_eq!(pending.values.lock().expect("OAuth pending").len(), 1);
}
#[tokio::test]
async fn openai_reauthorization_pending_payload_contains_only_stable_account_identity() {
let accounts = Arc::new(MemoryAccountStore::default());
accounts
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_pending_reauth".to_owned(),
name: "pending reauthorization".to_owned(),
secret: secret("pending-reauth-access"),
verified_account: profile("chatgpt-pending-reauth"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let pending = Arc::new(TestOAuthPending::default());
let config = valid_config();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(accounts, Arc::clone(&pending)),
)
.await
.expect("OpenAI bundle");
let context = MutationContext {
actor: MutationActor::AdminApiKey,
request_id: "request-pending-reauth".to_owned(),
};
bundle
.admin_provider()
.start_authorization(PendingAuthorizationMutation::new(
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new("acct_pending_reauth").expect("account id"),
},
AuthorizationOwnerBinding::from_context(&context),
))
.await
.expect("start reauthorization");
let values = pending.values.lock().expect("OAuth pending");
let (_, payload, _, _) = values.values().next().expect("stored pending");
let document = payload.expose_to_provider();
let mutation = document
.get("mutation")
.and_then(Value::as_object)
.expect("pending mutation");
let target = mutation
.get("target")
.and_then(Value::as_object)
.expect("pending target");
assert_eq!(
mutation.get("schema_version").and_then(Value::as_u64),
Some(3)
);
assert_eq!(
document
.get("reauthorization_account_id")
.and_then(Value::as_str),
Some("acct_pending_reauth")
);
assert!(
document
.get("reauthorization_credential_revision")
.is_none()
);
assert_eq!(
target.get("kind").and_then(Value::as_str),
Some("reauthorize")
);
assert_eq!(
target.get("account_id").and_then(Value::as_str),
Some("acct_pending_reauth")
);
assert!(target.get("expected_credential_revision").is_none());
}
#[tokio::test]
async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export() {
let store = Arc::new(MemoryAccountStore::default());
@@ -567,10 +660,20 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
.await
.expect_err("invalid import");
assert_eq!(import_error.kind(), ProviderAdminErrorKind::Invalid);
let mut stale_record = record.clone();
stale_record.name = "stale name".to_owned();
stale_record.email = None;
stale_record.plan_type = None;
stale_record.credential_revision = Revision::new(99).expect("stale revision");
stale_record.has_refresh_token = false;
stale_record.access_token_expires_at = None;
stale_record.next_refresh_at = None;
stale_record.enabled = false;
stale_record.availability = AdminAccountAvailability::Banned;
stale_record.cooldown_until = Some(Utc::now() + chrono::Duration::minutes(5));
let rotation_error = admin
.prepare_rotation(PrepareCredentialRotation {
account: record.clone(),
expected_credential_revision: record.credential_revision,
account: stale_record,
provider_material: ProviderDocument::new(OpaqueProviderData::new(Map::new())),
})
.await
@@ -799,6 +902,30 @@ impl ProviderCooldownPort for TestCooldown {
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
fn put_scoped_if_later(
&self,
_cooldown: ProviderScopedCooldown,
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
fn read_scoped<'a>(
&'a self,
_account_id: &'a ProviderAccountId,
_scope: &'a ProviderCooldownScope,
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>> {
Box::pin(async { Ok(None) })
}
fn clear_scoped<'a>(
&'a self,
_account_id: &'a ProviderAccountId,
_scope: &'a ProviderCooldownScope,
_through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
}
struct TestRuntimePolicy;
@@ -822,7 +949,7 @@ struct TestOAuthPending {
}
type PendingKey = (String, String);
type PendingValue = (String, OpaqueProviderData, SystemTime);
type PendingValue = (String, OpaqueProviderData, SystemTime, Option<String>);
impl OAuthPendingFlowPort for TestOAuthPending {
fn put_if_absent(
@@ -844,36 +971,97 @@ impl OAuthPendingFlowPort for TestOAuthPending {
flow.owner().expose_to_store().to_owned(),
flow.payload().clone(),
SystemTime::now() + flow.ttl(),
None,
),
);
Ok(OAuthPendingPutOutcome::Stored)
})
}
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a gateway_core::provider_ports::OAuthPendingBinding,
owner: &'a gateway_core::provider_ports::OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>> {
claim: &'a gateway_core::provider_ports::OAuthPendingBinding,
_claim_ttl: Duration,
) -> BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, expires_at)) = values.get(&key) else {
return Ok(OAuthPendingTakeOutcome::NotFound);
let Some((stored_owner, payload, expires_at, stored_claim)) = values.get_mut(&key)
else {
return Ok(OAuthPendingClaimOutcome::NotFound);
};
if expires_at <= &SystemTime::now() {
if *expires_at <= SystemTime::now() {
values.remove(&key);
return Ok(OAuthPendingTakeOutcome::NotFound);
return Ok(OAuthPendingClaimOutcome::NotFound);
}
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingTakeOutcome::OwnerMismatch);
return Ok(OAuthPendingClaimOutcome::OwnerMismatch);
}
let (_, payload, _) = values.remove(&key).expect("checked pending");
Ok(OAuthPendingTakeOutcome::Taken(payload))
if stored_claim.is_some() {
return Ok(OAuthPendingClaimOutcome::InProgress);
}
*stored_claim = Some(claim.expose_to_store().to_owned());
Ok(OAuthPendingClaimOutcome::Claimed(payload.clone()))
})
}
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a gateway_core::provider_ports::OAuthPendingBinding,
owner: &'a gateway_core::provider_ports::OAuthPendingBinding,
claim: &'a gateway_core::provider_ports::OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, _, stored_claim)) = values.get_mut(&key) else {
return Ok(OAuthPendingReleaseOutcome::NotFound);
};
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingReleaseOutcome::OwnerMismatch);
}
if stored_claim.as_deref() != Some(claim.expose_to_store()) {
return Ok(OAuthPendingReleaseOutcome::ClaimMismatch);
}
*stored_claim = None;
Ok(OAuthPendingReleaseOutcome::Released)
})
}
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a gateway_core::provider_ports::OAuthPendingBinding,
owner: &'a gateway_core::provider_ports::OAuthPendingBinding,
claim: &'a gateway_core::provider_ports::OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, _, stored_claim)) = values.get(&key) else {
return Ok(OAuthPendingConsumeOutcome::NotFound);
};
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingConsumeOutcome::OwnerMismatch);
}
if stored_claim.as_deref() != Some(claim.expose_to_store()) {
return Ok(OAuthPendingConsumeOutcome::ClaimMismatch);
}
values.remove(&key);
Ok(OAuthPendingConsumeOutcome::Consumed)
})
}
}
@@ -11,7 +11,7 @@ use ed25519_dalek::pkcs8::EncodePrivateKey as _;
use futures::future::BoxFuture;
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, LoadedCredential, PlaintextCredential,
ProviderAccount, ProviderAccountId,
ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest,
@@ -456,7 +456,6 @@ async fn manual_refresh_fixture(
available,
});
let service = CodexCredentialAdminService::new(
store.repository(),
refresher.clone(),
Arc::new(ManualVerifier),
leases.clone(),
@@ -478,11 +477,12 @@ async fn manual_refresh_prepares_revision_fenced_rotation_without_store_mutation
let (store, refresher, leases, service) =
manual_refresh_fixture(Ok(refreshed_tokens()), true).await;
let account_id = ProviderAccountId::new("acct_manual_refresh").expect("account id");
let current = store
.load_current_credential(&account_id)
.await
.expect("current credential");
let prepared = service
.manual_refresh(
account_id.clone(),
CredentialRevision::new(1).expect("revision"),
)
.manual_refresh(current)
.await
.expect("manual refresh");
@@ -528,24 +528,16 @@ async fn manual_refresh_prepares_revision_fenced_rotation_without_store_mutation
}
#[tokio::test]
async fn manual_refresh_stale_revision_and_missing_lease_fail_before_exchange() {
let (_, stale_refresher, _, stale_service) =
manual_refresh_fixture(Ok(refreshed_tokens()), true).await;
let account_id = ProviderAccountId::new("acct_manual_refresh").expect("account id");
let stale = stale_service
.manual_refresh(
account_id.clone(),
CredentialRevision::new(2).expect("revision"),
)
.await
.expect_err("stale revision");
assert_eq!(stale, CodexCredentialAdminError::RevisionConflict);
assert!(stale_refresher.seen.lock().expect("seen tokens").is_empty());
let (_, unavailable_refresher, _, unavailable_service) =
async fn manual_refresh_missing_lease_fails_before_exchange() {
let (store, unavailable_refresher, _, unavailable_service) =
manual_refresh_fixture(Ok(refreshed_tokens()), false).await;
let account_id = ProviderAccountId::new("acct_manual_refresh").expect("account id");
let current = store
.load_current_credential(&account_id)
.await
.expect("current credential");
let unavailable = unavailable_service
.manual_refresh(account_id, CredentialRevision::new(1).expect("revision"))
.manual_refresh(current)
.await
.expect_err("lease unavailable");
assert_eq!(
@@ -606,9 +598,7 @@ impl CodexAccountIdentityVerifier for ImportVerifier {
pub(super) fn import_service(
refresher: Arc<ManualRefresher>,
) -> provider_openai::credential::CodexCredentialAdminService {
let store = Arc::new(MemoryAccountStore::default());
CodexCredentialAdminService::new(
store.repository(),
refresher,
Arc::new(ImportVerifier),
Arc::new(ManualLeases {
@@ -628,7 +618,7 @@ pub(super) fn unused_import_refresher() -> Arc<ManualRefresher> {
}
#[tokio::test]
async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
async fn oauth_import_uses_only_tokens_and_verified_identity() {
let refresher = unused_import_refresher();
let prepared = import_service(refresher.clone())
.prepare_import_document(serde_json::json!({
@@ -645,15 +635,16 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
"accessTokenExpiresAt": "2100-01-01T00:00:00+00:00",
"status": "disabled",
"addedAt": "2026-07-18T10:47:01+08:00",
"updatedAt": "2026-07-19T11:00:00+08:00"
"updatedAt": "2026-07-19T11:00:00+08:00",
"unrelated": { "value": true }
}]
}))
.await
.expect("CPR import");
assert_eq!(prepared.accounts().len(), 1);
let account = &prepared.accounts()[0];
assert_eq!(account.account.id().as_str(), "acct_cpr_import");
assert!(!account.account.enabled());
assert!(account.account.id().as_str().starts_with("acct_"));
assert!(account.account.enabled());
assert_eq!(account.account.upstream_account_id(), Some("chatgpt-cpr"));
let runtime = CodexCredentialCodec::decode(&account.credential).expect("credential");
assert_eq!(
@@ -666,33 +657,6 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
"token-cpr"
);
assert!(refresher.seen.lock().expect("seen tokens").is_empty());
let error = import_service(unused_import_refresher())
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{"token": "token-cpr", "unexpected": true}]
}))
.await
.expect_err("unknown CPR account key");
assert_eq!(error, CodexCredentialAdminError::InvalidInput);
}
#[tokio::test]
async fn cpr_import_uses_verified_token_identity_instead_of_stale_export_projections() {
let error = import_service(unused_import_refresher())
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_verified_identity",
"accountId": "user-chatgpt-cpr",
"userId": "user-chatgpt-cpr",
"token": "token-cpr"
}]
}))
.await
.expect_err("stale document identity must not override authenticated identity");
assert_eq!(error, CodexCredentialAdminError::IdentityRejected);
}
#[tokio::test]
@@ -744,7 +708,7 @@ async fn cpr_batch_allows_distinct_users_in_the_same_workspace() {
}
#[tokio::test]
async fn credential_bundle_and_auth_document_normalize_to_the_same_core_accounts() {
async fn oauth_import_extracts_tokens_without_parsing_source_projections() {
let bundle = import_service(unused_import_refresher())
.prepare_import_document(serde_json::json!({
"exported_at": "2026-07-03T15:46:38.717Z",
@@ -754,7 +718,7 @@ async fn credential_bundle_and_auth_document_normalize_to_the_same_core_accounts
"platform": "openai",
"type": "oauth",
"credentials": {
"at": "token-bundle",
"access_token": "token-bundle",
"refresh_token": "refresh-bundle",
"chatgpt_account_id": "chatgpt-bundle",
"chatgpt_user_id": "user-chatgpt-bundle"
@@ -832,7 +796,6 @@ async fn real_cpr_fixture_import_contract() {
&std::fs::read(fixture).expect("read CPR fixture"),
)
.expect("parse CPR fixture");
let store = Arc::new(MemoryAccountStore::default());
let refresher = Arc::new(
provider_openai::credential::token_client::official_openai_token_client()
.expect("official token client"),
@@ -860,7 +823,6 @@ async fn real_cpr_fixture_import_contract() {
);
let verifier = Arc::new(CodexAccountIdentityService::new(signed, accounts));
let service = CodexCredentialAdminService::new(
store.repository(),
refresher,
verifier,
Arc::new(ManualLeases {
@@ -7,7 +7,8 @@ use chrono::Utc;
use futures::executor::block_on;
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountFeedbackStats, AccountSelectionPolicy,
ProviderAccountId, RotationStrategy,
AccountStateChange, OpaqueProviderData, ProviderAccountId, ProviderAccountStore as _,
QuotaObservation, RotationStrategy,
};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
@@ -23,6 +24,7 @@ use provider_openai::credential::{
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use secrecy::ExposeSecret;
use serde_json::json;
use url::Url;
use crate::support::{
@@ -603,7 +605,7 @@ fn credential_expired_failure_marks_unified_account_expired() {
}
#[test]
fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() {
fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
@@ -629,7 +631,7 @@ fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() {
.expect("record rate-limit failure");
let account = store.account("acct_primary").expect("account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
@@ -637,6 +639,171 @@ fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() {
);
}
#[test]
fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() + 24 * 60 * 60;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist exhausted quota");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[test]
fn rate_limited_failure_ignores_exhaustion_after_its_reset_time() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() - 1;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist expired quota exhaustion");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Cooldown
);
}
#[test]
fn rate_limited_failure_preserves_terminal_account_state() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
block_on(store.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Invalid,
reason: Some("identity_verification_required".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
}))
.expect("mark account invalid");
let current = store.account("acct_primary").expect("invalid account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&current,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Invalid
);
}
#[test]
fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let first = store.account("acct_first").expect("first account");
let second = store.account("acct_second").expect("second account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(async {
let (first_result, second_result) = futures::join!(
selector.record_failure(
&first,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
),
selector.record_failure(
&second,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
),
);
first_result.expect("record first account failure");
second_result.expect("record second account failure");
});
assert_eq!(
[
store
.account("acct_first")
.expect("persisted first account")
.availability(),
store
.account("acct_second")
.expect("persisted second account")
.availability(),
],
[AccountAvailability::Cooldown; 2]
);
}
#[test]
fn selector_keeps_a_quota_exhausted_account_eligible_when_configured_not_to_skip_it() {
let store = Arc::new(MemoryAccountStore::default());
@@ -264,6 +264,55 @@ async fn account_identity_rejects_imported_workspace_conflict() {
assert_eq!(error, CodexIdentityVerificationError::Rejected);
}
#[tokio::test]
async fn import_identity_uses_signed_claims_when_usage_rejects_the_request() {
let (signed, _) = verifier();
let service = CodexAccountIdentityService::new(
Arc::new(signed),
Arc::new(StaticAccountSource {
result: Err(CodexIdentityVerificationError::Rejected),
}),
);
let profile = service
.verify_import(&secret(
&claims(get_current_timestamp() + 3_600),
Some("refresh-token"),
))
.await
.expect("signed import identity");
assert_eq!(profile.chatgpt_account_id, "account-signed");
assert_eq!(profile.chatgpt_user_id, "user-signed");
}
#[tokio::test]
async fn import_identity_keeps_signed_workspace_when_usage_returns_a_different_workspace() {
let (signed, _) = verifier();
let service = CodexAccountIdentityService::new(
Arc::new(signed),
Arc::new(StaticAccountSource {
result: Ok(CodexAuthenticatedAccount {
chatgpt_account_id: "account-default".to_owned(),
chatgpt_user_id: "user-default".to_owned(),
email: None,
plan_type: None,
}),
}),
);
let profile = service
.verify_import(&secret(
&claims(get_current_timestamp() + 3_600),
Some("refresh-token"),
))
.await
.expect("signed import identity");
assert_eq!(profile.chatgpt_account_id, "account-signed");
assert_eq!(profile.chatgpt_user_id, "user-signed");
}
#[tokio::test]
async fn account_identity_keeps_usage_unavailable_out_of_ready_profile() {
let (signed, _) = verifier();
@@ -5,7 +5,7 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwnerBinding, PendingAuthorizationMutation,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::model::{MutationActor, MutationContext};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::routing::ProviderKind;
use provider_openai::credential::token_client::{
@@ -15,10 +15,10 @@ use provider_openai::credential::token_client::{
use provider_openai::credential::{
CodexAccountIdentityVerifier, CodexCredentialAdmin, CodexIdentityExpectation,
CodexIdentityVerification, CodexIdentityVerificationError, CodexOAuthAdmin,
CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthPendingStore,
CodexOAuthPendingStoreError, CodexOAuthSecret, CodexPendingAuthorization,
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ImportCodexOAuthCredential,
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthPendingClaimOutcome,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthSecret,
CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential,
ImportCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
};
use secrecy::{ExposeSecret, SecretString};
use sha2::{Digest as _, Sha256};
@@ -28,7 +28,24 @@ use crate::support::{MemoryAccountStore, profile, runtime_policy, secret};
#[derive(Default)]
struct PendingStore {
value: Mutex<Option<StoredCodexPendingAuthorization>>,
value: Mutex<Option<(StoredCodexPendingAuthorization, Option<String>)>>,
}
fn duplicate_pending(
pending: &StoredCodexPendingAuthorization,
) -> Result<CodexPendingAuthorization, CodexOAuthPendingStoreError> {
CodexPendingAuthorization::from_stored(StoredCodexPendingAuthorization {
flow_id: pending.flow_id.clone(),
owner_ref: pending.owner_ref.clone(),
started_request_ref: pending.started_request_ref.clone(),
name: pending.name.clone(),
expires_at: pending.expires_at,
state: pending.state.clone(),
nonce: pending.nonce.clone(),
code_verifier: pending.code_verifier.clone(),
reauthorization_account_id: pending.reauthorization_account_id.clone(),
mutation: pending.mutation.clone(),
})
}
#[async_trait]
@@ -41,42 +58,85 @@ impl CodexOAuthPendingStore for PendingStore {
if value.is_some() {
return Err(CodexOAuthPendingStoreError::Conflict);
}
*value = Some(StoredCodexPendingAuthorization {
flow_id: pending.flow_id().to_owned(),
owner_ref: pending.owner_ref().to_owned(),
started_request_ref: pending.started_request_ref().to_owned(),
name: pending.name().to_owned(),
expires_at: pending.expires_at(),
state: pending.state().clone(),
nonce: pending.nonce().clone(),
code_verifier: pending.code_verifier().clone(),
reauthorization_account_id: pending
.reauthorization()
.map(|target| target.account_id().to_string()),
reauthorization_credential_revision: pending
.reauthorization()
.map(|target| target.credential_revision().get()),
mutation: pending.mutation().clone(),
});
*value = Some((
StoredCodexPendingAuthorization {
flow_id: pending.flow_id().to_owned(),
owner_ref: pending.owner_ref().to_owned(),
started_request_ref: pending.started_request_ref().to_owned(),
name: pending.name().to_owned(),
expires_at: pending.expires_at(),
state: pending.state().clone(),
nonce: pending.nonce().clone(),
code_verifier: pending.code_verifier().clone(),
reauthorization_account_id: pending.reauthorization().map(ToString::to_string),
mutation: pending.mutation().clone(),
},
None,
));
Ok(())
}
async fn take(
async fn claim(
&self,
owner_ref: &str,
flow_id: &str,
) -> Result<Option<CodexPendingAuthorization>, CodexOAuthPendingStoreError> {
claim_ref: &str,
_claim_ttl: std::time::Duration,
) -> Result<CodexOAuthPendingClaimOutcome, CodexOAuthPendingStoreError> {
let mut value = self.value.lock().expect("pending lock");
if value
.as_ref()
.is_some_and(|pending| pending.owner_ref == owner_ref && pending.flow_id == flow_id)
{
return value
.take()
.map(CodexPendingAuthorization::from_stored)
.transpose();
let Some((pending, claim)) = value.as_mut() else {
return Ok(CodexOAuthPendingClaimOutcome::NotFound);
};
if pending.owner_ref != owner_ref || pending.flow_id != flow_id {
return Ok(CodexOAuthPendingClaimOutcome::NotFound);
}
Ok(None)
if claim.is_some() {
return Ok(CodexOAuthPendingClaimOutcome::InProgress);
}
*claim = Some(claim_ref.to_owned());
duplicate_pending(pending)
.map(Box::new)
.map(CodexOAuthPendingClaimOutcome::Claimed)
}
async fn release_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let mut value = self.value.lock().expect("pending lock");
let Some((pending, claim)) = value.as_mut() else {
return Ok(false);
};
if pending.owner_ref != owner_ref
|| pending.flow_id != flow_id
|| claim.as_deref() != Some(claim_ref)
{
return Ok(false);
}
*claim = None;
Ok(true)
}
async fn consume_claim(
&self,
owner_ref: &str,
flow_id: &str,
claim_ref: &str,
) -> Result<bool, CodexOAuthPendingStoreError> {
let mut value = self.value.lock().expect("pending lock");
let Some((pending, claim)) = value.as_ref() else {
return Ok(false);
};
if pending.owner_ref != owner_ref
|| pending.flow_id != flow_id
|| claim.as_deref() != Some(claim_ref)
{
return Ok(false);
}
*value = None;
Ok(true)
}
}
@@ -186,7 +246,6 @@ fn reauthorization_mutation(request_id: &str, account_id: &str) -> PendingAuthor
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id).expect("account id"),
expected_credential_revision: Revision::new(1).expect("revision"),
},
AuthorizationOwnerBinding::from_context(&context),
)
@@ -313,7 +372,7 @@ async fn completion_rejects_repeated_code_callback_parameter() {
.await
.expect_err("repeated code");
assert_eq!(error, CodexOAuthAdminError::UpstreamRejected);
assert_eq!(error, CodexOAuthAdminError::CallbackRejected);
}
#[tokio::test]
@@ -337,7 +396,7 @@ async fn completion_rejects_repeated_state_callback_parameter() {
.await
.expect_err("repeated state");
assert_eq!(error, CodexOAuthAdminError::UpstreamRejected);
assert_eq!(error, CodexOAuthAdminError::CallbackRejected);
}
struct UnavailableVerifier;
@@ -389,7 +448,7 @@ async fn completion_classifies_unavailable_identity_verification_as_upstream_una
}
#[tokio::test]
async fn callback_state_mismatch_fails_closed_and_consumes_flow() {
async fn callback_state_mismatch_releases_flow_for_retry() {
let service = service();
let started = started(&service).await;
let flow_id = started.flow_id;
@@ -405,14 +464,48 @@ async fn callback_state_mismatch_fails_closed_and_consumes_flow() {
.complete_authorization(command())
.await
.expect_err("bad state"),
CodexOAuthAdminError::UpstreamRejected
CodexOAuthAdminError::CallbackRejected
);
assert_eq!(
service
.complete_authorization(command())
.await
.expect_err("one shot"),
CodexOAuthAdminError::NotFound
.expect_err("retryable callback rejection"),
CodexOAuthAdminError::CallbackRejected
);
}
#[tokio::test]
async fn token_exchange_rejection_releases_flow_for_retry() {
let service = service();
let started = started(&service).await;
let flow_id = started.flow_id;
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let command = || CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: flow_id.clone(),
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=rejected-authorization-code&state={state}"
)),
};
assert_eq!(
service
.complete_authorization(command())
.await
.expect_err("token exchange rejection"),
CodexOAuthAdminError::TokenRejected
);
assert_eq!(
service
.complete_authorization(command())
.await
.expect_err("retryable token exchange rejection"),
CodexOAuthAdminError::TokenRejected
);
}
@@ -430,7 +523,7 @@ fn oauth_command_debug_redacts_owner_flow_and_callback() {
}
#[tokio::test]
async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotation() {
async fn reauthorization_uses_revision_advanced_after_start_for_prepared_rotation() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
@@ -458,6 +551,8 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
})
.await
.expect("start reauthorization");
let advanced_revision = store.advance_credential_revision("acct_oauth_reauth").await;
assert_eq!(advanced_revision.get(), 2);
let state = Url::parse(&started.authorization_url)
.expect("authorization URL")
.query_pairs()
@@ -481,7 +576,7 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
prepared.credential.account_id().as_str(),
"acct_oauth_reauth"
);
assert_eq!(prepared.credential.expected_revision().get(), 1);
assert_eq!(prepared.credential.expected_revision().get(), 2);
let runtime =
provider_openai::credential::CodexCredentialCodec::decode(prepared.credential.credential())
.expect("prepared credential");
@@ -501,7 +596,7 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
.expect("unchanged account")
.revision()
.get(),
1
2
);
}
@@ -446,18 +446,248 @@ async fn synchronize_without_accounts_is_a_noop_before_network_io() {
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_ready_accounts() {
async fn periodic_quota_synchronization_skips_fresh_ready_accounts() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_ready").await;
let account = store
.account("acct_periodic_ready")
.expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist fresh quota");
let summary = blocked_network_quota_service(&store)
.synchronize()
.await
.expect("ready account must not query upstream");
.expect("fresh ready account must not query upstream");
assert_eq!(summary, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_quota() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_stuck_cooldown").await;
let account = store
.account("acct_periodic_stuck_cooldown")
.expect("created account");
let reset_at = Utc::now().timestamp() + 24 * 60 * 60;
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account cooldown");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist exhausted quota");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(429))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service.synchronize().await.expect("quota synchronization");
assert_eq!(
store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_limit() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_exhausted").await;
let account = store
.account("acct_periodic_exhausted")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(429))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service.synchronize().await.expect("quota synchronization");
assert_eq!(
store
.account("acct_periodic_exhausted")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhausted_quota() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_temporary_cooldown").await;
let account = store
.account("acct_periodic_temporary_cooldown")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account cooldown");
let summary = blocked_network_quota_service(&store)
.synchronize()
.await
.expect("ordinary cooldown must not query upstream");
assert_eq!(summary, CodexQuotaSyncSummary::default());
assert_eq!(
store
.account("acct_periodic_temporary_cooldown")
.expect("cooldown account")
.availability(),
AccountAvailability::Cooldown
);
}
#[tokio::test]
async fn periodic_quota_synchronization_limits_stale_ready_account_scans() {
let store = Arc::new(MemoryAccountStore::default());
for index in 0..9 {
let account_id = format!("acct_periodic_stale_ready_{index}");
create_account(&store, &account_id).await;
let account = store.account(&account_id).expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::UNIX_EPOCH),
})
.await
.expect("persist stale quota");
}
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "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(),
);
let summary = service.synchronize().await.expect("quota synchronization");
assert_eq!(summary.exhausted, 8);
assert_eq!(
store
.account("acct_periodic_stale_ready_0")
.expect("scanned account")
.availability(),
AccountAvailability::QuotaExhausted
);
assert_eq!(
store
.account("acct_periodic_stale_ready_8")
.expect("deferred account")
.availability(),
AccountAvailability::Ready
);
}
#[tokio::test]
async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
let store = Arc::new(MemoryAccountStore::default());
@@ -511,6 +741,90 @@ async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
assert_eq!(summary, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_its_reset_is_due() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_reset_due").await;
let account = store
.account("acct_periodic_reset_due")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist future reset");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let first = service.synchronize().await.expect("initial refresh");
assert_eq!(first.transient, 1);
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": 1}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist due reset");
let reset_recheck = service.synchronize().await.expect("reset recheck");
let repeated = service
.synchronize()
.await
.expect("throttled reset recheck");
assert_eq!(reset_recheck.transient, 1);
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn invalid_quota_json_for_one_account_does_not_abort_the_batch() {
let store = Arc::new(MemoryAccountStore::default());
@@ -596,6 +596,75 @@ async fn transient_refresh_defers_without_invalidating_account() {
assert_eq!(refresher.seen.lock().expect("seen tokens lock").len(), 1);
}
#[tokio::test]
async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
let store = Arc::new(MemoryAccountStore::default());
seed_due_account(&store).await;
let original = store.account("acct_refresh").expect("seeded account");
let mut expired_profile = profile("chatgpt-acct_refresh");
expired_profile.access_token_expires_at =
Some(chrono::Utc::now() - chrono::Duration::minutes(5));
store
.repository()
.rotate_oauth_secret(RotateCodexCredential {
account_id: "acct_refresh".to_owned(),
expected_credential_revision: original.revision().get(),
secret: secret("old-access"),
verified_account: expired_profile,
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::hours(12)),
})
.await
.expect("seed expired access token with deferred refresh");
let deferred = store.account("acct_refresh").expect("deferred account");
store
.repository()
.apply_state(
&deferred,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
SystemTime::now(),
)
.await
.expect("mark account quota exhausted");
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
let service = CodexCredentialRefreshService::new(
store.repository(),
refresher.clone(),
Arc::new(VerifiedIdentity),
Arc::new(RefreshLeases {
available: true,
requests: Mutex::new(Vec::new()),
}),
Arc::new(CountingCredentialState::default()),
runtime_policy(),
);
let outcomes = service.refresh_due().await.expect("refresh due");
assert!(matches!(
outcomes.as_slice(),
[CodexCredentialRefreshOutcome::Refreshed {
account_id,
credential_revision: 3,
}] if account_id == "acct_refresh"
));
assert_eq!(
refresher.seen.lock().expect("seen tokens lock").as_slice(),
["rt-old-access"]
);
let refreshed = store.account("acct_refresh").expect("refreshed account");
assert_eq!(
refreshed.availability(),
AccountAvailability::QuotaExhausted
);
assert!(
refreshed
.access_token_expires_at()
.is_some_and(|expires_at| expires_at > SystemTime::now())
);
}
#[tokio::test]
async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() {
let (store, refresher, service) = setup(Ok(success_tokens()), false).await;
@@ -4,13 +4,17 @@ use std::sync::Arc;
use std::time::{Duration, SystemTime};
use chrono::Utc;
use gateway_core::engine::credential::{AccountFeedbackStats, ProviderAccountId};
use futures::StreamExt;
use gateway_core::engine::credential::{
AccountAvailability, AccountFeedbackStats, ProviderAccountId,
};
use gateway_core::engine::provider::{Provider as _, ProviderRequest};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId,
AccountAttemptContext, AttemptContext, CancellationToken, ContinuationAttempt, ModelRequestId,
RequestAttemptContext, UpstreamSendState,
};
use gateway_core::error::ProviderErrorKind;
use gateway_core::event::GatewayEvent;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
@@ -25,6 +29,12 @@ 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 tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
sync::oneshot,
time::timeout,
};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
@@ -122,6 +132,10 @@ fn provider_with_affinity_and_base_url(
}
async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
create_account_with_enabled(store, id, true).await;
}
async fn create_account_with_enabled(store: &Arc<MemoryAccountStore>, id: &str, enabled: bool) {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: id.to_owned(),
@@ -129,7 +143,7 @@ async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
secret: secret(&format!("at-{id}")),
verified_account: profile(&format!("chatgpt-{id}")),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
enabled,
})
.await;
}
@@ -147,6 +161,19 @@ fn generate_operation() -> Operation {
))
}
fn http_generate_operation() -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest {
let provider = ProviderKind::new(provider_name).expect("provider");
let upstream_model = UpstreamModelId::new("gpt-5.4").expect("upstream model");
@@ -185,6 +212,97 @@ fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext
)
}
fn diagnostic_context(request_id: &str, account_id: &str) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::diagnostic(
BTreeSet::new(),
ProviderAccountId::new(account_id).expect("account id"),
None,
),
None,
CancellationToken::new(),
)
}
async fn paused_chunked_sse_server(
first_chunk: String,
second_chunk: String,
) -> (
String,
oneshot::Sender<()>,
oneshot::Receiver<()>,
tokio::task::JoinHandle<()>,
) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind chunked SSE listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let (release_sender, release_receiver) = oneshot::channel();
let (first_chunk_sender, first_chunk_sent) = oneshot::channel();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept chunked SSE request");
read_http_request(&mut stream).await;
stream
.write_all(
b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n",
)
.await
.expect("write chunked SSE headers");
write_http_chunk(&mut stream, &first_chunk).await;
let _ = first_chunk_sender.send(());
let _ = release_receiver.await;
if !second_chunk.is_empty() {
write_http_chunk(&mut stream, &second_chunk).await;
}
stream
.write_all(b"0\r\n\r\n")
.await
.expect("terminate chunked SSE response");
});
(base_url, release_sender, first_chunk_sent, server)
}
async fn read_http_request(stream: &mut TcpStream) {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
loop {
let read = stream.read(&mut buffer).await.expect("read HTTP request");
if read == 0 {
return;
}
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
return;
}
}
}
async fn write_http_chunk(stream: &mut TcpStream, body: &str) {
stream
.write_all(format!("{:X}\r\n", body.len()).as_bytes())
.await
.expect("write HTTP chunk size");
stream
.write_all(body.as_bytes())
.await
.expect("write HTTP chunk body");
stream
.write_all(b"\r\n")
.await
.expect("terminate HTTP chunk");
stream.flush().await.expect("flush HTTP chunk");
}
#[tokio::test]
async fn openai_provider_rejects_a_foreign_provider_candidate_before_account_selection() {
let store = Arc::new(MemoryAccountStore::default());
@@ -271,6 +389,25 @@ async fn generate_without_an_eligible_openai_account_fails_before_network_io() {
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn disabled_account_is_excluded_from_normal_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
create_account_with_enabled(&store, "acct_disabled_scheduling", false).await;
let result = provider(&store)
.execute(
planned_request("openai", generate_operation()),
context("req_disabled_scheduling", CancellationToken::new()),
)
.await;
let Err(error) = result else {
panic!("disabled account must not be scheduled normally")
};
assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount);
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn opaque_provider_options_do_not_change_openai_account_selection() {
let store = Arc::new(MemoryAccountStore::default());
@@ -334,6 +471,478 @@ async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key()
assert!(keys[0].bytes().all(|byte| byte.is_ascii_hexdigit()));
}
#[tokio::test]
async fn completed_response_persists_session_affinity_before_stream_consumer_stops_polling() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_completed_affinity").await;
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_affinity\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
("prompt_cache_key".to_owned(), json!("affinity-key")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
);
let mut stream =
provider_with_affinity_and_base_url(&store, Arc::clone(&affinity), server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
context("req_completed_affinity", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
while let Some(event) = stream.next().await {
let event = event.expect("provider event");
if event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)))
{
break;
}
}
drop(stream);
assert_eq!(affinity.binding_count(), 1);
}
#[tokio::test]
async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_atomic_failure";
create_account(&store, account_id).await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_atomic_failure\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"retry_after_seconds\":17,\"response\":{\"id\":\"resp_atomic_failure\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"atomic upstream failure\"}}}\n\n"
)
.to_owned(),
String::new(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_atomic_failure", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let mut failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("response.failed must produce a typed failure"),
}
};
assert_eq!(visible_before_error, 0);
assert_eq!(failure.kind(), ProviderErrorKind::RateLimited);
assert_eq!(failure.send_state(), UpstreamSendState::Sent);
assert_eq!(failure.upstream_status(), Some(429));
assert!(failure.replay_is_safe());
let events = failure.take_atomic_client_events();
assert_eq!(
events
.iter()
.filter_map(|event| event.wire_event()?.event_type())
.collect::<Vec<_>>(),
vec!["response.created", "response.failed"]
);
let account = store.account(account_id).expect("rate-limited account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
let _ = release.send(());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first_event_latency").await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_first_event\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n"
)
.to_owned(),
concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_first_event\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"later failure\"}}}\n\n"
)
.to_owned(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_first_event_latency", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let first_event = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("first chunk must not wait for the withheld failure")
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
break next;
}
};
assert_eq!(
first_event.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
);
release.send(()).expect("release second upstream chunk");
let mut later_wire_types = Vec::new();
let failure = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("later failure must arrive after release")
.expect("provider stream must return typed failure");
match next {
Ok(event) => {
if let Some(event_type) = event.wire_event().and_then(|wire| wire.event_type()) {
later_wire_types.push(event_type.to_owned());
}
}
Err(error) => break error,
}
};
assert_eq!(later_wire_types, vec!["response.failed"]);
assert!(!failure.replay_is_safe());
assert!(!failure.has_atomic_client_events());
server.await.expect("chunked SSE server");
}
#[tokio::test]
async fn continuation_should_hold_created_until_semantic_output() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_continuation_prefetch").await;
let (base_url, release, first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_continuation\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n"
)
.to_owned(),
concat!(
"event: response.content_part.added\n",
"data: {\"type\":\"response.content_part.added\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\"}}\n\n",
"event: response.output_text.delta\n",
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_continuation\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)
.to_owned(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_continuation_prefetch", CancellationToken::new())
.with_continuation_attempt(ContinuationAttempt::Native),
)
.await
.expect("prepare provider stream");
let mut first_chunk_sent = Box::pin(first_chunk_sent);
loop {
tokio::select! {
sent = &mut first_chunk_sent => {
sent.expect("first upstream chunk");
break;
}
next = stream.next() => {
let event = next
.expect("provider stream must stay open")
.expect("provider event");
assert!(!event.has_client_event(), "created was exposed before the first chunk barrier");
}
}
}
let blocked = timeout(Duration::from_millis(100), async {
loop {
let next = stream
.next()
.await
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
return next;
}
}
})
.await;
assert!(blocked.is_err());
release.send(()).expect("release semantic output chunk");
let first_event = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("semantic output must release the prefetched batch")
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
break next;
}
};
assert_eq!(
first_event.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
);
while let Some(event) = stream.next().await {
event.expect("continuation response must complete");
}
server.await.expect("chunked SSE server");
}
#[tokio::test]
async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_bare_atomic_failure").await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_bare_failure\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"bare failure\"}}}\n\n"
)
.to_owned(),
String::new(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_bare_atomic_failure", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let mut failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("bare response.failed must produce a typed failure"),
}
};
assert_eq!(visible_before_error, 0);
assert!(failure.replay_is_safe());
let events = failure.take_atomic_client_events();
assert_eq!(events.len(), 1);
assert_eq!(
events[0].wire_event().and_then(|wire| wire.event_type()),
Some("response.failed")
);
assert!(events[0]
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Started(meta) if meta.response_id() == "resp_bare_failure")));
let _ = release.send(());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn continuation_prefetch_over_64_kib_without_a_boundary_should_fail_as_protocol_error() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_prefetch_limit").await;
let padding = "x".repeat(64 * 1024);
let body = format!(
"event: response.created\ndata: {}\n\n",
json!({
"type": "response.created",
"response": {
"id": "resp_prefetch_limit",
"model": "gpt-5.4",
"status": "in_progress",
"padding": padding,
}
})
);
assert!(body.len() > 64 * 1024);
let (base_url, release, _first_chunk_sent, server) =
paused_chunked_sse_server(body, String::new()).await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_prefetch_limit", CancellationToken::new())
.with_continuation_attempt(ContinuationAttempt::Native),
)
.await
.expect("prepare provider stream");
let mut visible_before_error = 0;
let failure = loop {
match stream.next().await {
Some(Ok(event)) => visible_before_error += usize::from(event.has_client_event()),
Some(Err(error)) => break error,
None => panic!("oversized prefetch must fail"),
}
};
assert_eq!(visible_before_error, 0);
assert_eq!(failure.kind(), ProviderErrorKind::Protocol);
assert!(!failure.has_atomic_client_events());
let _ = release.send(());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn response_failed_after_semantic_output_is_exposed_and_not_replay_safe() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_semantic_failure").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_semantic_failure\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
"event: response.content_part.added\n",
"data: {\"type\":\"response.content_part.added\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\"}}\n\n",
"event: response.output_text.delta\n",
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_semantic_failure\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"failure after output\"}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_semantic_failure", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut wire_types = Vec::new();
let failure = loop {
match stream.next().await {
Some(Ok(event)) => {
if let Some(event_type) = event.wire_event().and_then(|wire| wire.event_type()) {
wire_types.push(event_type.to_owned());
}
}
Some(Err(error)) => break error,
None => panic!("response.failed must produce a typed failure"),
}
};
assert_eq!(
wire_types,
vec![
"response.created",
"response.content_part.added",
"response.output_text.delta",
"response.failed"
]
);
assert!(!failure.replay_is_safe());
assert!(!failure.has_atomic_client_events());
}
#[tokio::test]
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");
store
.repository()
.apply_state(
&account,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
SystemTime::now(),
)
.await
.expect("seed quota-exhausted state");
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.insert_header("x-codex-active-limit", "codex")
.insert_header("x-codex-primary-used-percent", "100")
.insert_header("x-codex-primary-window-minutes", "300")
.insert_header("x-codex-primary-reset-at", "1900000000")
.insert_header("x-codex-limit-reached", "true")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_disabled_diagnostic\",\"model\":\"gpt-5.4\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_disabled_diagnostic\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream =
provider_with_affinity_and_base_url(&store, Arc::clone(&affinity), server.uri())
.execute(
planned_request("openai", http_generate_operation()),
diagnostic_context("req_disabled_diagnostic", account_id),
)
.await
.expect("disabled diagnostic should prepare a fixed-account stream");
let mut completed = false;
while let Some(event) = stream.next().await {
let event = event.expect("disabled diagnostic upstream response");
completed |= event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)));
}
assert!(completed);
let account = store
.account(account_id)
.expect("disabled test account after test");
assert!(!account.enabled());
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(!store.has_quota(account_id));
assert_eq!(affinity.binding_count(), 0);
}
#[test]
fn request_observation_reads_openai_metadata_without_changing_the_operation() {
let store = Arc::new(MemoryAccountStore::default());
@@ -67,9 +67,53 @@ impl MemoryAccountStore {
.map(|stored| stored.account.clone())
}
pub(crate) async fn advance_credential_revision(&self, id: &str) -> CredentialRevision {
let account_id = ProviderAccountId::new(id).expect("test account id");
let current = self
.load_current_credential(&account_id)
.await
.expect("current test credential");
let account = current.account;
let expected_revision = account.revision();
let update = CredentialCasUpdate::new(
account_id.clone(),
expected_revision,
ProviderAccountUpdate {
account_id,
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
plan_type: account.plan_type().map(str::to_owned),
},
current.credential,
account.has_refresh_token(),
account.access_token_expires_at(),
account.next_refresh_at(),
)
.expect("valid test credential update");
match self
.compare_and_swap_credential(update)
.await
.expect("advance test credential revision")
{
CredentialCasOutcome::Updated(revision) => revision,
CredentialCasOutcome::Conflict => panic!("test credential revision conflicted"),
}
}
pub(crate) fn quota_reads(&self) -> usize {
self.quota_reads.load(Ordering::SeqCst)
}
pub(crate) fn has_quota(&self, id: &str) -> bool {
let Some(id) = ProviderAccountId::new(id).ok() else {
return false;
};
self.accounts
.lock()
.expect("account store lock")
.get(&id)
.is_some_and(|stored| stored.quota.is_some())
}
}
#[async_trait]
@@ -130,14 +174,22 @@ impl ProviderAccountStore for MemoryAccountStore {
&self,
account: &ProviderAccountId,
expected_revision: CredentialRevision,
) -> Result<LoadedCredential, StoreError> {
let loaded = self.load_current_credential(account).await?;
if loaded.account.revision() != expected_revision {
return Err(store_error(StoreErrorKind::Conflict));
}
Ok(loaded)
}
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, StoreError> {
let accounts = self.accounts.lock().expect("account store lock");
let stored = accounts
.get(account)
.ok_or_else(|| store_error(StoreErrorKind::InvalidData))?;
if stored.account.revision() != expected_revision {
return Err(store_error(StoreErrorKind::Conflict));
}
Ok(LoadedCredential {
account: stored.account.clone(),
credential: stored.credential.clone(),
@@ -440,6 +492,10 @@ impl MemorySessionAffinity {
pub(crate) fn lookup_keys(&self) -> Vec<String> {
self.lookups.lock().expect("session affinity lock").clone()
}
pub(crate) fn binding_count(&self) -> usize {
self.bindings.lock().expect("session affinity lock").len()
}
}
impl ProviderSessionAffinityPort for MemorySessionAffinity {
@@ -528,6 +528,40 @@ fn raw_sse_failure_should_keep_the_original_frame_before_reporting_the_typed_fai
assert!(!format!("{failure:?}").contains("raw failure marker"));
}
#[test]
fn bare_response_failed_should_project_started_identity_with_the_fallback_model() {
let raw = concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_bare_failed\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"bare failure\"}}}\n\n",
);
let failure = CodexCanonicalDecoder::new("fallback-model")
.push(raw.as_bytes())
.expect_err("bare response.failed remains a typed lifecycle failure");
assert!(matches!(
canonical_facts(failure.events()).as_slice(),
[GatewayEvent::Started(metadata)]
if metadata.response_id() == "resp_bare_failed"
&& metadata.model() == "fallback-model"
));
}
#[test]
fn response_failed_without_identity_should_keep_the_upstream_typed_failure() {
let raw = concat!(
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"missing identity\"}}}\n\n",
);
let failure = CodexCanonicalDecoder::new("fallback")
.push(raw.as_bytes())
.expect_err("missing identity must not replace the upstream lifecycle failure");
assert!(matches!(failure.error(), CodexCanonicalError::Upstream(_)));
assert!(canonical_facts(failure.events()).is_empty());
}
#[test]
fn decoder_should_preserve_same_chunk_output_before_typed_failure() {
let marker = "same-chunk-secret-marker";
@@ -143,3 +143,19 @@ fn structured_expiry_code_classifies_without_an_auth_status() {
CodexFailureCategory::CredentialExpired
);
}
#[test]
fn structured_quota_signal_takes_precedence_over_http_429() {
assert_eq!(
classify(
429,
r#"{"error":{"code":"insufficient_quota","message":"rate limit reached"}}"#
),
CodexFailureCategory::QuotaExhausted
);
}
#[test]
fn bare_http_429_remains_a_temporary_rate_limit() {
assert_eq!(classify(429, ""), CodexFailureCategory::RateLimited);
}
+270 -139
View File
@@ -2,7 +2,7 @@
use std::collections::BTreeSet;
use std::sync::Arc;
use std::time::SystemTime;
use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
@@ -14,8 +14,8 @@ use gateway_admin::model::observability::{
ProviderBillingInput,
};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwner, AuthorizationOwnerBinding,
AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard,
AuthorizationCommitGuard, AuthorizationMutationTarget, AuthorizationOwner,
AuthorizationOwnerBinding, AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedAuthorizationCredential,
PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation,
@@ -23,7 +23,7 @@ use gateway_admin::model::provider_credentials::{
ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::model::{AdminError, MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
@@ -33,9 +33,9 @@ use gateway_core::engine::credential::{
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError,
ProviderStoreErrorKind,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use serde::Deserialize;
@@ -58,10 +58,78 @@ use crate::credential::{
use crate::transport::profile::{GrokCliReleaseSnapshot, GrokCliReleaseStatus};
use crate::transport::{GROK_CLI_BASE_URL, XAI_PROVIDER_NAME, grok_billing_breakdown};
const PENDING_SCHEMA_VERSION: u64 = 2;
const PENDING_SCHEMA_VERSION: u64 = 3;
const PENDING_TTL: TimeDelta = TimeDelta::minutes(30);
const PENDING_CLAIM_TTL: Duration = Duration::from_secs(90);
const MAX_PENDING_TEXT_BYTES: usize = 512;
struct ClaimedXaiAuthorization {
authorization: StoredXaiAuthorization,
owner_ref: String,
claim_ref: String,
}
struct XaiAuthorizationCommitGuard {
pending: Arc<dyn OAuthPendingFlowPort>,
provider_kind: ProviderKind,
flow: OAuthPendingBinding,
owner: OAuthPendingBinding,
claim: OAuthPendingBinding,
}
impl XaiAuthorizationCommitGuard {
fn new(
pending: Arc<dyn OAuthPendingFlowPort>,
provider_kind: ProviderKind,
flow_id: &str,
owner_ref: &str,
claim_ref: &str,
) -> Result<Self, ProviderAdminError> {
Ok(Self {
pending,
provider_kind,
flow: binding(flow_id)?,
owner: binding(owner_ref)?,
claim: binding(claim_ref)?,
})
}
}
#[async_trait]
impl AuthorizationCommitGuard for XaiAuthorizationCommitGuard {
async fn commit(self: Box<Self>) -> Result<(), AdminError> {
match self
.pending
.consume_claim(&self.provider_kind, &self.flow, &self.owner, &self.claim)
.await
.map_err(map_pending_claim_settlement_error)?
{
OAuthPendingConsumeOutcome::Consumed => Ok(()),
OAuthPendingConsumeOutcome::NotFound
| OAuthPendingConsumeOutcome::OwnerMismatch
| OAuthPendingConsumeOutcome::ClaimMismatch => Err(AdminError::conflict(
"xAI OAuth pending claim is no longer current",
)),
}
}
async fn abort(self: Box<Self>) -> Result<(), AdminError> {
match self
.pending
.release_claim(&self.provider_kind, &self.flow, &self.owner, &self.claim)
.await
.map_err(map_pending_claim_settlement_error)?
{
OAuthPendingReleaseOutcome::Released => Ok(()),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => Err(AdminError::unavailable(
"xAI OAuth pending claim could not be released",
)),
}
}
}
pub(crate) struct XaiAdminProvider {
provider_kind: ProviderKind,
wire_profile: XaiWireProfileState,
@@ -172,31 +240,81 @@ impl XaiAdminProvider {
}
}
async fn take_pending(
async fn claim_pending(
&self,
context: &MutationContext,
flow_id: &str,
) -> Result<StoredXaiAuthorization, ProviderAdminError> {
) -> Result<ClaimedXaiAuthorization, ProviderAdminError> {
let flow = binding(flow_id)?;
let owner_binding = AuthorizationOwnerBinding::from_context(context);
let owner_ref = owner_ref(owner_binding.owner());
let owner = binding(&owner_ref)?;
let claim_ref = random_flow_id()?;
let claim = binding(&claim_ref)?;
match self
.pending
.take_if_owner(&self.provider_kind, &flow, &owner)
.claim_if_owner(
&self.provider_kind,
&flow,
&owner,
&claim,
PENDING_CLAIM_TTL,
)
.await
.map_err(map_provider_store_error)?
{
OAuthPendingTakeOutcome::Taken(payload) => {
let stored = decode_pending(payload, &self.oauth_config)?;
if stored.flow_id != flow_id || stored.owner_ref != owner_ref {
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
Ok(stored)
OAuthPendingClaimOutcome::Claimed(payload) => {
let stored = match decode_pending(payload, &self.oauth_config) {
Ok(stored) if stored.flow_id == flow_id && stored.owner_ref == owner_ref => {
stored
}
Ok(_) => {
self.release_pending_claim(flow_id, &owner_ref, &claim_ref)
.await?;
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
Err(error) => {
self.release_pending_claim(flow_id, &owner_ref, &claim_ref)
.await?;
return Err(error);
}
};
Ok(ClaimedXaiAuthorization {
authorization: stored,
owner_ref,
claim_ref,
})
}
OAuthPendingTakeOutcome::NotFound | OAuthPendingTakeOutcome::OwnerMismatch => {
OAuthPendingClaimOutcome::NotFound | OAuthPendingClaimOutcome::OwnerMismatch => {
Err(provider_error(ProviderAdminErrorKind::NotFound))
}
OAuthPendingClaimOutcome::InProgress => {
Err(provider_error(ProviderAdminErrorKind::Conflict))
}
}
}
async fn release_pending_claim(
&self,
flow_id: &str,
owner_ref: &str,
claim_ref: &str,
) -> Result<(), ProviderAdminError> {
let flow = binding(flow_id)?;
let owner = binding(owner_ref)?;
let claim = binding(claim_ref)?;
match self
.pending
.release_claim(&self.provider_kind, &flow, &owner, &claim)
.await
.map_err(map_provider_store_error)?
{
OAuthPendingReleaseOutcome::Released => Ok(()),
OAuthPendingReleaseOutcome::NotFound
| OAuthPendingReleaseOutcome::OwnerMismatch
| OAuthPendingReleaseOutcome::ClaimMismatch => {
Err(provider_error(ProviderAdminErrorKind::Unavailable))
}
}
}
}
@@ -402,16 +520,10 @@ impl ProviderAdmin for XaiAdminProvider {
if pending.provider_kind() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
if let AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} = pending.target()
{
let revision = CredentialRevision::new(expected_credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
if let AuthorizationMutationTarget::Reauthorize { account_id } = pending.target() {
let current = self
.accounts
.load_credential(account_id, revision)
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind {
@@ -433,76 +545,38 @@ impl ProviderAdmin for XaiAdminProvider {
&self,
command: CompleteAuthorization,
) -> Result<PreparedAuthorizationCommit, ProviderAdminError> {
let stored = self
.take_pending(&command.context, &command.flow_id)
let ClaimedXaiAuthorization {
authorization: stored,
owner_ref,
claim_ref,
} = self
.claim_pending(&command.context, &command.flow_id)
.await?;
let authorization = PendingAuthorization::from_server_state(
&self.oauth_config,
&SecretValue::new(stored.server_state),
)
.map_err(map_oauth_error)?;
let callback = callback(&command.callback_url)?;
let grant = authorization
.accept_callback(callback)
.map_err(map_oauth_error)?;
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
let tokens = self
.oauth
.exchange_authorization_code(&discovery, grant)
.await
.map_err(map_oauth_error)?;
let policy = self
.runtime_policy
.load_refresh_policy()
.await
.map_err(map_provider_store_error)?;
let credential = match stored.mutation.target() {
AuthorizationMutationTarget::Create { name } => {
let prepared = GrokCredentialAdmin
.prepare_verified_account(
&VerifiedGrokAccount {
account_id: ProviderAccountId::new(format!(
"acct_{}",
Uuid::now_v7().simple()
))
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
name: name.clone(),
email: None,
upstream_account_id: None,
plan_type: None,
tokens,
enabled: true,
},
policy,
)
.map_err(map_repository_error)?;
PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?)
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
let revision = CredentialRevision::new(expected_credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?;
let current = self
.accounts
.load_credential(account_id, revision)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::NotFound));
}
let prepared = verified_rotation(current, tokens, policy)?;
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
prepared,
self.provider_kind.clone(),
)?)
let authorization_guard = match XaiAuthorizationCommitGuard::new(
Arc::clone(&self.pending),
self.provider_kind.clone(),
&command.flow_id,
&owner_ref,
&claim_ref,
) {
Ok(guard) => guard,
Err(error) => {
self.release_pending_claim(&command.flow_id, &owner_ref, &claim_ref)
.await?;
return Err(error);
}
};
Ok(PreparedAuthorizationCommit {
pending: stored.mutation,
credential,
})
let completed = self
.complete_claimed_authorization(stored, &command.callback_url)
.await;
match completed {
Ok(completed) => Ok(completed.with_authorization_guard(Box::new(authorization_guard))),
Err(error) => {
self.release_pending_claim(&command.flow_id, &owner_ref, &claim_ref)
.await?;
Err(error)
}
}
}
async fn prepare_rotation(
@@ -510,16 +584,11 @@ impl ProviderAdmin for XaiAdminProvider {
command: PrepareCredentialRotation,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
validate_account_record(&command.account, &self.provider_kind)?;
if command.account.credential_revision != command.expected_credential_revision {
return Err(provider_error(ProviderAdminErrorKind::Conflict));
}
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let revision = CredentialRevision::new(command.expected_credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let current = self
.accounts
.load_credential(&account_id, revision)
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
@@ -573,15 +642,17 @@ impl ProviderAdmin for XaiAdminProvider {
validate_account_record(&command.account, &self.provider_kind)?;
let account_id = ProviderAccountId::new(command.account.id.clone())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let current = self.account(&account_id).await?;
if !account_matches_record(&current, &command.account) {
let current = self
.accounts
.load_current_credential(&account_id)
.await
.map_err(map_store_error)?;
if !account_matches_record(&current.account, &command.account) {
return Err(provider_error(ProviderAdminErrorKind::Conflict));
}
let revision = CredentialRevision::new(command.account.credential_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let prepared = self
.refresh
.prepare_manual_refresh(&account_id, revision)
.prepare_manual_refresh(current)
.await
.map_err(map_refresh_error)?;
prepared_rotation(prepared, command.account.provider_kind)
@@ -689,6 +760,77 @@ impl ProviderAdmin for XaiAdminProvider {
}
}
impl XaiAdminProvider {
async fn complete_claimed_authorization(
&self,
stored: StoredXaiAuthorization,
callback_url: &str,
) -> Result<PreparedAuthorizationCommit, ProviderAdminError> {
let authorization = PendingAuthorization::from_server_state(
&self.oauth_config,
&SecretValue::new(stored.server_state),
)
.map_err(map_oauth_error)?;
let callback = callback(callback_url)?;
let grant = authorization
.accept_callback(callback)
.map_err(map_oauth_error)?;
let discovery = self.oauth.discover().await.map_err(map_oauth_error)?;
let tokens = self
.oauth
.exchange_authorization_code(&discovery, grant)
.await
.map_err(map_oauth_error)?;
let policy = self
.runtime_policy
.load_refresh_policy()
.await
.map_err(map_provider_store_error)?;
let credential = match stored.mutation.target() {
AuthorizationMutationTarget::Create { name } => {
let prepared = GrokCredentialAdmin
.prepare_verified_account(
&VerifiedGrokAccount {
account_id: ProviderAccountId::new(format!(
"acct_{}",
Uuid::now_v7().simple()
))
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
name: name.clone(),
email: None,
upstream_account_id: None,
plan_type: None,
tokens,
enabled: true,
},
policy,
)
.map_err(map_repository_error)?;
PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?)
}
AuthorizationMutationTarget::Reauthorize { account_id } => {
let current = self
.accounts
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::NotFound));
}
let prepared = verified_rotation(current, tokens, policy)?;
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
prepared,
self.provider_kind.clone(),
)?)
}
};
Ok(PreparedAuthorizationCommit::new(
stored.mutation,
credential,
))
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RotationDocument {
@@ -889,19 +1031,9 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool {
account.id().as_str() == record.id
&& account.provider() == &record.provider_kind
&& account.revision().get() == record.credential_revision.get()
&& account.name() == record.name
&& account.email() == record.email.as_deref()
&& account.upstream_user_id() == record.upstream_user_id
&& account.upstream_account_id() == record.upstream_account_id.as_deref()
&& account.plan_type() == record.plan_type.as_deref()
&& account.enabled() == record.enabled
&& account.availability() == record.availability
&& account.cooldown_until().map(DateTime::<Utc>::from) == record.cooldown_until
&& account.access_token_expires_at().map(DateTime::<Utc>::from)
== record.access_token_expires_at
&& account.next_refresh_at().map(DateTime::<Utc>::from) == record.next_refresh_at
&& account.has_refresh_token() == record.has_refresh_token
&& account.authentication_kind() == record.authentication_kind
}
fn project_quota(
@@ -1053,13 +1185,8 @@ struct PendingMutationDocument {
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum PendingTargetDocument {
Create {
name: String,
},
Reauthorize {
account_id: String,
expected_credential_revision: u64,
},
Create { name: String },
Reauthorize { account_id: String },
}
#[derive(Deserialize)]
@@ -1124,19 +1251,12 @@ fn encode_target(target: &AuthorizationMutationTarget) -> Map<String, Value> {
document.insert("kind".to_owned(), Value::String("create".to_owned()));
document.insert("name".to_owned(), Value::String(name.clone()));
}
AuthorizationMutationTarget::Reauthorize {
account_id,
expected_credential_revision,
} => {
AuthorizationMutationTarget::Reauthorize { account_id } => {
document.insert("kind".to_owned(), Value::String("reauthorize".to_owned()));
document.insert(
"account_id".to_owned(),
Value::String(account_id.to_string()),
);
document.insert(
"expected_credential_revision".to_owned(),
Value::Number(Number::from(expected_credential_revision.get())),
);
}
}
document
@@ -1202,15 +1322,12 @@ fn decode_mutation(
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let target = match document.target {
PendingTargetDocument::Create { name } => AuthorizationMutationTarget::Create { name },
PendingTargetDocument::Reauthorize {
account_id,
expected_credential_revision,
} => AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
expected_credential_revision: Revision::new(expected_credential_revision)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
},
PendingTargetDocument::Reauthorize { account_id } => {
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
}
}
};
let actor = match document.owner {
PendingOwnerDocument::AdminSession { admin_user_id } => {
@@ -1342,6 +1459,20 @@ fn map_provider_store_error(error: ProviderStoreError) -> ProviderAdminError {
})
}
fn map_pending_claim_settlement_error(error: ProviderStoreError) -> AdminError {
match error.kind() {
ProviderStoreErrorKind::InvalidData => {
AdminError::internal("xAI OAuth pending claim is invalid")
}
ProviderStoreErrorKind::Conflict => {
AdminError::conflict("xAI OAuth pending claim conflicts with current state")
}
ProviderStoreErrorKind::Unavailable => {
AdminError::unavailable("xAI OAuth pending claim store is unavailable")
}
}
}
fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError {
provider_error(match error.kind() {
StoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict,
@@ -9,7 +9,7 @@ use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::{StreamExt as _, stream};
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, ProviderAccountId,
AccountAvailability, CredentialRevision, LoadedCredential, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderCredentialStatePort, ProviderLeaseAcquisition,
@@ -343,16 +343,13 @@ impl GrokCredentialRefreshService {
self.refresh_one_with_policy(credential, policy).await
}
/// 手工刷新一个指定 revision;只返回 Provider 验证后的 CAS command,不写 Store。
/// 手工刷新一次原子读取的当前 credential;只返回 Provider 验证后的 CAS command,不写 Store。
pub async fn prepare_manual_refresh(
&self,
account_id: &ProviderAccountId,
expected_revision: CredentialRevision,
current: LoadedCredential,
) -> Result<PreparedGrokCredentialRotation, GrokCredentialRefreshError> {
let current = self
.repository
.load_managed(account_id, expected_revision)
.await?;
let account_id = current.account.id().clone();
let expected_revision = current.account.revision();
let loaded = super::repository::loaded_from_core(current.clone())?;
let policy = self.runtime_policy.load_refresh_policy().await?;
if loaded
@@ -407,7 +404,7 @@ impl GrokCredentialRefreshService {
return Err(GrokCredentialRefreshError::InvalidRefreshResponse);
}
let (access_token_expires_at, next_refresh_at) =
refreshed_deadlines(account_id, tokens.expires_in, policy)
refreshed_deadlines(&account_id, tokens.expires_in, policy)
.ok_or(GrokCredentialRefreshError::InvalidRefreshResponse)?;
let prepared = GrokCredentialAdmin
.prepare_rotation(&RotateManagedGrokCredential {
@@ -463,21 +463,6 @@ impl GrokCredentialRepository {
loaded_from_core(loaded)
}
pub(crate) async fn load_managed(
&self,
account_id: &ProviderAccountId,
revision: CredentialRevision,
) -> Result<LoadedCredential, GrokCredentialRepositoryError> {
let loaded = self
.store
.load_credential(account_id, revision)
.await
.map_err(map_rotation_load_error)?;
ensure_xai(&loaded.account)?;
decode_secret(&loaded.credential)?;
Ok(loaded)
}
/// 读取 xAI Provider 的全部账号;credential 逐行按 revision fence 加载。
pub(crate) async fn list_loaded_for_provider(
&self,
@@ -528,8 +513,12 @@ impl GrokCredentialRepository {
&self,
account_id: &ProviderAccountId,
) -> Result<LoadedGrokCredential, GrokCredentialRepositoryError> {
let account = self.ensure_xai_account(account_id).await?;
self.load(account_id, account.revision()).await
let loaded = self
.store
.load_current_credential(account_id)
.await
.map_err(map_store_error)?;
loaded_from_core(loaded)
}
pub(crate) async fn replace_quota(
@@ -6,14 +6,15 @@ use std::time::{Duration, SystemTime};
use chrono::Utc;
use gateway_core::engine::credential::{
AccountCandidate, AccountFeedbackStats, AccountSelectionContext, AccountSelector,
ProviderAccount, ProviderAccountId,
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderLeaseAcquisition, ProviderLeasePort,
ProviderLeaseRequest, ProviderSchedulingLeaseRequest,
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderLeaseAcquisition,
ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingLeaseRequest,
ProviderScopedCooldown,
};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService};
use super::repository::GrokCredentialRepository;
@@ -28,6 +29,22 @@ const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
const STREAM_INTERRUPTION_COOLDOWN: Duration = Duration::from_secs(30);
const UNAUTHORIZED_COOLDOWN: Duration = Duration::from_secs(60);
const PAYMENT_REQUIRED_COOLDOWN: Duration = Duration::from_secs(60);
const MODEL_ACCESS_DENIED_COOLDOWN: Duration = Duration::from_secs(5 * 60);
const MODEL_QUOTA_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
const MAX_MODEL_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RuntimeCooldownScope {
Account,
Model,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct RuntimeCooldown {
until: SystemTime,
scope: RuntimeCooldownScope,
}
/// 仅经 Core account port、TTL catalog cache 和 Redis lease 选择一个 OAuth session。
pub struct GrokAccountSessionSelector {
@@ -66,39 +83,48 @@ impl GrokAccountSessionSelector {
&self,
request: GrokSessionSelection,
) -> Result<SelectedGrokSession, GrokSessionSelectorError> {
let diagnostic = request.availability() == AccountAvailabilityPolicy::BypassForDiagnostic;
let accounts = self
.repository
.list_accounts_for_provider()
.await
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
self.quota.prepare_scheduling(&accounts).await;
// 目录支持度按套餐 scope 去重后并发批量读取,避免逐账号串行往返。
let account_scopes = accounts
.iter()
.map(|account| GrokCatalogScope::for_account(account).ok())
.collect::<Vec<_>>();
let unique_scopes = account_scopes
.iter()
.flatten()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let support_reads = futures::future::join_all(unique_scopes.iter().map(|scope| {
self.catalog_cache
.observed_model_support(scope, request.upstream_model().as_str())
}))
.await;
let support_by_scope = unique_scopes
.into_iter()
.zip(support_reads)
.collect::<std::collections::BTreeMap<_, _>>();
let mut catalog_eligible = Vec::new();
for (account, scope) in accounts.into_iter().zip(account_scopes) {
let unsupported = scope
.is_some_and(|scope| matches!(support_by_scope.get(&scope), Some(Ok(Some(false)))));
if !unsupported {
catalog_eligible.push(account);
}
if !diagnostic {
self.quota.prepare_scheduling(&accounts).await;
}
let catalog_eligible = if diagnostic {
accounts
} else {
// 目录支持度按套餐 scope 去重后并发批量读取,避免逐账号串行往返。
let account_scopes = accounts
.iter()
.map(|account| GrokCatalogScope::for_account(account).ok())
.collect::<Vec<_>>();
let unique_scopes = account_scopes
.iter()
.flatten()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let support_reads = futures::future::join_all(unique_scopes.iter().map(|scope| {
self.catalog_cache
.observed_model_support(scope, request.upstream_model().as_str())
}))
.await;
let support_by_scope = unique_scopes
.into_iter()
.zip(support_reads)
.collect::<std::collections::BTreeMap<_, _>>();
accounts
.into_iter()
.zip(account_scopes)
.filter_map(|(account, scope)| {
let unsupported = scope.is_some_and(|scope| {
matches!(support_by_scope.get(&scope), Some(Ok(Some(false))))
});
(!unsupported).then_some(account)
})
.collect()
};
if catalog_eligible.is_empty() {
return Err(GrokSessionSelectorError::NoEligibleSession);
}
@@ -112,7 +138,12 @@ impl GrokAccountSessionSelector {
.load_state(&self.provider_kind, &account_ids)
.await
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
let runtime_cooldowns = self.runtime_cooldowns(&catalog_eligible).await;
let runtime_cooldowns = if diagnostic {
std::collections::BTreeMap::new()
} else {
self.runtime_cooldowns(&catalog_eligible, request.upstream_model())
.await
};
let mut candidates = catalog_eligible
.into_iter()
.filter(|account| !runtime_cooldowns.contains_key(account.id()))
@@ -147,12 +178,14 @@ impl GrokAccountSessionSelector {
excluded_accounts: request.excluded_accounts().clone(),
preferred_account: request.required_account().cloned().or(affinity_account),
round_robin_cursor: scheduling.round_robin_cursor(),
availability: request.availability(),
};
let mut capacity_denied = false;
let mut retry_after = None;
while let Some(selected) = AccountSelector.select(&candidates, &context) {
let selected_id = selected.account.id().clone();
let selected_revision = selected.account.revision();
let allows_account_state_mutation = !diagnostic || selected.account.enabled();
let lease = self
.scheduling
.try_acquire(ProviderLeaseRequest::Scheduling(
@@ -186,15 +219,16 @@ impl GrokAccountSessionSelector {
.load(&selected_id, selected_revision)
.await
.map_err(|_| GrokSessionSelectorError::InvalidSession)?;
if loaded
.refresh_token_expires_at
.is_some_and(|expires_at| expires_at <= Utc::now())
if !diagnostic
&& loaded
.refresh_token_expires_at
.is_some_and(|expires_at| expires_at <= Utc::now())
{
return Err(GrokSessionSelectorError::InvalidSession);
}
let binding = GrokSessionBinding::new(selected_id.as_str())
.map_err(|_| GrokSessionSelectorError::InvalidSession)?;
return SelectedGrokSession::new(
let session = SelectedGrokSession::new(
selected_id,
selected_revision,
loaded.access_token,
@@ -206,23 +240,38 @@ impl GrokAccountSessionSelector {
binding,
guard,
)
.map_err(|_| GrokSessionSelectorError::InvalidSession);
.map_err(|_| GrokSessionSelectorError::InvalidSession)?;
return Ok(if allows_account_state_mutation {
session
} else {
session.without_account_state_mutation()
});
}
if capacity_denied {
return Err(GrokSessionSelectorError::CapacityUnavailable { retry_after });
}
// 钉死账号时只看该账号,否则取最早退出 cooldown 的候选。
let cooled_until = match request.required_account() {
let cooled = match request.required_account() {
Some(required) => runtime_cooldowns.get(required).copied(),
None => runtime_cooldowns.values().copied().min(),
None => runtime_cooldowns
.values()
.copied()
.min_by_key(|item| item.until),
};
if let Some(until) = cooled_until {
return Err(GrokSessionSelectorError::AccountCoolingDown {
retry_after: until
.duration_since(SystemTime::now())
.ok()
.filter(|remaining| !remaining.is_zero()),
if let Some(cooled) = cooled {
let retry_after = cooled
.until
.duration_since(SystemTime::now())
.ok()
.filter(|remaining| !remaining.is_zero());
return Err(match cooled.scope {
RuntimeCooldownScope::Account => {
GrokSessionSelectorError::AccountCoolingDown { retry_after }
}
RuntimeCooldownScope::Model => {
GrokSessionSelectorError::ModelCoolingDown { retry_after }
}
});
}
Err(GrokSessionSelectorError::NoEligibleSession)
@@ -231,28 +280,56 @@ impl GrokAccountSessionSelector {
async fn runtime_cooldowns(
&self,
accounts: &[ProviderAccount],
) -> std::collections::BTreeMap<ProviderAccountId, SystemTime> {
upstream_model: &UpstreamModelId,
) -> std::collections::BTreeMap<ProviderAccountId, RuntimeCooldown> {
let now = SystemTime::now();
// 冷却状态并发批量读取,避免逐账号串行往返。
let reads = futures::future::join_all(
let model_scope = ProviderCooldownScope::upstream_model(upstream_model.clone());
// 账号级和当前模型级冷却都并发批量读取,避免逐账号串行往返。
let account_reads = futures::future::join_all(
accounts
.iter()
.map(|account| self.cooldowns.read(account.id())),
)
.await;
);
let model_reads = futures::future::join_all(
accounts
.iter()
.map(|account| self.cooldowns.read_scoped(account.id(), &model_scope)),
);
let (account_reads, model_reads) = futures::join!(account_reads, model_reads);
let mut cooled = std::collections::BTreeMap::new();
for (account, read) in accounts.iter().zip(reads) {
let Ok(Some(cooldown)) = read else {
continue;
};
if cooldown.credential_revision() != account.revision() {
if cooldown.credential_revision() < account.revision() {
let _ = self.cooldowns.clear(account.id(), account.revision()).await;
for ((account, account_read), model_read) in
accounts.iter().zip(account_reads).zip(model_reads)
{
if let Ok(Some(cooldown)) = account_read {
if cooldown.credential_revision() != account.revision() {
if cooldown.credential_revision() < account.revision() {
let _ = self.cooldowns.clear(account.id(), account.revision()).await;
}
} else if cooldown.until() > now {
insert_runtime_cooldown(
&mut cooled,
account.id(),
cooldown.until(),
RuntimeCooldownScope::Account,
);
}
continue;
}
if cooldown.until() > now {
cooled.insert(account.id().clone(), cooldown.until());
if let Ok(Some(cooldown)) = model_read {
if cooldown.credential_revision() != account.revision() {
if cooldown.credential_revision() < account.revision() {
let _ = self
.cooldowns
.clear_scoped(account.id(), &model_scope, account.revision())
.await;
}
} else if cooldown.until() > now {
insert_runtime_cooldown(
&mut cooled,
account.id(),
cooldown.until(),
RuntimeCooldownScope::Model,
);
}
}
}
cooled
@@ -271,6 +348,26 @@ impl GrokAccountSessionSelector {
))
.await;
}
async fn record_model_runtime_cooldown(
&self,
session: &SelectedGrokSession,
upstream_model: UpstreamModelId,
duration: Duration,
) {
let cooldown_until = SystemTime::now()
.checked_add(duration)
.unwrap_or_else(SystemTime::now);
let _ = self
.cooldowns
.put_scoped_if_later(ProviderScopedCooldown::new(
session.account_id().clone(),
session.credential_revision(),
ProviderCooldownScope::upstream_model(upstream_model),
cooldown_until,
))
.await;
}
}
impl fmt::Debug for GrokAccountSessionSelector {
@@ -297,6 +394,9 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
failure: GrokCredentialFailure,
) -> GrokCredentialFeedbackFuture<'a> {
Box::pin(async move {
if !session.allows_account_state_mutation() {
return;
}
let observed_at = Utc::now();
let persistent = match failure {
GrokCredentialFailure::Unauthorized => {
@@ -309,6 +409,43 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
"upstream_quota_exhausted",
None,
),
GrokCredentialFailure::FreeQuotaExhausted => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_free_quota_exhausted",
None,
),
GrokCredentialFailure::PaymentRequired { retry_after } => {
let retry_after = bounded_cooldown(
retry_after,
PAYMENT_REQUIRED_COOLDOWN,
MAX_RATE_LIMIT_COOLDOWN,
);
self.record_runtime_cooldown(session, retry_after).await;
return;
}
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model,
retry_after,
} => {
let retry_after =
bounded_cooldown(retry_after, MODEL_QUOTA_COOLDOWN, MAX_MODEL_COOLDOWN);
self.record_model_runtime_cooldown(session, upstream_model, retry_after)
.await;
return;
}
GrokCredentialFailure::ModelAccessDenied {
upstream_model,
retry_after,
} => {
let retry_after = bounded_cooldown(
retry_after,
MODEL_ACCESS_DENIED_COOLDOWN,
MAX_MODEL_COOLDOWN,
);
self.record_model_runtime_cooldown(session, upstream_model, retry_after)
.await;
return;
}
GrokCredentialFailure::RateLimited { retry_after } => {
let retry_after = retry_after
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
@@ -336,6 +473,51 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
.await;
})
}
fn record_success<'a>(
&'a self,
session: &'a SelectedGrokSession,
) -> GrokCredentialFeedbackFuture<'a> {
Box::pin(async move {
if !session.allows_account_state_mutation() {
return;
}
let Ok(current) = self.repository.load_current(session.account_id()).await else {
return;
};
let account = &current.account;
if account.revision() != session.credential_revision()
|| !account.enabled()
|| !matches!(
account.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
)
{
return;
}
if let Err(error) = self
.repository
.update_state(&UpdateGrokCredentialState {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: GrokCredentialAvailability::Ready,
availability_reason: None,
cooldown_until: None,
observed_at: Utc::now(),
})
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
"xAI account state recovery after successful upstream response failed"
);
}
})
}
}
fn minimum_retry_after(left: Option<Duration>, right: Option<Duration>) -> Option<Duration> {
@@ -345,3 +527,29 @@ fn minimum_retry_after(left: Option<Duration>, right: Option<Duration>) -> Optio
(None, None) => None,
}
}
fn bounded_cooldown(
retry_after: Option<Duration>,
default: Duration,
maximum: Duration,
) -> Duration {
retry_after
.filter(|duration| !duration.is_zero())
.unwrap_or(default)
.min(maximum)
}
fn insert_runtime_cooldown(
cooldowns: &mut std::collections::BTreeMap<ProviderAccountId, RuntimeCooldown>,
account_id: &ProviderAccountId,
until: SystemTime,
scope: RuntimeCooldownScope,
) {
let replacement = RuntimeCooldown { until, scope };
if cooldowns
.get(account_id)
.is_none_or(|current| current.until < until)
{
cooldowns.insert(account_id.clone(), replacement);
}
}
+184 -30
View File
@@ -9,8 +9,8 @@ use base64::Engine as _;
use futures::{StreamExt, future::BoxFuture};
use gateway_core::engine::continuation::ContinuationBinding;
use gateway_core::engine::credential::{
AccountAvailability, AccountFeedbackStats, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
AccountAvailability, AccountAvailabilityPolicy, AccountFeedbackStats, ProviderAccount,
ProviderAccountId, ProviderAccountStore,
};
use gateway_core::engine::provider::{
EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration,
@@ -229,6 +229,7 @@ impl GrokBuildProvider {
output_items: BTreeMap::new(),
});
let selected = Arc::new(selected);
let allows_account_state_mutation = selected.allows_account_state_mutation();
let metadata = provider_call_metadata(candidate, &selected)?;
let events = cold_http_sse_stream(
Arc::clone(&self.selector),
@@ -245,8 +246,12 @@ impl GrokBuildProvider {
session_capture,
},
);
Ok(ProviderStream::new(metadata, events, selected)
.with_account_feedback(Arc::clone(&self.account_feedback)))
let stream = ProviderStream::new(metadata, events, selected);
Ok(if allows_account_state_mutation {
stream.with_account_feedback(Arc::clone(&self.account_feedback))
} else {
stream
})
}
async fn execute_compaction(
@@ -288,6 +293,7 @@ impl GrokBuildProvider {
)
.await?,
);
let allows_account_state_mutation = selected.allows_account_state_mutation();
let metadata = provider_call_metadata(candidate, &selected)?;
let events = cold_compaction_http_sse_stream(
Arc::clone(&self.selector),
@@ -304,8 +310,12 @@ impl GrokBuildProvider {
session: Arc::clone(&selected),
},
);
Ok(ProviderStream::new(metadata, events, selected)
.with_account_feedback(Arc::clone(&self.account_feedback)))
let stream = ProviderStream::new(metadata, events, selected);
Ok(if allows_account_state_mutation {
stream.with_account_feedback(Arc::clone(&self.account_feedback))
} else {
stream
})
}
}
@@ -324,6 +334,11 @@ async fn select_grok_session(
context.account_selection_policy(),
context.deadline(),
)
.with_availability_policy(if context.is_diagnostic_required_account() {
AccountAvailabilityPolicy::BypassForDiagnostic
} else {
AccountAvailabilityPolicy::Enforce
})
.with_affinity(affinity);
let selection_deadline = remaining(context.deadline())
.ok_or_else(|| provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent))?;
@@ -912,6 +927,7 @@ async fn start_grok_inference(
transport: &dyn GrokInferenceTransport,
credential_recovery: &dyn GrokCredentialRecovery,
request: GrokInferenceRequest,
upstream_model: &UpstreamModelId,
context: &AttemptContext,
session: &SelectedGrokSession,
) -> Result<AcceptedGrokInference, GrokInferenceStartFailure> {
@@ -950,6 +966,7 @@ async fn start_grok_inference(
InferenceBoundary::Response(Ok(response)) => response,
InferenceBoundary::Response(Err(error)) => {
let observation = xai_error_observation(&error).ok();
let credential_failure = transport_credential_failure(&error, upstream_model);
let error =
map_continuation_failure(context, map_transport_error_for_context(error, context));
let error = recover_or_record_failure(
@@ -957,6 +974,7 @@ async fn start_grok_inference(
credential_recovery,
session,
error,
credential_failure,
context.credential_recovery_attempted(),
)
.await;
@@ -978,6 +996,7 @@ async fn next_grok_chunk(
body: &mut GrokInferenceChunkStream,
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
upstream_model: &UpstreamModelId,
context: &AttemptContext,
) -> Result<Option<bytes::Bytes>, ProviderError> {
let Some(stream_deadline) = remaining(context.deadline()) else {
@@ -999,10 +1018,12 @@ async fn next_grok_chunk(
)),
chunk = body.next() => match chunk {
Some(Ok(chunk)) => Ok(Some(chunk)),
Some(Err(error)) => {
let error = map_stream_error(error);
Err(record_stream_failure(selector, session, error).await)
}
Some(Err(error)) => Err(map_and_record_stream_transport_failure(
selector,
session,
error,
upstream_model,
).await),
None => Ok(None),
},
}
@@ -1052,6 +1073,7 @@ fn cold_compaction_http_sse_stream(
transport.as_ref(),
credential_recovery.as_ref(),
inference_request,
&upstream_model,
&context,
&session,
)
@@ -1077,6 +1099,7 @@ fn cold_compaction_http_sse_stream(
&mut body,
selector.as_ref(),
&session,
&upstream_model,
&context,
)
.await
@@ -1108,6 +1131,7 @@ fn cold_compaction_http_sse_stream(
let summary = summary.finish().map_err(map_compaction_decode_error)?;
let started = facts.started.ok_or_else(protocol_sent)?;
let upstream_completed = facts.completed.is_some();
let completed = facts.completed.unwrap_or_else(|| started.clone());
let (created, output_done, terminal) = crate::transport::compaction::compaction_wire_events(
&started,
@@ -1120,6 +1144,9 @@ fn cold_compaction_http_sse_stream(
.map_err(|_| protocol_sent())?
.into_parts();
ensure_sent_context(&context)?;
if upstream_completed && session.allows_account_state_mutation() {
selector.record_success(&session).await;
}
yield ProviderEvent::canonical_with_wire(vec![GatewayEvent::Started(started)], created);
yield ProviderEvent::wire(output_done);
let mut terminal_facts = facts.accounting;
@@ -1268,6 +1295,7 @@ fn cold_http_sse_stream(
InferenceBoundary::Response(Ok(response)) => response,
InferenceBoundary::Response(Err(error)) => {
let observation = xai_error_observation(&error)?;
let credential_failure = transport_credential_failure(&error, &upstream_model);
let error =
map_continuation_failure(&context, map_transport_error_for_context(error, &context));
let error = recover_or_record_failure(
@@ -1275,6 +1303,7 @@ fn cold_http_sse_stream(
credential_recovery.as_ref(),
&session,
error,
credential_failure,
context.credential_recovery_attempted(),
)
.await;
@@ -1309,10 +1338,12 @@ fn cold_http_sse_stream(
)),
chunk = body.next() => match chunk {
Some(Ok(chunk)) => Ok(Some(chunk)),
Some(Err(error)) => {
let error = map_stream_error(error);
Err(record_stream_failure(selector.as_ref(), &session, error).await)
},
Some(Err(error)) => Err(map_and_record_stream_transport_failure(
selector.as_ref(),
&session,
error,
&upstream_model,
).await),
None => Ok(None),
},
}?;
@@ -1333,6 +1364,9 @@ fn cold_http_sse_stream(
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
attach_xai_session_update(&mut events, &mut session_capture)?;
if completed && session.allows_account_state_mutation() {
selector.record_success(&session).await;
}
for event in events {
ensure_sent_context(&context)?;
yield event;
@@ -1350,7 +1384,14 @@ fn cold_http_sse_stream(
return;
}
};
let completed = final_events
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
attach_xai_session_update(&mut final_events, &mut session_capture)?;
if completed && session.allows_account_state_mutation() {
selector.record_success(&session).await;
}
for event in final_events {
ensure_sent_context(&context)?;
yield event;
@@ -1449,6 +1490,9 @@ async fn record_failure(
session: &SelectedGrokSession,
error: ProviderError,
) -> ProviderError {
if !session.allows_account_state_mutation() {
return error;
}
let failure = match error.kind() {
ProviderErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized),
ProviderErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited {
@@ -1463,11 +1507,26 @@ async fn record_failure(
error
}
async fn record_credential_failure(
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
error: ProviderError,
failure: GrokCredentialFailure,
) -> ProviderError {
if session.allows_account_state_mutation() {
selector.record_failure(session, failure).await;
}
error
}
async fn record_stream_failure(
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
error: ProviderError,
) -> ProviderError {
if !session.allows_account_state_mutation() {
return error;
}
if matches!(
error.kind(),
ProviderErrorKind::Unauthorized
@@ -1489,8 +1548,12 @@ async fn recover_or_record_failure(
recovery: &dyn GrokCredentialRecovery,
session: &SelectedGrokSession,
error: ProviderError,
credential_failure: Option<GrokCredentialFailure>,
recovery_attempted: bool,
) -> ProviderError {
if !session.allows_account_state_mutation() {
return error;
}
if error.requires_credential_recovery() && !recovery_attempted {
return match recovery
.recover_unauthorized(session.account_id(), session.credential_revision())
@@ -1498,12 +1561,68 @@ async fn recover_or_record_failure(
{
GrokCredentialRecoveryOutcome::Recovered => error.with_same_account_retry(),
GrokCredentialRecoveryOutcome::Rejected => error,
GrokCredentialRecoveryOutcome::Unavailable => {
record_failure(selector, session, error).await
}
GrokCredentialRecoveryOutcome::Unavailable => match credential_failure {
Some(failure) => record_credential_failure(selector, session, error, failure).await,
None => error,
},
};
}
record_failure(selector, session, error).await
match credential_failure {
Some(failure) => record_credential_failure(selector, session, error, failure).await,
None => error,
}
}
fn transport_credential_failure(
error: &GrokInferenceTransportError,
upstream_model: &UpstreamModelId,
) -> Option<GrokCredentialFailure> {
match error.kind() {
GrokInferenceTransportErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized),
GrokInferenceTransportErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited {
retry_after: error.retry_after(),
}),
GrokInferenceTransportErrorKind::QuotaExhausted => {
Some(GrokCredentialFailure::QuotaExhausted)
}
GrokInferenceTransportErrorKind::FreeQuotaExhausted => {
Some(GrokCredentialFailure::FreeQuotaExhausted)
}
GrokInferenceTransportErrorKind::PaymentRequired => {
Some(GrokCredentialFailure::PaymentRequired {
retry_after: error.retry_after(),
})
}
GrokInferenceTransportErrorKind::ModelQuotaExhausted => {
Some(GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: upstream_model.clone(),
retry_after: error.retry_after(),
})
}
GrokInferenceTransportErrorKind::ModelAccessDenied => {
Some(GrokCredentialFailure::ModelAccessDenied {
upstream_model: upstream_model.clone(),
retry_after: error.retry_after(),
})
}
_ => None,
}
}
async fn map_and_record_stream_transport_failure(
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
error: GrokInferenceTransportError,
upstream_model: &UpstreamModelId,
) -> ProviderError {
let request_scoped = error.kind() == GrokInferenceTransportErrorKind::SafetyRejected;
let credential_failure = transport_credential_failure(&error, upstream_model);
let error = map_stream_error(error);
match credential_failure {
Some(failure) => record_credential_failure(selector, session, error, failure).await,
None if request_scoped => error,
None => record_stream_failure(selector, session, error).await,
}
}
fn map_continuation_failure(context: &AttemptContext, error: ProviderError) -> ProviderError {
@@ -1590,6 +1709,11 @@ fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError {
cooling_down_message(retry_after),
"account_cooling_down",
),
GrokSessionSelectorError::ModelCoolingDown { retry_after } => (
retry_after,
model_cooling_down_message(retry_after),
"model_cooling_down",
),
GrokSessionSelectorError::CapacityUnavailable { retry_after } => (
retry_after,
"account is at its concurrency or request-interval limit".to_owned(),
@@ -1627,10 +1751,20 @@ fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError {
fn cooling_down_message(retry_after: Option<Duration>) -> String {
match retry_after {
Some(retry_after) => format!(
"account is cooling down after an interrupted upstream stream; retry in {}s",
"account is cooling down after an upstream failure; retry in {}s",
retry_after.as_secs().saturating_add(1)
),
None => "account is cooling down after an interrupted upstream stream".to_owned(),
None => "account is cooling down after an upstream failure".to_owned(),
}
}
fn model_cooling_down_message(retry_after: Option<Duration>) -> String {
match retry_after {
Some(retry_after) => format!(
"all eligible accounts are cooling down for this model; retry in {}s",
retry_after.as_secs().saturating_add(1)
),
None => "all eligible accounts are cooling down for this model".to_owned(),
}
}
@@ -1652,13 +1786,19 @@ fn map_transport_error_with_state(
forced_send_state: Option<UpstreamSendState>,
allow_explicit_replay: bool,
) -> ProviderError {
let kind = match error.kind() {
let transport_kind = error.kind();
let kind = match transport_kind {
GrokInferenceTransportErrorKind::InvalidRequest => ProviderErrorKind::InvalidRequest,
GrokInferenceTransportErrorKind::Unsupported => ProviderErrorKind::Unsupported,
GrokInferenceTransportErrorKind::Unauthorized => ProviderErrorKind::Unauthorized,
GrokInferenceTransportErrorKind::PermissionDenied => ProviderErrorKind::PermissionDenied,
GrokInferenceTransportErrorKind::PermissionDenied
| GrokInferenceTransportErrorKind::ModelAccessDenied
| GrokInferenceTransportErrorKind::PaymentRequired
| GrokInferenceTransportErrorKind::SafetyRejected => ProviderErrorKind::PermissionDenied,
GrokInferenceTransportErrorKind::RateLimited => ProviderErrorKind::RateLimited,
GrokInferenceTransportErrorKind::QuotaExhausted => ProviderErrorKind::QuotaExhausted,
GrokInferenceTransportErrorKind::QuotaExhausted
| GrokInferenceTransportErrorKind::FreeQuotaExhausted
| GrokInferenceTransportErrorKind::ModelQuotaExhausted => ProviderErrorKind::QuotaExhausted,
GrokInferenceTransportErrorKind::Timeout => ProviderErrorKind::Timeout,
GrokInferenceTransportErrorKind::Transport => ProviderErrorKind::Transport,
GrokInferenceTransportErrorKind::Protocol => ProviderErrorKind::Protocol,
@@ -1673,7 +1813,7 @@ fn map_transport_error_with_state(
mapped = mapped.with_status(status);
if allow_explicit_replay
&& forced_send_state.is_none()
&& explicit_rejection_is_replay_safe(kind, status)
&& explicit_rejection_is_replay_safe(transport_kind, status)
{
mapped = mapped.with_replay_safe();
}
@@ -1687,6 +1827,9 @@ fn map_transport_error_with_state(
if let Some(code) = error.upstream_code().cloned() {
mapped = mapped.with_upstream_code(code);
}
if let Some(detail) = error.client_visible_upstream_error().cloned() {
mapped = mapped.with_client_visible_upstream_error(detail);
}
if error.requires_credential_recovery() {
mapped = mapped.with_credential_recovery().with_replay_safe();
}
@@ -1696,14 +1839,25 @@ fn map_transport_error_with_state(
mapped
}
fn explicit_rejection_is_replay_safe(kind: ProviderErrorKind, status: u16) -> bool {
fn explicit_rejection_is_replay_safe(kind: GrokInferenceTransportErrorKind, status: u16) -> bool {
matches!(
(kind, status),
(ProviderErrorKind::Unauthorized, 401)
| (ProviderErrorKind::QuotaExhausted, 402)
| (ProviderErrorKind::QuotaExhausted, 403)
| (ProviderErrorKind::QuotaExhausted, 429)
| (ProviderErrorKind::RateLimited, 429)
(GrokInferenceTransportErrorKind::Unauthorized, 401)
| (
GrokInferenceTransportErrorKind::QuotaExhausted,
402 | 403 | 429
)
| (
GrokInferenceTransportErrorKind::FreeQuotaExhausted,
402 | 403 | 429
)
| (
GrokInferenceTransportErrorKind::ModelQuotaExhausted,
402 | 403 | 429
)
| (GrokInferenceTransportErrorKind::PaymentRequired, 402)
| (GrokInferenceTransportErrorKind::ModelAccessDenied, 403)
| (GrokInferenceTransportErrorKind::RateLimited, 429)
)
}
@@ -19,6 +19,7 @@ use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder};
use serde_json::Value;
use super::request::{GrokResponseTransform, GrokResponsesRequest};
use super::scrub_account_fingerprints;
const CONTENTS_PER_OUTPUT: u32 = 1_024;
const LONG_CONTEXT_THRESHOLD: u64 = 200_000;
@@ -820,53 +821,6 @@ fn upstream_error_field<'a>(value: &'a Value, field: &str) -> Option<&'a str> {
.and_then(Value::as_str)
}
const UUID_TEXT_LEN: usize = 36;
/// 把 message 中的 UUID 替换为占位符,控制字符归一为空格,
/// 使文本满足客户端可见错误的约束且不携带可定位账号的标识。
fn scrub_account_fingerprints(message: &str) -> String {
let bytes = message.as_bytes();
let mut scrubbed = String::with_capacity(message.len());
let mut index = 0;
while index < message.len() {
if uuid_at(bytes, index) {
scrubbed.push_str("[redacted]");
index += UUID_TEXT_LEN;
continue;
}
let character = message[index..].chars().next().expect("char boundary");
scrubbed.push(if character.is_control() {
' '
} else {
character
});
index += character.len_utf8();
}
scrubbed
}
fn uuid_at(bytes: &[u8], index: usize) -> bool {
let Some(candidate) = bytes.get(index..index + UUID_TEXT_LEN) else {
return false;
};
if index > 0 && bytes[index - 1].is_ascii_alphanumeric() {
return false;
}
if bytes
.get(index + UUID_TEXT_LEN)
.is_some_and(u8::is_ascii_alphanumeric)
{
return false;
}
candidate.iter().enumerate().all(|(offset, byte)| {
if matches!(offset, 8 | 13 | 18 | 23) {
*byte == b'-'
} else {
byte.is_ascii_hexdigit()
}
})
}
fn protocol_error(_error: impl std::fmt::Debug) -> ProviderError {
protocol_error_marker()
}
@@ -54,3 +54,50 @@ pub use transport::{
GrokInferenceTransportError, GrokInferenceTransportErrorKind, GrokInferenceTransportFuture,
GrokInferenceTransportMetrics,
};
const UUID_TEXT_LEN: usize = 36;
/// 把 message 中的 UUID 替换为占位符,控制字符归一为空格,
/// 使文本满足客户端可见错误的约束且不携带可定位账号的标识。
fn scrub_account_fingerprints(message: &str) -> String {
let bytes = message.as_bytes();
let mut scrubbed = String::with_capacity(message.len());
let mut index = 0;
while index < message.len() {
if uuid_at(bytes, index) {
scrubbed.push_str("[redacted]");
index += UUID_TEXT_LEN;
continue;
}
let character = message[index..].chars().next().expect("char boundary");
scrubbed.push(if character.is_control() {
' '
} else {
character
});
index += character.len_utf8();
}
scrubbed
}
fn uuid_at(bytes: &[u8], index: usize) -> bool {
let Some(candidate) = bytes.get(index..index + UUID_TEXT_LEN) else {
return false;
};
if index > 0 && bytes[index - 1].is_ascii_alphanumeric() {
return false;
}
if bytes
.get(index + UUID_TEXT_LEN)
.is_some_and(u8::is_ascii_alphanumeric)
{
return false;
}
candidate.iter().enumerate().all(|(offset, byte)| {
if matches!(offset, 8 | 13 | 18 | 23) {
*byte == b'-'
} else {
byte.is_ascii_hexdigit()
}
})
}
@@ -30,9 +30,11 @@ use crate::{
OAuthHttpResponse, OAuthHttpTransport, TransportFailure, TransportFailureKind, TransportFuture,
};
use gateway_core::engine::UpstreamSendState;
use gateway_core::error::SafeUpstreamValue;
use gateway_core::error::{ClientVisibleUpstreamError, SafeUpstreamValue};
use gateway_core::event::UpstreamHttpVersion;
use super::scrub_account_fingerprints;
pub(crate) const OFFICIAL_OAUTH_HOST: &str = "auth.x.ai";
const OFFICIAL_INFERENCE_HOST: &str = "cli-chat-proxy.grok.com";
const OFFICIAL_INFERENCE_PATH: &str = "/v1/responses";
@@ -990,29 +992,22 @@ async fn classify_inference_status(response: Response) -> GrokInferenceTransport
Ok(BoundedBody::TooLarge) | Err(_) => Vec::new(),
};
let metadata = inference_error_metadata(&body);
let credential_recovery_required = status == StatusCode::UNAUTHORIZED
|| (status == StatusCode::FORBIDDEN && forbidden_requires_recovery(&metadata));
let kind = match status {
StatusCode::BAD_REQUEST
| StatusCode::NOT_FOUND
| StatusCode::CONFLICT
| StatusCode::UNPROCESSABLE_ENTITY => GrokInferenceTransportErrorKind::InvalidRequest,
StatusCode::UNAUTHORIZED => GrokInferenceTransportErrorKind::Unauthorized,
StatusCode::PAYMENT_REQUIRED => GrokInferenceTransportErrorKind::QuotaExhausted,
StatusCode::FORBIDDEN => classify_forbidden(&metadata),
StatusCode::PAYMENT_REQUIRED => classify_payment_required(&metadata),
StatusCode::FORBIDDEN => classify_forbidden(&metadata, &body),
StatusCode::REQUEST_TIMEOUT | StatusCode::GATEWAY_TIMEOUT => {
GrokInferenceTransportErrorKind::Timeout
}
StatusCode::TOO_MANY_REQUESTS => {
if quota_exhausted(&metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else {
GrokInferenceTransportErrorKind::RateLimited
}
}
StatusCode::TOO_MANY_REQUESTS => classify_too_many_requests(&metadata),
status if status.is_server_error() => GrokInferenceTransportErrorKind::Unavailable,
_ => GrokInferenceTransportErrorKind::Protocol,
};
let credential_recovery_required = kind == GrokInferenceTransportErrorKind::Unauthorized;
let mut error = GrokInferenceTransportError::new(kind, UpstreamSendState::Sent)
.with_status(status_code)
.with_response_facts(http_version, request_id)
@@ -1022,6 +1017,15 @@ async fn classify_inference_status(response: Response) -> GrokInferenceTransport
} else {
metadata.code.as_deref().and_then(normalize_failure_code)
};
if let Some(message) = metadata.client_message.as_deref()
&& let Ok(detail) = ClientVisibleUpstreamError::new(
scrub_account_fingerprints(message),
upstream_code.clone(),
metadata.error_type.clone(),
)
{
error = error.with_client_visible_upstream_error(detail);
}
if let Some(code) = upstream_code.and_then(|code| SafeUpstreamValue::new(code).ok()) {
error = error.with_upstream_code(code);
}
@@ -1039,6 +1043,8 @@ struct InferenceErrorMetadata {
code: Option<String>,
error_type: Option<String>,
message: Option<String>,
// 仅结构化 JSON 的 message 可进入客户端协议;纯文本正文只参与内部分类。
client_message: Option<String>,
}
fn inference_error_metadata(body: &[u8]) -> InferenceErrorMetadata {
@@ -1053,6 +1059,9 @@ fn inference_error_metadata(body: &[u8]) -> InferenceErrorMetadata {
return InferenceErrorMetadata::default();
};
let nested = root.get("error").and_then(Value::as_object);
let message = nested
.and_then(|error| first_string(error, &["message", "error"]))
.or_else(|| first_string(root, &["message", "error"]));
InferenceErrorMetadata {
code: nested
.and_then(|error| first_string(error, &["code", "error_code"]))
@@ -1060,9 +1069,8 @@ fn inference_error_metadata(body: &[u8]) -> InferenceErrorMetadata {
error_type: nested
.and_then(|error| first_string(error, &["type", "error_type"]))
.or_else(|| first_string(root, &["type", "error_type"])),
message: nested
.and_then(|error| first_string(error, &["message", "error"]))
.or_else(|| first_string(root, &["message", "error"])),
client_message: message.clone(),
message,
}
}
@@ -1075,22 +1083,53 @@ fn first_string(object: &serde_json::Map<String, Value>, fields: &[&str]) -> Opt
.map(ToOwned::to_owned)
}
fn classify_forbidden(metadata: &InferenceErrorMetadata) -> GrokInferenceTransportErrorKind {
let text = inference_metadata_text(metadata);
if credential_rejected(&text) {
GrokInferenceTransportErrorKind::Unauthorized
} else if quota_exhausted(metadata) {
fn classify_payment_required(metadata: &InferenceErrorMetadata) -> GrokInferenceTransportErrorKind {
if model_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
} else if paid_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else if free_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::FreeQuotaExhausted
} else {
GrokInferenceTransportErrorKind::PaymentRequired
}
}
fn classify_forbidden(
metadata: &InferenceErrorMetadata,
body: &[u8],
) -> GrokInferenceTransportErrorKind {
let text = inference_metadata_text(metadata);
let body_text = String::from_utf8_lossy(body).to_ascii_lowercase();
if safety_rejected(&text) || safety_rejected(&body_text) {
GrokInferenceTransportErrorKind::SafetyRejected
} else if credential_rejected(&text) {
GrokInferenceTransportErrorKind::Unauthorized
} else if model_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
} else if paid_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else if free_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::FreeQuotaExhausted
} else if model_access_denied(metadata) {
GrokInferenceTransportErrorKind::ModelAccessDenied
} else {
GrokInferenceTransportErrorKind::PermissionDenied
}
}
fn forbidden_requires_recovery(metadata: &InferenceErrorMetadata) -> bool {
let text = inference_metadata_text(metadata);
credential_rejected(&text)
|| text.contains("access to the chat endpoint is denied")
|| text.trim_matches([' ', '.', '!', '\t', '\r', '\n']) == "access denied"
fn classify_too_many_requests(
metadata: &InferenceErrorMetadata,
) -> GrokInferenceTransportErrorKind {
if model_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
} else if paid_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else if free_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::FreeQuotaExhausted
} else {
GrokInferenceTransportErrorKind::RateLimited
}
}
fn reasoning_decode_failed(metadata: &InferenceErrorMetadata) -> bool {
@@ -1100,7 +1139,17 @@ fn reasoning_decode_failed(metadata: &InferenceErrorMetadata) -> bool {
})
}
fn quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
fn paid_quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
contains_any(
&inference_metadata_text(metadata),
&[
"personal-team-blocked:spending-limit",
"personal_team_blocked_spending_limit",
],
)
}
fn free_quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
contains_any(
&inference_metadata_text(metadata),
&[
@@ -1110,12 +1159,29 @@ fn quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
"free_usage_exhausted",
"used all the included free usage",
"used all your free usage",
"personal-team-blocked:spending-limit",
"personal_team_blocked_spending_limit",
],
)
}
fn model_quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
inference_metadata_text(metadata).contains("used all the included free usage for model")
}
fn model_access_denied(metadata: &InferenceErrorMetadata) -> bool {
metadata.message.as_deref().is_some_and(|message| {
let message = message.to_ascii_lowercase();
message.contains("access to the chat endpoint is denied")
|| message.trim_matches([' ', '.', '!', '\t', '\r', '\n']) == "access denied"
})
}
fn safety_rejected(text: &str) -> bool {
contains_any(
text,
&["content violates usage guidelines", "safety_check_type_"],
)
}
fn inference_metadata_text(metadata: &InferenceErrorMetadata) -> String {
[
metadata.code.as_deref().unwrap_or_default(),
@@ -5,7 +5,7 @@ use std::pin::Pin;
use std::time::{Duration, SystemTime};
use gateway_core::engine::credential::{
AccountSelectionPolicy, CredentialRevision, ProviderAccountId,
AccountAvailabilityPolicy, AccountSelectionPolicy, CredentialRevision, ProviderAccountId,
};
use gateway_core::engine::provider::ProviderResource;
use gateway_core::routing::UpstreamModelId;
@@ -87,6 +87,7 @@ pub struct SelectedGrokSession {
user_id: SecretValue,
email: Option<SecretValue>,
binding: GrokSessionBinding,
allows_account_state_mutation: bool,
_guard: Box<dyn GrokSessionLeaseGuard>,
}
@@ -120,10 +121,18 @@ impl SelectedGrokSession {
user_id,
email,
binding,
allows_account_state_mutation: true,
_guard: Box::new(guard),
})
}
/// 标记禁用账号的管理端诊断只读取真实上游结果,不回写账号状态。
#[must_use]
pub(crate) const fn without_account_state_mutation(mut self) -> Self {
self.allows_account_state_mutation = false;
self
}
/// 返回选中的账号 ID。
#[must_use]
pub const fn account_id(&self) -> &ProviderAccountId {
@@ -168,6 +177,11 @@ impl SelectedGrokSession {
pub const fn binding(&self) -> &GrokSessionBinding {
&self.binding
}
#[must_use]
pub(crate) const fn allows_account_state_mutation(&self) -> bool {
self.allows_account_state_mutation
}
}
impl fmt::Debug for SelectedGrokSession {
@@ -180,6 +194,10 @@ impl fmt::Debug for SelectedGrokSession {
.field("user_id", &"[REDACTED]")
.field("email", &self.email.as_ref().map(|_| "[REDACTED]"))
.field("binding", &self.binding)
.field(
"allows_account_state_mutation",
&self.allows_account_state_mutation,
)
.field("guard", &"[LEASE]")
.finish()
}
@@ -192,6 +210,7 @@ pub struct GrokSessionSelection {
excluded_accounts: BTreeSet<ProviderAccountId>,
required_account: Option<ProviderAccountId>,
account_selection_policy: AccountSelectionPolicy,
availability: AccountAvailabilityPolicy,
affinity: Option<GrokSessionAffinityKey>,
deadline: SystemTime,
}
@@ -211,6 +230,7 @@ impl GrokSessionSelection {
excluded_accounts,
required_account,
account_selection_policy,
availability: AccountAvailabilityPolicy::Enforce,
affinity: None,
deadline,
}
@@ -223,6 +243,16 @@ impl GrokSessionSelection {
self
}
/// 为固定账号的管理端诊断指定本地可用性判定策略。
#[must_use]
pub(crate) const fn with_availability_policy(
mut self,
availability: AccountAvailabilityPolicy,
) -> Self {
self.availability = availability;
self
}
/// 返回冻结的上游模型。
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
@@ -247,6 +277,11 @@ impl GrokSessionSelection {
self.account_selection_policy
}
#[must_use]
pub const fn availability(&self) -> AccountAvailabilityPolicy {
self.availability
}
/// 返回不含原始客户端会话值的账号亲和键。
#[must_use]
pub const fn affinity(&self) -> Option<&GrokSessionAffinityKey> {
@@ -265,8 +300,8 @@ pub type GrokSessionSelectorFuture<'a> = Pin<
Box<dyn Future<Output = Result<SelectedGrokSession, GrokSessionSelectorError>> + Send + 'a>,
>;
/// 凭据维度的上游失败,可安全落库为可用性反馈。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// 已脱敏的上游失败反馈;selector 决定账号持久状态或细粒度运行时 cooldown。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GrokCredentialFailure {
/// 选中的 OAuth token 被拒绝。
Unauthorized,
@@ -275,8 +310,25 @@ pub enum GrokCredentialFailure {
/// 已解析并脱敏的延迟;不保留上游原始 header。
retry_after: Option<Duration>,
},
/// 选中的 OAuth 账号配额耗尽。
/// 选中的 OAuth 账号付费额度耗尽。
QuotaExhausted,
/// 选中的 OAuth 账号订阅级免费额度耗尽。
FreeQuotaExhausted,
/// HTTP 402 没有给出可证明账号额度耗尽的稳定信号。
PaymentRequired {
/// 已解析并脱敏的延迟;缺失时由 selector 使用短暂默认值。
retry_after: Option<Duration>,
},
/// 选中账号只耗尽了指定模型的额度。
ModelQuotaExhausted {
upstream_model: UpstreamModelId,
retry_after: Option<Duration>,
},
/// 选中账号只缺少指定模型的访问权限。
ModelAccessDenied {
upstream_model: UpstreamModelId,
retry_after: Option<Duration>,
},
/// 成功的 SSE 响应在终止 Responses 事件前结束。
StreamInterrupted,
}
@@ -290,12 +342,22 @@ pub trait GrokSessionSelector: Send + Sync {
/// 执行一次选择,不做 Provider 内部回退。
fn select(&self, request: GrokSessionSelection) -> GrokSessionSelectorFuture<'_>;
/// 持久化一条已分类的失败,不重试也不替换原错误。
/// 记录一条已分类的失败反馈,不重试也不替换原错误。
fn record_failure<'a>(
&'a self,
session: &'a SelectedGrokSession,
failure: GrokCredentialFailure,
) -> GrokCredentialFeedbackFuture<'a>;
/// 成功完成真实上游请求后收敛账号的可恢复状态。
///
/// 默认实现让只关心选择或失败反馈的测试替身无需模拟持久状态写入。
fn record_success<'a>(
&'a self,
_: &'a SelectedGrokSession,
) -> GrokCredentialFeedbackFuture<'a> {
Box::pin(async {})
}
}
/// 不含密钥的选择器失败。
@@ -316,6 +378,12 @@ pub enum GrokSessionSelectorError {
/// 最早恢复的会话剩余冷却时长。
retry_after: Option<Duration>,
},
/// 每个可用会话都只针对当前模型处于运行时冷却窗口内。
#[error("Grok Build model is cooling down for the available accounts")]
ModelCoolingDown {
/// 最早恢复的账号+模型组合剩余冷却时长。
retry_after: Option<Duration>,
},
/// 会话元数据或 Provider 持有的明文密钥非法。
#[error("Grok Build session data is invalid")]
InvalidSession,
@@ -5,7 +5,7 @@ use std::time::Duration;
use futures::Stream;
use gateway_core::engine::UpstreamSendState;
use gateway_core::error::SafeUpstreamValue;
use gateway_core::error::{ClientVisibleUpstreamError, SafeUpstreamValue};
use gateway_core::event::UpstreamHttpVersion;
use url::Url;
use zeroize::Zeroizing;
@@ -278,8 +278,18 @@ pub enum GrokInferenceTransportErrorKind {
PermissionDenied,
/// 会话级或 Provider 级限流。
RateLimited,
/// 账号额度或配额耗尽。
/// 账号付费额度或 spending limit 耗尽。
QuotaExhausted,
/// 账号订阅级免费额度耗尽。
FreeQuotaExhausted,
/// 当前账号只耗尽了指定模型的免费额度。
ModelQuotaExhausted,
/// 当前账号只缺少指定模型的访问权限。
ModelAccessDenied,
/// 上游要求付费,但没有给出可证明账号额度耗尽的稳定信号。
PaymentRequired,
/// 当前请求被内容安全策略拒绝,不代表账号或模型不可用。
SafetyRejected,
/// 截止时间或 transport 超时。
Timeout,
/// 网络/TLS/连接失败。
@@ -301,7 +311,8 @@ pub struct GrokInferenceTransportError {
retry_after: Option<Duration>,
http_version: Option<UpstreamHttpVersion>,
request_id: Option<SafeUpstreamValue>,
upstream_code: Option<SafeUpstreamValue>,
upstream_code: Option<Box<SafeUpstreamValue>>,
client_visible_upstream_error: Option<Box<ClientVisibleUpstreamError>>,
transport_metrics: GrokInferenceTransportMetrics,
credential_recovery_required: bool,
sensitive_context_redacted: bool,
@@ -319,6 +330,7 @@ impl GrokInferenceTransportError {
http_version: None,
request_id: None,
upstream_code: None,
client_visible_upstream_error: None,
transport_metrics: GrokInferenceTransportMetrics {
headers_ms: None,
client_cache_status: None,
@@ -359,7 +371,14 @@ impl GrokInferenceTransportError {
/// 附着从错误 JSON 中提取并清洗后的稳定机器码。
#[must_use]
pub fn with_upstream_code(mut self, code: SafeUpstreamValue) -> Self {
self.upstream_code = Some(code);
self.upstream_code = Some(Box::new(code));
self
}
/// 附着仅供原客户端协议展示的结构化上游错误。
#[must_use]
pub fn with_client_visible_upstream_error(mut self, error: ClientVisibleUpstreamError) -> Self {
self.client_visible_upstream_error = Some(Box::new(error));
self
}
@@ -420,8 +439,13 @@ impl GrokInferenceTransportError {
}
#[must_use]
pub const fn upstream_code(&self) -> Option<&SafeUpstreamValue> {
self.upstream_code.as_ref()
pub fn upstream_code(&self) -> Option<&SafeUpstreamValue> {
self.upstream_code.as_deref()
}
#[must_use]
pub fn client_visible_upstream_error(&self) -> Option<&ClientVisibleUpstreamError> {
self.client_visible_upstream_error.as_deref()
}
#[must_use]
@@ -452,6 +476,13 @@ impl fmt::Debug for GrokInferenceTransportError {
.field("http_version", &self.http_version)
.field("request_id", &self.request_id)
.field("upstream_code", &self.upstream_code)
.field(
"client_visible_upstream_error",
&self
.client_visible_upstream_error
.as_ref()
.map(|_| "<present>"),
)
.field("transport_metrics", &self.transport_metrics)
.field(
"credential_recovery_required",
+104 -18
View File
@@ -30,11 +30,12 @@ use gateway_core::event::GatewayEvent;
use gateway_core::operation::{Operation, OperationKind};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown,
ProviderCooldownPort, ProviderCredentialState, ProviderCredentialStatePort,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCooldownScope, ProviderCredentialState, ProviderCredentialStatePort,
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts,
ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderStoreError, ProviderStorePorts,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
@@ -372,7 +373,7 @@ async fn xai_admin_provider_restores_full_pending_envelope_and_binds_owner() {
.await
.expect_err("restored pending must validate callback state");
assert_eq!(invalid_callback.kind(), ProviderAdminErrorKind::Invalid);
assert_eq!(pending.len(), 0);
assert_eq!(pending.len(), 1);
let mismatched_flow = "pending-envelope-mismatch";
let owner_ref = admin_session_owner_ref("admin-owner");
@@ -399,7 +400,7 @@ async fn xai_admin_provider_restores_full_pending_envelope_and_binds_owner() {
.await
.expect_err("flow binding mismatch");
assert_eq!(mismatched_envelope.kind(), ProviderAdminErrorKind::Invalid);
assert_eq!(pending.len(), 0);
assert_eq!(pending.len(), 2);
}
#[tokio::test]
@@ -508,7 +509,6 @@ async fn xai_admin_provider_rejects_unprepared_mutations_before_store_commit() {
let rotation_error = admin
.prepare_rotation(PrepareCredentialRotation {
account: record.clone(),
expected_credential_revision: record.credential_revision,
provider_material: ProviderDocument::new(OpaqueProviderData::new(Map::new())),
})
.await
@@ -542,7 +542,7 @@ fn pending_payload(flow_id: &str, owner_ref: &str) -> OpaqueProviderData {
.expose()
.to_owned();
let value = json!({
"schema_version": 1,
"schema_version": 3,
"flow_id": flow_id,
"owner_ref": owner_ref,
"expires_at": (Utc::now() + chrono::Duration::minutes(30)).to_rfc3339(),
@@ -796,6 +796,30 @@ impl ProviderCooldownPort for TestCooldown {
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
fn put_scoped_if_later(
&self,
_cooldown: ProviderScopedCooldown,
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
fn read_scoped<'a>(
&'a self,
_account_id: &'a ProviderAccountId,
_scope: &'a ProviderCooldownScope,
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>> {
Box::pin(async { Ok(None) })
}
fn clear_scoped<'a>(
&'a self,
_account_id: &'a ProviderAccountId,
_scope: &'a ProviderCooldownScope,
_through_revision: CredentialRevision,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
}
struct TestRuntimePolicy;
@@ -819,7 +843,7 @@ struct TestOAuthPending {
}
type PendingKey = (String, String);
type PendingValue = (String, OpaqueProviderData, SystemTime);
type PendingValue = (String, OpaqueProviderData, SystemTime, Option<String>);
impl TestOAuthPending {
fn insert(&self, flow_id: &str, owner: String, payload: OpaqueProviderData) {
@@ -829,6 +853,7 @@ impl TestOAuthPending {
owner,
payload,
SystemTime::now() + Duration::from_secs(1_800),
None,
),
);
}
@@ -858,36 +883,97 @@ impl OAuthPendingFlowPort for TestOAuthPending {
flow.owner().expose_to_store().to_owned(),
flow.payload().clone(),
SystemTime::now() + flow.ttl(),
None,
),
);
Ok(OAuthPendingPutOutcome::Stored)
})
}
fn take_if_owner<'a>(
fn claim_if_owner<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingTakeOutcome, ProviderStoreError>> {
claim: &'a OAuthPendingBinding,
_claim_ttl: Duration,
) -> BoxFuture<'a, Result<OAuthPendingClaimOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, expires_at)) = values.get(&key) else {
return Ok(OAuthPendingTakeOutcome::NotFound);
let Some((stored_owner, payload, expires_at, stored_claim)) = values.get_mut(&key)
else {
return Ok(OAuthPendingClaimOutcome::NotFound);
};
if expires_at <= &SystemTime::now() {
if *expires_at <= SystemTime::now() {
values.remove(&key);
return Ok(OAuthPendingTakeOutcome::NotFound);
return Ok(OAuthPendingClaimOutcome::NotFound);
}
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingTakeOutcome::OwnerMismatch);
return Ok(OAuthPendingClaimOutcome::OwnerMismatch);
}
let (_, payload, _) = values.remove(&key).expect("checked pending");
Ok(OAuthPendingTakeOutcome::Taken(payload))
if stored_claim.is_some() {
return Ok(OAuthPendingClaimOutcome::InProgress);
}
*stored_claim = Some(claim.expose_to_store().to_owned());
Ok(OAuthPendingClaimOutcome::Claimed(payload.clone()))
})
}
fn release_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingReleaseOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, _, stored_claim)) = values.get_mut(&key) else {
return Ok(OAuthPendingReleaseOutcome::NotFound);
};
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingReleaseOutcome::OwnerMismatch);
}
if stored_claim.as_deref() != Some(claim.expose_to_store()) {
return Ok(OAuthPendingReleaseOutcome::ClaimMismatch);
}
*stored_claim = None;
Ok(OAuthPendingReleaseOutcome::Released)
})
}
fn consume_claim<'a>(
&'a self,
provider_kind: &'a ProviderKind,
flow: &'a OAuthPendingBinding,
owner: &'a OAuthPendingBinding,
claim: &'a OAuthPendingBinding,
) -> BoxFuture<'a, Result<OAuthPendingConsumeOutcome, ProviderStoreError>> {
Box::pin(async move {
let key = (
provider_kind.as_str().to_owned(),
flow.expose_to_store().to_owned(),
);
let mut values = self.values.lock().expect("OAuth pending");
let Some((stored_owner, _, _, stored_claim)) = values.get(&key) else {
return Ok(OAuthPendingConsumeOutcome::NotFound);
};
if stored_owner != owner.expose_to_store() {
return Ok(OAuthPendingConsumeOutcome::OwnerMismatch);
}
if stored_claim.as_deref() != Some(claim.expose_to_store()) {
return Ok(OAuthPendingConsumeOutcome::ClaimMismatch);
}
values.remove(&key);
Ok(OAuthPendingConsumeOutcome::Consumed)
})
}
}
@@ -381,9 +381,13 @@ async fn manual_refresh_returns_prepared_rotation_without_writing_store() {
let id = input.account_id.clone();
let (store, _, refresher, service) =
fixture(input, [Ok(success_tokens(Some("manual-refresh")))], true).await;
let current = store
.load_current_credential(&id)
.await
.expect("load current credential");
let prepared = service
.prepare_manual_refresh(&id, CredentialRevision::new(1).expect("revision"))
.prepare_manual_refresh(current)
.await
.expect("prepare manual refresh");
assert_eq!(store.account(&id).expect("account").revision().get(), 1);
@@ -401,20 +405,35 @@ async fn manual_refresh_returns_prepared_rotation_without_writing_store() {
}
#[tokio::test]
async fn manual_refresh_rejects_stale_revision_before_upstream_call() {
async fn manual_refresh_preserves_snapshot_revision_as_the_final_cas_fence() {
let input = due_input("manual-stale");
let id = input.account_id.clone();
let (_, _, refresher, service) = fixture(input, [], true).await;
let (store, _, refresher, service) = fixture(
input,
[Ok(success_tokens(Some("manual-stale-refresh")))],
true,
)
.await;
let current = store
.load_current_credential(&id)
.await
.expect("load current credential");
force_due(&store, &id).await;
let prepared = service
.prepare_manual_refresh(current)
.await
.expect("prepare from current snapshot");
let (_, credential, _guard) = prepared.into_parts();
assert!(matches!(
service
.prepare_manual_refresh(&id, CredentialRevision::new(2).expect("revision"))
.await,
Err(GrokCredentialRefreshError::Repository(
provider_xai::GrokCredentialRepositoryError::StaleCredentialRevision
))
store
.compare_and_swap_credential(credential)
.await
.expect("attempt stale prepared rotation"),
CredentialCasOutcome::Conflict
));
assert_eq!(refresher.prepare_calls.load(Ordering::SeqCst), 0);
assert_eq!(refresher.prepare_calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
@@ -423,11 +442,13 @@ async fn manual_refresh_preserves_failure_class_without_store_write() {
let id = input.account_id.clone();
let (store, _, _, service) =
fixture(input, [Err(GrokRefreshFailure::InvalidGrant)], true).await;
let current = store
.load_current_credential(&id)
.await
.expect("load current credential");
assert!(matches!(
service
.prepare_manual_refresh(&id, CredentialRevision::new(1).expect("revision"))
.await,
service.prepare_manual_refresh(current).await,
Err(GrokCredentialRefreshError::ManualFailure(
GrokRefreshFailure::InvalidGrant
))
@@ -10,9 +10,10 @@ use gateway_core::engine::credential::{
QuotaWriteOutcome, RotationStrategy,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingState,
ProviderStoreError,
ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderSchedulingState, ProviderStoreError,
};
use gateway_core::routing::UpstreamModelId;
use provider_xai::{
GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport,
GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture,
@@ -119,7 +120,8 @@ impl SelectorFixture {
.replace(GrokPlanCatalog::new(
GrokCatalogScope::for_account(&account).expect("catalog scope"),
Utc::now(),
GrokCredentialCatalogSeed::new(["grok-4.5"], None).expect("catalog"),
GrokCredentialCatalogSeed::new(["grok-4.5", "grok-4.6"], None)
.expect("catalog"),
))
.await
.expect("cache catalog");
@@ -183,9 +185,32 @@ impl SelectorFixture {
excluded: BTreeSet<ProviderAccountId>,
required_account: Option<ProviderAccountId>,
strategy: RotationStrategy,
) -> GrokSessionSelection {
self.request_for_model_with_policy("grok-4.5", excluded, required_account, strategy)
}
fn request_for_model(
&self,
upstream_model: &str,
required_account: Option<ProviderAccountId>,
) -> GrokSessionSelection {
self.request_for_model_with_policy(
upstream_model,
BTreeSet::new(),
required_account,
RotationStrategy::Smart,
)
}
fn request_for_model_with_policy(
&self,
upstream_model: &str,
excluded: BTreeSet<ProviderAccountId>,
required_account: Option<ProviderAccountId>,
strategy: RotationStrategy,
) -> GrokSessionSelection {
GrokSessionSelection::new(
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
UpstreamModelId::new(upstream_model).expect("model"),
excluded,
required_account,
AccountSelectionPolicy::new(
@@ -387,6 +412,131 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account
assert_ne!(next.account_id(), &selected);
}
#[tokio::test]
async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown() {
let fixture = SelectorFixture::new(&["payment-required", "available"]).await;
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
let selected = session.account_id().clone();
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::PaymentRequired {
retry_after: Some(Duration::from_secs(5)),
},
)
.await;
assert_eq!(
fixture
.store
.account(&selected)
.expect("selected account")
.availability(),
AccountAvailability::Ready
);
assert!(fixture.cooldowns.cooldown(&selected).is_some());
let next = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("another account remains available");
assert_ne!(next.account_id(), &selected);
}
#[tokio::test]
async fn model_quota_feedback_blocks_only_the_failed_account_and_model() {
let fixture = SelectorFixture::new(&["model-quota", "available"]).await;
let failed_model = UpstreamModelId::new("grok-4.5").expect("failed model");
let session = fixture
.selector
.select(fixture.request_for_model(failed_model.as_str(), None))
.await
.expect("session");
let selected = session.account_id().clone();
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: failed_model.clone(),
retry_after: Some(Duration::from_secs(5)),
},
)
.await;
assert_eq!(
fixture
.store
.account(&selected)
.expect("selected account")
.availability(),
AccountAvailability::Ready
);
let scope = ProviderCooldownScope::upstream_model(failed_model);
assert!(
fixture
.cooldowns
.scoped_cooldown(&selected, &scope)
.is_some()
);
let unrelated_model = fixture
.selector
.select(fixture.request_for_model("grok-4.6", Some(selected.clone())))
.await
.expect("same account remains valid for another model");
assert_eq!(unrelated_model.account_id(), &selected);
let failed_model_retry = fixture
.selector
.select(fixture.request_for_model("grok-4.5", None))
.await
.expect("another account serves the failed model");
assert_ne!(failed_model_retry.account_id(), &selected);
}
#[tokio::test]
async fn model_access_feedback_reports_model_cooldown_without_blocking_the_account() {
let fixture = SelectorFixture::new(&["model-access"]).await;
let model = UpstreamModelId::new("grok-4.5").expect("model");
let session = fixture
.selector
.select(fixture.request_for_model(model.as_str(), None))
.await
.expect("session");
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::ModelAccessDenied {
upstream_model: model,
retry_after: None,
},
)
.await;
assert_eq!(
fixture
.store
.account(session.account_id())
.expect("account")
.availability(),
AccountAvailability::Ready
);
assert!(matches!(
fixture
.selector
.select(fixture.request_for_model("grok-4.5", None))
.await,
Err(GrokSessionSelectorError::ModelCoolingDown {
retry_after: Some(_)
})
));
}
#[tokio::test]
async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting_account_state() {
let fixture = SelectorFixture::new(&["stream-interrupted"]).await;
@@ -18,7 +18,8 @@ use gateway_core::engine::{
UpstreamSendState,
};
use gateway_core::error::{
ContinuationFailure, ProviderError, ProviderErrorKind, SafeUpstreamValue,
ClientVisibleUpstreamError, ContinuationFailure, ProviderError, ProviderErrorKind,
SafeUpstreamValue,
};
use gateway_core::event::{GatewayEvent, UpstreamHttpVersion};
use gateway_core::operation::{
@@ -1295,6 +1296,157 @@ async fn explicit_quota_http_429_marks_provider_error_replay_safe() {
assert!(error.replay_is_safe());
}
#[tokio::test]
async fn unknown_http_402_is_retryable_without_account_quota_feedback() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::PaymentRequired,
UpstreamSendState::Sent,
)
.with_status(402),
);
let provider = provider(selector.clone(), transport).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::PermissionDenied);
assert!(error.replay_is_safe());
assert_eq!(
selector.feedback.lock().expect("feedback").as_slice(),
&[GrokCredentialFailure::PaymentRequired { retry_after: None }]
);
}
#[tokio::test]
async fn model_quota_transport_failure_preserves_model_scoped_feedback() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::ModelQuotaExhausted,
UpstreamSendState::Sent,
)
.with_status(403),
);
let provider = provider(selector.clone(), transport).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted);
assert!(error.replay_is_safe());
assert_eq!(
selector.feedback.lock().expect("feedback").as_slice(),
&[GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: UpstreamModelId::new("grok-4.5").expect("model"),
retry_after: None,
}]
);
}
#[tokio::test]
async fn model_access_denial_is_retryable_without_credential_recovery() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::ModelAccessDenied,
UpstreamSendState::Sent,
)
.with_status(403),
);
let provider = provider(selector.clone(), transport).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::PermissionDenied);
assert!(error.replay_is_safe());
assert_eq!(
selector.feedback.lock().expect("feedback").as_slice(),
&[GrokCredentialFailure::ModelAccessDenied {
upstream_model: UpstreamModelId::new("grok-4.5").expect("model"),
retry_after: None,
}]
);
}
#[tokio::test]
async fn safety_rejection_does_not_rotate_or_mutate_account_state() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::SafetyRejected,
UpstreamSendState::Sent,
)
.with_status(403),
);
let provider = provider(selector.clone(), transport).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::PermissionDenied);
assert!(!error.replay_is_safe());
assert!(selector.feedback.lock().expect("feedback").is_empty());
}
#[tokio::test]
async fn transport_error_should_preserve_client_visible_upstream_detail() {
let detail = ClientVisibleUpstreamError::new(
"You have run out of credits",
Some("personal_team_blocked_spending_limit".to_owned()),
Some("insufficient_quota".to_owned()),
)
.expect("safe upstream detail");
let error = mapped_transport_error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::QuotaExhausted,
UpstreamSendState::Sent,
)
.with_status(402)
.with_client_visible_upstream_error(detail),
false,
)
.await;
let detail = error
.client_visible_upstream_error()
.expect("mapped client-visible detail");
assert_eq!(
(detail.message(), detail.code(), detail.error_type()),
(
"You have run out of credits",
Some("personal_team_blocked_spending_limit"),
Some("insufficient_quota"),
)
);
}
#[tokio::test]
async fn explicit_http_408_does_not_mark_provider_error_replay_safe() {
let error = mapped_transport_error(
+74 -6
View File
@@ -15,9 +15,10 @@ use gateway_core::engine::credential::{
};
use gateway_core::error::{StoreError, StoreErrorKind};
use gateway_core::provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionExclusionPort,
ProviderSessionExclusions, ProviderStoreError,
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderSessionAffinityKey,
ProviderSessionAffinityPort, ProviderSessionExclusionPort, ProviderSessionExclusions,
ProviderStoreError,
};
use gateway_core::routing::ProviderKind;
use provider_xai::{
@@ -283,11 +284,19 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
account: &ProviderAccountId,
expected_revision: CredentialRevision,
) -> Result<LoadedCredential, StoreError> {
let accounts = lock(&self.accounts);
let stored = accounts.get(account).ok_or_else(invalid)?;
if stored.account.revision() != expected_revision {
let loaded = self.load_current_credential(account).await?;
if loaded.account.revision() != expected_revision {
return Err(conflict());
}
Ok(loaded)
}
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, StoreError> {
let accounts = lock(&self.accounts);
let stored = accounts.get(account).ok_or_else(invalid)?;
Ok(LoadedCredential {
account: stored.account.clone(),
credential: stored.credential.clone(),
@@ -553,12 +562,24 @@ pub struct TestSessionExclusions;
#[derive(Default)]
pub struct MemoryCooldownPort {
cooldowns: Mutex<BTreeMap<ProviderAccountId, ProviderCooldown>>,
scoped_cooldowns:
Mutex<BTreeMap<(ProviderAccountId, ProviderCooldownScope), ProviderScopedCooldown>>,
}
impl MemoryCooldownPort {
pub fn cooldown(&self, account_id: &ProviderAccountId) -> Option<ProviderCooldown> {
lock(&self.cooldowns).get(account_id).cloned()
}
pub fn scoped_cooldown(
&self,
account_id: &ProviderAccountId,
scope: &ProviderCooldownScope,
) -> Option<ProviderScopedCooldown> {
lock(&self.scoped_cooldowns)
.get(&(account_id.clone(), scope.clone()))
.cloned()
}
}
impl ProviderSessionAffinityPort for TestSessionAffinity {
@@ -666,6 +687,53 @@ impl ProviderCooldownPort for MemoryCooldownPort {
Ok(should_remove)
})
}
fn put_scoped_if_later(
&self,
cooldown: ProviderScopedCooldown,
) -> futures::future::BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let key = (cooldown.account_id().clone(), cooldown.scope().clone());
let mut cooldowns = lock(&self.scoped_cooldowns);
let should_write = cooldowns.get(&key).is_none_or(|current| {
current.credential_revision() < cooldown.credential_revision()
|| (current.credential_revision() == cooldown.credential_revision()
&& current.until() < cooldown.until())
});
if should_write {
cooldowns.insert(key, cooldown);
}
Ok(should_write)
})
}
fn read_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
) -> futures::future::BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>>
{
Box::pin(async move { Ok(self.scoped_cooldown(account_id, scope)) })
}
fn clear_scoped<'a>(
&'a self,
account_id: &'a ProviderAccountId,
scope: &'a ProviderCooldownScope,
through_revision: CredentialRevision,
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let key = (account_id.clone(), scope.clone());
let mut cooldowns = lock(&self.scoped_cooldowns);
let should_remove = cooldowns
.get(&key)
.is_some_and(|current| current.credential_revision() <= through_revision);
if should_remove {
cooldowns.remove(&key);
}
Ok(should_remove)
})
}
}
impl ProviderRuntimePolicyPort for StaticRuntimePolicy {
@@ -257,7 +257,7 @@ async fn inference_transport_should_classify_http_failures_without_retaining_bod
let cases = [
(400, GrokInferenceTransportErrorKind::InvalidRequest),
(401, GrokInferenceTransportErrorKind::Unauthorized),
(402, GrokInferenceTransportErrorKind::QuotaExhausted),
(402, GrokInferenceTransportErrorKind::PaymentRequired),
(403, GrokInferenceTransportErrorKind::PermissionDenied),
(408, GrokInferenceTransportErrorKind::Timeout),
(429, GrokInferenceTransportErrorKind::RateLimited),
@@ -302,10 +302,82 @@ async fn inference_transport_should_classify_http_failures_without_retaining_bod
Some(GrokInferenceClientCacheStatus::Miss)
);
assert!(error.transport_metrics().headers_ms().is_some());
assert!(error.client_visible_upstream_error().is_none());
assert!(!rendered.contains(&secret));
}
}
#[tokio::test]
async fn inference_transport_should_expose_safe_flat_json_error_details() {
let server = MockServer::start().await;
let message = "You have run out of credits or need a Grok subscription";
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(402).set_body_json(json!({
"code": "personal-team-blocked:spending-limit",
"error": message,
})))
.expect(1)
.mount(&server)
.await;
let origin = Url::parse(&server.uri()).expect("wiremock origin");
let error = inference_transport(&origin)
.execute(inference_request(&origin))
.await
.expect_err("spending limit response must fail");
let detail = error
.client_visible_upstream_error()
.expect("structured error detail");
let rendered = format!("{error:?}\n{error}");
assert_eq!(
error.kind(),
GrokInferenceTransportErrorKind::QuotaExhausted
);
assert_eq!(
(detail.message(), detail.code(), detail.error_type()),
(message, Some("personal_team_blocked_spending_limit"), None,)
);
assert!(!rendered.contains(message));
}
#[tokio::test]
async fn inference_transport_should_scrub_nested_json_error_details() {
let server = MockServer::start().await;
let account_fingerprint = "123e4567-e89b-12d3-a456-426614174000";
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(429).set_body_json(json!({
"error": {
"code": "rate-limit:burst",
"type": "rate_limit_error",
"message": format!("team {account_fingerprint}\nrate limited"),
}
})))
.expect(1)
.mount(&server)
.await;
let origin = Url::parse(&server.uri()).expect("wiremock origin");
let error = inference_transport(&origin)
.execute(inference_request(&origin))
.await
.expect_err("rate limit response must fail");
let detail = error
.client_visible_upstream_error()
.expect("structured error detail");
let rendered = format!("{error:?}\n{error}");
assert_eq!(
(detail.message(), detail.code(), detail.error_type()),
(
"team [redacted] rate limited",
Some("rate_limit_burst"),
Some("rate_limit_error"),
)
);
assert!(!rendered.contains(account_fingerprint));
}
#[tokio::test]
async fn inference_transport_should_scope_forbidden_failures_from_safe_metadata() {
let cases = [
@@ -317,14 +389,20 @@ async fn inference_transport_should_scope_forbidden_failures_from_safe_metadata(
),
(
json!({"error": {"code": "usage_exhausted", "message": "used all the included free usage for model"}}),
GrokInferenceTransportErrorKind::QuotaExhausted,
GrokInferenceTransportErrorKind::ModelQuotaExhausted,
false,
Some("usage_exhausted"),
),
(
json!({"error": {"code": "permission_denied", "message": "access to the chat endpoint is denied"}}),
GrokInferenceTransportErrorKind::PermissionDenied,
true,
GrokInferenceTransportErrorKind::ModelAccessDenied,
false,
Some("permission_denied"),
),
(
json!({"error": {"code": "permission_denied", "message": "request rejected", "details": {"check": "SAFETY_CHECK_TYPE_INPUT"}}}),
GrokInferenceTransportErrorKind::SafetyRejected,
false,
Some("permission_denied"),
),
(
@@ -359,7 +437,62 @@ async fn inference_transport_should_scope_forbidden_failures_from_safe_metadata(
}
#[tokio::test]
async fn inference_transport_should_classify_free_usage_429_as_quota_exhausted() {
async fn inference_transport_should_keep_unknown_402_out_of_account_quota_state() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(402).set_body_json(json!({
"error": {
"code": "payment_required",
"type": "billing_error",
"message": "Payment is required for this request"
}
})))
.expect(1)
.mount(&server)
.await;
let origin = Url::parse(&server.uri()).expect("wiremock origin");
let error = inference_transport(&origin)
.execute(inference_request(&origin))
.await
.expect_err("payment-required response must fail");
assert_eq!(
(error.kind(), error.requires_credential_recovery()),
(GrokInferenceTransportErrorKind::PaymentRequired, false)
);
}
#[tokio::test]
async fn inference_transport_should_classify_model_quota_for_403_and_429() {
for status in [403, 429] {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(status).set_body_json(json!({
"error": {
"code": "subscription:free-usage-exhausted",
"message": "You have used all the included free usage for model grok-4.5"
}
})))
.expect(1)
.mount(&server)
.await;
let origin = Url::parse(&server.uri()).expect("wiremock origin");
let error = inference_transport(&origin)
.execute(inference_request(&origin))
.await
.expect_err("model quota response must fail");
assert_eq!(
error.kind(),
GrokInferenceTransportErrorKind::ModelQuotaExhausted
);
}
}
#[tokio::test]
async fn inference_transport_should_classify_free_usage_429_as_account_free_quota_exhausted() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
@@ -380,7 +513,7 @@ async fn inference_transport_should_classify_free_usage_429_as_quota_exhausted()
assert_eq!(
error.kind(),
GrokInferenceTransportErrorKind::QuotaExhausted
GrokInferenceTransportErrorKind::FreeQuotaExhausted
);
assert_eq!(error.status(), Some(429));
assert_eq!(
+2 -6
View File
@@ -47,8 +47,6 @@ Client Key 固定绑定一个 Provider:`openai` 或 `xai`。同一次请求不
Admin API 不暴露全局配置版本,mutation 请求也不要求客户端提供全局配置版本。会改变路由快照或安全配置的写入由后端在事务内推进内部配置 revision,并用于快照发布与审计。
单个 credential 的轮换或重新授权仍使用 `expectedCredentialRevision`,它只保护该 credential 本身,不是全局配置锁。
## 2. 健康检查
| 方法 | 路由 | 鉴权 | 说明 |
@@ -150,7 +148,6 @@ OpenAI rotation 请求字段为:
{
"provider": "openai",
"accountId": "acct_...",
"expectedCredentialRevision": 1,
"accessToken": "...",
"refreshToken": "..."
}
@@ -162,12 +159,11 @@ OAuth start 使用:
{
"provider": "openai",
"name": "account name",
"accountId": null,
"expectedCredentialRevision": null
"accountId": null
}
```
重新授权已有账号时,`accountId` 与 `expectedCredentialRevision` 必须同时提供。
重新授权已有账号时只需提供 `accountId`。
## 6. Client Key
-3
View File
@@ -105,8 +105,6 @@ export interface Account {
tokenRefreshing: boolean
availability: string
enabled: boolean
credentialRevision: number
stateRevision: number | null
accessTokenExpiresAt: string | null
accessTokenExpiresAtDisplay: string | null
refreshTokenExpiresAt: string | null
@@ -160,7 +158,6 @@ export interface AccountImportResponse {
export interface AccountMutationResponse {
accountId: string
credentialRevision?: number
}
export interface AccountDeletionResponse {
@@ -72,7 +72,6 @@ export function useAccountOnboarding(options: {
...(account
? {
accountId: account.id,
expectedCredentialRevision: account.credentialRevision,
}
: {}),
})
+1 -1
View File
@@ -1 +1 @@
version: "3.0.16"
version: "3.0.20"