Compare commits

...
12 Commits
Author SHA1 Message Date
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
zyycnandCodex 5168f6a924 fix(openai): replace stale passive rate-limit snapshots
Treat each passive response-header observation as a complete rate-limit bucket snapshot, preventing stale secondary windows from surviving into a newer primary-only update.\n\nExtract duplicate OAuth callback parameter rejection to keep the release quality gate warning-free.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 23:12:17 +08:00
zyycnandCodex 3a407c635a chore: bump version to 3.0.16
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 22:39:31 +08:00
zyycnandCodex 6901e0ed45 fix(openai): accept standard OAuth callback parameters
Accept standard OAuth callback parameters while keeping one-shot flow and duplicate security parameter rejection.

Classify unavailable identity verification correctly and add redacted OAuth/import failure diagnostics.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 22:39:01 +08:00
zyycnandClaude Fable 5 9758c420fa refactor(availability): converge account availability on the gateway-core enum
- gateway-admin: drop the mirrored 7-variant enum; re-export the core type
- gateway-api: serialize wire availability via core as_str (bytes unchanged)
- gateway-store: drop ProviderAccountAvailability and four identity mappers;
  rows parse through core's codec with the same error string; redis state
  validation reuses core parse instead of a literal 7-string set
- providers/openai+xai: drop core<->admin identity mappers made moot by the
  admin re-export
- providers/xai: a single exhaustive From<GrokCredentialAvailability> impl
  is now the only Grok->core translation point

No behavior change: DB strings, wire JSON, and the AccountStatus fold
(including Invalid->Expired display) are byte-identical. 1244 tests green.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-27 20:40:47 +08:00
ZyyandCodex 6d2bf06ac7 chore: bump version to 3.0.15
Co-authored-by: Codex <noreply@openai.com>
2026-07-27 17:29:36 +08:00
ZyyandCodex b8faaf250f fix(frontend): align dashboard account usage metrics
Add a metric layout variant for account usage windows so dashboard account rows align request bars with adjacent recent-usage text without changing the shared compact account-table presentation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 17:22:28 +08:00
ZyyandCodex 40137b1205 fix(backend): preserve safe upstream error details
Keep admin account connection tests routed through the provider probe even when the stored account is already marked unavailable so clicking test can refresh state from the upstream response.

Preserve provider-declared safe upstream details when OpenAI-compatible API errors collapse local provider unavailability to the stable no-provider contract.

Co-authored-by: Codex <noreply@openai.com>
2026-07-27 17:09:59 +08:00
53 changed files with 3384 additions and 1321 deletions
+1
View File
@@ -1194,6 +1194,7 @@ dependencies = [
"subtle",
"thiserror",
"tokio",
"tracing",
"uuid",
]
+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]
@@ -9,17 +9,8 @@ use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision, observability::TimeRange};
/// Provider 账号当前可用性。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountAvailability {
Unknown,
Ready,
Cooldown,
QuotaExhausted,
Expired,
Banned,
Invalid,
}
/// Provider 账号当前可用性;权威定义与字符串编解码在 gateway-core,此处仅复用。
pub use gateway_core::engine::credential::AccountAvailability;
/// 管理页使用的归一化账号状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -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},
};
@@ -310,6 +311,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 +335,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 +362,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 +405,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(())
}
}
@@ -151,18 +151,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) {
(
@@ -184,11 +193,20 @@ fn validate_authorization_commit(
_ => 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 +215,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)
}
@@ -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,7 +58,8 @@ 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>>,
@@ -72,7 +73,8 @@ 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()),
@@ -113,6 +115,13 @@ 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;
}
fn export_inputs(&self) -> Vec<ProviderExportCredentialInput> {
self.export_inputs.lock().expect("export inputs").clone()
}
@@ -246,6 +255,11 @@ 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))
@@ -260,9 +274,16 @@ impl ProviderAdmin for FakeProviderAdmin {
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,
})
}
@@ -347,9 +368,13 @@ pub(super) struct FakeAccountStore {
impl FakeAccountStore {
pub(super) fn new(kind: &str, events: EventLog) -> Arc<Self> {
Self::with_account(account_record(kind), events)
}
fn with_account(account: AccountRecord, events: EventLog) -> Arc<Self> {
Arc::new(Self {
events,
account: account_record(kind),
account,
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
quota_window_usage: Mutex::new(Vec::new()),
@@ -660,6 +685,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());
@@ -686,9 +756,11 @@ fn provider_registry_should_reject_duplicate_kind() {
}
#[tokio::test]
async fn connection_test_should_preserve_safe_provider_error_details() {
async fn connection_test_should_probe_unavailable_account() {
let provider = FakeProviderAdmin::new("xai", events());
let store = FakeAccountStore::new("xai", events());
let mut account = account_record("xai");
account.availability = AccountAvailability::QuotaExhausted;
let store = FakeAccountStore::with_account(account, events());
let services =
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
@@ -709,7 +781,7 @@ async fn connection_test_should_preserve_safe_provider_error_details() {
message,
provider_error_code: Some(code),
provider_error_type: Some(error_type),
..
account_status: gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
}) if message == "included usage exhausted"
&& code == "usage_exhausted"
&& error_type == "invalid_request_error"
@@ -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();
@@ -18,9 +18,9 @@ use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
AdminError as AdminServiceError, AdminErrorKind, PageSize, Revision,
accounts::{
AccountAvailability, AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost,
AccountListQuery, AccountModelUsage, AccountSort, AccountSortField,
AccountStatus as DomainAccountStatus, AccountUsage, SortDirection,
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
AccountUsage, SortDirection,
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
@@ -1407,7 +1407,7 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
status: status.clone(),
display_status: status,
token_refreshing: false,
availability: account_availability_name(account.availability).to_owned(),
availability: account.availability.as_str().to_owned(),
enabled: account.enabled,
credential_revision: account.credential_revision.get(),
state_revision: None,
@@ -1752,18 +1752,6 @@ fn account_status_name(status: DomainAccountStatus) -> &'static str {
}
}
fn account_availability_name(availability: AccountAvailability) -> &'static str {
match availability {
AccountAvailability::Unknown => "unknown",
AccountAvailability::Ready => "ready",
AccountAvailability::Cooldown => "cooldown",
AccountAvailability::QuotaExhausted => "quota_exhausted",
AccountAvailability::Expired => "expired",
AccountAvailability::Banned => "banned",
AccountAvailability::Invalid => "invalid",
}
}
fn format_number(value: u64) -> String {
let text = value.to_string();
let mut output = String::with_capacity(text.len() + text.len() / 3);
@@ -75,10 +75,14 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
if error.kind() == ProviderErrorKind::Unavailable
&& error.send_state() == UpstreamSendState::NotSent =>
{
GatewayError::new(
let gateway = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream provider is currently available for this request",
)
);
match error.client_visible_upstream_error() {
Some(upstream) => gateway.with_client_visible_upstream_error(upstream.clone()),
None => gateway,
}
}
EngineError::Provider(error) => GatewayError::from_provider(error),
EngineError::Cancelled => {
@@ -135,6 +135,36 @@ fn local_provider_capacity_exhaustion_should_map_to_service_unavailable() {
);
}
#[test]
fn local_provider_capacity_exhaustion_should_preserve_safe_provider_detail() {
let detail = ClientVisibleUpstreamError::new(
"no account is eligible for the requested model",
Some("no_eligible_account".to_owned()),
Some("account_unavailable_error".to_owned()),
)
.expect("safe provider detail");
let error = EngineError::Provider(
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
.with_client_visible_upstream_error(detail),
);
let gateway = gateway_error_from_engine(&error);
assert_eq!(gateway.kind(), GatewayErrorKind::NoAvailableProvider);
assert_eq!(
gateway.client_message(),
"no account is eligible for the requested model"
);
assert_eq!(gateway.client_error_code(), Some("no_eligible_account"));
assert_eq!(
gateway.client_error_type(),
Some("account_unavailable_error")
);
assert_eq!(
gateway.safe_message(),
"no upstream provider is currently available for this request"
);
}
#[test]
fn no_eligible_provider_account_should_map_to_service_unavailable() {
let error = EngineError::Provider(ProviderError::new(
@@ -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,
@@ -158,7 +216,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>,
@@ -483,27 +541,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(),
)
@@ -818,7 +885,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
@@ -896,7 +963,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
@@ -901,6 +901,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 +1121,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(),
)
@@ -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()
@@ -778,35 +778,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 集合。
@@ -1202,7 +1202,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 +1210,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 +1225,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]
@@ -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 =
@@ -12,8 +12,7 @@ use gateway_admin::{
model::{
MutationContext, Revision as AdminRevision,
accounts::{
AccountAvailability as AdminAccountAvailability, AccountCost,
AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountCost, AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
AccountSummary, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
@@ -35,8 +34,8 @@ use gateway_admin::{
use sqlx::{PgPool, Postgres, Row, Transaction};
use gateway_core::engine::credential::{
AccountAvailability as CoreAccountAvailability, AccountStateChange, CredentialCasOutcome,
CredentialCasUpdate, CredentialRevision as CoreCredentialRevision, LoadedCredential,
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision as CoreCredentialRevision, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountStore, ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation,
@@ -121,43 +120,9 @@ impl ProviderAccountAdminScope {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderAccountAvailability {
Unknown,
Ready,
Cooldown,
QuotaExhausted,
Expired,
Banned,
Invalid,
}
impl ProviderAccountAvailability {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unknown => "unknown",
Self::Ready => "ready",
Self::Cooldown => "cooldown",
Self::QuotaExhausted => "quota_exhausted",
Self::Expired => "expired",
Self::Banned => "banned",
Self::Invalid => "invalid",
}
}
fn parse(value: &str) -> StoreResult<Self> {
match value {
"unknown" => Ok(Self::Unknown),
"ready" => Ok(Self::Ready),
"cooldown" => Ok(Self::Cooldown),
"quota_exhausted" => Ok(Self::QuotaExhausted),
"expired" => Ok(Self::Expired),
"banned" => Ok(Self::Banned),
"invalid" => Ok(Self::Invalid),
_ => Err(invalid("unknown availability value")),
}
}
/// 解析数据库 availability 列;字符串映射本体唯一归属 gateway-core。
fn parse_availability(value: &str) -> StoreResult<AccountAvailability> {
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -175,7 +140,7 @@ pub struct ProviderAccountSummary {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
@@ -220,7 +185,7 @@ pub struct NewProviderAccount {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
@@ -259,9 +224,7 @@ impl NewProviderAccount {
if !self.has_refresh_token && self.next_refresh_at.is_some() {
return Err(invalid("next_refresh_at requires a refresh token"));
}
if (self.availability == ProviderAccountAvailability::Cooldown)
!= self.cooldown_until.is_some()
{
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown_until must be present exactly for cooldown availability",
));
@@ -394,7 +357,7 @@ impl fmt::Debug for ProviderCredentialUpdate {
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderAccountObservation {
pub account_id: String,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub provider_quota_json: Option<JsonObject>,
@@ -406,7 +369,7 @@ pub struct ProviderAccountObservation {
pub struct ProviderAccountStateUpdate {
pub account_id: String,
pub expected_revision: Revision,
pub availability: ProviderAccountAvailability,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
@@ -415,9 +378,7 @@ pub struct ProviderAccountStateUpdate {
impl ProviderAccountStateUpdate {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == ProviderAccountAvailability::Cooldown)
!= self.cooldown_until.is_some()
{
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
@@ -429,9 +390,7 @@ impl ProviderAccountStateUpdate {
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == ProviderAccountAvailability::Cooldown)
!= self.cooldown_until.is_some()
{
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
@@ -1225,9 +1184,10 @@ impl AccountStore for PgAdminAccountStore {
.into_iter()
.filter(|account| account.provider_kind == provider_kind.as_str())
.filter(|account| {
query.availability.as_ref().is_none_or(|expected| {
expected.matches(admin_account_availability(account.availability))
})
query
.availability
.as_ref()
.is_none_or(|expected| expected.matches(account.availability))
})
.filter(|account| {
query
@@ -1602,12 +1562,12 @@ fn admin_account_status(
) -> AdminAccountStatus {
if !account.enabled {
AdminAccountStatus::Disabled
} else if account.availability == ProviderAccountAvailability::Banned {
} else if account.availability == AccountAvailability::Banned {
AdminAccountStatus::Banned
} else if account.availability == ProviderAccountAvailability::QuotaExhausted {
} else if account.availability == AccountAvailability::QuotaExhausted {
AdminAccountStatus::QuotaExhausted
} else if account.availability == ProviderAccountAvailability::Expired
|| account.availability == ProviderAccountAvailability::Invalid
} else if account.availability == AccountAvailability::Expired
|| account.availability == AccountAvailability::Invalid
|| account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now)
@@ -1732,7 +1692,7 @@ fn admin_account_record(summary: ProviderAccountSummary) -> AdminStoreResult<Acc
access_token_expires_at: summary.access_token_expires_at,
next_refresh_at: summary.next_refresh_at,
enabled: summary.enabled,
availability: admin_account_availability(summary.availability),
availability: summary.availability,
availability_reason: summary.availability_reason,
cooldown_until: summary.cooldown_until,
availability_observed_at: summary.availability_observed_at,
@@ -1757,7 +1717,7 @@ fn prepared_account(credential: PreparedCredentialCreate) -> StoreResult<NewProv
access_token_expires_at: credential.access_token_expires_at,
next_refresh_at: credential.next_refresh_at,
enabled: credential.enabled,
availability: provider_account_availability(credential.availability),
availability: credential.availability,
availability_reason: credential.availability_reason,
cooldown_until: credential.cooldown_until,
availability_observed_at: credential.availability_observed_at,
@@ -1772,34 +1732,6 @@ fn provider_document_json(document: ProviderDocument) -> StoreResult<JsonObject>
)
}
const fn provider_account_availability(
availability: AdminAccountAvailability,
) -> ProviderAccountAvailability {
match availability {
AdminAccountAvailability::Unknown => ProviderAccountAvailability::Unknown,
AdminAccountAvailability::Ready => ProviderAccountAvailability::Ready,
AdminAccountAvailability::Cooldown => ProviderAccountAvailability::Cooldown,
AdminAccountAvailability::QuotaExhausted => ProviderAccountAvailability::QuotaExhausted,
AdminAccountAvailability::Expired => ProviderAccountAvailability::Expired,
AdminAccountAvailability::Banned => ProviderAccountAvailability::Banned,
AdminAccountAvailability::Invalid => ProviderAccountAvailability::Invalid,
}
}
const fn admin_account_availability(
availability: ProviderAccountAvailability,
) -> AdminAccountAvailability {
match availability {
ProviderAccountAvailability::Unknown => AdminAccountAvailability::Unknown,
ProviderAccountAvailability::Ready => AdminAccountAvailability::Ready,
ProviderAccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
ProviderAccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
ProviderAccountAvailability::Expired => AdminAccountAvailability::Expired,
ProviderAccountAvailability::Banned => AdminAccountAvailability::Banned,
ProviderAccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
fn admin_account_usage(usage: ProviderAccountUsageObservation) -> AdminStoreResult<AccountUsage> {
Ok(AccountUsage {
account_id: usage.account_id,
@@ -2214,7 +2146,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(DateTime::<Utc>::from),
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.account.enabled(),
availability: availability_from_core(account.account.availability()),
availability: account.account.availability(),
availability_reason: None,
cooldown_until: account.account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: Utc::now(),
@@ -2416,7 +2348,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
account_id: change.account_id.as_str().to_owned(),
expected_revision: Revision::new(change.expected_revision.get())
.map_err(core_store_error)?,
availability: availability_from_core(change.availability),
availability: change.availability,
availability_reason: change.reason,
cooldown_until: change.cooldown_until.map(DateTime::<Utc>::from),
availability_observed_at: DateTime::<Utc>::from(change.observed_at),
@@ -2527,7 +2459,7 @@ fn core_account_from_summary(
)
.with_runtime_state(
summary.enabled,
availability_to_core(summary.availability),
summary.availability,
summary.cooldown_until.map(Into::into),
)
.with_refresh_schedule(
@@ -2536,30 +2468,6 @@ fn core_account_from_summary(
))
}
const fn availability_to_core(value: ProviderAccountAvailability) -> CoreAccountAvailability {
match value {
ProviderAccountAvailability::Unknown => CoreAccountAvailability::Unknown,
ProviderAccountAvailability::Ready => CoreAccountAvailability::Ready,
ProviderAccountAvailability::Cooldown => CoreAccountAvailability::Cooldown,
ProviderAccountAvailability::QuotaExhausted => CoreAccountAvailability::QuotaExhausted,
ProviderAccountAvailability::Expired => CoreAccountAvailability::Expired,
ProviderAccountAvailability::Banned => CoreAccountAvailability::Banned,
ProviderAccountAvailability::Invalid => CoreAccountAvailability::Invalid,
}
}
const fn availability_from_core(value: CoreAccountAvailability) -> ProviderAccountAvailability {
match value {
CoreAccountAvailability::Unknown => ProviderAccountAvailability::Unknown,
CoreAccountAvailability::Ready => ProviderAccountAvailability::Ready,
CoreAccountAvailability::Cooldown => ProviderAccountAvailability::Cooldown,
CoreAccountAvailability::QuotaExhausted => ProviderAccountAvailability::QuotaExhausted,
CoreAccountAvailability::Expired => ProviderAccountAvailability::Expired,
CoreAccountAvailability::Banned => ProviderAccountAvailability::Banned,
CoreAccountAvailability::Invalid => ProviderAccountAvailability::Invalid,
}
}
fn core_store_error(error: StoreError) -> CoreStoreError {
let kind = match error {
StoreError::Unavailable { .. } => CoreStoreErrorKind::Unavailable,
@@ -2600,7 +2508,7 @@ fn account_summary_from_row(row: sqlx::postgres::PgRow) -> StoreResult<ProviderA
access_token_expires_at: get(&row, "access_token_expires_at")?,
next_refresh_at: get(&row, "next_refresh_at")?,
enabled: get(&row, "enabled")?,
availability: ProviderAccountAvailability::parse(&availability)?,
availability: parse_availability(&availability)?,
availability_reason: get(&row, "availability_reason")?,
cooldown_until: get(&row, "cooldown_until")?,
availability_observed_at: get(&row, "availability_observed_at")?,
@@ -55,10 +55,7 @@ impl CredentialStateCache {
"provider_account_id",
&self.provider_account_id,
)?;
if !matches!(
self.availability.as_str(),
"unknown" | "ready" | "cooldown" | "quota_exhausted" | "expired" | "banned" | "invalid"
) {
if AccountAvailability::parse(&self.availability).is_none() {
return Err(invalid("availability is not registered"));
}
Ok(())
@@ -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(
@@ -17,9 +17,9 @@ use gateway_admin::{
ports::store::AccountStore,
};
use gateway_core::engine::credential::{
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, OpaqueProviderData,
PlaintextCredential, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
QuotaObservation, QuotaWriteOutcome,
AccountAvailability, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
};
use gateway_core::routing::ProviderKind;
use gateway_store::{
@@ -27,9 +27,9 @@ use gateway_store::{
postgres::{
AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts,
NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountAvailability,
ProviderAccountObservation, ProviderAccountRepository, ProviderCredentialUpdate,
RotateProviderAccount, SetProviderAccountEnabled, UpdateProviderAccount,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountObservation,
ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount,
SetProviderAccountEnabled, UpdateProviderAccount,
},
};
use serde_json::json;
@@ -123,15 +123,15 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
let mut beta = account("acct_beta", "user-beta");
beta.provider_kind = "xai".to_owned();
beta.email = Some("beta@example.invalid".to_owned());
beta.availability = ProviderAccountAvailability::Banned;
beta.availability = AccountAvailability::Banned;
let mut charlie = account("acct_charlie", "user-charlie");
charlie.email = Some("charlie@example.invalid".to_owned());
let mut invalid = account("acct_invalid", "user-invalid");
invalid.email = Some("invalid@example.invalid".to_owned());
invalid.availability = ProviderAccountAvailability::Invalid;
invalid.availability = AccountAvailability::Invalid;
let mut quota_exhausted = account("acct_quota_exhausted", "user-quota-exhausted");
quota_exhausted.email = Some("quota-exhausted@example.invalid".to_owned());
quota_exhausted.availability = ProviderAccountAvailability::QuotaExhausted;
quota_exhausted.availability = AccountAvailability::QuotaExhausted;
for account in [alpha, beta, charlie, invalid, quota_exhausted] {
repository
.insert_provider_account(account)
@@ -292,10 +292,10 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
);
credential.provider_kind = "xai".to_owned();
credential.availability = match index {
0 => ProviderAccountAvailability::Expired,
1 => ProviderAccountAvailability::Banned,
2 => ProviderAccountAvailability::Invalid,
_ => ProviderAccountAvailability::Ready,
0 => AccountAvailability::Expired,
1 => AccountAvailability::Banned,
2 => AccountAvailability::Invalid,
_ => AccountAvailability::Ready,
};
repository
.insert_provider_account(credential)
@@ -644,7 +644,7 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
let mut updated = account("acct_admin_reimport", "user-admin-upsert");
updated.name = "updated import".to_owned();
updated.provider_credentials_json = credential_json("updated-import-secret");
updated.availability = ProviderAccountAvailability::Banned;
updated.availability = AccountAvailability::Banned;
updated.availability_reason = Some("upstream_account_deactivated".to_owned());
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
@@ -716,7 +716,7 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
access_token_expires_at: Some(Utc::now() + TimeDelta::minutes(5)),
next_refresh_at: None,
enabled: true,
availability: ProviderAccountAvailability::Ready,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: Utc::now(),
@@ -821,7 +821,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
};
let mut cooldown_account = account("acct_admin_b", "user-admin-b");
cooldown_account.availability = ProviderAccountAvailability::Cooldown;
cooldown_account.availability = AccountAvailability::Cooldown;
cooldown_account.cooldown_until = Some(Utc::now() + TimeDelta::minutes(10));
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
@@ -961,7 +961,7 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
access_token_expires_at: Some(Utc::now() + TimeDelta::hours(1)),
next_refresh_at: None,
enabled: true,
availability: ProviderAccountAvailability::Ready,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: Utc::now(),
@@ -1094,7 +1094,7 @@ fn provider_credentials_are_redacted_from_debug() {
fn cooldown_observation_requires_expiry() {
let observation = ProviderAccountObservation {
account_id: "account-1".to_owned(),
availability: ProviderAccountAvailability::Cooldown,
availability: AccountAvailability::Cooldown,
availability_reason: None,
cooldown_until: None,
provider_quota_json: None,
@@ -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
);
}
+192 -68
View File
@@ -5,9 +5,7 @@ use std::time::SystemTime;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_admin::model::accounts::{
AccountAvailability as AdminAccountAvailability, AccountRecord,
};
use gateway_admin::model::accounts::AccountRecord;
use gateway_admin::model::observability::{
CalculatedBillingBreakdown, CurrencyCost, DashboardDesktopRelease, DashboardWireAttribute,
DashboardWireProfile, DashboardWireTarget, DecimalAmount, DesktopReleaseStatus,
@@ -27,15 +25,15 @@ use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
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};
@@ -50,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;
@@ -253,11 +252,22 @@ impl ProviderAdmin for OpenAiAdminProvider {
command.document.into_provider_data().into_inner(),
))
.await
.inspect_err(|error| {
log_import_failure("prepare_document", credential_admin_error_code(*error));
})
.map_err(map_credential_admin_error)?;
let observed_at = Utc::now();
let mut credentials = Vec::with_capacity(prepared.accounts().len());
for account in prepared.into_accounts() {
let account = self.preserve_existing_installation_id(account).await?;
let account = self
.preserve_existing_installation_id(account)
.await
.inspect_err(|error| {
log_import_failure(
"preserve_installation_id",
provider_admin_error_code(error.kind()),
);
})?;
credentials.push(prepared_create(account, observed_at)?);
}
Ok(PreparedCredentialImport {
@@ -289,6 +299,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
&self,
command: CompleteAuthorization,
) -> Result<PreparedAuthorizationCommit, ProviderAdminError> {
let request_id = command.context.request_id.clone();
let binding = AuthorizationOwnerBinding::from_context(&command.context);
let completed = self
.oauth
@@ -298,30 +309,46 @@ impl ProviderAdmin for OpenAiAdminProvider {
callback_url: SecretString::from(command.callback_url),
})
.await
.inspect_err(|error| {
tracing::warn!(
request_id = %request_id,
oauth_stage = "complete_authorization",
oauth_error = oauth_error_code(error),
"OpenAI OAuth authorization completion failed"
);
})
.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(
@@ -519,7 +546,7 @@ fn prepared_create(
account,
credential,
} = prepared;
let availability = admin_availability(account.availability());
let availability = account.availability();
Ok(PreparedCredentialCreate {
account_id: account.id().clone(),
provider_kind: account.provider().clone(),
@@ -639,7 +666,7 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
)
.with_runtime_state(
account.enabled,
core_availability(account.availability),
account.availability,
account.cooldown_until.map(SystemTime::from),
)
.with_refresh_schedule(
@@ -698,7 +725,7 @@ fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) ->
&& account.plan_type() == record.plan_type.as_deref()
&& account.authentication_kind() == record.authentication_kind
&& account.enabled() == record.enabled
&& admin_availability(account.availability()) == record.availability
&& 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
@@ -706,30 +733,6 @@ fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) ->
&& account.has_refresh_token() == record.has_refresh_token
}
const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability {
match value {
AccountAvailability::Unknown => AdminAccountAvailability::Unknown,
AccountAvailability::Ready => AdminAccountAvailability::Ready,
AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
AccountAvailability::Expired => AdminAccountAvailability::Expired,
AccountAvailability::Banned => AdminAccountAvailability::Banned,
AccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
const fn core_availability(value: AdminAccountAvailability) -> AccountAvailability {
match value {
AdminAccountAvailability::Unknown => AccountAvailability::Unknown,
AdminAccountAvailability::Ready => AccountAvailability::Ready,
AdminAccountAvailability::Cooldown => AccountAvailability::Cooldown,
AdminAccountAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted,
AdminAccountAvailability::Expired => AccountAvailability::Expired,
AdminAccountAvailability::Banned => AccountAvailability::Banned,
AdminAccountAvailability::Invalid => AccountAvailability::Invalid,
}
}
fn empty_quota() -> ProviderQuota {
ProviderQuota {
observed_at: None,
@@ -903,21 +906,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),
}
}
}
@@ -1238,12 +1306,52 @@ fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdmin
})
}
const fn credential_admin_error_code(error: CodexCredentialAdminError) -> &'static str {
match error {
CodexCredentialAdminError::InvalidInput => "invalid_input",
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",
CodexCredentialAdminError::AccountBanned => "account_banned",
CodexCredentialAdminError::RefreshUnavailable => "refresh_unavailable",
CodexCredentialAdminError::RefreshAmbiguous => "refresh_ambiguous",
CodexCredentialAdminError::IdentityRejected => "identity_rejected",
CodexCredentialAdminError::IdentityUnavailable => "identity_unavailable",
}
}
const fn provider_admin_error_code(kind: ProviderAdminErrorKind) -> &'static str {
match kind {
ProviderAdminErrorKind::Invalid => "invalid",
ProviderAdminErrorKind::Unsupported => "unsupported",
ProviderAdminErrorKind::NotFound => "not_found",
ProviderAdminErrorKind::Conflict => "conflict",
ProviderAdminErrorKind::Unavailable => "unavailable",
ProviderAdminErrorKind::Internal => "internal",
}
}
fn log_import_failure(stage: &'static str, error: &'static str) {
tracing::warn!(
import_stage = stage,
import_error = error,
"OpenAI credential import preparation failed"
);
}
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 => {
@@ -1252,6 +1360,22 @@ fn map_oauth_error(error: CodexOAuthAdminError) -> ProviderAdminError {
})
}
const fn oauth_error_code(error: &CodexOAuthAdminError) -> &'static str {
match error {
CodexOAuthAdminError::InvalidInput => "invalid_input",
CodexOAuthAdminError::NotFound => "not_found",
CodexOAuthAdminError::Conflict => "conflict",
CodexOAuthAdminError::FlowExpired => "flow_expired",
CodexOAuthAdminError::CallbackRejected => "callback_rejected",
CodexOAuthAdminError::TokenRejected => "token_rejected",
CodexOAuthAdminError::IdentityRejected => "identity_rejected",
CodexOAuthAdminError::UpstreamUnavailable => "upstream_unavailable",
CodexOAuthAdminError::Ambiguous => "ambiguous",
CodexOAuthAdminError::StorageUnavailable => "storage_unavailable",
CodexOAuthAdminError::Credential => "credential",
}
}
fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
use CodexCredentialQuotaError as Error;
provider_admin_error(match error {
@@ -42,32 +42,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 +107,6 @@ enum ParsedCodexAuthentication {
OAuth {
access_token: Option<String>,
refresh_token: Option<String>,
id_token: Option<String>,
},
AgentIdentity {
runtime_id: String,
@@ -148,7 +121,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 +128,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: _,
@@ -720,7 +691,7 @@ impl CodexCredentialAdminService {
Ok(prepared)
}
/// 按文档结构严格判别账号布局,并归一到唯一 `NewProviderAccount` 写入路径。
/// 从 OAuth token 对补全账号事实,并归一到唯一 `NewProviderAccount` 写入路径。
pub async fn prepare_import_document(
&self,
payload: Value,
@@ -741,19 +712,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 +750,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 +784,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 +809,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 +835,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))
}
@@ -1058,67 +1022,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 +1115,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 +1167,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 +1196,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,9 +47,10 @@ pub use identity::{
};
pub use oauth::{
CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthAuthorizationStarted,
CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthReauthorizationTarget,
CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthAuthorization,
CompletedCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexOAuthReauthorizationTarget, CodexPendingAuthorization, CompleteCodexOAuthAuthorization,
CompletedCodexOAuthAuthorization, CompletedCodexOAuthCredential, StartCodexOAuthAuthorization,
StoredCodexPendingAuthorization,
};
pub use quota::{
CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService,
@@ -1,13 +1,16 @@
//! 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,
@@ -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;
@@ -89,6 +93,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 +121,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
@@ -284,11 +301,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 +340,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 +379,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(
@@ -412,7 +495,38 @@ 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
@@ -483,10 +597,7 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
.map_err(map_admin_error)?,
)
};
Ok(CompletedCodexOAuthAuthorization {
mutation,
credential,
})
Ok((mutation, credential))
}
}
@@ -533,38 +644,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 +717,7 @@ impl CodexOAuthAdminService {
})
.await
.map_err(map_exchange_error)?;
Ok((pending, tokens.secret, tokens.id_token))
Ok((tokens.secret, tokens.id_token))
}
}
@@ -622,32 +765,44 @@ 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;
for (key, value) in url.query_pairs() {
match key.as_ref() {
"code" if code.is_none() => code = Some(value.into_owned()),
"state" if state.is_none() => state = Some(value.into_owned()),
_ => return Err(CodexOAuthAdminError::UpstreamRejected),
"code" => set_unique_callback_parameter(&mut code, value.into_owned())?,
"state" => set_unique_callback_parameter(&mut state, value.into_owned())?,
// OAuth 服务可能附加 `scope`、`iss` 等标准参数。
// 固定回调地址与唯一的 `code`/`state` 才是安全边界,无关参数不参与校验。
_ => {}
}
}
let code = code.filter(|value| valid_secret(value));
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),
}
}
fn set_unique_callback_parameter(
target: &mut Option<String>,
value: String,
) -> Result<(), CodexOAuthAdminError> {
if target.replace(value).is_some() {
return Err(CodexOAuthAdminError::CallbackRejected);
}
Ok(())
}
fn random_secret() -> Result<String, CodexOAuthAdminError> {
let mut bytes = [0_u8; 32];
getrandom::fill(&mut bytes).map_err(|_| CodexOAuthAdminError::StorageUnavailable)?;
@@ -710,16 +865,33 @@ 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,
}
}
fn map_identity_error(_: CodexIdentityVerificationError) -> CodexOAuthAdminError {
CodexOAuthAdminError::UpstreamRejected
fn map_identity_error(error: CodexIdentityVerificationError) -> CodexOAuthAdminError {
match error {
CodexIdentityVerificationError::Rejected => CodexOAuthAdminError::IdentityRejected,
CodexIdentityVerificationError::Unavailable => CodexOAuthAdminError::UpstreamUnavailable,
}
}
fn map_admin_error(_: CodexCredentialAdminError) -> CodexOAuthAdminError {
@@ -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 主动额度刷新的调度策略。
///
@@ -235,6 +238,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 +250,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 +418,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 +437,160 @@ 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 kind = match account.availability() {
AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_observation(&account, observation),
},
AccountAvailability::Ready
if observations.is_some() && quota_observation_is_stale(&account, observation, now) =>
{
CodexQuotaRefreshKind::Ready
}
_ => return None,
};
Some(CodexQuotaRefreshCandidate { account, kind })
}
fn quota_reset_at_from_observation(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
) -> Option<SystemTime> {
observation
.filter(|observation| observation.expected_revision == account.revision())
.and_then(quota_snapshot_from_observation)
.and_then(|snapshot| snapshot.fact().resets_at())
.map(SystemTime::from)
}
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,
@@ -519,7 +677,26 @@ impl CodexCredentialQuotaService {
let 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
}
};
let accounts =
self.scheduling
.reserve_periodic_refreshes(accounts, observations.as_ref(), now);
if accounts.is_empty() {
return Ok(summary);
}
@@ -918,12 +1095,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(
@@ -973,13 +1149,11 @@ fn merge_passive_quota(
.and_then(|value| value.as_array().cloned())
.unwrap_or_default();
for (limit_id, details) in &rate_limits.limits {
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
continue;
};
if Some(limit_id.as_str()) == active_limit {
let mut value = quota
.remove("rate_limit")
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
merge_rate_limit_details(&mut value, details);
quota.insert("rate_limit".to_owned(), Value::Object(value));
quota.insert("rate_limit".to_owned(), Value::Object(rate_limit));
continue;
}
let index = additional.iter().position(|item| {
@@ -1000,12 +1174,7 @@ fn merge_passive_quota(
if let Some(name) = details.limit_name.as_ref() {
item.insert("limit_name".to_owned(), Value::String(name.clone()));
}
let mut value = item
.remove("rate_limit")
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
merge_rate_limit_details(&mut value, details);
item.insert("rate_limit".to_owned(), Value::Object(value));
item.insert("rate_limit".to_owned(), Value::Object(rate_limit));
match index {
Some(index) => additional[index] = Value::Object(item),
None => additional.push(Value::Object(item)),
@@ -1018,12 +1187,14 @@ fn merge_passive_quota(
quota
}
fn merge_rate_limit_details(target: &mut Map<String, Value>, details: &RateLimitDetails) {
fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option<Map<String, Value>> {
// 响应头的窗口属于同一次上游观测;跨响应合并会制造不存在的额度窗口。
let mut snapshot = Map::new();
if let Some(allowed) = details.allowed {
target.insert("allowed".to_owned(), Value::Bool(allowed));
snapshot.insert("allowed".to_owned(), Value::Bool(allowed));
}
if let Some(limit_reached) = details.limit_reached {
target.insert("limit_reached".to_owned(), Value::Bool(limit_reached));
snapshot.insert("limit_reached".to_owned(), Value::Bool(limit_reached));
}
for (field, window) in [
("primary_window", details.primary),
@@ -1032,28 +1203,29 @@ fn merge_rate_limit_details(target: &mut Map<String, Value>, details: &RateLimit
let Some(window) = window else {
continue;
};
let mut value = target
.remove(field)
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
merge_rate_limit_window(&mut value, window);
target.insert(field.to_owned(), Value::Object(value));
snapshot.insert(
field.to_owned(),
Value::Object(passive_rate_limit_window(window)),
);
}
(!snapshot.is_empty()).then_some(snapshot)
}
fn merge_rate_limit_window(target: &mut Map<String, Value>, window: RateLimitWindow) {
fn passive_rate_limit_window(window: RateLimitWindow) -> Map<String, Value> {
let mut snapshot = Map::new();
if let Some(number) = serde_json::Number::from_f64(window.used_percent) {
target.insert("used_percent".to_owned(), Value::Number(number));
snapshot.insert("used_percent".to_owned(), Value::Number(number));
}
if let Some(seconds) = window
.window_minutes
.and_then(|minutes| minutes.checked_mul(60))
{
target.insert("limit_window_seconds".to_owned(), Value::from(seconds));
snapshot.insert("limit_window_seconds".to_owned(), Value::from(seconds));
}
if let Some(reset_at) = window.reset_at {
target.insert("reset_at".to_owned(), Value::from(reset_at));
snapshot.insert("reset_at".to_owned(), Value::from(reset_at));
}
snapshot
}
fn quota_success_availability(
@@ -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,
});
}
@@ -537,7 +554,7 @@ impl CodexCredentialSelector {
CodexAccountFailure::RateLimited { retry_after } => {
let until = now.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
(
AccountAvailability::QuotaExhausted,
AccountAvailability::Cooldown,
Some("rate_limited".to_owned()),
until,
CookieRecovery::None,
@@ -584,6 +601,7 @@ impl CodexCredentialSelector {
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 +644,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 +855,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 +889,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 {
+95 -100
View File
@@ -376,6 +376,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 +392,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 +662,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
@@ -686,7 +700,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 +751,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 +762,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 +786,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 +797,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,
@@ -847,20 +850,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 +868,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)) =
@@ -906,33 +905,14 @@ 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 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,
)
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,6 +920,19 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
return;
}
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()?);
}
for event in events {
yield event;
}
@@ -948,12 +941,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
return;
}
if completed {
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
return;
}
}
@@ -977,19 +964,13 @@ 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(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
failure,
)
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let (Some(capture), Some(updates)) =
@@ -1001,15 +982,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,6 +996,15 @@ 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 allows_account_state_mutation && completed && terminal_failure.is_none() {
// 同上:尾部 finish() 也可能产出 completed,亲和记录必须先于任何下游 yield。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
@@ -1033,14 +1015,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure.error)?;
return;
}
if completed {
selector
.record_success(&active_account, session_affinity_key.as_ref())
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
})
}
@@ -1781,24 +1755,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 +1798,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!(
+74 -12
View File
@@ -33,7 +33,8 @@ 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,
@@ -339,7 +340,7 @@ 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")
@@ -822,7 +823,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 +845,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)
})
}
}
@@ -628,7 +628,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 +645,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 +667,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 +718,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 +728,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"
@@ -603,7 +603,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 +629,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()
@@ -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();
@@ -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,25 @@ 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(),
reauthorization_credential_revision: pending.reauthorization_credential_revision,
mutation: pending.mutation.clone(),
})
}
#[async_trait]
@@ -41,42 +59,90 @@ 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(|target| target.account_id().to_string()),
reauthorization_credential_revision: pending
.reauthorization()
.map(|target| target.credential_revision().get()),
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)
}
}
@@ -266,7 +332,130 @@ async fn completion_returns_verified_core_account_without_writing_store() {
}
#[tokio::test]
async fn callback_state_mismatch_fails_closed_and_consumes_flow() {
async fn completion_accepts_unrelated_oauth_callback_parameters() {
let service = service();
let started = started(&service).await;
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 prepared = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&state={state}&scope=openid%20profile&iss=https%3A%2F%2Fauth.openai.com"
)),
})
.await
.expect("complete OAuth with standard callback parameters");
assert!(matches!(
prepared.credential,
CompletedCodexOAuthCredential::Create(_)
));
}
#[tokio::test]
async fn completion_rejects_repeated_code_callback_parameter() {
let service = service();
let started = started(&service).await;
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 error = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&code=another-code&state={state}"
)),
})
.await
.expect_err("repeated code");
assert_eq!(error, CodexOAuthAdminError::CallbackRejected);
}
#[tokio::test]
async fn completion_rejects_repeated_state_callback_parameter() {
let service = service();
let started = started(&service).await;
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 error = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&state={state}&state=another-state-value"
)),
})
.await
.expect_err("repeated state");
assert_eq!(error, CodexOAuthAdminError::CallbackRejected);
}
struct UnavailableVerifier;
#[async_trait]
impl CodexAccountIdentityVerifier for UnavailableVerifier {
async fn verify(
&self,
_: &CodexOAuthSecret,
_: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError> {
Err(CodexIdentityVerificationError::Unavailable)
}
async fn verify_authorization(
&self,
_: &CodexOAuthSecret,
_: &SecretString,
_: &SecretString,
_: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError> {
Err(CodexIdentityVerificationError::Unavailable)
}
}
#[tokio::test]
async fn completion_classifies_unavailable_identity_verification_as_upstream_unavailable() {
let store = Arc::new(MemoryAccountStore::default());
let service = service_with_store(store, Arc::new(UnavailableVerifier));
let started = started(&service).await;
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 error = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
callback_url: SecretString::from(format!(
"http://localhost:1455/auth/callback?code=authorization-code&state={state}"
)),
})
.await
.expect_err("identity verification unavailable");
assert_eq!(error, CodexOAuthAdminError::UpstreamUnavailable);
}
#[tokio::test]
async fn callback_state_mismatch_releases_flow_for_retry() {
let service = service();
let started = started(&service).await;
let flow_id = started.flow_id;
@@ -282,14 +471,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
);
}
@@ -361,6 +361,73 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
);
}
#[tokio::test]
async fn passive_rate_limit_headers_replace_stale_secondary_window() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_passive_quota_snapshot").await;
let account = store
.account("acct_passive_quota_snapshot")
.expect("account");
let baseline = json!({
"rate_limit": {
"primary_window": {
"used_percent": 91,
"reset_at": 1_900_000_000,
"limit_window_seconds": 2_592_000
},
"secondary_window": {
"used_percent": 43,
"reset_at": 1_900_000_000,
"limit_window_seconds": 2_592_000
}
}
});
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
baseline.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist baseline quota");
let service = quota_service(&store);
let headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "92".to_owned()),
(
"x-codex-primary-window-minutes".to_owned(),
"43200".to_owned(),
),
(
"x-codex-primary-reset-at".to_owned(),
"1900000000".to_owned(),
),
];
assert!(
service
.synchronize_passive_headers(&account, &headers)
.await
.expect("passive quota")
);
let snapshot = service
.read_account(account.id())
.await
.expect("read quota")
.expect("quota snapshot");
let monthly = snapshot
.windows()
.iter()
.filter(|window| window.kind() == CodexQuotaWindowKind::Monthly)
.collect::<Vec<_>>();
assert_eq!(monthly.len(), 1);
assert_eq!(monthly[0].used_percent(), Some(92.0));
}
#[tokio::test]
async fn synchronize_without_accounts_is_a_noop_before_network_io() {
let store = Arc::new(MemoryAccountStore::default());
@@ -379,18 +446,106 @@ 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_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());
@@ -444,6 +599,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,
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::{
@@ -122,6 +126,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 +137,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 +155,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 +206,25 @@ 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(),
)
}
#[tokio::test]
async fn openai_provider_rejects_a_foreign_provider_candidate_before_account_selection() {
let store = Arc::new(MemoryAccountStore::default());
@@ -271,6 +311,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 +393,129 @@ 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 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());
@@ -70,6 +70,17 @@ impl MemoryAccountStore {
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]
@@ -440,6 +451,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 {
+260 -117
View File
@@ -2,22 +2,20 @@
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};
use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::model::accounts::{
AccountAvailability as AdminAccountAvailability, AccountRecord,
};
use gateway_admin::model::accounts::AccountRecord;
use gateway_admin::model::observability::{
CalculatedBillingBreakdown, CurrencyCost, DashboardDesktopRelease, DashboardWireAttribute,
DashboardWireProfile, DashboardWireTarget, DecimalAmount, DesktopReleaseStatus,
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,
@@ -25,20 +23,19 @@ 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::{
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
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;
@@ -63,8 +60,76 @@ use crate::transport::{GROK_CLI_BASE_URL, XAI_PROVIDER_NAME, grok_billing_breakd
const PENDING_SCHEMA_VERSION: u64 = 2;
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,
@@ -175,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))
}
}
}
}
@@ -436,76 +551,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(
@@ -692,6 +769,82 @@ 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,
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(),
)?)
}
};
Ok(PreparedAuthorizationCommit::new(
stored.mutation,
credential,
))
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RotationDocument {
@@ -773,7 +926,7 @@ fn prepared_create(
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
next_refresh_at: account.next_refresh_at().map(Into::into),
enabled: account.enabled(),
availability: admin_availability(account.availability()),
availability: account.availability(),
availability_reason: None,
cooldown_until: account.cooldown_until().map(Into::into),
availability_observed_at: observed_at,
@@ -880,7 +1033,7 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
)
.with_runtime_state(
account.enabled,
core_availability(account.availability),
account.availability,
account.cooldown_until.map(Into::into),
)
.with_refresh_schedule(
@@ -899,7 +1052,7 @@ fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) ->
&& account.upstream_account_id() == record.upstream_account_id.as_deref()
&& account.plan_type() == record.plan_type.as_deref()
&& account.enabled() == record.enabled
&& admin_availability(account.availability()) == record.availability
&& 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
@@ -907,30 +1060,6 @@ fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) ->
&& account.has_refresh_token() == record.has_refresh_token
}
const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability {
match value {
AccountAvailability::Unknown => AdminAccountAvailability::Unknown,
AccountAvailability::Ready => AdminAccountAvailability::Ready,
AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
AccountAvailability::Expired => AdminAccountAvailability::Expired,
AccountAvailability::Banned => AdminAccountAvailability::Banned,
AccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
const fn core_availability(value: AdminAccountAvailability) -> AccountAvailability {
match value {
AdminAccountAvailability::Unknown => AccountAvailability::Unknown,
AdminAccountAvailability::Ready => AccountAvailability::Ready,
AdminAccountAvailability::Cooldown => AccountAvailability::Cooldown,
AdminAccountAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted,
AdminAccountAvailability::Expired => AccountAvailability::Expired,
AdminAccountAvailability::Banned => AccountAvailability::Banned,
AdminAccountAvailability::Invalid => AccountAvailability::Invalid,
}
}
fn project_quota(
snapshot: Option<GrokQuotaSnapshot>,
refresh_token_expires_at: Option<DateTime<Utc>>,
@@ -1369,6 +1498,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,
@@ -6,10 +6,10 @@ use std::time::SystemTime;
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::{
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
LoadedCredential, NewProviderAccount, OpaqueProviderData, PlaintextCredential, ProviderAccount,
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation,
QuotaWriteOutcome,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::provider_ports::ProviderRefreshPolicy;
@@ -170,7 +170,7 @@ impl GrokCredentialAdmin {
)
.with_runtime_state(
input.enabled,
availability(input.initial_availability),
input.initial_availability.into(),
input.initial_cooldown_until.map(to_system_time),
)
.with_refresh_schedule(true, Some(to_system_time(input.next_refresh_at)));
@@ -440,7 +440,7 @@ impl GrokCredentialRepository {
.apply_state_change(AccountStateChange {
account_id: input.account_id.clone(),
expected_revision: input.expected_revision,
availability: availability(input.availability),
availability: input.availability.into(),
reason: input.availability_reason.clone(),
cooldown_until: input.cooldown_until.map(to_system_time),
observed_at: to_system_time(input.observed_at),
@@ -839,18 +839,6 @@ fn validate_cooldown(
Ok(())
}
fn availability(value: GrokCredentialAvailability) -> AccountAvailability {
match value {
GrokCredentialAvailability::Unknown => AccountAvailability::Unknown,
GrokCredentialAvailability::Ready => AccountAvailability::Ready,
GrokCredentialAvailability::Cooldown => AccountAvailability::Cooldown,
GrokCredentialAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted,
GrokCredentialAvailability::Expired => AccountAvailability::Expired,
GrokCredentialAvailability::Banned => AccountAvailability::Banned,
GrokCredentialAvailability::Invalid => AccountAvailability::Invalid,
}
}
fn ensure_xai(account: &ProviderAccount) -> Result<(), GrokCredentialRepositoryError> {
if account.provider().as_str() == XAI_PROVIDER_KIND
&& account.authentication_kind() == XAI_AUTHENTICATION_KIND_OAUTH
@@ -6,8 +6,8 @@ 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,
@@ -66,39 +66,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 +121,11 @@ 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).await
};
let mut candidates = catalog_eligible
.into_iter()
.filter(|account| !runtime_cooldowns.contains_key(account.id()))
@@ -147,12 +160,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 +201,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,7 +222,12 @@ 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 {
@@ -297,6 +318,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 => {
@@ -336,6 +360,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> {
@@ -4,8 +4,8 @@ use std::fmt;
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::{
CredentialCasUpdate, CredentialRevision, LoadedCredential, ProviderAccountId,
ProviderAccountUpdate,
AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential,
ProviderAccountId, ProviderAccountUpdate,
};
use gateway_core::provider_ports::ProviderLeaseGuard;
use serde::{Deserialize, Serialize};
@@ -75,6 +75,21 @@ pub enum GrokCredentialAvailability {
Invalid,
}
/// Grok 可用性语义翻入内核标准态的唯一汇入点。
impl From<GrokCredentialAvailability> for AccountAvailability {
fn from(value: GrokCredentialAvailability) -> Self {
match value {
GrokCredentialAvailability::Unknown => Self::Unknown,
GrokCredentialAvailability::Ready => Self::Ready,
GrokCredentialAvailability::Cooldown => Self::Cooldown,
GrokCredentialAvailability::QuotaExhausted => Self::QuotaExhausted,
GrokCredentialAvailability::Expired => Self::Expired,
GrokCredentialAvailability::Banned => Self::Banned,
GrokCredentialAvailability::Invalid => Self::Invalid,
}
}
}
/// 创建一个明文 OAuth Provider account。
pub struct CreateGrokCredential {
pub account_id: ProviderAccountId,
+44 -6
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))?;
@@ -1108,6 +1123,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 +1136,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;
@@ -1333,6 +1352,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 +1372,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 +1478,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 {
@@ -1468,6 +1500,9 @@ async fn record_stream_failure(
session: &SelectedGrokSession,
error: ProviderError,
) -> ProviderError {
if !session.allows_account_state_mutation() {
return error;
}
if matches!(
error.kind(),
ProviderErrorKind::Unauthorized
@@ -1491,6 +1526,9 @@ async fn recover_or_record_failure(
error: ProviderError,
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())
@@ -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> {
@@ -296,6 +331,16 @@ pub trait GrokSessionSelector: Send + Sync {
session: &'a SelectedGrokSession,
failure: GrokCredentialFailure,
) -> GrokCredentialFeedbackFuture<'a>;
/// 成功完成真实上游请求后收敛账号的可恢复状态。
///
/// 默认实现让只关心选择或失败反馈的测试替身无需模拟持久状态写入。
fn record_success<'a>(
&'a self,
_: &'a SelectedGrokSession,
) -> GrokCredentialFeedbackFuture<'a> {
Box::pin(async {})
}
}
/// 不含密钥的选择器失败。
+82 -33
View File
@@ -6,9 +6,7 @@ use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, Utc};
use futures::{StreamExt as _, future::BoxFuture};
use gateway_admin::model::accounts::{
AccountAvailability as AdminAccountAvailability, AccountRecord,
};
use gateway_admin::model::accounts::AccountRecord;
use gateway_admin::model::observability::{
CurrencyCost, DesktopReleaseStatus, ProviderBillingInput,
};
@@ -32,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,
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts,
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
ProviderStoreError, ProviderStorePorts,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
@@ -374,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");
@@ -401,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]
@@ -620,7 +619,7 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.enabled(),
availability: admin_availability(account.availability()),
availability: account.availability(),
availability_reason: None,
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: now,
@@ -630,18 +629,6 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
}
}
const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability {
match value {
AccountAvailability::Unknown => AdminAccountAvailability::Unknown,
AccountAvailability::Ready => AdminAccountAvailability::Ready,
AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown,
AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted,
AccountAvailability::Expired => AdminAccountAvailability::Expired,
AccountAvailability::Banned => AdminAccountAvailability::Banned,
AccountAvailability::Invalid => AdminAccountAvailability::Invalid,
}
}
struct TestLeases;
impl ProviderLeasePort for TestLeases {
@@ -833,7 +820,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) {
@@ -843,6 +830,7 @@ impl TestOAuthPending {
owner,
payload,
SystemTime::now() + Duration::from_secs(1_800),
None,
),
);
}
@@ -872,36 +860,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)
})
}
}
@@ -8,7 +8,7 @@ import {
quotaWindowPercentTextClass,
} from '../constants'
type UsageWindowVariant = 'compact' | 'detail'
type UsageWindowVariant = 'compact' | 'detail' | 'metric'
type UsageWindowMode = 'quota' | 'local' | 'unknown'
const props = withDefaults(
@@ -24,6 +24,7 @@ const props = withDefaults(
)
const isCompact = computed(() => props.variant === 'compact')
const isMetric = computed(() => props.variant === 'metric')
const localUsage = computed<AccountLocalUsage | null>(() => {
const value = props.window?.localUsage
if (!value || typeof value !== 'object' || Array.isArray(value))
@@ -55,32 +56,56 @@ const quotaLocalUsageVisible = computed(() =>
&& localUsage.value.totalTokens > 0,
)
function variantValue(metric: string, compact: string, detail: string) {
if (isMetric.value)
return metric
if (isCompact.value)
return compact
return detail
}
const rootClass = computed(() =>
isCompact.value ? 'flex min-w-0 flex-col' : 'rounded-lg bg-(--cp-bg-subtle) p-2',
variantValue(
'grid min-w-0 grid-rows-[11px_14px] gap-1',
'flex min-w-0 flex-col',
'rounded-lg bg-(--cp-bg-subtle) p-2',
),
)
const headerClass = computed(() =>
isCompact.value
? 'mb-1 flex items-center justify-between gap-2 text-[11px] leading-none font-bold'
: 'flex items-center justify-between gap-3 text-[12px] font-bold',
variantValue(
'flex items-center justify-between gap-2 text-[11px] leading-none font-bold',
'mb-1 flex items-center justify-between gap-2 text-[11px] leading-none font-bold',
'flex items-center justify-between gap-3 text-[12px] font-bold',
),
)
const labelClass = computed(() =>
isCompact.value ? 'text-(--cp-text-muted)' : 'text-(--cp-text-secondary)',
isCompact.value || isMetric.value ? 'text-(--cp-text-muted)' : 'text-(--cp-text-secondary)',
)
const valueClass = computed(() =>
isCompact.value
? 'text-[10px] leading-none font-heavy'
: 'text-[12px] font-heavy',
variantValue(
'text-[11px] leading-none font-bold',
'text-[10px] leading-none font-heavy',
'text-[12px] font-heavy',
),
)
const trackOffsetClass = computed(() =>
variantValue('self-center', 'mt-auto', 'mt-2'),
)
const trackShapeClass = computed(() =>
isCompact.value
? 'h-1 w-full overflow-hidden rounded-full'
: 'h-2 overflow-hidden rounded-full',
variantValue(
'h-1.5 w-full overflow-hidden rounded-full',
'h-1 w-full overflow-hidden rounded-full',
'h-2 overflow-hidden rounded-full',
),
)
const trackClass = computed(() => `${trackShapeClass.value} bg-(--cp-default-border)`)
const barStyle = computed(() => {
if (!props.window)
return undefined
return quotaWindowBarStyle(props.window, isCompact.value ? '4px' : '8px')
return quotaWindowBarStyle(
props.window,
variantValue('6px', '4px', '8px'),
)
})
const barClass = computed(() => props.window ? quotaWindowBarClass(props.window) : undefined)
const percentTextClass = computed(() =>
@@ -162,7 +187,7 @@ const requestBars = computed(() => {
</span>
</div>
<div
:class="[trackClass, !isCompact ? 'mt-2' : 'mt-auto']"
:class="[trackClass, trackOffsetClass]"
role="progressbar"
:aria-label="window.labelDisplay"
aria-valuemin="0"
@@ -198,7 +223,7 @@ const requestBars = computed(() => {
</div>
<div
class="flex items-stretch gap-px"
:class="[trackShapeClass, !isCompact ? 'mt-2' : 'mt-auto']"
:class="[trackShapeClass, trackOffsetClass]"
role="img"
:aria-label="requestTimelineTitle"
:title="requestTimelineTitle"
@@ -272,7 +272,7 @@ const statusBars = computed(() => {
<AccountUsageWindow
:window="account.usageWindow"
:show-local-value="false"
variant="compact"
variant="metric"
/>
</span>
</article>
+1 -1
View File
@@ -1 +1 @@
version: "3.0.14"
version: "3.0.18"