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

fix: resolve maintainability audit findings

Browse Source 
Serialize snapshot publication and preserve cancellation-safe execution cleanup.
Move the read-only catalog contract to routing and reuse xAI egress and JWKS state.
Unify admin post-commit observations, retain account UI state across remounts, and remove unused contracts and brittle tests.
Document regression evidence, integration verification, and remaining coverage gaps.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-12 15:17:22 +08:00
co-authored by Codex
parent 881e78975e
commit bfca754714
57 changed files with 2977 additions and 1364 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/gateway/tests/main.rs
+3 -22
View File
@@ -1,6 +1,9 @@
mod architecture;
mod bootstrap;
// 组合根是否混入业务策略由代码审查判断，行数和标识符/注释关键词不能证明职责越界。
// 依赖 DAG、公开模块面和模块镜像继续由 architecture 检查，源码纪律由下方测试维护。
use std::{
collections::BTreeSet,
fs,
@@ -75,28 +78,6 @@ fn workspace_production_files_have_no_hidden_modules_or_test_hooks() {
}
}
#[test]
fn bootstrap_owns_only_bundle_wiring() {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src");
let bootstrap = fs::read_to_string(root.join("bootstrap.rs")).expect("read bootstrap");
for forbidden in [
"sqlx::",
"redis::",
"Repository",
"impl Provider",
"impl ExecutionService",
"tokio::spawn",
"access_token",
"refresh_token",
] {
assert!(
!bootstrap.contains(forbidden),
"bootstrap owns `{forbidden}`"
);
}
assert!(bootstrap.lines().count() <= 300);
}
fn rust_files(root: &Path) -> BTreeSet<PathBuf> {
let mut result = BTreeSet::new();
collect_rust_files(root, root, &mut result);
backend/crates/gateway-admin/src/model/provider_credentials.rs
+1 -51
View File
@@ -15,7 +15,7 @@ use serde_json::{Map, Value};
use uuid::Uuid;
use super::{
AdminError, MutationActor, MutationContext, PageSize, Revision,
AdminError, MutationActor, MutationContext, Revision,
accounts::{
AccountImportSettings, AccountRecord, AccountSummary, AccountUsage, CredentialState,
},
@@ -49,56 +49,6 @@ impl fmt::Debug for ProviderDocument {
}
}
/// Credential 列表稳定游标。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialCursor {
pub created_at: DateTime<Utc>,
pub account_id: ProviderAccountId,
}
/// Provider credential 列表查询。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialListQuery {
pub credential_state: Option<CredentialStateFilter>,
pub enabled: Option<bool>,
pub window: CredentialListWindow,
}
/// Credential 状态筛选。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CredentialStateFilter {
Exact(CredentialState),
AnyOf(Vec<CredentialState>),
}
impl CredentialStateFilter {
#[must_use]
pub fn matches(&self, credential_state: CredentialState) -> bool {
match self {
Self::Exact(expected) => *expected == credential_state,
Self::AnyOf(expected) => expected.contains(&credential_state),
}
}
}
/// Credential 目录的集合窗口；完整列表与游标分页互斥。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CredentialListWindow {
All,
Page {
cursor: Option<CredentialCursor>,
page_size: PageSize,
},
}
/// Provider credential 列表。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialPage {
pub config_revision: Revision,
pub items: Vec<AccountRecord>,
pub next_cursor: Option<CredentialCursor>,
}
/// Provider credential 详情。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialDetails {
backend/crates/gateway-admin/src/ports/store.rs
+1 -8
View File
@@ -32,8 +32,7 @@ use crate::model::{
},
provider_credentials::{
AuthorizationCommit, CredentialDetails, CredentialImportCommit, CredentialImportResult,
CredentialListQuery, CredentialMutationResult, CredentialPage, CredentialRotationCommit,
ProviderExportCredentialInput,
CredentialMutationResult, CredentialRotationCommit, ProviderExportCredentialInput,
},
settings::{AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RuntimeSettings},
};
@@ -110,12 +109,6 @@ pub trait AccountStore: Send + Sync {
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>>;
async fn list_credentials(
&self,
provider_kind: &gateway_core::routing::ProviderKind,
query: CredentialListQuery,
) -> AdminStoreResult<CredentialPage>;
async fn credential_details(
&self,
provider_kind: &gateway_core::routing::ProviderKind,
backend/crates/gateway-admin/src/use_case/mod.rs
+40 -1
View File
@@ -13,6 +13,8 @@ pub mod settings;
pub mod system;
pub mod xai;
use std::sync::Arc;
use crate::{
model::{
AdminError, AdminErrorKind, MutationContext,
@@ -22,7 +24,7 @@ use crate::{
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
PendingAuthorizationMutation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
StartAuthorization,
ProviderQuotaRequest, StartAuthorization,
},
},
ports::{
@@ -97,6 +99,43 @@ async fn publish_committed(
Ok(())
}
async fn publish_credentials_and_observe_quota(
provider: &Arc<dyn ProviderAdmin>,
snapshot: &dyn SnapshotControl,
revision: crate::model::Revision,
account_ids: &[ProviderAccountId],
request_id: &str,
) -> Result<(), AdminError> {
provider.account_facts_changed(account_ids).await;
publish_committed(snapshot, revision).await?;
// 凭据已经提交；额度只是可重建的观察，不能拖住管理请求或回滚提交结果。
let provider = Arc::clone(provider);
let account_ids = account_ids.to_vec();
let request_id = request_id.to_owned();
tokio::spawn(async move {
for account_id in account_ids {
if let Err(error) = provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await
{
tracing::warn!(
request_id,
provider = %provider.provider_kind().as_str(),
account_id = %account_id.as_str(),
quota_error = ?error.kind(),
"initial quota observation failed"
);
}
}
});
Ok(())
}
async fn required_credential(
accounts: &dyn AccountStore,
provider_kind: &ProviderKind,
backend/crates/gateway-admin/src/use_case/openai.rs
+21 -69
View File
@@ -3,17 +3,16 @@
use std::sync::Arc;
use async_trait::async_trait;
use gateway_core::{account::ProviderAccountId, runtime::SnapshotControl};
use gateway_core::runtime::SnapshotControl;
use crate::{
model::{
AdminError,
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
ImportCredentials, PrepareCredentialImport, PrepareCredentialRotation,
ProviderQuotaRequest, RotateCredential, StartAuthorization,
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
CredentialMutationResult, ImportCredentials, PrepareCredentialImport,
PrepareCredentialRotation, RotateCredential, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -21,18 +20,14 @@ use crate::{
use super::{
commit_authorization, commit_credential_rotation, delete_credentials, map_provider_error,
map_store_error, pending_authorization, publish_committed, required_credential,
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
map_store_error, pending_authorization, publish_committed,
publish_credentials_and_observe_quota, required_credential, validate_authorization_commit,
validate_prepared_import, validate_prepared_rotation,
};
/// OpenAI 固定管理路由消费的服务。
#[async_trait]
pub trait OpenAiService: Send + Sync {
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError>;
async fn details(
&self,
account_id: &ProviderAccountId,
) -> Result<CredentialDetails, AdminError>;
async fn import_document(
&self,
command: ImportCredentials,
@@ -77,55 +72,10 @@ impl DefaultOpenAiService {
snapshot,
}
}
fn observe_initial_quotas(&self, account_ids: &[ProviderAccountId], request_id: &str) {
let provider = Arc::clone(&self.provider);
let account_ids = account_ids.to_vec();
let request_id = request_id.to_owned();
tokio::spawn(async move {
for account_id in account_ids {
if let Err(error) = 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]
impl OpenAiService for DefaultOpenAiService {
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError> {
self.accounts
.list_credentials(self.provider.provider_kind(), query)
.await
.map_err(|error| map_store_error(error, "OpenAI credential"))
}
async fn details(
&self,
account_id: &ProviderAccountId,
) -> Result<CredentialDetails, AdminError> {
required_credential(
self.accounts.as_ref(),
self.provider.provider_kind(),
account_id,
"OpenAI credential",
)
.await
}
async fn import_document(
&self,
command: ImportCredentials,
@@ -167,12 +117,14 @@ impl OpenAiService for DefaultOpenAiService {
.await
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
drop(proxy_reservation);
self.provider
.account_facts_changed(&result.credential_ids)
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
// 导入完成后立即做一次观察；失败不回滚已经提交的 credential。
self.observe_initial_quotas(&result.credential_ids, &context.request_id);
publish_credentials_and_observe_quota(
&self.provider,
self.snapshot.as_ref(),
result.config_revision,
&result.credential_ids,
&context.request_id,
)
.await?;
Ok(result)
}
@@ -220,14 +172,14 @@ impl OpenAiService for DefaultOpenAiService {
"OpenAI authorization",
)
.await?;
self.provider
.account_facts_changed(std::slice::from_ref(&result.account_id))
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
self.observe_initial_quotas(
publish_credentials_and_observe_quota(
&self.provider,
self.snapshot.as_ref(),
result.config_revision,
std::slice::from_ref(&result.account_id),
&context.request_id,
);
)
.await?;
Ok(result)
}
backend/crates/gateway-admin/src/use_case/xai.rs
+20 -25
View File
@@ -11,8 +11,8 @@ use crate::{
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
CredentialListQuery, CredentialMutationResult, CredentialPage, ImportCredentials,
PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
CredentialMutationResult, ImportCredentials, PrepareCredentialImport,
StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -20,14 +20,13 @@ use crate::{
use super::{
commit_authorization, delete_credentials, map_provider_error, map_store_error,
pending_authorization, publish_committed, validate_authorization_commit,
validate_prepared_import,
pending_authorization, publish_committed, publish_credentials_and_observe_quota,
validate_authorization_commit, validate_prepared_import,
};
/// xAI 固定管理路由消费的服务。
#[async_trait]
pub trait XaiService: Send + Sync {
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError>;
async fn import_document(
&self,
command: ImportCredentials,
@@ -72,13 +71,6 @@ impl DefaultXaiService {
#[async_trait]
impl XaiService for DefaultXaiService {
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError> {
self.accounts
.list_credentials(self.provider.provider_kind(), query)
.await
.map_err(|error| map_store_error(error, "xAI credential"))
}
async fn import_document(
&self,
command: ImportCredentials,
@@ -120,18 +112,14 @@ impl XaiService for DefaultXaiService {
.await
.map_err(|error| map_store_error(error, "xAI credential import"))?;
drop(proxy_reservation);
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;
}
publish_credentials_and_observe_quota(
&self.provider,
self.snapshot.as_ref(),
result.config_revision,
&result.credential_ids,
&context.request_id,
)
.await?;
Ok(result)
}
@@ -179,7 +167,14 @@ impl XaiService for DefaultXaiService {
"xAI authorization",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
publish_credentials_and_observe_quota(
&self.provider,
self.snapshot.as_ref(),
result.config_revision,
std::slice::from_ref(&result.account_id),
&context.request_id,
)
.await?;
Ok(result)
}
backend/crates/gateway-admin/tests/use_case/accounts.rs
+32 -21
View File
@@ -33,10 +33,9 @@ use gateway_admin::{
AuthorizationCommit, AuthorizationCommitGuard, AuthorizationCredentialCommit,
AuthorizationMutationTarget, AuthorizationStarted, CompleteAuthorization,
ConsumeProviderResetCredit, CredentialCommitGuard, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit,
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
@@ -71,6 +70,8 @@ pub(super) struct FakeProviderAdmin {
export_inputs: Mutex<Vec<ProviderExportCredentialInput>>,
import_account_ids: Mutex<Vec<String>>,
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
quota_started: tokio::sync::Notify,
quota_gate: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
quota: Mutex<ProviderQuota>,
quota_refresh_account: Mutex<Option<(Arc<FakeAccountStore>, AccountRecord)>>,
current_credential_revision: Mutex<Revision>,
@@ -89,6 +90,8 @@ impl FakeProviderAdmin {
export_inputs: Mutex::new(Vec::new()),
import_account_ids: Mutex::new(vec!["acct_prepared".to_owned()]),
quota_requests: Mutex::new(Vec::new()),
quota_started: tokio::sync::Notify::new(),
quota_gate: Mutex::new(None),
quota: Mutex::new(empty_quota()),
quota_refresh_account: Mutex::new(None),
current_credential_revision: Mutex::new(revision(1)),
@@ -138,6 +141,26 @@ impl FakeProviderAdmin {
.clone()
}
pub(super) fn pause_next_quota(&self) -> tokio::sync::oneshot::Sender<()> {
let (sender, receiver) = tokio::sync::oneshot::channel();
*self.quota_gate.lock().expect("quota gate") = Some(receiver);
sender
}
pub(super) async fn wait_for_quota_requests(&self, expected: usize) {
tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
let notified = self.quota_started.notified();
if self.quota_requests().len() >= expected {
return;
}
notified.await;
}
})
.await
.expect("quota observation started");
}
fn reset_credit_commands(&self) -> Vec<ConsumeProviderResetCredit> {
self.reset_credit_commands
.lock()
@@ -377,6 +400,11 @@ impl ProviderAdmin for FakeProviderAdmin {
.lock()
.expect("provider quota requests")
.push(request);
self.quota_started.notify_one();
let gate = self.quota_gate.lock().expect("quota gate").take();
if let Some(gate) = gate {
gate.await.expect("release paused quota");
}
if let Some(kind) = self
.quota_failure
.lock()
@@ -612,23 +640,6 @@ impl AccountStore for FakeAccountStore {
.clone())
}
async fn list_credentials(
&self,
provider_kind: &ProviderKind,
_: CredentialListQuery,
) -> AdminStoreResult<CredentialPage> {
self.record("store.list_credentials");
let accounts = self.accounts.lock().expect("accounts").clone();
Ok(CredentialPage {
config_revision: revision(1),
items: accounts
.into_iter()
.filter(|account| &account.provider_kind == provider_kind)
.collect(),
next_cursor: None,
})
}
async fn credential_details(
&self,
provider_kind: &ProviderKind,
backend/crates/gateway-admin/tests/use_case/mod.rs
+5 -14
View File
@@ -45,12 +45,11 @@ use gateway_admin::{
},
provider_credentials::{
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput,
ProviderModels, ProviderQuota,
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
PreparedCredentialImport, PreparedCredentialRotation, ProviderExport,
ProviderExportCredentialInput, ProviderModels, ProviderQuota,
},
settings::{AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RuntimeSettings},
system::{SystemOperationAccepted, SystemUpdateDetail, SystemUpdateStatus, SystemVersion},
@@ -372,14 +371,6 @@ impl AccountStore for UnavailableStore {
Err(unavailable("account quota window usage"))
}
async fn list_credentials(
&self,
_: &ProviderKind,
_: CredentialListQuery,
) -> AdminStoreResult<CredentialPage> {
Err(unavailable("credentials"))
}
async fn credential_details(
&self,
_: &ProviderKind,
backend/crates/gateway-admin/tests/use_case/openai.rs
+7 -17
View File
@@ -58,7 +58,7 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
.await
.expect("import credential");
await_quota_requests(&provider, 1).await;
provider.wait_for_quota_requests(1).await;
assert_eq!(
recorded(&events),
@@ -142,7 +142,7 @@ async fn openai_import_should_refresh_quota_for_every_imported_account() {
.await
.expect("import credentials");
await_quota_requests(&provider, 2).await;
provider.wait_for_quota_requests(2).await;
assert_eq!(
result.credential_ids,
@@ -187,7 +187,7 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
.await
.expect("committed import remains successful");
await_quota_requests(&provider, 1).await;
provider.wait_for_quota_requests(1).await;
assert_eq!(
result.credential_ids,
@@ -224,7 +224,7 @@ async fn openai_authorization_create_should_observe_initial_quota() {
.await
.expect("complete authorization");
await_quota_requests(&provider, 1).await;
provider.wait_for_quota_requests(1).await;
assert_eq!(
recorded(&events),
@@ -280,7 +280,7 @@ async fn openai_authorization_create_should_remain_successful_when_initial_quota
.await
.expect("quota observation must not fail authorization");
await_quota_requests(&provider, 1).await;
provider.wait_for_quota_requests(1).await;
assert_eq!(result.account_id.as_str(), "acct_prepared");
assert_eq!(provider.quota_requests().len(), 1);
@@ -322,7 +322,7 @@ async fn openai_authorization_store_failure_should_release_claim_for_retry() {
.await
.expect("released claim must permit retry");
await_quota_requests(&provider, 1).await;
provider.wait_for_quota_requests(1).await;
assert_eq!(
recorded(&events),
@@ -406,7 +406,7 @@ async fn openai_reauthorization_should_commit_after_credential_revision_advances
.await
.expect("complete reauthorization");
await_quota_requests(&provider, 1).await;
provider.wait_for_quota_requests(1).await;
assert_eq!(result.credential_revision, Some(revision(3)));
@@ -434,16 +434,6 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
.await
}
async fn await_quota_requests(provider: &FakeProviderAdmin, expected: usize) {
for _ in 0..100 {
if provider.quota_requests().len() == expected {
return;
}
tokio::task::yield_now().await;
}
assert_eq!(provider.quota_requests().len(), expected);
}
fn deletion(account_id: &str) -> CredentialDeletion {
CredentialDeletion {
context: context("request-openai"),
backend/crates/gateway-admin/tests/use_case/xai.rs
+116 -1
View File
@@ -4,7 +4,9 @@ use gateway_core::account::ProviderAccountId;
use gateway_admin::{
AdminServices,
model::provider_credentials::{ImportCredentials, ProviderQuotaRequest},
model::provider_credentials::{
CompleteAuthorization, ImportCredentials, ProviderQuotaRequest, StartAuthorization,
},
ports::provider::ProviderAdminErrorKind,
};
@@ -28,11 +30,13 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
.await
.expect("import xAI credential");
provider.wait_for_quota_requests(1).await;
assert_eq!(
recorded(&events),
[
"provider.prepare_import",
"store.commit_import",
"provider.account_facts_changed",
"provider.quota"
]
);
@@ -70,6 +74,7 @@ async fn xai_import_should_refresh_quota_for_every_imported_account() {
.await
.expect("import credentials");
provider.wait_for_quota_requests(2).await;
assert_eq!(
provider.quota_requests(),
[
@@ -106,10 +111,120 @@ async fn xai_import_should_remain_successful_when_quota_refresh_fails() {
.await
.expect("committed import remains successful");
provider.wait_for_quota_requests(1).await;
assert_eq!(result.credential_ids.len(), 1);
assert_eq!(provider.quota_requests().len(), 1);
}
#[tokio::test]
async fn xai_authorization_should_observe_quota_without_affecting_commit() {
for quota_failure in [None, Some(ProviderAdminErrorKind::Unavailable)] {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
if let Some(kind) = quota_failure {
provider.fail_next_quota(kind);
}
let store = FakeAccountStore::new("xai", events.clone());
let services = service(provider.clone(), store.clone()).await;
services
.xai()
.start_authorization(StartAuthorization {
outbound_proxy: None,
context: context("oauth-xai-create"),
name: "new xAI credential".to_owned(),
reauthorization: None,
})
.await
.expect("start authorization");
let result = services
.xai()
.complete_authorization(CompleteAuthorization {
settings: None,
context: context("oauth-xai-create"),
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");
provider.wait_for_quota_requests(1).await;
assert_eq!(
recorded(&events),
[
"provider.start_authorization",
"provider.complete_authorization",
"store.commit_authorization",
"provider.account_facts_changed",
"provider.quota",
]
);
assert_eq!(
provider.quota_requests(),
[ProviderQuotaRequest {
account_id: result.account_id,
refresh: true,
rolling_usage: None,
}]
);
assert_eq!(store.audit_requests(), ["oauth-xai-create"]);
}
}
#[tokio::test]
async fn credential_import_should_return_while_initial_quota_is_pending() {
for kind in ["openai", "xai"] {
let events = events();
let provider = FakeProviderAdmin::new(kind, events.clone());
let quota_release = provider.pause_next_quota();
let store = FakeAccountStore::new(kind, events);
let services = service(provider.clone(), store).await;
let command = ImportCredentials {
outbound_proxy_id: None,
settings: None,
context: context("import-slow-quota"),
document: document(),
};
let result = tokio::time::timeout(std::time::Duration::from_secs(5), async {
match kind {
"openai" => services.openai().import_document(command).await,
_ => services.xai().import_document(command).await,
}
})
.await
.expect("committed import must not wait for quota")
.expect("import credentials");
provider.wait_for_quota_requests(1).await;
assert_eq!(result.credential_ids.len(), 1);
quota_release.send(()).expect("finish quota observation");
}
}
#[tokio::test]
async fn xai_import_store_failure_should_not_observe_quota_or_invalidate_provider_facts() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events.clone());
store.fail_next_commit();
let services = service(provider.clone(), store).await;
services
.xai()
.import_document(ImportCredentials {
outbound_proxy_id: None,
settings: None,
context: context("import-xai-store-failure"),
document: document(),
})
.await
.expect_err("store commit fails");
assert_eq!(
recorded(&events),
["provider.prepare_import", "store.commit_import"]
);
assert!(provider.quota_requests().is_empty());
}
async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>) -> AdminServices {
super::AdminHarness::new()
.provider(provider)
backend/crates/gateway-api/tests/admin/mod.rs
+5 -14
View File
@@ -40,12 +40,11 @@ use gateway_admin::{
},
provider_credentials::{
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput,
ProviderModels, ProviderQuota,
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
PreparedCredentialImport, PreparedCredentialRotation, ProviderExport,
ProviderExportCredentialInput, ProviderModels, ProviderQuota,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ModelMappings, ReplaceRuntimeSettings,
@@ -738,14 +737,6 @@ impl AccountStore for UnusedStore {
Err(unavailable("account quota window usage"))
}
async fn list_credentials(
&self,
_: &ProviderKind,
_: CredentialListQuery,
) -> AdminStoreResult<CredentialPage> {
Err(unavailable("credential list"))
}
async fn credential_details(
&self,
_: &ProviderKind,
backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs
+378 -6
View File
@@ -7,6 +7,7 @@ use std::sync::{
};
use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use axum::{
body::Body,
http::{HeaderMap, HeaderValue, Request, StatusCode, header::AUTHORIZATION},
@@ -17,18 +18,40 @@ use gateway_api::openai::responses::{
DecodedResponsesRequest, OpenAiRequestHeaders, ResponseCreateFrameError,
decode_response_create_with_context,
};
use gateway_core::account::ProviderAccountId;
use gateway_core::engine::admission::{
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
ClientAdmissionRequest, ClientAdmissionRestoreResult,
};
use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
AuthenticatedClient, ClientAuthenticationError, DefaultExecutionService, ExecutionService,
ExecutionSession, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind};
use gateway_core::engine::provider::{
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
};
use gateway_core::engine::{
AttemptContext, AttemptRecord, CommitRequirement, CoordinatedEvent, EngineError,
ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
RecoveryReport,
};
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind, StoreError};
use gateway_core::event::{
GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta,
};
use gateway_core::lifecycle::{
CancellationToken, ConnectionDraining, ConnectionGuard, ConnectionLifecycle,
};
use gateway_core::metering::{Decimal, ProviderReportedCost};
use gateway_core::operation::Operation;
use gateway_core::routing::PublicModelId;
use gateway_core::upstream::UpstreamSendState;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ProviderCatalogGeneration, ProviderKind, ProviderModelCapabilities, PublicModelId,
};
use gateway_core::runtime::RuntimeSnapshotHandle;
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use gateway_protocol::openai::codex_responses_request_semantics;
use serde_json::{Value, json};
use tokio_tungstenite::connect_async;
@@ -37,7 +60,10 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tower::ServiceExt;
use super::decode_response_create;
use crate::openai::{api_router, authenticated_client, models::ModelsExecution};
use crate::openai::{
EmptyWorkerHealth, IgnoredClientApiKeyUsage, UnusedContinuation, api_router,
authenticated_client, models::ModelsExecution, snapshot,
};
fn decode_response_create_with_turn_header(
payload: Value,
@@ -811,6 +837,352 @@ async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit(
const TEST_WEBSOCKET_KEY: &str = "AAAAAAAAAAAAAAAAAAAAAA==";
#[derive(Default)]
struct SettlementPorts {
request_id: Mutex<Option<ModelRequestId>>,
finalizations: Mutex<Vec<ModelRequestFinalization>>,
active: AtomicBool,
settlement_calls: AtomicUsize,
settlement_started: tokio::sync::Notify,
settlement_gate: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
charges: Mutex<Vec<ClientBudgetCharge>>,
releases: AtomicUsize,
released: tokio::sync::Notify,
}
#[async_trait]
impl ExecutionStore for SettlementPorts {
async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> {
*self.request_id.lock().unwrap() = Some(request.id);
Ok(())
}
async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> {
Ok(())
}
async fn mark_send_state(
&self,
_: &ModelRequestId,
_: UpstreamSendState,
) -> Result<(), StoreError> {
Ok(())
}
async fn mark_downstream_committed(
&self,
_: &ModelRequestId,
_: SystemTime,
_: Option<u16>,
) -> Result<(), StoreError> {
Ok(())
}
async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> {
Ok(())
}
async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> {
Ok(())
}
async fn finalize_model_request(
&self,
finalization: ModelRequestFinalization,
) -> Result<(), StoreError> {
self.finalizations.lock().unwrap().push(finalization);
Ok(())
}
async fn recover_expired(&self, _: SystemTime) -> Result<RecoveryReport, StoreError> {
unreachable!("the WebSocket test does not recover requests")
}
}
impl ClientAdmissionPort for SettlementPorts {
fn admit(
&self,
_: ClientAdmissionRequest,
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
Box::pin(async {
assert!(!self.active.swap(true, Ordering::SeqCst));
Ok(ClientAdmissionDecision::Granted)
})
}
fn release<'a>(
&'a self,
_: &'a ClientApiKeyId,
request_id: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async move {
assert_eq!(
self.charges.lock().unwrap().len(),
1,
"settlement must finish first"
);
assert_eq!(self.request_id.lock().unwrap().as_ref(), Some(request_id));
self.releases.fetch_add(1, Ordering::SeqCst);
let active = self.active.swap(false, Ordering::SeqCst);
self.released.notify_one();
Ok(active)
})
}
fn restore(
&self,
_: ClientAdmissionRecovery,
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
Box::pin(async { unreachable!("the WebSocket test does not restore admission") })
}
}
impl ClientBudgetPort for SettlementPorts {
fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
Box::pin(async { Ok(()) })
}
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
Box::pin(async move {
self.settlement_calls.fetch_add(1, Ordering::SeqCst);
assert!(self.active.load(Ordering::SeqCst));
let gate = self
.settlement_gate
.lock()
.unwrap()
.take()
.expect("one settlement");
self.settlement_started.notify_one();
gate.await
.expect("release settlement after client disconnect");
assert!(self.active.load(Ordering::SeqCst));
self.charges.lock().unwrap().push(charge);
Ok(())
})
}
}
impl ProviderCircuitPort for SettlementPorts {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct ChargedWebSocketProvider;
fn websocket_charge() -> Decimal {
"1.25".parse().unwrap()
}
#[async_trait]
impl Provider for ChargedWebSocketProvider {
fn name(&self) -> &'static str {
"openai"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::default()
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
async fn execute(
&self,
request: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().unwrap().clone(),
ProviderAccountId::new("acct_api_test").unwrap(),
UpstreamTransport::new("websocket").unwrap(),
);
let response = ResponseMeta::new("resp_ws_settlement", "model-a");
let events = [
ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(response.clone())],
ProtocolWireEvent::json("openai", Some("response.created".to_owned()), json!({
"type": "response.created",
"response": {"id": "resp_ws_settlement", "model": "model-a", "status": "in_progress"}
})).unwrap(),
),
ProviderEvent::canonical_with_wire(
vec![
GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(websocket_charge().scaled()).unwrap(),
),
GatewayEvent::Completed(response),
],
ProtocolWireEvent::json("openai", Some("response.completed".to_owned()), json!({
"type": "response.completed",
"response": {"id": "resp_ws_settlement", "model": "model-a", "status": "completed"}
})).unwrap(),
),
];
Ok(ProviderStream::new(
metadata,
futures::stream::iter(events.map(Ok)),
(),
))
}
}
#[derive(Default)]
struct SettlementConnectionLifecycle {
closed: Arc<tokio::sync::Notify>,
}
struct SettlementConnectionGuard(Arc<tokio::sync::Notify>);
impl ConnectionGuard for SettlementConnectionGuard {}
impl Drop for SettlementConnectionGuard {
fn drop(&mut self) {
self.0.notify_one();
}
}
impl ConnectionLifecycle for SettlementConnectionLifecycle {
fn try_register(&self) -> Result<Box<dyn ConnectionGuard>, ConnectionDraining> {
Ok(Box::new(SettlementConnectionGuard(self.closed.clone())))
}
fn cancellation(&self) -> CancellationToken {
CancellationToken::new()
}
fn is_draining(&self) -> bool {
false
}
}
#[tokio::test]
async fn websocket_disconnect_during_core_settlement_finishes_charge_before_releasing_admission() {
let (release_settlement, gate) = tokio::sync::oneshot::channel();
let ports = Arc::new(SettlementPorts {
settlement_gate: Mutex::new(Some(gate)),
..Default::default()
});
let execution = Arc::new(
DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot("sk_ws_settlement", "openai")),
ports.clone(),
ProviderRegistry::new([Arc::new(ChargedWebSocketProvider) as Arc<dyn Provider>])
.unwrap(),
ports.clone(),
ports.clone(),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
)
.with_budget(ports.clone()),
);
let lifecycle = Arc::new(SettlementConnectionLifecycle::default());
let admin = crate::admin::AdminTestFixture::new().await;
let app = gateway_api::initialize(
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: Vec::new(),
request_timeout_seconds: None,
request_id_header: "x-request-id".to_owned(),
},
execution,
admin.services,
Vec::new(),
Arc::new(EmptyWorkerHealth),
lifecycle.clone(),
)
.unwrap()
.router();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.unwrap();
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer sk_ws_settlement"),
);
let (mut socket, response) = connect_async(request).await.unwrap();
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
socket
.send(ClientMessage::Text(
json!({
"type": "response.create", "model": "model-a", "input": "hello"
})
.to_string()
.into(),
))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(5), ports.settlement_started.notified())
.await
.expect("production Core reached settlement");
{
let finalizations = ports.finalizations.lock().unwrap();
assert_eq!(finalizations.len(), 1);
assert_eq!(
finalizations[0].outcome,
gateway_core::engine::ExecutionOutcome::Succeeded,
"unexpected execution error: {:?}",
finalizations[0].error,
);
}
assert!(ports.charges.lock().unwrap().is_empty());
assert_eq!(ports.releases.load(Ordering::SeqCst), 0);
// 等到生产连接 guard 释放，保证 forward 的 select 已因断连取消 next_event；
// 在此之前不能打开结算屏障，否则只会验证正常完成而错过取消窗口。
socket.close(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(5), lifecycle.closed.notified())
.await
.expect("production WebSocket handler exited");
assert!(ports.active.load(Ordering::SeqCst));
release_settlement
.send(())
.expect("the original settlement future survived disconnect");
tokio::time::timeout(Duration::from_secs(5), ports.released.notified())
.await
.expect("detached Core cleanup released admission");
assert!(!ports.active.load(Ordering::SeqCst));
assert_eq!(ports.releases.load(Ordering::SeqCst), 1);
assert_eq!(ports.settlement_calls.load(Ordering::SeqCst), 1);
{
let charges = ports.charges.lock().unwrap();
assert_eq!(charges.len(), 1);
assert_eq!(charges[0].amount_usd, websocket_charge());
assert_eq!(charges[0].key_id.as_str(), "key_api_test");
assert_eq!(
Some(&charges[0].request_id),
ports.request_id.lock().unwrap().as_ref()
);
}
server.abort();
let _ = server.await;
}
fn upgrade_request(authorization: &str) -> Request<Body> {
Request::get("/v1/responses")
.header(AUTHORIZATION, authorization)
backend/crates/gateway-core/src/engine/execution.rs
+37 -72
View File
@@ -173,8 +173,11 @@ pub trait ExecutionSession: Send {
&mut self,
client_status_code: u16,
) -> BoxFuture<'_, Result<(), EngineError>>;
/// 执行已终结且结算、准入释放均已返回，协议层才可以放弃清理责任。
/// 结算失败仍由 Store 保留费用重试，释放失败仍按租约 TTL 收敛。
fn is_finalized(&self) -> bool;
fn cancel(&self);
/// 将会话交给宿主持续驱动清理；取消请求事件的等待不会丢弃已开始的结算。
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()>;
}
@@ -978,8 +981,9 @@ impl AdmissionLease {
}
struct DefaultExecutionSession {
core: Option<ResponseExecutionSession<dyn ExecutionStore>>,
core: ResponseExecutionSession<dyn ExecutionStore>,
admission: Option<AdmissionLease>,
cleanup: Option<BoxFuture<'static, ()>>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
observed_provider_outcomes: usize,
@@ -996,8 +1000,9 @@ impl DefaultExecutionSession {
budget: Option<Arc<dyn ClientBudgetPort>>,
) -> Self {
Self {
core: Some(core),
core,
admission: Some(admission),
cleanup: None,
circuits,
continuation,
observed_provider_outcomes: 0,
@@ -1006,33 +1011,29 @@ impl DefaultExecutionSession {
}
}
fn core_mut(
&mut self,
) -> Result<&mut ResponseExecutionSession<dyn ExecutionStore>, EngineError> {
self.core.as_mut().ok_or(EngineError::InvalidDeliveryState)
}
async fn settle_if_finalized(&mut self) {
if self
.core
.as_ref()
.is_some_and(ResponseExecutionSession::is_finalized)
if self.core.is_finalized()
&& let Some(admission) = self.admission.take()
{
if let Some(budget) = self.budget.take()
&& let Some(core) = &self.core
{
settle_budget(budget.as_ref(), core.budget_charge()).await;
}
admission.release().await;
let budget = self.budget.take();
let charge = self.core.budget_charge();
// 在首次 await 前把完整清理责任留在会话内。事件等待被取消后，后续 poll
// 或 detach 继续同一个 future，既不丢失费用，也不重启已完成的结算。
self.cleanup = Some(Box::pin(async move {
if let Some(budget) = budget {
settle_budget(budget.as_ref(), charge).await;
}
admission.release().await;
}));
}
if let Some(cleanup) = self.cleanup.as_mut() {
cleanup.await;
self.cleanup = None;
}
}
async fn observe_provider_outcomes(&mut self) {
let Some(core) = self.core.as_ref() else {
return;
};
let outcomes = core.provider_attempt_outcomes();
let outcomes = self.core.provider_attempt_outcomes();
let new_outcomes = outcomes
.get(self.observed_provider_outcomes..)
.unwrap_or_default()
@@ -1048,11 +1049,7 @@ impl DefaultExecutionSession {
let Some(state) = state else {
return;
};
let Some(pin) = self
.core
.as_ref()
.and_then(|core| core.native_continuation_pin(state))
else {
let Some(pin) = self.core.native_continuation_pin(state) else {
return;
};
self.continuation_recorded = true;
@@ -1060,47 +1057,27 @@ impl DefaultExecutionSession {
}
async fn finalize_detached(&mut self) {
let Some(mut core) = self.core.take() else {
return;
};
core.cancel();
if !core.is_finalized()
&& let Err(error) = core.cancel_and_finalize().await
{
if let Err(error) = self.core.cancel_and_finalize().await {
tracing::warn!(%error, "Detached execution 终态收敛失败");
}
let pending = core
.provider_attempt_outcomes()
.get(self.observed_provider_outcomes..)
.unwrap_or_default();
publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await;
if let Some(budget) = self.budget.take() {
settle_budget(budget.as_ref(), core.budget_charge()).await;
}
if let Some(admission) = self.admission.take() {
admission.release().await;
}
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
}
}
impl Drop for DefaultExecutionSession {
fn drop(&mut self) {
if let Some(core) = &self.core {
core.cancel();
}
self.core.cancel();
}
}
impl ExecutionSession for DefaultExecutionSession {
fn trace(&self) -> crate::diagnostics::TraceContext {
self.core
.as_ref()
.map(ResponseExecutionSession::trace)
.unwrap_or_default()
self.core.trace()
}
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
let result = self.core_mut()?.next_event().await;
let result = self.core.next_event().await;
if let Ok(Some(event)) = result.as_ref() {
self.record_continuation(event.session_update()).await;
}
@@ -1112,7 +1089,7 @@ impl ExecutionSession for DefaultExecutionSession {
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
Box::pin(async move {
let result = self.core_mut()?.collect_uncommitted().await;
let result = self.core.collect_uncommitted().await;
if let Ok(events) = result.as_ref() {
let state = events.iter().find_map(ProviderEvent::session_update);
self.record_continuation(state).await;
@@ -1124,16 +1101,11 @@ impl ExecutionSession for DefaultExecutionSession {
}
fn response_headers(&self) -> &[ProviderResponseHeader] {
self.core
.as_ref()
.map(ResponseExecutionSession::response_headers)
.unwrap_or_default()
self.core.response_headers()
}
fn response_status_code(&self) -> Option<u16> {
self.core
.as_ref()
.and_then(ResponseExecutionSession::response_status_code)
self.core.response_status_code()
}
fn commit_downstream(
@@ -1141,7 +1113,7 @@ impl ExecutionSession for DefaultExecutionSession {
client_status_code: Option<u16>,
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
let result = self.core_mut()?.commit_downstream(client_status_code).await;
let result = self.core.commit_downstream(client_status_code).await;
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
result
@@ -1153,10 +1125,7 @@ impl ExecutionSession for DefaultExecutionSession {
client_status_code: u16,
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
let result = self
.core_mut()?
.record_client_status(client_status_code)
.await;
let result = self.core.record_client_status(client_status_code).await;
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
result
@@ -1164,15 +1133,11 @@ impl ExecutionSession for DefaultExecutionSession {
}
fn is_finalized(&self) -> bool {
self.core
.as_ref()
.is_none_or(ResponseExecutionSession::is_finalized)
self.core.is_finalized() && self.admission.is_none() && self.cleanup.is_none()
}
fn cancel(&self) {
if let Some(core) = &self.core {
core.cancel();
}
self.core.cancel();
}
fn detach_finalize(mut self: Box<Self>) -> BoxFuture<'static, ()> {
backend/crates/gateway-core/src/engine/provider.rs
+26 -103
View File
@@ -8,7 +8,7 @@ use std::task::{Context, Poll};
use std::time::Instant;
use async_trait::async_trait;
use futures::Stream;
use futures::{Stream, future::BoxFuture};
use thiserror::Error;
use crate::account::{
@@ -21,7 +21,8 @@ use crate::identity::ProviderKind;
use crate::operation::Operation;
use crate::policy::ClientApiKeyId;
use crate::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, PublicModelId, UpstreamModelId,
ProviderCandidate, ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
ProviderModelCapabilities, PublicModelId, UpstreamModelId,
};
use crate::upstream::OpaqueUpstreamValue;
use crate::upstream::{UpstreamSendState, UpstreamTransport};
@@ -437,64 +438,6 @@ pub struct ContinuationRequestObservation {
pub requested: bool,
}
/// Provider 实时目录编译后的单模型能力。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderModelCapabilities {
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
presentation: Option<ModelPresentation>,
}
/// Provider 实时目录成功发布后的进程内单调代次。
///
/// 代次只表达“目录内容已经变化”，不承载模型、ETag 或 Provider 私有数据。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderCatalogGeneration(u64);
impl ProviderCatalogGeneration {
#[must_use]
pub const fn new(value: u64) -> Self {
Self(value)
}
#[must_use]
pub const fn get(self) -> u64 {
self.0
}
}
impl ProviderModelCapabilities {
#[must_use]
pub const fn new(upstream_model: UpstreamModelId, capabilities: ModelCapabilities) -> Self {
Self {
upstream_model,
capabilities,
presentation: None,
}
}
#[must_use]
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
self.presentation = Some(presentation);
self
}
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
}
#[must_use]
pub const fn capabilities(&self) -> &ModelCapabilities {
&self.capabilities
}
#[must_use]
pub const fn presentation(&self) -> Option<&ModelPresentation> {
self.presentation.as_ref()
}
}
impl ProviderRequest {
/// 绑定 operation 与请求计划中冻结的 Provider 候选。
#[must_use]
@@ -586,15 +529,6 @@ pub enum RegistryError {
},
}
/// Registry 实时目录查询错误。
#[derive(Debug, Error)]
pub enum ProviderCatalogError {
#[error("provider `{provider}` is not registered")]
NotRegistered { provider: String },
#[error("provider model catalog query failed")]
Query(#[source] ProviderError),
}
/// 唯一保存 `Arc<dyn Provider>` 的异构注册表。
#[derive(Default)]
pub struct ProviderRegistryBuilder {
@@ -665,11 +599,6 @@ impl ProviderRegistry {
self.providers.get(provider)
}
/// 返回全部已注册 Provider。
pub fn provider_kinds(&self) -> impl Iterator<Item = &ProviderKind> {
self.providers.keys()
}
#[must_use]
pub fn request_observation(
&self,
@@ -684,35 +613,6 @@ impl ProviderRegistry {
})
}
/// 读取全部 Provider 的目录代次，用于 Core 快照稳定性校验与对账。
#[must_use]
pub fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
self.providers
.iter()
.map(|(kind, provider)| (kind.clone(), provider.catalog_generation()))
.collect()
}
/// 通过编译期 Provider adapter 查询并编译实时能力目录。
///
/// # Errors
///
/// Provider 未注册或目录查询失败时返回错误。
pub async fn query_model_capabilities(
&self,
provider_kind: &ProviderKind,
) -> Result<Vec<ProviderModelCapabilities>, ProviderCatalogError> {
let provider = self.providers.get(provider_kind).ok_or_else(|| {
ProviderCatalogError::NotRegistered {
provider: provider_kind.as_str().to_owned(),
}
})?;
provider
.query_model_capabilities()
.await
.map_err(ProviderCatalogError::Query)
}
/// 判断 Provider 是否已注册。
#[must_use]
pub fn contains(&self, provider: &ProviderKind) -> bool {
@@ -731,3 +631,26 @@ impl ProviderRegistry {
self.providers.is_empty()
}
}
impl ProviderCatalogPort for ProviderRegistry {
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
self.providers
.iter()
.map(|(kind, provider)| (kind.clone(), provider.catalog_generation()))
.collect()
}
fn query_model_capabilities(
&self,
provider_kind: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
let provider = self.providers.get(provider_kind);
Box::pin(async move {
provider
.ok_or(ProviderCatalogUnavailable)?
.query_model_capabilities()
.await
.map_err(|_| ProviderCatalogUnavailable)
})
}
}
backend/crates/gateway-core/src/lib.rs
+1 -1
View File
@@ -140,7 +140,7 @@ pub async fn initialize(
.map_err(|_| CoreError::AdmissionRecoveryUnavailable)?;
let compiler = Arc::new(RuntimeSnapshotCompiler::new(
Arc::clone(&ports.snapshots),
providers.clone(),
Arc::new(providers.clone()),
));
let initial = compiler
.compile()
backend/crates/gateway-core/src/routing/catalog.rs
+86
View File
@@ -0,0 +1,86 @@
//! 快照编译需要的 Provider 目录事实，不暴露执行或账号选择能力。
use std::collections::BTreeMap;
use futures::future::BoxFuture;
use crate::identity::ProviderKind;
use super::{ModelCapabilities, ModelPresentation, UpstreamModelId};
/// Provider 实时目录编译后的单模型能力。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderModelCapabilities {
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
presentation: Option<ModelPresentation>,
}
impl ProviderModelCapabilities {
#[must_use]
pub const fn new(upstream_model: UpstreamModelId, capabilities: ModelCapabilities) -> Self {
Self {
upstream_model,
capabilities,
presentation: None,
}
}
#[must_use]
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
self.presentation = Some(presentation);
self
}
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
}
#[must_use]
pub const fn capabilities(&self) -> &ModelCapabilities {
&self.capabilities
}
#[must_use]
pub const fn presentation(&self) -> Option<&ModelPresentation> {
self.presentation.as_ref()
}
}
/// Provider 实时目录成功发布后的进程内单调代次。
///
/// 代次只表达“目录内容已经变化”，不承载模型、ETag 或 Provider 私有数据。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderCatalogGeneration(u64);
impl ProviderCatalogGeneration {
#[must_use]
pub const fn new(value: u64) -> Self {
Self(value)
}
#[must_use]
pub const fn get(self) -> u64 {
self.0
}
}
/// 目录未知，不能据此判定某个模型不存在；不携带执行错误或上游协议细节。
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("provider model catalog is unavailable")]
pub struct ProviderCatalogUnavailable;
/// 快照编译与对账使用的对象安全目录端口。
pub trait ProviderCatalogPort: Send + Sync {
/// 返回全部已注册 Provider 的目录代次；注册集合在初始化后保持不变。
///
/// 即使某个目录暂时不可读，也必须保留它的 Provider 与最近成功发布的代次。
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration>;
/// 查询指定 Provider 的模型事实；成功的空列表表示已知无模型，失败表示未知。
fn query_model_capabilities(
&self,
provider: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>>;
}
backend/crates/gateway-core/src/routing/mod.rs
+5
View File
@@ -1,5 +1,6 @@
//! Provider、模型目录、精确模型映射与请求级候选计划。
mod catalog;
pub mod snapshot;
pub use crate::account::scope::{
@@ -7,6 +8,10 @@ pub use crate::account::scope::{
FrozenAccountScope, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory,
};
pub use crate::identity::ProviderKind;
pub use catalog::{
ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
ProviderModelCapabilities,
};
pub use snapshot::RuntimeSnapshot;
use std::collections::{BTreeMap, BTreeSet};
backend/crates/gateway-core/src/routing/snapshot.rs
+18 -13
View File
@@ -8,7 +8,6 @@ use std::time::Duration;
use futures::future::BoxFuture;
use crate::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
use crate::operation::Operation;
use crate::policy::{
ClientApiKeyId, ClientPolicy, CodexClientMinVersions, CodexClientVersion,
@@ -18,8 +17,9 @@ use crate::validation::RoutingError;
use super::{
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
ProviderCandidate, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, UpstreamModelId,
ProviderCandidate, ProviderCatalogGeneration, ProviderCatalogPort, ProviderKind, ProviderModel,
PublicModelId, RoutingContext, RoutingGroupSnapshot, RoutingPlan, RuntimeAccount,
RuntimeAccountDirectory, UpstreamModelId,
};
const MAXIMUM_CATALOG_STABILITY_ATTEMPTS: usize = 4;
@@ -212,13 +212,16 @@ pub enum RuntimeSnapshotCompileError {
#[derive(Clone)]
pub struct RuntimeSnapshotCompiler {
store: Arc<dyn SnapshotStorePort>,
providers: ProviderRegistry,
catalogs: Arc<dyn ProviderCatalogPort>,
}
impl RuntimeSnapshotCompiler {
#[must_use]
pub const fn new(store: Arc<dyn SnapshotStorePort>, providers: ProviderRegistry) -> Self {
Self { store, providers }
pub const fn new(
store: Arc<dyn SnapshotStorePort>,
catalogs: Arc<dyn ProviderCatalogPort>,
) -> Self {
Self { store, catalogs }
}
pub(crate) fn store(&self) -> Arc<dyn SnapshotStorePort> {
@@ -228,13 +231,13 @@ impl RuntimeSnapshotCompiler {
pub(crate) fn provider_catalog_generations(
&self,
) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
self.providers.catalog_generations()
self.catalogs.catalog_generations()
}
/// 读取一个 revision，并为已注册 Provider 查询实时模型目录。
pub async fn compile(&self) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS {
let catalog_generations = self.providers.catalog_generations();
let catalog_generations = self.catalogs.catalog_generations();
let facts = self
.store
.load_snapshot_facts()
@@ -243,8 +246,10 @@ impl RuntimeSnapshotCompiler {
if facts.config_revision != facts.observed_current_revision {
return Err(RuntimeSnapshotCompileError::RevisionChanged);
}
let snapshot = compile_runtime_snapshot(facts, &self.providers).await?;
let observed_generations = self.providers.catalog_generations();
let provider_kinds = catalog_generations.keys().cloned().collect();
let snapshot =
compile_runtime_snapshot(facts, self.catalogs.as_ref(), provider_kinds).await?;
let observed_generations = self.catalogs.catalog_generations();
if catalog_generations == observed_generations {
return Ok(snapshot.with_provider_catalog_generations(observed_generations));
}
@@ -255,16 +260,16 @@ impl RuntimeSnapshotCompiler {
async fn compile_runtime_snapshot(
facts: SnapshotFacts,
providers: &ProviderRegistry,
catalogs: &dyn ProviderCatalogPort,
provider_kinds: Vec<ProviderKind>,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
let registered_providers = provider_kinds.iter().cloned().collect::<BTreeSet<_>>();
// 目录查询失败表示未知；查询成功后，即使为空，也必须与“已知缺少模型”区分。
let mut provider_models = Vec::new();
let mut known_provider_catalogs = BTreeSet::new();
for provider in &provider_kinds {
let Ok(models) = providers.query_model_capabilities(provider).await else {
let Ok(models) = catalogs.query_model_capabilities(provider).await else {
continue;
};
known_provider_catalogs.insert(provider.clone());
backend/crates/gateway-core/src/runtime/mod.rs
+57 -45
View File
@@ -6,17 +6,17 @@ use std::sync::{Arc, RwLock};
use std::time::Duration;
use futures::future::BoxFuture;
use futures::lock::Mutex;
use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased};
use futures_timer::Delay;
use crate::engine::provider::ProviderCatalogGeneration;
use crate::health::{HealthProbe, HealthState};
use crate::identity::ProviderKind;
use crate::lifecycle::CancellationToken;
use crate::routing::ConfigRevision;
use crate::routing::snapshot::{
RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotStorePort,
RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler,
};
use crate::routing::{ConfigRevision, ProviderCatalogGeneration};
use crate::task::{
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable,
@@ -135,11 +135,17 @@ pub struct RuntimeSnapshotPublisher {
compiler: Arc<RuntimeSnapshotCompiler>,
snapshots: RuntimeSnapshotHandle,
subscriptions: Arc<dyn SnapshotSubscriptionPort>,
refresh_lock: Arc<Mutex<()>>,
}
enum RefreshMode {
Required,
Reconcile,
}
impl RuntimeSnapshotPublisher {
#[must_use]
pub const fn new(
pub fn new(
compiler: Arc<RuntimeSnapshotCompiler>,
snapshots: RuntimeSnapshotHandle,
subscriptions: Arc<dyn SnapshotSubscriptionPort>,
@@ -148,21 +154,55 @@ impl RuntimeSnapshotPublisher {
compiler,
snapshots,
subscriptions,
refresh_lock: Arc::new(Mutex::new(())),
}
}
/// 重新编译并原子替换本进程快照。
/// 串行重编译并替换本进程快照；无法确认配置时暂停新请求。
pub async fn refresh(&self) -> Result<ConfigRevision, RuntimeSnapshotCompileError> {
let snapshot = self.compiler.compile().await?;
self.refresh_with_mode(RefreshMode::Required).await
}
async fn refresh_with_mode(
&self,
mode: RefreshMode,
) -> Result<ConfigRevision, RuntimeSnapshotCompileError> {
// 所有入口共享从读取事实到发布或暂停的完整临界区；只锁最后的替换会让
// 旧编译覆盖新授权，或让晚到的失败暂停新快照。请求读取不等待此锁。
let _refresh = self.refresh_lock.lock().await;
let configuration_changed = if matches!(mode, RefreshMode::Reconcile) {
let persisted_revision = self
.compiler
.store()
.current_config_revision()
.await
.map_err(|_| {
self.snapshots.suspend();
RuntimeSnapshotCompileError::StoreUnavailable
})?;
let changed = runtime_revision_needs_refresh(
self.published_revision().map(ConfigRevision::get),
persisted_revision.get(),
);
if !changed && !self.provider_catalogs_need_refresh() {
return Ok(persisted_revision);
}
changed
} else {
true
};
let snapshot = self.compiler.compile().await.inspect_err(|_| {
// 仅目录代次变化时保留旧快照，下一周期继续对账；配置缺失、变化
// 或持久 revision 回退均须 fail closed，不按 revision 大小丢弃刷新。
if configuration_changed {
self.snapshots.suspend();
}
})?;
let revision = snapshot.revision();
self.snapshots.publish(snapshot);
Ok(revision)
}
pub fn suspend(&self) {
self.snapshots.suspend();
}
#[must_use]
pub fn published_revision(&self) -> Option<ConfigRevision> {
self.snapshots.revision()
@@ -176,9 +216,7 @@ impl RuntimeSnapshotPublisher {
/// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚。
async fn publish_committed_inner(&self, committed_revision: ConfigRevision) {
if self.refresh().await.is_err() {
self.suspend();
}
let _ = self.refresh().await;
let _ = self
.subscriptions
.publish_snapshot_revision(committed_revision)
@@ -204,7 +242,6 @@ impl RuntimeSnapshotPublisher {
schedule,
lease: None,
task: Box::new(RuntimeSnapshotReconciliationTask {
store: self.compiler.store(),
publisher: self.clone(),
}),
},
@@ -238,7 +275,6 @@ impl SnapshotControl for RuntimeSnapshotPublisher {
}
struct RuntimeSnapshotReconciliationTask {
store: Arc<dyn SnapshotStorePort>,
publisher: RuntimeSnapshotPublisher,
}
@@ -248,33 +284,11 @@ impl ScheduledTask for RuntimeSnapshotReconciliationTask {
_context: WorkerCycleContext,
) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let persisted_revision = match self.store.current_config_revision().await {
Ok(revision) => revision,
Err(_) => {
self.publisher.suspend();
return Err(WorkerTaskError::safe(
"runtime snapshot revision is unavailable",
));
}
};
let configuration_changed = runtime_revision_needs_refresh(
self.publisher.published_revision().map(ConfigRevision::get),
persisted_revision.get(),
);
if !configuration_changed && !self.publisher.provider_catalogs_need_refresh() {
return Ok(());
}
if self.publisher.refresh().await.is_err() {
// 已提交的配置无法确认时必须 fail closed；仅目录重编译失败时继续
// 服务旧的不可变快照，单调代次会让下一周期再次尝试。
if configuration_changed {
self.publisher.suspend();
}
return Err(WorkerTaskError::safe(
"runtime snapshot reconciliation failed",
));
}
Ok(())
self.publisher
.refresh_with_mode(RefreshMode::Reconcile)
.await
.map(|_| ())
.map_err(|_| WorkerTaskError::safe("runtime snapshot reconciliation failed"))
})
}
}
@@ -314,9 +328,7 @@ impl DaemonTask for RuntimeSnapshotSubscriptionTask {
};
match notified {
Some(Ok(_)) => {
if self.publisher.refresh().await.is_err() {
self.publisher.suspend();
}
let _ = self.publisher.refresh().await;
}
Some(Err(_)) | None => break,
}
backend/crates/gateway-core/tests/engine/coordinator.rs
+4 -4
View File
@@ -16,8 +16,7 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::execution::gateway_error_from_engine;
use gateway_core::engine::provider::{
Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities,
ProviderRegistry, ProviderRequest, ProviderStream,
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
};
use gateway_core::engine::{
AttemptContext, AttemptCoordinator, AttemptRecord, AttemptTransport, CommitRequirement,
@@ -41,8 +40,9 @@ use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, Provi
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
AccountRoutingSnapshot, ClientRoutingScope, ConfigRevision, FrozenAccountScope,
ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
ModelCapabilities, ProviderCatalogGeneration, ProviderKind, ProviderModel,
ProviderModelCapabilities, PublicModelId, RoutingContext, RoutingPlan, RuntimeAccount,
RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use serde_json::{Map, Value, json};
backend/crates/gateway-core/tests/engine/execution.rs
+368 -13
View File
@@ -2,7 +2,7 @@ use std::{
collections::{BTreeMap, BTreeSet},
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
atomic::{AtomicBool, AtomicUsize, Ordering},
},
time::{Duration, Instant, SystemTime},
};
@@ -14,6 +14,8 @@ use gateway_core::error::GatewayError;
struct Admissions {
active: Arc<AtomicBool>,
limits: Mutex<Vec<RateLimits>>,
release_gate: Mutex<Option<oneshot::Receiver<()>>>,
releases: AtomicUsize,
}
impl ClientAdmissionPort for Admissions {
@@ -32,7 +34,14 @@ impl ClientAdmissionPort for Admissions {
_: &'a ClientApiKeyId,
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) })
Box::pin(async {
self.releases.fetch_add(1, Ordering::SeqCst);
let gate = self.release_gate.lock().unwrap().take();
if let Some(gate) = gate {
gate.await.expect("release gate");
}
Ok(self.active.swap(false, Ordering::SeqCst))
})
}
fn restore(
&self,
@@ -42,10 +51,14 @@ impl ClientAdmissionPort for Admissions {
}
}
#[derive(Default)]
struct Budget {
reject: bool,
active: Arc<AtomicBool>,
charges: Mutex<Vec<ClientBudgetCharge>>,
settlement_gate: Mutex<Option<oneshot::Receiver<()>>>,
settlements: AtomicUsize,
fail_settlement: bool,
}
impl ClientBudgetPort for Budget {
@@ -64,12 +77,22 @@ impl ClientBudgetPort for Budget {
}
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
Box::pin(async move {
self.settlements.fetch_add(1, Ordering::SeqCst);
assert!(
self.active.load(Ordering::SeqCst),
"settle before releasing the concurrent request slot"
);
let gate = self.settlement_gate.lock().unwrap().take();
if let Some(gate) = gate {
gate.await.expect("settlement gate");
}
assert!(self.active.load(Ordering::SeqCst));
self.charges.lock().unwrap().push(charge);
Ok(())
if self.fail_settlement {
Err(ClientBudgetError)
} else {
Ok(())
}
})
}
}
@@ -111,6 +134,7 @@ fn budget_rejection_releases_client_concurrency_without_creating_a_charge() {
reject: true,
active: admissions.active.clone(),
charges: Mutex::default(),
..Default::default()
});
let service = service(admissions.clone(), budget.clone());
for transport in [
@@ -200,6 +224,7 @@ fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all
reject: false,
active: admissions.active.clone(),
charges: Mutex::default(),
..Default::default()
});
let service = service(admissions.clone(), budget.clone());
for transport in [
@@ -232,9 +257,336 @@ fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all
);
}
struct ChargedProvider {
fail: bool,
}
fn known_charge() -> Decimal {
"1.25".parse().unwrap()
}
#[async_trait]
impl Provider for ChargedProvider {
fn name(&self) -> &'static str {
"openai"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::default()
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
async fn execute(
&self,
request: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().unwrap().clone(),
ProviderAccountId::new("acct_start").unwrap(),
UpstreamTransport::new("websocket").unwrap(),
);
let response = ResponseMeta::new("resp_cleanup", "gpt-start");
let terminal = if self.fail {
Err(ProviderError::new(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::Sent,
))
} else {
Ok(GatewayEvent::Completed(response.clone()).into())
};
let events: Vec<Result<ProviderEvent, ProviderError>> = vec![
Ok(GatewayEvent::Started(response).into()),
Ok(GatewayEvent::ProviderCost(
ProviderReportedCost::from_usd_ticks(known_charge().scaled()).unwrap(),
)
.into()),
terminal,
];
Ok(ProviderStream::new(
metadata,
futures::stream::iter(events),
(),
))
}
}
fn charged_service(
admissions: Arc<Admissions>,
budget: Option<Arc<Budget>>,
fail: bool,
) -> DefaultExecutionService {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::new([Arc::new(ChargedProvider { fail }) as Arc<dyn Provider>]).unwrap(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
match budget {
Some(budget) => service.with_budget(budget),
None => service,
}
}
async fn consume_charged_prefix(session: &mut dyn ExecutionSession) {
let first = session.next_event().await.unwrap().unwrap();
assert_eq!(
first.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
session.commit_downstream(None).await.unwrap();
let cost = session.next_event().await.unwrap().unwrap();
assert_eq!(
cost.commit_requirement(),
CommitRequirement::AlreadyCommitted
);
}
fn assert_cleanup_completed(admissions: &Admissions, budget: &Budget, request_id: &ModelRequestId) {
assert!(!admissions.active.load(Ordering::SeqCst));
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
let charges = budget.charges.lock().unwrap();
assert_eq!(charges.len(), 1);
assert_eq!(charges[0].request_id, *request_id);
assert_eq!(charges[0].key_id.as_str(), "key_start_test");
assert_eq!(charges[0].amount_usd, known_charge());
}
#[test]
fn normal_completion_settles_known_charge_once_for_all_transports() {
block_on(async {
for transport in [
ClientTransport::HttpJson,
ClientTransport::HttpSse,
ClientTransport::WebSocket,
] {
let admissions = Arc::new(Admissions::default());
let budget = Arc::new(Budget {
active: admissions.active.clone(),
..Default::default()
});
let service = charged_service(admissions.clone(), Some(budget.clone()), false);
let mut started = service.start(request(&service, transport)).await.unwrap();
if transport == ClientTransport::HttpJson {
assert_eq!(
started.session.collect_uncommitted().await.unwrap().len(),
3
);
assert!(!started.session.is_finalized());
started.session.commit_downstream(Some(200)).await.unwrap();
} else {
consume_charged_prefix(started.session.as_mut()).await;
assert!(started.session.next_event().await.unwrap().is_some());
}
assert!(started.session.next_event().await.unwrap().is_none());
assert!(started.session.is_finalized());
started.session.detach_finalize().await;
assert_cleanup_completed(&admissions, &budget, &started.request_id);
}
});
}
#[test]
fn settlement_survives_cancelled_event_wait_and_resumes_without_restart() {
block_on(async {
for fail in [false, true] {
let admissions = Arc::new(Admissions::default());
let (release, gate) = oneshot::channel();
let budget = Arc::new(Budget {
active: admissions.active.clone(),
settlement_gate: Mutex::new(Some(gate)),
..Default::default()
});
let service = charged_service(admissions.clone(), Some(budget.clone()), fail);
let mut started = service
.start(request(&service, ClientTransport::WebSocket))
.await
.unwrap();
consume_charged_prefix(started.session.as_mut()).await;
if !fail {
assert!(started.session.next_event().await.unwrap().is_some());
}
for _ in 0..2 {
let mut terminal = started.session.next_event();
assert!(futures::poll!(terminal.as_mut()).is_pending());
drop(terminal);
assert!(!started.session.is_finalized());
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
assert!(budget.charges.lock().unwrap().is_empty());
assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
}
release
.send(())
.expect("the original settlement is still alive");
assert!(started.session.next_event().await.unwrap().is_none());
assert!(started.session.is_finalized());
started.session.detach_finalize().await;
assert_cleanup_completed(&admissions, &budget, &started.request_id);
}
});
}
#[test]
fn detached_cleanup_settles_cancelled_execution_and_resumes_existing_settlement() {
block_on(async {
for cancel_event_wait in [false, true] {
let admissions = Arc::new(Admissions::default());
let (release, gate) = oneshot::channel();
let budget = Arc::new(Budget {
active: admissions.active.clone(),
settlement_gate: Mutex::new(Some(gate)),
..Default::default()
});
let service = charged_service(admissions.clone(), Some(budget.clone()), false);
let mut started = service
.start(request(&service, ClientTransport::WebSocket))
.await
.unwrap();
consume_charged_prefix(started.session.as_mut()).await;
if cancel_event_wait {
started.session.cancel();
let mut terminal = started.session.next_event();
assert!(futures::poll!(terminal.as_mut()).is_pending());
drop(terminal);
assert!(!started.session.is_finalized());
}
let mut detached = started.session.detach_finalize();
assert!(futures::poll!(detached.as_mut()).is_pending());
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
release
.send(())
.expect("detached cleanup retains settlement");
detached.await;
assert_cleanup_completed(&admissions, &budget, &started.request_id);
}
});
}
#[test]
fn cancelled_release_resumes_without_repeating_settlement_with_or_without_budget() {
block_on(async {
for with_budget in [false, true] {
let (release, gate) = oneshot::channel();
let admissions = Arc::new(Admissions {
release_gate: Mutex::new(Some(gate)),
..Default::default()
});
let budget = Arc::new(Budget {
active: admissions.active.clone(),
..Default::default()
});
let service = charged_service(
admissions.clone(),
with_budget.then(|| budget.clone()),
false,
);
let mut started = service
.start(request(&service, ClientTransport::HttpSse))
.await
.unwrap();
consume_charged_prefix(started.session.as_mut()).await;
assert!(started.session.next_event().await.unwrap().is_some());
let mut terminal = started.session.next_event();
assert!(futures::poll!(terminal.as_mut()).is_pending());
drop(terminal);
assert!(!started.session.is_finalized());
assert!(admissions.active.load(Ordering::SeqCst));
assert_eq!(
budget.settlements.load(Ordering::SeqCst),
usize::from(with_budget)
);
let mut detached = started.session.detach_finalize();
assert!(futures::poll!(detached.as_mut()).is_pending());
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
release
.send(())
.expect("the original release is still alive");
detached.await;
if with_budget {
assert_cleanup_completed(&admissions, &budget, &started.request_id);
} else {
assert!(!admissions.active.load(Ordering::SeqCst));
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
assert!(budget.charges.lock().unwrap().is_empty());
}
}
});
}
#[test]
fn cancelled_buffered_commit_keeps_settlement_for_detached_cleanup() {
block_on(async {
let admissions = Arc::new(Admissions::default());
let (release, gate) = oneshot::channel();
let budget = Arc::new(Budget {
active: admissions.active.clone(),
settlement_gate: Mutex::new(Some(gate)),
..Default::default()
});
let service = charged_service(admissions.clone(), Some(budget.clone()), false);
let mut started = service
.start(request(&service, ClientTransport::HttpJson))
.await
.unwrap();
assert_eq!(
started.session.collect_uncommitted().await.unwrap().len(),
3
);
let mut commit = started.session.commit_downstream(Some(200));
assert!(futures::poll!(commit.as_mut()).is_pending());
drop(commit);
assert!(!started.session.is_finalized());
release
.send(())
.expect("buffered commit retains settlement");
started.session.detach_finalize().await;
assert_cleanup_completed(&admissions, &budget, &started.request_id);
});
}
#[test]
fn settlement_failure_keeps_provider_error_and_releases_concurrency_once() {
block_on(async {
for fail_settlement in [false, true] {
let admissions = Arc::new(Admissions::default());
let budget = Arc::new(Budget {
active: admissions.active.clone(),
fail_settlement,
..Default::default()
});
let service = charged_service(admissions.clone(), Some(budget.clone()), true);
let mut started = service
.start(request(&service, ClientTransport::WebSocket))
.await
.unwrap();
consume_charged_prefix(started.session.as_mut()).await;
assert!(matches!(
started.session.next_event().await,
Err(EngineError::Provider(error)) if error.kind() == ProviderErrorKind::InvalidRequest
));
assert!(started.session.is_finalized());
started.session.detach_finalize().await;
// Store 端口已接管精确费用后，结算错误不能改写 Provider 错误或触发第二次结算。
assert_cleanup_completed(&admissions, &budget, &started.request_id);
}
});
}
use async_trait::async_trait;
use bytes::Bytes;
use futures::{executor::block_on, future::BoxFuture};
use futures::{channel::oneshot, executor::block_on, future::BoxFuture};
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use gateway_core::engine::admission::{
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
@@ -245,31 +597,34 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::execution::{
ClientApiKeyUsageSink, ClientTransport, DefaultExecutionService, ExecutionRequestMetadata,
ExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
StartExecution, StartProviderExecution, provider_failure_affects_circuit,
ExecutionService, ExecutionSession, ProviderCircuitDecision, ProviderCircuitError,
ProviderCircuitPort, StartExecution, StartProviderExecution, provider_failure_affects_circuit,
};
use gateway_core::engine::probe::{AccountProbe, AccountProbeErrorSource, AccountProbeRequest};
use gateway_core::engine::provider::{
Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities,
ProviderRegistry, ProviderRequest, ProviderRequestObservation, ProviderStream,
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderRequestObservation,
ProviderStream,
};
use gateway_core::engine::{
AttemptContext, AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport,
AttemptContext, AttemptRecord, CommitRequirement, EngineError, ExecutionStore,
IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest, ProbeFailure,
RecoveryReport,
};
use gateway_core::error::{
ClientVisibleUpstreamResponse, GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError,
StoreErrorKind,
};
use gateway_core::event::{GatewayEvent, ProviderEvent, ResponseMeta};
use gateway_core::metering::{Decimal, ProviderReportedCost};
use gateway_core::operation::{
GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProtocolPayload,
RawJsonPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
UpstreamModelId,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
ProviderCatalogGeneration, ProviderKind, ProviderModel, ProviderModelCapabilities,
PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
use gateway_core::runtime::RuntimeSnapshotHandle;
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
backend/crates/gateway-core/tests/engine/provider.rs
+81 -4
View File
@@ -11,8 +11,8 @@ use futures_timer::Delay;
use gateway_core::account::{AccountFeedbackStats, ProviderAccountId};
use gateway_core::engine::AttemptContext;
use gateway_core::engine::provider::{
EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration,
ProviderModelCapabilities, ProviderRegistry, ProviderRequest, ProviderStream, RegistryError,
EventStream, Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
RegistryError,
};
use gateway_core::error::{ProviderError, ProviderErrorKind};
use gateway_core::event::{
@@ -20,7 +20,10 @@ use gateway_core::event::{
TextDelta,
};
use gateway_core::operation::OperationKind;
use gateway_core::routing::{ModelCapabilities, ProviderKind, UpstreamModelId};
use gateway_core::routing::{
ModelCapabilities, ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
ProviderKind, ProviderModelCapabilities, UpstreamModelId,
};
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use serde_json::json;
@@ -84,9 +87,14 @@ fn registry_should_query_provider_compiled_model_capabilities() {
.expect("provider");
let registry = builder.build();
let provider = ProviderKind::new("openai").expect("provider kind");
let models = futures::executor::block_on(registry.query_model_capabilities(&provider))
let catalog: &dyn ProviderCatalogPort = &registry;
let models = futures::executor::block_on(catalog.query_model_capabilities(&provider))
.expect("live catalog");
assert_eq!(
catalog.catalog_generations(),
std::collections::BTreeMap::from([(provider, ProviderCatalogGeneration::default())]),
);
assert_eq!(models[0].upstream_model().as_str(), "live-model");
assert!(
models[0]
@@ -98,6 +106,75 @@ fn registry_should_query_provider_compiled_model_capabilities() {
);
}
struct CatalogProvider(Result<Vec<ProviderModelCapabilities>, ProviderErrorKind>);
#[async_trait]
impl Provider for CatalogProvider {
fn name(&self) -> &'static str {
"alpha"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::new(7)
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
self.0
.clone()
.map_err(|kind| ProviderError::new(kind, UpstreamSendState::NotSent))
}
async fn execute(
&self,
_: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
panic!("目录端口不得触发执行")
}
}
#[test]
fn registry_catalog_port_should_preserve_unknown_and_known_empty_catalogs() {
let provider = ProviderKind::new("alpha").expect("provider");
for (response, expected) in [
(
Err(ProviderErrorKind::Unavailable),
Err(ProviderCatalogUnavailable),
),
(Ok(Vec::new()), Ok(Vec::new())),
] {
let registry =
ProviderRegistry::new([Arc::new(CatalogProvider(response)) as Arc<dyn Provider>])
.expect("registry");
let catalog: &dyn ProviderCatalogPort = &registry;
assert_eq!(
futures::executor::block_on(catalog.query_model_capabilities(&provider)),
expected,
);
assert_eq!(
catalog.catalog_generations().get(&provider),
Some(&ProviderCatalogGeneration::new(7)),
"目录不可读也不能丢失注册身份和最近成功发布的代次",
);
}
}
#[test]
fn registry_catalog_port_should_report_unregistered_provider_as_unknown() {
let registry = ProviderRegistry::default();
let catalog: &dyn ProviderCatalogPort = &registry;
let provider = ProviderKind::new("missing").expect("provider");
assert_eq!(
futures::executor::block_on(catalog.query_model_capabilities(&provider)),
Err(ProviderCatalogUnavailable),
);
assert!(catalog.catalog_generations().is_empty());
}
struct DropLease(Arc<AtomicBool>);
impl Drop for DropLease {
backend/crates/gateway-core/tests/routing/snapshot.rs
+64 -91
View File
@@ -2,16 +2,9 @@ use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use futures::executor::block_on;
use futures::future::BoxFuture;
use gateway_core::engine::AttemptContext;
use gateway_core::engine::provider::{
Provider, ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRegistry,
ProviderRequest, ProviderStream,
};
use gateway_core::error::{ProviderError, ProviderErrorKind};
use gateway_core::operation::OperationKind;
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::{
@@ -20,10 +13,10 @@ use gateway_core::routing::snapshot::{
SnapshotProviderAccountFacts, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ModelPresentation, ProviderKind, PublicModelId,
UpstreamModelId,
ConfigRevision, ModelCapabilities, ModelPresentation, ProviderCatalogGeneration,
ProviderCatalogPort, ProviderCatalogUnavailable, ProviderKind, ProviderModelCapabilities,
PublicModelId, UpstreamModelId,
};
use gateway_core::upstream::UpstreamSendState;
#[derive(Clone)]
struct TestSnapshotStore {
@@ -51,90 +44,72 @@ impl SnapshotStorePort for TestSnapshotStore {
}
}
struct PublishingCatalogProvider {
struct PublishingCatalog {
generation: AtomicU64,
queries: AtomicUsize,
}
enum TestCatalogProvider {
enum TestCatalog {
NoProviders,
Unavailable,
Empty,
}
#[async_trait]
impl Provider for TestCatalogProvider {
fn name(&self) -> &'static str {
"alpha"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::new(0)
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
impl ProviderCatalogPort for TestCatalog {
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
match self {
Self::Unavailable => Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
)),
Self::Empty => Ok(Vec::new()),
Self::NoProviders => BTreeMap::new(),
Self::Unavailable | Self::Empty => catalog_generations(0),
}
}
async fn execute(
fn query_model_capabilities(
&self,
_: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
_: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
Box::pin(async move {
match self {
Self::NoProviders | Self::Unavailable => Err(ProviderCatalogUnavailable),
Self::Empty => Ok(Vec::new()),
}
})
}
}
#[async_trait]
impl Provider for PublishingCatalogProvider {
fn name(&self) -> &'static str {
"alpha"
fn catalog_generations(generation: u64) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
BTreeMap::from([(
ProviderKind::new("alpha").expect("provider"),
ProviderCatalogGeneration::new(generation),
)])
}
impl ProviderCatalogPort for PublishingCatalog {
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
catalog_generations(self.generation.load(Ordering::SeqCst))
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::new(self.generation.load(Ordering::SeqCst))
}
async fn query_model_capabilities(
fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
self.generation.store(1, Ordering::SeqCst);
}
Ok(vec![
ProviderModelCapabilities::new(
UpstreamModelId::new("upstream-model").expect("model"),
ModelCapabilities::new(
std::collections::BTreeSet::from([OperationKind::Generate]),
_: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
Box::pin(async move {
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
self.generation.store(1, Ordering::SeqCst);
}
Ok(vec![
ProviderModelCapabilities::new(
UpstreamModelId::new("upstream-model").expect("model"),
ModelCapabilities::new(
std::collections::BTreeSet::from([OperationKind::Generate]),
None,
),
)
.with_presentation(ModelPresentation::new(
Some("Upstream Model".to_owned()),
None,
),
)
.with_presentation(ModelPresentation::new(
Some("Upstream Model".to_owned()),
None,
)),
])
}
async fn execute(
&self,
_: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
)),
])
})
}
}
@@ -150,11 +125,10 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
#[test]
fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
let providers =
ProviderRegistry::new([Arc::new(TestCatalogProvider::Unavailable) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
let compiler = RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))),
Arc::new(TestCatalog::Unavailable),
);
let snapshot = block_on(compiler.compile()).expect("compile snapshot");
let provider = ProviderKind::new("alpha").expect("provider");
@@ -170,11 +144,10 @@ fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
#[test]
fn known_empty_catalog_should_report_model_not_found() {
let providers =
ProviderRegistry::new([Arc::new(TestCatalogProvider::Empty) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
let compiler = RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))),
Arc::new(TestCatalog::Empty),
);
let snapshot = block_on(compiler.compile()).expect("compile empty catalog");
let error = snapshot
.plan(
@@ -199,18 +172,18 @@ fn known_empty_catalog_should_report_model_not_found() {
#[test]
fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
let provider = Arc::new(PublishingCatalogProvider {
let catalog = Arc::new(PublishingCatalog {
generation: AtomicU64::new(0),
queries: AtomicUsize::new(0),
});
let providers =
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
let compiler = RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))),
catalog.clone(),
);
let snapshot = block_on(compiler.compile()).expect("stable catalog snapshot");
assert_eq!(provider.queries.load(Ordering::SeqCst), 2);
assert_eq!(catalog.queries.load(Ordering::SeqCst), 2);
assert_eq!(
snapshot
.provider_catalog_generations()
@@ -311,7 +284,7 @@ fn facts_revision(facts: &SnapshotFacts) -> ConfigRevision {
}
fn compiler(store: Arc<dyn SnapshotStorePort>) -> RuntimeSnapshotCompiler {
RuntimeSnapshotCompiler::new(store, ProviderRegistry::default())
RuntimeSnapshotCompiler::new(store, Arc::new(TestCatalog::NoProviders))
}
fn revision(value: u64) -> ConfigRevision {
backend/crates/gateway-core/tests/runtime/mod.rs
+445 -8
View File
@@ -1,30 +1,42 @@
use std::collections::BTreeMap;
use std::num::NonZeroU32;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use futures::channel::{mpsc, oneshot};
use futures::executor::block_on;
use futures::future::BoxFuture;
use gateway_core::account::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::provider::ProviderRegistry;
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotProviderAccountFacts, SnapshotSettingsFacts,
SnapshotStoreError, SnapshotStorePort,
};
use gateway_core::routing::{ConfigRevision, RuntimeSnapshot};
use gateway_core::routing::{
AccountGroupId, ConfigRevision, ProviderCatalogGeneration, ProviderCatalogPort,
ProviderCatalogUnavailable, ProviderKind, ProviderModelCapabilities, RuntimeSnapshot,
};
use gateway_core::runtime::{
RuntimeSnapshotHandle, RuntimeSnapshotPublisher, SnapshotControl, SnapshotRevisionStream,
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
};
use gateway_core::task::WorkerKind;
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerKind, WorkerRunnable,
};
type ReadGate<T> = Arc<Mutex<Option<oneshot::Receiver<Result<T, SnapshotStoreError>>>>>;
#[derive(Clone)]
struct TestSnapshotStore {
facts: Arc<Mutex<Result<SnapshotFacts, SnapshotStoreError>>>,
current_revision: Arc<Mutex<Result<ConfigRevision, SnapshotStoreError>>>,
next_load: ReadGate<SnapshotFacts>,
next_revision: ReadGate<ConfigRevision>,
loads: Arc<AtomicUsize>,
}
impl TestSnapshotStore {
@@ -36,23 +48,45 @@ impl TestSnapshotStore {
Self {
facts: Arc::new(Mutex::new(facts)),
current_revision: Arc::new(Mutex::new(current_revision)),
next_load: Arc::default(),
next_revision: Arc::default(),
loads: Arc::default(),
}
}
}
impl SnapshotStorePort for TestSnapshotStore {
fn load_snapshot_facts(&self) -> BoxFuture<'_, Result<SnapshotFacts, SnapshotStoreError>> {
Box::pin(async move { self.facts.lock().expect("facts lock").clone() })
Box::pin(async move {
self.loads.fetch_add(1, Ordering::SeqCst);
let gate = self.next_load.lock().expect("load gate lock").take();
if let Some(gate) = gate {
return gate.await.expect("release load");
}
self.facts.lock().expect("facts lock").clone()
})
}
fn current_config_revision(&self) -> BoxFuture<'_, Result<ConfigRevision, SnapshotStoreError>> {
Box::pin(async move { self.current_revision.lock().expect("revision lock").clone() })
Box::pin(async move {
let gate = self
.next_revision
.lock()
.expect("revision gate lock")
.take();
if let Some(gate) = gate {
return gate.await.expect("release revision read");
}
self.current_revision.lock().expect("revision lock").clone()
})
}
}
#[derive(Default)]
struct TestSnapshotSubscriptions {
published: Mutex<Vec<ConfigRevision>>,
stream: Mutex<Option<SnapshotRevisionStream>>,
next_publish: Mutex<Option<oneshot::Receiver<()>>>,
}
impl SnapshotSubscriptionPort for TestSnapshotSubscriptions {
@@ -61,6 +95,10 @@ impl SnapshotSubscriptionPort for TestSnapshotSubscriptions {
revision: ConfigRevision,
) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>> {
Box::pin(async move {
let gate = self.next_publish.lock().expect("publish gate lock").take();
if let Some(gate) = gate {
gate.await.expect("release notification");
}
self.published
.lock()
.expect("published lock")
@@ -72,7 +110,14 @@ impl SnapshotSubscriptionPort for TestSnapshotSubscriptions {
fn subscribe_snapshot_revisions(
&self,
) -> BoxFuture<'_, Result<SnapshotRevisionStream, SnapshotSubscriptionError>> {
Box::pin(async move { Ok(Box::pin(futures::stream::empty()) as SnapshotRevisionStream) })
Box::pin(async move {
Ok(self
.stream
.lock()
.expect("subscription stream lock")
.take()
.unwrap_or_else(|| Box::pin(futures::stream::pending())))
})
}
}
@@ -187,6 +232,398 @@ fn publisher_should_contribute_reconciliation_and_subscription_workers() {
);
}
#[test]
fn overlapping_refreshes_should_not_restore_revoked_account_scope() {
block_on(async {
let (release, gate) = oneshot::channel();
let catalog = Arc::new(TestCatalog::default());
*catalog.next_query.lock().expect("catalog gate lock") = Some(gate);
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(catalog_compiler(store.clone(), catalog.clone())),
handle.clone(),
Arc::new(TestSnapshotSubscriptions::default()),
);
let other_publisher = publisher.clone();
// 旧刷新已读完 revision 10 的完整事实，停在目录查询；不用 sleep 碰调度概率。
let mut old = Box::pin(publisher.refresh());
assert!(futures::poll!(old.as_mut()).is_pending());
assert_eq!(catalog.queries.load(Ordering::SeqCst), 1);
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
*store.current_revision.lock().expect("revision lock") = Ok(revision(11));
let mut new = Box::pin(other_publisher.refresh());
assert!(futures::poll!(new.as_mut()).is_pending());
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
assert_eq!(handle.revision(), Some(revision(9)));
release.send(()).expect("release old catalog query");
assert_eq!(old.await.expect("old refresh"), revision(10));
let frozen_old = handle.acquire().expect("old request snapshot");
assert!(removed_account_allowed(&frozen_old));
assert_eq!(new.await.expect("new refresh"), revision(11));
let current = handle.acquire().expect("new request snapshot");
assert!(!removed_account_allowed(&current));
assert_eq!(current.revision(), revision(11));
assert!(removed_account_allowed(&frozen_old));
assert_eq!(frozen_old.revision(), revision(10));
});
}
#[test]
fn failed_committed_refresh_should_suspend_before_next_publication() {
block_on(async {
let (release, gate) = oneshot::channel();
let store = Arc::new(TestSnapshotStore::new(Ok(facts(11, 11))));
*store.next_load.lock().expect("load gate lock") = Some(gate);
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
let subscriptions = Arc::new(TestSnapshotSubscriptions::default());
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(compiler(store.clone())),
handle.clone(),
subscriptions.clone(),
);
let other_publisher = publisher.clone();
let mut old = publisher.publish_committed(revision(10));
assert!(futures::poll!(old.as_mut()).is_pending());
let mut new = other_publisher.publish_committed(revision(11));
assert!(futures::poll!(new.as_mut()).is_pending());
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
release
.send(Err(SnapshotStoreError::unavailable()))
.expect("release failed facts read");
old.await;
assert!(handle.acquire().is_err());
new.await;
assert_eq!(handle.revision(), Some(revision(11)));
assert_eq!(
*subscriptions.published.lock().expect("published lock"),
vec![revision(10), revision(11)],
);
});
}
#[test]
fn reconciliation_revision_failure_should_not_suspend_later_committed_refresh() {
block_on(async {
let (release, gate) = oneshot::channel();
let store = Arc::new(TestSnapshotStore::new(Ok(facts(11, 11))));
*store.next_revision.lock().expect("revision gate lock") = Some(gate);
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(compiler(store.clone())),
handle.clone(),
Arc::new(TestSnapshotSubscriptions::default()),
);
let (task, context) = reconciliation_task(&publisher);
let mut reconcile = task.run_cycle(context);
assert!(futures::poll!(reconcile.as_mut()).is_pending());
let mut committed = publisher.publish_committed(revision(11));
assert!(futures::poll!(committed.as_mut()).is_pending());
assert_eq!(store.loads.load(Ordering::SeqCst), 0);
release
.send(Err(SnapshotStoreError::unavailable()))
.expect("release failed revision read");
assert!(reconcile.await.is_err());
assert!(handle.acquire().is_err());
committed.await;
assert_eq!(handle.revision(), Some(revision(11)));
});
}
#[test]
fn subscription_refresh_should_wait_for_inflight_compile_and_reload_authoritative_facts() {
block_on(async {
let (release, gate) = oneshot::channel();
let catalog = Arc::new(TestCatalog::default());
*catalog.next_query.lock().expect("catalog gate lock") = Some(gate);
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
let (notifications, stream) = mpsc::unbounded();
let subscriptions = Arc::new(TestSnapshotSubscriptions {
stream: Mutex::new(Some(Box::pin(stream))),
..TestSnapshotSubscriptions::default()
});
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(catalog_compiler(store.clone(), catalog)),
handle.clone(),
subscriptions,
);
let task = publisher
.worker_contributions()
.expect("workers")
.into_iter()
.find_map(|contribution| match contribution {
WorkerContribution::Registration(registration) => match registration.runnable {
WorkerRunnable::Daemon { task, .. } => Some(task),
WorkerRunnable::Scheduled { .. } => None,
},
WorkerContribution::Disabled { .. } => None,
})
.expect("subscription task");
let cancellation = CancellationToken::new();
let mut daemon = task.run(cancellation.clone());
let mut old = Box::pin(publisher.refresh());
assert!(futures::poll!(old.as_mut()).is_pending());
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
// 通知只是提示，即使版本过期也要在取得发布权后重新读取 Store。
notifications
.unbounded_send(Ok(revision(3)))
.expect("notify");
assert!(futures::poll!(daemon.as_mut()).is_pending());
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
assert_eq!(handle.revision(), Some(revision(9)));
release.send(()).expect("release old catalog query");
old.await.expect("old refresh");
assert!(futures::poll!(daemon.as_mut()).is_pending());
let current = handle.acquire().expect("subscription snapshot");
assert_eq!(current.revision(), revision(11));
assert!(!removed_account_allowed(&current));
cancellation.cancel();
daemon.await.expect("stop subscription");
});
}
#[test]
fn reconciliation_should_fail_closed_on_persisted_revision_rollback_and_recover() {
block_on(async {
let store = Arc::new(TestSnapshotStore::new(Ok(facts(7, 7))));
*store.facts.lock().expect("facts lock") = Err(SnapshotStoreError::unavailable());
let handle = RuntimeSnapshotHandle::new(empty_snapshot(8));
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(compiler(store.clone())),
handle.clone(),
Arc::new(TestSnapshotSubscriptions::default()),
);
let (task, context) = reconciliation_task(&publisher);
assert!(task.run_cycle(context.clone()).await.is_err());
assert!(handle.acquire().is_err());
*store.facts.lock().expect("facts lock") = Ok(facts(7, 7));
task.run_cycle(context.clone()).await.expect("recover");
assert_eq!(handle.revision(), Some(revision(7)));
// 有效的持久回退也必须发布，不能用 revision 数值单调性掩盖竞争。
*store.facts.lock().expect("facts lock") = Ok(facts(6, 6));
*store.current_revision.lock().expect("revision lock") = Ok(revision(6));
task.run_cycle(context).await.expect("publish rollback");
assert_eq!(handle.revision(), Some(revision(6)));
});
}
#[test]
fn catalog_only_reconciliation_should_keep_snapshot_on_failure_and_retry_generation() {
block_on(async {
let catalog = Arc::new(TestCatalog::default());
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
let handle = RuntimeSnapshotHandle::default();
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(catalog_compiler(store.clone(), catalog.clone())),
handle.clone(),
Arc::new(TestSnapshotSubscriptions::default()),
);
publisher.refresh().await.expect("initial snapshot");
let frozen = handle.acquire().expect("frozen snapshot");
let (task, context) = reconciliation_task(&publisher);
task.run_cycle(context.clone()).await.expect("unchanged");
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
assert!(Arc::ptr_eq(
&frozen,
&handle.acquire().expect("unchanged snapshot")
));
catalog.generation.store(1, Ordering::SeqCst);
*store.facts.lock().expect("facts lock") = Err(SnapshotStoreError::unavailable());
assert!(task.run_cycle(context.clone()).await.is_err());
assert!(Arc::ptr_eq(
&frozen,
&handle.acquire().expect("retained snapshot")
));
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(10, true));
task.run_cycle(context)
.await
.expect("retry catalog refresh");
let current = handle.acquire().expect("refreshed catalog");
assert_eq!(current.revision(), revision(10));
assert_eq!(
current
.provider_catalog_generations()
.values()
.copied()
.collect::<Vec<_>>(),
vec![ProviderCatalogGeneration::new(1)],
);
assert_eq!(
frozen
.provider_catalog_generations()
.values()
.copied()
.collect::<Vec<_>>(),
vec![ProviderCatalogGeneration::new(0)],
);
});
}
#[test]
fn dropping_inflight_refresh_should_release_publication_for_next_refresh() {
block_on(async {
let (release, gate) = oneshot::channel();
let store = Arc::new(TestSnapshotStore::new(Ok(facts(11, 11))));
*store.next_load.lock().expect("load gate lock") = Some(gate);
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(compiler(store)),
handle.clone(),
Arc::new(TestSnapshotSubscriptions::default()),
);
let mut cancelled = Box::pin(publisher.refresh());
assert!(futures::poll!(cancelled.as_mut()).is_pending());
let mut next = Box::pin(publisher.refresh());
assert!(futures::poll!(next.as_mut()).is_pending());
drop(cancelled);
assert!(release.send(Ok(facts(10, 10))).is_err());
assert_eq!(next.await.expect("next refresh"), revision(11));
assert_eq!(handle.revision(), Some(revision(11)));
});
}
#[test]
fn pending_notification_should_not_block_next_local_refresh() {
block_on(async {
let (release, gate) = oneshot::channel();
let store = Arc::new(TestSnapshotStore::new(Ok(facts(10, 10))));
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
let subscriptions = Arc::new(TestSnapshotSubscriptions {
next_publish: Mutex::new(Some(gate)),
..TestSnapshotSubscriptions::default()
});
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(compiler(store.clone())),
handle.clone(),
subscriptions,
);
let mut committed = publisher.publish_committed(revision(10));
assert!(futures::poll!(committed.as_mut()).is_pending());
assert_eq!(handle.revision(), Some(revision(10)));
*store.facts.lock().expect("facts lock") = Ok(facts(11, 11));
let mut next = Box::pin(publisher.refresh());
assert_eq!(
futures::poll!(next.as_mut()),
std::task::Poll::Ready(Ok(revision(11))),
);
release.send(()).expect("release notification");
committed.await;
assert_eq!(handle.revision(), Some(revision(11)));
});
}
fn reconciliation_task(
publisher: &RuntimeSnapshotPublisher,
) -> (Box<dyn ScheduledTask>, WorkerCycleContext) {
publisher
.worker_contributions()
.expect("workers")
.into_iter()
.find_map(|contribution| match contribution {
WorkerContribution::Registration(registration) => match registration.runnable {
WorkerRunnable::Scheduled { task, .. } => Some((
task,
WorkerCycleContext::new(registration.id, None, CancellationToken::new()),
)),
WorkerRunnable::Daemon { .. } => None,
},
WorkerContribution::Disabled { .. } => None,
})
.expect("reconciliation task")
}
#[derive(Default)]
struct TestCatalog {
generation: AtomicU64,
queries: AtomicUsize,
next_query: Mutex<Option<oneshot::Receiver<()>>>,
}
impl ProviderCatalogPort for TestCatalog {
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
BTreeMap::from([(
ProviderKind::new("audit").expect("provider"),
ProviderCatalogGeneration::new(self.generation.load(Ordering::SeqCst)),
)])
}
fn query_model_capabilities(
&self,
_: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
Box::pin(async move {
self.queries.fetch_add(1, Ordering::SeqCst);
let gate = self.next_query.lock().expect("catalog gate lock").take();
if let Some(gate) = gate {
gate.await.expect("release catalog query");
}
Ok(Vec::new())
})
}
}
fn catalog_compiler(
store: Arc<dyn SnapshotStorePort>,
catalog: Arc<TestCatalog>,
) -> RuntimeSnapshotCompiler {
RuntimeSnapshotCompiler::new(store, catalog)
}
fn scoped_facts(value: u64, allow_removed: bool) -> SnapshotFacts {
let group = AccountGroupId::new(format!("grp_{:032x}", 1)).expect("group");
let keep = ProviderAccountId::new("acct_audit_keep").expect("kept account");
let removed = ProviderAccountId::new("acct_audit_removed").expect("removed account");
let mut memberships = vec![SnapshotAccountGroupMemberFacts::new(
group.clone(),
keep.clone(),
)];
if allow_removed {
memberships.push(SnapshotAccountGroupMemberFacts::new(
group.clone(),
removed.clone(),
));
}
SnapshotFacts::new(
revision(value),
revision(value),
SnapshotSettingsFacts::new(3, 0, "smart", BTreeMap::new(), None, None),
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_audit_synthetic").expect("key ID"),
PlaintextClientApiKey::new("sk_audit_synthetic_not_a_real_key").expect("synthetic key"),
vec![group.clone()],
RateLimits::unlimited(),
)],
vec![SnapshotAccountGroupFacts::new(
group,
"Synthetic group".to_owned(),
true,
)],
vec![
SnapshotProviderAccountFacts::new(keep, "audit"),
SnapshotProviderAccountFacts::new(removed, "audit"),
],
memberships,
)
}
fn removed_account_allowed(snapshot: &RuntimeSnapshot) -> bool {
snapshot
.client_policies()
.next()
.expect("client policy")
.account_scope()
.allows(&ProviderAccountId::new("acct_audit_removed").expect("removed account"))
}
fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts {
SnapshotFacts::new(
revision(config_revision),
@@ -212,7 +649,7 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
}
fn compiler(store: Arc<dyn SnapshotStorePort>) -> RuntimeSnapshotCompiler {
RuntimeSnapshotCompiler::new(store, ProviderRegistry::default())
catalog_compiler(store, Arc::new(TestCatalog::default()))
}
fn empty_snapshot(value: u64) -> RuntimeSnapshot {
backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs
-81
View File
@@ -476,87 +476,6 @@ impl AccountStore for PgAdminAccountStore {
self.usage_by_windows(windows).await
}
async fn list_credentials(
&self,
provider_kind: &ProviderKind,
query: CredentialListQuery,
) -> AdminStoreResult<CredentialPage> {
let (control_plane, accounts) = futures::try_join!(
self.control_plane.load_control_plane(),
self.accounts
.list_provider_accounts(Some(provider_kind.as_str()), true),
)
.map_err(|error| admin_store_error(ENTITY, error))?;
let mut accounts = accounts
.into_iter()
.filter(|account| account.provider_kind == provider_kind.as_str())
.filter(|account| {
query
.credential_state
.as_ref()
.is_none_or(|expected| expected.matches(account.credential_state))
})
.filter(|account| {
query
.enabled
.is_none_or(|enabled| account.enabled == enabled)
})
.filter(|account| {
let CredentialListWindow::Page {
cursor: Some(cursor),
..
} = &query.window
else {
return true;
};
account.created_at > cursor.created_at
|| (account.created_at == cursor.created_at
&& account.id.as_str() > cursor.account_id.as_str())
})
.collect::<Vec<_>>();
if matches!(&query.window, CredentialListWindow::Page { .. }) {
accounts.sort_by(|left, right| {
left.created_at
.cmp(&right.created_at)
.then_with(|| left.id.cmp(&right.id))
});
}
let next_cursor = match query.window {
CredentialListWindow::All => None,
CredentialListWindow::Page { page_size, .. } => {
let page_size = usize::from(page_size.get());
let has_more = accounts.len() > page_size;
accounts.truncate(page_size);
has_more
.then(|| accounts.last())
.flatten()
.map(|account| {
Ok(CredentialCursor {
created_at: account.created_at,
account_id: CoreProviderAccountId::new(account.id.clone()).map_err(
|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted Provider account ID is invalid",
)
},
)?,
})
})
.transpose()?
}
};
Ok(CredentialPage {
config_revision: admin_revision(control_plane.settings.config_revision)?,
items: accounts
.into_iter()
.map(admin_account_record)
.collect::<AdminStoreResult<Vec<_>>>()?,
next_cursor,
})
}
async fn credential_details(
&self,
provider_kind: &ProviderKind,
backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs
+2 -3
View File
@@ -21,9 +21,8 @@ use gateway_admin::{
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
},
provider_credentials::{
AuthorizationCommit, AuthorizationCredentialCommit, CredentialCursor,
CredentialDetails, CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialListWindow, CredentialMutationResult, CredentialPage,
AuthorizationCommit, AuthorizationCredentialCommit, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
CredentialRotationCommit, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotationFacts, ProviderDocument, ProviderExportCredentialInput,
},
backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs
+1 -79
View File
@@ -16,8 +16,7 @@ use gateway_admin::{
observability::TimeRange,
provider_credentials::{
AuthorizationCommit, AuthorizationCredentialCommit, AuthorizationMutationTarget,
AuthorizationOwnerBinding, CredentialListQuery, CredentialListWindow,
CredentialStateFilter, PendingAuthorizationMutation, PreparedCredentialCreate,
AuthorizationOwnerBinding, PendingAuthorizationMutation, PreparedCredentialCreate,
ProviderDocument,
},
},
@@ -776,83 +775,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
database.close().await;
}
#[tokio::test]
async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collections() {
let Some(database) = TestDatabase::create("provider_credential_terminal_windows").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
for index in 0..205_u16 {
let mut credential = account(
&format!("acct_xai_{index:03}"),
&format!("user-xai-{index:03}"),
);
credential.provider_kind = "xai".to_owned();
credential.credential_state = match index {
0 => CredentialState::Expired,
1 => CredentialState::Banned,
2 => CredentialState::Invalid,
_ => CredentialState::Ready,
};
repository
.insert_provider_account(credential)
.await
.expect("insert xAI credential fixture");
}
let store = admin_account_store(&database.pool);
let provider = ProviderKind::new("xai").expect("xAI Provider kind");
let complete = store
.list_credentials(
&provider,
CredentialListQuery {
credential_state: None,
enabled: None,
window: CredentialListWindow::All,
},
)
.await
.expect("unpaged credential collection");
assert_eq!(complete.items.len(), 205);
assert!(complete.next_cursor.is_none());
let page = store
.list_credentials(
&provider,
CredentialListQuery {
credential_state: None,
enabled: None,
window: CredentialListWindow::Page {
cursor: None,
page_size: PageSize::new(200).expect("page size"),
},
},
)
.await
.expect("paged credential collection");
assert_eq!(page.items.len(), 200);
assert!(page.next_cursor.is_some());
let invalid_group = store
.list_credentials(
&provider,
CredentialListQuery {
credential_state: Some(CredentialStateFilter::AnyOf(vec![
CredentialState::Expired,
CredentialState::Banned,
CredentialState::Invalid,
])),
enabled: None,
window: CredentialListWindow::All,
},
)
.await
.expect("grouped invalid credential filter");
assert_eq!(invalid_group.items.len(), 3);
database.close().await;
}
#[tokio::test]
async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs() {
let Some(database) = TestDatabase::create("provider_account_terminal_usage").await else {
backend/crates/providers/openai/src/credential/catalog.rs
+1 -2
View File
@@ -9,11 +9,10 @@ use chrono::{DateTime, Utc};
use gateway_core::account::{
CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
};
use gateway_core::engine::provider::ProviderCatalogGeneration;
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use gateway_core::routing::{ProviderCatalogGeneration, ProviderKind, UpstreamModelId};
use secrecy::ExposeSecret;
use thiserror::Error;
use tokio::sync::Notify;
backend/crates/providers/openai/src/provider/mod.rs
+3 -3
View File
@@ -12,8 +12,7 @@ use futures::{StreamExt, future::BoxFuture};
use gateway_core::account::{AccountFeedbackStats, ProviderAccount};
use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope};
use gateway_core::engine::provider::{
ContinuationRequestObservation, EventStream, Provider, ProviderCallMetadata,
ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRequest,
ContinuationRequestObservation, EventStream, Provider, ProviderCallMetadata, ProviderRequest,
ProviderRequestObservation, ProviderSelectionObservation, ProviderStream,
};
use gateway_core::engine::{AttemptContext, AttemptTransport, ContinuationAttempt};
@@ -34,7 +33,8 @@ use gateway_core::operation::{
};
use gateway_core::provider_ports::ProviderSessionAffinityKey;
use gateway_core::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId,
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderCatalogGeneration,
ProviderKind, ProviderModelCapabilities, UpstreamModelId,
};
use gateway_core::task::{
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
backend/crates/providers/openai/tests/transport/http_client.rs
+15 -28
View File
@@ -1,34 +1,9 @@
use super::*;
#[test]
fn reqwest_client_should_keep_http2_fallback_hot_while_idle() {
let source = std::fs::read_to_string(
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/transport/client.rs"),
)
.expect("transport client source");
for setting in [
".pool_idle_timeout(None::<Duration>)",
".http2_keep_alive_interval(Duration::from_secs(30))",
".http2_keep_alive_timeout(Duration::from_secs(5))",
".http2_keep_alive_while_idle(true)",
] {
assert!(source.contains(setting), "missing HTTP/2 setting {setting}");
}
assert!(!source.contains(".http1_only()"));
}
#[test]
fn endpoints_should_join_backend_paths() {
assert_eq!(
provider_openai::transport::endpoint_url("https://api.example.com/", "/codex/responses"),
"https://api.example.com/codex/responses"
);
}
#[test]
fn custom_ca_should_report_environment_cache_key_consistently() {
const CASE_ENV: &str = "CODEX_PROXY_TEST_CUSTOM_CA_CACHE_KEY_CASE";
const CASE_COMPLETED: &str = "custom-ca-cache-key-case-completed:";
const SSL_CERT_PATH: &str = "/tmp/codex-proxy-ssl-cert-file.pem";
const CODEX_CA_PATH: &str = "/tmp/codex-proxy-codex-ca-certificate.pem";
@@ -50,6 +25,7 @@ fn custom_ca_should_report_environment_cache_key_consistently() {
provider_openai::transport::tls::custom_ca_env_cache_key(),
expected
);
println!("\n{CASE_COMPLETED}{case}");
return;
}
@@ -81,7 +57,7 @@ fn custom_ca_should_report_environment_cache_key_consistently() {
let mut command = Command::new(&current_exe);
command
.arg("--exact")
.arg("custom_ca_should_report_environment_cache_key_consistently")
.arg("transport::http_client::custom_ca_should_report_environment_cache_key_consistently")
.arg("--nocapture")
.env(CASE_ENV, case)
.env_remove(provider_openai::transport::tls::CODEX_CA_CERT_ENV)
@@ -100,6 +76,14 @@ fn custom_ca_should_report_environment_cache_key_consistently() {
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
// libtest 匹配零个测试也返回成功；完成标记证明该环境分支执行了生产断言。
let stdout = String::from_utf8_lossy(&output.stdout);
let completed = format!("{CASE_COMPLETED}{case}");
assert_eq!(
stdout.lines().filter(|line| *line == completed).count(),
1,
"isolated custom CA case {case} did not complete exactly once\nstdout:\n{stdout}"
);
}
}
@@ -563,8 +547,11 @@ async fn build_reqwest_client_should_reuse_cached_connection_pool() {
});
let url = format!("http://{addr}/reuse");
let client = reqwest::Client::builder().no_proxy().build().unwrap();
// 前后分别取得生产缓存中的客户端，避免只验证单个 reqwest client 自身能复用连接。
// 本地服务使用 HTTP/1.1；该断言不代表已验证 HTTP/2 的 idle ping 或保活间隔。
let client = provider_openai::transport::build_reqwest_client().unwrap();
client.get(&url).send().await.unwrap().text().await.unwrap();
let client = provider_openai::transport::build_reqwest_client().unwrap();
client.get(&url).send().await.unwrap().text().await.unwrap();
assert!(server.await.unwrap());
backend/crates/providers/xai/Cargo.toml
+1 -1
View File
@@ -34,7 +34,7 @@ zeroize = { version = "1.8.2", features = ["derive"] }
[dev-dependencies]
p256 = { version = "0.14.0", features = ["ecdsa", "pkcs8"] }
tokio = { version = "1.52.3", features = ["macros", "rt"] }
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread"] }
wiremock = "0.6.5"
[[test]]
backend/crates/providers/xai/src/credential/catalog.rs
+1 -38
View File
@@ -11,11 +11,10 @@ use gateway_core::account::{
AccountQuotaSignals, CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccount,
ProviderAccountId, QuotaObservation, QuotaState,
};
use gateway_core::engine::provider::ProviderCatalogGeneration;
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope,
};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{ProviderCatalogGeneration, ProviderKind};
use tokio::sync::Mutex;
use super::repository::{GrokCredentialRepository, LoadedGrokCredential};
@@ -902,42 +901,6 @@ impl GrokCredentialCatalogService {
ProviderCatalogGeneration::new(published.generation)
}
pub async fn fetch_seed(
&self,
access_token: SecretValue,
user_id: SecretValue,
email: Option<SecretValue>,
) -> Result<GrokCredentialCatalogSeed, GrokCredentialCatalogError> {
let session =
GrokModelCatalogSession::new(access_token, user_id, email, self.wire_profile.clone())
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?;
let snapshot = self
.client
.fetch(&session)
.await
.map_err(|_| GrokCredentialCatalogError::Upstream)?;
GrokCredentialCatalogSeed::from_snapshot(&snapshot)
}
pub async fn cache_seed(
&self,
account_id: &ProviderAccountId,
seed: GrokCredentialCatalogSeed,
) -> Result<(), GrokCredentialCatalogError> {
let account = self
.repository
.load_current(account_id)
.await
.map_err(|_| GrokCredentialCatalogError::Store)?
.account;
let scope = GrokCatalogScope::for_account(&account)
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?;
self.cache
.replace(GrokPlanCatalog::new(scope, Utc::now(), seed))
.await
.map_err(|_| GrokCredentialCatalogError::Cache)
}
/// 优先读取套餐目录 cache；缺失时才用当前账号所属套餐的有限候选集实时填充。
pub async fn cached_or_refresh_account_catalog(
&self,
backend/crates/providers/xai/src/credential/oidc_verifier.rs
+57 -35
View File
@@ -36,12 +36,19 @@ const MAX_AUDIENCES: usize = 16;
/// 使用官方 JWKS 与 user-info endpoint 的严格生产 verifier。
pub struct ReqwestOidcTokenVerifier {
client: Client,
egress: crate::transport::egress::EgressCache<
Option<gateway_core::account::OutboundProxy>,
OidcEgress,
>,
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
jwks_cache: Mutex<JwksCache>,
cache_ttl: Duration,
}
struct OidcEgress {
client: Client,
jwks_cache: Mutex<JwksCache>,
}
impl ReqwestOidcTokenVerifier {
/// 构建 verifier；使用显式 endpoint policy 与 JWKS cache TTL，只允许访问
/// `auth.x.ai` 官方验证端点。
@@ -53,11 +60,16 @@ impl ReqwestOidcTokenVerifier {
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
cache_ttl: Duration,
) -> Result<Self, crate::transport::GrokReqwestTransportBuildError> {
let client = endpoint_policy.build_oauth_client(Some(REQUEST_TIMEOUT))?;
let client = endpoint_policy.build_oauth_client(Some(REQUEST_TIMEOUT), None)?;
Ok(Self {
client,
egress: crate::transport::egress::EgressCache::with_entry(
None,
OidcEgress {
client,
jwks_cache: Mutex::new(JwksCache::default()),
},
),
endpoint_policy,
jwks_cache: Mutex::new(JwksCache::default()),
cache_ttl,
})
}
@@ -66,12 +78,14 @@ impl ReqwestOidcTokenVerifier {
&self,
context: TokenVerificationContext<'_>,
candidate: TokenCandidate<'_>,
egress: &OidcEgress,
) -> Result<VerificationEvidence, VerificationFailure> {
self.validate_context(&context)?;
match context.flow() {
VerificationFlow::AuthorizationCode => self.verify_id_token(&context, candidate).await,
VerificationFlow::AuthorizationCode => {
self.verify_id_token(&context, candidate, egress).await
}
VerificationFlow::CredentialImport | VerificationFlow::CredentialImportRefreshed => {
self.verify_userinfo(&context, candidate).await
self.verify_userinfo(&context, candidate, egress).await
}
}
}
@@ -119,9 +133,10 @@ impl ReqwestOidcTokenVerifier {
&self,
context: &TokenVerificationContext<'_>,
candidate: TokenCandidate<'_>,
egress: &OidcEgress,
) -> Result<VerificationEvidence, VerificationFailure> {
let subject = self
.verify_signed_id_token(context, &candidate, true)
.verify_signed_id_token(context, &candidate, true, egress)
.await?;
Ok(VerificationEvidence::id_token(subject))
}
@@ -131,6 +146,7 @@ impl ReqwestOidcTokenVerifier {
context: &TokenVerificationContext<'_>,
candidate: &TokenCandidate<'_>,
require_nonce: bool,
egress: &OidcEgress,
) -> Result<String, VerificationFailure> {
let token = candidate
.id_token()
@@ -145,7 +161,7 @@ impl ReqwestOidcTokenVerifier {
return Err(VerificationFailure::Rejected);
}
let kid = header.kid.as_deref().ok_or(VerificationFailure::Rejected)?;
let key = self.decoding_key(context.jwks_uri(), kid).await?;
let key = self.decoding_key(context.jwks_uri(), kid, egress).await?;
let mut validation = Validation::new(Algorithm::ES256);
validation.leeway = 0;
@@ -187,8 +203,11 @@ impl ReqwestOidcTokenVerifier {
&self,
context: &TokenVerificationContext<'_>,
candidate: TokenCandidate<'_>,
egress: &OidcEgress,
) -> Result<VerificationEvidence, VerificationFailure> {
let subject = self.fetch_userinfo_subject(context, &candidate).await?;
let subject = self
.fetch_userinfo_subject(context, &candidate, egress)
.await?;
Ok(VerificationEvidence::user_info(subject))
}
@@ -196,12 +215,13 @@ impl ReqwestOidcTokenVerifier {
&self,
context: &TokenVerificationContext<'_>,
candidate: &TokenCandidate<'_>,
egress: &OidcEgress,
) -> Result<String, VerificationFailure> {
let access_token = candidate.access_token().expose();
if !valid_bearer_token(access_token) {
return Err(VerificationFailure::Rejected);
}
let response = self
let response = egress
.client
.get(context.userinfo_endpoint().clone())
.header(ACCEPT, "application/json")
@@ -232,13 +252,14 @@ impl ReqwestOidcTokenVerifier {
&self,
jwks_uri: &Url,
kid: &str,
egress: &OidcEgress,
) -> Result<DecodingKey, VerificationFailure> {
if !valid_key_id(kid) {
return Err(VerificationFailure::Rejected);
}
// 持锁完成获取，保证冷启动和轮换时每个进程只有一个在途 JWKS 请求。
let mut cache = self.jwks_cache.lock().await;
// 每个出口持锁获取，冷启动和轮换共用单飞；不同出口互不阻塞。
let mut cache = egress.jwks_cache.lock().await;
let cache_is_fresh = cache.uri.as_ref() == Some(jwks_uri)
&& cache
.fetched_at
@@ -249,7 +270,7 @@ impl ReqwestOidcTokenVerifier {
}
// fresh cache 未命中 kid 时只强制刷新一次，以支持官方密钥轮换。
let keys = self.fetch_jwks(jwks_uri).await?;
let keys = self.fetch_jwks(jwks_uri, &egress.client).await?;
cache.replace(jwks_uri.clone(), keys);
return cache
.keys
@@ -259,7 +280,7 @@ impl ReqwestOidcTokenVerifier {
}
// 过期缓存不会在网络或协议失败时作为 stale fallback 使用。
let keys = self.fetch_jwks(jwks_uri).await?;
let keys = self.fetch_jwks(jwks_uri, &egress.client).await?;
cache.replace(jwks_uri.clone(), keys);
cache
.keys
@@ -271,12 +292,12 @@ impl ReqwestOidcTokenVerifier {
async fn fetch_jwks(
&self,
jwks_uri: &Url,
client: &Client,
) -> Result<HashMap<String, CachedJwk>, VerificationFailure> {
if !self.endpoint_policy.validate_jwks(jwks_uri) {
return Err(VerificationFailure::Rejected);
}
let response = self
.client
let response = client
.get(jwks_uri.clone())
.header(ACCEPT, "application/json")
.send()
@@ -304,23 +325,24 @@ impl TokenVerifier for ReqwestOidcTokenVerifier {
candidate: TokenCandidate<'a>,
) -> VerificationFuture<'a> {
Box::pin(async move {
if let Some(proxy) = context.outbound_proxy() {
let client = crate::transport::network::account_proxy_client(
&self.client,
Some(proxy),
Some(REQUEST_TIMEOUT),
)
.map_err(|_| VerificationFailure::Unavailable)?;
let scoped = Self {
client,
endpoint_policy: Arc::clone(&self.endpoint_policy),
jwks_cache: Mutex::new(JwksCache::default()),
cache_ttl: self.cache_ttl,
};
scoped.verify_inner(context, candidate).await
} else {
self.verify_inner(context, candidate).await
}
self.validate_context(&context)?;
let entry = self
.egress
.entry(context.outbound_proxy().cloned())
.ok_or(VerificationFailure::Unavailable)?;
let egress = entry
.get_or_try_init(|| async {
let client = self
.endpoint_policy
.build_oauth_client(Some(REQUEST_TIMEOUT), context.outbound_proxy())
.map_err(|_| VerificationFailure::Unavailable)?;
Ok::<_, VerificationFailure>(OidcEgress {
client,
jwks_cache: Mutex::new(JwksCache::default()),
})
})
.await?;
self.verify_inner(context, candidate, egress).await
})
}
}
backend/crates/providers/xai/src/credential/refresh.rs
+3 -14
View File
@@ -718,9 +718,6 @@ impl GrokCredentialRefreshService {
let account_id = credential.account_id.clone();
let access_expires_at = refreshed_access_token_expiry(tokens.expires_in)
.ok_or(GrokCredentialRefreshError::InvalidRefreshResponse)?;
let access_token = tokens.access_token.clone();
let subject = credential.subject.clone();
let email = credential.email.clone();
let record = match self
.repository
.rotate_oauth_credential(&RotateGrokCredential {
@@ -754,17 +751,9 @@ impl GrokCredentialRefreshService {
Err(error) => return Err(error.into()),
};
if let Ok(seed) = self
.catalog
.fetch_seed(
access_token,
SecretValue::new(subject),
email.map(SecretValue::new),
)
.await
{
let _ = self.catalog.cache_seed(&account_id, seed).await;
}
// 凭据提交后重新读取账号事实；刷新期间切换出口不会改变 credential CAS，
// 但后续目录请求必须使用已提交的当前出口，不能复用旧的 token-only 请求。
let _ = self.catalog.refresh_account_catalog(&account_id).await;
let _ = self
.credential_state
.clear_refresh_backoff(&account_id)
backend/crates/providers/xai/src/credential/verification.rs
+7 -2
View File
@@ -112,7 +112,10 @@ impl<'a> TokenVerificationContext<'a> {
}
}
pub(crate) fn with_outbound_proxy(
/// 绑定 OAuth owner 选择的账号出口，供公开的 [`TokenVerifier`] 实现遵守同一代理合同。
/// 出口只决定路由，不改变 issuer、JWKS、签名和 nonce 的信任校验。
#[must_use]
pub fn with_outbound_proxy(
mut self,
proxy: Option<&'a gateway_core::account::OutboundProxy>,
) -> Self {
@@ -120,7 +123,9 @@ impl<'a> TokenVerificationContext<'a> {
self
}
pub(crate) fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
/// 返回本次验证的显式出口；`None` 表示直连，不得继承环境代理。
#[must_use]
pub fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
self.outbound_proxy
}
backend/crates/providers/xai/src/provider/mod.rs
+3 -4
View File
@@ -12,8 +12,7 @@ use gateway_core::account::{
};
use gateway_core::engine::continuation::ContinuationBinding;
use gateway_core::engine::provider::{
EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration,
ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation,
EventStream, Provider, ProviderCallMetadata, ProviderRequest, ProviderRequestObservation,
ProviderSelectionObservation, ProviderStream,
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt};
@@ -29,8 +28,8 @@ use gateway_core::operation::{
Feature, GenerateRequest, Operation, OperationKind, ProviderSessionState,
};
use gateway_core::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, SupportLevel,
UpstreamModelId,
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderCatalogGeneration,
ProviderKind, ProviderModelCapabilities, SupportLevel, UpstreamModelId,
};
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDefinitionError, WorkerId,
backend/crates/providers/xai/src/transport/egress.rs
+53
View File
@@ -0,0 +1,53 @@
//! 出口状态的有界生命周期；client 初始化与 JWKS 单飞共享同一淘汰边界。
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use tokio::sync::OnceCell;
pub(crate) const MAX_CACHED_EGRESS_STATES: usize = 64;
type Entries<K, V> = VecDeque<(K, Arc<OnceCell<V>>)>;
pub(crate) struct EgressCache<K, V> {
entries: Mutex<Entries<K, V>>,
}
impl<K: PartialEq, V> EgressCache<K, V> {
pub(crate) fn new() -> Self {
Self {
entries: Mutex::new(VecDeque::new()),
}
}
pub(crate) fn with_entry(key: K, value: V) -> Self {
Self {
entries: Mutex::new(VecDeque::from([(
key,
Arc::new(OnceCell::new_with(Some(value))),
)])),
}
}
/// 初始化或 JWKS 验证期间持有 cell，避免淘汰后产生第二个单飞 owner。
/// 推理初始化完成后仅保留 client clone；淘汰不会中断旧出口的在途请求。
pub(crate) fn entry(&self, key: K) -> Option<Arc<OnceCell<V>>> {
let mut entries = self.entries.lock().ok()?;
if let Some(index) = entries.iter().position(|(candidate, _)| candidate == &key) {
let entry = entries.remove(index)?;
let cell = Arc::clone(&entry.1);
entries.push_back(entry);
return Some(cell);
}
if entries.len() == MAX_CACHED_EGRESS_STATES {
// 满载且全部在用时拒绝新出口；不扩容，也不破坏既有出口的单飞。
let idle = entries
.iter()
.position(|(_, value)| Arc::strong_count(value) == 1)?;
entries.remove(idle);
}
let cell = Arc::new(OnceCell::new());
entries.push_back((key, Arc::clone(&cell)));
Some(cell)
}
}
backend/crates/providers/xai/src/transport/mod.rs
+1
View File
@@ -4,6 +4,7 @@ pub(crate) mod canonical;
pub(crate) mod catalog;
pub(crate) mod compaction;
pub(crate) mod config;
pub(crate) mod egress;
pub(crate) mod headers;
mod inference;
pub(crate) mod network;
backend/crates/providers/xai/src/transport/network.rs
+116 -155
View File
@@ -1,7 +1,6 @@
//! Explicit account egress with redirects and business retries disabled.
use gateway_core::diagnostics::{StreamCapture, StreamFormat, TraceContext};
use std::collections::{HashMap, VecDeque};
use std::fmt;
use std::io;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -34,7 +33,9 @@ use gateway_core::error::{ClientVisibleUpstreamError, OpaqueUpstreamValue, Provi
use gateway_core::event::UpstreamHttpVersion;
use gateway_core::upstream::UpstreamSendState;
use super::egress::{EgressCache, MAX_CACHED_EGRESS_STATES};
use super::{GrokQuotaFailureKind, classify_grok_quota_failure, scrub_account_fingerprints};
use gateway_core::account::OutboundProxy;
pub(crate) const OFFICIAL_OAUTH_HOST: &str = "auth.x.ai";
const OFFICIAL_INFERENCE_HOST: &str = "cli-chat-proxy.grok.com";
@@ -167,10 +168,12 @@ pub trait GrokEndpointPolicy: fmt::Debug + Send + Sync {
fn build_oauth_client(
&self,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError>;
fn build_inference_client(
&self,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError>;
fn validate_oauth(&self, url: &Url) -> bool;
fn validate_inference(&self, url: &Url) -> bool;
@@ -187,15 +190,21 @@ impl GrokEndpointPolicy for OfficialGrokEndpointPolicy {
fn build_oauth_client(
&self,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
build_official_client(GrokDnsResolutionPolicy::official_oauth(), timeout)
build_official_client(GrokDnsResolutionPolicy::official_oauth(), timeout, proxy)
}
fn build_inference_client(
&self,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
build_official_client(GrokDnsResolutionPolicy::official_inference(), timeout)
build_official_client(
GrokDnsResolutionPolicy::official_inference(),
timeout,
proxy,
)
}
fn validate_oauth(&self, url: &Url) -> bool {
@@ -229,7 +238,7 @@ impl GrokEndpointPolicy for OfficialGrokEndpointPolicy {
/// 官方 OAuth HTTP transport。只允许 `auth.x.ai:443`。
pub struct ReqwestOAuthTransport {
client: Client,
clients: EgressCache<Option<OutboundProxy>, Client>,
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
}
@@ -242,9 +251,9 @@ impl ReqwestOAuthTransport {
pub fn new(
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
) -> Result<Self, GrokReqwestTransportBuildError> {
let client = endpoint_policy.build_oauth_client(Some(OAUTH_REQUEST_TIMEOUT))?;
let client = endpoint_policy.build_oauth_client(Some(OAUTH_REQUEST_TIMEOUT), None)?;
Ok(Self {
client,
clients: EgressCache::with_entry(None, client),
endpoint_policy,
})
}
@@ -262,18 +271,24 @@ impl fmt::Debug for ReqwestOAuthTransport {
impl OAuthHttpTransport for ReqwestOAuthTransport {
fn execute(&self, request: OAuthHttpRequest) -> TransportFuture<'_> {
let client = self.client.clone();
let endpoint_policy = self.endpoint_policy.clone();
Box::pin(async move {
if !endpoint_policy.validate_oauth(request.url()) {
return Err(TransportFailure::new(TransportFailureKind::NotSent));
}
let client = account_proxy_client(
&client,
request.outbound_proxy.as_ref(),
Some(OAUTH_REQUEST_TIMEOUT),
)
.map_err(|_| TransportFailure::new(TransportFailureKind::NotSent))?;
let entry = self
.clients
.entry(request.outbound_proxy.clone())
.ok_or_else(|| TransportFailure::new(TransportFailureKind::NotSent))?;
let client = entry
.get_or_try_init(|| async {
endpoint_policy.build_oauth_client(
Some(OAUTH_REQUEST_TIMEOUT),
request.outbound_proxy.as_ref(),
)
})
.await
.map_err(|_| TransportFailure::new(TransportFailureKind::NotSent))?;
let mut builder = match request.method() {
HttpMethod::Get => client.get(request.url().clone()),
HttpMethod::Post => client.post(request.url().clone()),
@@ -305,13 +320,14 @@ impl OAuthHttpTransport for ReqwestOAuthTransport {
/// 官方 Grok Responses HTTP SSE transport。
pub struct ReqwestGrokInferenceTransport {
clients: Mutex<BoundInferenceClients>,
clients: EgressCache<GrokSessionBinding, Client>,
unbound_client: Mutex<Option<Client>>,
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
}
impl ReqwestGrokInferenceTransport {
/// 单个进程缓存的账号隔离推理连接池上限。
pub const MAX_CACHED_ACCOUNT_CLIENTS: usize = 64;
pub const MAX_CACHED_ACCOUNT_CLIENTS: usize = MAX_CACHED_EGRESS_STATES;
/// 构建只允许官方 CLI proxy 的生产 transport。
///
@@ -321,47 +337,44 @@ impl ReqwestGrokInferenceTransport {
pub fn new(
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
) -> Result<Self, GrokReqwestTransportBuildError> {
let unbound_client = endpoint_policy.build_inference_client(None)?;
let unbound_client = endpoint_policy.build_inference_client(None, None)?;
Ok(Self {
clients: Mutex::new(BoundInferenceClients::new(unbound_client)),
clients: EgressCache::new(),
unbound_client: Mutex::new(Some(unbound_client)),
endpoint_policy,
})
}
fn client_for(
async fn client_for(
&self,
binding: &GrokSessionBinding,
) -> Result<(Client, GrokInferenceClientCacheStatus), GrokInferenceTransportError> {
let unbound_client = {
let mut clients = self
.clients
.lock()
.map_err(|_| inference_client_pool_unavailable())?;
if let Some(client) = clients.get(binding) {
return Ok((client, GrokInferenceClientCacheStatus::Hit));
}
if binding.outbound_proxy().is_none() {
clients.take_unbound()
} else {
None
}
};
let client = match unbound_client {
Some(client) => client,
None => self
.endpoint_policy
.build_inference_client(None)
.map_err(|_| inference_client_pool_unavailable_for_cache_miss())?,
};
let client = account_proxy_client(&client, binding.outbound_proxy(), None)
.map_err(|_| inference_client_pool_unavailable_for_cache_miss())?;
let client = self
let entry = self
.clients
.lock()
.map_err(|_| inference_client_pool_unavailable_for_cache_miss())?
.insert_if_absent(binding, client);
Ok((client, GrokInferenceClientCacheStatus::Miss))
.entry(binding.clone())
.ok_or_else(inference_client_pool_unavailable_for_cache_miss)?;
let status = if entry.initialized() {
GrokInferenceClientCacheStatus::Hit
} else {
GrokInferenceClientCacheStatus::Miss
};
let client = entry
.get_or_try_init(|| async {
if binding.outbound_proxy().is_none()
&& let Some(client) = self
.unbound_client
.lock()
.map_err(|_| GrokReqwestTransportBuildError::ClientInitialization)?
.take()
{
return Ok(client);
}
self.endpoint_policy
.build_inference_client(None, binding.outbound_proxy())
})
.await
.map_err(|_| inference_client_pool_unavailable_for_cache_miss())?;
Ok((client.clone(), status))
}
}
@@ -396,7 +409,7 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport {
.map(|header| (header.name(), header.value().expose().as_bytes())),
);
trace.capture("upstream.request.body", request.body());
let (client, client_cache_status) = self.client_for(request.binding())?;
let (client, client_cache_status) = self.client_for(request.binding()).await?;
let mut builder = client
.post(request.endpoint().clone())
.body(request.body().to_vec());
@@ -475,62 +488,6 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport {
}
}
struct BoundInferenceClients {
by_binding: HashMap<GrokSessionBinding, Client>,
least_recently_used: VecDeque<GrokSessionBinding>,
unbound_client: Option<Client>,
}
impl BoundInferenceClients {
fn new(unbound_client: Client) -> Self {
Self {
by_binding: HashMap::with_capacity(
ReqwestGrokInferenceTransport::MAX_CACHED_ACCOUNT_CLIENTS,
),
least_recently_used: VecDeque::with_capacity(
ReqwestGrokInferenceTransport::MAX_CACHED_ACCOUNT_CLIENTS,
),
unbound_client: Some(unbound_client),
}
}
fn get(&mut self, binding: &GrokSessionBinding) -> Option<Client> {
if let Some(client) = self.by_binding.get(binding).cloned() {
self.record_use(binding);
return Some(client);
}
None
}
fn take_unbound(&mut self) -> Option<Client> {
self.unbound_client.take()
}
fn insert_if_absent(&mut self, binding: &GrokSessionBinding, client: Client) -> Client {
if let Some(existing) = self.get(binding) {
return existing;
}
self.insert(binding.clone(), client.clone());
client
}
fn insert(&mut self, binding: GrokSessionBinding, client: Client) {
if self.by_binding.len() == ReqwestGrokInferenceTransport::MAX_CACHED_ACCOUNT_CLIENTS
&& let Some(expired_binding) = self.least_recently_used.pop_front()
{
self.by_binding.remove(&expired_binding);
}
self.least_recently_used.push_back(binding.clone());
self.by_binding.insert(binding, client);
}
fn record_use(&mut self, binding: &GrokSessionBinding) {
self.least_recently_used
.retain(|candidate| candidate != binding);
self.least_recently_used.push_back(binding.clone());
}
}
fn inference_client_pool_unavailable() -> GrokInferenceTransportError {
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::Unavailable,
@@ -563,7 +520,7 @@ fn inference_transport_metrics(
/// 官方 Grok CLI proxy 模型目录 GET transport。
pub struct ReqwestGrokModelCatalogTransport {
client: Client,
clients: EgressCache<Option<OutboundProxy>, Client>,
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
}
@@ -576,9 +533,9 @@ impl ReqwestGrokModelCatalogTransport {
pub fn new(
endpoint_policy: Arc<dyn GrokEndpointPolicy>,
) -> Result<Self, GrokReqwestTransportBuildError> {
let client = endpoint_policy.build_inference_client(Some(OAUTH_REQUEST_TIMEOUT))?;
let client = endpoint_policy.build_inference_client(Some(OAUTH_REQUEST_TIMEOUT), None)?;
Ok(Self {
client,
clients: EgressCache::with_entry(None, client),
endpoint_policy,
})
}
@@ -596,7 +553,6 @@ impl fmt::Debug for ReqwestGrokModelCatalogTransport {
impl GrokModelCatalogTransport for ReqwestGrokModelCatalogTransport {
fn execute(&self, request: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> {
let client = self.client.clone();
let endpoint_policy = self.endpoint_policy.clone();
Box::pin(async move {
if !endpoint_policy.validate_model_catalog(request.endpoint()) {
@@ -604,14 +560,27 @@ impl GrokModelCatalogTransport for ReqwestGrokModelCatalogTransport {
GrokModelCatalogTransportErrorKind::Protocol,
));
}
let client = account_proxy_client(
&client,
request.outbound_proxy.as_ref(),
Some(OAUTH_REQUEST_TIMEOUT),
)
.map_err(|_| {
GrokModelCatalogTransportError::new(GrokModelCatalogTransportErrorKind::Protocol)
})?;
let entry = self
.clients
.entry(request.outbound_proxy.clone())
.ok_or_else(|| {
GrokModelCatalogTransportError::new(
GrokModelCatalogTransportErrorKind::Transport,
)
})?;
let client = entry
.get_or_try_init(|| async {
endpoint_policy.build_inference_client(
Some(OAUTH_REQUEST_TIMEOUT),
request.outbound_proxy.as_ref(),
)
})
.await
.map_err(|_| {
GrokModelCatalogTransportError::new(
GrokModelCatalogTransportErrorKind::Protocol,
)
})?;
let mut builder = client.get(request.endpoint().clone());
for header in request.headers() {
builder = builder.header(header.name(), header.value().expose());
@@ -655,7 +624,6 @@ impl GrokModelCatalogTransport for ReqwestGrokModelCatalogTransport {
impl GrokBillingTransport for ReqwestGrokModelCatalogTransport {
fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> {
let client = self.client.clone();
let endpoint_policy = self.endpoint_policy.clone();
Box::pin(async move {
let endpoint = endpoint_policy
@@ -663,12 +631,23 @@ impl GrokBillingTransport for ReqwestGrokModelCatalogTransport {
.ok_or_else(|| {
GrokBillingTransportError::new(GrokBillingTransportErrorKind::Protocol)
})?;
let client = account_proxy_client(
&client,
request.outbound_proxy.as_ref(),
Some(OAUTH_REQUEST_TIMEOUT),
)
.map_err(|_| GrokBillingTransportError::new(GrokBillingTransportErrorKind::Protocol))?;
let entry = self
.clients
.entry(request.outbound_proxy.clone())
.ok_or_else(|| {
GrokBillingTransportError::new(GrokBillingTransportErrorKind::Transport)
})?;
let client = entry
.get_or_try_init(|| async {
endpoint_policy.build_inference_client(
Some(OAUTH_REQUEST_TIMEOUT),
request.outbound_proxy.as_ref(),
)
})
.await
.map_err(|_| {
GrokBillingTransportError::new(GrokBillingTransportErrorKind::Protocol)
})?;
let mut builder = client.get(endpoint);
for header in request.headers() {
builder = builder.header(header.name(), header.value().expose());
@@ -702,6 +681,7 @@ impl GrokBillingTransport for ReqwestGrokModelCatalogTransport {
fn build_official_client(
dns_policy: GrokDnsResolutionPolicy,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
let mut builder = Client::builder()
.redirect(Policy::none())
@@ -714,42 +694,23 @@ fn build_official_client(
.http2_keep_alive_timeout(HTTP2_KEEP_ALIVE_TIMEOUT)
.http2_keep_alive_while_idle(true)
.tcp_nodelay(true)
.https_only(true)
.dns_resolver(Arc::new(StrictDnsResolver::new(dns_policy)?));
if let Some(timeout) = timeout {
builder = builder.timeout(timeout);
}
builder
.build()
.map_err(|_| GrokReqwestTransportBuildError::ClientInitialization)
}
pub(crate) fn account_proxy_client(
base: &Client,
proxy: Option<&gateway_core::account::OutboundProxy>,
timeout: Option<Duration>,
) -> Result<Client, GrokReqwestTransportBuildError> {
let Some(proxy) = proxy else {
return Ok(base.clone());
};
let invalid = || GrokReqwestTransportBuildError::ClientInitialization;
let mut builder = Client::builder()
.no_proxy()
.redirect(Policy::none())
.proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| invalid())?)
.connect_timeout(CONNECT_TIMEOUT)
.pool_idle_timeout(POOL_IDLE_TIMEOUT)
.pool_max_idle_per_host(POOL_MAX_IDLE_PER_HOST)
.http2_adaptive_window(true)
.http2_keep_alive_interval(HTTP2_KEEP_ALIVE_INTERVAL)
.http2_keep_alive_timeout(HTTP2_KEEP_ALIVE_TIMEOUT)
.http2_keep_alive_while_idle(true)
.tcp_nodelay(true)
.https_only(true);
if let Some(proxy) = proxy {
// 保留 reqwest 的代理解析语义（socks5 本地解析、socks5h 远端解析）；
// 只允许官方 host 的直连 resolver 不能用于解析代理端点。
builder = builder.proxy(
reqwest::Proxy::all(proxy.expose_url())
.map_err(|_| GrokReqwestTransportBuildError::ClientInitialization)?,
);
} else {
builder = builder.dns_resolver(Arc::new(StrictDnsResolver::new(dns_policy)?));
}
if let Some(timeout) = timeout {
builder = builder.timeout(timeout);
}
builder.build().map_err(|_| invalid())
builder
.build()
.map_err(|_| GrokReqwestTransportBuildError::ClientInitialization)
}
fn valid_official_url(url: &Url, host: &str, path: Option<&str>) -> bool {
backend/crates/providers/xai/tests/admin.rs
+96 -3
View File
@@ -38,16 +38,109 @@ use gateway_core::task::{
WorkerContribution, WorkerCycleContext, WorkerKind, WorkerRunnable, WorkerTaskError,
};
use provider_xai::{
DiscoveryDocument, GrokOAuthConfig, PendingAuthorization, RedirectUriAllowlist,
DiscoveryDocument, GrokAccountSessionSelector, GrokCredentialFailure, GrokCredentialRepository,
GrokOAuthConfig, GrokSessionBinding, GrokSessionSelector, OfficialGrokEndpointPolicy,
PendingAuthorization, RedirectUriAllowlist, ReqwestGrokModelCatalogTransport, SecretValue,
SelectedGrokSession,
};
use serde_json::{Map, Value, json};
use sha2::{Digest as _, Sha256};
use crate::support::{
MemoryProviderAccountStore, TestSessionAffinity, TestSessionExclusions, create_input,
seed_input, xai_config,
MemoryCooldownPort, MemoryGrokCatalogCache, MemoryProviderAccountStore, TestSessionAffinity,
TestSessionExclusions, account_id, create_input, seed_input, xai_config,
};
#[tokio::test]
async fn account_unavailable_clears_real_selector_cooldowns_only_for_deleted_account() {
let store = MemoryProviderAccountStore::shared();
let cooldowns = Arc::new(MemoryCooldownPort::default());
let ports = ProviderStorePorts::new(
store.clone(),
Arc::new(TestLeases),
Arc::new(TestSessionAffinity),
Arc::new(TestSessionExclusions),
Arc::new(TestCatalogCache::default()),
Arc::new(TestArtifactProfiles),
Arc::new(TestCredentialState),
cooldowns.clone(),
Arc::new(TestRuntimePolicy),
Arc::new(TestOAuthPending::default()),
);
let repository = GrokCredentialRepository::new(store.clone());
let selector = GrokAccountSessionSelector::new(
ProviderKind::new("xai").unwrap(),
repository.clone(),
MemoryGrokCatalogCache::shared(),
Arc::new(crate::support::grok_quota_service(
repository,
Arc::new(
ReqwestGrokModelCatalogTransport::new(Arc::new(OfficialGrokEndpointPolicy))
.unwrap(),
),
)),
ports.leases(),
cooldowns.clone(),
ports.account_feedback(),
);
let bundle = provider_xai::initialize(xai_config(), ports).await.unwrap();
let models = ["grok-4.5", "grok-4.6"].map(|id| UpstreamModelId::new(id).unwrap());
for suffix in ["deleted-cooldown", "retained-cooldown"] {
let input = create_input(suffix, &format!("subject-{suffix}"));
seed_input(&store, &input).await.unwrap();
let session = SelectedGrokSession::new(
input.account_id,
CredentialRevision::new(1).unwrap(),
input.secret.access_token,
SecretValue::new(input.account.subject),
None,
GrokSessionBinding::new(suffix).unwrap(),
(),
)
.unwrap();
// cooldown 必须来自生产 selector 对真实反馈类型的处理，不能由 mock 手工伪造。
selector
.record_failure(
&session,
GrokCredentialFailure::RateLimited {
retry_after: Some(Duration::from_secs(60)),
},
)
.await;
for model in &models {
selector
.record_failure(
&session,
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: model.clone(),
retry_after: None,
},
)
.await;
assert!(
cooldowns
.scoped_cooldown(
session.account_id(),
&ProviderCooldownScope::upstream_model(model.clone())
)
.is_some()
);
}
assert!(cooldowns.cooldown(session.account_id()).is_some());
}
let deleted = account_id("deleted-cooldown");
store.delete_account(&deleted).await.unwrap();
bundle.admin_provider().account_unavailable(&deleted).await;
assert!(cooldowns.cooldown(&deleted).is_none());
let retained = account_id("retained-cooldown");
assert!(cooldowns.cooldown(&retained).is_some());
for model in models {
let scope = ProviderCooldownScope::upstream_model(model);
assert!(cooldowns.scoped_cooldown(&deleted, &scope).is_none());
assert!(cooldowns.scoped_cooldown(&retained, &scope).is_some());
}
}
#[tokio::test]
async fn xai_bundle_exposes_core_admin_and_drains_worker_contributions_once() {
let mut bundle = provider_xai::initialize(xai_config(), provider_ports())
backend/crates/providers/xai/tests/credential/catalog.rs
+1 -21
View File
@@ -19,7 +19,7 @@ use provider_xai::{
GrokCredentialCatalogSeed, GrokCredentialRepository, GrokModelCatalogRequest,
GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind,
GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokPlanCatalog,
GrokQuotaError, SecretValue,
GrokQuotaError,
};
use crate::support::{
@@ -643,26 +643,6 @@ fn seed_rejects_duplicates_and_supports_exact_membership() {
assert!(!seed.permits("grok-4"));
}
#[tokio::test]
async fn fetch_seed_rejects_non_header_safe_identity() {
let (_, repository) = repository_with_accounts(&[]).await;
let service = crate::support::grok_catalog_service(
repository,
QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]),
MemoryGrokCatalogCache::shared(),
);
assert!(matches!(
service
.fetch_seed(
SecretValue::new("access"),
SecretValue::new("非-ascii"),
None,
)
.await,
Err(GrokCredentialCatalogError::InvalidCredentialData)
));
}
#[tokio::test]
async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fields() {
let (store, repository) = repository_with_accounts(&[("quota", "subject-quota")]).await;
backend/crates/providers/xai/tests/credential/oidc_verifier.rs
+129
View File
@@ -23,6 +23,135 @@ use crate::support::loopback_endpoint_policy;
const CACHE_TTL: Duration = Duration::from_secs(60 * 60);
#[tokio::test]
async fn proxy_jwks_cache_evicts_idle_exits_at_the_fixed_bound() {
let origin = MockServer::start().await;
let exit = MockServer::start().await;
let endpoint = endpoints(&origin);
let capacity = provider_xai::ReqwestGrokInferenceTransport::MAX_CACHED_ACCOUNT_CLIENTS;
let key = TestKey::new(36, "bounded-proxy-key");
mount_jwks(&exit, key.jwks(), (capacity + 2) as u64).await;
let verifier = verifier(&endpoint.origin, CACHE_TTL);
let nonce = SecretValue::new("bounded-proxy-nonce");
let algorithms = vec!["ES256".to_owned()];
let access = SecretValue::new("access-token");
let token = SecretValue::new(key.sign(&valid_claims(&nonce), Algorithm::ES256));
let proxies = (0..=capacity)
.map(|index| {
let mut url = Url::parse(&exit.uri()).unwrap();
url.set_username(&format!("synthetic-exit-{index}"))
.unwrap();
gateway_core::account::OutboundProxy::parse(url.as_str()).unwrap()
})
.collect::<Vec<_>>();
for index in (0..=capacity).chain([capacity, 0]) {
verifier
.verify(
auth_context(&endpoint, &algorithms, &nonce)
.with_outbound_proxy(Some(&proxies[index])),
TokenCandidate::new(&access, Some(&token), None),
)
.await
.unwrap();
}
// 65 个出口各取一次，最近出口命中，最旧出口被淘汰后重新获取。
assert_eq!(exit.received_requests().await.unwrap().len(), capacity + 2);
assert!(origin.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn proxy_jwks_cache_single_flights_and_isolates_each_exit() {
let origin = MockServer::start().await;
let exit_a = MockServer::start().await;
let exit_b = MockServer::start().await;
let endpoint = endpoints(&origin);
let proxy_a = gateway_core::account::OutboundProxy::parse(&exit_a.uri()).unwrap();
let proxy_b = gateway_core::account::OutboundProxy::parse(&exit_b.uri()).unwrap();
let nonce = SecretValue::new("proxy-bound-nonce");
let algorithms = vec!["ES256".to_owned()];
let access = SecretValue::new("access-token");
let verifier = verifier(&endpoint.origin, CACHE_TTL);
// 同一 kid 配不同密钥：复用别的出口缓存会直接使验签失败。
for (server, proxy, seed) in [
(&exit_a, Some(&proxy_a), 31),
(&exit_b, Some(&proxy_b), 32),
(&origin, None, 33),
] {
let key = TestKey::new(seed, "shared-kid");
Mock::given(method("GET"))
.and(path("/.well-known/jwks.json"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(key.jwks())
.set_delay(Duration::from_millis(25)),
)
.expect(1)
.mount(server)
.await;
let token = SecretValue::new(key.sign(&valid_claims(&nonce), Algorithm::ES256));
let verify = || {
verifier.verify(
auth_context(&endpoint, &algorithms, &nonce).with_outbound_proxy(proxy),
TokenCandidate::new(&access, Some(&token), None),
)
};
let (first, second) = tokio::join!(verify(), verify());
assert_eq!(first.unwrap().subject(), "user-123");
assert_eq!(second.unwrap().subject(), "user-123");
assert!(verify().await.is_ok(), "同出口后续验证必须命中缓存");
}
}
#[tokio::test]
async fn proxy_jwks_rotation_and_expired_cache_remain_fail_closed() {
for ttl in [CACHE_TTL, Duration::ZERO] {
let origin = MockServer::start().await;
let exit = MockServer::start().await;
let endpoint = endpoints(&origin);
let proxy = gateway_core::account::OutboundProxy::parse(&exit.uri()).unwrap();
let verifier = verifier(&endpoint.origin, ttl);
let old = TestKey::new(34, "old-proxy-key");
let new = TestKey::new(35, "rotated-proxy-key");
mount_jwks_once_then(&exit, old.jwks(), new.jwks()).await;
let nonce = SecretValue::new("rotation-proxy-nonce");
let algorithms = vec!["ES256".to_owned()];
let access = SecretValue::new("access-token");
for key in [&old, &new] {
let token = SecretValue::new(key.sign(&valid_claims(&nonce), Algorithm::ES256));
assert!(
verifier
.verify(
auth_context(&endpoint, &algorithms, &nonce)
.with_outbound_proxy(Some(&proxy)),
TokenCandidate::new(&access, Some(&token), None),
)
.await
.is_ok()
);
}
exit.verify().await;
exit.reset().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(503))
.expect(1)
.mount(&exit)
.await;
// TTL 到期或未知 kid 均不能在获取失败后借用旧缓存。
let rejected_key = if ttl.is_zero() { &new } else { &old };
let token = SecretValue::new(rejected_key.sign(&valid_claims(&nonce), Algorithm::ES256));
assert!(matches!(
verifier
.verify(
auth_context(&endpoint, &algorithms, &nonce).with_outbound_proxy(Some(&proxy)),
TokenCandidate::new(&access, Some(&token), None),
)
.await,
Err(VerificationFailure::Unavailable)
));
assert!(origin.received_requests().await.unwrap().is_empty());
}
}
#[tokio::test]
async fn authorization_code_should_validate_es256_and_cache_jwks() {
let server = MockServer::start().await;
backend/crates/providers/xai/tests/credential/refresh.rs
+154 -4
View File
@@ -7,7 +7,7 @@ use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use chrono::Utc;
use gateway_core::account::{
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState,
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState, OutboundProxy,
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
};
use gateway_core::provider_ports::{
@@ -18,20 +18,170 @@ use gateway_core::provider_ports::{
use provider_xai::{
GrokCredentialCatalogCache, GrokCredentialRecovery, GrokCredentialRecoveryOutcome,
GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService,
GrokCredentialRefresher, GrokCredentialRepository, GrokModelCatalogRequest,
GrokCredentialRefresher, GrokCredentialRepository, GrokEndpointPolicy, GrokModelCatalogRequest,
GrokModelCatalogTransport, GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse,
GrokRefreshFailure, GrokRefreshTokens, SecretValue,
GrokRefreshFailure, GrokRefreshTokens, GrokReqwestTransportBuildError,
OfficialGrokEndpointPolicy, ReqwestGrokModelCatalogTransport, SecretValue,
};
use crate::support::{
MemoryGrokCatalogCache, MemoryProviderAccountStore, create_input, credential_object,
runtime_policy, seed_input,
rejecting_account_proxy, runtime_policy, seed_input,
};
const OFFICIAL_FIXTURE: &[u8] =
include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json");
const OAUTH_BACKOFF_ATTEMPTS: u32 = 5;
#[tokio::test]
async fn refreshed_catalog_uses_committed_account_exit_including_proxy_changes() {
// 0: 保留 A；1: 刷新在途时 A→B；2: 刷新在途时 A→直连。
for selected_exit in 0..3 {
let (proxy_a, capture_a) = rejecting_account_proxy().await;
let (proxy_b, capture_b) = rejecting_account_proxy().await;
let (direct_sink, capture_direct) = rejecting_account_proxy().await;
let current_proxy = [Some(proxy_a.clone()), Some(proxy_b), None][selected_exit].clone();
let store = MemoryProviderAccountStore::shared();
let input = due_input("postcommit-egress");
seed_input(&store, &input).await.unwrap();
store.set_outbound_proxy(&input.account_id, Some(proxy_a.clone()));
let repository = GrokCredentialRepository::new(store.clone());
let refresher = Arc::new(SwitchingProxyRefresher {
store: Arc::clone(&store),
account_id: input.account_id.clone(),
initial_proxy: proxy_a,
current_proxy: current_proxy.clone(),
});
let catalog = Arc::new(crate::support::grok_catalog_service(
repository.clone(),
Arc::new(
ReqwestGrokModelCatalogTransport::new(Arc::new(CatalogExitPolicy { direct_sink }))
.unwrap(),
),
MemoryGrokCatalogCache::shared(),
));
let service = GrokCredentialRefreshService::new(
repository,
refresher,
catalog,
Arc::new(TestRefreshLeases {
available: true,
calls: AtomicUsize::new(0),
}),
Arc::new(CountingCredentialState::default()),
runtime_policy(),
);
let outcomes = tokio::time::timeout(Duration::from_secs(5), service.refresh_due())
.await
.unwrap()
.unwrap();
assert!(matches!(outcomes.as_slice(), [
GrokCredentialRefreshOutcome::Refreshed { credential_revision, .. }
] if credential_revision.get() == 2));
let committed = store.account(&input.account_id).unwrap();
assert_eq!(committed.outbound_proxy(), current_proxy.as_ref());
assert_eq!(committed.revision().get(), 2);
for (index, mut capture) in [capture_a, capture_b, capture_direct]
.into_iter()
.enumerate()
{
if index == selected_exit {
let connect = tokio::time::timeout(Duration::from_secs(5), &mut capture)
.await
.unwrap()
.unwrap();
assert!(connect.starts_with("CONNECT cli-chat-proxy.grok.com:443 HTTP/1.1"));
} else {
assert!(!capture.is_finished(), "目录请求不得到达旧出口或误走直连");
capture.abort();
}
}
}
}
struct SwitchingProxyRefresher {
store: Arc<MemoryProviderAccountStore>,
account_id: ProviderAccountId,
initial_proxy: OutboundProxy,
current_proxy: Option<OutboundProxy>,
}
#[async_trait]
impl GrokCredentialRefresher for SwitchingProxyRefresher {
async fn prepare_cycle(&self) -> Result<(), GrokRefreshFailure> {
Ok(())
}
async fn refresh(&self, _: &SecretValue) -> Result<GrokRefreshTokens, GrokRefreshFailure> {
panic!("refresh service must forward the selected account proxy")
}
async fn refresh_with_proxy(
&self,
_: &SecretValue,
proxy: Option<&OutboundProxy>,
) -> Result<GrokRefreshTokens, GrokRefreshFailure> {
assert_eq!(proxy, Some(&self.initial_proxy));
// 精确合成网络刷新已开始、credential CAS 尚未提交的出口修改。
self.store
.set_outbound_proxy(&self.account_id, self.current_proxy.clone());
assert_eq!(
self.store
.account(&self.account_id)
.unwrap()
.revision()
.get(),
1
);
Ok(success_tokens(Some("new-refresh")))
}
}
#[derive(Debug)]
struct CatalogExitPolicy {
direct_sink: OutboundProxy,
}
impl GrokEndpointPolicy for CatalogExitPolicy {
fn build_oauth_client(
&self,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<reqwest::Client, GrokReqwestTransportBuildError> {
self.build_inference_client(timeout, proxy)
}
fn build_inference_client(
&self,
timeout: Option<Duration>,
proxy: Option<&OutboundProxy>,
) -> Result<reqwest::Client, GrokReqwestTransportBuildError> {
// “直连”也终止于本地 sink，旧实现漏代理时不会触达真实上游。
// 显式代理分支仍使用生产构造器及官方 HTTPS/CONNECT 目标。
OfficialGrokEndpointPolicy
.build_inference_client(timeout, Some(proxy.unwrap_or(&self.direct_sink)))
}
fn validate_oauth(&self, url: &url::Url) -> bool {
OfficialGrokEndpointPolicy.validate_oauth(url)
}
fn validate_inference(&self, url: &url::Url) -> bool {
OfficialGrokEndpointPolicy.validate_inference(url)
}
fn validate_model_catalog(&self, url: &url::Url) -> bool {
OfficialGrokEndpointPolicy.validate_model_catalog(url)
}
fn route_billing(&self, url: &url::Url) -> Option<url::Url> {
OfficialGrokEndpointPolicy.route_billing(url)
}
fn validate_jwks(&self, url: &url::Url) -> bool {
OfficialGrokEndpointPolicy.validate_jwks(url)
}
fn validate_userinfo(&self, url: &url::Url) -> bool {
OfficialGrokEndpointPolicy.validate_userinfo(url)
}
}
struct StaticCatalogTransport;
impl GrokModelCatalogTransport for StaticCatalogTransport {
backend/crates/providers/xai/tests/credential/selector.rs
+2 -63
View File
@@ -12,9 +12,8 @@ use gateway_core::account::{
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderCooldownPort, ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort,
ProviderLeaseRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
ProviderStoreError,
ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderSchedulingLeaseRequest, ProviderSchedulingState, ProviderStoreError,
};
use gateway_core::routing::{
ClientRoutingScope, FrozenAccountScope, RuntimeAccount, RuntimeAccountDirectory,
@@ -1032,66 +1031,6 @@ async fn quota_reset_strategy_uses_provider_reported_earliest_reset() {
assert_eq!(session.account_id(), &account_id("zzz-earlier-reset"));
}
#[tokio::test]
async fn deleting_account_clears_all_account_and_model_scoped_cooldowns() {
// 删除账号后其 account/model scoped cooldown key 全部清除。
let fixture = SelectorFixture::new(&["clear-cooldowns", "available"]).await;
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
let selected = session.account_id().clone();
let model = UpstreamModelId::new("grok-4.5").expect("model");
// 写 account cooldown + model scoped cooldown。
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::RateLimited {
retry_after: Some(Duration::from_secs(60)),
},
)
.await;
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: model.clone(),
retry_after: Some(Duration::from_secs(3600)),
},
)
.await;
assert!(fixture.cooldowns.cooldown(&selected).is_some());
assert!(
fixture
.cooldowns
.scoped_cooldown(
&selected,
&ProviderCooldownScope::upstream_model(model.clone())
)
.is_some()
);
// 删除账号（clear_all 清全部）。
fixture
.cooldowns
.clear_all(&selected)
.await
.expect("clear all cooldowns");
assert!(fixture.cooldowns.cooldown(&selected).is_none());
assert!(
fixture
.cooldowns
.scoped_cooldown(
&selected,
&ProviderCooldownScope::upstream_model(model.clone())
)
.is_none()
);
}
#[tokio::test]
async fn bare_402_cooldown_expires_and_account_recovers_without_persisted_exhaustion() {
// bare 402 只写短期 account cooldown，不持久化
backend/crates/providers/xai/tests/support.rs
+54 -3
View File
@@ -115,7 +115,10 @@ impl LoopbackGrokEndpointPolicy {
&& url.fragment().is_none()
}
fn build_client(timeout: Option<Duration>) -> Result<Client, GrokReqwestTransportBuildError> {
fn build_client(
timeout: Option<Duration>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
let mut builder = Client::builder()
.redirect(Policy::none())
.no_proxy()
@@ -124,6 +127,12 @@ impl LoopbackGrokEndpointPolicy {
if let Some(timeout) = timeout {
builder = builder.timeout(timeout);
}
if let Some(proxy) = proxy {
builder = builder.proxy(
reqwest::Proxy::all(proxy.expose_url())
.map_err(|_| GrokReqwestTransportBuildError::ClientInitialization)?,
);
}
builder
.build()
.map_err(|_| GrokReqwestTransportBuildError::ClientInitialization)
@@ -134,15 +143,17 @@ impl GrokEndpointPolicy for LoopbackGrokEndpointPolicy {
fn build_oauth_client(
&self,
timeout: Option<Duration>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
Self::build_client(timeout)
Self::build_client(timeout, proxy)
}
fn build_inference_client(
&self,
timeout: Option<Duration>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
Self::build_client(timeout)
Self::build_client(timeout, proxy)
}
fn validate_oauth(&self, url: &Url) -> bool {
@@ -193,6 +204,35 @@ pub fn loopback_endpoint_policy(origin: &Url) -> Arc<dyn GrokEndpointPolicy> {
Arc::new(LoopbackGrokEndpointPolicy::for_origin(origin))
}
pub async fn rejecting_account_proxy() -> (
gateway_core::account::OutboundProxy,
tokio::task::JoinHandle<String>,
) {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = gateway_core::account::OutboundProxy::parse(&format!(
"http://user:pass@{}",
listener.local_addr().unwrap()
))
.unwrap();
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let mut header = Vec::new();
while !header.ends_with(b"\r\n\r\n") {
header.push(stream.read_u8().await.unwrap());
assert!(header.len() < 8192);
}
stream
.write_all(
b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
)
.await
.unwrap();
String::from_utf8(header).unwrap()
});
(proxy, task)
}
#[derive(Clone)]
struct StoredAccount {
account: ProviderAccount,
@@ -232,6 +272,16 @@ impl MemoryProviderAccountStore {
.with_scheduling(concurrency_limit, weight);
}
pub fn set_outbound_proxy(
&self,
id: &ProviderAccountId,
proxy: Option<gateway_core::account::OutboundProxy>,
) {
let mut accounts = lock(&self.accounts);
let stored = accounts.get_mut(id).expect("seeded account");
stored.account = stored.account.clone().with_outbound_proxy(proxy);
}
pub fn last_error_message(&self, id: &ProviderAccountId) -> Option<String> {
lock(&self.accounts)
.get(id)
@@ -619,6 +669,7 @@ fn rebuild_account(previous: &ProviderAccount, replacement: AccountReplacement)
)
.with_scheduling(previous.concurrency_limit(), previous.weight())
.with_refresh_schedule(replacement.has_refresh_token, replacement.next_refresh_at)
.with_outbound_proxy(previous.outbound_proxy().cloned())
}
#[derive(Default)]
backend/crates/providers/xai/tests/transport/catalog.rs
+13
View File
@@ -640,6 +640,19 @@ async fn client_should_enforce_hard_limit_for_injected_transport_too() {
));
}
#[test]
fn session_rejects_non_header_safe_identity() {
assert!(matches!(
GrokModelCatalogSession::new(
SecretValue::new("oauth-access"),
SecretValue::new("非-ascii"),
None,
crate::support::xai_wire_profile(),
),
Err(provider_xai::GrokModelCatalogSessionError::InvalidHeaderData)
));
}
#[test]
fn session_debug_should_redact_oauth_and_identity_values() {
let debug = format!("{:?}", session(Some("person@example.com")));
backend/crates/providers/xai/tests/transport/egress.rs
+125
View File
@@ -0,0 +1,125 @@
//! 通过生产辅助请求 transport 观察出口 cache 的真实连接生命周期。
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use gateway_core::account::OutboundProxy;
use provider_xai::{
GrokBillingClient, GrokModelCatalogSession, ReqwestGrokInferenceTransport,
ReqwestGrokModelCatalogTransport, SecretValue,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::task::{JoinHandle, JoinSet};
use url::Url;
use crate::support::{loopback_endpoint_policy, xai_wire_profile};
#[tokio::test]
async fn auxiliary_pool_reuses_connections_and_evicts_the_least_recent_exit() {
let exit = KeepAliveExit::start().await;
let transport = Arc::new(
ReqwestGrokModelCatalogTransport::new(loopback_endpoint_policy(&exit.origin)).unwrap(),
);
let client = GrokBillingClient::new(transport);
let capacity = ReqwestGrokInferenceTransport::MAX_CACHED_ACCOUNT_CLIENTS;
let proxies = (0..=capacity)
.map(|index| {
let mut url = exit.origin.clone();
url.set_username(&format!("synthetic-exit-{index}"))
.unwrap();
OutboundProxy::parse(url.as_str()).unwrap()
})
.collect::<Vec<_>>();
for proxy in &proxies {
fetch_billing(&client, proxy).await;
}
assert_eq!(exit.connections.load(Ordering::SeqCst), capacity + 1);
fetch_billing(&client, proxies.last().unwrap()).await;
assert_eq!(
exit.connections.load(Ordering::SeqCst),
capacity + 1,
"命中出口应复用实际 TCP 连接，而不只是复用构造计数"
);
fetch_billing(&client, &proxies[0]).await;
assert_eq!(
exit.connections.load(Ordering::SeqCst),
capacity + 2,
"最旧出口已淘汰，重新请求必须建立新连接"
);
}
async fn fetch_billing(client: &GrokBillingClient, proxy: &OutboundProxy) {
let session = GrokModelCatalogSession::new(
SecretValue::new("synthetic-access"),
SecretValue::new("synthetic-user"),
None,
xai_wire_profile(),
)
.unwrap()
.with_outbound_proxy(Some(proxy.clone()));
tokio::time::timeout(Duration::from_secs(5), client.fetch(&session))
.await
.unwrap()
.unwrap();
}
struct KeepAliveExit {
origin: Url,
connections: Arc<AtomicUsize>,
task: JoinHandle<()>,
}
impl KeepAliveExit {
async fn start() -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let origin = Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap();
let connections = Arc::new(AtomicUsize::new(0));
let accepted = Arc::clone(&connections);
let task = tokio::spawn(async move {
let mut requests = JoinSet::new();
loop {
tokio::select! {
connection = listener.accept() => {
let (mut socket, _) = connection.unwrap();
accepted.fetch_add(1, Ordering::SeqCst);
requests.spawn(async move {
loop {
let mut header = Vec::new();
while !header.ends_with(b"\r\n\r\n") {
let Ok(byte) = socket.read_u8().await else { return; };
header.push(byte);
assert!(header.len() < 16 * 1024);
}
assert!(header.starts_with(b"GET "));
let body = r#"{"config":{"creditUsagePercent":25}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}",
body.len(),
);
if socket.write_all(response.as_bytes()).await.is_err() {
return;
}
}
});
}
result = requests.join_next(), if !requests.is_empty() => {
result.unwrap().unwrap();
}
}
}
});
Self {
origin,
connections,
task,
}
}
}
impl Drop for KeepAliveExit {
fn drop(&mut self) {
self.task.abort();
}
}
backend/crates/providers/xai/tests/transport/mod.rs
+1
View File
@@ -2,6 +2,7 @@ mod canonical;
mod catalog;
mod compaction;
mod config;
mod egress;
mod headers;
mod network;
mod profile;
backend/crates/providers/xai/tests/transport/network.rs
+130 -33
View File
@@ -24,36 +24,7 @@ use provider_xai::{
ReqwestOAuthTransport, SecretValue,
};
use crate::support::loopback_endpoint_policy;
async fn rejecting_account_proxy() -> (
gateway_core::account::OutboundProxy,
tokio::task::JoinHandle<String>,
) {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = gateway_core::account::OutboundProxy::parse(&format!(
"http://user:pass@{}",
listener.local_addr().unwrap()
))
.unwrap();
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let mut header = Vec::new();
while !header.ends_with(b"\r\n\r\n") {
header.push(stream.read_u8().await.unwrap());
assert!(header.len() < 8192);
}
stream
.write_all(
b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
)
.await
.unwrap();
String::from_utf8(header).unwrap()
});
(proxy, task)
}
use crate::support::{loopback_endpoint_policy, rejecting_account_proxy};
#[tokio::test]
async fn inference_and_oauth_connect_through_each_accounts_authenticated_proxy() {
@@ -234,6 +205,130 @@ async fn inference_transport_should_reuse_one_client_only_within_the_same_bindin
assert_eq!(client_builds.load(Ordering::SeqCst), 2);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn inference_proxy_client_is_built_once_for_concurrent_same_binding() {
let origin_server = MockServer::start().await;
let proxy_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(
ResponseTemplate::new(200).set_body_raw("data: [DONE]\n\n", "text/event-stream"),
)
.expect(2)
.mount(&proxy_server)
.await;
let origin = Url::parse(&origin_server.uri()).unwrap();
let proxy = gateway_core::account::OutboundProxy::parse(&proxy_server.uri()).unwrap();
let builds = Arc::new(AtomicUsize::new(0));
let transport = Arc::new(
ReqwestGrokInferenceTransport::new(Arc::new(CountingEndpointPolicy {
inner: loopback_endpoint_policy(&origin),
inference_client_builds: Arc::clone(&builds),
concurrency: Some(Arc::new(BuildConcurrency::new(Duration::from_millis(100)))),
}))
.unwrap(),
);
let start = Arc::new(tokio::sync::Barrier::new(2));
let tasks = (0..2)
.map(|_| {
let transport = Arc::clone(&transport);
let start = Arc::clone(&start);
let binding = GrokSessionBinding::new("same-proxy-account")
.unwrap()
.with_outbound_proxy(Some(proxy.clone()));
let request = GrokInferenceRequest::new(
origin.join("v1/responses").unwrap(),
Vec::new(),
b"{}".to_vec(),
binding,
);
tokio::spawn(async move {
start.wait().await;
transport
.execute(request)
.await
.unwrap()
.into_body()
.collect::<Vec<_>>()
.await
})
})
.collect::<Vec<_>>();
for task in tasks {
assert!(task.await.unwrap().iter().all(Result::is_ok));
}
// 启动时的直连 client + 一个代理 client；不再先构建并丢弃另一个直连 client。
assert_eq!(builds.load(Ordering::SeqCst), 2);
assert!(origin_server.received_requests().await.unwrap().is_empty());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auxiliary_proxy_clients_single_flight_and_clearing_restores_direct() {
let origin_server = MockServer::start().await;
let exit_a = MockServer::start().await;
let exit_b = MockServer::start().await;
for (server, expected) in [(&origin_server, 1), (&exit_a, 3), (&exit_b, 1)] {
Mock::given(method("GET"))
.and(path("/v1/billing"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(json!({"config":{"creditUsagePercent":25}})),
)
.expect(expected)
.mount(server)
.await;
}
let origin = Url::parse(&origin_server.uri()).unwrap();
let builds = Arc::new(AtomicUsize::new(0));
let transport = Arc::new(
ReqwestGrokModelCatalogTransport::new(Arc::new(CountingEndpointPolicy {
inner: loopback_endpoint_policy(&origin),
inference_client_builds: Arc::clone(&builds),
concurrency: Some(Arc::new(BuildConcurrency::new(Duration::from_millis(100)))),
}))
.unwrap(),
);
let proxy_a = gateway_core::account::OutboundProxy::parse(&exit_a.uri()).unwrap();
let proxy_b = gateway_core::account::OutboundProxy::parse(&exit_b.uri()).unwrap();
let start = Arc::new(tokio::sync::Barrier::new(2));
let tasks = (0..2)
.map(|_| {
let client = GrokBillingClient::new(transport.clone());
let start = Arc::clone(&start);
let proxy = proxy_a.clone();
tokio::spawn(async move {
let session = GrokModelCatalogSession::new(
SecretValue::new("access-token"),
SecretValue::new("user-id"),
None,
crate::support::xai_wire_profile(),
)
.unwrap()
.with_outbound_proxy(Some(proxy));
start.wait().await;
client.fetch(&session).await.unwrap();
})
})
.collect::<Vec<_>>();
for task in tasks {
task.await.unwrap();
}
assert_eq!(builds.load(Ordering::SeqCst), 2, "同出口并发仅初始化一次");
let client = GrokBillingClient::new(transport);
for (proxy, expected_builds) in [(Some(proxy_a), 2), (Some(proxy_b), 3), (None, 3)] {
let session = GrokModelCatalogSession::new(
SecretValue::new("access-token"),
SecretValue::new("user-id"),
None,
crate::support::xai_wire_profile(),
)
.unwrap()
.with_outbound_proxy(proxy);
client.fetch(&session).await.unwrap();
assert_eq!(builds.load(Ordering::SeqCst), expected_builds);
}
}
#[tokio::test]
async fn inference_transport_should_evict_the_least_recent_binding_at_the_fixed_capacity() {
let server = MockServer::start().await;
@@ -1089,22 +1184,24 @@ impl GrokEndpointPolicy for CountingEndpointPolicy {
fn build_oauth_client(
&self,
timeout: Option<Duration>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
self.inner.build_oauth_client(timeout)
self.inner.build_oauth_client(timeout, proxy)
}
fn build_inference_client(
&self,
timeout: Option<Duration>,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, GrokReqwestTransportBuildError> {
self.inference_client_builds.fetch_add(1, Ordering::SeqCst);
let Some(concurrency) = &self.concurrency else {
return self.inner.build_inference_client(timeout);
return self.inner.build_inference_client(timeout, proxy);
};
let active = concurrency.active.fetch_add(1, Ordering::SeqCst) + 1;
concurrency.max_active.fetch_max(active, Ordering::SeqCst);
std::thread::sleep(concurrency.delay);
let result = self.inner.build_inference_client(timeout);
let result = self.inner.build_inference_client(timeout, proxy);
concurrency.active.fetch_sub(1, Ordering::SeqCst);
result
}
backend/migrations/README.md
+5
View File
@@ -43,3 +43,8 @@ export CPR_TEST_REDIS_URL='redis://:<password>@127.0.0.1:6379'
测试自建随机 schema / key 前缀做隔离，但仍应使用专用开发或测试实例，不要连接生产库。
凭据从自己的部署配置或 CI Secret 中读取，不要粘贴未脱敏的 `docker inspect` 输出。
环境变量未设置导致的跳过不算数据库测试通过。
运行包含 `StoreBundle` 初始化的完整集成测试时，两条测试 URL 都须包含密码，且密码满足 Store
启动配置的 48 位十六进制要求；仅能连接数据库并不代表该初始化合同通过。专用服务使用对应测试密码，
并在测试进程中清除 `CPR_DATABASE_URL`、`CPR_REDIS_URL`、`CPR_DATABASE_PASSWORD` 和
`CPR_REDIS_PASSWORD`，避免启动配置被部署环境覆盖。
docs/api.md
+4 -3
View File
@@ -397,9 +397,10 @@ OAuth start 使用：
`refreshToken`。刷新响应中的三个 token 字段均按官方语义独立轮换：返回新值时替换，省略时分别保留
现值。重新授权也保留这些回调保护，但只轮换目标账号的 token。回调地址只承载 `code`/`state`，
不以 host/path 形式作为拒绝条件。
- 账号文件导入和首次 OAuth 创建在 credential 提交后立即尝试一次额度观测。观测失败只记录告警，
不回滚已提交的账号；重新授权和手工或后台 RT 刷新只更新 token，不隐式等同于手工额度刷新，也不更新
既有账号资料或 OAuth principal。
- 账号文件导入和 OAuth complete（包括重新授权）在 credential 提交后后台尝试一次额度观测，不等待
观测完成才返回成功。观测失败只记录告警，不回滚已提交的账号；手工或后台 RT 刷新只更新 token，
不隐式等同于手工额度刷新，也不更新既有账号资料或 OAuth principal。xAI 导入与 OAuth complete
使用相同的提交后观察流程。
- OAuth pending flow 先取得带过期时间的独占 claim，只有账号事务提交成功后才消费。失败会释放 claim，
但上游 authorization code 本身通常只能交换一次；已完成过 token exchange 时应重新创建 OAuth flow。
- `GET /accounts/quota` 只读取最后一次落库快照；`POST /accounts/quota/refresh` 才访问上游。access token
docs/architecture.md
+20 -4
View File
@@ -110,13 +110,17 @@ flowchart LR
| `upstream` | 跨 Engine、Event、Error 与 Provider 共用的 transport 名称、发送状态和不透明上游值 |
| `lifecycle` | 取消信号、连接注册与 drain 合同 |
| `engine` | attempt、发送/提交屏障、执行编排和持久化调用时序 |
| `routing` | 冻结路由事实、请求计划以及运行时快照的表示与编译 |
| `routing` | 冻结路由事实、请求计划、Provider 只读目录合同以及运行时快照的表示与编译 |
| `runtime` | 当前快照的发布、读取、revision 订阅与周期对账任务 |
`event` 通过 `validation` / `upstream` 使用基础值，不再依赖承载原始事件的执行错误；账号值对象和
选择策略通过 `identity` / `account::scope` 使用身份与范围，不依赖路由计划。`routing` 和 `error` 的
既有公开类型路径保留 re-export，定义与 Core 内部使用均归属上述 owner。架构测试约束这些叶子依赖。
Provider 模型能力、目录代次与 `ProviderCatalogPort` 由 `routing::catalog` 定义。快照编译和对账只消费
该只读合同，不反向依赖执行注册表；现有 `ProviderRegistry` 直接实现目录端口，仍只维护一份 Provider
注册集合。已知空目录与查询失败的未知目录保持不同语义，目录替身无需实现请求执行。
`engine::observation` 统一维护单次响应的用量、费用、时间和响应 ID，并负责重试前清理；协调器继续
独占发送、提交、重试和终结顺序。Provider 上报费用优先于本地估算，丢弃的 attempt 不得污染最终计量。
Client Key 费用账本独立累计各次 attempt 的实际费用，不能因请求重试而清空已产生的费用或未知计费状态。
@@ -191,6 +195,9 @@ Core 只理解 `Operation`、能力要求、Provider 候选、稳定错误和 ca
连接池按出口隔离。修改代理推进配置 revision，使后续请求使用新出口；已执行请求可沿原连接完成。
代理变更不推进 `credential_revision`，不会使进行中的令牌刷新因凭据版本冲突而丢失结果。
xAI 的推理连接池继续按账号绑定隔离；OAuth、目录/计费辅助请求在各自 transport 内复用有界出口
client，OIDC 的 JWKS 缓存与单飞归属对应出口状态。自动刷新提交凭据后的目录预热重新读取当前账号
出口，不从 token-only 路径绕过代理；JWKS 过期或获取失败仍不使用 stale fallback。
代理独立保存和测试，通过 `outboundProxyId` 绑定账号；账号保留解析后的 URL 供 Provider 使用。
连接配置改变时在同一事务内同步关联账号并清除旧测试结果，已绑定的代理不可删除。
@@ -282,6 +289,8 @@ Core 负责准入与结算时序，Store 持久化费用账本，Admin 负责限
换号和内部重试复用同一名额，Key 与账号并发上限同时生效；预算拒绝或启动失败释放名额。
完成、失败、取消和断连使用同一结算路径，在结束时以网关请求 ID 幂等累计已取得费用。
执行会话持有结算与准入释放的同一个清理 future；协议等待被取消后，后续驱动或宿主 detached finalizer
继续原清理，不重复结算。只有执行终态、结算与释放均已返回，协议层才可视为 finalized 并放弃清理责任。
缺少用量或价格、无法取得 USD 费用的尝试按零累计，不创建待核账记录，也不会阻断 Key；
内部重试中已经取得的费用仍须累计。发送状态只决定重放是否安全，不能据此推断费用。
请求日志继续保留真实错误、用量与费用来源，账本按零累计不表示上游实际免费。
@@ -300,8 +309,9 @@ HTTP validation
-> Admin use case
-> Provider prepare/verify when needed
-> PostgreSQL transaction + audit
-> invalidate affected Provider-derived facts when needed
-> publish committed runtime snapshot
-> notify Provider-derived caches
-> best-effort observations when needed
```
会改变路由快照或安全配置的 mutation 在同一 PostgreSQL 事务中提交业务事实、推进内部
@@ -312,6 +322,10 @@ HTTP validation
revision；credential 轮换只推进账号自己的 `credential_revision`。Redis 通知用于缩短收敛延迟，
PostgreSQL 周期对账才是正确性基础。
同一 publisher 的管理提交、订阅通知与周期对账共享编译到发布/暂停的临界区，避免旧结果覆盖新授权或
晚到的失败暂停新快照；请求读取不等待刷新锁。对账仍允许持久 revision 回退，不能只比较版本大小。
配置不可确认时暂停新请求；仅目录 generation 变化且编译失败时继续使用旧快照，留待下次对账。
## 8. 状态所有权
| 状态 | 唯一权威 | 说明 |
@@ -320,7 +334,7 @@ PostgreSQL 周期对账才是正确性基础。
| Client Key 金额窗口与费用事件 | PostgreSQL | 准入与幂等结算的权威账本，独立于请求观测与日志保留策略 |
| admission、lease、cooldown、circuit、会话亲和、continuation、OAuth pending、目录 cache | Redis | 可重建、可过期的协调状态 |
| 日志、OAuth 恢复记录、在线更新状态、备份暂存 | `.runtime/` | 部署节点本地运行文件 |
| 重置卡库存与消费结果 | OpenAI upstream | 后端不建立本地卡库存；前端只保留当前浏览器会话的最近查询 |
| 重置卡库存与消费结果 | OpenAI upstream | 后端不建立本地卡库存；前端按账号保留当前浏览器会话的最近查询、未决消费幂等键与发送锁，展开行卸载不会清空 |
| Provider 公开模型与请求画像 | Provider/runtime cache | 由官方目录或发布源刷新，不写成第二份业务配置 |
| Windows 安装包临时直链 | Host 进程内短缓存 | 按需解析、严格校验、到期前丢弃；不写 PostgreSQL/Redis，也不代理包字节 |
@@ -346,7 +360,9 @@ credential 与 quota 是两组独立事实：credential refresh 不等于 quota
原工具类型及对应 item ID，`call_id` 始终保持配对。转换仅处理协议字段，超限或转换失败终止流。
默认 `store: false` 的续接使用现有会话 owner 重放完整历史；上一轮要求上游存储且指令未变时，
原生续接移除本轮 `instructions`，避免 Grok 拒绝与 `previous_response_id` 同传。
- 新账号导入和首次 OAuth 在 credential 提交后尽力读取一次额度；失败只留下观测，不回滚账号事务。
- 账号导入和 OAuth complete（包括重新授权）在 credential 提交、Provider 事实失效及快照发布后，
由 Admin 共用流程后台读取一次额度；不等待观测完成才返回管理请求，失败记录告警但不回滚账号事务。
手工和后台 credential refresh 仍不隐式刷新 quota。
- quota refresh、正常推理返回的 rate-limit headers 和后台健康任务汇入同一额度事实；套餐只用于展示与
目录 cache 隔离，不创建套餐专属状态机。
frontend/src/views/accounts/composables/useAccountConnectionTest.ts
+34 -29
View File
@@ -8,6 +8,7 @@ import { getAccountModels, refreshAccountModels } from '@/api'
import { API_BASE_URL } from '@/api/constants'
import { toast } from '@/components/base/BaseToast'
import { useIdSet } from '@/composables/useIdSet'
import { useRequestState } from '@/composables/useRequestState'
import { errorMessage, withMinimumDuration } from '@/utils/async'
import { formatDateTime, formatTime } from '@/utils/date'
@@ -162,8 +163,10 @@ export function useAccountConnectionTest(options: { reload: () => Promise<unknow
const connectionTestFinishedAt = shallowRef('')
const connectionTestDurationMs = shallowRef<number | null>(null)
const testingConnections = useIdSet<string>()
const loadingConnectionTestModels = shallowRef(false)
const refreshingConnectionTestModels = shallowRef(false)
const modelsRequest = useRequestState()
const modelsRequestMode = shallowRef<'load' | 'refresh'>('load')
const loadingConnectionTestModels = computed(() => modelsRequest.loading.value && modelsRequestMode.value === 'load')
const refreshingConnectionTestModels = computed(() => modelsRequest.loading.value && modelsRequestMode.value === 'refresh')
const connectionTestSelectedModel = shallowRef('')
const connectionTestModelOptions = ref<ConnectionTestModelOption[]>([])
const connectionTestStreamUrl = shallowRef<string>()
@@ -396,25 +399,39 @@ export function useAccountConnectionTest(options: { reload: () => Promise<unknow
clearConnectionTestRun()
}
async function loadConnectionTestModels(account = testingAccount.value) {
if (!account?.id)
async function loadConnectionTestModels(account = testingAccount.value, refresh = false) {
if (!account?.id || !showConnectionTestModal.value || account.id !== testingAccount.value?.id)
return
loadingConnectionTestModels.value = true
// 初次查询与主动刷新属于同一弹窗会话，旧请求不得覆盖新账号或新一轮查询。
const requestId = modelsRequest.start()
modelsRequestMode.value = refresh ? 'refresh' : 'load'
connectionTestError.value = ''
try {
const result = await getAccountModels({ accountId: account.id })
applyConnectionTestModels(result)
if (!connectionTestSelectedModel.value) {
const result = await (refresh ? refreshAccountModels : getAccountModels)({ accountId: account.id })
if (!modelsRequest.isCurrent(requestId))
return
applyConnectionTestModels(result, refresh)
if (refresh) {
toast.success(`已刷新 ${connectionTestModelOptions.value.length} 个上游模型`)
}
else if (!connectionTestSelectedModel.value) {
connectionTestError.value = '没有可测试模型'
}
}
catch (error: unknown) {
connectionTestError.value = errorMessage(error, '加载测试模型失败')
connectionTestModelOptions.value = []
connectionTestSelectedModel.value = ''
if (!modelsRequest.isCurrent(requestId))
return
connectionTestError.value = errorMessage(error, refresh ? '刷新上游模型失败' : '加载测试模型失败')
if (refresh) {
toast.error(connectionTestError.value)
}
else {
connectionTestModelOptions.value = []
connectionTestSelectedModel.value = ''
}
}
finally {
loadingConnectionTestModels.value = false
modelsRequest.finish(requestId)
}
}
@@ -435,22 +452,9 @@ export function useAccountConnectionTest(options: { reload: () => Promise<unknow
}
async function handleRefreshConnectionTestModels(account = testingAccount.value) {
if (!account?.id || refreshingConnectionTestModels.value)
if (modelsRequest.loading.value)
return
refreshingConnectionTestModels.value = true
connectionTestError.value = ''
try {
const result = await refreshAccountModels({ accountId: account.id })
applyConnectionTestModels(result, true)
toast.success(`已刷新 ${connectionTestModelOptions.value.length} 个上游模型`)
}
catch (error: unknown) {
connectionTestError.value = errorMessage(error, '刷新上游模型失败')
toast.error(connectionTestError.value)
}
finally {
refreshingConnectionTestModels.value = false
}
await loadConnectionTestModels(account, true)
}
async function handleTestConnection(account = testingAccount.value) {
@@ -524,11 +528,12 @@ export function useAccountConnectionTest(options: { reload: () => Promise<unknow
failConnectionTest('测试连接已断开')
})
watch(showConnectionTestModal, (open) => {
watch([showConnectionTestModal, () => testingAccount.value?.id], ([open]) => {
modelsRequest.invalidate()
if (!open) {
abortConnectionTest()
}
})
}, { flush: 'sync' })
onBeforeUnmount(() => {
abortConnectionTest()
frontend/src/views/accounts/composables/useAccountMutations.ts
+3 -3
View File
@@ -23,7 +23,7 @@ export function useAccountMutations(options: {
accounts: Ref<AccountRow[]>
selectedIds: Ref<Set<string>>
reload: () => Promise<unknown>
replaceAccount: (account: AccountRow) => boolean
replaceAccount: (account: AccountRow) => Promise<boolean>
}) {
const loadAccounts = options.reload
const { downloadJson } = useDownload()
@@ -174,7 +174,7 @@ export function useAccountMutations(options: {
await refreshingQuotaAccounts.run(accountId, async () => {
try {
const result = await withMinimumDuration(() => refreshAccountQuota({ accountId }))
const remainsVisible = options.replaceAccount(result.account)
const remainsVisible = await options.replaceAccount(result.account)
if (!remainsVisible) {
const selectedIds = new Set(options.selectedIds.value)
selectedIds.delete(accountId)
@@ -192,7 +192,7 @@ export function useAccountMutations(options: {
await recoveringAccounts.run(accountId, async () => {
try {
const result = await withMinimumDuration(() => recoverAccount({ accountId }))
const remainsVisible = options.replaceAccount(result.account)
const remainsVisible = await options.replaceAccount(result.account)
if (!remainsVisible) {
const selectedIds = new Set(options.selectedIds.value)
selectedIds.delete(accountId)
frontend/src/views/accounts/composables/useAccountResetCredits.ts
+114 -103
View File
@@ -1,5 +1,5 @@
import type { Account, AccountResetCredit } from '@/api'
import { computed, shallowRef, watch } from 'vue'
import { computed, shallowReactive, shallowRef, watch } from 'vue'
import {
consumeAccountResetCredit,
@@ -23,23 +23,75 @@ interface ResetCreditsSnapshot {
availableCount: number
}
// 只保留本次前端会话中由用户主动查询得到的结果，避免账号列表渲染触发上游请求。
const snapshotsByAccountId = new Map<string, ResetCreditsSnapshot>()
interface ResetCreditsSession {
accountId: string
snapshot: ResetCreditsSnapshot | null
pendingOperation: PendingResetCreditOperation | null
consuming: boolean
loading: boolean
loadError: string
loadSequence: number
accountUpdatedListeners: Set<(account: Account) => void>
}
// 库存仍以主动查询的上游结果为准；未决操作和消费锁必须跨展开行卸载存续。
const sessionsByAccountId = new Map<string, ResetCreditsSession>()
function getResetCreditsSession(accountId: string) {
let session = sessionsByAccountId.get(accountId)
if (!session) {
session = shallowReactive<ResetCreditsSession>({
accountId,
snapshot: null,
pendingOperation: null,
consuming: false,
loading: false,
loadError: '',
loadSequence: 0,
accountUpdatedListeners: new Set(),
})
sessionsByAccountId.set(accountId, session)
}
return session
}
async function loadSessionCredits(session: ResetCreditsSession) {
const sequence = ++session.loadSequence
session.loading = true
session.loadError = ''
try {
const result = await getAccountResetCredits({ accountId: session.accountId })
if (sequence !== session.loadSequence)
return
session.snapshot = {
credits: result.credits,
availableCount: Math.max(0, result.availableCount),
}
}
catch (error: unknown) {
if (sequence === session.loadSequence)
session.loadError = errorMessage(error, '重置卡查询失败')
}
finally {
if (sequence === session.loadSequence)
session.loading = false
}
}
export function useAccountResetCredits(options: {
accountId: () => string
onAccountUpdated: (account: Account) => void
}) {
const credits = shallowRef<AccountResetCredit[]>([])
const availableCount = shallowRef(0)
const hasSnapshot = shallowRef(false)
const loading = shallowRef(false)
const consuming = shallowRef(false)
const loadError = shallowRef('')
const session = shallowRef(getResetCreditsSession(options.accountId()))
const credits = computed(() => session.value.snapshot?.credits ?? [])
const availableCount = computed(() => session.value.snapshot?.availableCount ?? 0)
const hasSnapshot = computed(() => session.value.snapshot !== null)
const loading = computed(() => session.value.loading)
const consuming = computed(() => session.value.consuming)
const loadError = computed(() => session.value.loadError)
const showConfirm = shallowRef(false)
const selectedCreditId = shallowRef('')
const pendingOperation = shallowRef<PendingResetCreditOperation | null>(null)
let loadSequence = 0
const pendingOperation = computed(() => session.value.pendingOperation)
const availableCredits = computed(() =>
credits.value.filter(credit => credit.status === 'available'),
@@ -55,12 +107,9 @@ export function useAccountResetCredits(options: {
ambiguous.value || (availableCount.value > 0 && selectedCredit.value !== undefined),
)
function reconcileSelectedCredit(accountId: string) {
if (accountId !== options.accountId())
return
function reconcileSelectedCredit() {
const operation = pendingOperation.value
if (operation?.accountId === accountId) {
if (operation) {
selectedCreditId.value = operation.creditId
return
}
@@ -69,24 +118,6 @@ export function useAccountResetCredits(options: {
selectedCreditId.value = ''
}
function applySnapshot(accountId: string, snapshot: ResetCreditsSnapshot) {
snapshotsByAccountId.set(accountId, snapshot)
if (accountId !== options.accountId())
return
credits.value = snapshot.credits
availableCount.value = snapshot.availableCount
hasSnapshot.value = true
reconcileSelectedCredit(accountId)
}
function restoreSnapshot(accountId: string) {
const snapshot = snapshotsByAccountId.get(accountId)
credits.value = snapshot?.credits ?? []
availableCount.value = snapshot?.availableCount ?? 0
hasSnapshot.value = snapshot !== undefined
reconcileSelectedCredit(accountId)
}
function selectCredit(creditId: string) {
if (pendingOperation.value || consuming.value)
return
@@ -95,77 +126,52 @@ export function useAccountResetCredits(options: {
selectedCreditId.value = creditId
}
function applyConfirmedConsumption(operation: PendingResetCreditOperation) {
const snapshot = snapshotsByAccountId.get(operation.accountId)
function applyConfirmedConsumption(target: ResetCreditsSession, operation: PendingResetCreditOperation) {
const snapshot = target.snapshot
if (!snapshot)
return
const creditIndex = snapshot.credits.findIndex(credit => credit.id === operation.creditId)
const nextCredits = creditIndex < 0
? snapshot.credits
: snapshot.credits.filter((_, index) => index !== creditIndex)
applySnapshot(operation.accountId, {
target.snapshot = {
credits: nextCredits,
availableCount: Math.max(0, snapshot.availableCount - 1),
})
}
async function loadCredits() {
const accountId = options.accountId()
const sequence = ++loadSequence
loading.value = true
loadError.value = ''
try {
const result = await getAccountResetCredits({ accountId })
if (sequence !== loadSequence || accountId !== options.accountId())
return
applySnapshot(accountId, {
credits: result.credits,
availableCount: Math.max(0, result.availableCount),
})
}
catch (error: unknown) {
if (sequence !== loadSequence || accountId !== options.accountId())
return
loadError.value = errorMessage(error, '重置卡查询失败')
}
finally {
if (sequence === loadSequence)
loading.value = false
}
}
function requestConsume() {
const accountId = options.accountId()
const existing = pendingOperation.value
if (!existing || existing.accountId !== accountId) {
if (availableCount.value <= 0)
return
const credit = selectedCredit.value
if (!credit)
return
pendingOperation.value = {
accountId,
creditId: credit.id,
credit,
redeemRequestId: generateRedeemRequestId(),
hasTransportFailure: false,
}
}
if (consuming.value || loading.value || !canRequestConsume.value)
return
showConfirm.value = true
}
function cancelConsume() {
if (consuming.value)
return
showConfirm.value = false
if (!pendingOperation.value?.hasTransportFailure)
pendingOperation.value = null
}
async function confirmConsume(): Promise<boolean> {
const operation = pendingOperation.value
if (!operation || consuming.value)
const target = session.value
if (!showConfirm.value || target.consuming || target.loading)
return false
consuming.value = true
const credit = selectedCredit.value
// 确认发送时才建立操作；仅打开或取消确认弹窗不占用账号的消费状态。
const operation = target.pendingOperation ?? (credit && availableCount.value > 0
? {
accountId: target.accountId,
creditId: credit.id,
credit,
redeemRequestId: generateRedeemRequestId(),
hasTransportFailure: false,
}
: null)
if (!operation)
return false
target.pendingOperation = operation
target.consuming = true
try {
const result = await consumeAccountResetCredit({
accountId: operation.accountId,
@@ -174,23 +180,23 @@ export function useAccountResetCredits(options: {
})
const confirmed = result.code === 'reset'
|| (result.code === 'already_redeemed' && operation.hasTransportFailure)
pendingOperation.value = null
showConfirm.value = false
target.pendingOperation = null
if (!confirmed) {
toast.error(resetResultMessage(result.code))
await loadCredits()
await loadSessionCredits(target)
return false
}
const successMessage = result.code === 'already_redeemed'
? '上次重置已完成'
: '额度已重置'
applyConfirmedConsumption(operation)
await loadCredits()
applyConfirmedConsumption(target, operation)
await loadSessionCredits(target)
try {
const quota = await refreshAccountQuota({ accountId: operation.accountId })
if (operation.accountId === options.accountId())
options.onAccountUpdated(quota.account)
// 原组件可能已经卸载或换号，只通知仍订阅该账号的实例。
for (const listener of target.accountUpdatedListeners)
listener(quota.account)
toast.success(successMessage)
}
catch (error: unknown) {
@@ -202,40 +208,45 @@ export function useAccountResetCredits(options: {
return true
}
catch (error: unknown) {
showConfirm.value = false
if (isAmbiguousConsumeError(error)) {
pendingOperation.value = {
target.pendingOperation = {
...operation,
hasTransportFailure: true,
}
toast.warning('消费结果暂不确定；重试会复用同一个请求标识', { duration: 5000 })
}
else {
pendingOperation.value = null
target.pendingOperation = null
toast.error(errorMessage(error, '额度重置失败'))
await loadCredits()
await loadSessionCredits(target)
}
return false
}
finally {
consuming.value = false
target.consuming = false
}
}
watch(
options.accountId,
(accountId) => {
loadSequence += 1
pendingOperation.value = null
(accountId, _, onCleanup) => {
const target = getResetCreditsSession(accountId)
session.value = target
selectedCreditId.value = ''
restoreSnapshot(accountId)
loading.value = false
loadError.value = ''
showConfirm.value = false
const listener = (account: Account) => options.onAccountUpdated(account)
target.accountUpdatedListeners.add(listener)
onCleanup(() => target.accountUpdatedListeners.delete(listener))
},
{ immediate: true },
{ immediate: true, flush: 'sync' },
)
watch([credits, pendingOperation], reconcileSelectedCredit, { immediate: true, flush: 'sync' })
watch(consuming, (isConsuming) => {
if (!isConsuming)
showConfirm.value = false
}, { flush: 'sync' })
return {
credits,
availableCredits,
@@ -249,7 +260,7 @@ export function useAccountResetCredits(options: {
loadError,
ambiguous,
showConfirm,
loadCredits,
loadCredits: () => loadSessionCredits(session.value),
selectCredit,
requestConsume,
cancelConsume,
frontend/src/views/accounts/composables/useAccountsQuery.ts
+6 -46
View File
@@ -8,16 +8,6 @@ import { usePagedQuery } from '@/composables/usePagedQuery'
import { errorMessage } from '@/utils/async'
type AccountRow = Awaited<ReturnType<typeof getAccounts>>['items'][number]
type AccountSummary = Awaited<ReturnType<typeof getAccounts>>['summary']
type AccountSummaryStatusKey = Exclude<keyof AccountSummary, 'total'>
const summaryKeyByStatus = {
normal: 'normal',
quota_exhausted: 'quotaExhausted',
rate_limited: 'rateLimited',
disabled: 'disabled',
error: 'error',
} as const satisfies Record<AccountRow['status'], AccountSummaryStatusKey>
export function useAccountsQuery() {
const searchQuery = shallowRef('')
@@ -78,42 +68,12 @@ export function useAccountsQuery() {
void query.execute()
}
function replaceAccount(updated: AccountRow) {
const index = query.items.value.findIndex(account => account.id === updated.id)
if (index < 0)
return false
const current = query.items.value[index]
if (!current)
return false
updateAccountSummary(current.status, updated.status)
if (statusQuery.value && statusQuery.value !== updated.status) {
query.items.value = query.items.value.filter(account => account.id !== updated.id)
query.total.value = Math.max(0, query.total.value - 1)
return false
}
const accounts = [...query.items.value]
accounts[index] = updated
query.items.value = accounts
return true
}
function updateAccountSummary(
previousStatus: AccountRow['status'],
nextStatus: AccountRow['status'],
) {
if (previousStatus === nextStatus)
return
const previousKey = summaryKeyByStatus[previousStatus]
const nextKey = summaryKeyByStatus[nextStatus]
const summary = { ...accountSummary.value }
summary[previousKey] = Math.max(0, summary[previousKey] - 1)
summary[nextKey] += 1
accountSummary.value = summary
async function replaceAccount(updated: AccountRow) {
// 状态变更可能影响筛选、排序及全局概览，统一回读并复用分页查询的末页回退。
query.invalidate()
if (!await query.execute())
return true // 回读失败或被新查询取代时，不依据旧页面取消选择。
return query.items.value.some(account => account.id === updated.id)
}
watchDebounced(
frontend/src/views/accounts/index.vue
+1 -1
View File
@@ -295,7 +295,7 @@ const {
<AccountQuotaPanel
:account="row"
:refreshing="refreshingQuotaAccountIds.has(row.id)"
@account-updated="replaceAccount"
@account-updated="void replaceAccount($event)"
@refresh-quota="handleRefreshQuota"
/>
<AccountUsagePanel :account="row" />