From bfca754714efa78cd098cac8504466636195d12e Mon Sep 17 00:00:00 2001 From: Zyy Date: Sat, 12 Sep 2026 15:16:20 +0800 Subject: [PATCH] fix: resolve maintainability audit findings 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 --- backend/apps/gateway/tests/main.rs | 25 +- .../src/model/provider_credentials.rs | 52 +- .../crates/gateway-admin/src/ports/store.rs | 9 +- .../crates/gateway-admin/src/use_case/mod.rs | 41 +- .../gateway-admin/src/use_case/openai.rs | 90 +--- .../crates/gateway-admin/src/use_case/xai.rs | 45 +- .../gateway-admin/tests/use_case/accounts.rs | 53 +- .../gateway-admin/tests/use_case/mod.rs | 19 +- .../gateway-admin/tests/use_case/openai.rs | 24 +- .../gateway-admin/tests/use_case/xai.rs | 117 ++++- backend/crates/gateway-api/tests/admin/mod.rs | 19 +- .../tests/openai/responses/websocket/mod.rs | 384 ++++++++++++++- .../gateway-core/src/engine/execution.rs | 109 ++--- .../gateway-core/src/engine/provider.rs | 129 +---- backend/crates/gateway-core/src/lib.rs | 2 +- .../gateway-core/src/routing/catalog.rs | 86 ++++ .../crates/gateway-core/src/routing/mod.rs | 5 + .../gateway-core/src/routing/snapshot.rs | 31 +- .../crates/gateway-core/src/runtime/mod.rs | 102 ++-- .../gateway-core/tests/engine/coordinator.rs | 8 +- .../gateway-core/tests/engine/execution.rs | 381 ++++++++++++++- .../gateway-core/tests/engine/provider.rs | 85 +++- .../gateway-core/tests/routing/snapshot.rs | 155 +++--- .../crates/gateway-core/tests/runtime/mod.rs | 453 +++++++++++++++++- .../provider_accounts/admin_adapter.rs | 81 ---- .../src/postgres/provider_accounts/mod.rs | 5 +- .../tests/postgres/provider_accounts/mod.rs | 80 +--- .../openai/src/credential/catalog.rs | 3 +- .../providers/openai/src/provider/mod.rs | 6 +- .../openai/tests/transport/http_client.rs | 43 +- backend/crates/providers/xai/Cargo.toml | 2 +- .../providers/xai/src/credential/catalog.rs | 39 +- .../xai/src/credential/oidc_verifier.rs | 92 ++-- .../providers/xai/src/credential/refresh.rs | 17 +- .../xai/src/credential/verification.rs | 9 +- .../crates/providers/xai/src/provider/mod.rs | 7 +- .../providers/xai/src/transport/egress.rs | 53 ++ .../crates/providers/xai/src/transport/mod.rs | 1 + .../providers/xai/src/transport/network.rs | 271 +++++------ backend/crates/providers/xai/tests/admin.rs | 99 +++- .../providers/xai/tests/credential/catalog.rs | 22 +- .../xai/tests/credential/oidc_verifier.rs | 129 +++++ .../providers/xai/tests/credential/refresh.rs | 158 +++++- .../xai/tests/credential/selector.rs | 65 +-- backend/crates/providers/xai/tests/support.rs | 57 ++- .../providers/xai/tests/transport/catalog.rs | 13 + .../providers/xai/tests/transport/egress.rs | 125 +++++ .../providers/xai/tests/transport/mod.rs | 1 + .../providers/xai/tests/transport/network.rs | 163 +++++-- backend/migrations/README.md | 5 + docs/api.md | 7 +- docs/architecture.md | 24 +- .../composables/useAccountConnectionTest.ts | 63 +-- .../composables/useAccountMutations.ts | 6 +- .../composables/useAccountResetCredits.ts | 217 +++++---- .../accounts/composables/useAccountsQuery.ts | 52 +- frontend/src/views/accounts/index.vue | 2 +- 57 files changed, 2977 insertions(+), 1364 deletions(-) create mode 100644 backend/crates/gateway-core/src/routing/catalog.rs create mode 100644 backend/crates/providers/xai/src/transport/egress.rs create mode 100644 backend/crates/providers/xai/tests/transport/egress.rs diff --git a/backend/apps/gateway/tests/main.rs b/backend/apps/gateway/tests/main.rs index a2b2023bc..c99574358 100644 --- a/backend/apps/gateway/tests/main.rs +++ b/backend/apps/gateway/tests/main.rs @@ -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 { let mut result = BTreeSet::new(); collect_rust_files(root, root, &mut result); diff --git a/backend/crates/gateway-admin/src/model/provider_credentials.rs b/backend/crates/gateway-admin/src/model/provider_credentials.rs index d88ffb823..8f03248b9 100644 --- a/backend/crates/gateway-admin/src/model/provider_credentials.rs +++ b/backend/crates/gateway-admin/src/model/provider_credentials.rs @@ -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, - pub account_id: ProviderAccountId, -} - -/// Provider credential 列表查询。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CredentialListQuery { - pub credential_state: Option, - pub enabled: Option, - pub window: CredentialListWindow, -} - -/// Credential 状态筛选。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum CredentialStateFilter { - Exact(CredentialState), - AnyOf(Vec), -} - -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, - page_size: PageSize, - }, -} - -/// Provider credential 列表。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CredentialPage { - pub config_revision: Revision, - pub items: Vec, - pub next_cursor: Option, -} - /// Provider credential 详情。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct CredentialDetails { diff --git a/backend/crates/gateway-admin/src/ports/store.rs b/backend/crates/gateway-admin/src/ports/store.rs index 94f169957..3026a674c 100644 --- a/backend/crates/gateway-admin/src/ports/store.rs +++ b/backend/crates/gateway-admin/src/ports/store.rs @@ -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>; - async fn list_credentials( - &self, - provider_kind: &gateway_core::routing::ProviderKind, - query: CredentialListQuery, - ) -> AdminStoreResult; - async fn credential_details( &self, provider_kind: &gateway_core::routing::ProviderKind, diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index 6afc9288a..7dc40c274 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -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, + 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, diff --git a/backend/crates/gateway-admin/src/use_case/openai.rs b/backend/crates/gateway-admin/src/use_case/openai.rs index cf2bc1e59..3722e0bbb 100644 --- a/backend/crates/gateway-admin/src/use_case/openai.rs +++ b/backend/crates/gateway-admin/src/use_case/openai.rs @@ -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; - async fn details( - &self, - account_id: &ProviderAccountId, - ) -> Result; 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 { - 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 { - 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) } diff --git a/backend/crates/gateway-admin/src/use_case/xai.rs b/backend/crates/gateway-admin/src/use_case/xai.rs index b042ce6a4..babfae0cf 100644 --- a/backend/crates/gateway-admin/src/use_case/xai.rs +++ b/backend/crates/gateway-admin/src/use_case/xai.rs @@ -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; 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 { - 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) } diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index ea51306a0..bba8ae933 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -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>, import_account_ids: Mutex>, quota_requests: Mutex>, + quota_started: tokio::sync::Notify, + quota_gate: Mutex>>, quota: Mutex, quota_refresh_account: Mutex, AccountRecord)>>, current_credential_revision: Mutex, @@ -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 { 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 { - 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, diff --git a/backend/crates/gateway-admin/tests/use_case/mod.rs b/backend/crates/gateway-admin/tests/use_case/mod.rs index b5b693b50..49d79bab5 100644 --- a/backend/crates/gateway-admin/tests/use_case/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/mod.rs @@ -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 { - Err(unavailable("credentials")) - } - async fn credential_details( &self, _: &ProviderKind, diff --git a/backend/crates/gateway-admin/tests/use_case/openai.rs b/backend/crates/gateway-admin/tests/use_case/openai.rs index bee8fff33..45b19d51a 100644 --- a/backend/crates/gateway-admin/tests/use_case/openai.rs +++ b/backend/crates/gateway-admin/tests/use_case/openai.rs @@ -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, store: Arc) .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"), diff --git a/backend/crates/gateway-admin/tests/use_case/xai.rs b/backend/crates/gateway-admin/tests/use_case/xai.rs index 1164c5c50..12f5bc01f 100644 --- a/backend/crates/gateway-admin/tests/use_case/xai.rs +++ b/backend/crates/gateway-admin/tests/use_case/xai.rs @@ -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, store: Arc) -> AdminServices { super::AdminHarness::new() .provider(provider) diff --git a/backend/crates/gateway-api/tests/admin/mod.rs b/backend/crates/gateway-api/tests/admin/mod.rs index 18c5c6e5f..039fe9516 100644 --- a/backend/crates/gateway-api/tests/admin/mod.rs +++ b/backend/crates/gateway-api/tests/admin/mod.rs @@ -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 { - Err(unavailable("credential list")) - } - async fn credential_details( &self, _: &ProviderKind, diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs index 8795b1568..115a2c582 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs @@ -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>, + finalizations: Mutex>, + active: AtomicBool, + settlement_calls: AtomicUsize, + settlement_started: tokio::sync::Notify, + settlement_gate: Mutex>>, + charges: Mutex>, + 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, + ) -> 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 { + unreachable!("the WebSocket test does not recover requests") + } +} + +impl ClientAdmissionPort for SettlementPorts { + fn admit( + &self, + _: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result> { + 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> { + 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> { + 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> { + 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, ProviderError> { + Ok(Vec::new()) + } + + async fn execute( + &self, + request: ProviderRequest, + _: AttemptContext, + ) -> Result { + 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, +} + +struct SettlementConnectionGuard(Arc); + +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, 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]) + .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 { Request::get("/v1/responses") .header(AUTHORIZATION, authorization) diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index 4ad8bb939..1d60f6046 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -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) -> BoxFuture<'static, ()>; } @@ -978,8 +981,9 @@ impl AdmissionLease { } struct DefaultExecutionSession { - core: Option>, + core: ResponseExecutionSession, admission: Option, + cleanup: Option>, circuits: Arc, continuation: Arc, observed_provider_outcomes: usize, @@ -996,8 +1000,9 @@ impl DefaultExecutionSession { budget: Option>, ) -> 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, 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, 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, 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 { - 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, ) -> 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) -> BoxFuture<'static, ()> { diff --git a/backend/crates/gateway-core/src/engine/provider.rs b/backend/crates/gateway-core/src/engine/provider.rs index 14c37e691..e593ce02b 100644 --- a/backend/crates/gateway-core/src/engine/provider.rs +++ b/backend/crates/gateway-core/src/engine/provider.rs @@ -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, -} - -/// 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` 的异构注册表。 #[derive(Default)] pub struct ProviderRegistryBuilder { @@ -665,11 +599,6 @@ impl ProviderRegistry { self.providers.get(provider) } - /// 返回全部已注册 Provider。 - pub fn provider_kinds(&self) -> impl Iterator { - 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 { - 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, 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 { + self.providers + .iter() + .map(|(kind, provider)| (kind.clone(), provider.catalog_generation())) + .collect() + } + + fn query_model_capabilities( + &self, + provider_kind: &ProviderKind, + ) -> BoxFuture<'_, Result, ProviderCatalogUnavailable>> { + let provider = self.providers.get(provider_kind); + Box::pin(async move { + provider + .ok_or(ProviderCatalogUnavailable)? + .query_model_capabilities() + .await + .map_err(|_| ProviderCatalogUnavailable) + }) + } +} diff --git a/backend/crates/gateway-core/src/lib.rs b/backend/crates/gateway-core/src/lib.rs index ca9fe8723..6ae636131 100644 --- a/backend/crates/gateway-core/src/lib.rs +++ b/backend/crates/gateway-core/src/lib.rs @@ -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() diff --git a/backend/crates/gateway-core/src/routing/catalog.rs b/backend/crates/gateway-core/src/routing/catalog.rs new file mode 100644 index 000000000..263f1f36e --- /dev/null +++ b/backend/crates/gateway-core/src/routing/catalog.rs @@ -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, +} + +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; + + /// 查询指定 Provider 的模型事实；成功的空列表表示已知无模型，失败表示未知。 + fn query_model_capabilities( + &self, + provider: &ProviderKind, + ) -> BoxFuture<'_, Result, ProviderCatalogUnavailable>>; +} diff --git a/backend/crates/gateway-core/src/routing/mod.rs b/backend/crates/gateway-core/src/routing/mod.rs index ed57e0534..7923f4df2 100644 --- a/backend/crates/gateway-core/src/routing/mod.rs +++ b/backend/crates/gateway-core/src/routing/mod.rs @@ -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}; diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index 250fc5333..85e2abfca 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -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, - providers: ProviderRegistry, + catalogs: Arc, } impl RuntimeSnapshotCompiler { #[must_use] - pub const fn new(store: Arc, providers: ProviderRegistry) -> Self { - Self { store, providers } + pub const fn new( + store: Arc, + catalogs: Arc, + ) -> Self { + Self { store, catalogs } } pub(crate) fn store(&self) -> Arc { @@ -228,13 +231,13 @@ impl RuntimeSnapshotCompiler { pub(crate) fn provider_catalog_generations( &self, ) -> BTreeMap { - self.providers.catalog_generations() + self.catalogs.catalog_generations() } /// 读取一个 revision，并为已注册 Provider 查询实时模型目录。 pub async fn compile(&self) -> Result { 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, ) -> Result { - let provider_kinds = providers.provider_kinds().cloned().collect::>(); let registered_providers = provider_kinds.iter().cloned().collect::>(); // 目录查询失败表示未知；查询成功后，即使为空，也必须与“已知缺少模型”区分。 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()); diff --git a/backend/crates/gateway-core/src/runtime/mod.rs b/backend/crates/gateway-core/src/runtime/mod.rs index f8066e456..5595709f4 100644 --- a/backend/crates/gateway-core/src/runtime/mod.rs +++ b/backend/crates/gateway-core/src/runtime/mod.rs @@ -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, snapshots: RuntimeSnapshotHandle, subscriptions: Arc, + refresh_lock: Arc>, +} + +enum RefreshMode { + Required, + Reconcile, } impl RuntimeSnapshotPublisher { #[must_use] - pub const fn new( + pub fn new( compiler: Arc, snapshots: RuntimeSnapshotHandle, subscriptions: Arc, @@ -148,21 +154,55 @@ impl RuntimeSnapshotPublisher { compiler, snapshots, subscriptions, + refresh_lock: Arc::new(Mutex::new(())), } } - /// 重新编译并原子替换本进程快照。 + /// 串行重编译并替换本进程快照；无法确认配置时暂停新请求。 pub async fn refresh(&self) -> Result { - let snapshot = self.compiler.compile().await?; + self.refresh_with_mode(RefreshMode::Required).await + } + + async fn refresh_with_mode( + &self, + mode: RefreshMode, + ) -> Result { + // 所有入口共享从读取事实到发布或暂停的完整临界区；只锁最后的替换会让 + // 旧编译覆盖新授权，或让晚到的失败暂停新快照。请求读取不等待此锁。 + 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 { 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, 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, } diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 9838c8f8c..4edb984bf 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -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}; diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs index 6362421ed..689a3abdf 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -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, limits: Mutex>, + release_gate: Mutex>>, + releases: AtomicUsize, } impl ClientAdmissionPort for Admissions { @@ -32,7 +34,14 @@ impl ClientAdmissionPort for Admissions { _: &'a ClientApiKeyId, _: &'a ModelRequestId, ) -> BoxFuture<'a, Result> { - 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, charges: Mutex>, + settlement_gate: Mutex>>, + 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, ProviderError> { + Ok(Vec::new()) + } + + async fn execute( + &self, + request: ProviderRequest, + _: AttemptContext, + ) -> Result { + 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> = 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, + budget: Option>, + fail: bool, +) -> DefaultExecutionService { + let service = DefaultExecutionService::new( + RuntimeSnapshotHandle::new(start_snapshot()), + Arc::new(TrackingExecutionStore::default()), + ProviderRegistry::new([Arc::new(ChargedProvider { fail }) as Arc]).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}; diff --git a/backend/crates/gateway-core/tests/engine/provider.rs b/backend/crates/gateway-core/tests/engine/provider.rs index 58802171b..baa8a1f78 100644 --- a/backend/crates/gateway-core/tests/engine/provider.rs +++ b/backend/crates/gateway-core/tests/engine/provider.rs @@ -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 = ®istry; + 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, 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, ProviderError> { + self.0 + .clone() + .map_err(|kind| ProviderError::new(kind, UpstreamSendState::NotSent)) + } + + async fn execute( + &self, + _: ProviderRequest, + _: AttemptContext, + ) -> Result { + 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]) + .expect("registry"); + let catalog: &dyn ProviderCatalogPort = ®istry; + + 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 = ®istry; + 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); impl Drop for DropLease { diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs index 41f6d8060..7eca394e0 100644 --- a/backend/crates/gateway-core/tests/routing/snapshot.rs +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -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, ProviderError> { +impl ProviderCatalogPort for TestCatalog { + fn catalog_generations(&self) -> BTreeMap { 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 { - Err(ProviderError::new( - ProviderErrorKind::Unavailable, - UpstreamSendState::NotSent, - )) + _: &ProviderKind, + ) -> BoxFuture<'_, Result, 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 { + BTreeMap::from([( + ProviderKind::new("alpha").expect("provider"), + ProviderCatalogGeneration::new(generation), + )]) +} + +impl ProviderCatalogPort for PublishingCatalog { + fn catalog_generations(&self) -> BTreeMap { + 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, 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, 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 { - 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]) - .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]) - .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]).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) -> RuntimeSnapshotCompiler { - RuntimeSnapshotCompiler::new(store, ProviderRegistry::default()) + RuntimeSnapshotCompiler::new(store, Arc::new(TestCatalog::NoProviders)) } fn revision(value: u64) -> ConfigRevision { diff --git a/backend/crates/gateway-core/tests/runtime/mod.rs b/backend/crates/gateway-core/tests/runtime/mod.rs index 34b5cc1dd..a24ad5797 100644 --- a/backend/crates/gateway-core/tests/runtime/mod.rs +++ b/backend/crates/gateway-core/tests/runtime/mod.rs @@ -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 = Arc>>>>; #[derive(Clone)] struct TestSnapshotStore { facts: Arc>>, current_revision: Arc>>, + next_load: ReadGate, + next_revision: ReadGate, + loads: Arc, } 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> { - 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> { - 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>, + stream: Mutex>, + next_publish: Mutex>>, } 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> { - 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(¤t)); + 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(¤t)); + 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![ProviderCatalogGeneration::new(1)], + ); + assert_eq!( + frozen + .provider_catalog_generations() + .values() + .copied() + .collect::>(), + 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, 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>>, +} + +impl ProviderCatalogPort for TestCatalog { + fn catalog_generations(&self) -> BTreeMap { + BTreeMap::from([( + ProviderKind::new("audit").expect("provider"), + ProviderCatalogGeneration::new(self.generation.load(Ordering::SeqCst)), + )]) + } + + fn query_model_capabilities( + &self, + _: &ProviderKind, + ) -> BoxFuture<'_, Result, 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, + catalog: Arc, +) -> 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) -> RuntimeSnapshotCompiler { - RuntimeSnapshotCompiler::new(store, ProviderRegistry::default()) + catalog_compiler(store, Arc::new(TestCatalog::default())) } fn empty_snapshot(value: u64) -> RuntimeSnapshot { diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs index 2b98ec716..d087c9dbf 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs @@ -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 { - 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::>(); - 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::>>()?, - next_cursor, - }) - } - async fn credential_details( &self, provider_kind: &ProviderKind, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs index 0a2e6782f..ab0f24497 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs @@ -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, }, diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs index 44945719a..b8f6b872f 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs @@ -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 { diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index 555cba692..5781202e6 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -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; diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 21cb2c457..2f27b344a 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -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, diff --git a/backend/crates/providers/openai/tests/transport/http_client.rs b/backend/crates/providers/openai/tests/transport/http_client.rs index 8550b75f2..d7f0b14b2 100644 --- a/backend/crates/providers/openai/tests/transport/http_client.rs +++ b/backend/crates/providers/openai/tests/transport/http_client.rs @@ -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::)", - ".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(¤t_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()); diff --git a/backend/crates/providers/xai/Cargo.toml b/backend/crates/providers/xai/Cargo.toml index 57e7f2dbf..e3599a864 100644 --- a/backend/crates/providers/xai/Cargo.toml +++ b/backend/crates/providers/xai/Cargo.toml @@ -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]] diff --git a/backend/crates/providers/xai/src/credential/catalog.rs b/backend/crates/providers/xai/src/credential/catalog.rs index dd22792be..02ea546f0 100644 --- a/backend/crates/providers/xai/src/credential/catalog.rs +++ b/backend/crates/providers/xai/src/credential/catalog.rs @@ -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, - ) -> Result { - 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, diff --git a/backend/crates/providers/xai/src/credential/oidc_verifier.rs b/backend/crates/providers/xai/src/credential/oidc_verifier.rs index 2ebe7d9e9..e91d67c0f 100644 --- a/backend/crates/providers/xai/src/credential/oidc_verifier.rs +++ b/backend/crates/providers/xai/src/credential/oidc_verifier.rs @@ -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, + OidcEgress, + >, endpoint_policy: Arc, - jwks_cache: Mutex, cache_ttl: Duration, } +struct OidcEgress { + client: Client, + jwks_cache: Mutex, +} + impl ReqwestOidcTokenVerifier { /// 构建 verifier；使用显式 endpoint policy 与 JWKS cache TTL，只允许访问 /// `auth.x.ai` 官方验证端点。 @@ -53,11 +60,16 @@ impl ReqwestOidcTokenVerifier { endpoint_policy: Arc, cache_ttl: Duration, ) -> Result { - 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 { - 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 { 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 { 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 { - 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 { 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 { 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, 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 }) } } diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index eb46a4372..ee7ec128f 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -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) diff --git a/backend/crates/providers/xai/src/credential/verification.rs b/backend/crates/providers/xai/src/credential/verification.rs index 8ab1497d8..0da5d7d8b 100644 --- a/backend/crates/providers/xai/src/credential/verification.rs +++ b/backend/crates/providers/xai/src/credential/verification.rs @@ -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 } diff --git a/backend/crates/providers/xai/src/provider/mod.rs b/backend/crates/providers/xai/src/provider/mod.rs index 241fd1511..81cde653f 100644 --- a/backend/crates/providers/xai/src/provider/mod.rs +++ b/backend/crates/providers/xai/src/provider/mod.rs @@ -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, diff --git a/backend/crates/providers/xai/src/transport/egress.rs b/backend/crates/providers/xai/src/transport/egress.rs new file mode 100644 index 000000000..1229ffecc --- /dev/null +++ b/backend/crates/providers/xai/src/transport/egress.rs @@ -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 = VecDeque<(K, Arc>)>; + +pub(crate) struct EgressCache { + entries: Mutex>, +} + +impl EgressCache { + 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>> { + 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) + } +} diff --git a/backend/crates/providers/xai/src/transport/mod.rs b/backend/crates/providers/xai/src/transport/mod.rs index dcb22405e..7836a0102 100644 --- a/backend/crates/providers/xai/src/transport/mod.rs +++ b/backend/crates/providers/xai/src/transport/mod.rs @@ -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; diff --git a/backend/crates/providers/xai/src/transport/network.rs b/backend/crates/providers/xai/src/transport/network.rs index b85bf31c3..fcfb5d347 100644 --- a/backend/crates/providers/xai/src/transport/network.rs +++ b/backend/crates/providers/xai/src/transport/network.rs @@ -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, + proxy: Option<&OutboundProxy>, ) -> Result; fn build_inference_client( &self, timeout: Option, + proxy: Option<&OutboundProxy>, ) -> Result; 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, + proxy: Option<&OutboundProxy>, ) -> Result { - build_official_client(GrokDnsResolutionPolicy::official_oauth(), timeout) + build_official_client(GrokDnsResolutionPolicy::official_oauth(), timeout, proxy) } fn build_inference_client( &self, timeout: Option, + proxy: Option<&OutboundProxy>, ) -> Result { - 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, Client>, endpoint_policy: Arc, } @@ -242,9 +251,9 @@ impl ReqwestOAuthTransport { pub fn new( endpoint_policy: Arc, ) -> Result { - 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, + clients: EgressCache, + unbound_client: Mutex>, endpoint_policy: Arc, } 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, ) -> Result { - 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, - least_recently_used: VecDeque, - unbound_client: Option, -} - -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 { - if let Some(client) = self.by_binding.get(binding).cloned() { - self.record_use(binding); - return Some(client); - } - None - } - - fn take_unbound(&mut self) -> Option { - 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, Client>, endpoint_policy: Arc, } @@ -576,9 +533,9 @@ impl ReqwestGrokModelCatalogTransport { pub fn new( endpoint_policy: Arc, ) -> Result { - 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, + proxy: Option<&OutboundProxy>, ) -> Result { 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, -) -> Result { - 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 { diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index a76d8aaf6..e4ae5dc9d 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -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()) diff --git a/backend/crates/providers/xai/tests/credential/catalog.rs b/backend/crates/providers/xai/tests/credential/catalog.rs index 723600ac0..952b15fc4 100644 --- a/backend/crates/providers/xai/tests/credential/catalog.rs +++ b/backend/crates/providers/xai/tests/credential/catalog.rs @@ -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; diff --git a/backend/crates/providers/xai/tests/credential/oidc_verifier.rs b/backend/crates/providers/xai/tests/credential/oidc_verifier.rs index b6d2463a0..342fe0c80 100644 --- a/backend/crates/providers/xai/tests/credential/oidc_verifier.rs +++ b/backend/crates/providers/xai/tests/credential/oidc_verifier.rs @@ -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::>(); + 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; diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index c71856444..8f7acf30f 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -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, + account_id: ProviderAccountId, + initial_proxy: OutboundProxy, + current_proxy: Option, +} + +#[async_trait] +impl GrokCredentialRefresher for SwitchingProxyRefresher { + async fn prepare_cycle(&self) -> Result<(), GrokRefreshFailure> { + Ok(()) + } + + async fn refresh(&self, _: &SecretValue) -> Result { + panic!("refresh service must forward the selected account proxy") + } + + async fn refresh_with_proxy( + &self, + _: &SecretValue, + proxy: Option<&OutboundProxy>, + ) -> Result { + 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, + proxy: Option<&OutboundProxy>, + ) -> Result { + self.build_inference_client(timeout, proxy) + } + + fn build_inference_client( + &self, + timeout: Option, + proxy: Option<&OutboundProxy>, + ) -> Result { + // “直连”也终止于本地 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 { + 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 { diff --git a/backend/crates/providers/xai/tests/credential/selector.rs b/backend/crates/providers/xai/tests/credential/selector.rs index fb90c814d..85ce4b228 100644 --- a/backend/crates/providers/xai/tests/credential/selector.rs +++ b/backend/crates/providers/xai/tests/credential/selector.rs @@ -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，不持久化 diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index b0028266a..a3cfb0948 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -115,7 +115,10 @@ impl LoopbackGrokEndpointPolicy { && url.fragment().is_none() } - fn build_client(timeout: Option) -> Result { + fn build_client( + timeout: Option, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { 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, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result { - Self::build_client(timeout) + Self::build_client(timeout, proxy) } fn build_inference_client( &self, timeout: Option, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result { - 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 { Arc::new(LoopbackGrokEndpointPolicy::for_origin(origin)) } +pub async fn rejecting_account_proxy() -> ( + gateway_core::account::OutboundProxy, + tokio::task::JoinHandle, +) { + 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, + ) { + 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 { 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)] diff --git a/backend/crates/providers/xai/tests/transport/catalog.rs b/backend/crates/providers/xai/tests/transport/catalog.rs index fa66cdfb5..fb629e9d6 100644 --- a/backend/crates/providers/xai/tests/transport/catalog.rs +++ b/backend/crates/providers/xai/tests/transport/catalog.rs @@ -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"))); diff --git a/backend/crates/providers/xai/tests/transport/egress.rs b/backend/crates/providers/xai/tests/transport/egress.rs new file mode 100644 index 000000000..5777fed58 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/egress.rs @@ -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::>(); + 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, + 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(); + } +} diff --git a/backend/crates/providers/xai/tests/transport/mod.rs b/backend/crates/providers/xai/tests/transport/mod.rs index 9dc53d132..bb9548173 100644 --- a/backend/crates/providers/xai/tests/transport/mod.rs +++ b/backend/crates/providers/xai/tests/transport/mod.rs @@ -2,6 +2,7 @@ mod canonical; mod catalog; mod compaction; mod config; +mod egress; mod headers; mod network; mod profile; diff --git a/backend/crates/providers/xai/tests/transport/network.rs b/backend/crates/providers/xai/tests/transport/network.rs index 4d390a446..dcb8872c2 100644 --- a/backend/crates/providers/xai/tests/transport/network.rs +++ b/backend/crates/providers/xai/tests/transport/network.rs @@ -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, -) { - 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::>() + .await + }) + }) + .collect::>(); + 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::>(); + 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, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result { - self.inner.build_oauth_client(timeout) + self.inner.build_oauth_client(timeout, proxy) } fn build_inference_client( &self, timeout: Option, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result { 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 } diff --git a/backend/migrations/README.md b/backend/migrations/README.md index 84c0067c4..434e90f5c 100644 --- a/backend/migrations/README.md +++ b/backend/migrations/README.md @@ -43,3 +43,8 @@ export CPR_TEST_REDIS_URL='redis://:@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`，避免启动配置被部署环境覆盖。 diff --git a/docs/api.md b/docs/api.md index 4f4e247b0..0063ff958 100644 --- a/docs/api.md +++ b/docs/api.md @@ -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 diff --git a/docs/architecture.md b/docs/architecture.md index 2f49f0399..272e8eada 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -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 隔离，不创建套餐专属状态机。 diff --git a/frontend/src/views/accounts/composables/useAccountConnectionTest.ts b/frontend/src/views/accounts/composables/useAccountConnectionTest.ts index 1f48897a6..63cd93559 100644 --- a/frontend/src/views/accounts/composables/useAccountConnectionTest.ts +++ b/frontend/src/views/accounts/composables/useAccountConnectionTest.ts @@ -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(null) const testingConnections = useIdSet() - 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([]) const connectionTestStreamUrl = shallowRef() @@ -396,25 +399,39 @@ export function useAccountConnectionTest(options: { reload: () => Promise Promise Promise { + watch([showConnectionTestModal, () => testingAccount.value?.id], ([open]) => { + modelsRequest.invalidate() if (!open) { abortConnectionTest() } - }) + }, { flush: 'sync' }) onBeforeUnmount(() => { abortConnectionTest() diff --git a/frontend/src/views/accounts/composables/useAccountMutations.ts b/frontend/src/views/accounts/composables/useAccountMutations.ts index ee247a03e..6a87bd46c 100644 --- a/frontend/src/views/accounts/composables/useAccountMutations.ts +++ b/frontend/src/views/accounts/composables/useAccountMutations.ts @@ -23,7 +23,7 @@ export function useAccountMutations(options: { accounts: Ref selectedIds: Ref> reload: () => Promise - replaceAccount: (account: AccountRow) => boolean + replaceAccount: (account: AccountRow) => Promise }) { 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) diff --git a/frontend/src/views/accounts/composables/useAccountResetCredits.ts b/frontend/src/views/accounts/composables/useAccountResetCredits.ts index 7620296c5..8a586bcab 100644 --- a/frontend/src/views/accounts/composables/useAccountResetCredits.ts +++ b/frontend/src/views/accounts/composables/useAccountResetCredits.ts @@ -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() +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() + +function getResetCreditsSession(accountId: string) { + let session = sessionsByAccountId.get(accountId) + if (!session) { + session = shallowReactive({ + 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([]) - 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(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 { - 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, diff --git a/frontend/src/views/accounts/composables/useAccountsQuery.ts b/frontend/src/views/accounts/composables/useAccountsQuery.ts index 44ca161fe..4228bbfe3 100644 --- a/frontend/src/views/accounts/composables/useAccountsQuery.ts +++ b/frontend/src/views/accounts/composables/useAccountsQuery.ts @@ -8,16 +8,6 @@ import { usePagedQuery } from '@/composables/usePagedQuery' import { errorMessage } from '@/utils/async' type AccountRow = Awaited>['items'][number] -type AccountSummary = Awaited>['summary'] -type AccountSummaryStatusKey = Exclude - -const summaryKeyByStatus = { - normal: 'normal', - quota_exhausted: 'quotaExhausted', - rate_limited: 'rateLimited', - disabled: 'disabled', - error: 'error', -} as const satisfies Record 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( diff --git a/frontend/src/views/accounts/index.vue b/frontend/src/views/accounts/index.vue index 8a5cce259..85de9e137 100644 --- a/frontend/src/views/accounts/index.vue +++ b/frontend/src/views/accounts/index.vue @@ -295,7 +295,7 @@ const {