diff --git a/backend/apps/gateway/src/bootstrap/command.rs b/backend/apps/gateway/src/bootstrap/command.rs index 354db1ba0..058cdc832 100644 --- a/backend/apps/gateway/src/bootstrap/command.rs +++ b/backend/apps/gateway/src/bootstrap/command.rs @@ -93,7 +93,7 @@ async fn execute( Some(plugin_runtime.clone()), Some(plugin_runtime.observer_registry()), Some(plugin_runtime.policy_registry()), - Some(plugin_runtime.middleware_registry()), + Some(plugin_runtime.execution_registry()), ); let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control()); or_shutdown!(plugin_runtime, management.bind(&plugin_runtime)); @@ -110,7 +110,7 @@ async fn execute( providers.admin, std::sync::Arc::new(gateway_host::pricing::ModelsDevPricing), ); - let mut services = gateway_admin::service::Registry::new(service_middleware( + let mut services = gateway_admin::public_service::Registry::new(service_middleware( core.snapshots(), &plugin_runtime, )); diff --git a/backend/apps/gateway/src/bootstrap/plugins.rs b/backend/apps/gateway/src/bootstrap/plugins.rs index d6cc8e6f2..00ee71caa 100644 --- a/backend/apps/gateway/src/bootstrap/plugins.rs +++ b/backend/apps/gateway/src/bootstrap/plugins.rs @@ -72,11 +72,11 @@ impl ManagementPorts { pub(super) fn service_middleware( snapshots: RuntimeSnapshotHandle, runtime: &PluginRuntime, -) -> gateway_admin::service::PlanSource { - let middleware = runtime.middleware_registry(); +) -> gateway_admin::public_service::PlanSource { + let middleware = runtime.execution_registry(); Arc::new(move || { let snapshot = snapshots.snapshot_for_diagnostics()?; - middleware.resolve(snapshot.extensions()?) + middleware.middleware(snapshot.extensions()?) }) } diff --git a/backend/apps/gateway/src/bootstrap/server.rs b/backend/apps/gateway/src/bootstrap/server.rs index 569f032ec..28cae2fcc 100644 --- a/backend/apps/gateway/src/bootstrap/server.rs +++ b/backend/apps/gateway/src/bootstrap/server.rs @@ -28,7 +28,7 @@ pub async fn run() -> Result<(), BootstrapError> { Some(plugin_runtime.clone()), Some(plugin_runtime.observer_registry()), Some(plugin_runtime.policy_registry()), - Some(plugin_runtime.middleware_registry()), + Some(plugin_runtime.execution_registry()), Some(plugin_runtime.frontend_authentication_registry()), ); let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control()); @@ -122,14 +122,14 @@ pub async fn run() -> Result<(), BootstrapError> { ) .with_middleware({ let snapshots = core.snapshots(); - let middleware = plugin_runtime.middleware_registry(); + let middleware = plugin_runtime.execution_registry(); move |snapshot| { if let Some(snapshot) = snapshot { - return middleware.resolve(snapshot.extensions()?); + return middleware.middleware(snapshot.extensions()?); } // 数据面未就绪时，管理与诊断路由仍可使用发布候选中的插件 let diagnostic = snapshots.snapshot_for_diagnostics()?; - middleware.resolve(diagnostic.extensions()?) + middleware.middleware(diagnostic.extensions()?) } }); let http_dispatcher = api.dispatcher(); diff --git a/backend/apps/gateway/tests/architecture.rs b/backend/apps/gateway/tests/architecture.rs index 79e31f542..cfafa8811 100644 --- a/backend/apps/gateway/tests/architecture.rs +++ b/backend/apps/gateway/tests/architecture.rs @@ -79,18 +79,25 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() { path if path.starts_with("policy/") => { Some(&["account", "identity", "policy", "validation"]) } - "settings/values.rs" | "settings/compiled.rs" => { - Some(&["account", "concurrency", "identity", "metering", "policy"]) - } + path if path.starts_with("settings/") => Some(&[ + "account", + "concurrency", + "identity", + "metering", + "policy", + "settings", + ]), path if path.starts_with("account/") => Some(&["account", "identity", "validation"]), _ => None, }; let Some(allowed) = allowed else { continue }; let source = fs::read_to_string(root.join(&relative)).expect("read core source"); let syntax = syn::parse_file(&source).expect("parse core source"); - let mut references = CrateReferences::default(); + let mut references = CrateReferences::in_module(&relative); references.visit_file(&syntax); - for dependency in references.0 { + // 相对路径中的 self 指向当前 owner，不是跨 owner 依赖 + references.owners.remove(&references.module_path[0]); + for dependency in references.owners { assert!( allowed.contains(&dependency.as_str()), "{} must not depend on crate::{dependency}", @@ -100,49 +107,155 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() { } } +#[test] +fn core_owner_paths_form_a_dag() { + let root = backend_root().join("crates/gateway-core/src"); + let mut graph = BTreeMap::>::new(); + for relative in super::rust_files(&root) { + if relative == Path::new("lib.rs") { + continue; + } + let mut references = CrateReferences::in_module(&relative); + let owner = references.module_path[0].clone(); + let source = fs::read_to_string(root.join(&relative)).expect("read core source"); + references.visit_file(&syn::parse_file(&source).expect("parse core source")); + references.owners.remove(&owner); + graph.entry(owner).or_default().extend(references.owners); + } + for owner in graph.keys() { + let mut pending = graph[owner].iter().cloned().collect::>(); + let mut visited = BTreeSet::new(); + while let Some(dependency) = pending.pop() { + assert_ne!( + &dependency, owner, + "Core owner dependency cycle reaches {owner}" + ); + if visited.insert(dependency.clone()) + && let Some(next) = graph.get(&dependency) + { + pending.extend(next.iter().cloned()); + } + } + } +} + +#[test] +fn core_reference_scan_resolves_groups_reexports_and_relative_paths() { + let syntax = syn::parse_file( + r#" + pub use crate::{engine::AttemptContext, settings::SettingsValues as Values}; + use super::super::runtime as publication; + use self::compiled::Value; + // crate::ignored::Comment + const EXAMPLE: &str = "crate::ignored::Literal"; + "#, + ) + .expect("parse reference fixture"); + let mut references = CrateReferences::in_module(Path::new("routing/snapshot.rs")); + references.visit_file(&syntax); + assert_eq!( + references.owners, + BTreeSet::from([ + "engine".to_owned(), + "settings".to_owned(), + "runtime".to_owned(), + "routing".to_owned(), + ]) + ); +} + #[derive(Default)] -struct CrateReferences(BTreeSet); +struct CrateReferences { + owners: BTreeSet, + module_path: Vec, +} impl<'ast> Visit<'ast> for CrateReferences { fn visit_path(&mut self, path: &'ast syn::Path) { - let mut segments = path.segments.iter(); - if segments - .next() - .is_some_and(|segment| segment.ident == "crate") - && let Some(owner) = segments.next() - { - self.0.insert(owner.ident.to_string()); - } + self.record( + &path + .segments + .iter() + .map(|part| part.ident.to_string()) + .collect::>(), + ); syn::visit::visit_path(self, path); } fn visit_item_use(&mut self, item: &'ast syn::ItemUse) { - self.use_tree(&item.tree, false); + self.use_tree(&item.tree, &mut Vec::new()); } } impl CrateReferences { - fn use_tree(&mut self, tree: &syn::UseTree, crate_root: bool) { + fn in_module(relative: &Path) -> Self { + let mut module = relative + .with_extension("") + .components() + .map(|part| { + part.as_os_str() + .to_str() + .expect("UTF-8 module path") + .to_owned() + }) + .collect::>(); + if module.last().is_some_and(|part| part == "mod") { + module.pop(); + } + Self { + owners: BTreeSet::new(), + module_path: module, + } + } + + fn record(&mut self, path: &[String]) { + let Some(first) = path.first() else { return }; + let owner = if first == "crate" { + path.get(1).cloned() + } else if first == "self" || first == "super" { + let mut module = self.module_path.clone(); + let mut rest = path; + while rest + .first() + .is_some_and(|part| part == "super" || part == "self") + { + if rest[0] == "super" { + module.pop(); + } + rest = &rest[1..]; + } + module.first().or_else(|| rest.first()).cloned() + } else { + None + }; + if let Some(owner) = owner { + self.owners.insert(owner); + } + } + + fn use_tree(&mut self, tree: &syn::UseTree, prefix: &mut Vec) { match tree { - syn::UseTree::Path(path) if crate_root => { - self.0.insert(path.ident.to_string()); + syn::UseTree::Path(path) => { + prefix.push(path.ident.to_string()); + self.use_tree(&path.tree, prefix); + prefix.pop(); } - syn::UseTree::Name(name) if crate_root => { - self.0.insert(name.ident.to_string()); + syn::UseTree::Name(name) => { + prefix.push(name.ident.to_string()); + self.record(prefix); + prefix.pop(); } - syn::UseTree::Rename(rename) if crate_root => { - self.0.insert(rename.ident.to_string()); + syn::UseTree::Rename(rename) => { + prefix.push(rename.ident.to_string()); + self.record(prefix); + prefix.pop(); } - syn::UseTree::Path(path) => self.use_tree(&path.tree, path.ident == "crate"), syn::UseTree::Group(group) => { for item in &group.items { - self.use_tree(item, crate_root); + self.use_tree(item, prefix); } } - syn::UseTree::Glob(_) if crate_root => { - self.0.insert("*".to_owned()); - } - _ => {} + syn::UseTree::Glob(_) => self.record(prefix), } } } @@ -221,7 +334,7 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[ "crates/providers/openai", &["config", "credential", "transport"], ), - ("crates/providers/xai", &["credential", "transport"]), + ("crates/providers/xai", &["transport"]), ]; /// 不对应单一生产模块、而是校验 crate/workspace 整体契约的根级测试场景 diff --git a/backend/crates/gateway-admin/src/lib.rs b/backend/crates/gateway-admin/src/lib.rs index 4b5df63bb..9a89db5e4 100644 --- a/backend/crates/gateway-admin/src/lib.rs +++ b/backend/crates/gateway-admin/src/lib.rs @@ -20,7 +20,7 @@ pub mod backup; pub mod freeze_recovery; pub mod model; pub mod ports; -pub mod service; +pub mod public_service; mod use_case; pub use use_case::plugins::{PluginDistributionPorts, PluginManagementService, PluginsService}; @@ -205,7 +205,7 @@ pub enum AdminConfigError { #[derive(Clone)] pub struct AdminServices { timezone: gateway_core::time::DeploymentTimeZone, - public_services: Arc, + public_services: Arc, plugins: Arc, plugin_management: Arc, proxies: Arc, @@ -231,7 +231,7 @@ impl AdminServices { self.timezone } - pub fn public_services(&self) -> Arc { + pub fn public_services(&self) -> Arc { self.public_services.clone() } @@ -331,7 +331,7 @@ impl AdminBundle { /// 组合根提供给控制面的运行能力；与配置和存储端口分别传入 pub struct AdminRuntimePorts { pub timezone: gateway_core::time::DeploymentTimeZone, - pub service_middleware: service::PlanSource, + pub service_middleware: public_service::PlanSource, pub plugin_preparation: Arc, pub plugin_management: Arc, pub published_snapshot: gateway_core::runtime::RuntimeSnapshotHandle, @@ -466,7 +466,7 @@ async fn initialize_inner( registry.clone(), pricing_source, ); - let mut public_services = service::Registry::new(service_middleware); + let mut public_services = public_service::Registry::new(service_middleware); public_services.register_settings(&settings)?; let services = AdminServices { timezone, diff --git a/backend/crates/gateway-admin/src/ports/plugin_management.rs b/backend/crates/gateway-admin/src/ports/plugin_management.rs index c6a898875..9ac9f1e75 100644 --- a/backend/crates/gateway-admin/src/ports/plugin_management.rs +++ b/backend/crates/gateway-admin/src/ports/plugin_management.rs @@ -1,7 +1,7 @@ //! 插件管理页面、资源访问与授权回调的运行时端口 use async_trait::async_trait; -use gateway_core::runtime::extensions::ExtensionSetReference; +use gateway_core::routing::extensions::ExtensionSetReference; use crate::model::{ AdminError, diff --git a/backend/crates/gateway-admin/src/ports/plugins.rs b/backend/crates/gateway-admin/src/ports/plugins.rs index b1a71f7ee..5648a4b72 100644 --- a/backend/crates/gateway-admin/src/ports/plugins.rs +++ b/backend/crates/gateway-admin/src/ports/plugins.rs @@ -26,7 +26,7 @@ use crate::model::{ PluginCompatibilityRequirements, PluginSource, }, }; -use gateway_core::runtime::extensions::ExtensionSetReference; +use gateway_core::routing::extensions::ExtensionSetReference; /// Admin 准备完整候选并保活到提交后的发布结束；进程状态不能代替持久启用状态 #[async_trait] diff --git a/backend/crates/gateway-admin/src/service/mod.rs b/backend/crates/gateway-admin/src/public_service/mod.rs similarity index 98% rename from backend/crates/gateway-admin/src/service/mod.rs rename to backend/crates/gateway-admin/src/public_service/mod.rs index 382d0fccb..a82142cee 100644 --- a/backend/crates/gateway-admin/src/service/mod.rs +++ b/backend/crates/gateway-admin/src/public_service/mod.rs @@ -8,11 +8,11 @@ use std::{future::Future, sync::Arc}; use futures::future::BoxFuture; use gateway_core::{ - engine::middleware::FrozenMiddlewarePlan, - middleware::{ - compose, + engine::middleware::{ + FrozenMiddlewarePlan, service::{Context, Error, Value}, }, + middleware::compose, }; use crate::model::{AdminError, AdminErrorKind}; diff --git a/backend/crates/gateway-admin/src/service/registry.rs b/backend/crates/gateway-admin/src/public_service/registry.rs similarity index 97% rename from backend/crates/gateway-admin/src/service/registry.rs rename to backend/crates/gateway-admin/src/public_service/registry.rs index 760d123d1..e6b1d8cf7 100644 --- a/backend/crates/gateway-admin/src/service/registry.rs +++ b/backend/crates/gateway-admin/src/public_service/registry.rs @@ -2,7 +2,7 @@ use crate::model::AdminError; use futures::future::BoxFuture; -use gateway_core::middleware::service::{Error, Value}; +use gateway_core::engine::middleware::service::{Error, Value}; use serde::{Serialize, de::DeserializeOwned}; use std::{collections::BTreeMap, sync::Arc}; diff --git a/backend/crates/gateway-admin/src/service/settings.rs b/backend/crates/gateway-admin/src/public_service/settings.rs similarity index 100% rename from backend/crates/gateway-admin/src/service/settings.rs rename to backend/crates/gateway-admin/src/public_service/settings.rs diff --git a/backend/crates/gateway-admin/src/use_case/credential_mutation.rs b/backend/crates/gateway-admin/src/use_case/credential_mutation.rs new file mode 100644 index 000000000..d5cc1d816 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/credential_mutation.rs @@ -0,0 +1,371 @@ +//! 凭据变更共享的绑定校验、事务结算、失效与提交后观察 + +use crate::{ + model::{ + AdminError, MutationContext, + accounts::DeleteAccounts, + provider_credentials::{ + AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult, + CredentialDetails, CredentialMutationResult, CredentialRotationCommit, + PendingAuthorizationMutation, PreparedAuthorizationCommit, + PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation, + ProviderQuotaRequest, StartAuthorization, + }, + }, + ports::{provider::ProviderAdmin, store::AccountStore}, +}; +use gateway_core::{account::ProviderAccountId, routing::ProviderKind, runtime::SnapshotControl}; + +use super::{map_store_error, publish_committed}; + +pub(super) async fn publish_credentials_and_observe_quota( + provider: &std::sync::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 = provider.clone(); + 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(()) +} + +pub(super) async fn required_credential( + accounts: &dyn AccountStore, + provider_kind: &ProviderKind, + account_id: &ProviderAccountId, + resource: &'static str, +) -> Result { + accounts + .credential_details(provider_kind, account_id) + .await + .map_err(|error| map_store_error(error, resource))? + .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) +} + +pub(super) async fn required_plugin_credential( + accounts: &dyn AccountStore, + account_id: &ProviderAccountId, + resource: &'static str, +) -> Result { + accounts + .credential_details_by_id(account_id) + .await + .map_err(|error| map_store_error(error, resource))? + .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) +} + +pub(super) async fn import_proxy_binding( + proxies: &dyn crate::ports::proxy::ProxyStore, + id: Option<&str>, +) -> Result, AdminError> { + let Some(id) = id else { return Ok(None) }; + let reservation = proxies + .reserve_import(id) + .await + .map_err(|error| map_store_error(error, "proxy"))?; + Ok(Some(reservation)) +} + +pub(super) async fn pending_authorization( + accounts: &dyn AccountStore, + proxies: &dyn crate::ports::proxy::ProxyStore, + provider_kind: &ProviderKind, + command: &StartAuthorization, + resource: &'static str, +) -> Result { + let proxy; + let mut proxy_id = None; + let target = match &command.reauthorization { + Some(account_id) => { + proxy = required_credential(accounts, provider_kind, account_id, resource) + .await? + .credential + .outbound_proxy; + AuthorizationMutationTarget::Reauthorize { + account_id: account_id.clone(), + } + } + None => { + use crate::model::proxies::AccountProxySelection; + proxy = match &command.outbound_proxy { + Some(AccountProxySelection::Url(value)) => Some(value.clone()), + Some(AccountProxySelection::Saved(id)) => { + let record = proxies + .get(id) + .await + .map_err(|error| map_store_error(error, "proxy"))?; + proxy_id = Some(record.id); + Some(record.proxy) + } + Some(AccountProxySelection::Direct) | None => None, + }; + AuthorizationMutationTarget::Create { + name: command.name.clone(), + } + } + }; + Ok(PendingAuthorizationMutation::new( + provider_kind.clone(), + target, + crate::model::provider_credentials::AuthorizationOwnerBinding::from_context( + &command.context, + ), + ) + .with_outbound_proxy(proxy) + .with_outbound_proxy_id(proxy_id)) +} + +pub(super) fn validate_prepared_import( + provider_kind: &ProviderKind, + prepared: &PreparedCredentialImport, + _resource: &'static str, +) -> Result<(), AdminError> { + if prepared.provider_kind != *provider_kind + || prepared + .credentials + .iter() + .any(|credential| credential.provider_kind != *provider_kind) + { + return Err(AdminError::conflict("Provider 准备结果与请求范围不一致")); + } + Ok(()) +} + +pub(super) fn validate_prepared_rotation( + account: &crate::model::accounts::AccountRecord, + prepared: &PreparedCredentialRotation, + _resource: &'static str, +) -> Result<(), AdminError> { + validate_prepared_rotation_facts(account, prepared.facts()) +} + +pub(super) fn validate_prepared_rotation_facts( + account: &crate::model::accounts::AccountRecord, + facts: &crate::model::provider_credentials::PreparedCredentialRotationFacts, +) -> Result<(), AdminError> { + if facts.account_id.as_str() != account.id.as_str() + || facts.provider_kind != account.provider_kind + { + return Err(AdminError::conflict("Provider 准备结果与当前凭据不一致")); + } + Ok(()) +} + +pub(super) async fn validate_authorization_commit( + provider_kind: &ProviderKind, + context: &MutationContext, + prepared: PreparedAuthorizationCommit, + resource: &'static str, +) -> Result { + if prepared.pending.provider_kind() != provider_kind + || !prepared.pending.owner_binding().matches_context(context) + { + let validation_error = AdminError::conflict("OAuth 授权绑定已失效，请重新发起授权"); + if let Err(error) = prepared.abort().await { + tracing::warn!( + resource, + settlement_error = %error, + "OAuth authorization claim release failed after preparation validation" + ); + return Err(error); + } + return Err(validation_error); + } + let matches_target = match (prepared.pending.target(), &prepared.credential) { + ( + AuthorizationMutationTarget::Create { .. }, + PreparedAuthorizationCredential::Create(credential), + ) => credential.provider_kind == *provider_kind, + ( + AuthorizationMutationTarget::Reauthorize { account_id }, + PreparedAuthorizationCredential::Reauthorize(credential), + ) => { + let facts = credential.facts(); + facts.provider_kind == *provider_kind && facts.account_id == *account_id + } + _ => false, + }; + if !matches_target { + let validation_error = AdminError::conflict("OAuth 凭据与待提交目标不一致"); + if let Err(error) = prepared.abort().await { + tracing::warn!( + resource, + settlement_error = %error, + "OAuth authorization claim release failed after preparation validation" + ); + return Err(error); + } + return Err(validation_error); + } + Ok(prepared) +} + +pub(super) async fn commit_authorization( + accounts: &dyn AccountStore, + prepared: PreparedAuthorizationCommit, + key: crate::model::provider_credentials::AuthorizationReceiptKey, + settings: Option, + context: &MutationContext, + resource: &'static str, +) -> Result { + if settings.is_some() + && matches!( + &prepared.credential, + PreparedAuthorizationCredential::Reauthorize(_) + ) + { + prepared.abort().await?; + return Err(AdminError::invalid("重新授权不能修改账号设置")); + } + let crate::model::provider_credentials::AuthorizationCommitSettlement { + command, + credential_guard, + authorization_guard, + } = prepared.into_commit(settings, key); + match accounts.commit_authorization(command, context).await { + Ok(result) => { + if let Some(guard) = credential_guard + && result.newly_committed + { + guard.finish(); + } + if let Some(guard) = authorization_guard + && let Err(error) = guard.commit().await + { + // 账号与回执已经原子提交，Redis 清理失败不能把确定的成功改写为失败 + tracing::warn!(resource, settlement_error = %error, "authorization committed but pending cleanup failed"); + } + Ok(result) + } + Err(error) => { + drop(credential_guard); + if let Some(guard) = authorization_guard + && let Err(settlement_error) = guard.abort().await + { + tracing::warn!( + resource, + store_error = %error, + settlement_error = %settlement_error, + "OAuth authorization claim release failed after Store commit failure" + ); + return Err(settlement_error); + } + Err(map_store_error(error, resource)) + } + } +} + +pub(super) async fn commit_credential_rotation( + accounts: &dyn AccountStore, + prepared: PreparedCredentialRotation, + settings: Option, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (facts, guard) = prepared.into_parts(); + match accounts + .commit_credential_rotation( + CredentialRotationCommit { + prepared: facts, + settings, + }, + context, + ) + .await + { + Ok(result) => { + guard.finish(); + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +pub(super) async fn commit_credential_refresh( + accounts: &dyn AccountStore, + prepared: PreparedCredentialRotation, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (facts, guard) = prepared.into_parts(); + match accounts + .commit_credential_refresh( + CredentialRotationCommit { + prepared: facts, + settings: None, + }, + context, + ) + .await + { + Ok(result) => { + guard.finish(); + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +pub(super) async fn delete_credentials( + accounts: &dyn AccountStore, + provider: &dyn ProviderAdmin, + command: CredentialDeletion, + resource: &'static str, +) -> Result { + for account_id in &command.account_ids { + required_credential(accounts, provider.provider_kind(), account_id, resource).await?; + } + let account_ids = command.account_ids; + let revision = accounts + .delete_accounts( + DeleteAccounts { + account_ids: account_ids + .iter() + .map(|account_id| account_id.as_str().to_owned()) + .collect(), + }, + &command.context, + ) + .await + .map_err(|error| map_store_error(error, resource))?; + for account_id in &account_ids { + provider.account_unavailable(account_id).await; + } + provider.account_facts_changed(&account_ids).await; + Ok(CredentialDeletionResult { + config_revision: revision, + account_ids, + }) +} diff --git a/backend/crates/gateway-admin/src/use_case/error.rs b/backend/crates/gateway-admin/src/use_case/error.rs new file mode 100644 index 000000000..08f9065ef --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/error.rs @@ -0,0 +1,67 @@ +//! 管理端口错误到用例错误的统一分类与来源保留 + +use crate::model::{AdminError, AdminErrorKind}; +use crate::ports::store::{AdminStoreError, AdminStoreErrorKind}; + +pub(super) fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError { + let kind = match error.kind() { + AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid, + AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound, + AdminStoreErrorKind::StaleRevision + | AdminStoreErrorKind::DuplicateName + | AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict, + AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable, + }; + let message = match kind { + AdminErrorKind::Invalid => "请求参数不合法", + AdminErrorKind::NotFound => "请求的资源不存在", + AdminErrorKind::Conflict => "当前资源状态冲突，请刷新后重试", + AdminErrorKind::Unavailable => "依赖服务暂不可用", + _ => "服务内部错误", + }; + AdminError::new(kind, message).with_source(OperationFailure { + operation: resource, + source: error.into(), + }) +} + +pub(super) fn map_provider_error( + error: crate::ports::provider::ProviderAdminError, + resource: &'static str, +) -> AdminError { + use crate::ports::provider::ProviderAdminErrorKind; + + let kind = match error.kind() { + ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid, + ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid, + ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound, + ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict, + ProviderAdminErrorKind::Ambiguous => AdminErrorKind::UpstreamResultUnknown, + ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable, + ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable, + ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway, + ProviderAdminErrorKind::Internal => AdminErrorKind::Internal, + }; + let message = match kind { + AdminErrorKind::Invalid => "Provider 请求不合法", + AdminErrorKind::NotFound => "Provider 资源不存在", + AdminErrorKind::Conflict => "Provider 资源状态冲突，请刷新后重试", + AdminErrorKind::BadGateway => "上游服务请求失败", + AdminErrorKind::UpstreamResultUnknown => "上游执行结果未知，请刷新状态后再决定是否重试", + AdminErrorKind::Unavailable => "Provider 服务暂不可用", + AdminErrorKind::Internal => "服务内部错误", + _ => "Provider 操作失败", + }; + AdminError::new(kind, error.public_message().unwrap_or(message)).with_source(OperationFailure { + operation: resource, + source: error.into(), + }) +} + +// 用例的操作名补充端口资源上下文，底层分类与原始原因仍由来源持有 +#[derive(Debug, thiserror::Error)] +#[error("admin operation failed: {operation}")] +struct OperationFailure { + operation: &'static str, + source: gateway_core::error::ErrorSource, +} diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index 8ed47dccd..877800c7d 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -19,450 +19,16 @@ pub mod proxies; pub mod settings; pub mod system; -use crate::{ - model::{ - AdminError, AdminErrorKind, MutationContext, - accounts::DeleteAccounts, - provider_credentials::{ - AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult, - CredentialDetails, CredentialMutationResult, CredentialRotationCommit, - PendingAuthorizationMutation, PreparedAuthorizationCommit, - PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation, - ProviderQuotaRequest, StartAuthorization, - }, - }, - ports::{ - provider::ProviderAdmin, - store::{AccountStore, AdminStoreError, AdminStoreErrorKind}, - }, +mod credential_mutation; +mod error; +mod publication; + +use credential_mutation::{ + commit_authorization, commit_credential_refresh, commit_credential_rotation, + delete_credentials, import_proxy_binding, pending_authorization, + publish_credentials_and_observe_quota, required_credential, required_plugin_credential, + validate_authorization_commit, validate_prepared_import, validate_prepared_rotation, + validate_prepared_rotation_facts, }; -use gateway_core::{ - account::ProviderAccountId, - routing::{ConfigRevision, ProviderKind}, - runtime::SnapshotControl, -}; - -fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError { - let kind = match error.kind() { - AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid, - AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound, - AdminStoreErrorKind::StaleRevision - | AdminStoreErrorKind::DuplicateName - | AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict, - AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable, - }; - let message = match kind { - AdminErrorKind::Invalid => "请求参数不合法", - AdminErrorKind::NotFound => "请求的资源不存在", - AdminErrorKind::Conflict => "当前资源状态冲突，请刷新后重试", - AdminErrorKind::Unavailable => "依赖服务暂不可用", - _ => "服务内部错误", - }; - AdminError::new(kind, message).with_source(OperationFailure { - operation: resource, - source: error.into(), - }) -} - -fn map_provider_error( - error: crate::ports::provider::ProviderAdminError, - resource: &'static str, -) -> AdminError { - use crate::ports::provider::ProviderAdminErrorKind; - - let kind = match error.kind() { - ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid, - ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid, - ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound, - ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict, - ProviderAdminErrorKind::Ambiguous => AdminErrorKind::UpstreamResultUnknown, - ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable, - ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable, - ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway, - ProviderAdminErrorKind::Internal => AdminErrorKind::Internal, - }; - let message = match kind { - AdminErrorKind::Invalid => "Provider 请求不合法", - AdminErrorKind::NotFound => "Provider 资源不存在", - AdminErrorKind::Conflict => "Provider 资源状态冲突，请刷新后重试", - AdminErrorKind::BadGateway => "上游服务请求失败", - AdminErrorKind::UpstreamResultUnknown => "上游执行结果未知，请刷新状态后再决定是否重试", - AdminErrorKind::Unavailable => "Provider 服务暂不可用", - AdminErrorKind::Internal => "服务内部错误", - _ => "Provider 操作失败", - }; - AdminError::new(kind, error.public_message().unwrap_or(message)).with_source(OperationFailure { - operation: resource, - source: error.into(), - }) -} - -async fn publish_committed( - snapshot: &dyn SnapshotControl, - revision: crate::model::Revision, -) -> Result<(), AdminError> { - let revision = ConfigRevision::new(revision.get()) - .map_err(|_| AdminError::internal("已提交的配置版本不合法"))?; - snapshot.publish_committed(revision).await; - Ok(()) -} - -async fn publish_credentials_and_observe_quota( - provider: &std::sync::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 = provider.clone(); - 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, - account_id: &ProviderAccountId, - resource: &'static str, -) -> Result { - accounts - .credential_details(provider_kind, account_id) - .await - .map_err(|error| map_store_error(error, resource))? - .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) -} - -async fn required_plugin_credential( - accounts: &dyn AccountStore, - account_id: &ProviderAccountId, - resource: &'static str, -) -> Result { - accounts - .credential_details_by_id(account_id) - .await - .map_err(|error| map_store_error(error, resource))? - .ok_or_else(|| AdminError::not_found("Provider 凭据不存在")) -} - -async fn import_proxy_binding( - proxies: &dyn crate::ports::proxy::ProxyStore, - id: Option<&str>, -) -> Result, AdminError> { - let Some(id) = id else { return Ok(None) }; - let reservation = proxies - .reserve_import(id) - .await - .map_err(|error| map_store_error(error, "proxy"))?; - Ok(Some(reservation)) -} - -async fn pending_authorization( - accounts: &dyn AccountStore, - proxies: &dyn crate::ports::proxy::ProxyStore, - provider_kind: &ProviderKind, - command: &StartAuthorization, - resource: &'static str, -) -> Result { - let proxy; - let mut proxy_id = None; - let target = match &command.reauthorization { - Some(account_id) => { - proxy = required_credential(accounts, provider_kind, account_id, resource) - .await? - .credential - .outbound_proxy; - AuthorizationMutationTarget::Reauthorize { - account_id: account_id.clone(), - } - } - None => { - use crate::model::proxies::AccountProxySelection; - proxy = match &command.outbound_proxy { - Some(AccountProxySelection::Url(value)) => Some(value.clone()), - Some(AccountProxySelection::Saved(id)) => { - let record = proxies - .get(id) - .await - .map_err(|error| map_store_error(error, "proxy"))?; - proxy_id = Some(record.id); - Some(record.proxy) - } - Some(AccountProxySelection::Direct) | None => None, - }; - AuthorizationMutationTarget::Create { - name: command.name.clone(), - } - } - }; - Ok(PendingAuthorizationMutation::new( - provider_kind.clone(), - target, - crate::model::provider_credentials::AuthorizationOwnerBinding::from_context( - &command.context, - ), - ) - .with_outbound_proxy(proxy) - .with_outbound_proxy_id(proxy_id)) -} - -fn validate_prepared_import( - provider_kind: &ProviderKind, - prepared: &PreparedCredentialImport, - _resource: &'static str, -) -> Result<(), AdminError> { - if prepared.provider_kind != *provider_kind - || prepared - .credentials - .iter() - .any(|credential| credential.provider_kind != *provider_kind) - { - return Err(AdminError::conflict("Provider 准备结果与请求范围不一致")); - } - Ok(()) -} - -fn validate_prepared_rotation( - account: &crate::model::accounts::AccountRecord, - prepared: &PreparedCredentialRotation, - _resource: &'static str, -) -> Result<(), AdminError> { - validate_prepared_rotation_facts(account, prepared.facts()) -} - -fn validate_prepared_rotation_facts( - account: &crate::model::accounts::AccountRecord, - facts: &crate::model::provider_credentials::PreparedCredentialRotationFacts, -) -> Result<(), AdminError> { - if facts.account_id.as_str() != account.id.as_str() - || facts.provider_kind != account.provider_kind - { - return Err(AdminError::conflict("Provider 准备结果与当前凭据不一致")); - } - Ok(()) -} - -async fn validate_authorization_commit( - provider_kind: &ProviderKind, - context: &MutationContext, - prepared: PreparedAuthorizationCommit, - resource: &'static str, -) -> Result { - if prepared.pending.provider_kind() != provider_kind - || !prepared.pending.owner_binding().matches_context(context) - { - let validation_error = AdminError::conflict("OAuth 授权绑定已失效，请重新发起授权"); - if let Err(error) = prepared.abort().await { - tracing::warn!( - resource, - settlement_error = %error, - "OAuth authorization claim release failed after preparation validation" - ); - return Err(error); - } - return Err(validation_error); - } - let matches_target = match (prepared.pending.target(), &prepared.credential) { - ( - AuthorizationMutationTarget::Create { .. }, - PreparedAuthorizationCredential::Create(credential), - ) => credential.provider_kind == *provider_kind, - ( - AuthorizationMutationTarget::Reauthorize { account_id }, - PreparedAuthorizationCredential::Reauthorize(credential), - ) => { - let facts = credential.facts(); - facts.provider_kind == *provider_kind && facts.account_id == *account_id - } - _ => false, - }; - if !matches_target { - let validation_error = AdminError::conflict("OAuth 凭据与待提交目标不一致"); - if let Err(error) = prepared.abort().await { - tracing::warn!( - resource, - settlement_error = %error, - "OAuth authorization claim release failed after preparation validation" - ); - return Err(error); - } - return Err(validation_error); - } - Ok(prepared) -} - -async fn commit_authorization( - accounts: &dyn AccountStore, - prepared: PreparedAuthorizationCommit, - key: crate::model::provider_credentials::AuthorizationReceiptKey, - settings: Option, - context: &MutationContext, - resource: &'static str, -) -> Result { - if settings.is_some() - && matches!( - &prepared.credential, - PreparedAuthorizationCredential::Reauthorize(_) - ) - { - prepared.abort().await?; - return Err(AdminError::invalid("重新授权不能修改账号设置")); - } - let crate::model::provider_credentials::AuthorizationCommitSettlement { - command, - credential_guard, - authorization_guard, - } = prepared.into_commit(settings, key); - match accounts.commit_authorization(command, context).await { - Ok(result) => { - if let Some(guard) = credential_guard - && result.newly_committed - { - guard.finish(); - } - if let Some(guard) = authorization_guard - && let Err(error) = guard.commit().await - { - // 账号与回执已经原子提交，Redis 清理失败不能把确定的成功改写为失败 - tracing::warn!(resource, settlement_error = %error, "authorization committed but pending cleanup failed"); - } - Ok(result) - } - Err(error) => { - drop(credential_guard); - if let Some(guard) = authorization_guard - && let Err(settlement_error) = guard.abort().await - { - tracing::warn!( - resource, - store_error = %error, - settlement_error = %settlement_error, - "OAuth authorization claim release failed after Store commit failure" - ); - return Err(settlement_error); - } - Err(map_store_error(error, resource)) - } - } -} - -async fn commit_credential_rotation( - accounts: &dyn AccountStore, - prepared: PreparedCredentialRotation, - settings: Option, - context: &MutationContext, - resource: &'static str, -) -> Result { - let (facts, guard) = prepared.into_parts(); - match accounts - .commit_credential_rotation( - CredentialRotationCommit { - prepared: facts, - settings, - }, - context, - ) - .await - { - Ok(result) => { - guard.finish(); - Ok(result) - } - Err(error) => { - drop(guard); - Err(map_store_error(error, resource)) - } - } -} - -async fn commit_credential_refresh( - accounts: &dyn AccountStore, - prepared: PreparedCredentialRotation, - context: &MutationContext, - resource: &'static str, -) -> Result { - let (facts, guard) = prepared.into_parts(); - match accounts - .commit_credential_refresh( - CredentialRotationCommit { - prepared: facts, - settings: None, - }, - context, - ) - .await - { - Ok(result) => { - guard.finish(); - Ok(result) - } - Err(error) => { - drop(guard); - Err(map_store_error(error, resource)) - } - } -} - -async fn delete_credentials( - accounts: &dyn AccountStore, - provider: &dyn ProviderAdmin, - command: CredentialDeletion, - resource: &'static str, -) -> Result { - for account_id in &command.account_ids { - required_credential(accounts, provider.provider_kind(), account_id, resource).await?; - } - let account_ids = command.account_ids; - let revision = accounts - .delete_accounts( - DeleteAccounts { - account_ids: account_ids - .iter() - .map(|account_id| account_id.as_str().to_owned()) - .collect(), - }, - &command.context, - ) - .await - .map_err(|error| map_store_error(error, resource))?; - for account_id in &account_ids { - provider.account_unavailable(account_id).await; - } - provider.account_facts_changed(&account_ids).await; - Ok(CredentialDeletionResult { - config_revision: revision, - account_ids, - }) -} - -// 用例的操作名补充端口资源上下文，底层分类与原始原因仍由来源持有 -#[derive(Debug, thiserror::Error)] -#[error("admin operation failed: {operation}")] -struct OperationFailure { - operation: &'static str, - source: gateway_core::error::ErrorSource, -} +use error::{map_provider_error, map_store_error}; +use publication::publish_committed; diff --git a/backend/crates/gateway-admin/src/use_case/plugins/management.rs b/backend/crates/gateway-admin/src/use_case/plugins/management.rs index cb7b8f417..c8f0d7deb 100644 --- a/backend/crates/gateway-admin/src/use_case/plugins/management.rs +++ b/backend/crates/gateway-admin/src/use_case/plugins/management.rs @@ -2,7 +2,8 @@ use std::sync::Arc; -use gateway_core::runtime::{RuntimeSnapshotHandle, extensions::ExtensionSetReference}; +use gateway_core::routing::extensions::ExtensionSetReference; +use gateway_core::runtime::RuntimeSnapshotHandle; use crate::{ model::{ diff --git a/backend/crates/gateway-admin/src/use_case/publication.rs b/backend/crates/gateway-admin/src/use_case/publication.rs new file mode 100644 index 000000000..5bc2807c8 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/publication.rs @@ -0,0 +1,14 @@ +//! 管理事务提交后的运行快照发布入口 + +use crate::model::AdminError; +use gateway_core::{routing::ConfigRevision, runtime::SnapshotControl}; + +pub(super) async fn publish_committed( + snapshot: &dyn SnapshotControl, + revision: crate::model::Revision, +) -> Result<(), AdminError> { + let revision = ConfigRevision::new(revision.get()) + .map_err(|_| AdminError::internal("已提交的配置版本不合法"))?; + snapshot.publish_committed(revision).await; + Ok(()) +} diff --git a/backend/crates/gateway-admin/tests/use_case/mod.rs b/backend/crates/gateway-admin/tests/use_case/mod.rs index 7f8e617cb..12a9a7c25 100644 --- a/backend/crates/gateway-admin/tests/use_case/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/mod.rs @@ -107,7 +107,7 @@ pub(super) struct AdminHarness { plugin_store: Arc, plugin_inspector: Arc, client_key_verifier: Arc, - service_middleware: gateway_admin::service::PlanSource, + service_middleware: gateway_admin::public_service::PlanSource, } impl AdminHarness { @@ -194,7 +194,10 @@ impl AdminHarness { self } - pub(super) fn service_middleware(mut self, source: gateway_admin::service::PlanSource) -> Self { + pub(super) fn service_middleware( + mut self, + source: gateway_admin::public_service::PlanSource, + ) -> Self { self.service_middleware = source; self } diff --git a/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs b/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs index 3b1d6c598..859cb75f1 100644 --- a/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs +++ b/backend/crates/gateway-admin/tests/use_case/plugins/artifacts.rs @@ -24,12 +24,10 @@ use gateway_admin::{ store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult}, }, }; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::{ routing::ConfigRevision, - runtime::{ - RuntimeSnapshotHandle, SnapshotControl, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, - }, + runtime::{RuntimeSnapshotHandle, SnapshotControl}, }; use serde_json::json; @@ -117,7 +115,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for Fixture { &self, _: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot, _: Option, - _: Option<&gateway_core::runtime::extensions::ExtensionSetReference>, + _: Option<&gateway_core::routing::extensions::ExtensionSetReference>, ) -> Option< std::collections::BTreeMap< String, @@ -132,7 +130,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for Fixture { impl gateway_admin::ports::plugins::PluginStateLifecycle for Fixture { async fn activate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::instances::PluginInstance, ) -> Result<(), AdminError> { Ok(()) @@ -142,7 +140,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for Fixture { } async fn migrate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: gateway_admin::model::plugins::state::PluginStateTransition, ) -> Result<(), AdminError> { panic!("unexpected state migration") diff --git a/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs b/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs index b1305cf23..796f936fd 100644 --- a/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs +++ b/backend/crates/gateway-admin/tests/use_case/plugins/instances.rs @@ -40,12 +40,10 @@ use gateway_admin::{ store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult}, }, }; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::{ routing::ConfigRevision, - runtime::{ - RuntimeSnapshotHandle, SnapshotControl, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, - }, + runtime::{RuntimeSnapshotHandle, SnapshotControl}, }; use serde_json::json; diff --git a/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs b/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs index 072cee754..560654cb4 100644 --- a/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/plugins/mod.rs @@ -111,7 +111,7 @@ impl gateway_admin::ports::proxy::ProxyStore for TestPluginPorts { impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts { async fn validate_target( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, ) -> Result<(), AdminError> { Err(AdminError::not_found("unused plugin fixture")) @@ -119,7 +119,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor async fn start_callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: gateway_admin::model::plugins::management::StartPluginManagementCallback, _: &gateway_admin::model::auth::AdminRequestContext, @@ -129,7 +129,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: &str, _: gateway_admin::model::plugins::management::PluginManagementRequest, @@ -139,14 +139,14 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn views( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, ) -> Result, AdminError> { Ok(Vec::new()) } async fn resource( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: &str, _: bool, @@ -156,7 +156,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn handle( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: gateway_admin::model::plugins::management::PluginManagementRequest, ) -> Result @@ -171,7 +171,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts &self, _: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot, _: Option, - _: Option<&gateway_core::runtime::extensions::ExtensionSetReference>, + _: Option<&gateway_core::routing::extensions::ExtensionSetReference>, ) -> Option< std::collections::BTreeMap< String, @@ -186,7 +186,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { async fn activate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::instances::PluginInstance, ) -> Result<(), AdminError> { panic!("unexpected state activation") @@ -196,7 +196,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { } async fn migrate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: gateway_admin::model::plugins::state::PluginStateTransition, ) -> Result<(), AdminError> { panic!("unexpected state migration") @@ -222,7 +222,7 @@ impl gateway_admin::ports::plugins::PluginPreparation for TestPluginPorts { async fn prepare( &self, _: gateway_admin::model::plugins::instances::PluginInstanceSnapshot, - ) -> Result { + ) -> Result { Err(AdminError::invalid("unused plugin fixture")) } } diff --git a/backend/crates/gateway-admin/tests/use_case/service_contract.rs b/backend/crates/gateway-admin/tests/use_case/service_contract.rs index 3d34cf4a5..35cb471d7 100644 --- a/backend/crates/gateway-admin/tests/use_case/service_contract.rs +++ b/backend/crates/gateway-admin/tests/use_case/service_contract.rs @@ -247,7 +247,7 @@ fn sdk_settings_contract_matches_host_declarations() { _ => None, }) .collect(); - let register = source(root, "src/service/settings.rs") + let register = source(root, "src/public_service/settings.rs") .items .into_iter() .find_map(|item| match item { @@ -293,7 +293,7 @@ fn sdk_settings_contract_matches_host_declarations() { let generated = quote! { //! 宿主设置服务的操作标识与请求、响应数据合同 //! - //! 从宿主设置类型与 service/settings.rs 生成；更新命令见 SDK 维护说明 + //! 从宿主设置类型与 public_service/settings.rs 生成；更新命令见 SDK 维护说明 use super::Operation; use serde::{Deserialize, Serialize}; diff --git a/backend/crates/gateway-admin/tests/use_case/settings.rs b/backend/crates/gateway-admin/tests/use_case/settings.rs index e7dbbdc6c..517ca0448 100644 --- a/backend/crates/gateway-admin/tests/use_case/settings.rs +++ b/backend/crates/gateway-admin/tests/use_case/settings.rs @@ -28,7 +28,7 @@ struct PricingSettingsStore { #[tokio::test] async fn explicit_service_calls_freeze_empty_plan_and_cancel_before_entering_terminal() { - use gateway_admin::service::{Origin, Plan, scope}; + use gateway_admin::public_service::{Origin, Plan, scope}; use std::sync::atomic::{AtomicUsize, Ordering}; let resolutions = Arc::new(AtomicUsize::new(0)); diff --git a/backend/crates/gateway-api/src/lib.rs b/backend/crates/gateway-api/src/lib.rs index 0f0c4559f..2ccb140ac 100644 --- a/backend/crates/gateway-api/src/lib.rs +++ b/backend/crates/gateway-api/src/lib.rs @@ -138,7 +138,7 @@ impl ApiBundle { } /// 内部调用使用同一总路由；组合根持有强引用，Runtime 只保存 Weak - pub fn dispatcher(&self) -> Arc { + pub fn dispatcher(&self) -> Arc { Arc::new(middleware::RouterDispatcher { router: middleware::wrap( self.router.clone(), diff --git a/backend/crates/gateway-api/src/middleware/mod.rs b/backend/crates/gateway-api/src/middleware/mod.rs index 3e05fe103..8319a4e2c 100644 --- a/backend/crates/gateway-api/src/middleware/mod.rs +++ b/backend/crates/gateway-api/src/middleware/mod.rs @@ -17,23 +17,21 @@ use axum::{ http::{HeaderName, HeaderValue, Request, StatusCode}, response::{IntoResponse as _, Response}, }; -use gateway_core::{ - engine::middleware::FrozenMiddlewarePlan, - lifecycle::CancellationToken, - middleware::{ - compose, - http::{self as contract, Settings}, - }, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::http as contract; +use gateway_core::engine::middleware::http::Settings; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::middleware::compose; use http_body::{Body as HttpBody, Frame, SizeHint}; use http_body_util::BodyExt as _; use tokio::time::Instant; use tower::ServiceExt as _; tokio::task_local! { static CURRENT: contract::Context; } -tokio::task_local! { static SETTINGS: Option; } +tokio::task_local! { static SETTINGS: Option; } -pub(crate) fn current_settings() -> Option { +pub(crate) fn current_settings() -> Option +{ SETTINGS.try_with(Clone::clone).ok().flatten() } @@ -52,7 +50,7 @@ pub(crate) async fn scope( } pub(crate) type SettingsSource = - Arc Option + Send + Sync>; + Arc Option + Send + Sync>; pub(crate) struct RouterDispatcher { pub(crate) router: Router, @@ -60,7 +58,7 @@ pub(crate) struct RouterDispatcher { } impl contract::Dispatcher for RouterDispatcher { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { (self.settings)() } @@ -195,12 +193,12 @@ async fn handle( }); request.extensions_mut().insert(guard.clone()); request.extensions_mut().insert(context.clone()); - let service_context = gateway_admin::service::Origin { + let service_context = gateway_admin::public_service::Origin { request_id: context.request_id.clone(), call_id: context.call_id.clone(), cancellation: context.cancellation.clone(), extensions: context.extensions.clone(), - plan: gateway_admin::service::Plan::Frozen(plan.clone()), + plan: gateway_admin::public_service::Plan::Frozen(plan.clone()), }; request .extensions_mut() @@ -220,7 +218,7 @@ async fn handle( terminal_context, SETTINGS.scope( settings, - gateway_admin::service::scope( + gateway_admin::public_service::scope( service_context, dispatch(router, request.map(Body::new), started), ), diff --git a/backend/crates/gateway-api/src/middleware/websocket.rs b/backend/crates/gateway-api/src/middleware/websocket.rs index 487c85ea8..891c187a0 100644 --- a/backend/crates/gateway-api/src/middleware/websocket.rs +++ b/backend/crates/gateway-api/src/middleware/websocket.rs @@ -14,15 +14,14 @@ use futures::{ future::BoxFuture, stream::{SplitSink, SplitStream}, }; -use gateway_core::{ - engine::middleware::{FrozenMiddlewarePlan, MiddlewareError, MiddlewareHeader}, - lifecycle::CancellationToken, - middleware::{ - compose, - http::{self as http, upgrade}, - websocket as core, - }, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareHeader; +use gateway_core::engine::middleware::http; +use gateway_core::engine::middleware::http::upgrade; +use gateway_core::engine::middleware::websocket as core; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::middleware::compose; use http_body_util::BodyExt as _; pub(super) struct Upgrade { diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs new file mode 100644 index 000000000..4e9bd4a74 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/frame.rs @@ -0,0 +1,171 @@ +//! 下游连接的帧、写入结果、退出原因与队列预算 + +use std::{sync::Arc, time::Duration}; +use thiserror::Error; +use tokio::sync::OwnedSemaphorePermit; + +pub(super) const OUTBOUND_COMMAND_BUFFER: usize = 32; +pub(super) const INBOUND_EVENT_BUFFER: usize = 32; +pub(super) const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60); +pub(super) const DOWNSTREAM_CLOSE_TIMEOUT: Duration = Duration::from_secs(1); +pub(super) const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60); +pub(super) const DOWNSTREAM_PING_INTERVAL: Duration = Duration::from_secs(25); + +/// WebSocket pump 的写入与生命周期预算 +#[derive(Clone, Copy)] +pub struct ConnectionConfig { + pub(super) write_timeout: Duration, + pub(super) max_age: Duration, +} + +impl ConnectionConfig { + /// 生产环境固定预算 + pub const PRODUCTION: Self = Self { + write_timeout: DOWNSTREAM_WRITE_TIMEOUT, + max_age: CONNECTION_MAX_AGE, + }; +} + +/// 业务层可观察的客户端输入；Ping/Pong 始终由 pump 消费 +pub enum ConnectionEvent { + Text(String), + Binary, + Expired, + Exited(PumpExitReason), +} + +/// 接收队列与活动期间暂存的请求共用容量；移动帧不释放它占用的名额 +pub(in super::super) struct PendingConnectionEvent { + pub(in super::super) event: ConnectionEvent, + pub(super) _permit: Option, +} + +impl PendingConnectionEvent { + pub(super) fn exited(reason: PumpExitReason) -> Self { + Self { + event: ConnectionEvent::Exited(reason), + _permit: None, + } + } +} + +/// 下游写入阶段；名称刻意使用 write，而不是暗示客户端已消费的 delivery +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FramePhase { + Metadata, + First, + Data, + Terminal, + FirstAndTerminal, + Error, + ConnectionLimit, + Close, +} + +impl FramePhase { + pub(super) const fn as_str(self) -> &'static str { + match self { + Self::Metadata => "metadata", + Self::First => "first", + Self::Data => "data", + Self::Terminal => "terminal", + Self::FirstAndTerminal => "first_and_terminal", + Self::Error => "error", + Self::ConnectionLimit => "connection_limit", + Self::Close => "close", + } + } + + pub(super) const fn is_milestone(self) -> bool { + matches!( + self, + Self::First + | Self::Terminal + | Self::FirstAndTerminal + | Self::Error + | Self::ConnectionLimit + | Self::Close + ) + } +} + +/// 一次下游写入的请求归属与协议阶段 +#[derive(Clone)] +pub struct WriteContext { + pub(super) request_id: Option>, + pub(super) phase: FramePhase, +} + +impl WriteContext { + /// 创建归属于某个请求的写入上下文 + #[must_use] + pub fn request(request_id: &Arc, phase: FramePhase) -> Self { + Self { + request_id: Some(Arc::clone(request_id)), + phase, + } + } + + /// 创建连接级写入上下文 + #[must_use] + pub const fn connection(phase: FramePhase) -> Self { + Self { + request_id: None, + phase, + } + } +} + +/// WebSocket pump 停止的稳定原因 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PumpExitReason { + ClientClose, + PeerEof, + ReadError, + WriteError, + WriteTimeout, + LifecycleShutdown, + ConnectionMaxAge, + CoordinatorDropped, + InboundOverload, + ServerClose, + PumpStopped, + MiddlewareFailed, +} + +impl PumpExitReason { + pub(in super::super) const fn as_str(self) -> &'static str { + match self { + Self::ClientClose => "client_close", + Self::PeerEof => "peer_eof", + Self::ReadError => "read_error", + Self::WriteError => "write_error", + Self::WriteTimeout => "write_timeout", + Self::LifecycleShutdown => "lifecycle_shutdown", + Self::ConnectionMaxAge => "connection_max_age", + Self::CoordinatorDropped => "coordinator_dropped", + Self::InboundOverload => "inbound_overload", + Self::ServerClose => "server_close", + Self::PumpStopped => "pump_stopped", + Self::MiddlewareFailed => "middleware_failed", + } + } +} + +/// 下游 WebSocket 写入失败 +#[derive(Debug, Error)] +pub enum ConnectionWriteError { + #[error("downstream WebSocket pump is closed")] + Closed, + #[error("downstream WebSocket write timed out after {timeout:?}")] + Timeout { timeout: Duration }, + #[error("downstream WebSocket transport write failed: {message}")] + Transport { message: String }, +} + +/// 只有实际 transport 写入才记为 Written；插件丢弃消息不伪造写入成功 +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WriteOutcome { + Written, + Suppressed, +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs new file mode 100644 index 000000000..59d79562c --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs @@ -0,0 +1,362 @@ +//! 协调层持有的连接句柄与唯一 pump 启动入口 + +use super::{ + frame::*, + pump::run_pump, + state::{ConnectionCommand, ConnectionStats}, +}; +use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code}; +use futures::{Sink, Stream}; +use gateway_core::{ + engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader}, + lifecycle::CancellationToken, +}; +use std::{ + collections::VecDeque, + fmt, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; +use tokio::{ + sync::{mpsc, oneshot}, + task::JoinHandle, + time::{Instant, timeout}, +}; + +/// 协调层持有的单 owner WebSocket 连接句柄 +pub struct ResponsesWebSocketConnection { + connection_id: Arc, + opened_at: Instant, + expired: Arc, + commands: Option>, + incoming: mpsc::Receiver, + deferred: VecDeque, + exited: oneshot::Receiver, + pump_task: Option>, + stats: Arc, + config: ConnectionConfig, + exit_reason: Option, +} + +impl ResponsesWebSocketConnection { + pub(in super::super) fn new( + socket: WebSocket, + connection_id: String, + cancellation: CancellationToken, + middleware: Option, + headers: Arc<[MiddlewareHeader]>, + ) -> Self { + spawn_connection( + socket, + Arc::::from(connection_id), + cancellation, + ConnectionConfig::PRODUCTION, + middleware, + headers, + ) + } + + pub(in super::super) fn id(&self) -> &str { + &self.connection_id + } + + /// 返回连接是否已经达到生命周期上限 + #[must_use] + pub fn is_expired(&self) -> bool { + self.expired.load(Ordering::Acquire) + } + + pub(in super::super) fn age(&self) -> Duration { + self.opened_at.elapsed() + } + + /// 等待下一个需要业务层处理的客户端事件 + pub async fn next_event(&mut self) -> Option { + if self.exit_reason.is_none() { + match self.exited.try_recv() { + Ok(reason) => { + self.exit_reason = Some(reason); + } + Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { + self.exit_reason = Some(PumpExitReason::PumpStopped); + } + Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {} + } + } + if let Some(reason) = self.exit_reason { + return Some(ConnectionEvent::Exited(reason)); + } + if let Some(event) = self.deferred.pop_front() { + return Some(event.event); + } + self.next_active_event().await.map(|event| event.event) + } + + /// 活动响应期间只消费新到的帧，已排队的下一轮请求继续保持串行 + pub(in super::super) async fn next_active_event(&mut self) -> Option { + if let Some(reason) = self.exit_reason { + return Some(PendingConnectionEvent::exited(reason)); + } + let event = tokio::select! { + biased; + reason = &mut self.exited => { + Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped))) + } + event = self.incoming.recv() => event, + }; + if let Some(PendingConnectionEvent { + event: ConnectionEvent::Exited(reason), + .. + }) = event.as_ref() + { + self.exit_reason.get_or_insert(*reason); + } else if event.is_none() { + self.exit_reason.get_or_insert(PumpExitReason::PumpStopped); + } + event + } + + pub(in super::super) fn defer(&mut self, event: PendingConnectionEvent) { + self.deferred.push_back(event); + } + + /// 等待连接退出，不消费留给后续串行请求的业务帧 + pub async fn wait_for_exit(&mut self) -> PumpExitReason { + if let Some(reason) = self.exit_reason { + return reason; + } + let reason = (&mut self.exited) + .await + .unwrap_or(PumpExitReason::PumpStopped); + self.exit_reason = Some(reason); + reason + } + + /// 串行写入文本帧，并等待 pump 确认 transport 写入结果 + /// + /// # Errors + /// + /// pump 已关闭、写入超时或 transport 失败时返回稳定错误 + pub async fn send_text( + &mut self, + payload: String, + context: WriteContext, + ) -> Result { + self.send(Message::Text(payload.into()), context).await + } + + pub(in super::super) async fn close_policy( + &mut self, + reason: &'static str, + request_id: Option<&Arc>, + ) { + let context = request_id.map_or_else( + || WriteContext::connection(FramePhase::Close), + |request_id| WriteContext::request(request_id, FramePhase::Close), + ); + self.close( + CloseFrame { + code: close_code::POLICY, + reason: reason.into(), + }, + context, + PumpExitReason::ServerClose, + ) + .await; + } + + pub(in super::super) async fn close_for_connection_limit(&mut self, reason: &'static str) { + self.close( + CloseFrame { + code: close_code::NORMAL, + reason: reason.into(), + }, + WriteContext::connection(FramePhase::Close), + PumpExitReason::ConnectionMaxAge, + ) + .await; + } + + async fn close( + &mut self, + frame: CloseFrame, + context: WriteContext, + exit_reason: PumpExitReason, + ) { + if self + .send(Message::Close(Some(frame)), context) + .await + .is_ok() + { + self.exit_reason.get_or_insert(exit_reason); + } + } + + async fn send( + &mut self, + message: Message, + context: WriteContext, + ) -> Result { + let Some(commands) = self.commands.as_ref().cloned() else { + return Err(ConnectionWriteError::Closed); + }; + self.stats.observe_command_queue(&commands); + let (acknowledged, acknowledgement) = oneshot::channel(); + let started_at = Instant::now(); + let request_id = context.request_id.clone(); + let phase = context.phase; + let command = ConnectionCommand { + message, + context, + acknowledged, + }; + let write = async move { + commands + .send(command) + .await + .map_err(|_| ConnectionWriteError::Closed)?; + acknowledgement + .await + .map_err(|_| ConnectionWriteError::Closed)? + }; + let result = match timeout(self.config.write_timeout, write).await { + Ok(result) => result, + Err(_) => Err(ConnectionWriteError::Timeout { + timeout: self.config.write_timeout, + }), + }; + let duration = started_at.elapsed(); + match &result { + Ok(WriteOutcome::Written) if phase.is_milestone() => tracing::info!( + websocket_connection_id = %self.connection_id, + request_id = request_id.as_deref().unwrap_or(""), + frame_phase = phase.as_str(), + write_duration_ms = duration.as_millis(), + "Responses WebSocket frame write succeeded" + ), + Ok(outcome) => tracing::debug!( + websocket_connection_id = %self.connection_id, + request_id = request_id.as_deref().unwrap_or(""), + frame_phase = phase.as_str(), + write_duration_ms = duration.as_millis(), + ?outcome, + "Responses WebSocket frame processed" + ), + Err(error) => tracing::info!( + websocket_connection_id = %self.connection_id, + request_id = request_id.as_deref().unwrap_or(""), + frame_phase = phase.as_str(), + write_duration_ms = duration.as_millis(), + error = %error, + "Responses WebSocket frame write failed" + ), + } + if let Err(error) = &result { + let reason = match error { + ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, + ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, + ConnectionWriteError::Closed => PumpExitReason::PumpStopped, + }; + self.terminate(reason); + } + result + } + + fn terminate(&mut self, reason: PumpExitReason) { + self.exit_reason.get_or_insert(reason); + self.commands.take(); + if let Some(task) = self.pump_task.take() { + task.abort(); + } + } + + pub(in super::super) fn log_summary(&self, request_count: u64) { + let connection_age_ms = self.age().as_millis(); + let read_idle_ms = connection_age_ms + .saturating_sub(u128::from(self.stats.last_read_ms.load(Ordering::Relaxed))); + let write_idle_ms = connection_age_ms + .saturating_sub(u128::from(self.stats.last_write_ms.load(Ordering::Relaxed))); + tracing::info!( + websocket_connection_id = %self.connection_id, + request_count, + connection_age_ms, + pump_exit_reason = self + .exit_reason + .map_or("unknown", PumpExitReason::as_str), + command_queue_high_water = self + .stats + .command_queue_high_water + .load(Ordering::Relaxed), + command_backpressure_count = self + .stats + .command_backpressure_count + .load(Ordering::Relaxed), + ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed), + ping_written_count = self.stats.ping_written_count.load(Ordering::Relaxed), + pong_received_count = self.stats.pong_received_count.load(Ordering::Relaxed), + read_idle_ms, + write_idle_ms, + "Responses WebSocket disconnected" + ); + } +} + +impl Drop for ResponsesWebSocketConnection { + fn drop(&mut self) { + self.commands.take(); + if let Some(task) = self.pump_task.take() { + task.abort(); + } + } +} + +/// 为 WebSocket transport 启动包含冻结中间件计划的单 owner pump +pub fn spawn_connection( + socket: S, + connection_id: Arc, + cancellation: CancellationToken, + config: ConnectionConfig, + middleware: Option, + headers: Arc<[MiddlewareHeader]>, +) -> ResponsesWebSocketConnection +where + S: Stream> + Sink + Unpin + Send + 'static, + E: fmt::Display + Send + 'static, +{ + let opened_at = Instant::now(); + let expired = Arc::new(AtomicBool::new(false)); + let stats = Arc::new(ConnectionStats::default()); + let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER); + let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER); + let (exit_tx, exit_rx) = oneshot::channel(); + let pump_task = tokio::spawn(run_pump( + socket, + command_rx, + incoming_tx, + exit_tx, + Arc::clone(&connection_id), + cancellation, + opened_at, + Arc::clone(&expired), + Arc::clone(&stats), + config, + middleware, + headers, + )); + ResponsesWebSocketConnection { + connection_id, + opened_at, + expired, + commands: Some(command_tx), + incoming: incoming_rx, + deferred: VecDeque::new(), + exited: exit_rx, + pump_task: Some(pump_task), + stats, + config, + exit_reason: None, + } +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs index 4f2f64dae..44267c054 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/middleware.rs @@ -1,11 +1,13 @@ //! 插件消息组合在 transport 锁之外执行，主动发送复用唯一 writer -use super::{ConnectionWriteError, PumpContext, write_message}; +use super::{ConnectionWriteError, frame::DOWNSTREAM_CLOSE_TIMEOUT, state::PumpContext}; pub(super) use crate::middleware::websocket::CancelOnDrop; use crate::middleware::websocket::{from_core, into_core}; use axum::extract::ws::Message; -use futures::{Sink, future::BoxFuture}; -use gateway_core::{engine::middleware::MiddlewareError, middleware::websocket as core}; -use std::sync::Arc; +use futures::{Sink, SinkExt, future::BoxFuture}; +use gateway_core::engine::middleware::{MiddlewareError, websocket as core}; +use gateway_core::lifecycle::CancellationToken; +use std::{fmt, sync::Arc}; +use tokio::time::timeout; pub(super) struct Writer { socket: tokio::sync::Mutex, @@ -37,6 +39,36 @@ where }) => result.unwrap_or(Err(ConnectionWriteError::Timeout { timeout: self.context.config.write_timeout })), } } + + pub(super) async fn finish_peer_close( + &self, + cancellation: &CancellationToken, + ) -> Result<(), ConnectionWriteError> { + // 主动 Sender 也可能在 pump 之外持锁；先独占 transport 再取消业务， + // 有在途写或已取消的半帧时立即跳过，不能等待它释放锁后继续刷新 + let socket = self + .socket + .try_lock() + .ok() + .filter(|_| !self.context.cancellation.is_cancelled()); + self.context.cancellation.cancel(); + let Some(mut socket) = socket else { + return Err(ConnectionWriteError::Closed); + }; + // 收到 Close 后 transport 已排队应答；只刷新它，不能再发送一份 Close + // 业务调用已取消，清理只受宿主取消与独立短预算约束，不沿用业务写入的长超时 + tokio::select! { + biased; + () = cancellation.cancelled() => Err(ConnectionWriteError::Closed), + result = timeout(DOWNSTREAM_CLOSE_TIMEOUT, async { + socket.flush().await.map_err(|error| { + ConnectionWriteError::Transport { message: error.to_string() } + }) + }) => result.unwrap_or(Err(ConnectionWriteError::Timeout { + timeout: DOWNSTREAM_CLOSE_TIMEOUT, + })), + } + } } impl core::Sender for Writer where @@ -81,3 +113,25 @@ pub(super) async fn transform( .map(from_core) .transpose() } + +async fn write_message( + socket: &mut S, + message: Message, + context: &PumpContext, +) -> Result<(), ConnectionWriteError> +where + S: Sink + Unpin, + E: fmt::Display, +{ + tokio::select! { + biased; + () = context.cancellation.cancelled() => Err(ConnectionWriteError::Closed), + result = timeout(context.config.write_timeout, socket.send(message)) => { + match result { + Ok(Ok(())) => { context.stats.record_write(context.opened_at); Ok(()) } + Ok(Err(error)) => Err(ConnectionWriteError::Transport { message: error.to_string() }), + Err(_) => Err(ConnectionWriteError::Timeout { timeout: context.config.write_timeout }), + } + } + } +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs index 42db1d2fb..981bd2117 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/mod.rs @@ -1,806 +1,13 @@ -//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界 +//! 下游客户端 WebSocket 的连接合同、句柄与单 owner pump +mod frame; +mod handle; mod middleware; +mod pump; +mod state; -use std::{ - collections::VecDeque, - fmt, - sync::{ - Arc, - atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, - }, - time::Duration, +pub use frame::{ + ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase, PumpExitReason, + WriteContext, WriteOutcome, }; - -use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code}; -use futures::{Sink, SinkExt, Stream, StreamExt}; -use gateway_core::{ - engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader}, - lifecycle::CancellationToken, - middleware::websocket as contract, -}; -use thiserror::Error; -use tokio::{ - sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}, - task::JoinHandle, - time::{Instant, timeout}, -}; - -const OUTBOUND_COMMAND_BUFFER: usize = 32; -const INBOUND_EVENT_BUFFER: usize = 32; -const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60); -const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60); -const DOWNSTREAM_PING_INTERVAL: Duration = Duration::from_secs(25); - -/// WebSocket pump 的写入与生命周期预算 -#[derive(Clone, Copy)] -pub struct ConnectionConfig { - write_timeout: Duration, - max_age: Duration, -} - -impl ConnectionConfig { - /// 生产环境固定预算 - pub const PRODUCTION: Self = Self { - write_timeout: DOWNSTREAM_WRITE_TIMEOUT, - max_age: CONNECTION_MAX_AGE, - }; -} - -/// 业务层可观察的客户端输入；Ping/Pong 始终由 pump 消费 -pub enum ConnectionEvent { - Text(String), - Binary, - Expired, - Exited(PumpExitReason), -} - -/// 接收队列与活动期间暂存的请求共用容量；移动帧不释放它占用的名额 -pub(super) struct PendingConnectionEvent { - pub(super) event: ConnectionEvent, - _permit: Option, -} - -impl PendingConnectionEvent { - fn exited(reason: PumpExitReason) -> Self { - Self { - event: ConnectionEvent::Exited(reason), - _permit: None, - } - } -} - -/// 下游写入阶段；名称刻意使用 write，而不是暗示客户端已消费的 delivery -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum FramePhase { - Metadata, - First, - Data, - Terminal, - FirstAndTerminal, - Error, - ConnectionLimit, - Close, -} - -impl FramePhase { - const fn as_str(self) -> &'static str { - match self { - Self::Metadata => "metadata", - Self::First => "first", - Self::Data => "data", - Self::Terminal => "terminal", - Self::FirstAndTerminal => "first_and_terminal", - Self::Error => "error", - Self::ConnectionLimit => "connection_limit", - Self::Close => "close", - } - } - - const fn is_milestone(self) -> bool { - matches!( - self, - Self::First - | Self::Terminal - | Self::FirstAndTerminal - | Self::Error - | Self::ConnectionLimit - | Self::Close - ) - } -} - -/// 一次下游写入的请求归属与协议阶段 -#[derive(Clone)] -pub struct WriteContext { - request_id: Option>, - phase: FramePhase, -} - -impl WriteContext { - /// 创建归属于某个请求的写入上下文 - #[must_use] - pub fn request(request_id: &Arc, phase: FramePhase) -> Self { - Self { - request_id: Some(Arc::clone(request_id)), - phase, - } - } - - /// 创建连接级写入上下文 - #[must_use] - pub const fn connection(phase: FramePhase) -> Self { - Self { - request_id: None, - phase, - } - } -} - -/// WebSocket pump 停止的稳定原因 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PumpExitReason { - ClientClose, - PeerEof, - ReadError, - WriteError, - WriteTimeout, - LifecycleShutdown, - ConnectionMaxAge, - CoordinatorDropped, - InboundOverload, - ServerClose, - PumpStopped, - MiddlewareFailed, -} - -impl PumpExitReason { - pub(super) const fn as_str(self) -> &'static str { - match self { - Self::ClientClose => "client_close", - Self::PeerEof => "peer_eof", - Self::ReadError => "read_error", - Self::WriteError => "write_error", - Self::WriteTimeout => "write_timeout", - Self::LifecycleShutdown => "lifecycle_shutdown", - Self::ConnectionMaxAge => "connection_max_age", - Self::CoordinatorDropped => "coordinator_dropped", - Self::InboundOverload => "inbound_overload", - Self::ServerClose => "server_close", - Self::PumpStopped => "pump_stopped", - Self::MiddlewareFailed => "middleware_failed", - } - } -} - -/// 下游 WebSocket 写入失败 -#[derive(Debug, Error)] -pub enum ConnectionWriteError { - #[error("downstream WebSocket pump is closed")] - Closed, - #[error("downstream WebSocket write timed out after {timeout:?}")] - Timeout { timeout: Duration }, - #[error("downstream WebSocket transport write failed: {message}")] - Transport { message: String }, -} - -/// 只有实际 transport 写入才记为 Written；插件丢弃消息不伪造写入成功 -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum WriteOutcome { - Written, - Suppressed, -} - -struct ConnectionCommand { - message: Message, - context: WriteContext, - acknowledged: oneshot::Sender>, -} - -#[derive(Default)] -struct ConnectionStats { - command_queue_high_water: AtomicUsize, - command_backpressure_count: AtomicU64, - ping_received_count: AtomicU64, - ping_written_count: AtomicU64, - pong_received_count: AtomicU64, - last_read_ms: AtomicU64, - last_write_ms: AtomicU64, -} - -impl ConnectionStats { - fn record_read(&self, opened_at: Instant) { - self.last_read_ms.store( - u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), - Ordering::Relaxed, - ); - } - - fn record_write(&self, opened_at: Instant) { - self.last_write_ms.store( - u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), - Ordering::Relaxed, - ); - } - - fn observe_command_queue(&self, sender: &mpsc::Sender) { - let queued = OUTBOUND_COMMAND_BUFFER - .saturating_sub(sender.capacity()) - .saturating_add(1) - .min(OUTBOUND_COMMAND_BUFFER); - self.command_queue_high_water - .fetch_max(queued, Ordering::Relaxed); - if sender.capacity() == 0 { - self.command_backpressure_count - .fetch_add(1, Ordering::Relaxed); - } - } -} - -/// 协调层持有的单 owner WebSocket 连接句柄 -pub struct ResponsesWebSocketConnection { - connection_id: Arc, - opened_at: Instant, - expired: Arc, - commands: Option>, - incoming: mpsc::Receiver, - deferred: VecDeque, - exited: oneshot::Receiver, - pump_task: Option>, - stats: Arc, - config: ConnectionConfig, - exit_reason: Option, -} - -impl ResponsesWebSocketConnection { - pub(super) fn new( - socket: WebSocket, - connection_id: String, - cancellation: CancellationToken, - middleware: Option, - headers: Arc<[MiddlewareHeader]>, - ) -> Self { - spawn_connection( - socket, - Arc::::from(connection_id), - cancellation, - ConnectionConfig::PRODUCTION, - middleware, - headers, - ) - } - - pub(super) fn id(&self) -> &str { - &self.connection_id - } - - /// 返回连接是否已经达到生命周期上限 - #[must_use] - pub fn is_expired(&self) -> bool { - self.expired.load(Ordering::Acquire) - } - - pub(super) fn age(&self) -> Duration { - self.opened_at.elapsed() - } - - /// 等待下一个需要业务层处理的客户端事件 - pub async fn next_event(&mut self) -> Option { - if self.exit_reason.is_none() { - match self.exited.try_recv() { - Ok(reason) => { - self.exit_reason = Some(reason); - } - Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { - self.exit_reason = Some(PumpExitReason::PumpStopped); - } - Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {} - } - } - if let Some(reason) = self.exit_reason { - return Some(ConnectionEvent::Exited(reason)); - } - if let Some(event) = self.deferred.pop_front() { - return Some(event.event); - } - self.next_active_event().await.map(|event| event.event) - } - - /// 活动响应期间只消费新到的帧，已排队的下一轮请求继续保持串行 - pub(super) async fn next_active_event(&mut self) -> Option { - if let Some(reason) = self.exit_reason { - return Some(PendingConnectionEvent::exited(reason)); - } - let event = tokio::select! { - biased; - reason = &mut self.exited => { - Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped))) - } - event = self.incoming.recv() => event, - }; - if let Some(PendingConnectionEvent { - event: ConnectionEvent::Exited(reason), - .. - }) = event.as_ref() - { - self.exit_reason.get_or_insert(*reason); - } else if event.is_none() { - self.exit_reason.get_or_insert(PumpExitReason::PumpStopped); - } - event - } - - pub(super) fn defer(&mut self, event: PendingConnectionEvent) { - self.deferred.push_back(event); - } - - /// 等待连接退出，不消费留给后续串行请求的业务帧 - pub async fn wait_for_exit(&mut self) -> PumpExitReason { - if let Some(reason) = self.exit_reason { - return reason; - } - let reason = (&mut self.exited) - .await - .unwrap_or(PumpExitReason::PumpStopped); - self.exit_reason = Some(reason); - reason - } - - /// 串行写入文本帧，并等待 pump 确认 transport 写入结果 - /// - /// # Errors - /// - /// pump 已关闭、写入超时或 transport 失败时返回稳定错误 - pub async fn send_text( - &mut self, - payload: String, - context: WriteContext, - ) -> Result { - self.send(Message::Text(payload.into()), context).await - } - - pub(super) async fn close_policy( - &mut self, - reason: &'static str, - request_id: Option<&Arc>, - ) { - let context = request_id.map_or_else( - || WriteContext::connection(FramePhase::Close), - |request_id| WriteContext::request(request_id, FramePhase::Close), - ); - self.close( - CloseFrame { - code: close_code::POLICY, - reason: reason.into(), - }, - context, - PumpExitReason::ServerClose, - ) - .await; - } - - pub(super) async fn close_for_connection_limit(&mut self, reason: &'static str) { - self.close( - CloseFrame { - code: close_code::NORMAL, - reason: reason.into(), - }, - WriteContext::connection(FramePhase::Close), - PumpExitReason::ConnectionMaxAge, - ) - .await; - } - - async fn close( - &mut self, - frame: CloseFrame, - context: WriteContext, - exit_reason: PumpExitReason, - ) { - if self - .send(Message::Close(Some(frame)), context) - .await - .is_ok() - { - self.exit_reason.get_or_insert(exit_reason); - } - } - - async fn send( - &mut self, - message: Message, - context: WriteContext, - ) -> Result { - let Some(commands) = self.commands.as_ref().cloned() else { - return Err(ConnectionWriteError::Closed); - }; - self.stats.observe_command_queue(&commands); - let (acknowledged, acknowledgement) = oneshot::channel(); - let started_at = Instant::now(); - let request_id = context.request_id.clone(); - let phase = context.phase; - let command = ConnectionCommand { - message, - context, - acknowledged, - }; - let write = async move { - commands - .send(command) - .await - .map_err(|_| ConnectionWriteError::Closed)?; - acknowledgement - .await - .map_err(|_| ConnectionWriteError::Closed)? - }; - let result = match timeout(self.config.write_timeout, write).await { - Ok(result) => result, - Err(_) => Err(ConnectionWriteError::Timeout { - timeout: self.config.write_timeout, - }), - }; - let duration = started_at.elapsed(); - match &result { - Ok(WriteOutcome::Written) if phase.is_milestone() => tracing::info!( - websocket_connection_id = %self.connection_id, - request_id = request_id.as_deref().unwrap_or(""), - frame_phase = phase.as_str(), - write_duration_ms = duration.as_millis(), - "Responses WebSocket frame write succeeded" - ), - Ok(outcome) => tracing::debug!( - websocket_connection_id = %self.connection_id, - request_id = request_id.as_deref().unwrap_or(""), - frame_phase = phase.as_str(), - write_duration_ms = duration.as_millis(), - ?outcome, - "Responses WebSocket frame processed" - ), - Err(error) => tracing::info!( - websocket_connection_id = %self.connection_id, - request_id = request_id.as_deref().unwrap_or(""), - frame_phase = phase.as_str(), - write_duration_ms = duration.as_millis(), - error = %error, - "Responses WebSocket frame write failed" - ), - } - if let Err(error) = &result { - let reason = match error { - ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, - ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, - ConnectionWriteError::Closed => PumpExitReason::PumpStopped, - }; - self.terminate(reason); - } - result - } - - fn terminate(&mut self, reason: PumpExitReason) { - self.exit_reason.get_or_insert(reason); - self.commands.take(); - if let Some(task) = self.pump_task.take() { - task.abort(); - } - } - - pub(super) fn log_summary(&self, request_count: u64) { - let connection_age_ms = self.age().as_millis(); - let read_idle_ms = connection_age_ms - .saturating_sub(u128::from(self.stats.last_read_ms.load(Ordering::Relaxed))); - let write_idle_ms = connection_age_ms - .saturating_sub(u128::from(self.stats.last_write_ms.load(Ordering::Relaxed))); - tracing::info!( - websocket_connection_id = %self.connection_id, - request_count, - connection_age_ms, - pump_exit_reason = self - .exit_reason - .map_or("unknown", PumpExitReason::as_str), - command_queue_high_water = self - .stats - .command_queue_high_water - .load(Ordering::Relaxed), - command_backpressure_count = self - .stats - .command_backpressure_count - .load(Ordering::Relaxed), - ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed), - ping_written_count = self.stats.ping_written_count.load(Ordering::Relaxed), - pong_received_count = self.stats.pong_received_count.load(Ordering::Relaxed), - read_idle_ms, - write_idle_ms, - "Responses WebSocket disconnected" - ); - } -} - -impl Drop for ResponsesWebSocketConnection { - fn drop(&mut self) { - self.commands.take(); - if let Some(task) = self.pump_task.take() { - task.abort(); - } - } -} - -/// 为 WebSocket transport 启动包含冻结中间件计划的单 owner pump -pub fn spawn_connection( - socket: S, - connection_id: Arc, - cancellation: CancellationToken, - config: ConnectionConfig, - middleware: Option, - headers: Arc<[MiddlewareHeader]>, -) -> ResponsesWebSocketConnection -where - S: Stream> + Sink + Unpin + Send + 'static, - E: fmt::Display + Send + 'static, -{ - let opened_at = Instant::now(); - let expired = Arc::new(AtomicBool::new(false)); - let stats = Arc::new(ConnectionStats::default()); - let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER); - let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER); - let (exit_tx, exit_rx) = oneshot::channel(); - let pump_task = tokio::spawn(run_pump( - socket, - command_rx, - incoming_tx, - exit_tx, - Arc::clone(&connection_id), - cancellation, - opened_at, - Arc::clone(&expired), - Arc::clone(&stats), - config, - middleware, - headers, - )); - ResponsesWebSocketConnection { - connection_id, - opened_at, - expired, - commands: Some(command_tx), - incoming: incoming_rx, - deferred: VecDeque::new(), - exited: exit_rx, - pump_task: Some(pump_task), - stats, - config, - exit_reason: None, - } -} - -#[expect(clippy::too_many_arguments)] -async fn run_pump( - socket: S, - commands: mpsc::Receiver, - incoming: mpsc::Sender, - exited: oneshot::Sender, - connection_id: Arc, - cancellation: CancellationToken, - opened_at: Instant, - expired: Arc, - stats: Arc, - config: ConnectionConfig, - middleware: Option, - headers: Arc<[MiddlewareHeader]>, -) where - S: Stream> + Sink + Unpin + Send + 'static, - E: fmt::Display + Send + 'static, -{ - let context = Arc::new(PumpContext { - connection_id, - cancellation: cancellation.child_token(), - opened_at, - expired, - stats, - config, - middleware, - headers, - }); - let _lifetime = middleware::CancelOnDrop(Some(context.cancellation.clone())); - let (writer, mut reader) = socket.split(); - let writer = Arc::new(middleware::Writer::new(writer, context.clone())); - let sender: Arc = writer.clone(); - // 同一个受管任务同时驱动两侧；等待写入不能阻塞入站控制、关闭或取消 - let reason = tokio::select! { - biased; - reason = read_connection(&mut reader, &incoming, sender.clone(), &context) => reason, - reason = write_connection(writer.clone(), commands, &context) => reason, - }; - context.cancellation.cancel(); - let _ = exited.send(reason); - tracing::debug!(websocket_connection_id = %context.connection_id, pump_exit_reason = reason.as_str(), "Responses WebSocket pump exited"); -} - -struct PumpContext { - middleware: Option, - headers: Arc<[MiddlewareHeader]>, - connection_id: Arc, - cancellation: CancellationToken, - opened_at: Instant, - expired: Arc, - stats: Arc, - config: ConnectionConfig, -} - -async fn read_connection( - socket: &mut S, - incoming: &mpsc::Sender, - sender: Arc, - context: &PumpContext, -) -> PumpExitReason -where - S: Stream> + Unpin, - E: fmt::Display, -{ - let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER)); - let deadline = tokio::time::sleep_until(context.opened_at + context.config.max_age); - tokio::pin!(deadline); - let mut deadline_elapsed = false; - loop { - tokio::select! { - biased; - () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, - () = &mut deadline, if !deadline_elapsed => { - deadline_elapsed = true; - context.expired.store(true, Ordering::Release); - // 到期只阻止下一轮，当前响应仍由协调层按原有合同收尾 - match emit_incoming(incoming, &event_slots, ConnectionEvent::Expired) { - Ok(()) | Err(PumpExitReason::InboundOverload) => {} - Err(reason) => return reason, - } - } - message = socket.next() => { - if matches!(&message, Some(Ok(_))) { context.stats.record_read(context.opened_at); } - let message = match message { - Some(Ok(message)) => match middleware::transform(message, contract::Direction::Incoming, sender.clone(), context).await { - Ok(Some(message)) => Some(Ok(message)), - Ok(None) => continue, - Err(_) => return PumpExitReason::MiddlewareFailed, - }, - other => other, - }; - let event = match message { - Some(Ok(Message::Text(payload))) => Some(ConnectionEvent::Text(payload.to_string())), - Some(Ok(Message::Binary(_))) => Some(ConnectionEvent::Binary), - Some(Ok(Message::Ping(_))) => { - // Axum/tungstenite 在继续读取时自动刷新 Pong，不能再手工发送一份 - context.stats.ping_received_count.fetch_add(1, Ordering::Relaxed); - None - } - Some(Ok(Message::Pong(_))) => { - context.stats.pong_received_count.fetch_add(1, Ordering::Relaxed); - None - } - Some(Ok(Message::Close(_))) => return PumpExitReason::ClientClose, - Some(Err(error)) => { - tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket receive failed"); - return PumpExitReason::ReadError; - } - None => return PumpExitReason::PeerEof, - }; - if let Some(event) = event && let Err(reason) = emit_incoming(incoming, &event_slots, event) { return reason; } - } - } - } -} - -async fn write_connection( - writer: Arc>, - mut commands: mpsc::Receiver, - context: &PumpContext, -) -> PumpExitReason -where - S: Sink + Unpin + Send + 'static, - E: fmt::Display + Send + 'static, -{ - let mut heartbeat = tokio::time::interval_at( - context.opened_at + DOWNSTREAM_PING_INTERVAL, - DOWNSTREAM_PING_INTERVAL, - ); - heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut ping_sequence = 0_u64; - loop { - let message = tokio::select! { - biased; - () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, - command = commands.recv() => { - let Some(command) = command else { return PumpExitReason::CoordinatorDropped; }; - let ConnectionCommand { message, context: write_context, acknowledged } = command; - let closing = matches!(message, Message::Close(_)); - let result = match middleware::transform(message, contract::Direction::Outgoing, writer.clone(), context).await { - Ok(Some(message)) => writer.write(message).await.map(|()| WriteOutcome::Written), - Ok(None) => Ok(WriteOutcome::Suppressed), - Err(error) => Err(ConnectionWriteError::Transport { message: error.to_string() }), - }; - let reason = result.as_ref().err().map(|error| write_exit_reason(error, context)); - if let Err(error) = &result { - tracing::debug!(websocket_connection_id = %context.connection_id, request_id = write_context.request_id.as_deref().unwrap_or(""), frame_phase = write_context.phase.as_str(), error = %error, "Responses WebSocket pump transport write failed"); - } - let _ = acknowledged.send(result); - if let Some(reason) = reason { return reason; } - if closing { return PumpExitReason::ServerClose; } - continue; - } - _ = heartbeat.tick() => { - ping_sequence = ping_sequence.wrapping_add(1); - Message::Ping(ping_sequence.to_be_bytes().to_vec().into()) - } - }; - let message = match middleware::transform( - message, - contract::Direction::Outgoing, - writer.clone(), - context, - ) - .await - { - Ok(Some(message)) => message, - Ok(None) => continue, - Err(_) => return PumpExitReason::MiddlewareFailed, - }; - if let Err(error) = writer.write(message).await { - tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket control write failed"); - return write_exit_reason(&error, context); - } - context - .stats - .ping_written_count - .fetch_add(1, Ordering::Relaxed); - } -} - -async fn write_message( - socket: &mut S, - message: Message, - context: &PumpContext, -) -> Result<(), ConnectionWriteError> -where - S: Sink + Unpin, - E: fmt::Display, -{ - tokio::select! { - biased; - () = context.cancellation.cancelled() => Err(ConnectionWriteError::Closed), - result = timeout(context.config.write_timeout, socket.send(message)) => { - match result { - Ok(Ok(())) => { context.stats.record_write(context.opened_at); Ok(()) } - Ok(Err(error)) => Err(ConnectionWriteError::Transport { message: error.to_string() }), - Err(_) => Err(ConnectionWriteError::Timeout { timeout: context.config.write_timeout }), - } - } - } -} - -fn write_exit_reason(error: &ConnectionWriteError, context: &PumpContext) -> PumpExitReason { - match error { - ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, - ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, - ConnectionWriteError::Closed if context.cancellation.is_cancelled() => { - PumpExitReason::LifecycleShutdown - } - ConnectionWriteError::Closed => PumpExitReason::PumpStopped, - } -} - -fn emit_incoming( - incoming: &mpsc::Sender, - slots: &Arc, - event: ConnectionEvent, -) -> Result<(), PumpExitReason> { - let permit = Arc::clone(slots) - .try_acquire_owned() - .map_err(|_| PumpExitReason::InboundOverload)?; - incoming - .try_send(PendingConnectionEvent { - event, - _permit: Some(permit), - }) - .map_err(|error| match error { - mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload, - mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped, - }) -} +pub use handle::{ResponsesWebSocketConnection, spawn_connection}; diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs new file mode 100644 index 000000000..14168917a --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs @@ -0,0 +1,230 @@ +//! 单 owner 驱动连接读写、心跳与有界入站队列 + +use super::{ + frame::*, + middleware, + state::{ConnectionCommand, ConnectionStats, PumpContext}, +}; +use axum::extract::ws::Message; +use futures::{Sink, Stream, StreamExt}; +use gateway_core::{ + engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader, websocket as contract}, + lifecycle::CancellationToken, +}; +use std::{ + fmt, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; +use tokio::{ + sync::{Semaphore, mpsc, oneshot}, + time::Instant, +}; + +#[expect(clippy::too_many_arguments)] +pub(super) async fn run_pump( + socket: S, + commands: mpsc::Receiver, + incoming: mpsc::Sender, + exited: oneshot::Sender, + connection_id: Arc, + cancellation: CancellationToken, + opened_at: Instant, + expired: Arc, + stats: Arc, + config: ConnectionConfig, + middleware: Option, + headers: Arc<[MiddlewareHeader]>, +) where + S: Stream> + Sink + Unpin + Send + 'static, + E: fmt::Display + Send + 'static, +{ + let context = Arc::new(PumpContext { + connection_id, + cancellation: cancellation.child_token(), + opened_at, + expired, + stats, + config, + middleware, + headers, + }); + let _lifetime = middleware::CancelOnDrop(Some(context.cancellation.clone())); + let (writer, mut reader) = socket.split(); + let writer = Arc::new(middleware::Writer::new(writer, context.clone())); + let sender: Arc = writer.clone(); + // 同一个受管任务同时驱动两侧；等待写入不能阻塞入站控制、关闭或取消 + let reason = tokio::select! { + biased; + reason = read_connection(&mut reader, &incoming, sender.clone(), &context) => reason, + reason = write_connection(writer.clone(), commands, &context) => reason, + }; + if matches!(reason, PumpExitReason::ClientClose) { + if let Err(error) = writer.finish_peer_close(&cancellation).await { + tracing::debug!(websocket_connection_id = %context.connection_id, %error, "Responses WebSocket close acknowledgement could not be flushed"); + } + } else { + context.cancellation.cancel(); + } + let _ = exited.send(reason); + tracing::debug!(websocket_connection_id = %context.connection_id, pump_exit_reason = reason.as_str(), "Responses WebSocket pump exited"); +} + +async fn read_connection( + socket: &mut S, + incoming: &mpsc::Sender, + sender: Arc, + context: &PumpContext, +) -> PumpExitReason +where + S: Stream> + Unpin, + E: fmt::Display, +{ + let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER)); + let deadline = tokio::time::sleep_until(context.opened_at + context.config.max_age); + tokio::pin!(deadline); + let mut deadline_elapsed = false; + loop { + tokio::select! { + biased; + () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, + () = &mut deadline, if !deadline_elapsed => { + deadline_elapsed = true; + context.expired.store(true, Ordering::Release); + // 到期只阻止下一轮，当前响应仍由协调层按原有合同收尾 + match emit_incoming(incoming, &event_slots, ConnectionEvent::Expired) { + Ok(()) | Err(PumpExitReason::InboundOverload) => {} + Err(reason) => return reason, + } + } + message = socket.next() => { + if matches!(&message, Some(Ok(_))) { context.stats.record_read(context.opened_at); } + let message = match message { + Some(Ok(message)) => match middleware::transform(message, contract::Direction::Incoming, sender.clone(), context).await { + Ok(Some(message)) => Some(Ok(message)), + Ok(None) => continue, + Err(_) => return PumpExitReason::MiddlewareFailed, + }, + other => other, + }; + let event = match message { + Some(Ok(Message::Text(payload))) => Some(ConnectionEvent::Text(payload.to_string())), + Some(Ok(Message::Binary(_))) => Some(ConnectionEvent::Binary), + Some(Ok(Message::Ping(_))) => { + // Axum/tungstenite 在继续读取时自动刷新 Pong，不能再手工发送一份 + context.stats.ping_received_count.fetch_add(1, Ordering::Relaxed); + None + } + Some(Ok(Message::Pong(_))) => { + context.stats.pong_received_count.fetch_add(1, Ordering::Relaxed); + None + } + Some(Ok(Message::Close(_))) => return PumpExitReason::ClientClose, + Some(Err(error)) => { + tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket receive failed"); + return PumpExitReason::ReadError; + } + None => return PumpExitReason::PeerEof, + }; + if let Some(event) = event && let Err(reason) = emit_incoming(incoming, &event_slots, event) { return reason; } + } + } + } +} + +async fn write_connection( + writer: Arc>, + mut commands: mpsc::Receiver, + context: &PumpContext, +) -> PumpExitReason +where + S: Sink + Unpin + Send + 'static, + E: fmt::Display + Send + 'static, +{ + let mut heartbeat = tokio::time::interval_at( + context.opened_at + DOWNSTREAM_PING_INTERVAL, + DOWNSTREAM_PING_INTERVAL, + ); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut ping_sequence = 0_u64; + loop { + let message = tokio::select! { + biased; + () = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown, + command = commands.recv() => { + let Some(command) = command else { return PumpExitReason::CoordinatorDropped; }; + let ConnectionCommand { message, context: write_context, acknowledged } = command; + let closing = matches!(message, Message::Close(_)); + let result = match middleware::transform(message, contract::Direction::Outgoing, writer.clone(), context).await { + Ok(Some(message)) => writer.write(message).await.map(|()| WriteOutcome::Written), + Ok(None) => Ok(WriteOutcome::Suppressed), + Err(error) => Err(ConnectionWriteError::Transport { message: error.to_string() }), + }; + let reason = result.as_ref().err().map(|error| write_exit_reason(error, context)); + if let Err(error) = &result { + tracing::debug!(websocket_connection_id = %context.connection_id, request_id = write_context.request_id.as_deref().unwrap_or(""), frame_phase = write_context.phase.as_str(), error = %error, "Responses WebSocket pump transport write failed"); + } + let _ = acknowledged.send(result); + if let Some(reason) = reason { return reason; } + if closing { return PumpExitReason::ServerClose; } + continue; + } + _ = heartbeat.tick() => { + ping_sequence = ping_sequence.wrapping_add(1); + Message::Ping(ping_sequence.to_be_bytes().to_vec().into()) + } + }; + let message = match middleware::transform( + message, + contract::Direction::Outgoing, + writer.clone(), + context, + ) + .await + { + Ok(Some(message)) => message, + Ok(None) => continue, + Err(_) => return PumpExitReason::MiddlewareFailed, + }; + if let Err(error) = writer.write(message).await { + tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket control write failed"); + return write_exit_reason(&error, context); + } + context + .stats + .ping_written_count + .fetch_add(1, Ordering::Relaxed); + } +} + +fn write_exit_reason(error: &ConnectionWriteError, context: &PumpContext) -> PumpExitReason { + match error { + ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout, + ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError, + ConnectionWriteError::Closed if context.cancellation.is_cancelled() => { + PumpExitReason::LifecycleShutdown + } + ConnectionWriteError::Closed => PumpExitReason::PumpStopped, + } +} + +fn emit_incoming( + incoming: &mpsc::Sender, + slots: &Arc, + event: ConnectionEvent, +) -> Result<(), PumpExitReason> { + let permit = Arc::clone(slots) + .try_acquire_owned() + .map_err(|_| PumpExitReason::InboundOverload)?; + incoming + .try_send(PendingConnectionEvent { + event, + _permit: Some(permit), + }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload, + mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped, + }) +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs new file mode 100644 index 000000000..558ea7073 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/state.rs @@ -0,0 +1,75 @@ +//! 单连接共享的命令、运行预算与观测状态 + +use super::frame::{ + ConnectionConfig, ConnectionWriteError, OUTBOUND_COMMAND_BUFFER, WriteContext, WriteOutcome, +}; +use axum::extract::ws::Message; +use gateway_core::{ + engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader}, + lifecycle::CancellationToken, +}; +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, +}; +use tokio::{ + sync::{mpsc, oneshot}, + time::Instant, +}; + +pub(super) struct ConnectionCommand { + pub(super) message: Message, + pub(super) context: WriteContext, + pub(super) acknowledged: oneshot::Sender>, +} + +#[derive(Default)] +pub(super) struct ConnectionStats { + pub(super) command_queue_high_water: AtomicUsize, + pub(super) command_backpressure_count: AtomicU64, + pub(super) ping_received_count: AtomicU64, + pub(super) ping_written_count: AtomicU64, + pub(super) pong_received_count: AtomicU64, + pub(super) last_read_ms: AtomicU64, + pub(super) last_write_ms: AtomicU64, +} + +impl ConnectionStats { + pub(super) fn record_read(&self, opened_at: Instant) { + self.last_read_ms.store( + u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + } + + pub(super) fn record_write(&self, opened_at: Instant) { + self.last_write_ms.store( + u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + } + + pub(super) fn observe_command_queue(&self, sender: &mpsc::Sender) { + let queued = OUTBOUND_COMMAND_BUFFER + .saturating_sub(sender.capacity()) + .saturating_add(1) + .min(OUTBOUND_COMMAND_BUFFER); + self.command_queue_high_water + .fetch_max(queued, Ordering::Relaxed); + if sender.capacity() == 0 { + self.command_backpressure_count + .fetch_add(1, Ordering::Relaxed); + } + } +} + +pub(super) struct PumpContext { + pub(super) middleware: Option, + pub(super) headers: Arc<[MiddlewareHeader]>, + pub(super) connection_id: Arc, + pub(super) cancellation: CancellationToken, + pub(super) opened_at: Instant, + pub(super) expired: Arc, + pub(super) stats: Arc, + pub(super) config: ConnectionConfig, +} diff --git a/backend/crates/gateway-api/tests/admin/plugins/management.rs b/backend/crates/gateway-api/tests/admin/plugins/management.rs index 1954286b3..43bc95851 100644 --- a/backend/crates/gateway-api/tests/admin/plugins/management.rs +++ b/backend/crates/gateway-api/tests/admin/plugins/management.rs @@ -5,24 +5,31 @@ use axum::{ http::{Request, StatusCode}, }; use futures::future::BoxFuture; -use gateway_core::{ - engine::{ - execution::{ - AuthenticatedClient, ClientAuthenticationError, ExecutionService, - PreparedRootExecution, StartExecution, StartProviderExecution, StartedExecution, - }, - middleware::{ - FrozenMiddlewarePlan, MiddlewareBody, MiddlewareContext, MiddlewareError, - MiddlewareFrame, MiddlewareFraming, MiddlewareHeader, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - }, - error::GatewayError, - lifecycle::CancellationToken, - policy::ClientApiKeyId, - routing::PublicModelId, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::execution::AuthenticatedClient; +use gateway_core::engine::execution::ClientAuthenticationError; +use gateway_core::engine::execution::ExecutionService; +use gateway_core::engine::execution::PreparedRootExecution; +use gateway_core::engine::execution::StartExecution; +use gateway_core::engine::execution::StartProviderExecution; +use gateway_core::engine::execution::StartedExecution; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareBody; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareFrame; +use gateway_core::engine::middleware::MiddlewareFraming; +use gateway_core::engine::middleware::MiddlewareHeader; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::error::GatewayError; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::policy::ClientApiKeyId; +use gateway_core::routing::PublicModelId; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; use std::sync::{ Arc, Mutex, atomic::{AtomicBool, Ordering}, @@ -368,10 +375,11 @@ pub(super) fn resource_fixture( } async fn published_fixture() -> AdminTestFixture { - use gateway_core::{ - routing::{ConfigRevision, RuntimeSnapshot}, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, - }; + use gateway_core::routing::ConfigRevision; + use gateway_core::routing::RuntimeSnapshot; + use gateway_core::routing::extensions::ExtensionSetId; + use gateway_core::routing::extensions::ExtensionSetLease; + use gateway_core::routing::extensions::ExtensionSetReference; use std::sync::Arc; struct Lease; impl ExtensionSetLease for Lease { diff --git a/backend/crates/gateway-api/tests/admin/plugins/mod.rs b/backend/crates/gateway-api/tests/admin/plugins/mod.rs index da463be7c..b5b11361e 100644 --- a/backend/crates/gateway-api/tests/admin/plugins/mod.rs +++ b/backend/crates/gateway-api/tests/admin/plugins/mod.rs @@ -546,7 +546,7 @@ pub(super) struct TestPluginPorts { impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts { async fn validate_target( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, ) -> Result<(), AdminError> { Ok(()) @@ -554,7 +554,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor async fn start_callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: gateway_admin::model::plugins::management::StartPluginManagementCallback, _: &gateway_admin::model::auth::AdminRequestContext, @@ -564,7 +564,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn callback( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, _: &str, request: gateway_admin::model::plugins::management::PluginManagementRequest, @@ -576,14 +576,14 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn views( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, ) -> Result, AdminError> { Ok(Vec::new()) } async fn resource( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, path: &str, public: bool, @@ -596,7 +596,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor } async fn handle( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::management::PluginManagementTarget, request: gateway_admin::model::plugins::management::PluginManagementRequest, ) -> Result @@ -611,7 +611,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { async fn activate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: &gateway_admin::model::plugins::instances::PluginInstance, ) -> Result<(), AdminError> { Ok(()) @@ -621,7 +621,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts { } async fn migrate_state( &self, - _: &gateway_core::runtime::extensions::ExtensionSetReference, + _: &gateway_core::routing::extensions::ExtensionSetReference, _: gateway_admin::model::plugins::state::PluginStateTransition, ) -> Result<(), AdminError> { panic!("unexpected state migration") @@ -647,7 +647,7 @@ impl gateway_admin::ports::plugins::PluginPreparation for TestPluginPorts { async fn prepare( &self, _: gateway_admin::model::plugins::instances::PluginInstanceSnapshot, - ) -> Result { + ) -> Result { Err(AdminError::invalid("unused plugin fixture")) } } @@ -658,7 +658,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts &self, snapshot: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot, _: Option, - _: Option<&gateway_core::runtime::extensions::ExtensionSetReference>, + _: Option<&gateway_core::routing::extensions::ExtensionSetReference>, ) -> Option> { Some( diff --git a/backend/crates/gateway-api/tests/architecture.rs b/backend/crates/gateway-api/tests/architecture.rs index 24ca76d0a..32a05ed04 100644 --- a/backend/crates/gateway-api/tests/architecture.rs +++ b/backend/crates/gateway-api/tests/architecture.rs @@ -1,4 +1,4 @@ -//! 检查 API 层依赖、源码与测试模块边界，以及管理路由约定 +//! 检查 API 层依赖与管理路由约定，模块布局和测试镜像由 workspace 架构检查覆盖 use std::{fs, path::Path}; @@ -29,150 +29,6 @@ fn workspace_should_include_gateway_api() { assert!(workspace.contains("\"crates/gateway-api\"")); } -#[test] -fn source_tree_should_match_frozen_machine_manifest() { - let root = Path::new(env!("CARGO_MANIFEST_DIR")); - let mut actual = rust_files(&root.join("src")); - actual.sort(); - let mut expected = vec![ - "src/admin/account_groups.rs", - "src/admin/accounts/credentials.rs", - "src/admin/accounts/handlers.rs", - "src/admin/accounts/import_tasks.rs", - "src/admin/accounts/mod.rs", - "src/admin/accounts/presenter.rs", - "src/admin/accounts/wire.rs", - "src/admin/auth.rs", - "src/admin/backups.rs", - "src/admin/client_keys.rs", - "src/admin/diagnostics.rs", - "src/admin/extract.rs", - "src/admin/mod.rs", - "src/admin/observability/mod.rs", - "src/admin/observability/presenter.rs", - "src/admin/observability/query.rs", - "src/admin/observability/routes.rs", - "src/admin/observability/wire.rs", - "src/admin/plugins/artifacts.rs", - "src/admin/plugins/distribution.rs", - "src/admin/plugins/instances.rs", - "src/admin/plugins/management.rs", - "src/admin/plugins/mod.rs", - "src/admin/presenter.rs", - "src/admin/proxies.rs", - "src/admin/settings.rs", - "src/admin/system.rs", - "src/admin/wire.rs", - "src/auth.rs", - "src/health.rs", - "src/key_usage/mod.rs", - "src/key_usage/query.rs", - "src/key_usage/presenter.rs", - "src/lib.rs", - "src/middleware/headers.rs", - "src/middleware/mod.rs", - "src/middleware/websocket.rs", - "src/openai/middleware/http.rs", - "src/openai/middleware/mod.rs", - "src/openai/middleware/request.rs", - "src/openai/middleware/session.rs", - "src/openai/auth.rs", - "src/openai/error.rs", - "src/openai/images/http.rs", - "src/openai/images/mod.rs", - "src/openai/live/http.rs", - "src/openai/live/mod.rs", - "src/openai/live/websocket.rs", - "src/openai/mod.rs", - "src/openai/models.rs", - "src/openai/endpoint.rs", - "src/openai/responses/error.rs", - "src/openai/responses/http.rs", - "src/openai/responses/mod.rs", - "src/openai/responses/request.rs", - "src/openai/responses/response.rs", - "src/openai/responses/validation.rs", - "src/openai/responses/websocket/connection/mod.rs", - "src/openai/responses/websocket/connection/middleware.rs", - "src/openai/responses/websocket/forward.rs", - "src/openai/responses/websocket/mod.rs", - "src/openai/responses/websocket/protocol.rs", - "src/openai/router.rs", - "src/openai/search/http.rs", - "src/openai/search/mod.rs", - "src/openai/service.rs", - "src/openai/usage.rs", - "src/provider.rs", - "src/session_cookie.rs", - "src/time.rs", - ]; - expected.sort_unstable(); - - assert_eq!(actual, expected); -} - -#[test] -fn test_tree_should_match_frozen_rust_mirror() { - let root = Path::new(env!("CARGO_MANIFEST_DIR")); - let mut actual = rust_files(&root.join("tests")); - actual.sort(); - let mut expected = vec![ - "tests/admin/account_groups.rs", - "tests/admin/accounts/handlers.rs", - "tests/admin/accounts/import_tasks.rs", - "tests/admin/accounts/mod.rs", - "tests/admin/accounts/presenter.rs", - "tests/admin/auth.rs", - "tests/admin/client_keys.rs", - "tests/admin/errors.rs", - "tests/admin/error_diagnostics.rs", - "tests/admin/mod.rs", - "tests/admin/observability/mod.rs", - "tests/admin/observability/query.rs", - "tests/admin/observability/response.rs", - "tests/admin/plugins/management.rs", - "tests/admin/plugins/mod.rs", - "tests/admin/proxies.rs", - "tests/admin/settings.rs", - "tests/admin/system.rs", - "tests/admin/wire.rs", - "tests/architecture.rs", - "tests/auth.rs", - "tests/support.rs", - "tests/health.rs", - "tests/key_usage/mod.rs", - "tests/key_usage/fixtures.rs", - "tests/main.rs", - "tests/middleware/mod.rs", - "tests/middleware/websocket.rs", - "tests/openai/middleware/mod.rs", - "tests/openai/auth.rs", - "tests/openai/error.rs", - "tests/openai/images/mod.rs", - "tests/openai/mod.rs", - "tests/openai/models.rs", - "tests/openai/endpoint.rs", - "tests/openai/responses/http.rs", - "tests/openai/responses/mod.rs", - "tests/openai/responses/request.rs", - "tests/openai/responses/websocket/connection/mod.rs", - "tests/openai/responses/websocket/connection/middleware.rs", - "tests/openai/responses/websocket/forward.rs", - "tests/openai/responses/websocket/interrupt.rs", - "tests/openai/responses/websocket/mod.rs", - "tests/openai/responses/websocket/protocol.rs", - "tests/openai/live/mod.rs", - "tests/openai/live/websocket.rs", - "tests/openai/router.rs", - "tests/openai/search/mod.rs", - "tests/openai/usage.rs", - "tests/provider.rs", - ]; - expected.sort_unstable(); - - assert_eq!(actual, expected); -} - #[test] fn admin_routes_should_use_only_get_post_and_static_paths() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/admin"); @@ -192,20 +48,6 @@ fn admin_routes_should_use_only_get_post_and_static_paths() { ); } -fn rust_files(root: &Path) -> Vec { - let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); - all_files(root) - .into_iter() - .filter(|path| path.extension().and_then(|value| value.to_str()) == Some("rs")) - .map(|path| { - path.strip_prefix(manifest) - .expect("path below manifest") - .to_string_lossy() - .replace('\\', "/") - }) - .collect() -} - fn all_files(root: &Path) -> Vec { let mut files = Vec::new(); let mut pending = vec![root.to_path_buf()]; diff --git a/backend/crates/gateway-api/tests/middleware/mod.rs b/backend/crates/gateway-api/tests/middleware/mod.rs index 9721385dd..86817603f 100644 --- a/backend/crates/gateway-api/tests/middleware/mod.rs +++ b/backend/crates/gateway-api/tests/middleware/mod.rs @@ -16,15 +16,18 @@ use axum::{ }; use bytes::Bytes; use futures::{future::BoxFuture, stream}; -use gateway_core::{ - engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - lifecycle::CancellationToken, - middleware::http as contract, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::engine::middleware::http as contract; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; use http_body::Frame; use http_body_util::{BodyExt as _, StreamBody}; use tower::ServiceExt as _; @@ -70,12 +73,12 @@ impl MiddlewarePlan for Plan { } fn handle_websocket( &self, - context: gateway_core::middleware::websocket::Context, - message: gateway_core::middleware::websocket::Message, - next: gateway_core::middleware::websocket::Next, + context: gateway_core::engine::middleware::websocket::Context, + message: gateway_core::engine::middleware::websocket::Message, + next: gateway_core::engine::middleware::websocket::Next, ) -> BoxFuture< 'static, - Result, MiddlewareError>, + Result, MiddlewareError>, > { assert!(context.plan.is_some()); self.calls.fetch_add(1, Ordering::SeqCst); @@ -103,11 +106,13 @@ impl MiddlewarePlan for Plan { } fn handle_service( &self, - context: gateway_core::middleware::service::Context, + context: gateway_core::engine::middleware::service::Context, input: serde_json::Value, - next: gateway_core::middleware::service::Next, - ) -> BoxFuture<'static, Result> - { + next: gateway_core::engine::middleware::service::Next, + ) -> BoxFuture< + 'static, + Result, + > { self.calls.fetch_add(1, Ordering::SeqCst); assert_eq!(context.operation, "settings.load"); assert_eq!(context.request_id, "service-request"); @@ -146,7 +151,7 @@ impl MiddlewarePlan for Plan { response.extensions_mut().insert(contract::upgrade::Upgraded::new(status, Box::pin(async move { let session = session.await.map_err(|_| MiddlewareError::Fault)?; if idle { return std::future::pending().await; } - use gateway_core::middleware::websocket::{Kind, Message}; + use gateway_core::engine::middleware::websocket::{Kind, Message}; let receive = session.receive(); tokio::pin!(receive); tokio::select! { diff --git a/backend/crates/gateway-api/tests/openai/auth.rs b/backend/crates/gateway-api/tests/openai/auth.rs index 9cc8ffe81..298e0190d 100644 --- a/backend/crates/gateway-api/tests/openai/auth.rs +++ b/backend/crates/gateway-api/tests/openai/auth.rs @@ -374,9 +374,9 @@ fn chatgpt_remote_suffix_created_by_header_truncation_should_not_skip_the_gate() #[tokio::test] async fn http_settings_freeze_before_plan_resolution_and_apply_before_admission() { - use gateway_core::{ - engine::middleware::*, middleware::http as contract, runtime::extensions::*, - }; + use gateway_core::engine::middleware::http as contract; + use gateway_core::engine::middleware::*; + use gateway_core::routing::extensions::*; struct Lease; impl ExtensionSetLease for Lease { fn is_ready(&self) -> bool { diff --git a/backend/crates/gateway-api/tests/openai/middleware/mod.rs b/backend/crates/gateway-api/tests/openai/middleware/mod.rs index 4505a4c71..a61a3a182 100644 --- a/backend/crates/gateway-api/tests/openai/middleware/mod.rs +++ b/backend/crates/gateway-api/tests/openai/middleware/mod.rs @@ -10,15 +10,22 @@ use std::{ use bytes::Bytes; use futures::future::BoxFuture; -use gateway_core::{ - engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareBody, MiddlewareContext, MiddlewareError, MiddlewareFrame, - MiddlewareFraming, MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - operation::OperationKind, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareBody; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareFrame; +use gateway_core::engine::middleware::MiddlewareFraming; +use gateway_core::engine::middleware::MiddlewareHeader; +use gateway_core::engine::middleware::MiddlewareMount; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::operation::OperationKind; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; /// OpenAI API 入口共用的中间件夹具；授权和 RPC 由 Runtime 行为测试覆盖 #[derive(Debug, Default)] @@ -30,7 +37,8 @@ pub(super) struct RequestMiddleware { pub(super) response_actions: Mutex>, pub(super) response_headers: Vec, pub(super) short_circuit_json: Option, - pub(super) settings: Option, + pub(super) settings: + Option, } #[derive(Debug, Clone)] diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs index e54596e90..60ce032ea 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs @@ -2,14 +2,100 @@ use super::*; use bytes::Bytes; -use gateway_core::{ - engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan, - MiddlewareRequest, MiddlewareResponse, - }, - middleware::{http as http_contract, websocket as core}, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::engine::middleware::FrozenMiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareContext; +use gateway_core::engine::middleware::MiddlewareError; +use gateway_core::engine::middleware::MiddlewareNext; +use gateway_core::engine::middleware::MiddlewarePlan; +use gateway_core::engine::middleware::MiddlewareRequest; +use gateway_core::engine::middleware::MiddlewareResponse; +use gateway_core::engine::middleware::http as http_contract; +use gateway_core::engine::middleware::websocket as core; +use gateway_core::routing::extensions::ExtensionSetId; +use gateway_core::routing::extensions::ExtensionSetLease; +use gateway_core::routing::extensions::ExtensionSetReference; + +#[tokio::test(start_paused = true)] +async fn client_close_does_not_flush_a_retained_sender_with_an_inflight_write() { + #[derive(Default)] + struct RetainedSenderPlan(Mutex>>); + impl fmt::Debug for RetainedSenderPlan { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("RetainedSenderPlan") + } + } + impl MiddlewarePlan for RetainedSenderPlan { + fn has_websocket(&self) -> bool { + true + } + fn handle_websocket( + &self, + context: core::Context, + message: core::Message, + next: core::Next, + ) -> BoxFuture<'static, Result, MiddlewareError>> { + *self.0.lock().unwrap() = Some(context.sender); + next.run(message) + } + fn handle( + &self, + _: MiddlewareContext, + request: MiddlewareRequest, + next: MiddlewareNext, + ) -> BoxFuture<'static, Result> { + next.run(request) + } + } + + let plan = Arc::new(RetainedSenderPlan::default()); + let frozen = FrozenMiddlewarePlan::new( + plan.clone(), + ExtensionSetReference::new( + ExtensionSetId::new("retained-writer".into()).unwrap(), + Arc::new(Lease), + ), + ); + let (incoming, received) = unbounded_channel(); + let (written, mut output) = unbounded_channel(); + let flush_started = Arc::new(AtomicBool::new(false)); + let socket = TestSocket { + incoming: received, + written, + stall_writes: true, + flush_started: flush_started.clone(), + dropped: Arc::new(AtomicBool::new(false)), + }; + let mut connection = spawn_connection( + socket, + Arc::from("retained-writer"), + CancellationToken::new(), + ConnectionConfig::PRODUCTION, + Some(frozen), + Arc::from([]), + ); + incoming.send(Ok(Message::Text("capture".into()))).unwrap(); + assert!(matches!( + connection.next_event().await, + Some(ConnectionEvent::Text(_)) + )); + let sender = plan.0.lock().unwrap().clone().unwrap(); + let mut write = sender.send(core::Message { + kind: core::Kind::Text, + payload: Bytes::from_static(b"pending"), + }); + assert!(futures::poll!(write.as_mut()).is_pending()); + // 写 future 在 pump 的 select 之外持锁，而且暂不继续 poll 以响应取消 + incoming.send(Ok(Message::Close(None))).unwrap(); + let reason = tokio::time::timeout(Duration::from_millis(50), connection.wait_for_exit()) + .await + .expect("peer close must not wait for a retained sender"); + assert_eq!(reason, PumpExitReason::ClientClose); + assert!(!flush_started.load(Ordering::Acquire)); + assert!(write.await.is_err()); + assert!(matches!(output.try_recv(), Err(TryRecvError::Empty))); + plan.0.lock().unwrap().take(); +} + #[derive(Debug, Default)] struct Plan { block_terminal: bool, diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs index 4ce5dea46..4d6bc7c47 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs @@ -234,6 +234,7 @@ struct TestSocket { incoming: UnboundedReceiver>, written: UnboundedSender, stall_writes: bool, + flush_started: Arc, dropped: Arc, } @@ -268,7 +269,12 @@ impl Sink for TestSocket { } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) + self.flush_started.store(true, Ordering::Release); + if self.stall_writes { + Poll::Pending + } else { + Poll::Ready(Ok(())) + } } fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { @@ -286,6 +292,7 @@ struct PumpHarness { connection: ResponsesWebSocketConnection, incoming: UnboundedSender>, written: UnboundedReceiver, + flush_started: Arc, dropped: Arc, cancellation: CancellationToken, } @@ -294,11 +301,13 @@ fn test_connection(stall_writes: bool) -> PumpHarness { let (incoming_tx, incoming_rx) = unbounded_channel(); let (written_tx, written_rx) = unbounded_channel(); let dropped = Arc::new(AtomicBool::new(false)); + let flush_started = Arc::new(AtomicBool::new(false)); let cancellation = CancellationToken::new(); let socket = TestSocket { incoming: incoming_rx, written: written_tx, stall_writes, + flush_started: Arc::clone(&flush_started), dropped: Arc::clone(&dropped), }; let connection = spawn_connection( @@ -313,6 +322,7 @@ fn test_connection(stall_writes: bool) -> PumpHarness { connection, incoming: incoming_tx, written: written_rx, + flush_started, dropped, cancellation, } @@ -404,6 +414,62 @@ async fn client_close_should_take_priority_over_queued_business_frames() { )); } +#[tokio::test(start_paused = true)] +async fn client_close_flush_has_a_short_deadline_without_sending_another_frame() { + let PumpHarness { + mut connection, + incoming, + mut written, + flush_started, + dropped, + .. + } = test_connection(true); + incoming.send(Ok(Message::Close(None))).expect("close peer"); + let started = tokio::time::Instant::now(); + + let reason = tokio::time::timeout(Duration::from_secs(2), connection.wait_for_exit()) + .await + .expect("peer close cannot consume the ordinary write timeout"); + assert_eq!(reason, PumpExitReason::ClientClose); + assert_eq!(started.elapsed(), Duration::from_secs(1)); + assert!(flush_started.load(Ordering::Acquire)); + assert!(dropped.load(Ordering::Acquire)); + assert!( + written.recv().await.is_none(), + "do not send a duplicate Close" + ); +} + +#[tokio::test(start_paused = true)] +async fn lifecycle_shutdown_interrupts_a_pending_client_close_flush() { + let PumpHarness { + mut connection, + incoming, + flush_started, + cancellation, + dropped, + .. + } = test_connection(true); + incoming.send(Ok(Message::Close(None))).expect("close peer"); + for _ in 0..100 { + if flush_started.load(Ordering::Acquire) { + break; + } + tokio::task::yield_now().await; + } + assert!( + flush_started.load(Ordering::Acquire), + "close flush must start" + ); + cancellation.cancel(); + + let reason = tokio::time::timeout(Duration::from_millis(50), connection.wait_for_exit()) + .await + .expect("host cancellation must not wait for the close timeout"); + assert_eq!(reason, PumpExitReason::ClientClose); + assert!(dropped.load(Ordering::Acquire)); +} + #[tokio::test] async fn full_business_queue_should_not_block_lifecycle_shutdown() { let PumpHarness { @@ -988,6 +1054,40 @@ async fn requests_deferred_during_active_response_share_the_inbound_queue_limit( server.abort(); } +#[tokio::test] +async fn client_normal_close_after_completed_response_receives_normal_close_reply() { + let (trace, mut socket, server) = start_active_response().await; + trace.release_terminal.notify_one(); + loop { + let message = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("terminal response timeout") + .expect("response remains open") + .expect("valid response frame"); + if let ClientMessage::Text(text) = message { + let event: Value = serde_json::from_str(&text).expect("response JSON"); + if event["type"] == "response.completed" { + break; + } + } + } + socket + .close(Some(tokio_tungstenite::tungstenite::protocol::CloseFrame { + code: CloseCode::Normal, + reason: "client finished".into(), + })) + .await + .expect("send normal close"); + // close() 只发送本端关闭帧；必须读取对端应答，才能证明握手正常完成 + let reply = tokio::time::timeout(Duration::from_secs(2), socket.next()).await; + server.abort(); + assert!( + matches!(&reply, Ok(Some(Ok(ClientMessage::Close(Some(frame))))) + if frame.code == CloseCode::Normal && frame.reason == "client finished"), + "client must receive the close acknowledgement, got {reply:?}" + ); +} + #[tokio::test] async fn client_close_should_cancel_active_execution_without_starting_queued_requests() { let (trace, mut socket, server) = start_active_response().await; @@ -1009,6 +1109,12 @@ async fn client_close_should_cancel_active_execution_without_starting_queued_req .expect("cancel active execution promptly"); assert_eq!(trace.starts.load(Ordering::Acquire), 1); assert!(trace.cancelled.load(Ordering::Acquire)); + let reply = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("close reply timeout") + .expect("close reply") + .expect("clean close handshake"); + assert!(matches!(reply, ClientMessage::Close(None))); server.abort(); } @@ -1073,6 +1179,7 @@ async fn stalled_heartbeat_can_be_cancelled_without_waiting_for_write_timeout() written: _written, dropped, cancellation, + .. } = test_connection(true); tokio::time::advance(Duration::from_secs(25)).await; tokio::task::yield_now().await; diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs index c0c407169..f6c069a2c 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/forward.rs @@ -19,7 +19,7 @@ use gateway_core::error::{ ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ProviderError, ProviderErrorKind, }; use gateway_core::event::{ProtocolWireEvent, ProviderEvent, ProviderResponseHeader}; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::upstream::{OpaqueUpstreamValue, UpstreamSendState}; use serde_json::{Value, json}; use tokio::net::TcpStream; diff --git a/backend/crates/gateway-core/src/engine/authentication.rs b/backend/crates/gateway-core/src/engine/authentication.rs index d6a546ac2..c2659a30e 100644 --- a/backend/crates/gateway-core/src/engine/authentication.rs +++ b/backend/crates/gateway-core/src/engine/authentication.rs @@ -10,7 +10,7 @@ use futures::future::BoxFuture; use crate::{ policy::{ClientApiKeyId, PlaintextClientApiKey}, - runtime::extensions::{ExtensionSetId, ExtensionSetReference}, + routing::extensions::{ExtensionSetId, ExtensionSetReference}, }; pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024; @@ -18,24 +18,29 @@ pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024; /// API 从数据面 Authorization 头构造的有界认证信封 #[derive(Clone, PartialEq, Eq)] pub struct ClientAuthenticationRequest { - settings: Option, + settings: Option, authorization: Arc, native_bearer: Option, } impl ClientAuthenticationRequest { #[must_use] - pub fn with_settings(mut self, settings: crate::settings::RequestSettings) -> Self { + pub fn with_settings( + mut self, + settings: crate::routing::request_settings::RequestSettings, + ) -> Self { self.settings = Some(settings); self } #[must_use] - pub fn settings(&self) -> Option<&crate::settings::RequestSettings> { + pub fn settings(&self) -> Option<&crate::routing::request_settings::RequestSettings> { self.settings.as_ref() } - pub(crate) fn take_settings(&mut self) -> Option { + pub(crate) fn take_settings( + &mut self, + ) -> Option { self.settings.take() } diff --git a/backend/crates/gateway-core/src/engine/execution/admission.rs b/backend/crates/gateway-core/src/engine/execution/admission.rs new file mode 100644 index 000000000..0acb2075c --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/admission.rs @@ -0,0 +1,71 @@ +//! 请求准入租约及费用结算的释放边界 + +use super::failure::record_resource_failure; +use crate::{ + engine::{ + ModelRequestId, + admission::ClientAdmissionPort, + budget::{ClientBudgetCharge, ClientBudgetPort}, + }, + lifecycle::LeaseGuard, + policy::ClientApiKeyId, +}; +use std::sync::Arc; + +pub(super) struct AdmissionLease { + pub(super) renewal: Option>, + pub(super) armed: bool, + pub(super) port: Arc, + pub(super) diagnostics: Arc, + pub(super) client_api_key_id: ClientApiKeyId, + pub(super) model_request_id: ModelRequestId, +} + +pub(super) async fn settle_budget( + port: &dyn ClientBudgetPort, + charge: ClientBudgetCharge, + diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, +) { + let request_id = charge.request_id.clone(); + if let Err(error) = port.settle(charge).await { + record_resource_failure( + diagnostics, + &request_id, + "settle_client_budget", + "Client budget settlement failed; storage will retry on the next request", + error, + ) + .await; + } +} + +impl AdmissionLease { + pub(super) async fn release(mut self) { + self.renewal.take(); + if let Err(error) = self + .port + .release(&self.client_api_key_id, &self.model_request_id) + .await + { + record_resource_failure( + self.diagnostics.as_ref(), + &self.model_request_id, + "release_client_admission", + "Client admission release failed; lease TTL remains active", + error, + ) + .await; + } + self.armed = false; + } +} + +impl Drop for AdmissionLease { + fn drop(&mut self) { + self.renewal.take(); + if self.armed { + self.port + .abandon(&self.client_api_key_id, &self.model_request_id); + } + } +} diff --git a/backend/crates/gateway-core/src/engine/execution/contract.rs b/backend/crates/gateway-core/src/engine/execution/contract.rs new file mode 100644 index 000000000..432821eaa --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/contract.rs @@ -0,0 +1,480 @@ +//! 执行请求、认证事实与对外会话端口 + +use crate::{ + engine::{ + CoordinatedEvent, EngineError, ModelRequestId, + authentication::ClientAuthenticationRequest, + continuation::PreviousResponseId, + extensions::ExtensionCallScope, + middleware::{ + FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareContext, MiddlewareTarget, + }, + nested::ExecutionEffects, + }, + error::{GatewayError, GatewayErrorKind}, + event::{ProviderEvent, ProviderResponseHeader}, + identity::ProviderKind, + lifecycle::{CancellationToken, Deadline}, + operation::Operation, + policy::{ClientApiKeyId, ClientPolicy}, + routing::{ + ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, RuntimeSnapshot, + UpstreamModelId, request_settings::RequestSettings, + }, +}; +use futures::future::BoxFuture; +use std::{ + fmt, + net::IpAddr, + sync::Arc, + time::{Duration, SystemTime}, +}; +use uuid::Uuid; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientTransport { + HttpJson, + HttpSse, + WebSocket, + InternalProbe, + InternalPlugin, +} + +impl ClientTransport { + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::HttpJson => "http_json", + Self::HttpSse => "http_sse", + Self::WebSocket => "websocket", + Self::InternalProbe => "internal", + Self::InternalPlugin => "internal_plugin", + } + } +} + +/// API 解码后交给 Core 的稳定请求元数据 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecutionRequestMetadata { + pub protocol: String, + pub endpoint: String, + pub transport: ClientTransport, + pub stream: bool, + pub client_ip: Option, + pub user_agent: Option, + pub previous_response_id: Option, +} + +#[derive(Clone)] +pub struct AuthenticatedClient { + pub(super) settings: Option, + pub(super) snapshot: Arc, + pub(super) policy: ClientPolicy, + pub(super) authentication: Option, +} + +impl AuthenticatedClient { + pub(super) fn execution_timeout(&self) -> Option { + self.settings + .as_ref() + .and_then(|settings| settings.execution_timeout(self.policy.key_id())) + } + + #[must_use] + pub fn request_settings(&self) -> Option<&RequestSettings> { + self.settings.as_ref() + } + + /// 后续仍须重新认证；这里只更新下一次认证使用的冻结配置 + #[must_use] + pub fn with_request_settings(mut self, settings: RequestSettings) -> Self { + self.settings = Some(settings); + self + } + + #[must_use] + pub const fn snapshot(&self) -> &Arc { + &self.snapshot + } + + #[must_use] + pub const fn policy(&self) -> &ClientPolicy { + &self.policy + } +} + +impl fmt::Debug for AuthenticatedClient { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("AuthenticatedClient") + .field("key_id", &self.policy.key_id()) + .field("revision", &self.snapshot.revision()) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum ClientAuthenticationError { + #[error("client API key is invalid")] + InvalidKey, + #[error("runtime snapshot is unavailable")] + SnapshotUnavailable, + #[error("frontend authentication provider is unavailable")] + ProviderUnavailable, +} + +pub struct StartExecution { + pub client: AuthenticatedClient, + pub public_model: PublicModelId, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +/// 入口中间件开始前冻结的根请求身份与生命周期 +/// +/// 中间件短路时本值直接释放；调用 `next` 后必须原样交回 +/// [`ExecutionService::start_prepared`]，不能重新认证为另一个 Key +pub struct PreparedRootExecution { + pub(super) extension_scope: ExtensionCallScope, + pub(super) response_control: crate::engine::response_control::ResponseControl, + pub(super) client: AuthenticatedClient, + pub(super) request_id: ModelRequestId, + pub(super) started_at: SystemTime, + pub(super) deadline_at: Deadline, + pub(super) cancellation: CancellationToken, + pub(super) execution_effects: Arc, + pub(super) execution_effects_baseline: usize, +} + +impl PreparedRootExecution { + #[must_use] + pub fn request_settings(&self) -> RequestSettings { + self.client + .settings + .clone() + .unwrap_or_else(|| RequestSettings::new(self.client.snapshot.clone())) + .with_execution( + &self.client.policy, + self.deadline_at + .at() + .map(|at| duration_ms(at.duration_since(self.started_at).unwrap_or_default())), + ) + } + + /// 入口延续发起实例集合；保留已经交给调用方的取消句柄 + #[must_use] + pub fn with_extension_scope(mut self, scope: ExtensionCallScope) -> Self { + self.extension_scope = scope; + self + } + + #[must_use] + pub fn extension_scope(&self) -> ExtensionCallScope { + self.extension_scope.clone() + } + + /// 设置已经在改写时完成编译；此处只将同一份结果应用到冻结身份和生命周期 + pub fn apply_settings(&mut self, settings: &RequestSettings) -> Result<(), GatewayError> { + let deadline = settings + .execution_deadline(self.client.policy.key_id(), self.started_at) + .map_err(|_| { + GatewayError::new( + GatewayErrorKind::InvalidRequest, + "execution settings are invalid", + ) + })?; + self.client.policy = settings.apply_policy(self.client.policy.clone()); + self.client.snapshot = settings.snapshot(); + self.client.settings = Some(settings.clone()); + self.deadline_at = deadline; + Ok(()) + } + + #[must_use] + pub fn response_control(&self) -> crate::engine::response_control::ResponseControl { + self.response_control.clone() + } + + pub(super) fn new(client: AuthenticatedClient) -> Result { + let started_at = SystemTime::now(); + let deadline_at = Deadline::from_timeout(started_at, client.execution_timeout()) + .ok_or_else(|| { + GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") + })?; + let execution_effects = Arc::new(ExecutionEffects::default()); + let execution_effects_baseline = execution_effects.epoch(); + Ok(Self { + extension_scope: ExtensionCallScope::default(), + client, + response_control: crate::engine::response_control::ResponseControl::default(), + request_id: new_request_id()?, + started_at, + deadline_at, + cancellation: CancellationToken::new(), + execution_effects, + execution_effects_baseline, + }) + } + + #[must_use] + pub const fn request_id(&self) -> &ModelRequestId { + &self.request_id + } + + #[must_use] + pub const fn client(&self) -> &AuthenticatedClient { + &self.client + } + + #[must_use] + pub const fn started_at(&self) -> SystemTime { + self.started_at + } + + #[must_use] + pub const fn deadline_at(&self) -> Deadline { + self.deadline_at + } + + #[must_use] + pub fn cancellation(&self) -> CancellationToken { + self.cancellation.clone() + } + + /// 以冻结的入口身份补齐 Runtime-only binding 事实；调用位置事实由 API/Core 提供 + #[must_use] + pub fn middleware_context( + &self, + mut target: MiddlewareTarget, + extension_scope: ExtensionCallScope, + ) -> MiddlewareContext { + target.request_id = self.request_id.clone(); + let account_group_ids = self + .client + .policy + .account_scope() + .routing_snapshot() + .groups_snapshot() + .iter() + .map(|group| group.id().clone()) + .collect::>(); + MiddlewareContext::new( + target, + MiddlewareAuthority { + client_key_id: self.client.policy.key_id().clone(), + account_group_ids: Arc::from(account_group_ids), + cancellation: self.cancellation.clone(), + deadline: self.deadline_at, + extension_scope, + execution_effects: Some(Arc::clone(&self.execution_effects)), + }, + ) + } +} + +/// 入口中间件调用 `next` 后交给 Core 的已解码执行输入 +pub struct PreparedExecutionRequest { + pub public_model: PublicModelId, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +/// 启动一个由协议 adapter 明确绑定到 Provider 自有端点的请求 +pub struct StartProviderExecution { + pub client: AuthenticatedClient, + pub provider: ProviderKind, + /// Provider 自有模型端点可固定上游模型；非模型端点保持 `None` + pub upstream_model: Option, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +pub struct StartedExecution { + pub request_id: ModelRequestId, + pub created_at: SystemTime, + pub stream: bool, + pub session: Box, +} + +pub trait ExecutionSession: Send { + fn trace(&self) -> crate::diagnostics::TraceContext { + crate::diagnostics::TraceContext::default() + } + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>>; + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>>; + fn response_headers(&self) -> &[ProviderResponseHeader]; + fn response_status_code(&self) -> Option { + None + } + fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { + Err(EngineError::InvalidDeliveryState) + } + fn commit_downstream( + &mut self, + client_status_code: Option, + ) -> BoxFuture<'_, Result<(), EngineError>>; + fn record_client_status( + &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, ()>; +} + +pub trait ExecutionService: Send + Sync { + /// 未就绪时管理与健康接口仍可运行，模型认证按既有路径报告不可用 + fn request_settings(&self) -> Option { + None + } + + fn authenticate( + &self, + plaintext: &str, + ) -> Result; + fn authenticate_request( + &self, + request: ClientAuthenticationRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let plaintext = request + .native_bearer() + .ok_or(ClientAuthenticationError::InvalidKey)?; + self.authenticate(plaintext.expose_for_auth()) + }) + } + /// 只验证入口认证信封，不记录 Key 使用事实或执行推理准入 + fn verify_request( + &self, + _request: ClientAuthenticationRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Err(ClientAuthenticationError::InvalidKey) }) + } + fn public_models(&self, client: &AuthenticatedClient) -> Vec; + fn client_model_catalog<'a>( + &'a self, + _client: &'a AuthenticatedClient, + _protocol: &'a str, + _client_version: &'a str, + ) -> BoxFuture<'a, Result, ProviderCatalogUnavailable>> { + Box::pin(async { Err(ProviderCatalogUnavailable) }) + } + fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool; + + /// 重新鉴权并冻结本次根请求身份、请求 ID、deadline 与取消域 + fn prepare_execution( + &self, + client: AuthenticatedClient, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { PreparedRootExecution::new(client) }) + } + + /// 由已经验证管理会话与插件 models 域的宿主入口选择 Key + /// 不向页面提供明文，也不跳过普通请求的扩展、准入和计量 + fn prepare_plugin_execution( + &self, + _client_key_id: &ClientApiKeyId, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + Err(GatewayError::new( + GatewayErrorKind::Unauthorized, + "plugin model entry is unavailable", + )) + }) + } + + /// 为已经通过 `verify_request` 的只读入口建立中间件生命周期，不重复写 Key 使用事实 + fn prepare_verified_execution( + &self, + client: AuthenticatedClient, + ) -> Result { + PreparedRootExecution::new(client) + } + + /// 解析与准备阶段冻结的发布代次一致的中间件计划 + fn middleware_plan(&self, _prepared: &PreparedRootExecution) -> Option { + None + } + + /// 消费一次已准备身份并进入原有 Core 路由、准入、attempt 与结算路径 + fn start_prepared( + &self, + prepared: PreparedRootExecution, + request: PreparedExecutionRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + self.start(StartExecution { + client: prepared.client, + public_model: request.public_model, + operation: request.operation, + metadata: request.metadata, + }) + .await + }) + } + + /// 消费一次已准备身份，并进入 Provider 自有端点的原有准入与结算路径 + fn start_prepared_provider_endpoint( + &self, + prepared: PreparedRootExecution, + provider: ProviderKind, + upstream_model: Option, + operation: Operation, + metadata: ExecutionRequestMetadata, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + self.start_provider_endpoint(StartProviderExecution { + client: prepared.client, + provider, + upstream_model, + operation, + metadata, + }) + .await + }) + } + + fn start( + &self, + request: StartExecution, + ) -> BoxFuture<'_, Result>; + fn start_provider_endpoint( + &self, + request: StartProviderExecution, + ) -> BoxFuture<'_, Result>; + + /// 返回 Provider 注册表中首个声明 Live sideband 能力的网关。 + /// `None` 表示当前组合不含语音 sideband；协议 adapter 据此回退到稳定 501。 + fn live_gateway(&self) -> Option> { + None + } +} + +/// 只验证 Client API Key 并返回稳定 Key ID，不产生请求使用事实 +/// +/// 管理侧登录等非数据面场景必须使用该端口，避免把认证本身误记为一次 Key 使用 +pub trait ClientKeyVerifier: Send + Sync { + fn verify_client_key( + &self, + plaintext: &str, + ) -> Result; +} + +/// 成功认证后的 API Key 使用事实接收器 +/// +/// 认证仍是同步快照读取；实现必须自行异步、去重地持久化，不得阻塞客户端请求 +pub trait ClientApiKeyUsageSink: Send + Sync { + fn record_used(&self, key_id: &ClientApiKeyId); +} + +pub(super) fn duration_ms(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} + +pub(super) fn new_request_id() -> Result { + ModelRequestId::new(format!("req_{}", Uuid::now_v7().simple())) + .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID")) +} diff --git a/backend/crates/gateway-core/src/engine/execution/failure.rs b/backend/crates/gateway-core/src/engine/execution/failure.rs new file mode 100644 index 000000000..37149ec17 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/failure.rs @@ -0,0 +1,100 @@ +//! 执行边界的错误分类与资源释放诊断 + +use super::contract::ClientAuthenticationError; +use crate::{ + engine::{EngineError, ModelRequestId}, + error::{GatewayError, GatewayErrorKind}, +}; + +pub(super) fn authentication_gateway_error(error: ClientAuthenticationError) -> GatewayError { + match error { + ClientAuthenticationError::InvalidKey => { + GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid") + } + ClientAuthenticationError::SnapshotUnavailable => GatewayError::new( + GatewayErrorKind::Internal, + "runtime snapshot is unavailable", + ), + ClientAuthenticationError::ProviderUnavailable => GatewayError::new( + GatewayErrorKind::Internal, + "frontend authentication provider is unavailable", + ), + } +} + +pub(super) fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError { + match error { + crate::validation::RoutingError::ModelNotFound { + model, + mapped_model, + } => GatewayError::new( + GatewayErrorKind::ModelNotFound, + if model == mapped_model { + "the requested model was not found in the provider catalogs available to this API key; check the model name" + } else { + "the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping" + }, + ), + crate::validation::RoutingError::NoCapableProvider { .. } + | crate::validation::RoutingError::NoCapableProviderEndpoint { .. } + | crate::validation::RoutingError::EmptyAccountScope => GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "no provider can execute this request", + ), + crate::validation::RoutingError::UnsupportedProviderEndpoint { .. } => GatewayError::new( + GatewayErrorKind::Unsupported, + "the selected provider does not support this operation", + ), + _ => GatewayError::new( + GatewayErrorKind::Internal, + "runtime routing configuration is invalid", + ), + } +} + +pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { + match error { + EngineError::Cancelled => { + GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled") + } + EngineError::Deadline => { + GatewayError::new(GatewayErrorKind::Timeout, "request deadline elapsed") + } + EngineError::Provider(provider) => GatewayError::from_provider(provider), + EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => { + GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "no provider is available", + ) + } + EngineError::Store(source) => { + GatewayError::new(GatewayErrorKind::Internal, "request execution failed") + .with_source(source.clone()) + } + _ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"), + } +} + +pub(super) async fn record_resource_failure( + diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, + request_id: &ModelRequestId, + operation: &'static str, + message: &'static str, + source: impl Into, +) { + let mut failure = crate::diagnostics::OperationalFailure::new( + "core", + operation, + "resource_unavailable", + message, + ); + failure.correlation_id = Some(request_id.as_str().to_owned()); + failure.details = crate::error::ErrorDetails::capture(Some(&source.into()), None, false); + if diagnostics.record_failure(failure).await.is_err() { + tracing::warn!( + request_id = request_id.as_str(), + operation, + "resource diagnostic could not be recorded" + ); + } +} diff --git a/backend/crates/gateway-core/src/engine/execution/mod.rs b/backend/crates/gateway-core/src/engine/execution/mod.rs new file mode 100644 index 000000000..198340669 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/mod.rs @@ -0,0 +1,24 @@ +//! 数据面执行入口与请求、作用域、会话的单一所有权 + +mod admission; +mod contract; +mod failure; +pub use contract::{ + AuthenticatedClient, ClientApiKeyUsageSink, ClientAuthenticationError, ClientKeyVerifier, + ClientTransport, ExecutionRequestMetadata, ExecutionService, ExecutionSession, + PreparedExecutionRequest, PreparedRootExecution, StartExecution, StartProviderExecution, + StartedExecution, +}; +mod scope; +pub use scope::BoundModelExecutionContext; +mod service; +pub use failure::gateway_error_from_engine; +pub use service::DefaultExecutionService; +mod probe_store; +mod session; + +use std::time::Duration; +const DIAGNOSTIC_TIMEOUT: Duration = Duration::from_secs(10 * 60); +const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100); +const MAX_NESTED_EXECUTIONS: usize = 16; +const MAX_CONCURRENT_NESTED_EXECUTIONS: usize = 4; diff --git a/backend/crates/gateway-core/src/engine/execution/probe_store.rs b/backend/crates/gateway-core/src/engine/execution/probe_store.rs new file mode 100644 index 000000000..fe10c93d1 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/probe_store.rs @@ -0,0 +1,64 @@ +//! 诊断探测使用的非持久执行记录与观测事实 + +use crate::{ + engine::{ + AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, + ModelRequestId, NewModelRequest, ProviderAccountId, RecoveryReport, UpstreamSendState, + }, + error::StoreError, + identity::ProviderKind, + routing::UpstreamModelId, +}; +use std::time::SystemTime; + +pub(super) struct ProbeObservation { + pub(super) provider_kind: ProviderKind, + pub(super) account_id: ProviderAccountId, + pub(super) upstream_model: UpstreamModelId, +} + +pub(super) struct TransientExecutionStore; + +#[async_trait::async_trait] +impl ExecutionStore for TransientExecutionStore { + async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { + 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, _: ModelRequestFinalization) -> Result<(), StoreError> { + Ok(()) + } + + async fn recover_expired(&self, _: SystemTime) -> Result { + Ok(RecoveryReport::default()) + } +} diff --git a/backend/crates/gateway-core/src/engine/execution/scope.rs b/backend/crates/gateway-core/src/engine/execution/scope.rs new file mode 100644 index 000000000..deef09dda --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/scope.rs @@ -0,0 +1,301 @@ +//! 请求嵌套图、活动权限与作用域 guard + +use super::admission::{AdmissionLease, settle_budget}; +use super::contract::{AuthenticatedClient, ExecutionRequestMetadata}; +use super::{MAX_CONCURRENT_NESTED_EXECUTIONS, MAX_NESTED_EXECUTIONS}; +use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort}; +use crate::{ + concurrency::ConcurrencyWaitBudget, + engine::{ + ModelRequestId, ProviderAccountId, coordinator::CoordinationExtensions, + extensions::ExtensionCallScope, nested::ExecutionEffects, + }, + error::{GatewayError, GatewayErrorKind}, + identity::ProviderKind, + lifecycle::{CancellationToken, Deadline}, + operation::Operation, + policy::ClientApiKeyId, + routing::{ + FrozenAccountScope, PublicModelId, UpstreamModelId, request_settings::RequestSettings, + }, +}; +use std::{ + collections::BTreeMap, + fmt, + sync::{ + Arc, Mutex, Weak, + atomic::{AtomicUsize, Ordering}, + }, + time::SystemTime, +}; + +pub(super) enum ExecutionTarget { + Model(PublicModelId), + ProviderEndpoint { + provider: ProviderKind, + upstream_model: Option, + }, +} + +impl ExecutionTarget { + pub(super) fn public_model(&self) -> Option<&PublicModelId> { + match self { + Self::Model(model) => Some(model), + Self::ProviderEndpoint { .. } => None, + } + } + + pub(super) fn into_public_model(self) -> Option { + match self { + Self::Model(model) => Some(model), + Self::ProviderEndpoint { .. } => None, + } + } +} + +pub(super) struct PendingStartExecution { + pub(super) client: AuthenticatedClient, + pub(super) target: ExecutionTarget, + pub(super) operation: Operation, + pub(super) metadata: ExecutionRequestMetadata, +} + +pub(super) struct AuthorizedExecution { + pub(super) response_control: Option, + pub(super) account_scope: Arc, + pub(super) deadline_at: Deadline, + pub(super) cancellation: CancellationToken, + pub(super) extension_scope: ExtensionCallScope, + pub(super) required_provider: Option, + pub(super) required_account: Option, + pub(super) nested: Option, + pub(super) bound: Option, + pub(super) graph: Option>, + pub(super) execution_effects_baseline: usize, + pub(super) nested_permit: Option, +} + +pub(super) struct ExecutionStartGuard { + pub(super) admission: ExecutionAdmission, + pub(super) active_request: Option, + pub(super) concurrency_wait_budget: ConcurrencyWaitBudget, + pub(super) admission_decision_ms: Option, +} + +pub(super) struct PreparedExecutionStart { + pub(super) request_id: ModelRequestId, + pub(super) started_at: SystemTime, + pub(super) plan: crate::routing::RoutingPlan, + pub(super) extensions: CoordinationExtensions, + pub(super) authorization: AuthorizedExecution, + pub(super) guard: ExecutionStartGuard, +} + +pub(super) struct NestedExecutionFacts { + pub(super) parent_request_id: ModelRequestId, + pub(super) initiating_plugin_instance_id: String, +} + +pub(super) struct BoundExecutionFacts { + pub(super) initiating_plugin_instance_id: String, +} + +pub(super) struct NestedExecutionGraph { + pub(super) total_started: AtomicUsize, + pub(super) active: AtomicUsize, + pub(super) effects: Arc, +} + +impl NestedExecutionGraph { + pub(super) fn new(effects: Arc) -> Self { + Self { + total_started: AtomicUsize::new(0), + active: AtomicUsize::new(0), + effects, + } + } + + pub(super) fn acquire(self: &Arc) -> Result { + if self + .active + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |active| { + (active < MAX_CONCURRENT_NESTED_EXECUTIONS).then_some(active + 1) + }) + .is_err() + { + return Err(GatewayError::new( + GatewayErrorKind::ConcurrencyQueueFull, + "nested model execution concurrency is exhausted", + )); + } + if self + .total_started + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |total| { + (total < MAX_NESTED_EXECUTIONS).then_some(total + 1) + }) + .is_err() + { + self.active.fetch_sub(1, Ordering::AcqRel); + return Err(GatewayError::new( + GatewayErrorKind::PolicyDenied, + "nested model execution budget is exhausted", + )); + } + Ok(NestedExecutionPermit { + graph: Arc::clone(self), + armed: true, + }) + } +} + +pub(super) struct NestedExecutionPermit { + pub(super) graph: Arc, + pub(super) armed: bool, +} + +impl NestedExecutionPermit { + pub(super) fn release(mut self) { + if self.armed { + self.graph.active.fetch_sub(1, Ordering::AcqRel); + self.armed = false; + } + } +} + +impl Drop for NestedExecutionPermit { + fn drop(&mut self) { + if self.armed { + self.graph.active.fetch_sub(1, Ordering::AcqRel); + } + } +} + +pub(super) struct ActiveRequestAuthority { + pub(super) client: AuthenticatedClient, + pub(super) account_scope: Arc, + pub(super) deadline_at: Deadline, + pub(super) cancellation: CancellationToken, + pub(super) extension_scope: ExtensionCallScope, + pub(super) graph: Arc, +} + +pub(super) type ActiveRequestRegistry = + Arc>>>; + +pub(super) struct ActiveRequestLease { + pub(super) registry: ActiveRequestRegistry, + pub(super) request_id: ModelRequestId, + pub(super) authority: Arc, +} + +/// 管理页或 CLI 一次调用持有的显式模型执行身份 +/// +/// 结构不暴露 Key 明文；Runtime 只能把它原样交回 Core 发起模型请求 +#[derive(Clone)] +pub struct BoundModelExecutionContext { + pub(super) authority: Arc, +} + +impl BoundModelExecutionContext { + #[must_use] + pub fn request_settings(&self) -> Option<&RequestSettings> { + self.authority.client.request_settings() + } +} + +pub(super) struct BoundModelExecutionAuthority { + pub(super) client: AuthenticatedClient, + pub(super) account_scope: Arc, + pub(super) cancellation: CancellationToken, + pub(super) extension_scope: ExtensionCallScope, + pub(super) initiating_plugin_instance_id: String, +} + +impl fmt::Debug for BoundModelExecutionContext { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("BoundModelExecutionContext") + .field("key_id", &self.authority.client.policy.key_id()) + .field( + "initiating_plugin_instance_id", + &self.authority.initiating_plugin_instance_id, + ) + .finish_non_exhaustive() + } +} + +impl Drop for BoundModelExecutionAuthority { + fn drop(&mut self) { + self.cancellation.cancel(); + } +} + +impl ActiveRequestLease { + pub(super) fn release(self) { + drop(self); + } +} + +impl Drop for ActiveRequestLease { + fn drop(&mut self) { + self.authority.cancellation.cancel(); + let mut active = self + .registry + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if active + .get(&self.request_id) + .and_then(Weak::upgrade) + .is_some_and(|authority| Arc::ptr_eq(&authority, &self.authority)) + { + active.remove(&self.request_id); + } + } +} + +impl ExecutionStartGuard { + pub(super) async fn release_failed( + self, + budget: Option<&dyn ClientBudgetPort>, + request_id: ModelRequestId, + key_id: ClientApiKeyId, + diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, + ) { + if let Some(active_request) = self.active_request { + active_request.release(); + } + if let Some(budget) = budget { + settle_budget( + budget, + ClientBudgetCharge { + key_id, + request_id, + amount_usd: crate::metering::Decimal::ZERO, + completed_at: SystemTime::now(), + }, + diagnostics, + ) + .await; + } + self.admission.release().await; + } +} + +pub(super) enum ExecutionAdmission { + Client(AdmissionLease), + Nested(NestedExecutionPermit), + Bound(AdmissionLease, NestedExecutionPermit), +} + +impl ExecutionAdmission { + pub(super) async fn release(self) { + match self { + Self::Client(admission) => admission.release().await, + Self::Nested(permit) => permit.release(), + Self::Bound(admission, permit) => { + admission.release().await; + permit.release(); + } + } + } +} diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution/service.rs similarity index 64% rename from backend/crates/gateway-core/src/engine/execution.rs rename to backend/crates/gateway-core/src/engine/execution/service.rs index 29975b7f9..8cacc9b36 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution/service.rs @@ -1,739 +1,80 @@ -//! 数据面执行用例：认证、准入、路由、continuation 与会话生命周期 +//! 唯一执行服务的认证、准入、路由与调用编排 -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; -use std::net::IpAddr; -use std::sync::{ - Arc, Mutex, Weak, - atomic::{AtomicUsize, Ordering}, +use super::failure::{authentication_gateway_error, gateway_error_from_engine, map_routing_error}; + +use super::admission::{AdmissionLease, settle_budget}; +use super::contract::{ + AuthenticatedClient, ClientApiKeyUsageSink, ClientAuthenticationError, ClientKeyVerifier, + ClientTransport, ExecutionRequestMetadata, ExecutionService, PreparedExecutionRequest, + PreparedRootExecution, StartExecution, StartProviderExecution, StartedExecution, duration_ms, + new_request_id, }; -use std::time::{Duration, Instant, SystemTime}; - -use futures::{FutureExt, future::BoxFuture, pin_mut, select_biased}; -use futures_timer::Delay; -use uuid::Uuid; - -use crate::concurrency::{CapacityWait, ConcurrencyWaitBudget, ConcurrencyWaitQueue}; -use crate::engine::admission::{ - ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest, +use super::probe_store::{ProbeObservation, TransientExecutionStore}; +use super::scope::{ + ActiveRequestAuthority, ActiveRequestLease, ActiveRequestRegistry, AuthorizedExecution, + BoundExecutionFacts, BoundModelExecutionAuthority, BoundModelExecutionContext, + ExecutionStartGuard, ExecutionTarget, NestedExecutionFacts, NestedExecutionGraph, + PendingStartExecution, PreparedExecutionStart, }; -use crate::engine::authentication::{ - ClientAuthenticationRequest, FrontendAuthenticationDecision, - FrontendAuthenticationExtensionIndex, -}; -use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort}; -use crate::engine::continuation::{ - ContinuationBinding, NativeContinuationPin, NativeContinuationPort, - NativeContinuationStoreErrorKind, PreviousResponseId, -}; -use crate::engine::coordinator::{CoordinationExtensions, ResponseExecutionSession}; -use crate::engine::extensions::ExtensionCallScope; -use crate::engine::middleware::{ - FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareContext, MiddlewareExtensionIndex, - MiddlewareTarget, -}; -use crate::engine::nested::{ - AffinityLookupPort, AffinityLookupRequest, AffinityLookupResult, BoundModelExecutionBinding, - BoundModelExecutionRequest, ExecutionEffects, NestedModelExecutionPort, - NestedModelExecutionRequest, -}; -use crate::engine::observation::{ - FrozenRequestObservationContext, RequestObservationDispatch, RequestObservationScope, - RequestObserverExtensionIndex, -}; -use crate::engine::policy::{ - ModelRouteDecision, RequestPolicyContext, RequestPolicyExtensionIndex, -}; -use crate::engine::probe::{ - AccountProbe, AccountProbeError, AccountProbeErrorSource, AccountProbeRequest, - AccountProbeResult, AccountProbeUpstreamResponse, -}; -use crate::engine::provider::ProviderRegistry; -use crate::engine::{ - AttemptCoordinator, AttemptRecord, CoordinatedEvent, EngineError, ExecutionStore, - GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest, - ProbeFailure, ProviderAccountId, RecoveryReport, UpstreamSendState, -}; -use crate::error::{GatewayError, GatewayErrorKind, StoreError}; -use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader}; -use crate::identity::ProviderKind; -use crate::lifecycle::{CancellationToken, Deadline, LeaseGuard, REQUEST_LEASE_TTL}; -use crate::operation::{Operation, ProviderSessionState}; -use crate::policy::{ClientApiKeyId, ClientPolicy}; -use crate::provider_ports::{ProviderSessionAffinityPort, ProviderStoreErrorKind}; -use crate::routing::{ - FrozenAccountScope, ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, - RoutingContext, RuntimeSnapshot, UpstreamModelId, -}; -use crate::runtime::{RuntimeSnapshotHandle, RuntimeSnapshotPublisher}; -use crate::settings::RequestSettings; - -const DIAGNOSTIC_TIMEOUT: Duration = Duration::from_secs(10 * 60); -const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100); -const MAX_NESTED_EXECUTIONS: usize = 16; -const MAX_CONCURRENT_NESTED_EXECUTIONS: usize = 4; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ClientTransport { - HttpJson, - HttpSse, - WebSocket, - InternalProbe, - InternalPlugin, -} - -impl ClientTransport { - #[must_use] - pub const fn as_str(self) -> &'static str { - match self { - Self::HttpJson => "http_json", - Self::HttpSse => "http_sse", - Self::WebSocket => "websocket", - Self::InternalProbe => "internal", - Self::InternalPlugin => "internal_plugin", - } - } -} - -/// API 解码后交给 Core 的稳定请求元数据 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ExecutionRequestMetadata { - pub protocol: String, - pub endpoint: String, - pub transport: ClientTransport, - pub stream: bool, - pub client_ip: Option, - pub user_agent: Option, - pub previous_response_id: Option, -} - -#[derive(Clone)] -pub struct AuthenticatedClient { - settings: Option, - snapshot: Arc, - policy: ClientPolicy, - authentication: Option, -} - -impl AuthenticatedClient { - fn execution_timeout(&self) -> Option { - self.settings - .as_ref() - .and_then(|settings| settings.execution_timeout(self.policy.key_id())) - } - - #[must_use] - pub fn request_settings(&self) -> Option<&RequestSettings> { - self.settings.as_ref() - } - - /// 后续仍须重新认证；这里只更新下一次认证使用的冻结配置 - #[must_use] - pub fn with_request_settings(mut self, settings: RequestSettings) -> Self { - self.settings = Some(settings); - self - } - - #[must_use] - pub const fn snapshot(&self) -> &Arc { - &self.snapshot - } - - #[must_use] - pub const fn policy(&self) -> &ClientPolicy { - &self.policy - } -} - -impl fmt::Debug for AuthenticatedClient { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("AuthenticatedClient") - .field("key_id", &self.policy.key_id()) - .field("revision", &self.snapshot.revision()) - .finish_non_exhaustive() - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -pub enum ClientAuthenticationError { - #[error("client API key is invalid")] - InvalidKey, - #[error("runtime snapshot is unavailable")] - SnapshotUnavailable, - #[error("frontend authentication provider is unavailable")] - ProviderUnavailable, -} - -pub struct StartExecution { - pub client: AuthenticatedClient, - pub public_model: PublicModelId, - pub operation: Operation, - pub metadata: ExecutionRequestMetadata, -} - -/// 入口中间件开始前冻结的根请求身份与生命周期 -/// -/// 中间件短路时本值直接释放；调用 `next` 后必须原样交回 -/// [`ExecutionService::start_prepared`]，不能重新认证为另一个 Key -pub struct PreparedRootExecution { - extension_scope: ExtensionCallScope, - response_control: super::response_control::ResponseControl, - client: AuthenticatedClient, - request_id: ModelRequestId, - started_at: SystemTime, - deadline_at: Deadline, - cancellation: CancellationToken, - execution_effects: Arc, - execution_effects_baseline: usize, -} - -impl PreparedRootExecution { - #[must_use] - pub fn request_settings(&self) -> RequestSettings { - self.client - .settings - .clone() - .unwrap_or_else(|| RequestSettings::new(self.client.snapshot.clone())) - .with_execution( - &self.client.policy, - self.deadline_at - .at() - .map(|at| duration_ms(at.duration_since(self.started_at).unwrap_or_default())), - ) - } - - /// 入口延续发起实例集合；保留已经交给调用方的取消句柄 - #[must_use] - pub fn with_extension_scope(mut self, scope: ExtensionCallScope) -> Self { - self.extension_scope = scope; - self - } - - #[must_use] - pub fn extension_scope(&self) -> ExtensionCallScope { - self.extension_scope.clone() - } - - /// 设置已经在改写时完成编译；此处只将同一份结果应用到冻结身份和生命周期 - pub fn apply_settings(&mut self, settings: &RequestSettings) -> Result<(), GatewayError> { - let deadline = settings - .execution_deadline(self.client.policy.key_id(), self.started_at) - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::InvalidRequest, - "execution settings are invalid", - ) - })?; - self.client.policy = settings.apply_policy(self.client.policy.clone()); - self.client.snapshot = settings.snapshot(); - self.client.settings = Some(settings.clone()); - self.deadline_at = deadline; - Ok(()) - } - - #[must_use] - pub fn response_control(&self) -> super::response_control::ResponseControl { - self.response_control.clone() - } - - fn new(client: AuthenticatedClient) -> Result { - let started_at = SystemTime::now(); - let deadline_at = Deadline::from_timeout(started_at, client.execution_timeout()) - .ok_or_else(|| { - GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") - })?; - let execution_effects = Arc::new(ExecutionEffects::default()); - let execution_effects_baseline = execution_effects.epoch(); - Ok(Self { - extension_scope: ExtensionCallScope::default(), - client, - response_control: super::response_control::ResponseControl::default(), - request_id: new_request_id()?, - started_at, - deadline_at, - cancellation: CancellationToken::new(), - execution_effects, - execution_effects_baseline, - }) - } - - #[must_use] - pub const fn request_id(&self) -> &ModelRequestId { - &self.request_id - } - - #[must_use] - pub const fn client(&self) -> &AuthenticatedClient { - &self.client - } - - #[must_use] - pub const fn started_at(&self) -> SystemTime { - self.started_at - } - - #[must_use] - pub const fn deadline_at(&self) -> Deadline { - self.deadline_at - } - - #[must_use] - pub fn cancellation(&self) -> CancellationToken { - self.cancellation.clone() - } - - /// 以冻结的入口身份补齐 Runtime-only binding 事实；调用位置事实由 API/Core 提供 - #[must_use] - pub fn middleware_context( - &self, - mut target: MiddlewareTarget, - extension_scope: ExtensionCallScope, - ) -> MiddlewareContext { - target.request_id = self.request_id.clone(); - let account_group_ids = self - .client - .policy - .account_scope() - .routing_snapshot() - .groups_snapshot() - .iter() - .map(|group| group.id().clone()) - .collect::>(); - MiddlewareContext::new( - target, - MiddlewareAuthority { - client_key_id: self.client.policy.key_id().clone(), - account_group_ids: Arc::from(account_group_ids), - cancellation: self.cancellation.clone(), - deadline: self.deadline_at, - extension_scope, - execution_effects: Some(Arc::clone(&self.execution_effects)), - }, - ) - } -} - -/// 入口中间件调用 `next` 后交给 Core 的已解码执行输入 -pub struct PreparedExecutionRequest { - pub public_model: PublicModelId, - pub operation: Operation, - pub metadata: ExecutionRequestMetadata, -} - -/// 启动一个由协议 adapter 明确绑定到 Provider 自有端点的请求 -pub struct StartProviderExecution { - pub client: AuthenticatedClient, - pub provider: ProviderKind, - /// Provider 自有模型端点可固定上游模型；非模型端点保持 `None` - pub upstream_model: Option, - pub operation: Operation, - pub metadata: ExecutionRequestMetadata, -} - -enum ExecutionTarget { - Model(PublicModelId), - ProviderEndpoint { - provider: ProviderKind, - upstream_model: Option, +use super::{COORDINATION_TIMEOUT, DIAGNOSTIC_TIMEOUT}; +use super::{scope::ExecutionAdmission, session::DefaultExecutionSession}; +use crate::{ + concurrency::{CapacityWait, ConcurrencyWaitBudget, ConcurrencyWaitQueue}, + engine::{ + AttemptCoordinator, EngineError, ExecutionStore, GatewayEngine, ModelRequestId, + NewModelRequest, ProbeFailure, ProviderAccountId, UpstreamSendState, + admission::{ + ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, + ClientAdmissionRequest, + }, + authentication::{ + ClientAuthenticationRequest, FrontendAuthenticationDecision, + FrontendAuthenticationExtensionIndex, + }, + budget::{ClientBudgetCharge, ClientBudgetPort}, + continuation::{ + ContinuationBinding, NativeContinuationPin, NativeContinuationPort, + NativeContinuationStoreErrorKind, + }, + coordinator::CoordinationExtensions, + extensions::{ExecutionExtensionIndex, ExtensionCallScope}, + middleware::FrozenMiddlewarePlan, + nested::{ + AffinityLookupPort, AffinityLookupRequest, AffinityLookupResult, + BoundModelExecutionBinding, BoundModelExecutionRequest, ExecutionEffects, + NestedModelExecutionPort, NestedModelExecutionRequest, + }, + observation::{ + FrozenRequestObservationContext, RequestObservationDispatch, RequestObservationScope, + RequestObserverExtensionIndex, + }, + policy::{ModelRouteDecision, RequestPolicyContext, RequestPolicyExtensionIndex}, + probe::{ + AccountProbe, AccountProbeError, AccountProbeErrorSource, AccountProbeRequest, + AccountProbeResult, AccountProbeUpstreamResponse, + }, + provider::ProviderRegistry, }, -} - -impl ExecutionTarget { - fn public_model(&self) -> Option<&PublicModelId> { - match self { - Self::Model(model) => Some(model), - Self::ProviderEndpoint { .. } => None, - } - } - - fn into_public_model(self) -> Option { - match self { - Self::Model(model) => Some(model), - Self::ProviderEndpoint { .. } => None, - } - } -} - -struct PendingStartExecution { - client: AuthenticatedClient, - target: ExecutionTarget, - operation: Operation, - metadata: ExecutionRequestMetadata, -} - -struct AuthorizedExecution { - response_control: Option, - account_scope: Arc, - deadline_at: Deadline, - cancellation: CancellationToken, - extension_scope: ExtensionCallScope, - required_provider: Option, - required_account: Option, - nested: Option, - bound: Option, - graph: Option>, - execution_effects_baseline: usize, - nested_permit: Option, -} - -struct ExecutionStartGuard { - admission: ExecutionAdmission, - active_request: Option, - concurrency_wait_budget: ConcurrencyWaitBudget, - admission_decision_ms: Option, -} - -struct PreparedExecutionStart { - request_id: ModelRequestId, - started_at: SystemTime, - plan: crate::routing::RoutingPlan, - extensions: CoordinationExtensions, - authorization: AuthorizedExecution, - guard: ExecutionStartGuard, -} - -struct NestedExecutionFacts { - parent_request_id: ModelRequestId, - initiating_plugin_instance_id: String, -} - -struct BoundExecutionFacts { - initiating_plugin_instance_id: String, -} - -struct NestedExecutionGraph { - total_started: AtomicUsize, - active: AtomicUsize, - effects: Arc, -} - -impl NestedExecutionGraph { - fn new(effects: Arc) -> Self { - Self { - total_started: AtomicUsize::new(0), - active: AtomicUsize::new(0), - effects, - } - } - - fn acquire(self: &Arc) -> Result { - if self - .active - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |active| { - (active < MAX_CONCURRENT_NESTED_EXECUTIONS).then_some(active + 1) - }) - .is_err() - { - return Err(GatewayError::new( - GatewayErrorKind::ConcurrencyQueueFull, - "nested model execution concurrency is exhausted", - )); - } - if self - .total_started - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |total| { - (total < MAX_NESTED_EXECUTIONS).then_some(total + 1) - }) - .is_err() - { - self.active.fetch_sub(1, Ordering::AcqRel); - return Err(GatewayError::new( - GatewayErrorKind::PolicyDenied, - "nested model execution budget is exhausted", - )); - } - Ok(NestedExecutionPermit { - graph: Arc::clone(self), - armed: true, - }) - } -} - -struct NestedExecutionPermit { - graph: Arc, - armed: bool, -} - -impl NestedExecutionPermit { - fn release(mut self) { - if self.armed { - self.graph.active.fetch_sub(1, Ordering::AcqRel); - self.armed = false; - } - } -} - -impl Drop for NestedExecutionPermit { - fn drop(&mut self) { - if self.armed { - self.graph.active.fetch_sub(1, Ordering::AcqRel); - } - } -} - -struct ActiveRequestAuthority { - client: AuthenticatedClient, - account_scope: Arc, - deadline_at: Deadline, - cancellation: CancellationToken, - extension_scope: ExtensionCallScope, - graph: Arc, -} - -type ActiveRequestRegistry = Arc>>>; - -struct ActiveRequestLease { - registry: ActiveRequestRegistry, - request_id: ModelRequestId, - authority: Arc, -} - -/// 管理页或 CLI 一次调用持有的显式模型执行身份 -/// -/// 结构不暴露 Key 明文；Runtime 只能把它原样交回 Core 发起模型请求 -#[derive(Clone)] -pub struct BoundModelExecutionContext { - authority: Arc, -} - -impl BoundModelExecutionContext { - #[must_use] - pub fn request_settings(&self) -> Option<&RequestSettings> { - self.authority.client.request_settings() - } -} - -struct BoundModelExecutionAuthority { - client: AuthenticatedClient, - account_scope: Arc, - cancellation: CancellationToken, - extension_scope: ExtensionCallScope, - initiating_plugin_instance_id: String, -} - -impl fmt::Debug for BoundModelExecutionContext { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("BoundModelExecutionContext") - .field("key_id", &self.authority.client.policy.key_id()) - .field( - "initiating_plugin_instance_id", - &self.authority.initiating_plugin_instance_id, - ) - .finish_non_exhaustive() - } -} - -impl Drop for BoundModelExecutionAuthority { - fn drop(&mut self) { - self.cancellation.cancel(); - } -} - -impl ActiveRequestLease { - fn release(self) { - drop(self); - } -} - -impl Drop for ActiveRequestLease { - fn drop(&mut self) { - self.authority.cancellation.cancel(); - let mut active = self - .registry - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if active - .get(&self.request_id) - .and_then(Weak::upgrade) - .is_some_and(|authority| Arc::ptr_eq(&authority, &self.authority)) - { - active.remove(&self.request_id); - } - } -} - -pub struct StartedExecution { - pub request_id: ModelRequestId, - pub created_at: SystemTime, - pub stream: bool, - pub session: Box, -} - -pub trait ExecutionSession: Send { - fn trace(&self) -> crate::diagnostics::TraceContext { - crate::diagnostics::TraceContext::default() - } - fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>>; - fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>>; - fn response_headers(&self) -> &[ProviderResponseHeader]; - fn response_status_code(&self) -> Option { - None - } - fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { - Err(EngineError::InvalidDeliveryState) - } - fn commit_downstream( - &mut self, - client_status_code: Option, - ) -> BoxFuture<'_, Result<(), EngineError>>; - fn record_client_status( - &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, ()>; -} - -pub trait ExecutionService: Send + Sync { - /// 未就绪时管理与健康接口仍可运行，模型认证按既有路径报告不可用 - fn request_settings(&self) -> Option { - None - } - - fn authenticate( - &self, - plaintext: &str, - ) -> Result; - fn authenticate_request( - &self, - request: ClientAuthenticationRequest, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - let plaintext = request - .native_bearer() - .ok_or(ClientAuthenticationError::InvalidKey)?; - self.authenticate(plaintext.expose_for_auth()) - }) - } - /// 只验证入口认证信封，不记录 Key 使用事实或执行推理准入 - fn verify_request( - &self, - _request: ClientAuthenticationRequest, - ) -> BoxFuture<'_, Result> { - Box::pin(async { Err(ClientAuthenticationError::InvalidKey) }) - } - fn public_models(&self, client: &AuthenticatedClient) -> Vec; - fn client_model_catalog<'a>( - &'a self, - _client: &'a AuthenticatedClient, - _protocol: &'a str, - _client_version: &'a str, - ) -> BoxFuture<'a, Result, ProviderCatalogUnavailable>> { - Box::pin(async { Err(ProviderCatalogUnavailable) }) - } - fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool; - - /// 重新鉴权并冻结本次根请求身份、请求 ID、deadline 与取消域 - fn prepare_execution( - &self, - client: AuthenticatedClient, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { PreparedRootExecution::new(client) }) - } - - /// 由已经验证管理会话与插件 models 域的宿主入口选择 Key - /// 不向页面提供明文，也不跳过普通请求的扩展、准入和计量 - fn prepare_plugin_execution( - &self, - _client_key_id: &ClientApiKeyId, - ) -> BoxFuture<'_, Result> { - Box::pin(async { - Err(GatewayError::new( - GatewayErrorKind::Unauthorized, - "plugin model entry is unavailable", - )) - }) - } - - /// 为已经通过 `verify_request` 的只读入口建立中间件生命周期，不重复写 Key 使用事实 - fn prepare_verified_execution( - &self, - client: AuthenticatedClient, - ) -> Result { - PreparedRootExecution::new(client) - } - - /// 解析与准备阶段冻结的发布代次一致的中间件计划 - fn middleware_plan(&self, _prepared: &PreparedRootExecution) -> Option { - None - } - - /// 消费一次已准备身份并进入原有 Core 路由、准入、attempt 与结算路径 - fn start_prepared( - &self, - prepared: PreparedRootExecution, - request: PreparedExecutionRequest, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - self.start(StartExecution { - client: prepared.client, - public_model: request.public_model, - operation: request.operation, - metadata: request.metadata, - }) - .await - }) - } - - /// 消费一次已准备身份，并进入 Provider 自有端点的原有准入与结算路径 - fn start_prepared_provider_endpoint( - &self, - prepared: PreparedRootExecution, - provider: ProviderKind, - upstream_model: Option, - operation: Operation, - metadata: ExecutionRequestMetadata, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - self.start_provider_endpoint(StartProviderExecution { - client: prepared.client, - provider, - upstream_model, - operation, - metadata, - }) - .await - }) - } - - fn start( - &self, - request: StartExecution, - ) -> BoxFuture<'_, Result>; - fn start_provider_endpoint( - &self, - request: StartProviderExecution, - ) -> BoxFuture<'_, Result>; - - /// 返回 Provider 注册表中首个声明 Live sideband 能力的网关。 - /// `None` 表示当前组合不含语音 sideband；协议 adapter 据此回退到稳定 501。 - fn live_gateway(&self) -> Option> { - None - } -} - -/// 只验证 Client API Key 并返回稳定 Key ID，不产生请求使用事实 -/// -/// 管理侧登录等非数据面场景必须使用该端口，避免把认证本身误记为一次 Key 使用 -pub trait ClientKeyVerifier: Send + Sync { - fn verify_client_key( - &self, - plaintext: &str, - ) -> Result; -} - -/// 成功认证后的 API Key 使用事实接收器 -/// -/// 认证仍是同步快照读取；实现必须自行异步、去重地持久化，不得阻塞客户端请求 -pub trait ClientApiKeyUsageSink: Send + Sync { - fn record_used(&self, key_id: &ClientApiKeyId); -} + error::{GatewayError, GatewayErrorKind}, + event::GatewayEvent, + identity::ProviderKind, + lifecycle::{CancellationToken, Deadline, REQUEST_LEASE_TTL}, + operation::Operation, + policy::{ClientApiKeyId, ClientPolicy}, + provider_ports::{ProviderSessionAffinityPort, ProviderStoreErrorKind}, + routing::{ + FrozenAccountScope, ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, + RoutingContext, RuntimeSnapshot, UpstreamModelId, request_settings::RequestSettings, + }, + runtime::{RuntimeSnapshotHandle, RuntimeSnapshotPublisher}, +}; +use futures::{FutureExt as _, future::BoxFuture, pin_mut, select_biased}; +use futures_timer::Delay; +use std::{ + collections::BTreeSet, + sync::{Arc, Weak}, + time::{Instant, SystemTime}, +}; pub struct DefaultExecutionService { snapshots: RuntimeSnapshotHandle, @@ -748,7 +89,7 @@ pub struct DefaultExecutionService { budget: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, session_affinity: Option>, snapshot_refresh: Option>, @@ -778,7 +119,7 @@ impl DefaultExecutionService { budget: None, request_observers: None, request_policies: None, - middlewares: None, + execution_extensions: None, frontend_authentication: None, session_affinity: None, snapshot_refresh: None, @@ -805,8 +146,11 @@ impl DefaultExecutionService { } #[must_use] - pub fn with_middlewares(mut self, middlewares: MiddlewareExtensionIndex) -> Self { - self.middlewares = Some(middlewares); + pub fn with_execution_extensions( + mut self, + execution_extensions: ExecutionExtensionIndex, + ) -> Self { + self.execution_extensions = Some(execution_extensions); self } @@ -1069,11 +413,14 @@ impl DefaultExecutionService { .client .snapshot .extensions() - .and_then(|set| set.upstream_adapters()); - let middleware = self.middlewares.as_ref().and_then(|middlewares| { - let generation = request.client.snapshot.extensions()?.clone(); - middlewares.resolve(&generation) - }); + .and_then(|set| self.execution_extensions.as_ref()?.upstream_adapters(set)); + let middleware = self + .execution_extensions + .as_ref() + .and_then(|execution_extensions| { + let generation = request.client.snapshot.extensions()?.clone(); + execution_extensions.middleware(&generation) + }); let request_policy = self.request_policies.as_ref().and_then(|policies| { let generation = request.client.snapshot.extensions()?.clone(); let plan = policies.resolve(&generation)?; @@ -1303,7 +650,7 @@ impl DefaultExecutionService { failure_kind = error.kind().as_str(), "请求在路由或准入阶段被拒绝" ); - let rejection = super::EntryRejection { + let rejection = crate::engine::EntryRejection { request_id: request_id.clone(), client_key_id: budget_key_id.clone(), error: error.clone(), @@ -2140,58 +1487,6 @@ impl DefaultExecutionService { } } -struct ProbeObservation { - provider_kind: ProviderKind, - account_id: ProviderAccountId, - upstream_model: UpstreamModelId, -} - -struct TransientExecutionStore; - -#[async_trait::async_trait] -impl ExecutionStore for TransientExecutionStore { - async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { - 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, _: ModelRequestFinalization) -> Result<(), StoreError> { - Ok(()) - } - - async fn recover_expired(&self, _: SystemTime) -> Result { - Ok(RecoveryReport::default()) - } -} - impl ExecutionService for DefaultExecutionService { fn request_settings(&self) -> Option { self.snapshots.acquire().ok().map(RequestSettings::new) @@ -2328,7 +1623,7 @@ impl ExecutionService for DefaultExecutionService { fn middleware_plan(&self, prepared: &PreparedRootExecution) -> Option { let generation = prepared.client.snapshot.extensions()?.clone(); - self.middlewares.as_ref()?.resolve(&generation) + self.execution_extensions.as_ref()?.middleware(&generation) } fn prepare_plugin_execution( @@ -2484,305 +1779,6 @@ impl AccountProbe for DefaultExecutionService { } } -struct AdmissionLease { - renewal: Option>, - armed: bool, - port: Arc, - diagnostics: Arc, - client_api_key_id: ClientApiKeyId, - model_request_id: ModelRequestId, -} - -async fn settle_budget( - port: &dyn ClientBudgetPort, - charge: ClientBudgetCharge, - diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, -) { - let request_id = charge.request_id.clone(); - if let Err(error) = port.settle(charge).await { - record_resource_failure( - diagnostics, - &request_id, - "settle_client_budget", - "Client budget settlement failed; storage will retry on the next request", - error, - ) - .await; - } -} - -impl AdmissionLease { - async fn release(mut self) { - self.renewal.take(); - if let Err(error) = self - .port - .release(&self.client_api_key_id, &self.model_request_id) - .await - { - record_resource_failure( - self.diagnostics.as_ref(), - &self.model_request_id, - "release_client_admission", - "Client admission release failed; lease TTL remains active", - error, - ) - .await; - } - self.armed = false; - } -} - -impl Drop for AdmissionLease { - fn drop(&mut self) { - self.renewal.take(); - if self.armed { - self.port - .abandon(&self.client_api_key_id, &self.model_request_id); - } - } -} - -enum ExecutionAdmission { - Client(AdmissionLease), - Nested(NestedExecutionPermit), - Bound(AdmissionLease, NestedExecutionPermit), -} - -impl ExecutionAdmission { - async fn release(self) { - match self { - Self::Client(admission) => admission.release().await, - Self::Nested(permit) => permit.release(), - Self::Bound(admission, permit) => { - admission.release().await; - permit.release(); - } - } - } -} - -impl ExecutionStartGuard { - async fn release_failed( - self, - budget: Option<&dyn ClientBudgetPort>, - request_id: ModelRequestId, - key_id: ClientApiKeyId, - diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, - ) { - if let Some(active_request) = self.active_request { - active_request.release(); - } - if let Some(budget) = budget { - settle_budget( - budget, - ClientBudgetCharge { - key_id, - request_id, - amount_usd: crate::metering::Decimal::ZERO, - completed_at: SystemTime::now(), - }, - diagnostics, - ) - .await; - } - self.admission.release().await; - } -} - -struct DefaultExecutionSession { - core: ResponseExecutionSession, - admission: Option, - active_request: Option, - cleanup: Option>, - continuation: Arc, - continuation_recorded: bool, - budget: Option>, - diagnostics: Arc, -} - -impl DefaultExecutionSession { - fn new( - core: ResponseExecutionSession, - admission: ExecutionAdmission, - active_request: Option, - continuation: Arc, - budget: Option>, - diagnostics: Arc, - ) -> Self { - Self { - core, - admission: Some(admission), - active_request, - cleanup: None, - continuation, - continuation_recorded: false, - budget, - diagnostics, - } - } - - async fn settle_if_finalized(&mut self) { - if self.core.is_finalized() - && let Some(admission) = self.admission.take() - { - if let Some(active_request) = self.active_request.take() { - active_request.release(); - } - let budget = self.budget.take(); - let charge = self.core.budget_charge(); - let diagnostics = self.diagnostics.clone(); - // 在首次 await 前把完整清理责任留在会话内 - // 事件等待被取消后，后续 poll - // 或 detach 继续同一个 future，既不丢失费用，也不重启已完成的结算 - self.cleanup = Some(Box::pin(async move { - if let Some(budget) = budget { - settle_budget(budget.as_ref(), charge, diagnostics.as_ref()).await; - } - admission.release().await; - })); - } - if let Some(cleanup) = self.cleanup.as_mut() { - cleanup.await; - self.cleanup = None; - } - } - - async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) { - if self.continuation_recorded { - return; - } - let Some(state) = state else { - return; - }; - let Some(pin) = self.core.native_continuation_pin(state) else { - return; - }; - self.continuation_recorded = true; - record_native_continuation(self.continuation.as_ref(), pin).await; - } - - async fn finalize_detached(&mut self) { - if let Err(error) = self.core.cancel_and_finalize().await { - tracing::warn!(request_id = self.core.request_id().as_str(), operation = "finalize_detached_execution", %error, "Detached execution 终态收敛失败"); - } - self.settle_if_finalized().await; - } -} - -impl Drop for DefaultExecutionSession { - fn drop(&mut self) { - self.core.cancel(); - drop(self.active_request.take()); - } -} - -impl ExecutionSession for DefaultExecutionSession { - fn trace(&self) -> crate::diagnostics::TraceContext { - self.core.trace() - } - fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { - Box::pin(async move { - let result = self.core.next_event().await; - if let Ok(Some(event)) = result.as_ref() { - self.record_continuation(event.session_update()).await; - } - self.settle_if_finalized().await; - result - }) - } - - fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { - Box::pin(async move { - 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; - } - self.settle_if_finalized().await; - result - }) - } - - fn response_headers(&self) -> &[ProviderResponseHeader] { - self.core.response_headers() - } - - fn response_status_code(&self) -> Option { - self.core.response_status_code() - } - - fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { - self.core.discard_pending_delivery() - } - - fn commit_downstream( - &mut self, - client_status_code: Option, - ) -> BoxFuture<'_, Result<(), EngineError>> { - Box::pin(async move { - let result = self.core.commit_downstream(client_status_code).await; - self.settle_if_finalized().await; - result - }) - } - - fn record_client_status( - &mut self, - client_status_code: u16, - ) -> BoxFuture<'_, Result<(), EngineError>> { - Box::pin(async move { - let result = self.core.record_client_status(client_status_code).await; - self.settle_if_finalized().await; - result - }) - } - - fn is_finalized(&self) -> bool { - self.core.is_finalized() - && self.admission.is_none() - && self.active_request.is_none() - && self.cleanup.is_none() - } - - fn cancel(&self) { - self.core.cancel(); - } - - fn detach_finalize(mut self: Box) -> BoxFuture<'static, ()> { - Box::pin(async move { self.finalize_detached().await }) - } -} - -async fn record_native_continuation( - continuation: &dyn NativeContinuationPort, - pin: NativeContinuationPin, -) { - let provider = pin.provider().as_str().to_owned(); - let account = pin.account().as_str().to_owned(); - let record = continuation.record(pin).fuse(); - let timeout = Delay::new(COORDINATION_TIMEOUT).fuse(); - pin_mut!(record, timeout); - select_biased! { - result = record => { - if let Err(error) = result { - tracing::warn!( - provider = %provider, - account = %account, - %error, - "Continuation affinity 写入失败，后续请求将退化为外部续接" - ); - } - }, - _ = timeout => { - tracing::warn!( - provider = %provider, - account = %account, - "Continuation affinity 后台写入超时，已丢弃本次亲和记录" - ); - }, - } -} - fn constant_time_equal(left: &str, right: &str) -> bool { if left.len() != right.len() { return false; @@ -2877,61 +1873,6 @@ fn native_policy( .ok_or(ClientAuthenticationError::InvalidKey) } -fn duration_ms(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} - -fn new_request_id() -> Result { - ModelRequestId::new(format!("req_{}", Uuid::now_v7().simple())) - .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID")) -} - -fn authentication_gateway_error(error: ClientAuthenticationError) -> GatewayError { - match error { - ClientAuthenticationError::InvalidKey => { - GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid") - } - ClientAuthenticationError::SnapshotUnavailable => GatewayError::new( - GatewayErrorKind::Internal, - "runtime snapshot is unavailable", - ), - ClientAuthenticationError::ProviderUnavailable => GatewayError::new( - GatewayErrorKind::Internal, - "frontend authentication provider is unavailable", - ), - } -} - -fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError { - match error { - crate::validation::RoutingError::ModelNotFound { - model, - mapped_model, - } => GatewayError::new( - GatewayErrorKind::ModelNotFound, - if model == mapped_model { - "the requested model was not found in the provider catalogs available to this API key; check the model name" - } else { - "the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping" - }, - ), - crate::validation::RoutingError::NoCapableProvider { .. } - | crate::validation::RoutingError::NoCapableProviderEndpoint { .. } - | crate::validation::RoutingError::EmptyAccountScope => GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "no provider can execute this request", - ), - crate::validation::RoutingError::UnsupportedProviderEndpoint { .. } => GatewayError::new( - GatewayErrorKind::Unsupported, - "the selected provider does not support this operation", - ), - _ => GatewayError::new( - GatewayErrorKind::Internal, - "runtime routing configuration is invalid", - ), - } -} - fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeContinuationPin) { let Some(state) = pin.session_state() else { return; @@ -2943,50 +1884,3 @@ fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeCont } operation.set_provider_session_state(state.clone()); } - -pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { - match error { - EngineError::Cancelled => { - GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled") - } - EngineError::Deadline => { - GatewayError::new(GatewayErrorKind::Timeout, "request deadline elapsed") - } - EngineError::Provider(provider) => GatewayError::from_provider(provider), - EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => { - GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "no provider is available", - ) - } - EngineError::Store(source) => { - GatewayError::new(GatewayErrorKind::Internal, "request execution failed") - .with_source(source.clone()) - } - _ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"), - } -} - -async fn record_resource_failure( - diagnostics: &dyn crate::diagnostics::OperationalDiagnostics, - request_id: &ModelRequestId, - operation: &'static str, - message: &'static str, - source: impl Into, -) { - let mut failure = crate::diagnostics::OperationalFailure::new( - "core", - operation, - "resource_unavailable", - message, - ); - failure.correlation_id = Some(request_id.as_str().to_owned()); - failure.details = crate::error::ErrorDetails::capture(Some(&source.into()), None, false); - if diagnostics.record_failure(failure).await.is_err() { - tracing::warn!( - request_id = request_id.as_str(), - operation, - "resource diagnostic could not be recorded" - ); - } -} diff --git a/backend/crates/gateway-core/src/engine/execution/session.rs b/backend/crates/gateway-core/src/engine/execution/session.rs new file mode 100644 index 000000000..cd5570413 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution/session.rs @@ -0,0 +1,213 @@ +//! 执行会话的准入释放、费用结算与终态收敛 + +use super::COORDINATION_TIMEOUT; +use super::admission::settle_budget; +use super::contract::ExecutionSession; +use super::scope::{ActiveRequestLease, ExecutionAdmission}; +use crate::{ + engine::{ + CoordinatedEvent, EngineError, ExecutionStore, + budget::ClientBudgetPort, + continuation::{NativeContinuationPin, NativeContinuationPort}, + coordinator::ResponseExecutionSession, + }, + event::{ProviderEvent, ProviderResponseHeader}, + operation::ProviderSessionState, +}; +use futures::{FutureExt as _, future::BoxFuture, pin_mut, select_biased}; +use futures_timer::Delay; +use std::sync::Arc; + +pub(super) struct DefaultExecutionSession { + pub(super) core: ResponseExecutionSession, + pub(super) admission: Option, + pub(super) active_request: Option, + pub(super) cleanup: Option>, + pub(super) continuation: Arc, + pub(super) continuation_recorded: bool, + pub(super) budget: Option>, + pub(super) diagnostics: Arc, +} + +impl DefaultExecutionSession { + pub(super) fn new( + core: ResponseExecutionSession, + admission: ExecutionAdmission, + active_request: Option, + continuation: Arc, + budget: Option>, + diagnostics: Arc, + ) -> Self { + Self { + core, + admission: Some(admission), + active_request, + cleanup: None, + continuation, + continuation_recorded: false, + budget, + diagnostics, + } + } + + pub(super) async fn settle_if_finalized(&mut self) { + if self.core.is_finalized() + && let Some(admission) = self.admission.take() + { + if let Some(active_request) = self.active_request.take() { + active_request.release(); + } + let budget = self.budget.take(); + let charge = self.core.budget_charge(); + let diagnostics = self.diagnostics.clone(); + // 在首次 await 前把完整清理责任留在会话内 + // 事件等待被取消后，后续 poll + // 或 detach 继续同一个 future，既不丢失费用，也不重启已完成的结算 + self.cleanup = Some(Box::pin(async move { + if let Some(budget) = budget { + settle_budget(budget.as_ref(), charge, diagnostics.as_ref()).await; + } + admission.release().await; + })); + } + if let Some(cleanup) = self.cleanup.as_mut() { + cleanup.await; + self.cleanup = None; + } + } + + pub(super) async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) { + if self.continuation_recorded { + return; + } + let Some(state) = state else { + return; + }; + let Some(pin) = self.core.native_continuation_pin(state) else { + return; + }; + self.continuation_recorded = true; + record_native_continuation(self.continuation.as_ref(), pin).await; + } + + pub(super) async fn finalize_detached(&mut self) { + if let Err(error) = self.core.cancel_and_finalize().await { + tracing::warn!(request_id = self.core.request_id().as_str(), operation = "finalize_detached_execution", %error, "Detached execution 终态收敛失败"); + } + self.settle_if_finalized().await; + } +} + +impl Drop for DefaultExecutionSession { + fn drop(&mut self) { + self.core.cancel(); + drop(self.active_request.take()); + } +} + +impl ExecutionSession for DefaultExecutionSession { + fn trace(&self) -> crate::diagnostics::TraceContext { + self.core.trace() + } + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + let result = self.core.next_event().await; + if let Ok(Some(event)) = result.as_ref() { + self.record_continuation(event.session_update()).await; + } + self.settle_if_finalized().await; + result + }) + } + + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + 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; + } + self.settle_if_finalized().await; + result + }) + } + + fn response_headers(&self) -> &[ProviderResponseHeader] { + self.core.response_headers() + } + + fn response_status_code(&self) -> Option { + self.core.response_status_code() + } + + fn discard_pending_delivery(&mut self) -> Result<(), EngineError> { + self.core.discard_pending_delivery() + } + + fn commit_downstream( + &mut self, + client_status_code: Option, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + let result = self.core.commit_downstream(client_status_code).await; + self.settle_if_finalized().await; + result + }) + } + + fn record_client_status( + &mut self, + client_status_code: u16, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + let result = self.core.record_client_status(client_status_code).await; + self.settle_if_finalized().await; + result + }) + } + + fn is_finalized(&self) -> bool { + self.core.is_finalized() + && self.admission.is_none() + && self.active_request.is_none() + && self.cleanup.is_none() + } + + fn cancel(&self) { + self.core.cancel(); + } + + fn detach_finalize(mut self: Box) -> BoxFuture<'static, ()> { + Box::pin(async move { self.finalize_detached().await }) + } +} + +pub(super) async fn record_native_continuation( + continuation: &dyn NativeContinuationPort, + pin: NativeContinuationPin, +) { + let provider = pin.provider().as_str().to_owned(); + let account = pin.account().as_str().to_owned(); + let record = continuation.record(pin).fuse(); + let timeout = Delay::new(COORDINATION_TIMEOUT).fuse(); + pin_mut!(record, timeout); + select_biased! { + result = record => { + if let Err(error) = result { + tracing::warn!( + provider = %provider, + account = %account, + %error, + "Continuation affinity 写入失败，后续请求将退化为外部续接" + ); + } + }, + _ = timeout => { + tracing::warn!( + provider = %provider, + account = %account, + "Continuation affinity 后台写入超时，已丢弃本次亲和记录" + ); + }, + } +} diff --git a/backend/crates/gateway-core/src/engine/extensions.rs b/backend/crates/gateway-core/src/engine/extensions.rs index 1d058c3d3..e0726de4a 100644 --- a/backend/crates/gateway-core/src/engine/extensions.rs +++ b/backend/crates/gateway-core/src/engine/extensions.rs @@ -1,4 +1,4 @@ -//! 插件嵌套调用的递归作用域 +//! 执行扩展计划的弱引用索引、冻结解析与插件递归作用域 use std::{collections::BTreeSet, sync::Arc}; @@ -43,3 +43,82 @@ impl ExtensionCallScope { }) } } + +/// 同一代次的具体执行计划；不拥有发布状态，运行资源由对应发布集合保活 +#[derive(Debug, Default)] +pub struct ExecutionExtensionPlans { + pub middleware: Option>, + pub upstream_adapters: Option>, +} + +/// 同一发布代次不能被静默替换 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("execution extension generation is already registered")] +pub struct ExtensionRegistrationError; + +/// 仅以弱引用按冻结身份解析执行计划，不延长旧代次寿命或决定何时发布 +#[derive(Clone, Default)] +pub struct ExecutionExtensionIndex { + sets: Arc< + std::sync::RwLock< + std::collections::BTreeMap< + crate::routing::extensions::ExtensionSetId, + std::sync::Weak, + >, + >, + >, +} + +impl ExecutionExtensionIndex { + pub fn register( + &self, + id: crate::routing::extensions::ExtensionSetId, + plans: Arc, + ) -> Result, ExtensionRegistrationError> { + let mut sets = self + .sets + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + sets.retain(|_, plans| plans.strong_count() > 0); + if sets.contains_key(&id) { + return Err(ExtensionRegistrationError); + } + sets.insert(id, Arc::downgrade(&plans)); + Ok(plans) + } + + fn resolve( + &self, + generation: &crate::routing::extensions::ExtensionSetReference, + ) -> Option> { + self.sets + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(generation.id()) + .and_then(std::sync::Weak::upgrade) + } + + #[must_use] + pub fn middleware( + &self, + generation: &crate::routing::extensions::ExtensionSetReference, + ) -> Option { + let plan = self.resolve(generation)?.middleware.clone()?; + Some(super::middleware::FrozenMiddlewarePlan::new( + plan, + generation.clone(), + )) + } + + #[must_use] + pub fn upstream_adapters( + &self, + generation: &crate::routing::extensions::ExtensionSetReference, + ) -> Option { + let plan = self.resolve(generation)?.upstream_adapters.clone()?; + Some(super::upstream_adapter::FrozenUpstreamAdapterPlan::new( + plan, + generation.clone(), + )) + } +} diff --git a/backend/crates/gateway-core/src/engine/middleware.rs b/backend/crates/gateway-core/src/engine/middleware/contract.rs similarity index 80% rename from backend/crates/gateway-core/src/engine/middleware.rs rename to backend/crates/gateway-core/src/engine/middleware/contract.rs index 5cf19e0f8..386c338af 100644 --- a/backend/crates/gateway-core/src/engine/middleware.rs +++ b/backend/crates/gateway-core/src/engine/middleware/contract.rs @@ -5,10 +5,11 @@ //! Runtime/SDK 负责；身份、路由、账号租约、发送事实、重试、计量与结算仍由 Core //! 持有，不能通过本端口交给插件解释 -use std::collections::{BTreeMap, BTreeSet}; +use super::FrozenMiddlewarePlan; +use std::collections::BTreeSet; use std::fmt; use std::num::NonZeroU32; -use std::sync::{Arc, RwLock, Weak}; +use std::sync::Arc; use bytes::Bytes; use futures::future::BoxFuture; @@ -27,8 +28,7 @@ use crate::operation::{ CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload, }; use crate::policy::ClientApiKeyId; -use crate::runtime::extensions::{ExtensionSetId, ExtensionSetReference}; -use crate::settings::RequestSettings; +use crate::routing::request_settings::RequestSettings; /// 中间件在一次逻辑请求中的挂载位置 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] @@ -543,7 +543,7 @@ impl fmt::Debug for MiddlewareResponseEnvelope { /// Core 冻结的调用事实；完整业务身份可由插件读取 #[derive(Clone)] pub struct MiddlewareContext { - plan: Option, + pub(super) plan: Option, request_id: ModelRequestId, mount: MiddlewareMount, attempt_index: Option, @@ -799,175 +799,3 @@ impl From for MiddlewareError { /// 绑定当前调用且只能消费一次的宿主续体 pub type MiddlewareNext = crate::middleware::Next; - -/// 一个发布代次冻结的统一数据面中间件计划 -pub trait MiddlewarePlan: Send + Sync + fmt::Debug { - fn has_service(&self) -> bool { - false - } - fn handle_service( - &self, - _context: crate::middleware::service::Context, - input: crate::middleware::service::Value, - next: crate::middleware::service::Next, - ) -> BoxFuture< - 'static, - Result, - > { - next.run(input) - } - - fn has_websocket(&self) -> bool { - false - } - fn handle_websocket( - &self, - _context: crate::middleware::websocket::Context, - message: crate::middleware::websocket::Message, - next: crate::middleware::websocket::Next, - ) -> BoxFuture<'static, Result, MiddlewareError>> - { - next.run(message) - } - - /// HTTP 入口在认证、路由与正文解析之前组合，无匹配时保留直接路径 - fn has_http(&self) -> bool { - false - } - - fn handle_http( - &self, - _context: crate::middleware::http::Context, - request: crate::middleware::http::Request, - next: crate::middleware::http::Next, - ) -> BoxFuture<'static, Result> { - next.run(request) - } - - fn handle( - &self, - context: MiddlewareContext, - request: MiddlewareRequest, - next: MiddlewareNext, - ) -> BoxFuture<'static, Result>; -} - -/// 调用期间同时保活发布代次与中间件计划 -#[derive(Clone)] -pub struct FrozenMiddlewarePlan { - plan: Arc, - _generation: ExtensionSetReference, -} - -impl FrozenMiddlewarePlan { - #[must_use] - pub fn has_service(&self) -> bool { - self.plan.has_service() - } - pub fn handle_service( - &self, - context: crate::middleware::service::Context, - input: crate::middleware::service::Value, - next: crate::middleware::service::Next, - ) -> BoxFuture< - 'static, - Result, - > { - self.plan.handle_service(context, input, next) - } - - #[must_use] - pub fn has_websocket(&self) -> bool { - self.plan.has_websocket() - } - pub fn handle_websocket( - &self, - context: crate::middleware::websocket::Context, - message: crate::middleware::websocket::Message, - next: crate::middleware::websocket::Next, - ) -> BoxFuture<'static, Result, MiddlewareError>> - { - self.plan.handle_websocket(context, message, next) - } - - #[must_use] - pub fn has_http(&self) -> bool { - self.plan.has_http() - } - - pub fn handle_http( - &self, - mut context: crate::middleware::http::Context, - request: crate::middleware::http::Request, - next: crate::middleware::http::Next, - ) -> BoxFuture<'static, Result> { - context.plan = Some(self.clone()); - self.plan.handle_http(context, request, next) - } - - #[must_use] - pub fn new(plan: Arc, generation: ExtensionSetReference) -> Self { - Self { - plan, - _generation: generation, - } - } - - pub fn handle( - &self, - mut context: MiddlewareContext, - request: MiddlewareRequest, - next: MiddlewareNext, - ) -> BoxFuture<'static, Result> { - context.plan = Some(self.clone()); - self.plan.handle(context, request, next) - } -} - -impl fmt::Debug for FrozenMiddlewarePlan { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("FrozenMiddlewarePlan") - .finish_non_exhaustive() - } -} - -/// 同一发布代次不能被静默替换 -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -#[error("middleware generation is already registered")] -pub struct MiddlewareRegistrationError; - -/// 按请求冻结的发布集合解析中间件计划；索引不延长旧代次寿命 -#[derive(Clone, Default)] -pub struct MiddlewareExtensionIndex { - sets: Arc>>>, -} - -impl MiddlewareExtensionIndex { - pub fn register( - &self, - id: ExtensionSetId, - plan: Arc, - ) -> Result, MiddlewareRegistrationError> { - let mut sets = self - .sets - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - sets.retain(|_, plan| plan.strong_count() > 0); - if sets.contains_key(&id) { - return Err(MiddlewareRegistrationError); - } - sets.insert(id, Arc::downgrade(&plan)); - Ok(plan) - } - - #[must_use] - pub fn resolve(&self, generation: &ExtensionSetReference) -> Option { - self.sets - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .get(generation.id()) - .and_then(Weak::upgrade) - .map(|plan| FrozenMiddlewarePlan::new(plan, generation.clone())) - } -} diff --git a/backend/crates/gateway-core/src/middleware/http/mod.rs b/backend/crates/gateway-core/src/engine/middleware/http/mod.rs similarity index 86% rename from backend/crates/gateway-core/src/middleware/http/mod.rs rename to backend/crates/gateway-core/src/engine/middleware/http/mod.rs index 7743556a9..096c6c0e2 100644 --- a/backend/crates/gateway-core/src/middleware/http/mod.rs +++ b/backend/crates/gateway-core/src/engine/middleware/http/mod.rs @@ -13,7 +13,7 @@ use crate::{engine::middleware::MiddlewareError, lifecycle::CancellationToken}; pub type Body = UnsyncBoxBody>; pub type Request = http::Request; pub type Response = http::Response; -pub type Next = super::Next; +pub type Next = crate::middleware::Next; /// 路由前调用没有 Client Key；身份在默认认证流程中产生 #[derive(Clone, Debug)] @@ -29,7 +29,7 @@ pub struct Context { /// 内部子请求复用 API 总路由；端口不解释业务路径，也不经网络回环 pub trait Dispatcher: Send + Sync { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { None } @@ -44,7 +44,7 @@ pub trait Dispatcher: Send + Sync { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct Settings { pub timeout: Option, - pub runtime: Option, + pub runtime: Option, } #[must_use] diff --git a/backend/crates/gateway-core/src/middleware/http/upgrade.rs b/backend/crates/gateway-core/src/engine/middleware/http/upgrade.rs similarity index 93% rename from backend/crates/gateway-core/src/middleware/http/upgrade.rs rename to backend/crates/gateway-core/src/engine/middleware/http/upgrade.rs index e29829170..25b062181 100644 --- a/backend/crates/gateway-core/src/middleware/http/upgrade.rs +++ b/backend/crates/gateway-core/src/engine/middleware/http/upgrade.rs @@ -4,7 +4,7 @@ use std::sync::{Arc, Mutex}; use futures::future::{BoxFuture, Shared}; -use crate::{engine::middleware::MiddlewareError, middleware::websocket::Session}; +use crate::{engine::middleware::MiddlewareError, engine::middleware::websocket::Session}; pub type PendingSession = Shared, Arc>>>; diff --git a/backend/crates/gateway-core/src/engine/middleware/mod.rs b/backend/crates/gateway-core/src/engine/middleware/mod.rs new file mode 100644 index 000000000..c193773fa --- /dev/null +++ b/backend/crates/gateway-core/src/engine/middleware/mod.rs @@ -0,0 +1,15 @@ +//! 执行层的请求、HTTP、WebSocket 与公开服务中间件合同 + +mod contract; +pub mod http; +mod plan; +pub mod service; +pub mod websocket; + +pub use contract::{ + MiddlewareAuthority, MiddlewareBody, MiddlewareCapabilityDeclaration, MiddlewareContext, + MiddlewareError, MiddlewareFrame, MiddlewareFrameEnvelope, MiddlewareFraming, MiddlewareHeader, + MiddlewareMount, MiddlewareNext, MiddlewareRequest, MiddlewareResponse, + MiddlewareResponseEnvelope, MiddlewareTarget, +}; +pub use plan::{FrozenMiddlewarePlan, MiddlewarePlan}; diff --git a/backend/crates/gateway-core/src/engine/middleware/plan.rs b/backend/crates/gateway-core/src/engine/middleware/plan.rs new file mode 100644 index 000000000..d96ab8f86 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/middleware/plan.rs @@ -0,0 +1,152 @@ +//! 与发布代次一起冻结的协议及请求中间件计划 + +use super::{ + MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewareRequest, MiddlewareResponse, +}; +use crate::routing::extensions::ExtensionSetReference; +use futures::future::BoxFuture; +use std::{fmt, sync::Arc}; + +/// 一个发布代次冻结的统一数据面中间件计划 +pub trait MiddlewarePlan: Send + Sync + fmt::Debug { + fn has_service(&self) -> bool { + false + } + fn handle_service( + &self, + _context: crate::engine::middleware::service::Context, + input: crate::engine::middleware::service::Value, + next: crate::engine::middleware::service::Next, + ) -> BoxFuture< + 'static, + Result< + crate::engine::middleware::service::Value, + crate::engine::middleware::service::Error, + >, + > { + next.run(input) + } + + fn has_websocket(&self) -> bool { + false + } + fn handle_websocket( + &self, + _context: crate::engine::middleware::websocket::Context, + message: crate::engine::middleware::websocket::Message, + next: crate::engine::middleware::websocket::Next, + ) -> BoxFuture< + 'static, + Result, MiddlewareError>, + > { + next.run(message) + } + + /// HTTP 入口在认证、路由与正文解析之前组合，无匹配时保留直接路径 + fn has_http(&self) -> bool { + false + } + + fn handle_http( + &self, + _context: crate::engine::middleware::http::Context, + request: crate::engine::middleware::http::Request, + next: crate::engine::middleware::http::Next, + ) -> BoxFuture<'static, Result> + { + next.run(request) + } + + fn handle( + &self, + context: MiddlewareContext, + request: MiddlewareRequest, + next: MiddlewareNext, + ) -> BoxFuture<'static, Result>; +} + +/// 调用期间同时保活发布代次与中间件计划 +#[derive(Clone)] +pub struct FrozenMiddlewarePlan { + plan: Arc, + _generation: ExtensionSetReference, +} + +impl FrozenMiddlewarePlan { + #[must_use] + pub fn has_service(&self) -> bool { + self.plan.has_service() + } + pub fn handle_service( + &self, + context: crate::engine::middleware::service::Context, + input: crate::engine::middleware::service::Value, + next: crate::engine::middleware::service::Next, + ) -> BoxFuture< + 'static, + Result< + crate::engine::middleware::service::Value, + crate::engine::middleware::service::Error, + >, + > { + self.plan.handle_service(context, input, next) + } + + #[must_use] + pub fn has_websocket(&self) -> bool { + self.plan.has_websocket() + } + pub fn handle_websocket( + &self, + context: crate::engine::middleware::websocket::Context, + message: crate::engine::middleware::websocket::Message, + next: crate::engine::middleware::websocket::Next, + ) -> BoxFuture< + 'static, + Result, MiddlewareError>, + > { + self.plan.handle_websocket(context, message, next) + } + + #[must_use] + pub fn has_http(&self) -> bool { + self.plan.has_http() + } + + pub fn handle_http( + &self, + mut context: crate::engine::middleware::http::Context, + request: crate::engine::middleware::http::Request, + next: crate::engine::middleware::http::Next, + ) -> BoxFuture<'static, Result> + { + context.plan = Some(self.clone()); + self.plan.handle_http(context, request, next) + } + + #[must_use] + pub fn new(plan: Arc, generation: ExtensionSetReference) -> Self { + Self { + plan, + _generation: generation, + } + } + + pub fn handle( + &self, + mut context: MiddlewareContext, + request: MiddlewareRequest, + next: MiddlewareNext, + ) -> BoxFuture<'static, Result> { + context.plan = Some(self.clone()); + self.plan.handle(context, request, next) + } +} + +impl fmt::Debug for FrozenMiddlewarePlan { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("FrozenMiddlewarePlan") + .finish_non_exhaustive() + } +} diff --git a/backend/crates/gateway-core/src/middleware/service.rs b/backend/crates/gateway-core/src/engine/middleware/service.rs similarity index 95% rename from backend/crates/gateway-core/src/middleware/service.rs rename to backend/crates/gateway-core/src/engine/middleware/service.rs index 8e4d55500..eeffba953 100644 --- a/backend/crates/gateway-core/src/middleware/service.rs +++ b/backend/crates/gateway-core/src/engine/middleware/service.rs @@ -7,7 +7,7 @@ use crate::{ /// 值只在进入插件链时编码；操作的领域校验与副作用由终端所属服务执行 pub type Value = serde_json::Value; -pub type Next = super::Next; +pub type Next = crate::middleware::Next; #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, thiserror::Error)] #[error("{message}")] diff --git a/backend/crates/gateway-core/src/middleware/websocket.rs b/backend/crates/gateway-core/src/engine/middleware/websocket.rs similarity index 93% rename from backend/crates/gateway-core/src/middleware/websocket.rs rename to backend/crates/gateway-core/src/engine/middleware/websocket.rs index 16fc1cb38..2caba924c 100644 --- a/backend/crates/gateway-core/src/middleware/websocket.rs +++ b/backend/crates/gateway-core/src/engine/middleware/websocket.rs @@ -44,4 +44,4 @@ pub struct Context { pub cancellation: CancellationToken, pub sender: Arc, } -pub type Next = super::Next, MiddlewareError>; +pub type Next = crate::middleware::Next, MiddlewareError>; diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 6fd4b841b..34893182b 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -101,17 +101,7 @@ pub enum AttemptTrigger { AccountRetry, } -/// Provider 可为当前 attempt 选择的请求局部传输档位 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] -pub enum AttemptTransport { - /// 使用 Provider 的默认传输策略 - #[default] - Default, - /// 固定账号重试 Provider 的首选传输；序号由 Provider 的独立预算驱动 - Retry(NonZeroU32), - /// 使用 Provider 定义的备用传输 - Fallback, -} +pub use crate::upstream::AttemptTransport; impl AttemptTrigger { #[must_use] diff --git a/backend/crates/gateway-core/src/engine/nested.rs b/backend/crates/gateway-core/src/engine/nested.rs index 3ef1f35fb..82dede6a3 100644 --- a/backend/crates/gateway-core/src/engine/nested.rs +++ b/backend/crates/gateway-core/src/engine/nested.rs @@ -65,7 +65,7 @@ pub struct NestedModelExecutionRequest { /// Core 在绑定时冻结当前 /// Key 策略；每个实际模型请求仍独立经过准入、预算、账本与计费 pub struct BoundModelExecutionBinding { - pub settings: Option, + pub settings: Option, pub client_key_id: ClientApiKeyId, pub initiating_plugin_instance_id: String, /// 身份随父调用或连接回收；每个模型请求使用独立的执行期限 diff --git a/backend/crates/gateway-core/src/engine/observation.rs b/backend/crates/gateway-core/src/engine/observation.rs index 3dddea682..700b81e4d 100644 --- a/backend/crates/gateway-core/src/engine/observation.rs +++ b/backend/crates/gateway-core/src/engine/observation.rs @@ -22,8 +22,8 @@ use crate::{ metering::{CostEstimate, CostSource, Usage}, operation::OperationKind, policy::ClientApiKeyId, + routing::extensions::{ExtensionSetId, ExtensionSetReference}, routing::{AccountGroupId, ConfigRevision, PublicModelId}, - runtime::extensions::{ExtensionSetId, ExtensionSetReference}, upstream::UpstreamSendState, }; diff --git a/backend/crates/gateway-core/src/engine/policy.rs b/backend/crates/gateway-core/src/engine/policy.rs index 73d9935bd..6ee86cef8 100644 --- a/backend/crates/gateway-core/src/engine/policy.rs +++ b/backend/crates/gateway-core/src/engine/policy.rs @@ -19,8 +19,8 @@ use crate::{ identity::ProviderKind, operation::Operation, policy::ClientApiKeyId, + routing::extensions::{ExtensionSetId, ExtensionSetReference}, routing::{AccountGroupId, PublicModelId}, - runtime::extensions::{ExtensionSetId, ExtensionSetReference}, }; /// 模型路由插件的一次输入；正文仍保存在 `operation` 中，由 Runtime 按权限投影 diff --git a/backend/crates/gateway-core/src/engine/upstream_adapter.rs b/backend/crates/gateway-core/src/engine/upstream_adapter.rs index 1d1041aba..836dc1f02 100644 --- a/backend/crates/gateway-core/src/engine/upstream_adapter.rs +++ b/backend/crates/gateway-core/src/engine/upstream_adapter.rs @@ -15,7 +15,7 @@ use crate::{ metering::{CalculatedCost, Usage}, operation::Operation, routing::UpstreamModelId, - runtime::extensions::ExtensionSetReference, + routing::extensions::ExtensionSetReference, }; use super::{ diff --git a/backend/crates/gateway-core/src/error/mod.rs b/backend/crates/gateway-core/src/error/mod.rs index cc97f3ffa..2455a26e7 100644 --- a/backend/crates/gateway-core/src/error/mod.rs +++ b/backend/crates/gateway-core/src/error/mod.rs @@ -144,7 +144,7 @@ pub enum PreDeliveryRetry { /// 可靠 NotSent 的建连恢复；次数与时间由 Core 的请求级预算裁决 SameAccountConnectionRetry { /// Provider 明确指定实际失败的传输，避免 HTTP 恢复重新进入 WS - transport: crate::engine::AttemptTransport, + transport: crate::upstream::AttemptTransport, }, /// 固定本次账号，并要求 Provider 使用备用传输 SameAccountTransportFallback, @@ -666,7 +666,7 @@ impl ProviderError { } #[must_use] - pub fn with_connection_retry(mut self, transport: crate::engine::AttemptTransport) -> Self { + pub fn with_connection_retry(mut self, transport: crate::upstream::AttemptTransport) -> Self { self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountConnectionRetry { transport, })); diff --git a/backend/crates/gateway-core/src/lib.rs b/backend/crates/gateway-core/src/lib.rs index ff130bf22..7b206691f 100644 --- a/backend/crates/gateway-core/src/lib.rs +++ b/backend/crates/gateway-core/src/lib.rs @@ -128,7 +128,7 @@ pub struct CoreStartup { providers: ProviderRegistry, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, control: CoreControlPlaneStartup, } @@ -164,7 +164,7 @@ impl CoreStartup { self.providers, self.request_observers, self.request_policies, - self.middlewares, + self.execution_extensions, self.frontend_authentication, )); let execution: Arc = service.clone(); @@ -194,7 +194,7 @@ fn build_execution_service( providers: ProviderRegistry, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, ) -> DefaultExecutionService { let mut service = DefaultExecutionService::new( @@ -215,8 +215,8 @@ fn build_execution_service( if let Some(request_policies) = request_policies { service = service.with_request_policies(request_policies); } - if let Some(middlewares) = middlewares { - service = service.with_middlewares(middlewares); + if let Some(execution_extensions) = execution_extensions { + service = service.with_execution_extensions(execution_extensions); } if let Some(frontend_authentication) = frontend_authentication { service = service.with_frontend_authentication(frontend_authentication); @@ -276,10 +276,10 @@ impl CoreBundle { pub async fn initialize( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, ) -> Result { prepare( @@ -288,7 +288,7 @@ pub async fn initialize( extensions, request_observers, request_policies, - middlewares, + execution_extensions, frontend_authentication, ) .activate() @@ -300,10 +300,10 @@ pub async fn initialize( pub fn prepare( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, frontend_authentication: Option, ) -> CoreStartup { let control = prepare_control_plane(ports.clone(), providers.clone(), extensions); @@ -312,7 +312,7 @@ pub fn prepare( providers, request_observers, request_policies, - middlewares, + execution_extensions, frontend_authentication, control, } @@ -377,7 +377,7 @@ pub struct CoreCommandPlaneStartup { providers: ProviderRegistry, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, control: CoreControlPlaneStartup, } @@ -407,7 +407,7 @@ impl CoreCommandPlaneStartup { self.providers, self.request_observers, self.request_policies, - self.middlewares, + self.execution_extensions, None, ) .with_snapshot_refresh(Arc::clone(&control.publisher)), @@ -449,10 +449,10 @@ impl CoreCommandPlaneBundle { pub fn prepare_command_plane( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, request_observers: Option, request_policies: Option, - middlewares: Option, + execution_extensions: Option, ) -> CoreCommandPlaneStartup { let control = prepare_control_plane(ports.clone(), providers.clone(), extensions); CoreCommandPlaneStartup { @@ -460,7 +460,7 @@ pub fn prepare_command_plane( providers, request_observers, request_policies, - middlewares, + execution_extensions, control, } } @@ -470,7 +470,7 @@ pub fn prepare_command_plane( pub fn prepare_control_plane( ports: CoreStorePorts, providers: ProviderRegistry, - extensions: Option>, + extensions: Option>, ) -> CoreControlPlaneStartup { let mut compiler = RuntimeSnapshotCompiler::new(ports.snapshots, Arc::new(providers)); if let Some(extensions) = extensions { diff --git a/backend/crates/gateway-core/src/middleware/mod.rs b/backend/crates/gateway-core/src/middleware.rs similarity index 97% rename from backend/crates/gateway-core/src/middleware/mod.rs rename to backend/crates/gateway-core/src/middleware.rs index 7e8b3b818..c851004bd 100644 --- a/backend/crates/gateway-core/src/middleware/mod.rs +++ b/backend/crates/gateway-core/src/middleware.rs @@ -1,9 +1,5 @@ //! 与协议无关的洋葱组合器；输入、结果与错误保留调用方的领域类型 -pub mod http; -pub mod service; -pub mod websocket; - use futures::future::BoxFuture; /// 单次调用中的一层；上下文放在类型化输入中或由闭包捕获 diff --git a/backend/crates/gateway-core/src/provider_ports.rs b/backend/crates/gateway-core/src/provider_ports.rs deleted file mode 100644 index b9df9ffcb..000000000 --- a/backend/crates/gateway-core/src/provider_ports.rs +++ /dev/null @@ -1,1645 +0,0 @@ -//! Provider 运行时所需的中立存储能力 - -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; -use std::num::NonZeroU32; -use std::sync::Arc; -use std::time::{Duration, SystemTime}; - -use futures::future::BoxFuture; - -use crate::account::{ - AccountConcurrency, AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, - CredentialState, OpaqueProviderData, ProviderAccountId, ProviderAccountStore, -}; -use crate::identity::ProviderKind; -use crate::policy::ClientApiKeyId; -use crate::routing::{ConfigRevision, UpstreamModelId}; -use crate::validation::{IdentifierError, validate_text}; - -const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60); - -/// Provider 可据此决定是否重试，但看不到 SQL、Redis 或秘密原文 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ProviderStoreErrorKind { - Unavailable, - InvalidData, - Conflict, -} - -/// Provider 存储端口的脱敏错误 -#[derive(Debug, Clone, thiserror::Error)] -#[error("provider store {operation} failed: {kind:?}")] -pub struct ProviderStoreError { - kind: ProviderStoreErrorKind, - operation: &'static str, - source: Option, -} - -impl ProviderStoreError { - #[must_use] - pub const fn new(kind: ProviderStoreErrorKind, operation: &'static str) -> Self { - Self { - kind, - operation, - source: None, - } - } - - /// 包装基础设施失败，分类与原始来源分别保留 - #[must_use] - pub fn caused_by( - kind: ProviderStoreErrorKind, - operation: &'static str, - source: impl Into, - ) -> Self { - Self { - kind, - operation, - source: Some(source.into()), - } - } - - #[must_use] - pub const fn kind(&self) -> ProviderStoreErrorKind { - self.kind - } -} - -/// 一个 Provider 的完整可重建调度状态 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSchedulingState { - signals: BTreeMap, - round_robin_cursor: u64, -} - -impl ProviderSchedulingState { - #[must_use] - pub const fn new( - signals: BTreeMap, - round_robin_cursor: u64, - ) -> Self { - Self { - signals, - round_robin_cursor, - } - } - - #[must_use] - pub const fn signals(&self) -> &BTreeMap { - &self.signals - } - - #[must_use] - pub const fn round_robin_cursor(&self) -> u64 { - self.round_robin_cursor - } -} - -/// 请求级账号 lease 的全部中立事实 -#[derive(Debug, Clone)] -pub struct ProviderSchedulingLeaseRequest { - provider_kind: ProviderKind, - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - max_concurrent: AccountConcurrency, - request_interval: Duration, - deadline: crate::lifecycle::Deadline, - cancellation: crate::lifecycle::CancellationToken, -} - -impl ProviderSchedulingLeaseRequest { - #[must_use] - pub fn new( - provider_kind: ProviderKind, - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - max_concurrent: impl Into, - request_interval: Duration, - deadline: impl Into, - ) -> Self { - Self { - provider_kind, - account_id, - credential_revision, - max_concurrent: max_concurrent.into(), - request_interval, - deadline: deadline.into(), - cancellation: crate::lifecycle::CancellationToken::new(), - } - } - - #[must_use] - pub fn with_cancellation(mut self, cancellation: crate::lifecycle::CancellationToken) -> Self { - self.cancellation = cancellation; - self - } - - #[must_use] - pub fn cancellation(&self) -> crate::lifecycle::CancellationToken { - self.cancellation.clone() - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn max_concurrent(&self) -> AccountConcurrency { - self.max_concurrent - } - - #[must_use] - pub const fn request_interval(&self) -> Duration { - self.request_interval - } - - #[must_use] - pub const fn deadline(&self) -> crate::lifecycle::Deadline { - self.deadline - } -} - -/// Lease 生命周期由具体 Store guard 管理，Provider 只能持有 -pub trait ProviderLeaseGuard: Send + Sync + 'static {} - -impl ProviderLeaseGuard for T where T: Send + Sync + 'static {} - -pub enum ProviderLeaseAcquisition { - Acquired(Box), - Busy { retry_after: Option }, -} - -impl fmt::Debug for ProviderLeaseAcquisition { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Acquired(_) => formatter.write_str("Acquired([LEASE])"), - Self::Busy { retry_after } => formatter - .debug_struct("Busy") - .field("retry_after", retry_after) - .finish(), - } - } -} - -/// Provider 运行时会持有的三类 lease；刷新必须同时持有全局容量与账号互斥 lease -#[derive(Debug, Clone)] -pub enum ProviderLeaseRequest { - Scheduling(ProviderSchedulingLeaseRequest), - RefreshCapacity(ProviderRefreshCapacityRequest), - Refresh(ProviderRefreshLeaseRequest), -} - -pub trait ProviderLeasePort: Send + Sync { - fn load_state<'a>( - &'a self, - client_api_key_id: &'a ClientApiKeyId, - provider_kind: &'a ProviderKind, - accounts: &'a [ProviderAccountId], - ) -> BoxFuture<'a, Result>; - - fn try_acquire( - &self, - request: ProviderLeaseRequest, - ) -> BoxFuture<'_, Result>; - - /// 读取指定账号当前的在途请求数；只用于容量熔断的峰值证据， - /// 支持租约信号的存储实现覆盖，否则视为不可观测（空映射） - fn account_in_flight<'a>( - &'a self, - account_ids: &'a [ProviderAccountId], - ) -> BoxFuture<'a, Result, ProviderStoreError>> { - let _ = account_ids; - Box::pin(async move { Ok(BTreeMap::new()) }) - } -} - -/// Provider 从原始会话锚点派生的不可逆亲和键 -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ProviderSessionAffinityKey(String); - -impl ProviderSessionAffinityKey { - pub fn try_new(value: impl Into) -> Result { - let value = value.into(); - if value.is_empty() - || value.len() > 128 - || !value.bytes().all(|byte| { - byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'-' | b'_') - }) - { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate provider session affinity key", - )); - } - Ok(Self(value)) - } - - #[must_use] - pub fn expose_to_store(&self) -> &str { - &self.0 - } -} - -impl fmt::Debug for ProviderSessionAffinityKey { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("ProviderSessionAffinityKey([OPAQUE])") - } -} - -/// 会话绑定快照；版本区分同一账号的不同认领，防止过期和 A → B → A 后的旧写入 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSessionBinding { - account_id: ProviderAccountId, - revision: String, -} - -impl ProviderSessionBinding { - pub fn new( - account_id: ProviderAccountId, - revision: String, - ) -> Result { - if revision.len() != 32 || !revision.bytes().all(|byte| byte.is_ascii_hexdigit()) { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "decode provider session binding revision", - )); - } - Ok(Self { - account_id, - revision, - }) - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub fn revision(&self) -> &str { - &self.revision - } -} - -/// 已观测请求关联及其账号迁移权限，由 Provider 解释协议后写入 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSessionAlias { - pub session_key: ProviderSessionAffinityKey, - pub root_session_key: Option, - pub follow_only: bool, -} - -/// 客户端作用域内的会话账号绑定，原始会话身份由 Provider 哈希后传入 -pub trait ProviderSessionAffinityPort: Send + Sync { - fn load<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - /// 在发送前原子认领、续期或迁移；None 表示快照已变化，调用方必须释放租约重新选择 - /// expected 为 None 只允许首次认领，已有快照必须连同版本匹配 - fn compare_and_bind<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - expected: Option<&'a ProviderSessionBinding>, - account_id: &'a ProviderAccountId, - ttl: Duration, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - /// 显式观测到的请求关联只指向会话键，不缓存账号，迁移后仍读取当前绑定 - fn load_alias<'a>( - &'a self, - provider: &'a ProviderKind, - alias: &'a ProviderSessionAffinityKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - /// 只允许首次关联或同目标续期，冲突时禁止把同一轮次改指其他会话 - fn bind_alias<'a>( - &'a self, - provider: &'a ProviderKind, - alias: &'a ProviderSessionAffinityKey, - session: &'a ProviderSessionAlias, - ttl: Duration, - ) -> BoxFuture<'a, Result>; -} - -/// Provider 会话内已失败账号的可丢失排除集 -/// -/// Provider 自行派生会话键并决定何时写入或清理；Core 只承载调度所需的账号 ID -/// 与 compare-and-swap revision，不解释任一 Provider 协议字段 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderSessionExclusions { - excluded_accounts: BTreeSet, - revision: String, -} - -impl ProviderSessionExclusions { - #[must_use] - pub const fn new(excluded_accounts: BTreeSet, revision: String) -> Self { - Self { - excluded_accounts, - revision, - } - } - - #[must_use] - pub const fn excluded_accounts(&self) -> &BTreeSet { - &self.excluded_accounts - } - - #[must_use] - pub fn revision(&self) -> &str { - &self.revision - } -} - -/// 可丢失的 Provider 会话级账号排除状态 -/// -/// 该端口不接收协议正文；Provider 只能以不可逆会话键、账号 ID 和固定 TTL 操作 -pub trait ProviderSessionExclusionPort: Send + Sync { - fn load<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn record_failure<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - account_id: &'a ProviderAccountId, - ttl: Duration, - ) -> BoxFuture<'a, Result>; - - fn clear<'a>( - &'a self, - provider_kind: &'a ProviderKind, - key: &'a ProviderSessionAffinityKey, - expected_revision: &'a str, - ) -> BoxFuture<'a, Result>; -} - -/// 所有 Provider 共享的 OAuth refresh 并发容量 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ProviderRefreshCapacityRequest { - max_concurrent: NonZeroU32, -} - -impl ProviderRefreshCapacityRequest { - #[must_use] - pub const fn new(max_concurrent: NonZeroU32) -> Self { - Self { max_concurrent } - } - - #[must_use] - pub const fn max_concurrent(self) -> NonZeroU32 { - self.max_concurrent - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderRefreshLeaseRequest { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, -} - -impl ProviderRefreshLeaseRequest { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - ) -> Self { - Self { - account_id, - credential_revision, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } -} - -/// Provider 定义的 catalog cache 作用域 -/// -/// Core 不解释其值；例如 Provider 可以使用套餐、区域或产品线作为共享目录边界 -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ProviderCatalogScope(String); - -impl ProviderCatalogScope { - /// 创建一个稳定且可用于 Redis 隔离键的 Provider-owned 作用域 - /// - /// # Errors - /// - /// 空值、过长文本或控制字符会被拒绝 - pub fn new(value: impl Into) -> Result { - let value = value.into(); - validate_text(&value, 128, false, None)?; - Ok(Self(value)) - } - - /// 返回 Provider-owned 作用域文本 - #[must_use] - pub fn as_str(&self) -> &str { - &self.0 - } -} - -/// Opaque catalog cache 的 Provider 与 Provider-owned 作用域 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderCatalogCacheKey { - provider_kind: ProviderKind, - scope: ProviderCatalogScope, -} - -impl ProviderCatalogCacheKey { - #[must_use] - pub const fn new(provider_kind: ProviderKind, scope: ProviderCatalogScope) -> Self { - Self { - provider_kind, - scope, - } - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub const fn scope(&self) -> &ProviderCatalogScope { - &self.scope - } -} - -pub trait ProviderCatalogCachePort: Send + Sync { - fn replace<'a>( - &'a self, - key: &'a ProviderCatalogCacheKey, - catalog: &'a OpaqueProviderData, - ttl: Duration, - ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; - - fn read<'a>( - &'a self, - key: &'a ProviderCatalogCacheKey, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; -} - -/// Provider 从官方制品核验出的可重建请求画像 -/// -/// Core 只用单调制品序号约束覆盖顺序；具体版本字段由对应 Provider 放在 -/// `profile` 中解释 -/// 每个 Provider 的各制品分别保留一份最新画像 -#[derive(Clone, PartialEq)] -pub struct ProviderArtifactProfile { - provider_kind: ProviderKind, - artifact_key: String, - artifact_sequence: u64, - verified_at: SystemTime, - profile: OpaqueProviderData, -} - -impl ProviderArtifactProfile { - #[must_use] - pub const fn new( - provider_kind: ProviderKind, - artifact_key: String, - artifact_sequence: u64, - verified_at: SystemTime, - profile: OpaqueProviderData, - ) -> Self { - Self { - provider_kind, - artifact_key, - artifact_sequence, - verified_at, - profile, - } - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub fn artifact_key(&self) -> &str { - &self.artifact_key - } - - #[must_use] - pub const fn artifact_sequence(&self) -> u64 { - self.artifact_sequence - } - - #[must_use] - pub const fn verified_at(&self) -> SystemTime { - self.verified_at - } - - #[must_use] - pub const fn profile(&self) -> &OpaqueProviderData { - &self.profile - } -} - -impl fmt::Debug for ProviderArtifactProfile { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("ProviderArtifactProfile") - .field("provider_kind", &self.provider_kind) - .field("artifact_sequence", &self.artifact_sequence) - .field("verified_at", &self.verified_at) - .field("profile", &"[PROVIDER-OWNED]") - .finish() - } -} - -pub trait ProviderArtifactProfileCachePort: Send + Sync { - /// 覆盖同一 Provider、同一制品的固定 cache key - /// - /// 返回 `false` 表示 Store 已持有更高的制品序号；相同序号但内容不同必须返回 - /// [`ProviderStoreErrorKind::Conflict`]，不能静默改写已核验画像 - fn replace_if_newer( - &self, - profile: ProviderArtifactProfile, - ttl: Duration, - ) -> BoxFuture<'_, Result>; - - fn read<'a>( - &'a self, - provider_kind: &'a ProviderKind, - artifact_key: &'a str, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; -} - -/// Redis 中可重建的账号状态投影 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderCredentialState { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - enabled: bool, - credential_state: CredentialState, - observed_at: SystemTime, -} - -impl ProviderCredentialState { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - enabled: bool, - credential_state: CredentialState, - observed_at: SystemTime, - ) -> Self { - Self { - account_id, - credential_revision, - enabled, - credential_state, - observed_at, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn enabled(&self) -> bool { - self.enabled - } - - #[must_use] - pub const fn credential_state(&self) -> CredentialState { - self.credential_state - } - - #[must_use] - pub const fn observed_at(&self) -> SystemTime { - self.observed_at - } -} - -pub trait ProviderCredentialStatePort: Send + Sync { - fn replace( - &self, - state: ProviderCredentialState, - ) -> BoxFuture<'_, Result<(), ProviderStoreError>>; - - fn read<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn clear<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result>; - - /// 记录一次瞬态刷新失败并返回窗口内累计失败次数；每次失败刷新 TTL - fn record_refresh_backoff<'a>( - &'a self, - account_id: &'a ProviderAccountId, - window: Duration, - ) -> BoxFuture<'a, Result>; - - /// 凭据完整成功轮换后清零失败计数，退避窗口重新从 base 起步 - fn clear_refresh_backoff<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; -} - -/// 账号级 cooldown 的来源类别；调度状态统一按 `rate_limited` 处理 -pub use crate::account::AccountCooldownKind as ProviderCooldownKind; - -/// 可丢失的账号冷却事实，不进入持久状态；探测冻结到期后仍需确认恢复 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderCooldown { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - until: SystemTime, - kind: ProviderCooldownKind, -} - -impl ProviderCooldown { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - until: SystemTime, - ) -> Self { - Self::new_with_kind( - account_id, - credential_revision, - until, - ProviderCooldownKind::RateLimit, - ) - } - - #[must_use] - pub const fn new_with_kind( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - until: SystemTime, - kind: ProviderCooldownKind, - ) -> Self { - Self { - account_id, - credential_revision, - until, - kind, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn until(&self) -> SystemTime { - self.until - } - - #[must_use] - pub const fn scheduling_state(&self) -> crate::account::AccountCooldown { - crate::account::AccountCooldown { - until: self.until, - kind: self.kind, - } - } - - #[must_use] - pub const fn kind(&self) -> ProviderCooldownKind { - self.kind - } -} - -/// 可丢失 cooldown 的细粒度作用域 -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum ProviderCooldownScope { - /// 只阻止同一账号调用指定上游模型 - UpstreamModel(UpstreamModelId), -} - -impl ProviderCooldownScope { - #[must_use] - pub const fn upstream_model(model: UpstreamModelId) -> Self { - Self::UpstreamModel(model) - } - - #[must_use] - pub const fn kind(&self) -> &'static str { - match self { - Self::UpstreamModel(_) => "model", - } - } - - #[must_use] - pub fn value(&self) -> &str { - match self { - Self::UpstreamModel(model) => model.as_str(), - } - } -} - -/// 不进入账号持久状态的账号+作用域 cooldown -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderScopedCooldown { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - scope: ProviderCooldownScope, - until: SystemTime, -} - -impl ProviderScopedCooldown { - #[must_use] - pub const fn new( - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - scope: ProviderCooldownScope, - until: SystemTime, - ) -> Self { - Self { - account_id, - credential_revision, - scope, - until, - } - } - - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn scope(&self) -> &ProviderCooldownScope { - &self.scope - } - - #[must_use] - pub const fn until(&self) -> SystemTime { - self.until - } -} - -pub trait ProviderCooldownPort: Send + Sync { - fn put_if_later( - &self, - cooldown: ProviderCooldown, - ) -> BoxFuture<'_, Result>; - - fn read<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn clear<'a>( - &'a self, - account_id: &'a ProviderAccountId, - through_revision: CredentialRevision, - ) -> BoxFuture<'a, Result>; - - fn put_scoped_if_later( - &self, - cooldown: ProviderScopedCooldown, - ) -> BoxFuture<'_, Result>; - - fn read_scoped<'a>( - &'a self, - account_id: &'a ProviderAccountId, - scope: &'a ProviderCooldownScope, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; - - fn clear_scoped<'a>( - &'a self, - account_id: &'a ProviderAccountId, - scope: &'a ProviderCooldownScope, - through_revision: CredentialRevision, - ) -> BoxFuture<'a, Result>; - - /// 删除账号时清除该账号全部 account/model scope cooldown key - /// 生命周期与 credential revision 无关；不要求调用方持有 revision - fn clear_all<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result>; - - /// 记录一次容量类失败并返回滑动窗口内的累计次数，同时把观测到的账号 - /// 在途并发并入窗口峰值（`in_flight` 为 0 表示本次未观测，跳过峰值更新） - /// 只服务容量熔断触发器；调用频率受失败频率约束，不需要批量接口 - fn record_capacity_failure<'a>( - &'a self, - account_id: &'a ProviderAccountId, - window: Duration, - in_flight: u32, - ) -> BoxFuture<'a, Result>; - - /// 普通请求成功后原子清除临时限流及失败证据；必须保留任何容量冻结 - fn clear_after_success<'a>( - &'a self, - account_id: &'a ProviderAccountId, - through_revision: CredentialRevision, - ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; - - /// 读取窗口内观测到的在途并发峰值；无证据时返回 `None` - fn capacity_peak_in_flight<'a>( - &'a self, - account_id: &'a ProviderAccountId, - ) -> BoxFuture<'a, Result, ProviderStoreError>>; -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ProviderRefreshPolicy { - margin: Duration, - concurrency: NonZeroU32, -} - -impl ProviderRefreshPolicy { - pub fn try_new(margin: Duration, concurrency: NonZeroU32) -> Result { - if margin.is_zero() { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate refresh policy", - )); - } - Ok(Self { - margin, - concurrency, - }) - } - - #[must_use] - pub const fn margin(self) -> Duration { - self.margin - } - - #[must_use] - pub const fn concurrency(self) -> NonZeroU32 { - self.concurrency - } - - /// 判断 AT 是否已经进入当前配置的提前刷新窗口 - /// - /// `margin` 是运行时策略，不投影为账号的持久化时间字段；已过期 AT 也需要 - /// 尝试用 RT 恢复，因此视为到期 - #[must_use] - pub fn is_refresh_due( - self, - access_token_expires_at: SystemTime, - observed_at: SystemTime, - ) -> bool { - match access_token_expires_at.duration_since(observed_at) { - Ok(remaining) => remaining <= self.margin, - Err(_) => true, - } - } - - /// 按账号派生 `[0, margin]` 内的整秒错峰偏移。 - /// - /// 网关以固定周期扫描到期账号，到期时刻相同的账号会在同一轮齐刷， - /// 在 auth.openai.com 侧形成单 IP 批量刷新特征。偏移从账号 ID 的 - /// 稳定哈希派生（与退避扰动同源、独立 salt），跨扫描与重启保持不变； - /// 只会增大有效提前量，不会比配置的 margin 更贴近过期时刻。 - #[must_use] - pub fn refresh_stagger(self, account_id: &ProviderAccountId) -> Duration { - // margin 以整秒配置；秒级粒度在默认 300s 下对应约 10 个扫描槽位。 - let bound = u32::try_from(self.margin.as_secs()).unwrap_or(u32::MAX); - Duration::from_secs(u64::from(stable_factor( - account_id.as_str(), - "refresh-stagger", - 0, - bound, - ))) - } - - /// 含账号错峰偏移的到期判定；扫描路径先按 [`Self::staggered_refresh_bound`] - /// 取回候选超集，再用它在内存中收窄到本轮真正到期的账号。 - #[must_use] - pub fn is_refresh_due_staggered( - self, - account_id: &ProviderAccountId, - access_token_expires_at: SystemTime, - observed_at: SystemTime, - ) -> bool { - let staggered_margin = self.margin.saturating_add(self.refresh_stagger(account_id)); - match access_token_expires_at.duration_since(observed_at) { - Ok(remaining) => remaining <= staggered_margin, - Err(_) => true, - } - } - - /// 错峰候选窗口上界；覆盖 margin 与最大偏移之和。 - /// - /// 与 [`Self::refresh_stagger`] 的值域同址维护：偏移上限为 margin， - /// 因此 `2 × margin` 必然覆盖最大有效提前量。 - #[must_use] - pub fn staggered_refresh_bound(self) -> Duration { - self.margin.saturating_mul(2) - } -} - -/// 指数退避基准延迟；attempt=1 即为该值 -const REFRESH_BACKOFF_BASE_DELAY: Duration = Duration::from_secs(5); -/// 每多一次连续失败，基准延迟乘以该因子 -const REFRESH_BACKOFF_FACTOR: u32 = 3; -/// 退避延迟上限，避免连续失败时无限增长 -const REFRESH_BACKOFF_CAP: Duration = Duration::from_secs(300); - -/// 基于连续失败计数的指数退避重试时刻；复用 `provider_refresh_retry_at` 的稳定扰动 -/// -/// `attempt` 为窗口内累计失败次数（0 与 1 等价，均取基准延迟） -/// 延迟按 -/// `base * factor^(attempt-1)` 增长并封顶到 `REFRESH_BACKOFF_CAP` -pub fn provider_refresh_backoff_at( - account_id: &ProviderAccountId, - observed_at: SystemTime, - attempt: u32, - reason: &'static str, -) -> Result { - let exponent = attempt.saturating_sub(1); - let multiplier = REFRESH_BACKOFF_FACTOR.saturating_pow(exponent); - let scaled_seconds = REFRESH_BACKOFF_BASE_DELAY - .as_secs() - .saturating_mul(u64::from(multiplier)) - .min(REFRESH_BACKOFF_CAP.as_secs()); - let base_delay = Duration::from_secs(scaled_seconds); - provider_refresh_retry_at(account_id, observed_at, base_delay, reason) -} - -/// 临时失败后的持久重试时刻；稳定扰动避免多实例同频重试 -pub fn provider_refresh_retry_at( - account_id: &ProviderAccountId, - observed_at: SystemTime, - base_delay: Duration, - reason: &'static str, -) -> Result { - if base_delay.is_zero() { - return Err(invalid_refresh_policy("schedule refresh retry")); - } - let factor = stable_factor(account_id.as_str(), reason, 800, 1_200); - let millis = u64::try_from(base_delay.as_millis()) - .unwrap_or(u64::MAX) - .saturating_mul(u64::from(factor)) - .saturating_add(500) - / 1_000; - observed_at - .checked_add(Duration::from_millis(millis.max(1))) - .ok_or_else(|| invalid_refresh_policy("schedule refresh retry")) -} - -fn stable_factor(value: &str, salt: &str, minimum: u32, maximum: u32) -> u32 { - let mut hash = 0xcbf2_9ce4_8422_2325_u64; - for byte in value.bytes().chain([0]).chain(salt.bytes()) { - hash ^= u64::from(byte); - hash = hash.wrapping_mul(0x0000_0100_0000_01b3); - } - let width = u64::from(maximum - minimum) + 1; - minimum + u32::try_from(hash % width).unwrap_or_default() -} - -fn invalid_refresh_policy(operation: &'static str) -> ProviderStoreError { - ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) -} - -pub trait ProviderRuntimePolicyPort: Send + Sync { - /// 原子推进预热执行游标，重启或时钟回拨后不重复领取已消费的时刻 - fn claim_warmup_slot<'a>( - &'a self, - _timezone: crate::time::DeploymentTimeZone, - _slot: chrono::NaiveDateTime, - ) -> BoxFuture<'a, Result> { - Box::pin(async { - Err(ProviderStoreError::new( - ProviderStoreErrorKind::Unavailable, - "claim warmup slot", - )) - }) - } - /// 仅首次启动写入该 Provider 的默认选择；已保存的管理配置始终优先 - fn initialize_request_profile<'a>( - &'a self, - _provider: &'a ProviderKind, - initial: OpaqueProviderData, - ) -> BoxFuture<'a, Result> { - Box::pin(async move { Ok(initial) }) - } - - /// 读取候选配置版本实际引用的全局与 Client Key 画像配置 - /// - /// 实现必须在同一数据库快照内核对 revision，且只返回画像投影，不能读取 Key - /// 明文 - /// Provider 代次据此在发布前拒绝已失效的选择 - fn load_request_profile_configurations<'a>( - &'a self, - _revision: ConfigRevision, - _provider: &'a ProviderKind, - ) -> BoxFuture<'a, Result, ProviderStoreError>> { - Box::pin(async move { Ok(Vec::new()) }) - } - - fn load_refresh_policy( - &self, - ) -> BoxFuture<'_, Result>; - - /// 读取账号容量熔断策略；默认关闭，只有实现运行时设置的存储需要覆盖 - fn load_freeze_policy( - &self, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(ProviderFreezePolicy::disabled()) }) - } - - /// 读取账号模型预激活策略；默认关闭，只有实现运行时设置的存储需要覆盖 - fn load_warmup_policy( - &self, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(ProviderWarmupPolicy::disabled()) }) - } -} - -/// 账号容量熔断（自动冻结）策略；来源于 `runtime_settings`， -/// 由 Provider 触发路径与恢复 worker 共享同一份配置事实 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderFreezePolicy { - enabled: bool, - threshold: u32, - window: Duration, - freeze_duration: Duration, - probe_enabled: bool, - probe_model: Option, - adaptive_concurrency: bool, -} - -impl ProviderFreezePolicy { - /// 管理写入与运行策略构造共用的冻结约束 - pub fn validate_values( - threshold: u32, - window_seconds: u64, - duration_seconds: u64, - probe_model: Option<&str>, - ) -> Result<(), &'static str> { - for (valid, field) in [ - ( - (2..=1_000).contains(&threshold), - "account_auto_freeze_threshold", - ), - ( - (60..=3_600).contains(&window_seconds), - "account_auto_freeze_window_seconds", - ), - ( - (300..=604_800).contains(&duration_seconds), - "account_auto_freeze_duration_seconds", - ), - ( - valid_optional_probe_model(probe_model), - "account_auto_freeze_probe_model", - ), - ] { - if !valid { - return Err(field); - } - } - Ok(()) - } - - /// 边界与迁移 `0010_account_auto_freeze.sql` 的 check 约束一致； - /// store 层写入前已校验，这里兜底防御越界配置 - pub fn try_new( - enabled: bool, - threshold: u32, - window_seconds: u64, - freeze_duration_seconds: u64, - probe_enabled: bool, - probe_model: Option, - adaptive_concurrency: bool, - ) -> Result { - if Self::validate_values( - threshold, - window_seconds, - freeze_duration_seconds, - probe_model.as_deref(), - ) - .is_err() - { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate freeze policy", - )); - } - Ok(Self { - enabled, - threshold, - window: Duration::from_secs(window_seconds), - freeze_duration: Duration::from_secs(freeze_duration_seconds), - probe_enabled, - probe_model, - adaptive_concurrency, - }) - } - - /// 功能关闭时的全零策略；触发路径与 worker 都以此短路 - #[must_use] - pub const fn disabled() -> Self { - Self { - enabled: false, - threshold: 2, - window: Duration::from_secs(60), - freeze_duration: Duration::from_secs(300), - probe_enabled: false, - probe_model: None, - adaptive_concurrency: false, - } - } - - #[must_use] - pub const fn enabled(&self) -> bool { - self.enabled - } - - /// 窗口内触发冻结的请求级失败次数阈值 - #[must_use] - pub const fn threshold(&self) -> u32 { - self.threshold - } - - /// 失败计数滑动窗口；每次失败都会顺延窗口 - #[must_use] - pub const fn window(&self) -> Duration { - self.window - } - - /// 冻结时长；探测失败后的顺延也使用该值 - #[must_use] - pub const fn freeze_duration(&self) -> Duration { - self.freeze_duration - } - - #[must_use] - pub const fn probe_enabled(&self) -> bool { - self.probe_enabled - } - - /// 探测模型；`None` 表示由 worker 选择账号可用的第一个模型 - #[must_use] - pub fn probe_model(&self) -> Option<&str> { - self.probe_model.as_deref() - } - - #[must_use] - pub const fn adaptive_concurrency(&self) -> bool { - self.adaptive_concurrency - } -} - -/// 校验每日预激活时间格式，如 "08:00" 或 "08:00,13:00" -#[must_use] -pub fn valid_warmup_schedule_time(value: &str) -> bool { - if value.is_empty() - || value.len() > 255 - || value != value.trim() - || value.chars().any(char::is_control) - { - return false; - } - for part in value.split(',') { - let bytes = part.as_bytes(); - if bytes.len() != 5 || bytes[2] != b':' { - return false; - } - let Ok(hour) = part[0..2].parse::() else { - return false; - }; - let Ok(minute) = part[3..5].parse::() else { - return false; - }; - if hour > 23 || minute > 59 { - return false; - } - } - true -} - -/// 账号模型预激活（预热）策略；来源于 `runtime_settings` -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderWarmupPolicy { - enabled: bool, - schedule_time: String, - model: Option, -} - -impl ProviderWarmupPolicy { - /// 保存设置与创建预热策略使用相同的时间表和模型约束 - pub fn validate_values( - enabled: bool, - schedule_time: &str, - model: Option<&str>, - ) -> Result<(), &'static str> { - if !valid_warmup_schedule_time(schedule_time) { - return Err("account_warmup_schedule_time"); - } - if (enabled && model.is_none()) || !valid_optional_probe_model(model) { - return Err("account_warmup_model"); - } - Ok(()) - } - - pub fn try_new( - enabled: bool, - schedule_time: String, - model: Option, - ) -> Result { - if Self::validate_values(enabled, &schedule_time, model.as_deref()).is_err() { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate warmup policy", - )); - } - Ok(Self { - enabled, - schedule_time, - model, - }) - } - - #[must_use] - pub const fn disabled() -> Self { - Self { - enabled: false, - schedule_time: String::new(), - model: None, - } - } - - #[must_use] - pub const fn enabled(&self) -> bool { - self.enabled - } - - #[must_use] - pub fn schedule_time(&self) -> &str { - &self.schedule_time - } - - #[must_use] - pub fn model(&self) -> Option<&str> { - self.model.as_deref() - } - - /// 解析每日时间列表，返回如 `vec![(8, 0)]` - #[must_use] - pub fn scheduled_times(&self) -> Vec<(u32, u32)> { - self.schedule_time - .split(',') - .filter_map(|part| { - let part = part.trim(); - let mut iter = part.split(':'); - let hour = iter.next()?.parse::().ok()?; - let minute = iter.next()?.parse::().ok()?; - Some((hour, minute)) - }) - .collect() - } -} - -/// OAuth pending flow 的原始绑定只在 Provider 与 Store 边界内短暂存在 -#[derive(Clone, PartialEq, Eq)] -pub struct OAuthPendingBinding(String); - -impl OAuthPendingBinding { - pub fn try_new(value: impl Into) -> Result { - let value = value.into(); - if value.is_empty() || value.len() > 512 { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate OAuth pending binding", - )); - } - Ok(Self(value)) - } - - #[must_use] - pub fn expose_to_store(&self) -> &str { - &self.0 - } -} - -impl fmt::Debug for OAuthPendingBinding { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("OAuthPendingBinding([REDACTED])") - } -} - -#[derive(Clone, PartialEq)] -pub struct NewOAuthPendingFlow { - provider_kind: ProviderKind, - flow: OAuthPendingBinding, - owner: OAuthPendingBinding, - ttl: Duration, - payload: OpaqueProviderData, -} - -impl NewOAuthPendingFlow { - pub fn try_new( - provider_kind: ProviderKind, - flow: OAuthPendingBinding, - owner: OAuthPendingBinding, - ttl: Duration, - payload: OpaqueProviderData, - ) -> Result { - if ttl.is_zero() || ttl > MAX_PENDING_FLOW_TTL { - return Err(ProviderStoreError::new( - ProviderStoreErrorKind::InvalidData, - "validate OAuth pending TTL", - )); - } - Ok(Self { - provider_kind, - flow, - owner, - ttl, - payload, - }) - } - - #[must_use] - pub const fn provider_kind(&self) -> &ProviderKind { - &self.provider_kind - } - - #[must_use] - pub const fn flow(&self) -> &OAuthPendingBinding { - &self.flow - } - - #[must_use] - pub const fn owner(&self) -> &OAuthPendingBinding { - &self.owner - } - - #[must_use] - pub const fn ttl(&self) -> Duration { - self.ttl - } - - #[must_use] - pub const fn payload(&self) -> &OpaqueProviderData { - &self.payload - } -} - -impl fmt::Debug for NewOAuthPendingFlow { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("NewOAuthPendingFlow") - .field("provider_kind", &self.provider_kind) - .field("flow", &self.flow) - .field("owner", &self.owner) - .field("ttl", &self.ttl) - .field("payload", &"[PROVIDER-OWNED]") - .finish() - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum OAuthPendingPutOutcome { - Stored, - AlreadyExists, -} - -/// OAuth 回调处理取得临时 flow 的独占权结果 -/// -/// 与一次性消费不同，Provider 在上游换取令牌失败时可以释放 claim，让同一 flow -/// 使用新的回调地址重试；只有完整校验成功后才会消费 flow -#[derive(Clone, PartialEq)] -pub enum OAuthPendingClaimOutcome { - Claimed(OpaqueProviderData), - NotFound, - OwnerMismatch, - InProgress, -} - -impl fmt::Debug for OAuthPendingClaimOutcome { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Claimed(_) => formatter.write_str("Claimed([PROVIDER-OWNED])"), - Self::NotFound => formatter.write_str("NotFound"), - Self::OwnerMismatch => formatter.write_str("OwnerMismatch"), - Self::InProgress => formatter.write_str("InProgress"), - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum OAuthPendingReleaseOutcome { - Released, - NotFound, - OwnerMismatch, - ClaimMismatch, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum OAuthPendingConsumeOutcome { - Consumed, - NotFound, - OwnerMismatch, - ClaimMismatch, -} - -pub trait OAuthPendingFlowPort: Send + Sync { - fn put_if_absent( - &self, - flow: NewOAuthPendingFlow, - ) -> BoxFuture<'_, Result>; - - fn claim_if_owner<'a>( - &'a self, - provider_kind: &'a ProviderKind, - flow: &'a OAuthPendingBinding, - owner: &'a OAuthPendingBinding, - claim: &'a OAuthPendingBinding, - claim_ttl: Duration, - ) -> BoxFuture<'a, Result>; - - fn release_claim<'a>( - &'a self, - provider_kind: &'a ProviderKind, - flow: &'a OAuthPendingBinding, - owner: &'a OAuthPendingBinding, - claim: &'a OAuthPendingBinding, - ) -> BoxFuture<'a, Result>; - - fn consume_claim<'a>( - &'a self, - provider_kind: &'a ProviderKind, - flow: &'a OAuthPendingBinding, - owner: &'a OAuthPendingBinding, - claim: &'a OAuthPendingBinding, - ) -> BoxFuture<'a, Result>; -} - -/// Provider 只能按能力取用端口，无法取得 Redis client 或 repository 集合 -#[derive(Clone)] -pub struct ProviderStorePorts { - accounts: Arc, - leases: Arc, - session_affinity: Arc, - session_exclusions: Arc, - account_feedback: Arc, - catalog_cache: Arc, - artifact_profiles: Arc, - credential_state: Arc, - cooldowns: Arc, - runtime_policy: Arc, - oauth_pending: Arc, - diagnostics: Arc, -} - -impl ProviderStorePorts { - #[must_use] - // 每个参数代表独立能力端口；合并为单一配置对象会隐藏 Provider 能力边界 - #[expect(clippy::too_many_arguments)] - pub fn new( - accounts: Arc, - leases: Arc, - session_affinity: Arc, - session_exclusions: Arc, - catalog_cache: Arc, - artifact_profiles: Arc, - credential_state: Arc, - cooldowns: Arc, - runtime_policy: Arc, - oauth_pending: Arc, - diagnostics: Arc, - ) -> Self { - Self { - accounts, - leases, - session_affinity, - session_exclusions, - account_feedback: Arc::new(AccountFeedbackStats::default()), - catalog_cache, - artifact_profiles, - credential_state, - cooldowns, - runtime_policy, - oauth_pending, - diagnostics, - } - } - - #[must_use] - pub fn diagnostics(&self) -> Arc { - self.diagnostics.clone() - } - - #[must_use] - pub fn accounts(&self) -> Arc { - Arc::clone(&self.accounts) - } - - #[must_use] - pub fn leases(&self) -> Arc { - Arc::clone(&self.leases) - } - - #[must_use] - pub fn session_affinity(&self) -> Arc { - Arc::clone(&self.session_affinity) - } - - #[must_use] - pub fn session_exclusions(&self) -> Arc { - Arc::clone(&self.session_exclusions) - } - - #[must_use] - pub fn account_feedback(&self) -> Arc { - Arc::clone(&self.account_feedback) - } - - #[must_use] - pub fn catalog_cache(&self) -> Arc { - Arc::clone(&self.catalog_cache) - } - - #[must_use] - pub fn artifact_profiles(&self) -> Arc { - Arc::clone(&self.artifact_profiles) - } - - #[must_use] - pub fn credential_state(&self) -> Arc { - Arc::clone(&self.credential_state) - } - - #[must_use] - pub fn cooldowns(&self) -> Arc { - Arc::clone(&self.cooldowns) - } - - #[must_use] - pub fn runtime_policy(&self) -> Arc { - Arc::clone(&self.runtime_policy) - } - - #[must_use] - pub fn oauth_pending(&self) -> Arc { - Arc::clone(&self.oauth_pending) - } -} - -impl fmt::Debug for ProviderStorePorts { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("ProviderStorePorts([CAPABILITIES])") - } -} - -fn valid_optional_probe_model(model: Option<&str>) -> bool { - model.is_none_or(|model| { - !model.is_empty() - && model.len() <= 128 - && model == model.trim() - && !model.bytes().any(|byte| byte.is_ascii_control()) - }) -} diff --git a/backend/crates/gateway-core/src/provider_ports/affinity.rs b/backend/crates/gateway-core/src/provider_ports/affinity.rs new file mode 100644 index 000000000..b1fafa772 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/affinity.rs @@ -0,0 +1,174 @@ +//! Provider 会话亲和、绑定和排除事实 + +use super::{ProviderStoreError, ProviderStoreErrorKind}; +use crate::{account::ProviderAccountId, identity::ProviderKind}; +use futures::future::BoxFuture; +use std::{collections::BTreeSet, fmt, time::Duration}; + +/// Provider 从原始会话锚点派生的不可逆亲和键 +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProviderSessionAffinityKey(String); + +impl ProviderSessionAffinityKey { + pub fn try_new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() + || value.len() > 128 + || !value.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'-' | b'_') + }) + { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate provider session affinity key", + )); + } + Ok(Self(value)) + } + + #[must_use] + pub fn expose_to_store(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for ProviderSessionAffinityKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderSessionAffinityKey([OPAQUE])") + } +} + +/// 会话绑定快照；版本区分同一账号的不同认领，防止过期和 A → B → A 后的旧写入 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionBinding { + account_id: ProviderAccountId, + revision: String, +} + +impl ProviderSessionBinding { + pub fn new( + account_id: ProviderAccountId, + revision: String, + ) -> Result { + if revision.len() != 32 || !revision.bytes().all(|byte| byte.is_ascii_hexdigit()) { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "decode provider session binding revision", + )); + } + Ok(Self { + account_id, + revision, + }) + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub fn revision(&self) -> &str { + &self.revision + } +} + +/// 已观测请求关联及其账号迁移权限，由 Provider 解释协议后写入 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionAlias { + pub session_key: ProviderSessionAffinityKey, + pub root_session_key: Option, + pub follow_only: bool, +} + +/// 客户端作用域内的会话账号绑定，原始会话身份由 Provider 哈希后传入 +pub trait ProviderSessionAffinityPort: Send + Sync { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + /// 在发送前原子认领、续期或迁移；None 表示快照已变化，调用方必须释放租约重新选择 + /// expected 为 None 只允许首次认领，已有快照必须连同版本匹配 + fn compare_and_bind<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected: Option<&'a ProviderSessionBinding>, + account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + /// 显式观测到的请求关联只指向会话键，不缓存账号，迁移后仍读取当前绑定 + fn load_alias<'a>( + &'a self, + provider: &'a ProviderKind, + alias: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + /// 只允许首次关联或同目标续期，冲突时禁止把同一轮次改指其他会话 + fn bind_alias<'a>( + &'a self, + provider: &'a ProviderKind, + alias: &'a ProviderSessionAffinityKey, + session: &'a ProviderSessionAlias, + ttl: Duration, + ) -> BoxFuture<'a, Result>; +} + +/// Provider 会话内已失败账号的可丢失排除集 +/// +/// Provider 自行派生会话键并决定何时写入或清理；Core 只承载调度所需的账号 ID +/// 与 compare-and-swap revision，不解释任一 Provider 协议字段 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionExclusions { + excluded_accounts: BTreeSet, + revision: String, +} + +impl ProviderSessionExclusions { + #[must_use] + pub const fn new(excluded_accounts: BTreeSet, revision: String) -> Self { + Self { + excluded_accounts, + revision, + } + } + + #[must_use] + pub const fn excluded_accounts(&self) -> &BTreeSet { + &self.excluded_accounts + } + + #[must_use] + pub fn revision(&self) -> &str { + &self.revision + } +} + +/// 可丢失的 Provider 会话级账号排除状态 +/// +/// 该端口不接收协议正文；Provider 只能以不可逆会话键、账号 ID 和固定 TTL 操作 +pub trait ProviderSessionExclusionPort: Send + Sync { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn record_failure<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result>; + + fn clear<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_revision: &'a str, + ) -> BoxFuture<'a, Result>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/bundle.rs b/backend/crates/gateway-core/src/provider_ports/bundle.rs new file mode 100644 index 000000000..13ec981d0 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/bundle.rs @@ -0,0 +1,126 @@ +//! 按能力提供 Provider 所需的存储端口 + +use super::{ + OAuthPendingFlowPort, ProviderArtifactProfileCachePort, ProviderCatalogCachePort, + ProviderCooldownPort, ProviderCredentialStatePort, ProviderLeasePort, + ProviderRuntimePolicyPort, ProviderSessionAffinityPort, ProviderSessionExclusionPort, +}; +use crate::account::{AccountFeedbackStats, ProviderAccountStore}; +use std::{fmt, sync::Arc}; + +/// Provider 只能按能力取用端口，无法取得 Redis client 或 repository 集合 +#[derive(Clone)] +pub struct ProviderStorePorts { + accounts: Arc, + leases: Arc, + session_affinity: Arc, + session_exclusions: Arc, + account_feedback: Arc, + catalog_cache: Arc, + artifact_profiles: Arc, + credential_state: Arc, + cooldowns: Arc, + runtime_policy: Arc, + oauth_pending: Arc, + diagnostics: Arc, +} + +impl ProviderStorePorts { + #[must_use] + // 每个参数代表独立能力端口；合并为单一配置对象会隐藏 Provider 能力边界 + #[expect(clippy::too_many_arguments)] + pub fn new( + accounts: Arc, + leases: Arc, + session_affinity: Arc, + session_exclusions: Arc, + catalog_cache: Arc, + artifact_profiles: Arc, + credential_state: Arc, + cooldowns: Arc, + runtime_policy: Arc, + oauth_pending: Arc, + diagnostics: Arc, + ) -> Self { + Self { + accounts, + leases, + session_affinity, + session_exclusions, + account_feedback: Arc::new(AccountFeedbackStats::default()), + catalog_cache, + artifact_profiles, + credential_state, + cooldowns, + runtime_policy, + oauth_pending, + diagnostics, + } + } + + #[must_use] + pub fn diagnostics(&self) -> Arc { + self.diagnostics.clone() + } + + #[must_use] + pub fn accounts(&self) -> Arc { + Arc::clone(&self.accounts) + } + + #[must_use] + pub fn leases(&self) -> Arc { + Arc::clone(&self.leases) + } + + #[must_use] + pub fn session_affinity(&self) -> Arc { + Arc::clone(&self.session_affinity) + } + + #[must_use] + pub fn session_exclusions(&self) -> Arc { + Arc::clone(&self.session_exclusions) + } + + #[must_use] + pub fn account_feedback(&self) -> Arc { + Arc::clone(&self.account_feedback) + } + + #[must_use] + pub fn catalog_cache(&self) -> Arc { + Arc::clone(&self.catalog_cache) + } + + #[must_use] + pub fn artifact_profiles(&self) -> Arc { + Arc::clone(&self.artifact_profiles) + } + + #[must_use] + pub fn credential_state(&self) -> Arc { + Arc::clone(&self.credential_state) + } + + #[must_use] + pub fn cooldowns(&self) -> Arc { + Arc::clone(&self.cooldowns) + } + + #[must_use] + pub fn runtime_policy(&self) -> Arc { + Arc::clone(&self.runtime_policy) + } + + #[must_use] + pub fn oauth_pending(&self) -> Arc { + Arc::clone(&self.oauth_pending) + } +} + +impl fmt::Debug for ProviderStorePorts { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderStorePorts([CAPABILITIES])") + } +} diff --git a/backend/crates/gateway-core/src/provider_ports/catalog.rs b/backend/crates/gateway-core/src/provider_ports/catalog.rs new file mode 100644 index 000000000..8fc28a2be --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/catalog.rs @@ -0,0 +1,167 @@ +//! Provider 目录与身份画像缓存端口 + +use super::ProviderStoreError; +use crate::{ + account::OpaqueProviderData, + identity::ProviderKind, + validation::{IdentifierError, validate_text}, +}; +use futures::future::BoxFuture; +use std::{ + fmt, + time::{Duration, SystemTime}, +}; + +/// Provider 定义的 catalog cache 作用域 +/// +/// Core 不解释其值；例如 Provider 可以使用套餐、区域或产品线作为共享目录边界 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProviderCatalogScope(String); + +impl ProviderCatalogScope { + /// 创建一个稳定且可用于 Redis 隔离键的 Provider-owned 作用域 + /// + /// # Errors + /// + /// 空值、过长文本或控制字符会被拒绝 + pub fn new(value: impl Into) -> Result { + let value = value.into(); + validate_text(&value, 128, false, None)?; + Ok(Self(value)) + } + + /// 返回 Provider-owned 作用域文本 + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// Opaque catalog cache 的 Provider 与 Provider-owned 作用域 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCatalogCacheKey { + provider_kind: ProviderKind, + scope: ProviderCatalogScope, +} + +impl ProviderCatalogCacheKey { + #[must_use] + pub const fn new(provider_kind: ProviderKind, scope: ProviderCatalogScope) -> Self { + Self { + provider_kind, + scope, + } + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn scope(&self) -> &ProviderCatalogScope { + &self.scope + } +} + +pub trait ProviderCatalogCachePort: Send + Sync { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + ttl: Duration, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} + +/// Provider 从官方制品核验出的可重建请求画像 +/// +/// Core 只用单调制品序号约束覆盖顺序；具体版本字段由对应 Provider 放在 +/// `profile` 中解释 +/// 每个 Provider 的各制品分别保留一份最新画像 +#[derive(Clone, PartialEq)] +pub struct ProviderArtifactProfile { + provider_kind: ProviderKind, + artifact_key: String, + artifact_sequence: u64, + verified_at: SystemTime, + profile: OpaqueProviderData, +} + +impl ProviderArtifactProfile { + #[must_use] + pub const fn new( + provider_kind: ProviderKind, + artifact_key: String, + artifact_sequence: u64, + verified_at: SystemTime, + profile: OpaqueProviderData, + ) -> Self { + Self { + provider_kind, + artifact_key, + artifact_sequence, + verified_at, + profile, + } + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub fn artifact_key(&self) -> &str { + &self.artifact_key + } + + #[must_use] + pub const fn artifact_sequence(&self) -> u64 { + self.artifact_sequence + } + + #[must_use] + pub const fn verified_at(&self) -> SystemTime { + self.verified_at + } + + #[must_use] + pub const fn profile(&self) -> &OpaqueProviderData { + &self.profile + } +} + +impl fmt::Debug for ProviderArtifactProfile { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProviderArtifactProfile") + .field("provider_kind", &self.provider_kind) + .field("artifact_sequence", &self.artifact_sequence) + .field("verified_at", &self.verified_at) + .field("profile", &"[PROVIDER-OWNED]") + .finish() + } +} + +pub trait ProviderArtifactProfileCachePort: Send + Sync { + /// 覆盖同一 Provider、同一制品的固定 cache key + /// + /// 返回 `false` 表示 Store 已持有更高的制品序号；相同序号但内容不同必须返回 + /// [`super::ProviderStoreErrorKind::Conflict`]，不能静默改写已核验画像 + fn replace_if_newer( + &self, + profile: ProviderArtifactProfile, + ttl: Duration, + ) -> BoxFuture<'_, Result>; + + fn read<'a>( + &'a self, + provider_kind: &'a ProviderKind, + artifact_key: &'a str, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/cooldown.rs b/backend/crates/gateway-core/src/provider_ports/cooldown.rs new file mode 100644 index 000000000..99df144d5 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/cooldown.rs @@ -0,0 +1,219 @@ +//! Provider 冷却、容量窗口和运行状态端口 + +use super::ProviderStoreError; +/// 账号级 cooldown 的来源类别；调度状态统一按 `rate_limited` 处理 +pub use crate::account::AccountCooldownKind as ProviderCooldownKind; +use crate::{ + account::{CredentialRevision, ProviderAccountId}, + routing::UpstreamModelId, +}; +use futures::future::BoxFuture; +use std::time::{Duration, SystemTime}; + +/// 可丢失的账号冷却事实，不进入持久状态；探测冻结到期后仍需确认恢复 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCooldown { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + kind: ProviderCooldownKind, +} + +impl ProviderCooldown { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + ) -> Self { + Self::new_with_kind( + account_id, + credential_revision, + until, + ProviderCooldownKind::RateLimit, + ) + } + + #[must_use] + pub const fn new_with_kind( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + kind: ProviderCooldownKind, + ) -> Self { + Self { + account_id, + credential_revision, + until, + kind, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn until(&self) -> SystemTime { + self.until + } + + #[must_use] + pub const fn scheduling_state(&self) -> crate::account::AccountCooldown { + crate::account::AccountCooldown { + until: self.until, + kind: self.kind, + } + } + + #[must_use] + pub const fn kind(&self) -> ProviderCooldownKind { + self.kind + } +} + +/// 可丢失 cooldown 的细粒度作用域 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ProviderCooldownScope { + /// 只阻止同一账号调用指定上游模型 + UpstreamModel(UpstreamModelId), +} + +impl ProviderCooldownScope { + #[must_use] + pub const fn upstream_model(model: UpstreamModelId) -> Self { + Self::UpstreamModel(model) + } + + #[must_use] + pub const fn kind(&self) -> &'static str { + match self { + Self::UpstreamModel(_) => "model", + } + } + + #[must_use] + pub fn value(&self) -> &str { + match self { + Self::UpstreamModel(model) => model.as_str(), + } + } +} + +/// 不进入账号持久状态的账号+作用域 cooldown +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderScopedCooldown { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + scope: ProviderCooldownScope, + until: SystemTime, +} + +impl ProviderScopedCooldown { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + scope: ProviderCooldownScope, + until: SystemTime, + ) -> Self { + Self { + account_id, + credential_revision, + scope, + until, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn scope(&self) -> &ProviderCooldownScope { + &self.scope + } + + #[must_use] + pub const fn until(&self) -> SystemTime { + self.until + } +} + +pub trait ProviderCooldownPort: Send + Sync { + fn put_if_later( + &self, + cooldown: ProviderCooldown, + ) -> BoxFuture<'_, Result>; + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result>; + + fn put_scoped_if_later( + &self, + cooldown: ProviderScopedCooldown, + ) -> BoxFuture<'_, Result>; + + fn read_scoped<'a>( + &'a self, + account_id: &'a ProviderAccountId, + scope: &'a ProviderCooldownScope, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear_scoped<'a>( + &'a self, + account_id: &'a ProviderAccountId, + scope: &'a ProviderCooldownScope, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result>; + + /// 删除账号时清除该账号全部 account/model scope cooldown key + /// 生命周期与 credential revision 无关；不要求调用方持有 revision + fn clear_all<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result>; + + /// 记录一次容量类失败并返回滑动窗口内的累计次数，同时把观测到的账号 + /// 在途并发并入窗口峰值（`in_flight` 为 0 表示本次未观测，跳过峰值更新） + /// 只服务容量熔断触发器；调用频率受失败频率约束，不需要批量接口 + fn record_capacity_failure<'a>( + &'a self, + account_id: &'a ProviderAccountId, + window: Duration, + in_flight: u32, + ) -> BoxFuture<'a, Result>; + + /// 普通请求成功后原子清除临时限流及失败证据；必须保留任何容量冻结 + fn clear_after_success<'a>( + &'a self, + account_id: &'a ProviderAccountId, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; + + /// 读取窗口内观测到的在途并发峰值；无证据时返回 `None` + fn capacity_peak_in_flight<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/credential.rs b/backend/crates/gateway-core/src/provider_ports/credential.rs new file mode 100644 index 000000000..abbd1c75a --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/credential.rs @@ -0,0 +1,90 @@ +//! Provider 凭据运行状态读写端口 + +use super::ProviderStoreError; +use crate::account::{CredentialRevision, CredentialState, ProviderAccountId}; +use futures::future::BoxFuture; +use std::time::{Duration, SystemTime}; + +/// Redis 中可重建的账号状态投影 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCredentialState { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + enabled: bool, + credential_state: CredentialState, + observed_at: SystemTime, +} + +impl ProviderCredentialState { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + enabled: bool, + credential_state: CredentialState, + observed_at: SystemTime, + ) -> Self { + Self { + account_id, + credential_revision, + enabled, + credential_state, + observed_at, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + #[must_use] + pub const fn credential_state(&self) -> CredentialState { + self.credential_state + } + + #[must_use] + pub const fn observed_at(&self) -> SystemTime { + self.observed_at + } +} + +pub trait ProviderCredentialStatePort: Send + Sync { + fn replace( + &self, + state: ProviderCredentialState, + ) -> BoxFuture<'_, Result<(), ProviderStoreError>>; + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result>; + + /// 记录一次瞬态刷新失败并返回窗口内累计失败次数；每次失败刷新 TTL + fn record_refresh_backoff<'a>( + &'a self, + account_id: &'a ProviderAccountId, + window: Duration, + ) -> BoxFuture<'a, Result>; + + /// 凭据完整成功轮换后清零失败计数，退避窗口重新从 base 起步 + fn clear_refresh_backoff<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/error.rs b/backend/crates/gateway-core/src/provider_ports/error.rs new file mode 100644 index 000000000..4d4917be4 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/error.rs @@ -0,0 +1,48 @@ +//! Provider 存储端口的错误分类与原始来源 + +/// Provider 可据此决定是否重试，但看不到 SQL、Redis 或秘密原文 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderStoreErrorKind { + Unavailable, + InvalidData, + Conflict, +} + +/// Provider 存储端口的脱敏错误 +#[derive(Debug, Clone, thiserror::Error)] +#[error("provider store {operation} failed: {kind:?}")] +pub struct ProviderStoreError { + kind: ProviderStoreErrorKind, + operation: &'static str, + source: Option, +} + +impl ProviderStoreError { + #[must_use] + pub const fn new(kind: ProviderStoreErrorKind, operation: &'static str) -> Self { + Self { + kind, + operation, + source: None, + } + } + + /// 包装基础设施失败，分类与原始来源分别保留 + #[must_use] + pub fn caused_by( + kind: ProviderStoreErrorKind, + operation: &'static str, + source: impl Into, + ) -> Self { + Self { + kind, + operation, + source: Some(source.into()), + } + } + + #[must_use] + pub const fn kind(&self) -> ProviderStoreErrorKind { + self.kind + } +} diff --git a/backend/crates/gateway-core/src/provider_ports/lease.rs b/backend/crates/gateway-core/src/provider_ports/lease.rs new file mode 100644 index 000000000..6b3202085 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/lease.rs @@ -0,0 +1,216 @@ +//! Provider 调度与刷新租约的中立端口 + +use super::ProviderStoreError; +use crate::{ + account::{AccountConcurrency, AccountRuntimeSignals, CredentialRevision, ProviderAccountId}, + identity::ProviderKind, + policy::ClientApiKeyId, +}; +use futures::future::BoxFuture; +use std::{collections::BTreeMap, fmt, num::NonZeroU32, time::Duration}; + +/// 一个 Provider 的完整可重建调度状态 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSchedulingState { + signals: BTreeMap, + round_robin_cursor: u64, +} + +impl ProviderSchedulingState { + #[must_use] + pub const fn new( + signals: BTreeMap, + round_robin_cursor: u64, + ) -> Self { + Self { + signals, + round_robin_cursor, + } + } + + #[must_use] + pub const fn signals(&self) -> &BTreeMap { + &self.signals + } + + #[must_use] + pub const fn round_robin_cursor(&self) -> u64 { + self.round_robin_cursor + } +} + +/// 请求级账号 lease 的全部中立事实 +#[derive(Debug, Clone)] +pub struct ProviderSchedulingLeaseRequest { + provider_kind: ProviderKind, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + max_concurrent: AccountConcurrency, + request_interval: Duration, + deadline: crate::lifecycle::Deadline, + cancellation: crate::lifecycle::CancellationToken, +} + +impl ProviderSchedulingLeaseRequest { + #[must_use] + pub fn new( + provider_kind: ProviderKind, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + max_concurrent: impl Into, + request_interval: Duration, + deadline: impl Into, + ) -> Self { + Self { + provider_kind, + account_id, + credential_revision, + max_concurrent: max_concurrent.into(), + request_interval, + deadline: deadline.into(), + cancellation: crate::lifecycle::CancellationToken::new(), + } + } + + #[must_use] + pub fn with_cancellation(mut self, cancellation: crate::lifecycle::CancellationToken) -> Self { + self.cancellation = cancellation; + self + } + + #[must_use] + pub fn cancellation(&self) -> crate::lifecycle::CancellationToken { + self.cancellation.clone() + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn max_concurrent(&self) -> AccountConcurrency { + self.max_concurrent + } + + #[must_use] + pub const fn request_interval(&self) -> Duration { + self.request_interval + } + + #[must_use] + pub const fn deadline(&self) -> crate::lifecycle::Deadline { + self.deadline + } +} + +/// Lease 生命周期由具体 Store guard 管理，Provider 只能持有 +pub trait ProviderLeaseGuard: Send + Sync + 'static {} + +impl ProviderLeaseGuard for T where T: Send + Sync + 'static {} + +pub enum ProviderLeaseAcquisition { + Acquired(Box), + Busy { retry_after: Option }, +} + +impl fmt::Debug for ProviderLeaseAcquisition { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Acquired(_) => formatter.write_str("Acquired([LEASE])"), + Self::Busy { retry_after } => formatter + .debug_struct("Busy") + .field("retry_after", retry_after) + .finish(), + } + } +} + +/// Provider 运行时会持有的三类 lease；刷新必须同时持有全局容量与账号互斥 lease +#[derive(Debug, Clone)] +pub enum ProviderLeaseRequest { + Scheduling(ProviderSchedulingLeaseRequest), + RefreshCapacity(ProviderRefreshCapacityRequest), + Refresh(ProviderRefreshLeaseRequest), +} + +pub trait ProviderLeasePort: Send + Sync { + fn load_state<'a>( + &'a self, + client_api_key_id: &'a ClientApiKeyId, + provider_kind: &'a ProviderKind, + accounts: &'a [ProviderAccountId], + ) -> BoxFuture<'a, Result>; + + fn try_acquire( + &self, + request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result>; + + /// 读取指定账号当前的在途请求数；只用于容量熔断的峰值证据， + /// 支持租约信号的存储实现覆盖，否则视为不可观测（空映射） + fn account_in_flight<'a>( + &'a self, + account_ids: &'a [ProviderAccountId], + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + let _ = account_ids; + Box::pin(async move { Ok(BTreeMap::new()) }) + } +} + +/// 所有 Provider 共享的 OAuth refresh 并发容量 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderRefreshCapacityRequest { + max_concurrent: NonZeroU32, +} + +impl ProviderRefreshCapacityRequest { + #[must_use] + pub const fn new(max_concurrent: NonZeroU32) -> Self { + Self { max_concurrent } + } + + #[must_use] + pub const fn max_concurrent(self) -> NonZeroU32 { + self.max_concurrent + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderRefreshLeaseRequest { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, +} + +impl ProviderRefreshLeaseRequest { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + ) -> Self { + Self { + account_id, + credential_revision, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } +} diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs new file mode 100644 index 000000000..858c3e1b2 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -0,0 +1,41 @@ +//! Provider 运行时所需的存储合同，按能力组织端口与共享策略 + +mod error; +pub use error::{ProviderStoreError, ProviderStoreErrorKind}; +mod lease; +pub use lease::{ + ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, + ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderSchedulingLeaseRequest, + ProviderSchedulingState, +}; +mod affinity; +pub use affinity::{ + ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionAlias, + ProviderSessionBinding, ProviderSessionExclusionPort, ProviderSessionExclusions, +}; +mod catalog; +pub use catalog::{ + ProviderArtifactProfile, ProviderArtifactProfileCachePort, ProviderCatalogCacheKey, + ProviderCatalogCachePort, ProviderCatalogScope, +}; +mod credential; +pub use credential::{ProviderCredentialState, ProviderCredentialStatePort}; +mod cooldown; +pub use cooldown::{ + ProviderCooldown, ProviderCooldownKind, ProviderCooldownPort, ProviderCooldownScope, + ProviderScopedCooldown, +}; +mod refresh; +pub use refresh::{ProviderRefreshPolicy, provider_refresh_backoff_at, provider_refresh_retry_at}; +mod runtime_policy; +pub use runtime_policy::{ + ProviderFreezePolicy, ProviderRuntimePolicyPort, ProviderWarmupPolicy, + valid_warmup_schedule_time, +}; +mod oauth; +pub use oauth::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome, + OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome, +}; +mod bundle; +pub use bundle::ProviderStorePorts; diff --git a/backend/crates/gateway-core/src/provider_ports/oauth.rs b/backend/crates/gateway-core/src/provider_ports/oauth.rs new file mode 100644 index 000000000..a84a965ac --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/oauth.rs @@ -0,0 +1,184 @@ +//! OAuth 授权 pending flow 的领取与消费合同 + +use super::{ProviderStoreError, ProviderStoreErrorKind}; +use crate::{account::OpaqueProviderData, identity::ProviderKind}; +use futures::future::BoxFuture; +use std::{fmt, time::Duration}; + +const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60); + +/// OAuth pending flow 的原始绑定只在 Provider 与 Store 边界内短暂存在 +#[derive(Clone, PartialEq, Eq)] +pub struct OAuthPendingBinding(String); + +impl OAuthPendingBinding { + pub fn try_new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() || value.len() > 512 { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate OAuth pending binding", + )); + } + Ok(Self(value)) + } + + #[must_use] + pub fn expose_to_store(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for OAuthPendingBinding { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("OAuthPendingBinding([REDACTED])") + } +} + +#[derive(Clone, PartialEq)] +pub struct NewOAuthPendingFlow { + provider_kind: ProviderKind, + flow: OAuthPendingBinding, + owner: OAuthPendingBinding, + ttl: Duration, + payload: OpaqueProviderData, +} + +impl NewOAuthPendingFlow { + pub fn try_new( + provider_kind: ProviderKind, + flow: OAuthPendingBinding, + owner: OAuthPendingBinding, + ttl: Duration, + payload: OpaqueProviderData, + ) -> Result { + if ttl.is_zero() || ttl > MAX_PENDING_FLOW_TTL { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate OAuth pending TTL", + )); + } + Ok(Self { + provider_kind, + flow, + owner, + ttl, + payload, + }) + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn flow(&self) -> &OAuthPendingBinding { + &self.flow + } + + #[must_use] + pub const fn owner(&self) -> &OAuthPendingBinding { + &self.owner + } + + #[must_use] + pub const fn ttl(&self) -> Duration { + self.ttl + } + + #[must_use] + pub const fn payload(&self) -> &OpaqueProviderData { + &self.payload + } +} + +impl fmt::Debug for NewOAuthPendingFlow { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("NewOAuthPendingFlow") + .field("provider_kind", &self.provider_kind) + .field("flow", &self.flow) + .field("owner", &self.owner) + .field("ttl", &self.ttl) + .field("payload", &"[PROVIDER-OWNED]") + .finish() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingPutOutcome { + Stored, + AlreadyExists, +} + +/// OAuth 回调处理取得临时 flow 的独占权结果 +/// +/// 与一次性消费不同，Provider 在上游换取令牌失败时可以释放 claim，让同一 flow +/// 使用新的回调地址重试；只有完整校验成功后才会消费 flow +#[derive(Clone, PartialEq)] +pub enum OAuthPendingClaimOutcome { + Claimed(OpaqueProviderData), + NotFound, + OwnerMismatch, + InProgress, +} + +impl fmt::Debug for OAuthPendingClaimOutcome { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Claimed(_) => formatter.write_str("Claimed([PROVIDER-OWNED])"), + Self::NotFound => formatter.write_str("NotFound"), + Self::OwnerMismatch => formatter.write_str("OwnerMismatch"), + Self::InProgress => formatter.write_str("InProgress"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingReleaseOutcome { + Released, + NotFound, + OwnerMismatch, + ClaimMismatch, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingConsumeOutcome { + Consumed, + NotFound, + OwnerMismatch, + ClaimMismatch, +} + +pub trait OAuthPendingFlowPort: Send + Sync { + fn put_if_absent( + &self, + flow: NewOAuthPendingFlow, + ) -> BoxFuture<'_, Result>; + + fn claim_if_owner<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + claim: &'a OAuthPendingBinding, + claim_ttl: Duration, + ) -> BoxFuture<'a, Result>; + + fn release_claim<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + claim: &'a OAuthPendingBinding, + ) -> BoxFuture<'a, Result>; + + fn consume_claim<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + claim: &'a OAuthPendingBinding, + ) -> BoxFuture<'a, Result>; +} diff --git a/backend/crates/gateway-core/src/provider_ports/refresh.rs b/backend/crates/gateway-core/src/provider_ports/refresh.rs new file mode 100644 index 000000000..620288f47 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/refresh.rs @@ -0,0 +1,161 @@ +//! Provider 刷新时机、退避与重试策略 + +use super::{ProviderStoreError, ProviderStoreErrorKind}; +use crate::account::ProviderAccountId; +use std::{ + num::NonZeroU32, + time::{Duration, SystemTime}, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderRefreshPolicy { + margin: Duration, + concurrency: NonZeroU32, +} + +impl ProviderRefreshPolicy { + pub fn try_new(margin: Duration, concurrency: NonZeroU32) -> Result { + if margin.is_zero() { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate refresh policy", + )); + } + Ok(Self { + margin, + concurrency, + }) + } + + #[must_use] + pub const fn margin(self) -> Duration { + self.margin + } + + #[must_use] + pub const fn concurrency(self) -> NonZeroU32 { + self.concurrency + } + + /// 判断 AT 是否已经进入当前配置的提前刷新窗口 + /// + /// `margin` 是运行时策略，不投影为账号的持久化时间字段；已过期 AT 也需要 + /// 尝试用 RT 恢复，因此视为到期 + #[must_use] + pub fn is_refresh_due( + self, + access_token_expires_at: SystemTime, + observed_at: SystemTime, + ) -> bool { + match access_token_expires_at.duration_since(observed_at) { + Ok(remaining) => remaining <= self.margin, + Err(_) => true, + } + } + + /// 按账号派生 `[0, margin]` 内的整秒错峰偏移。 + /// + /// 网关以固定周期扫描到期账号，到期时刻相同的账号会在同一轮齐刷， + /// 在 auth.openai.com 侧形成单 IP 批量刷新特征。偏移从账号 ID 的 + /// 稳定哈希派生（与退避扰动同源、独立 salt），跨扫描与重启保持不变； + /// 只会增大有效提前量，不会比配置的 margin 更贴近过期时刻。 + #[must_use] + pub fn refresh_stagger(self, account_id: &ProviderAccountId) -> Duration { + // margin 以整秒配置；秒级粒度在默认 300s 下对应约 10 个扫描槽位。 + let bound = u32::try_from(self.margin.as_secs()).unwrap_or(u32::MAX); + Duration::from_secs(u64::from(stable_factor( + account_id.as_str(), + "refresh-stagger", + 0, + bound, + ))) + } + + /// 含账号错峰偏移的到期判定；扫描路径先按 [`Self::staggered_refresh_bound`] + /// 取回候选超集，再用它在内存中收窄到本轮真正到期的账号。 + #[must_use] + pub fn is_refresh_due_staggered( + self, + account_id: &ProviderAccountId, + access_token_expires_at: SystemTime, + observed_at: SystemTime, + ) -> bool { + let staggered_margin = self.margin.saturating_add(self.refresh_stagger(account_id)); + match access_token_expires_at.duration_since(observed_at) { + Ok(remaining) => remaining <= staggered_margin, + Err(_) => true, + } + } + + /// 错峰候选窗口上界；覆盖 margin 与最大偏移之和。 + /// + /// 与 [`Self::refresh_stagger`] 的值域同址维护：偏移上限为 margin， + /// 因此 `2 × margin` 必然覆盖最大有效提前量。 + #[must_use] + pub fn staggered_refresh_bound(self) -> Duration { + self.margin.saturating_mul(2) + } +} + +/// 指数退避基准延迟；attempt=1 即为该值 +const REFRESH_BACKOFF_BASE_DELAY: Duration = Duration::from_secs(5); +/// 每多一次连续失败，基准延迟乘以该因子 +const REFRESH_BACKOFF_FACTOR: u32 = 3; +/// 退避延迟上限，避免连续失败时无限增长 +const REFRESH_BACKOFF_CAP: Duration = Duration::from_secs(300); + +/// 基于连续失败计数的指数退避重试时刻；复用 `provider_refresh_retry_at` 的稳定扰动 +/// +/// `attempt` 为窗口内累计失败次数（0 与 1 等价，均取基准延迟） +/// 延迟按 +/// `base * factor^(attempt-1)` 增长并封顶到 `REFRESH_BACKOFF_CAP` +pub fn provider_refresh_backoff_at( + account_id: &ProviderAccountId, + observed_at: SystemTime, + attempt: u32, + reason: &'static str, +) -> Result { + let exponent = attempt.saturating_sub(1); + let multiplier = REFRESH_BACKOFF_FACTOR.saturating_pow(exponent); + let scaled_seconds = REFRESH_BACKOFF_BASE_DELAY + .as_secs() + .saturating_mul(u64::from(multiplier)) + .min(REFRESH_BACKOFF_CAP.as_secs()); + let base_delay = Duration::from_secs(scaled_seconds); + provider_refresh_retry_at(account_id, observed_at, base_delay, reason) +} + +/// 临时失败后的持久重试时刻；稳定扰动避免多实例同频重试 +pub fn provider_refresh_retry_at( + account_id: &ProviderAccountId, + observed_at: SystemTime, + base_delay: Duration, + reason: &'static str, +) -> Result { + if base_delay.is_zero() { + return Err(invalid_refresh_policy("schedule refresh retry")); + } + let factor = stable_factor(account_id.as_str(), reason, 800, 1_200); + let millis = u64::try_from(base_delay.as_millis()) + .unwrap_or(u64::MAX) + .saturating_mul(u64::from(factor)) + .saturating_add(500) + / 1_000; + observed_at + .checked_add(Duration::from_millis(millis.max(1))) + .ok_or_else(|| invalid_refresh_policy("schedule refresh retry")) +} + +fn stable_factor(value: &str, salt: &str, minimum: u32, maximum: u32) -> u32 { + let mut hash = 0xcbf2_9ce4_8422_2325_u64; + for byte in value.bytes().chain([0]).chain(salt.bytes()) { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + let width = u64::from(maximum - minimum) + 1; + minimum + u32::try_from(hash % width).unwrap_or_default() +} + +fn invalid_refresh_policy(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} diff --git a/backend/crates/gateway-core/src/provider_ports/runtime_policy.rs b/backend/crates/gateway-core/src/provider_ports/runtime_policy.rs new file mode 100644 index 000000000..f4d8255e5 --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/runtime_policy.rs @@ -0,0 +1,315 @@ +//! Provider 运行策略端口及冻结、预热配置 + +use super::{ProviderRefreshPolicy, ProviderStoreError, ProviderStoreErrorKind}; +use crate::{account::OpaqueProviderData, identity::ProviderKind, routing::ConfigRevision}; +use futures::future::BoxFuture; +use std::time::Duration; + +pub trait ProviderRuntimePolicyPort: Send + Sync { + /// 原子推进预热执行游标，重启或时钟回拨后不重复领取已消费的时刻 + fn claim_warmup_slot<'a>( + &'a self, + _timezone: crate::time::DeploymentTimeZone, + _slot: chrono::NaiveDateTime, + ) -> BoxFuture<'a, Result> { + Box::pin(async { + Err(ProviderStoreError::new( + ProviderStoreErrorKind::Unavailable, + "claim warmup slot", + )) + }) + } + /// 仅首次启动写入该 Provider 的默认选择；已保存的管理配置始终优先 + fn initialize_request_profile<'a>( + &'a self, + _provider: &'a ProviderKind, + initial: OpaqueProviderData, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(initial) }) + } + + /// 读取候选配置版本实际引用的全局与 Client Key 画像配置 + /// + /// 实现必须在同一数据库快照内核对 revision，且只返回画像投影，不能读取 Key + /// 明文 + /// Provider 代次据此在发布前拒绝已失效的选择 + fn load_request_profile_configurations<'a>( + &'a self, + _revision: ConfigRevision, + _provider: &'a ProviderKind, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { Ok(Vec::new()) }) + } + + fn load_refresh_policy( + &self, + ) -> BoxFuture<'_, Result>; + + /// 读取账号容量熔断策略；默认关闭，只有实现运行时设置的存储需要覆盖 + fn load_freeze_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(ProviderFreezePolicy::disabled()) }) + } + + /// 读取账号模型预激活策略；默认关闭，只有实现运行时设置的存储需要覆盖 + fn load_warmup_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(ProviderWarmupPolicy::disabled()) }) + } +} + +/// 账号容量熔断（自动冻结）策略；来源于 `runtime_settings`， +/// 由 Provider 触发路径与恢复 worker 共享同一份配置事实 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderFreezePolicy { + enabled: bool, + threshold: u32, + window: Duration, + freeze_duration: Duration, + probe_enabled: bool, + probe_model: Option, + adaptive_concurrency: bool, +} + +impl ProviderFreezePolicy { + /// 管理写入与运行策略构造共用的冻结约束 + pub fn validate_values( + threshold: u32, + window_seconds: u64, + duration_seconds: u64, + probe_model: Option<&str>, + ) -> Result<(), &'static str> { + for (valid, field) in [ + ( + (2..=1_000).contains(&threshold), + "account_auto_freeze_threshold", + ), + ( + (60..=3_600).contains(&window_seconds), + "account_auto_freeze_window_seconds", + ), + ( + (300..=604_800).contains(&duration_seconds), + "account_auto_freeze_duration_seconds", + ), + ( + valid_optional_probe_model(probe_model), + "account_auto_freeze_probe_model", + ), + ] { + if !valid { + return Err(field); + } + } + Ok(()) + } + + /// 边界与迁移 `0010_account_auto_freeze.sql` 的 check 约束一致； + /// store 层写入前已校验，这里兜底防御越界配置 + pub fn try_new( + enabled: bool, + threshold: u32, + window_seconds: u64, + freeze_duration_seconds: u64, + probe_enabled: bool, + probe_model: Option, + adaptive_concurrency: bool, + ) -> Result { + if Self::validate_values( + threshold, + window_seconds, + freeze_duration_seconds, + probe_model.as_deref(), + ) + .is_err() + { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate freeze policy", + )); + } + Ok(Self { + enabled, + threshold, + window: Duration::from_secs(window_seconds), + freeze_duration: Duration::from_secs(freeze_duration_seconds), + probe_enabled, + probe_model, + adaptive_concurrency, + }) + } + + /// 功能关闭时的全零策略；触发路径与 worker 都以此短路 + #[must_use] + pub const fn disabled() -> Self { + Self { + enabled: false, + threshold: 2, + window: Duration::from_secs(60), + freeze_duration: Duration::from_secs(300), + probe_enabled: false, + probe_model: None, + adaptive_concurrency: false, + } + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + /// 窗口内触发冻结的请求级失败次数阈值 + #[must_use] + pub const fn threshold(&self) -> u32 { + self.threshold + } + + /// 失败计数滑动窗口；每次失败都会顺延窗口 + #[must_use] + pub const fn window(&self) -> Duration { + self.window + } + + /// 冻结时长；探测失败后的顺延也使用该值 + #[must_use] + pub const fn freeze_duration(&self) -> Duration { + self.freeze_duration + } + + #[must_use] + pub const fn probe_enabled(&self) -> bool { + self.probe_enabled + } + + /// 探测模型；`None` 表示由 worker 选择账号可用的第一个模型 + #[must_use] + pub fn probe_model(&self) -> Option<&str> { + self.probe_model.as_deref() + } + + #[must_use] + pub const fn adaptive_concurrency(&self) -> bool { + self.adaptive_concurrency + } +} + +/// 校验每日预激活时间格式，如 "08:00" 或 "08:00,13:00" +#[must_use] +pub fn valid_warmup_schedule_time(value: &str) -> bool { + if value.is_empty() + || value.len() > 255 + || value != value.trim() + || value.chars().any(char::is_control) + { + return false; + } + for part in value.split(',') { + let bytes = part.as_bytes(); + if bytes.len() != 5 || bytes[2] != b':' { + return false; + } + let Ok(hour) = part[0..2].parse::() else { + return false; + }; + let Ok(minute) = part[3..5].parse::() else { + return false; + }; + if hour > 23 || minute > 59 { + return false; + } + } + true +} + +/// 账号模型预激活（预热）策略；来源于 `runtime_settings` +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderWarmupPolicy { + enabled: bool, + schedule_time: String, + model: Option, +} + +impl ProviderWarmupPolicy { + /// 保存设置与创建预热策略使用相同的时间表和模型约束 + pub fn validate_values( + enabled: bool, + schedule_time: &str, + model: Option<&str>, + ) -> Result<(), &'static str> { + if !valid_warmup_schedule_time(schedule_time) { + return Err("account_warmup_schedule_time"); + } + if (enabled && model.is_none()) || !valid_optional_probe_model(model) { + return Err("account_warmup_model"); + } + Ok(()) + } + + pub fn try_new( + enabled: bool, + schedule_time: String, + model: Option, + ) -> Result { + if Self::validate_values(enabled, &schedule_time, model.as_deref()).is_err() { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate warmup policy", + )); + } + Ok(Self { + enabled, + schedule_time, + model, + }) + } + + #[must_use] + pub const fn disabled() -> Self { + Self { + enabled: false, + schedule_time: String::new(), + model: None, + } + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + #[must_use] + pub fn schedule_time(&self) -> &str { + &self.schedule_time + } + + #[must_use] + pub fn model(&self) -> Option<&str> { + self.model.as_deref() + } + + /// 解析每日时间列表，返回如 `vec![(8, 0)]` + #[must_use] + pub fn scheduled_times(&self) -> Vec<(u32, u32)> { + self.schedule_time + .split(',') + .filter_map(|part| { + let part = part.trim(); + let mut iter = part.split(':'); + let hour = iter.next()?.parse::().ok()?; + let minute = iter.next()?.parse::().ok()?; + Some((hour, minute)) + }) + .collect() + } +} + +fn valid_optional_probe_model(model: Option<&str>) -> bool { + model.is_none_or(|model| { + !model.is_empty() + && model.len() <= 128 + && model == model.trim() + && !model.bytes().any(|byte| byte.is_ascii_control()) + }) +} diff --git a/backend/crates/gateway-core/src/runtime/extensions.rs b/backend/crates/gateway-core/src/routing/extensions.rs similarity index 82% rename from backend/crates/gateway-core/src/runtime/extensions.rs rename to backend/crates/gateway-core/src/routing/extensions.rs index 059cd26cd..e44169104 100644 --- a/backend/crates/gateway-core/src/runtime/extensions.rs +++ b/backend/crates/gateway-core/src/routing/extensions.rs @@ -32,13 +32,6 @@ impl ExtensionSetId { pub trait ExtensionSetLease: Send + Sync { fn is_ready(&self) -> bool; - /// 与发布集合一起冻结的上游适配计划，不能在请求中解析为新代次 - fn upstream_adapters( - &self, - ) -> Option> { - None - } - /// 与集合一起冻结的目录事实，读取不调用插件进程 fn model_aliases(&self) -> &[ContributedModelAlias] { &[] @@ -57,15 +50,6 @@ pub struct ExtensionSetReference { } impl ExtensionSetReference { - #[must_use] - pub fn upstream_adapters( - &self, - ) -> Option { - self.lease.upstream_adapters().map(|plan| { - crate::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new(plan, self.clone()) - }) - } - #[must_use] pub fn new(id: ExtensionSetId, lease: Arc) -> Self { Self { id, lease } diff --git a/backend/crates/gateway-core/src/routing/mod.rs b/backend/crates/gateway-core/src/routing/mod.rs index 2118b334d..532b212b5 100644 --- a/backend/crates/gateway-core/src/routing/mod.rs +++ b/backend/crates/gateway-core/src/routing/mod.rs @@ -1,6 +1,8 @@ //! Provider、模型目录、精确模型映射与请求级候选计划 mod catalog; +pub mod extensions; +pub mod request_settings; pub mod snapshot; pub use crate::account::scope::{ diff --git a/backend/crates/gateway-core/src/routing/request_settings.rs b/backend/crates/gateway-core/src/routing/request_settings.rs new file mode 100644 index 000000000..bb4837b57 --- /dev/null +++ b/backend/crates/gateway-core/src/routing/request_settings.rs @@ -0,0 +1,374 @@ +//! 冻结宿主基线与请求显式改写；派生快照只属于当前调用，不发布全局配置 + +use std::{ + collections::BTreeMap, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::account::FastMode; +use crate::policy::{ClientApiKeyId, ClientPolicy, ClientSettings, RateLimits}; +use crate::routing::RuntimeSnapshot; +use crate::settings::{InvalidSettings, SettingsValues}; + +/// 一次模型调用的有效设置；改写只影响当前请求，不发布配置或修改持久化 revision +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExecutionSettings { + pub runtime: SettingsValues, + pub fast_mode: FastMode, + pub client_limits: RateLimits, + /// null 表示不限制总时长；显式时限从请求开始计时，改写不重置计时原点 + pub timeout_ms: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SettingOverride { + pub instance_id: String, + pub order: u64, + pub value: T, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +struct ExecutionOverrides { + #[serde(skip)] + defaults: Arc, + client_key_id: String, + input: Arc, + fast_mode: Option>, + client_limits: Option>, + timeout_ms: Option>>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +struct HttpTimeout { + input_ms: Option, + change: Option>>, +} + +/// 共享值均不可变；同级子调用的修改不会回写父调用 +#[derive(Clone, Debug)] +pub struct RequestSettings { + baseline: Arc, + effective: Arc, + overrides: Arc>, + execution: Option>, + http_timeout: Option, + order: u64, +} + +impl PartialEq for RequestSettings { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.baseline, &other.baseline) + && self.overrides == other.overrides + && self.execution == other.execution + && self.http_timeout == other.http_timeout + } +} +impl Eq for RequestSettings {} + +impl RequestSettings { + #[must_use] + pub fn new(snapshot: Arc) -> Self { + Self { + baseline: snapshot.clone(), + effective: snapshot, + overrides: Arc::default(), + execution: None, + http_timeout: None, + order: 0, + } + } + + #[must_use] + pub fn snapshot(&self) -> Arc { + self.effective.clone() + } + + #[must_use] + pub fn values(&self) -> &SettingsValues { + self.effective.settings() + } + + /// 来源与有效值分开查询，不允许插件伪造其他实例的写入来源 + pub fn inspect(&self) -> Value { + serde_json::json!({ + "config_revision": self.baseline.revision().get(), + "host": self.baseline.settings(), + "overrides": self.overrides, + "execution": self.execution, + "http_timeout": self.http_timeout, + }) + } + + pub fn replace( + &self, + values: SettingsValues, + instance_id: &str, + ) -> Result { + self.replace_scoped(self.values(), values, instance_id) + } + + /// Key 画像等隐式作用域结果不记录成插件改写，避免随后切换 Key 时泄漏 + fn replace_scoped( + &self, + previous: &SettingsValues, + values: SettingsValues, + instance_id: &str, + ) -> Result { + if previous == &values { + return Ok(self.clone()); + } + let current = serde_json::to_value(previous).map_err(|_| InvalidSettings)?; + let next = serde_json::to_value(&values).map_err(|_| InvalidSettings)?; + let order = self.order.checked_add(1).ok_or(InvalidSettings)?; + let mut overrides = self.overrides.as_ref().clone(); + let effective_values = self.values(); + // 每个已声明设置项整体替换，不猜测画像对象、数组或 null 的深合并含义 + for (name, value) in next.as_object().ok_or(InvalidSettings)? { + if current.get(name) != Some(value) { + overrides.insert( + name.clone(), + SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: value.clone(), + }, + ); + } + } + // 只把本层实际修改的字段写回宿主层，作用域默认值不参与继承 + let values = if previous == effective_values { + values + } else { + apply_overrides(self.baseline.settings(), &overrides)? + }; + let effective = self + .effective + .resolve_settings(&values) + .map_err(|_| InvalidSettings)?; + Ok(Self { + baseline: self.baseline.clone(), + effective, + overrides: Arc::new(overrides), + execution: self.execution.clone(), + http_timeout: self.http_timeout.clone(), + order, + }) + } + + /// 持续会话的新请求读取当前宿主快照，只继承明确改写的字段 + pub fn rebase(&self, snapshot: Arc) -> Result { + if Arc::ptr_eq(&self.baseline, &snapshot) { + return Ok(self.clone()); + } + let effective = if self.overrides.is_empty() { + snapshot.clone() + } else { + snapshot + .resolve_settings(&apply_overrides(snapshot.settings(), &self.overrides)?) + .map_err(|_| InvalidSettings)? + }; + Ok(Self { + baseline: snapshot, + effective, + overrides: self.overrides.clone(), + execution: self.execution.clone(), + http_timeout: self.http_timeout.clone(), + order: self.order, + }) + } + + /// 认证、模型入口和插件视图共用解析规则，派生策略不改变 Key 默认值 + #[must_use] + pub fn apply_policy(&self, policy: ClientPolicy) -> ClientPolicy { + let values = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), None); + policy.with_settings( + values.runtime.request_profiles(), + values.fast_mode, + values.client_limits, + ) + } + + #[must_use] + pub fn with_execution(mut self, policy: &ClientPolicy, timeout_ms: Option) -> Self { + let input = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), timeout_ms); + let previous = self + .execution + .as_ref() + .filter(|scope| scope.client_key_id == policy.key_id().as_str()); + self.execution = Some(Arc::new(ExecutionOverrides { + defaults: policy.defaults().clone(), + client_key_id: policy.key_id().as_str().to_owned(), + input: Arc::new(input), + fast_mode: previous.and_then(|scope| scope.fast_mode.clone()), + client_limits: previous.and_then(|scope| scope.client_limits.clone()), + timeout_ms: previous.and_then(|scope| scope.timeout_ms.clone()), + })); + self + } + + #[must_use] + pub fn execution_values(&self) -> Option { + let scope = self.execution.as_ref()?; + Some(self.resolve_execution( + &scope.client_key_id, + &scope.defaults, + scope.input.timeout_ms, + )) + } + + fn resolve_execution( + &self, + key: &str, + defaults: &ClientSettings, + timeout_ms: Option, + ) -> ExecutionSettings { + let overrides = self + .execution + .as_ref() + .filter(|scope| scope.client_key_id == key); + let mut runtime = self.values().clone(); + // Key 画像覆盖宿主默认，插件显式覆盖整个设置项；此优先级只在这里解释 + if !self.overrides.contains_key("request_profiles") && !defaults.request_profiles.is_empty() + { + let mut profiles = runtime.request_profiles().clone(); + profiles.extend(defaults.request_profiles.clone()); + runtime = runtime.with_request_profiles(profiles); + } + ExecutionSettings { + runtime, + fast_mode: overrides + .and_then(|scope| scope.fast_mode.as_ref()) + .map_or(defaults.fast_mode, |change| change.value), + client_limits: overrides + .and_then(|scope| scope.client_limits.as_ref()) + .map_or(defaults.limits, |change| change.value), + timeout_ms: overrides + .and_then(|scope| scope.timeout_ms.as_ref()) + .map_or(timeout_ms, |change| change.value), + } + } + + pub(crate) fn execution_deadline( + &self, + key: &ClientApiKeyId, + started_at: SystemTime, + ) -> Result { + let scope = self + .execution + .as_ref() + .filter(|scope| scope.client_key_id == key.as_str()) + .ok_or(InvalidSettings)?; + let timeout_ms = scope + .timeout_ms + .as_ref() + .map_or(scope.input.timeout_ms, |change| change.value); + crate::lifecycle::Deadline::from_timeout(started_at, timeout_ms.map(Duration::from_millis)) + .ok_or(InvalidSettings) + } + + pub fn replace_execution( + &self, + values: &ExecutionSettings, + instance_id: &str, + ) -> Result { + let previous = self.execution_values().ok_or(InvalidSettings)?; + let mut updated = + self.replace_scoped(&previous.runtime, values.runtime.clone(), instance_id)?; + if previous.fast_mode == values.fast_mode + && previous.client_limits == values.client_limits + && previous.timeout_ms == values.timeout_ms + { + return Ok(updated); + } + let mut scope = self.execution.as_deref().cloned().ok_or(InvalidSettings)?; + let order = if updated.order == self.order { + self.order.checked_add(1).ok_or(InvalidSettings)? + } else { + updated.order + }; + if previous.fast_mode != values.fast_mode { + scope.fast_mode = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: values.fast_mode, + }); + } + if previous.client_limits != values.client_limits { + scope.client_limits = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: values.client_limits, + }); + } + if previous.timeout_ms != values.timeout_ms { + scope.timeout_ms = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order, + value: values.timeout_ms, + }); + } + updated.execution = Some(Arc::new(scope)); + updated.order = order; + Ok(updated) + } + + #[must_use] + pub fn execution_timeout(&self, key: &ClientApiKeyId) -> Option { + self.execution + .as_ref() + .filter(|scope| scope.client_key_id == key.as_str()) + .and_then(|scope| scope.timeout_ms.as_ref()) + .and_then(|change| change.value.map(Duration::from_millis)) + } + + #[must_use] + pub fn with_http_timeout(mut self, timeout_ms: Option) -> Self { + self.http_timeout.get_or_insert(HttpTimeout { + input_ms: timeout_ms, + change: None, + }); + self + } + + pub fn replace_http_timeout( + mut self, + timeout_ms: Option, + instance_id: &str, + ) -> Result { + let timeout = self.http_timeout.get_or_insert(HttpTimeout { + input_ms: None, + change: None, + }); + if timeout + .change + .as_ref() + .map_or(timeout.input_ms, |change| change.value) + != timeout_ms + { + self.order = self.order.checked_add(1).ok_or(InvalidSettings)?; + timeout.change = Some(SettingOverride { + instance_id: instance_id.to_owned(), + order: self.order, + value: timeout_ms, + }); + } + Ok(self) + } +} + +fn apply_overrides( + values: &SettingsValues, + overrides: &BTreeMap, +) -> Result { + let mut values = serde_json::to_value(values).map_err(|_| InvalidSettings)?; + let object = values.as_object_mut().ok_or(InvalidSettings)?; + for (name, change) in overrides { + object.insert(name.clone(), change.value.clone()); + } + serde_json::from_value(values).map_err(|_| InvalidSettings) +} diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index d059f3915..c5dd3fe50 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -220,7 +220,7 @@ pub enum RuntimeSnapshotCompileError { pub struct RuntimeSnapshotCompiler { store: Arc, catalogs: Arc, - extensions: Option>, + extensions: Option>, } impl RuntimeSnapshotCompiler { @@ -239,7 +239,7 @@ impl RuntimeSnapshotCompiler { #[must_use] pub fn with_extensions( mut self, - extensions: Arc, + extensions: Arc, ) -> Self { self.extensions = Some(extensions); self @@ -275,7 +275,7 @@ impl RuntimeSnapshotCompiler { // 同一配置 revision 的重试复用扩展候选，保留插件策略的发布一致性 let mut prepared_extensions: Option<( ConfigRevision, - crate::runtime::extensions::ExtensionSetReference, + crate::routing::extensions::ExtensionSetReference, )> = None; for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS { let facts = self @@ -524,7 +524,7 @@ async fn compile_runtime_snapshot( #[derive(Debug, Clone)] pub struct RuntimeSnapshot { settings: Arc, - extensions: Option, + extensions: Option, revision: ConfigRevision, providers: Arc>, provider_models: Arc>>, @@ -570,14 +570,14 @@ impl RuntimeSnapshot { #[must_use] pub fn with_extensions( mut self, - extensions: Option, + extensions: Option, ) -> Self { self.extensions = extensions; self } #[must_use] - pub const fn extensions(&self) -> Option<&crate::runtime::extensions::ExtensionSetReference> { + pub const fn extensions(&self) -> Option<&crate::routing::extensions::ExtensionSetReference> { self.extensions.as_ref() } diff --git a/backend/crates/gateway-core/src/runtime/mod.rs b/backend/crates/gateway-core/src/runtime/mod.rs index 5391cb6f7..d43baa5ce 100644 --- a/backend/crates/gateway-core/src/runtime/mod.rs +++ b/backend/crates/gateway-core/src/runtime/mod.rs @@ -1,542 +1,10 @@ //! 运行时快照的原子发布、健康状态与跨进程版本收敛 -pub mod extensions; +mod publications; +mod publisher; -use std::collections::BTreeMap; -use std::future::Future; -use std::pin::Pin; -use std::sync::{Arc, Mutex as SyncMutex, 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::health::{HealthProbe, HealthState}; -use crate::identity::ProviderKind; -use crate::lifecycle::CancellationToken; -use crate::routing::snapshot::{ - RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, +pub use publications::{RuntimeSnapshotHandle, RuntimeSnapshotUnavailable}; +pub use publisher::{ + RuntimeSnapshotPublisher, SnapshotControl, SnapshotRevisionStream, SnapshotSubscriptionError, + SnapshotSubscriptionPort, runtime_revision_needs_refresh, }; -use crate::routing::{ConfigRevision, ProviderCatalogGeneration}; -use crate::task::{ - DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, - WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, - WorkerSchedule, WorkerTaskError, -}; - -const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5); -const INITIAL_BACKOFF: Duration = Duration::from_secs(1); -const MAXIMUM_BACKOFF: Duration = Duration::from_secs(30); -const UNUSED_LEASE_TTL: Duration = Duration::from_secs(30); -const UNUSED_LEASE_RENEWAL: Duration = Duration::from_secs(10); - -/// 不泄漏订阅基础设施细节的通知错误 -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -#[error("runtime snapshot notification is unavailable")] -pub struct SnapshotSubscriptionError; - -impl SnapshotSubscriptionError { - #[must_use] - pub const fn unavailable() -> Self { - Self - } -} - -/// 可丢失的配置 revision 通知流；权威 revision 始终由 Store 端口读取 -pub type SnapshotRevisionStream = - Pin> + Send + 'static>>; - -/// 跨进程 revision 通知的基础设施中立端口 -pub trait SnapshotSubscriptionPort: Send + Sync { - fn publish_snapshot_revision( - &self, - revision: ConfigRevision, - ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>>; - - fn subscribe_snapshot_revisions( - &self, - ) -> BoxFuture<'_, Result>; -} - -/// 请求级冻结失败；此状态必须 fail closed -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -#[error("runtime snapshot is unavailable")] -pub struct RuntimeSnapshotUnavailable; - -/// RuntimeSnapshot 原子发布和请求级冻结句柄 -#[derive(Clone, Default)] -pub struct RuntimeSnapshotHandle { - current: Arc>>>, - publications: Arc>>>, -} - -impl RuntimeSnapshotHandle { - #[must_use] - pub fn new(initial: RuntimeSnapshot) -> Self { - Self { - current: Arc::new(RwLock::new(Some(Arc::new(initial)))), - publications: Default::default(), - } - } - - pub fn publish(&self, snapshot: RuntimeSnapshot) { - *write_unpoisoned(&self.current) = Some(Arc::new(snapshot)); - self.notify_publication(); - } - - pub fn suspend(&self) { - *write_unpoisoned(&self.current) = None; - self.notify_publication(); - } - - /// 合并发布与暂停通知；订阅者收到后读取当前快照，不依赖逐条投递 - pub fn subscribe_publications(&self) -> futures::channel::mpsc::Receiver<()> { - let (sender, receiver) = futures::channel::mpsc::channel(0); - self.publications - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(sender); - receiver - } - - fn notify_publication(&self) { - self.publications - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .retain_mut(|sender| match sender.try_send(()) { - Ok(()) => true, - Err(error) => error.is_full(), - }); - } - - #[must_use] - pub fn revision(&self) -> Option { - read_unpoisoned(&self.current) - .as_ref() - .map(|snapshot| snapshot.revision()) - } - - #[must_use] - pub fn provider_catalog_generations( - &self, - ) -> Option> { - read_unpoisoned(&self.current) - .as_ref() - .map(|snapshot| snapshot.provider_catalog_generations().clone()) - } - - /// 控制面只读诊断可观察尚未就绪的发布候选；数据面仍必须通过 [`Self::acquire`] - #[must_use] - pub fn snapshot_for_diagnostics(&self) -> Option> { - read_unpoisoned(&self.current).clone() - } - - /// 冻结当前 Arc；后续发布不改变已经开始的请求 - pub fn acquire(&self) -> Result, RuntimeSnapshotUnavailable> { - read_unpoisoned(&self.current) - .clone() - .filter(|snapshot| snapshot.extensions().is_none_or(|set| set.can_serve())) - .ok_or(RuntimeSnapshotUnavailable) - } -} - -impl HealthProbe for RuntimeSnapshotHandle { - fn name(&self) -> &'static str { - "runtime_snapshot" - } - - fn check(&self) -> BoxFuture<'_, HealthState> { - Box::pin(async move { - if self.acquire().is_ok() { - HealthState::Healthy - } else { - HealthState::Unhealthy("Runtime snapshot is unavailable".to_owned()) - } - }) - } -} - -/// Admin 提交配置后触发本进程刷新与跨进程通知的对象安全端口 -pub trait SnapshotControl: Send + Sync { - fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()>; -} - -#[derive(Default)] -struct RefreshPriorityState { - pending_commits: usize, - active_background: Option, -} - -struct CommittedRefreshGuard { - priority: Arc>, -} - -impl CommittedRefreshGuard { - fn begin(priority: Arc>) -> Self { - let active = { - let mut state = lock_unpoisoned(&priority); - state.pending_commits += 1; - state.active_background.clone() - }; - // 先登记待发布提交，再取消正在等待目录的后台刷新，避免后者抢先重入 - if let Some(active) = active { - active.cancel(); - } - Self { priority } - } -} - -impl Drop for CommittedRefreshGuard { - fn drop(&mut self) { - lock_unpoisoned(&self.priority).pending_commits -= 1; - } -} - -struct BackgroundRefreshGuard { - priority: Arc>, - cancellation: CancellationToken, -} - -impl BackgroundRefreshGuard { - fn begin(priority: Arc>) -> Option { - let mut state = lock_unpoisoned(&priority); - if state.pending_commits != 0 { - return None; - } - let cancellation = CancellationToken::new(); - state.active_background = Some(cancellation.clone()); - drop(state); - Some(Self { - priority, - cancellation, - }) - } -} - -impl Drop for BackgroundRefreshGuard { - fn drop(&mut self) { - lock_unpoisoned(&self.priority).active_background = None; - } -} - -// 仅中断未完成的后台读取或编译；配置提交负责发布最新事实并发送通知 -async fn preempt_for_committed_refresh( - cancellation: Option<&CancellationToken>, - operation: impl Future, -) -> Result { - let Some(cancellation) = cancellation else { - return Ok(operation.await); - }; - let cancelled = cancellation.cancelled().fuse(); - let operation = operation.fuse(); - pin_mut!(cancelled, operation); - select_biased! { - _ = cancelled => Err(RuntimeSnapshotCompileError::RevisionChanged), - result = operation => Ok(result), - } -} - -/// 配置提交后的本进程快照发布与跨进程失效通知 -#[derive(Clone)] -pub struct RuntimeSnapshotPublisher { - compiler: Arc, - snapshots: RuntimeSnapshotHandle, - subscriptions: Arc, - refresh_lock: Arc>, - refresh_priority: Arc>, -} - -enum RefreshMode { - Required, - Reconcile, - Committed, -} - -impl RuntimeSnapshotPublisher { - #[must_use] - pub fn new( - compiler: Arc, - snapshots: RuntimeSnapshotHandle, - subscriptions: Arc, - ) -> Self { - Self { - compiler, - snapshots, - subscriptions, - refresh_lock: Arc::new(Mutex::new(())), - refresh_priority: Arc::new(SyncMutex::new(RefreshPriorityState::default())), - } - } - - /// 串行重编译并替换本进程快照；无法确认配置时暂停新请求 - pub async fn refresh(&self) -> Result { - self.refresh_with_mode(RefreshMode::Required).await - } - - async fn refresh_with_mode( - &self, - mode: RefreshMode, - ) -> Result { - // 提交先登记并中断慢后台读取；发布/暂停仍在完整临界区内串行执行 - let _committed = matches!(mode, RefreshMode::Committed) - .then(|| CommittedRefreshGuard::begin(Arc::clone(&self.refresh_priority))); - let _refresh = self.refresh_lock.lock().await; - self.refresh_locked_with_priority(mode).await - } - - async fn refresh_locked_with_priority( - &self, - mode: RefreshMode, - ) -> Result { - if matches!(mode, RefreshMode::Committed) { - return self.refresh_locked(mode, None).await; - } - let background = BackgroundRefreshGuard::begin(Arc::clone(&self.refresh_priority)) - .ok_or(RuntimeSnapshotCompileError::RevisionChanged)?; - self.refresh_locked(mode, Some(&background.cancellation)) - .await - } - - async fn refresh_locked( - &self, - mode: RefreshMode, - cancellation: Option<&CancellationToken>, - ) -> Result { - let configuration_changed = if matches!(mode, RefreshMode::Reconcile) { - let persisted_revision = preempt_for_committed_refresh( - cancellation, - 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() - && self - .snapshots - .acquire() - .is_ok_and(|snapshot| snapshot.extensions().is_none_or(|set| set.is_ready())) - { - return Ok(persisted_revision); - } - changed - } else { - true - }; - let compiled = if matches!(mode, RefreshMode::Committed) { - match self.snapshots.snapshot_for_diagnostics() { - Some(previous) => self.compiler.compile_with_cached_catalog(&previous).await, - None => self.compiler.compile().await, - } - } else { - preempt_for_committed_refresh(cancellation, self.compiler.compile()).await? - }; - let snapshot = match compiled { - Ok(snapshot) => snapshot, - Err(error) => { - // 目录对账失败可继续服务旧快照；配置或权限变化无法确认时须暂停 - if configuration_changed { - self.snapshots.suspend(); - } - return Err(error); - } - }; - let revision = snapshot.revision(); - self.snapshots.publish(snapshot); - Ok(revision) - } - - #[must_use] - pub fn published_revision(&self) -> Option { - self.snapshots.revision() - } - - #[must_use] - fn provider_catalogs_need_refresh(&self) -> bool { - let Ok(snapshot) = self.snapshots.acquire() else { - return true; - }; - self.compiler.provider_catalog_generations() != *snapshot.provider_catalog_generations() - } - - /// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚 - async fn publish_committed_inner(&self, committed_revision: ConfigRevision) { - if let Err(error) = self.refresh_with_mode(RefreshMode::Committed).await { - tracing::warn!(?error, "committed runtime snapshot refresh failed"); - } - self.notify_committed_revision(committed_revision).await; - } - - async fn notify_committed_revision(&self, committed_revision: ConfigRevision) { - if let Err(error) = self - .subscriptions - .publish_snapshot_revision(committed_revision) - .await - { - tracing::warn!( - ?error, - revision = committed_revision.get(), - "runtime snapshot revision notification failed" - ); - } - } - - /// 交给 Host 的周期对账与长驻订阅任务 - pub fn worker_contributions(&self) -> Result, WorkerDefinitionError> { - let reconciliation_id = WorkerId::try_new( - WorkerKind::RuntimeSnapshotReconciliation, - "runtime_snapshot", - )?; - let schedule = WorkerSchedule::try_new( - RECONCILIATION_INTERVAL, - INITIAL_BACKOFF, - MAXIMUM_BACKOFF, - UNUSED_LEASE_TTL, - UNUSED_LEASE_RENEWAL, - )?; - let reconciliation = WorkerRegistration::try_new( - reconciliation_id, - WorkerRunnable::Scheduled { - schedule, - lease: None, - task: Box::new(RuntimeSnapshotReconciliationTask { - publisher: self.clone(), - }), - }, - )?; - let subscription_id = - WorkerId::try_new(WorkerKind::RuntimeChangeSubscription, "runtime_snapshot")?; - let restart = DaemonRestartPolicy::try_new(INITIAL_BACKOFF, MAXIMUM_BACKOFF)?; - let subscription = WorkerRegistration::try_new( - subscription_id, - WorkerRunnable::Daemon { - restart, - task: Box::new(RuntimeSnapshotSubscriptionTask { - subscriptions: Arc::clone(&self.subscriptions), - publisher: self.clone(), - }), - }, - )?; - Ok(vec![ - WorkerContribution::Registration(reconciliation), - WorkerContribution::Registration(subscription), - ]) - } -} - -impl SnapshotControl for RuntimeSnapshotPublisher { - fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()> { - Box::pin(async move { - self.publish_committed_inner(committed_revision).await; - }) - } -} - -struct RuntimeSnapshotReconciliationTask { - publisher: RuntimeSnapshotPublisher, -} - -impl ScheduledTask for RuntimeSnapshotReconciliationTask { - fn run_cycle( - &self, - _context: WorkerCycleContext, - ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { - Box::pin(async move { - self.publisher - .refresh_with_mode(RefreshMode::Reconcile) - .await - .map(|_| ()) - .map_err(|source| { - WorkerTaskError::safe("runtime snapshot reconciliation failed") - .with_source(source) - }) - }) - } -} - -struct RuntimeSnapshotSubscriptionTask { - subscriptions: Arc, - publisher: RuntimeSnapshotPublisher, -} - -impl DaemonTask for RuntimeSnapshotSubscriptionTask { - fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { - Box::pin(async move { - let mut retry_delay = INITIAL_BACKOFF; - loop { - if cancellation.is_cancelled() { - return Ok(()); - } - let subscription = self.subscriptions.subscribe_snapshot_revisions().await; - let mut subscription = match subscription { - Ok(subscription) => { - retry_delay = INITIAL_BACKOFF; - subscription - } - Err(_) => { - wait_or_cancel(&cancellation, retry_delay).await; - retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF); - continue; - } - }; - loop { - let cancelled = cancellation.cancelled().fuse(); - let next = subscription.next().fuse(); - pin_mut!(cancelled, next); - let notified = select_biased! { - _ = cancelled => return Ok(()), - next = next => next, - }; - match notified { - Some(Ok(_)) => { - let _ = self.publisher.refresh().await; - } - Some(Err(_)) | None => break, - } - } - } - }) - } -} - -async fn wait_or_cancel(cancellation: &CancellationToken, duration: Duration) { - let cancelled = cancellation.cancelled().fuse(); - let delay = Delay::new(duration).fuse(); - pin_mut!(cancelled, delay); - select_biased! { - _ = cancelled => {}, - _ = delay => {}, - } -} - -/// 当前发布版本与持久版本不一致时必须重载；缺失和回退同样 fail closed -#[must_use] -pub fn runtime_revision_needs_refresh( - published_revision: Option, - persisted_revision: u64, -) -> bool { - published_revision != Some(persisted_revision) -} - -fn read_unpoisoned(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { - lock.read() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -fn write_unpoisoned(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { - lock.write() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -fn lock_unpoisoned(lock: &SyncMutex) -> std::sync::MutexGuard<'_, T> { - lock.lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} diff --git a/backend/crates/gateway-core/src/runtime/publications.rs b/backend/crates/gateway-core/src/runtime/publications.rs new file mode 100644 index 000000000..693c68776 --- /dev/null +++ b/backend/crates/gateway-core/src/runtime/publications.rs @@ -0,0 +1,120 @@ +//! 当前路由快照的原子发布、冻结与健康投影 + +use crate::{ + health::{HealthProbe, HealthState}, + identity::ProviderKind, + routing::{ConfigRevision, ProviderCatalogGeneration, RuntimeSnapshot}, +}; +use futures::future::BoxFuture; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex as SyncMutex, RwLock}, +}; + +/// 请求级冻结失败；此状态必须 fail closed +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot is unavailable")] +pub struct RuntimeSnapshotUnavailable; + +/// RuntimeSnapshot 原子发布和请求级冻结句柄 +#[derive(Clone, Default)] +pub struct RuntimeSnapshotHandle { + current: Arc>>>, + publications: Arc>>>, +} + +impl RuntimeSnapshotHandle { + #[must_use] + pub fn new(initial: RuntimeSnapshot) -> Self { + Self { + current: Arc::new(RwLock::new(Some(Arc::new(initial)))), + publications: Default::default(), + } + } + + pub fn publish(&self, snapshot: RuntimeSnapshot) { + *write_unpoisoned(&self.current) = Some(Arc::new(snapshot)); + self.notify_publication(); + } + + pub fn suspend(&self) { + *write_unpoisoned(&self.current) = None; + self.notify_publication(); + } + + /// 合并发布与暂停通知；订阅者收到后读取当前快照，不依赖逐条投递 + pub fn subscribe_publications(&self) -> futures::channel::mpsc::Receiver<()> { + let (sender, receiver) = futures::channel::mpsc::channel(0); + self.publications + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(sender); + receiver + } + + fn notify_publication(&self) { + self.publications + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .retain_mut(|sender| match sender.try_send(()) { + Ok(()) => true, + Err(error) => error.is_full(), + }); + } + + #[must_use] + pub fn revision(&self) -> Option { + read_unpoisoned(&self.current) + .as_ref() + .map(|snapshot| snapshot.revision()) + } + + #[must_use] + pub fn provider_catalog_generations( + &self, + ) -> Option> { + read_unpoisoned(&self.current) + .as_ref() + .map(|snapshot| snapshot.provider_catalog_generations().clone()) + } + + /// 控制面只读诊断可观察尚未就绪的发布候选；数据面仍必须通过 [`Self::acquire`] + #[must_use] + pub fn snapshot_for_diagnostics(&self) -> Option> { + read_unpoisoned(&self.current).clone() + } + + /// 冻结当前 Arc；后续发布不改变已经开始的请求 + pub fn acquire(&self) -> Result, RuntimeSnapshotUnavailable> { + read_unpoisoned(&self.current) + .clone() + .filter(|snapshot| snapshot.extensions().is_none_or(|set| set.can_serve())) + .ok_or(RuntimeSnapshotUnavailable) + } +} + +impl HealthProbe for RuntimeSnapshotHandle { + fn name(&self) -> &'static str { + "runtime_snapshot" + } + + fn check(&self) -> BoxFuture<'_, HealthState> { + Box::pin(async move { + if self.acquire().is_ok() { + HealthState::Healthy + } else { + HealthState::Unhealthy("Runtime snapshot is unavailable".to_owned()) + } + }) + } +} + +fn read_unpoisoned(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { + lock.read() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn write_unpoisoned(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { + lock.write() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-core/src/runtime/publisher.rs b/backend/crates/gateway-core/src/runtime/publisher.rs new file mode 100644 index 000000000..6944e96a9 --- /dev/null +++ b/backend/crates/gateway-core/src/runtime/publisher.rs @@ -0,0 +1,429 @@ +//! 运行时快照的原子发布、健康状态与跨进程版本收敛 + +use super::publications::RuntimeSnapshotHandle; + +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex as SyncMutex}; +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::lifecycle::CancellationToken; +use crate::routing::ConfigRevision; +use crate::routing::snapshot::{RuntimeSnapshotCompileError, RuntimeSnapshotCompiler}; +use crate::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, + WorkerSchedule, WorkerTaskError, +}; + +const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5); +const INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const MAXIMUM_BACKOFF: Duration = Duration::from_secs(30); +const UNUSED_LEASE_TTL: Duration = Duration::from_secs(30); +const UNUSED_LEASE_RENEWAL: Duration = Duration::from_secs(10); + +/// 不泄漏订阅基础设施细节的通知错误 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot notification is unavailable")] +pub struct SnapshotSubscriptionError; + +impl SnapshotSubscriptionError { + #[must_use] + pub const fn unavailable() -> Self { + Self + } +} + +/// 可丢失的配置 revision 通知流；权威 revision 始终由 Store 端口读取 +pub type SnapshotRevisionStream = + Pin> + Send + 'static>>; + +/// 跨进程 revision 通知的基础设施中立端口 +pub trait SnapshotSubscriptionPort: Send + Sync { + fn publish_snapshot_revision( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>>; + + fn subscribe_snapshot_revisions( + &self, + ) -> BoxFuture<'_, Result>; +} + +/// Admin 提交配置后触发本进程刷新与跨进程通知的对象安全端口 +pub trait SnapshotControl: Send + Sync { + fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()>; +} + +#[derive(Default)] +struct RefreshPriorityState { + pending_commits: usize, + active_background: Option, +} + +struct CommittedRefreshGuard { + priority: Arc>, +} + +impl CommittedRefreshGuard { + fn begin(priority: Arc>) -> Self { + let active = { + let mut state = lock_unpoisoned(&priority); + state.pending_commits += 1; + state.active_background.clone() + }; + // 先登记待发布提交，再取消正在等待目录的后台刷新，避免后者抢先重入 + if let Some(active) = active { + active.cancel(); + } + Self { priority } + } +} + +impl Drop for CommittedRefreshGuard { + fn drop(&mut self) { + lock_unpoisoned(&self.priority).pending_commits -= 1; + } +} + +struct BackgroundRefreshGuard { + priority: Arc>, + cancellation: CancellationToken, +} + +impl BackgroundRefreshGuard { + fn begin(priority: Arc>) -> Option { + let mut state = lock_unpoisoned(&priority); + if state.pending_commits != 0 { + return None; + } + let cancellation = CancellationToken::new(); + state.active_background = Some(cancellation.clone()); + drop(state); + Some(Self { + priority, + cancellation, + }) + } +} + +impl Drop for BackgroundRefreshGuard { + fn drop(&mut self) { + lock_unpoisoned(&self.priority).active_background = None; + } +} + +// 仅中断未完成的后台读取或编译；配置提交负责发布最新事实并发送通知 +async fn preempt_for_committed_refresh( + cancellation: Option<&CancellationToken>, + operation: impl Future, +) -> Result { + let Some(cancellation) = cancellation else { + return Ok(operation.await); + }; + let cancelled = cancellation.cancelled().fuse(); + let operation = operation.fuse(); + pin_mut!(cancelled, operation); + select_biased! { + _ = cancelled => Err(RuntimeSnapshotCompileError::RevisionChanged), + result = operation => Ok(result), + } +} + +/// 配置提交后的本进程快照发布与跨进程失效通知 +#[derive(Clone)] +pub struct RuntimeSnapshotPublisher { + compiler: Arc, + snapshots: RuntimeSnapshotHandle, + subscriptions: Arc, + refresh_lock: Arc>, + refresh_priority: Arc>, +} + +enum RefreshMode { + Required, + Reconcile, + Committed, +} + +impl RuntimeSnapshotPublisher { + #[must_use] + pub fn new( + compiler: Arc, + snapshots: RuntimeSnapshotHandle, + subscriptions: Arc, + ) -> Self { + Self { + compiler, + snapshots, + subscriptions, + refresh_lock: Arc::new(Mutex::new(())), + refresh_priority: Arc::new(SyncMutex::new(RefreshPriorityState::default())), + } + } + + /// 串行重编译并替换本进程快照；无法确认配置时暂停新请求 + pub async fn refresh(&self) -> Result { + self.refresh_with_mode(RefreshMode::Required).await + } + + async fn refresh_with_mode( + &self, + mode: RefreshMode, + ) -> Result { + // 提交先登记并中断慢后台读取；发布/暂停仍在完整临界区内串行执行 + let _committed = matches!(mode, RefreshMode::Committed) + .then(|| CommittedRefreshGuard::begin(Arc::clone(&self.refresh_priority))); + let _refresh = self.refresh_lock.lock().await; + self.refresh_locked_with_priority(mode).await + } + + async fn refresh_locked_with_priority( + &self, + mode: RefreshMode, + ) -> Result { + if matches!(mode, RefreshMode::Committed) { + return self.refresh_locked(mode, None).await; + } + let background = BackgroundRefreshGuard::begin(Arc::clone(&self.refresh_priority)) + .ok_or(RuntimeSnapshotCompileError::RevisionChanged)?; + self.refresh_locked(mode, Some(&background.cancellation)) + .await + } + + async fn refresh_locked( + &self, + mode: RefreshMode, + cancellation: Option<&CancellationToken>, + ) -> Result { + let configuration_changed = if matches!(mode, RefreshMode::Reconcile) { + let persisted_revision = preempt_for_committed_refresh( + cancellation, + 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() + && self + .snapshots + .acquire() + .is_ok_and(|snapshot| snapshot.extensions().is_none_or(|set| set.is_ready())) + { + return Ok(persisted_revision); + } + changed + } else { + true + }; + let compiled = if matches!(mode, RefreshMode::Committed) { + match self.snapshots.snapshot_for_diagnostics() { + Some(previous) => self.compiler.compile_with_cached_catalog(&previous).await, + None => self.compiler.compile().await, + } + } else { + preempt_for_committed_refresh(cancellation, self.compiler.compile()).await? + }; + let snapshot = match compiled { + Ok(snapshot) => snapshot, + Err(error) => { + // 目录对账失败可继续服务旧快照；配置或权限变化无法确认时须暂停 + if configuration_changed { + self.snapshots.suspend(); + } + return Err(error); + } + }; + let revision = snapshot.revision(); + self.snapshots.publish(snapshot); + Ok(revision) + } + + #[must_use] + pub fn published_revision(&self) -> Option { + self.snapshots.revision() + } + + #[must_use] + fn provider_catalogs_need_refresh(&self) -> bool { + let Ok(snapshot) = self.snapshots.acquire() else { + return true; + }; + self.compiler.provider_catalog_generations() != *snapshot.provider_catalog_generations() + } + + /// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚 + async fn publish_committed_inner(&self, committed_revision: ConfigRevision) { + if let Err(error) = self.refresh_with_mode(RefreshMode::Committed).await { + tracing::warn!(?error, "committed runtime snapshot refresh failed"); + } + self.notify_committed_revision(committed_revision).await; + } + + async fn notify_committed_revision(&self, committed_revision: ConfigRevision) { + if let Err(error) = self + .subscriptions + .publish_snapshot_revision(committed_revision) + .await + { + tracing::warn!( + ?error, + revision = committed_revision.get(), + "runtime snapshot revision notification failed" + ); + } + } + + /// 交给 Host 的周期对账与长驻订阅任务 + pub fn worker_contributions(&self) -> Result, WorkerDefinitionError> { + let reconciliation_id = WorkerId::try_new( + WorkerKind::RuntimeSnapshotReconciliation, + "runtime_snapshot", + )?; + let schedule = WorkerSchedule::try_new( + RECONCILIATION_INTERVAL, + INITIAL_BACKOFF, + MAXIMUM_BACKOFF, + UNUSED_LEASE_TTL, + UNUSED_LEASE_RENEWAL, + )?; + let reconciliation = WorkerRegistration::try_new( + reconciliation_id, + WorkerRunnable::Scheduled { + schedule, + lease: None, + task: Box::new(RuntimeSnapshotReconciliationTask { + publisher: self.clone(), + }), + }, + )?; + let subscription_id = + WorkerId::try_new(WorkerKind::RuntimeChangeSubscription, "runtime_snapshot")?; + let restart = DaemonRestartPolicy::try_new(INITIAL_BACKOFF, MAXIMUM_BACKOFF)?; + let subscription = WorkerRegistration::try_new( + subscription_id, + WorkerRunnable::Daemon { + restart, + task: Box::new(RuntimeSnapshotSubscriptionTask { + subscriptions: Arc::clone(&self.subscriptions), + publisher: self.clone(), + }), + }, + )?; + Ok(vec![ + WorkerContribution::Registration(reconciliation), + WorkerContribution::Registration(subscription), + ]) + } +} + +impl SnapshotControl for RuntimeSnapshotPublisher { + fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()> { + Box::pin(async move { + self.publish_committed_inner(committed_revision).await; + }) + } +} + +struct RuntimeSnapshotReconciliationTask { + publisher: RuntimeSnapshotPublisher, +} + +impl ScheduledTask for RuntimeSnapshotReconciliationTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + self.publisher + .refresh_with_mode(RefreshMode::Reconcile) + .await + .map(|_| ()) + .map_err(|source| { + WorkerTaskError::safe("runtime snapshot reconciliation failed") + .with_source(source) + }) + }) + } +} + +struct RuntimeSnapshotSubscriptionTask { + subscriptions: Arc, + publisher: RuntimeSnapshotPublisher, +} + +impl DaemonTask for RuntimeSnapshotSubscriptionTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut retry_delay = INITIAL_BACKOFF; + loop { + if cancellation.is_cancelled() { + return Ok(()); + } + let subscription = self.subscriptions.subscribe_snapshot_revisions().await; + let mut subscription = match subscription { + Ok(subscription) => { + retry_delay = INITIAL_BACKOFF; + subscription + } + Err(_) => { + wait_or_cancel(&cancellation, retry_delay).await; + retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF); + continue; + } + }; + loop { + let cancelled = cancellation.cancelled().fuse(); + let next = subscription.next().fuse(); + pin_mut!(cancelled, next); + let notified = select_biased! { + _ = cancelled => return Ok(()), + next = next => next, + }; + match notified { + Some(Ok(_)) => { + let _ = self.publisher.refresh().await; + } + Some(Err(_)) | None => break, + } + } + } + }) + } +} + +async fn wait_or_cancel(cancellation: &CancellationToken, duration: Duration) { + let cancelled = cancellation.cancelled().fuse(); + let delay = Delay::new(duration).fuse(); + pin_mut!(cancelled, delay); + select_biased! { + _ = cancelled => {}, + _ = delay => {}, + } +} + +/// 当前发布版本与持久版本不一致时必须重载；缺失和回退同样 fail closed +#[must_use] +pub fn runtime_revision_needs_refresh( + published_revision: Option, + persisted_revision: u64, +) -> bool { + published_revision != Some(persisted_revision) +} + +fn lock_unpoisoned(lock: &SyncMutex) -> std::sync::MutexGuard<'_, T> { + lock.lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-core/src/settings/mod.rs b/backend/crates/gateway-core/src/settings/mod.rs index b521c09c0..10adbfb08 100644 --- a/backend/crates/gateway-core/src/settings/mod.rs +++ b/backend/crates/gateway-core/src/settings/mod.rs @@ -1,17 +1,4 @@ -//! 冻结宿主基线与请求显式改写；派生快照只属于当前调用，不发布全局配置 - -use std::{ - collections::BTreeMap, - sync::Arc, - time::{Duration, SystemTime}, -}; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::account::FastMode; -use crate::policy::{ClientApiKeyId, ClientPolicy, ClientSettings, RateLimits}; -use crate::routing::RuntimeSnapshot; +//! 运行设置的纯值、校验与编译结果 pub(crate) mod compiled; mod values; @@ -19,365 +6,6 @@ pub use values::{ SettingsValues, client_min_versions, response_body_limit, validate_request_limits, }; -/// 一次模型调用的有效设置；改写只影响当前请求，不发布配置或修改持久化 revision -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ExecutionSettings { - pub runtime: SettingsValues, - pub fast_mode: FastMode, - pub client_limits: RateLimits, - /// null 表示不限制总时长；显式时限从请求开始计时，改写不重置计时原点 - pub timeout_ms: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct SettingOverride { - pub instance_id: String, - pub order: u64, - pub value: T, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -struct ExecutionOverrides { - #[serde(skip)] - defaults: Arc, - client_key_id: String, - input: Arc, - fast_mode: Option>, - client_limits: Option>, - timeout_ms: Option>>, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -struct HttpTimeout { - input_ms: Option, - change: Option>>, -} - -/// 共享值均不可变；同级子调用的修改不会回写父调用 -#[derive(Clone, Debug)] -pub struct RequestSettings { - baseline: Arc, - effective: Arc, - overrides: Arc>, - execution: Option>, - http_timeout: Option, - order: u64, -} - -impl PartialEq for RequestSettings { - fn eq(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.baseline, &other.baseline) - && self.overrides == other.overrides - && self.execution == other.execution - && self.http_timeout == other.http_timeout - } -} -impl Eq for RequestSettings {} - #[derive(Debug, thiserror::Error)] #[error("request settings are invalid")] pub struct InvalidSettings; - -impl RequestSettings { - #[must_use] - pub fn new(snapshot: Arc) -> Self { - Self { - baseline: snapshot.clone(), - effective: snapshot, - overrides: Arc::default(), - execution: None, - http_timeout: None, - order: 0, - } - } - - #[must_use] - pub fn snapshot(&self) -> Arc { - self.effective.clone() - } - - #[must_use] - pub fn values(&self) -> &SettingsValues { - self.effective.settings() - } - - /// 来源与有效值分开查询，不允许插件伪造其他实例的写入来源 - pub fn inspect(&self) -> Value { - serde_json::json!({ - "config_revision": self.baseline.revision().get(), - "host": self.baseline.settings(), - "overrides": self.overrides, - "execution": self.execution, - "http_timeout": self.http_timeout, - }) - } - - pub fn replace( - &self, - values: SettingsValues, - instance_id: &str, - ) -> Result { - self.replace_scoped(self.values(), values, instance_id) - } - - /// Key 画像等隐式作用域结果不记录成插件改写，避免随后切换 Key 时泄漏 - fn replace_scoped( - &self, - previous: &SettingsValues, - values: SettingsValues, - instance_id: &str, - ) -> Result { - if previous == &values { - return Ok(self.clone()); - } - let current = serde_json::to_value(previous).map_err(|_| InvalidSettings)?; - let next = serde_json::to_value(&values).map_err(|_| InvalidSettings)?; - let order = self.order.checked_add(1).ok_or(InvalidSettings)?; - let mut overrides = self.overrides.as_ref().clone(); - let effective_values = self.values(); - // 每个已声明设置项整体替换，不猜测画像对象、数组或 null 的深合并含义 - for (name, value) in next.as_object().ok_or(InvalidSettings)? { - if current.get(name) != Some(value) { - overrides.insert( - name.clone(), - SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: value.clone(), - }, - ); - } - } - // 只把本层实际修改的字段写回宿主层，作用域默认值不参与继承 - let values = if previous == effective_values { - values - } else { - apply_overrides(self.baseline.settings(), &overrides)? - }; - let effective = self - .effective - .resolve_settings(&values) - .map_err(|_| InvalidSettings)?; - Ok(Self { - baseline: self.baseline.clone(), - effective, - overrides: Arc::new(overrides), - execution: self.execution.clone(), - http_timeout: self.http_timeout.clone(), - order, - }) - } - - /// 持续会话的新请求读取当前宿主快照，只继承明确改写的字段 - pub fn rebase(&self, snapshot: Arc) -> Result { - if Arc::ptr_eq(&self.baseline, &snapshot) { - return Ok(self.clone()); - } - let effective = if self.overrides.is_empty() { - snapshot.clone() - } else { - snapshot - .resolve_settings(&apply_overrides(snapshot.settings(), &self.overrides)?) - .map_err(|_| InvalidSettings)? - }; - Ok(Self { - baseline: snapshot, - effective, - overrides: self.overrides.clone(), - execution: self.execution.clone(), - http_timeout: self.http_timeout.clone(), - order: self.order, - }) - } - - /// 认证、模型入口和插件视图共用解析规则，派生策略不改变 Key 默认值 - #[must_use] - pub fn apply_policy(&self, policy: ClientPolicy) -> ClientPolicy { - let values = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), None); - policy.with_settings( - values.runtime.request_profiles(), - values.fast_mode, - values.client_limits, - ) - } - - #[must_use] - pub fn with_execution(mut self, policy: &ClientPolicy, timeout_ms: Option) -> Self { - let input = self.resolve_execution(policy.key_id().as_str(), policy.defaults(), timeout_ms); - let previous = self - .execution - .as_ref() - .filter(|scope| scope.client_key_id == policy.key_id().as_str()); - self.execution = Some(Arc::new(ExecutionOverrides { - defaults: policy.defaults().clone(), - client_key_id: policy.key_id().as_str().to_owned(), - input: Arc::new(input), - fast_mode: previous.and_then(|scope| scope.fast_mode.clone()), - client_limits: previous.and_then(|scope| scope.client_limits.clone()), - timeout_ms: previous.and_then(|scope| scope.timeout_ms.clone()), - })); - self - } - - #[must_use] - pub fn execution_values(&self) -> Option { - let scope = self.execution.as_ref()?; - Some(self.resolve_execution( - &scope.client_key_id, - &scope.defaults, - scope.input.timeout_ms, - )) - } - - fn resolve_execution( - &self, - key: &str, - defaults: &ClientSettings, - timeout_ms: Option, - ) -> ExecutionSettings { - let overrides = self - .execution - .as_ref() - .filter(|scope| scope.client_key_id == key); - let mut runtime = self.values().clone(); - // Key 画像覆盖宿主默认，插件显式覆盖整个设置项；此优先级只在这里解释 - if !self.overrides.contains_key("request_profiles") && !defaults.request_profiles.is_empty() - { - let mut profiles = runtime.request_profiles().clone(); - profiles.extend(defaults.request_profiles.clone()); - runtime = runtime.with_request_profiles(profiles); - } - ExecutionSettings { - runtime, - fast_mode: overrides - .and_then(|scope| scope.fast_mode.as_ref()) - .map_or(defaults.fast_mode, |change| change.value), - client_limits: overrides - .and_then(|scope| scope.client_limits.as_ref()) - .map_or(defaults.limits, |change| change.value), - timeout_ms: overrides - .and_then(|scope| scope.timeout_ms.as_ref()) - .map_or(timeout_ms, |change| change.value), - } - } - - pub(crate) fn execution_deadline( - &self, - key: &ClientApiKeyId, - started_at: SystemTime, - ) -> Result { - let scope = self - .execution - .as_ref() - .filter(|scope| scope.client_key_id == key.as_str()) - .ok_or(InvalidSettings)?; - let timeout_ms = scope - .timeout_ms - .as_ref() - .map_or(scope.input.timeout_ms, |change| change.value); - crate::lifecycle::Deadline::from_timeout(started_at, timeout_ms.map(Duration::from_millis)) - .ok_or(InvalidSettings) - } - - pub fn replace_execution( - &self, - values: &ExecutionSettings, - instance_id: &str, - ) -> Result { - let previous = self.execution_values().ok_or(InvalidSettings)?; - let mut updated = - self.replace_scoped(&previous.runtime, values.runtime.clone(), instance_id)?; - if previous.fast_mode == values.fast_mode - && previous.client_limits == values.client_limits - && previous.timeout_ms == values.timeout_ms - { - return Ok(updated); - } - let mut scope = self.execution.as_deref().cloned().ok_or(InvalidSettings)?; - let order = if updated.order == self.order { - self.order.checked_add(1).ok_or(InvalidSettings)? - } else { - updated.order - }; - if previous.fast_mode != values.fast_mode { - scope.fast_mode = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: values.fast_mode, - }); - } - if previous.client_limits != values.client_limits { - scope.client_limits = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: values.client_limits, - }); - } - if previous.timeout_ms != values.timeout_ms { - scope.timeout_ms = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order, - value: values.timeout_ms, - }); - } - updated.execution = Some(Arc::new(scope)); - updated.order = order; - Ok(updated) - } - - #[must_use] - pub fn execution_timeout(&self, key: &ClientApiKeyId) -> Option { - self.execution - .as_ref() - .filter(|scope| scope.client_key_id == key.as_str()) - .and_then(|scope| scope.timeout_ms.as_ref()) - .and_then(|change| change.value.map(Duration::from_millis)) - } - - #[must_use] - pub fn with_http_timeout(mut self, timeout_ms: Option) -> Self { - self.http_timeout.get_or_insert(HttpTimeout { - input_ms: timeout_ms, - change: None, - }); - self - } - - pub fn replace_http_timeout( - mut self, - timeout_ms: Option, - instance_id: &str, - ) -> Result { - let timeout = self.http_timeout.get_or_insert(HttpTimeout { - input_ms: None, - change: None, - }); - if timeout - .change - .as_ref() - .map_or(timeout.input_ms, |change| change.value) - != timeout_ms - { - self.order = self.order.checked_add(1).ok_or(InvalidSettings)?; - timeout.change = Some(SettingOverride { - instance_id: instance_id.to_owned(), - order: self.order, - value: timeout_ms, - }); - } - Ok(self) - } -} - -fn apply_overrides( - values: &SettingsValues, - overrides: &BTreeMap, -) -> Result { - let mut values = serde_json::to_value(values).map_err(|_| InvalidSettings)?; - let object = values.as_object_mut().ok_or(InvalidSettings)?; - for (name, change) in overrides { - object.insert(name.clone(), change.value.clone()); - } - serde_json::from_value(values).map_err(|_| InvalidSettings) -} diff --git a/backend/crates/gateway-core/src/upstream.rs b/backend/crates/gateway-core/src/upstream.rs index ec4db950c..9c17c6123 100644 --- a/backend/crates/gateway-core/src/upstream.rs +++ b/backend/crates/gateway-core/src/upstream.rs @@ -1,6 +1,6 @@ //! 上游调用跨 Engine、Event、Error 与 Provider 共享的中立边界事实 -use std::fmt; +use std::{fmt, num::NonZeroU32}; use crate::validation::{IdentifierError, validate_text}; @@ -73,3 +73,15 @@ impl fmt::Debug for OpaqueUpstreamValue { formatter.write_str("OpaqueUpstreamValue()") } } + +/// Provider 可为当前 attempt 选择的请求局部传输档位 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum AttemptTransport { + /// 使用 Provider 的默认传输策略 + #[default] + Default, + /// 固定账号重试 Provider 的首选传输；序号由 Provider 的独立预算驱动 + Retry(NonZeroU32), + /// 使用 Provider 定义的备用传输 + Fallback, +} diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution/mod.rs similarity index 99% rename from backend/crates/gateway-core/tests/engine/execution.rs rename to backend/crates/gateway-core/tests/engine/execution/mod.rs index 3ae6cce71..790f7b151 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution/mod.rs @@ -1,4 +1,5 @@ //! 执行服务的准入、预算、请求记录与终结行为测试 +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, @@ -2964,10 +2965,7 @@ use gateway_core::routing::{ PublicModelId, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId, }; -use gateway_core::runtime::{ - RuntimeSnapshotHandle, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, -}; +use gateway_core::runtime::RuntimeSnapshotHandle; use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; use serde_json::{Value, json}; @@ -3218,7 +3216,7 @@ impl Provider for NativeCatalogProvider { fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models() { #[derive(Debug)] struct Aliases(Vec); - impl gateway_core::runtime::extensions::ExtensionSetLease for Aliases { + impl gateway_core::routing::extensions::ExtensionSetLease for Aliases { fn is_ready(&self) -> bool { true } @@ -3239,7 +3237,7 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models( .with_settings(&settings) .unwrap() .with_extensions(Some(ExtensionSetReference::new( - gateway_core::runtime::extensions::ExtensionSetId::new("catalog-aliases".into()) + gateway_core::routing::extensions::ExtensionSetId::new("catalog-aliases".into()) .unwrap(), Arc::new(Aliases(vec![ gateway_core::routing::ContributedModelAlias { @@ -4500,7 +4498,7 @@ fn invalid_request_settings_leave_the_prepared_execution_unchanged() { #[test] fn execution_settings_must_match_the_prepared_key() { - use gateway_core::settings::RequestSettings; + use gateway_core::routing::request_settings::RequestSettings; let service = service(Arc::new(Admissions::default()), Arc::new(Budget::default())); let client = service.authenticate("sk_start_test").unwrap(); @@ -4595,7 +4593,7 @@ fn repeated_connection_failures_never_block_later_requests_for_the_provider() { #[test] fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() { block_on(async { - use gateway_core::settings::RequestSettings; + use gateway_core::routing::request_settings::RequestSettings; let snapshots = RuntimeSnapshotHandle::new(start_snapshot()); let service = DefaultExecutionService::new( snapshots.clone(), @@ -4707,7 +4705,7 @@ fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() { #[test] fn unchanged_and_precompiled_request_settings_reuse_the_frozen_snapshot() { - use gateway_core::settings::RequestSettings; + use gateway_core::routing::request_settings::RequestSettings; let snapshot = Arc::new(start_snapshot()); let baseline = RequestSettings::new(snapshot.clone()); @@ -4785,7 +4783,9 @@ fn unchanged_and_precompiled_request_settings_reuse_the_frozen_snapshot() { #[test] fn child_settings_recompute_key_scope_without_inheriting_parent_defaults() { block_on(async { - use gateway_core::{account::OpaqueProviderData, settings::RequestSettings}; + use gateway_core::{ + account::OpaqueProviderData, routing::request_settings::RequestSettings, + }; let provider = ProviderKind::new("openai").unwrap(); let profiles = |name: &str| { BTreeMap::from([( diff --git a/backend/crates/gateway-core/tests/engine/extensions.rs b/backend/crates/gateway-core/tests/engine/extensions.rs index 739fbb0da..4fead6150 100644 --- a/backend/crates/gateway-core/tests/engine/extensions.rs +++ b/backend/crates/gateway-core/tests/engine/extensions.rs @@ -1,8 +1,13 @@ -//! 执行测试共用的扩展代次引用与就绪租约替身 +//! 执行扩展索引的代次隔离、计划保活与租约替身 use std::sync::Arc; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::engine::extensions::{ExecutionExtensionIndex, ExecutionExtensionPlans}; +use gateway_core::engine::upstream_adapter::{UpstreamAdapter, UpstreamAdapterPlan}; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::{ + engine::AttemptContext, error::ProviderError, identity::ProviderKind, routing::UpstreamModelId, +}; struct ReadyLease; @@ -16,3 +21,71 @@ pub(super) fn reference(id: &str) -> ExtensionSetReference { let id = ExtensionSetId::new(id.into()).unwrap(); ExtensionSetReference::new(id, Arc::new(ReadyLease)) } + +#[derive(Debug)] +struct EmptyAdapterPlan; + +impl UpstreamAdapterPlan for EmptyAdapterPlan { + fn select( + &self, + _: &AttemptContext, + _: &ProviderKind, + _: &UpstreamModelId, + ) -> Result>, ProviderError> { + Ok(None) + } +} + +struct PublishedPlans(Arc); + +impl ExtensionSetLease for PublishedPlans { + fn is_ready(&self) -> bool { + self.0.upstream_adapters.is_some() + } +} + +#[test] +fn adapter_only_generation_stays_alive_until_the_last_frozen_plan_is_released() { + let index = ExecutionExtensionIndex::default(); + let id = ExtensionSetId::new("adapter-only".into()).unwrap(); + let plans = Arc::new(ExecutionExtensionPlans { + middleware: None, + upstream_adapters: Some(Arc::new(EmptyAdapterPlan)), + }); + let weak_plans = Arc::downgrade(&plans); + let published = Arc::new(PublishedPlans(index.register(id.clone(), plans).unwrap())); + let weak_publication = Arc::downgrade(&published); + let reference = ExtensionSetReference::new(id.clone(), published); + assert!(index.middleware(&reference).is_none()); + let frozen = index.upstream_adapters(&reference).unwrap(); + drop(reference); + + // 当前发布可以退场；在途计划仍保活整个代次，不能覆盖相同身份 + assert!(weak_publication.upgrade().is_some()); + assert!(weak_plans.upgrade().is_some()); + assert!(index.register(id.clone(), Arc::default()).is_err()); + let in_flight = frozen.clone(); + drop(frozen); + assert!(weak_publication.upgrade().is_some()); + drop(in_flight); + + assert!(weak_publication.upgrade().is_none()); + assert!(weak_plans.upgrade().is_none()); + let released = ExtensionSetReference::new(id.clone(), Arc::new(ReadyLease)); + assert!(index.upstream_adapters(&released).is_none()); + let _replacement = index.register(id, Arc::default()).unwrap(); +} + +#[test] +fn execution_index_does_not_own_registered_generations() { + let index = ExecutionExtensionIndex::default(); + let reference = reference("unpublished"); + let plans = Arc::new(ExecutionExtensionPlans::default()); + let weak = Arc::downgrade(&plans); + let owner = index.register(reference.id().clone(), plans).unwrap(); + assert!(weak.upgrade().is_some()); + drop(owner); + assert!(weak.upgrade().is_none()); + assert!(index.middleware(&reference).is_none()); + assert!(index.upstream_adapters(&reference).is_none()); +} diff --git a/backend/crates/gateway-core/tests/engine/middleware.rs b/backend/crates/gateway-core/tests/engine/middleware/mod.rs similarity index 95% rename from backend/crates/gateway-core/tests/engine/middleware.rs rename to backend/crates/gateway-core/tests/engine/middleware/mod.rs index 24bb2b665..8d9ddc034 100644 --- a/backend/crates/gateway-core/tests/engine/middleware.rs +++ b/backend/crates/gateway-core/tests/engine/middleware/mod.rs @@ -12,11 +12,13 @@ use gateway_core::account::ProviderAccountId; use gateway_core::account::scope::AccountGroupId; use gateway_core::engine::ModelRequestId; use gateway_core::engine::execution::ClientTransport; -use gateway_core::engine::extensions::ExtensionCallScope; +use gateway_core::engine::extensions::{ + ExecutionExtensionIndex, ExecutionExtensionPlans, ExtensionCallScope, +}; use gateway_core::engine::middleware::{ FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareBody, MiddlewareContext, MiddlewareError, - MiddlewareExtensionIndex, MiddlewareFrame, MiddlewareHeader, MiddlewareMount, MiddlewareNext, - MiddlewarePlan, MiddlewareRequest, MiddlewareResponse, MiddlewareTarget, + MiddlewareFrame, MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewarePlan, + MiddlewareRequest, MiddlewareResponse, MiddlewareTarget, }; use gateway_core::engine::provider::{ ProviderCallMetadata, ProviderMiddlewareTerminal, ProviderStream, execute_attempt_middleware, @@ -27,7 +29,7 @@ use gateway_core::lifecycle::CancellationToken; use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload}; use gateway_core::policy::ClientApiKeyId; use gateway_core::routing::UpstreamModelId; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::upstream::UpstreamTransport; use serde_json::{Value, json}; @@ -303,14 +305,17 @@ fn attempt_middleware_consumes_owned_next_once_and_preserves_host_envelopes() { futures::executor::block_on(async { let plan_calls = Arc::new(AtomicUsize::new(0)); let terminal_calls = Arc::new(AtomicUsize::new(0)); - let index = MiddlewareExtensionIndex::default(); + let index = ExecutionExtensionIndex::default(); let id = ExtensionSetId::new("middleware-generation".to_owned()).unwrap(); let reference = ExtensionSetReference::new(id.clone(), Arc::new(ReadyGeneration)); let owner = index .register( id.clone(), - Arc::new(RewritingPlan { - calls: Arc::clone(&plan_calls), + Arc::new(ExecutionExtensionPlans { + middleware: Some(Arc::new(RewritingPlan { + calls: Arc::clone(&plan_calls), + })), + upstream_adapters: None, }), ) .unwrap(); @@ -318,14 +323,17 @@ fn attempt_middleware_consumes_owned_next_once_and_preserves_host_envelopes() { index .register( id, - Arc::new(RewritingPlan { - calls: Arc::new(AtomicUsize::new(0)), + Arc::new(ExecutionExtensionPlans { + middleware: Some(Arc::new(RewritingPlan { + calls: Arc::new(AtomicUsize::new(0)), + })), + upstream_adapters: None }), ) .is_err(), "同一发布代次不能静默替换中间件计划", ); - let plan = index.resolve(&reference).unwrap(); + let plan = index.middleware(&reference).unwrap(); let terminal: ProviderMiddlewareTerminal = { let terminal_calls = Arc::clone(&terminal_calls); diff --git a/backend/crates/gateway-core/tests/engine/policy.rs b/backend/crates/gateway-core/tests/engine/policy.rs index dfd6a6792..e876a1e97 100644 --- a/backend/crates/gateway-core/tests/engine/policy.rs +++ b/backend/crates/gateway-core/tests/engine/policy.rs @@ -24,7 +24,7 @@ use gateway_core::{ }, policy::ClientApiKeyId, routing::ProviderKind, - runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, + routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, }; #[derive(Debug)] diff --git a/backend/crates/gateway-core/tests/main.rs b/backend/crates/gateway-core/tests/main.rs index 48d0a3cde..15edffcc2 100644 --- a/backend/crates/gateway-core/tests/main.rs +++ b/backend/crates/gateway-core/tests/main.rs @@ -16,7 +16,6 @@ mod policy; mod provider_ports; mod routing; mod runtime; -mod settings; mod task; mod upstream; diff --git a/backend/crates/gateway-core/tests/middleware/mod.rs b/backend/crates/gateway-core/tests/middleware.rs similarity index 100% rename from backend/crates/gateway-core/tests/middleware/mod.rs rename to backend/crates/gateway-core/tests/middleware.rs diff --git a/backend/crates/gateway-core/tests/provider_ports.rs b/backend/crates/gateway-core/tests/provider_ports/mod.rs similarity index 100% rename from backend/crates/gateway-core/tests/provider_ports.rs rename to backend/crates/gateway-core/tests/provider_ports/mod.rs diff --git a/backend/crates/gateway-core/tests/routing/mod.rs b/backend/crates/gateway-core/tests/routing/mod.rs index e0ad53344..92d21416c 100644 --- a/backend/crates/gateway-core/tests/routing/mod.rs +++ b/backend/crates/gateway-core/tests/routing/mod.rs @@ -18,6 +18,7 @@ use gateway_core::routing::{ UpstreamModelId, }; +mod request_settings; mod snapshot; fn settings() -> gateway_core::settings::SettingsValues { diff --git a/backend/crates/gateway-core/tests/settings/mod.rs b/backend/crates/gateway-core/tests/routing/request_settings.rs similarity index 99% rename from backend/crates/gateway-core/tests/settings/mod.rs rename to backend/crates/gateway-core/tests/routing/request_settings.rs index c3b6f0291..b574e985e 100644 --- a/backend/crates/gateway-core/tests/settings/mod.rs +++ b/backend/crates/gateway-core/tests/routing/request_settings.rs @@ -7,8 +7,9 @@ use gateway_core::{ account::OpaqueProviderData, identity::ProviderKind, policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits}, + routing::request_settings::RequestSettings, routing::{ClientRoutingScope, ConfigRevision, FrozenAccountScope, RuntimeSnapshot}, - settings::{RequestSettings, SettingsValues}, + settings::SettingsValues, }; use serde_json::json; diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs index f39bf6cb5..d76a16531 100644 --- a/backend/crates/gateway-core/tests/routing/snapshot.rs +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -10,6 +10,10 @@ use futures::future::BoxFuture; use gateway_core::account::FastMode; use gateway_core::operation::OperationKind; use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits}; +use gateway_core::routing::extensions::{ + ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetLease, + ExtensionSetReference, +}; use gateway_core::routing::snapshot::{ RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, SnapshotClientPolicyFacts, SnapshotFacts, @@ -20,10 +24,6 @@ use gateway_core::routing::{ ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable, ProviderKind, ProviderModelCapabilities, PublicModelId, UpstreamModelId, }; -use gateway_core::runtime::extensions::{ - ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetLease, - ExtensionSetReference, -}; use gateway_core::settings::SettingsValues; #[derive(Clone)] @@ -1084,7 +1084,9 @@ fn key_profiles_replace_whole_global_choice_and_previous_snapshot_stays_frozen() .unwrap() }; let values = |snapshot: &gateway_core::routing::RuntimeSnapshot| { - let settings = gateway_core::settings::RequestSettings::new(Arc::new(snapshot.clone())); + let settings = gateway_core::routing::request_settings::RequestSettings::new(Arc::new( + snapshot.clone(), + )); let mut values: Vec<_> = snapshot .client_policies() .map(|policy| { diff --git a/backend/crates/gateway-core/tests/runtime/mod.rs b/backend/crates/gateway-core/tests/runtime/mod.rs index 7915b1dbc..fc8652468 100644 --- a/backend/crates/gateway-core/tests/runtime/mod.rs +++ b/backend/crates/gateway-core/tests/runtime/mod.rs @@ -1,7 +1,7 @@ //! 验证运行时快照发布、版本对账与并发更新行为 use std::collections::BTreeMap; -mod extensions; +mod publications; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; diff --git a/backend/crates/gateway-core/tests/runtime/extensions.rs b/backend/crates/gateway-core/tests/runtime/publications.rs similarity index 96% rename from backend/crates/gateway-core/tests/runtime/extensions.rs rename to backend/crates/gateway-core/tests/runtime/publications.rs index b8488ebc1..43e832952 100644 --- a/backend/crates/gateway-core/tests/runtime/extensions.rs +++ b/backend/crates/gateway-core/tests/runtime/publications.rs @@ -5,9 +5,9 @@ use std::sync::{ atomic::{AtomicBool, AtomicUsize, Ordering}, }; -use gateway_core::runtime::{ - RuntimeSnapshotHandle, - extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, +use gateway_core::{ + routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, + runtime::RuntimeSnapshotHandle, }; struct Lease { diff --git a/backend/crates/gateway-host/src/lib.rs b/backend/crates/gateway-host/src/lib.rs index 91efed550..7ef96b4b7 100644 --- a/backend/crates/gateway-host/src/lib.rs +++ b/backend/crates/gateway-host/src/lib.rs @@ -59,7 +59,9 @@ pub async fn initialize(config: HostConfig) -> Result { let log_guard = initialize_logging(&config.logging, config.timezone)?; let cancellation = CancellationToken::new(); let connections = Arc::new(ConnectionTracker::new(cancellation.clone())); - let workers = WorkerSupervisor::new(cancellation.clone()); + // 关闭请求先排空 HTTP/WS，写泵仍需接收在途请求的最后写入 + // Worker 使用独立信号，只由 drain 结束后的 shutdown 或析构触发 + let workers = WorkerSupervisor::new(CancellationToken::new()); let system = Arc::new(ProcessSystemOperations::new( cancellation.clone(), config.system_update.clone(), @@ -190,6 +192,14 @@ impl HostBundle { } } +impl Drop for HostBundle { + fn drop(&mut self) { + // CLI、启动失败和 serve future 被丢弃也要通知公共能力停止 + // WorkerSupervisor 自己负责取消并终止剩余任务 + self.cancellation.cancel(); + } +} + #[derive(Debug, thiserror::Error)] pub enum HostError { #[error(transparent)] diff --git a/backend/crates/gateway-host/src/system_update/config.rs b/backend/crates/gateway-host/src/system_update/config.rs new file mode 100644 index 000000000..65fcb6cef --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/config.rs @@ -0,0 +1,165 @@ +//! 系统更新的部署配置、编译元数据与本地路径解析 + +use std::{ + env, fs, + path::{Path, PathBuf}, +}; + +use serde::Deserialize; + +use super::{OperationError, conflict, internal}; +use crate::config::ConfigError; + +const APP_BINARY_NAME: &str = "codex-proxy-rs"; +const DEFAULT_GITHUB_API_BASE: &str = "https://api.github.com/repos"; +const DEFAULT_UPDATE_REPOSITORY: &str = "zyycn/codex-proxy-rs"; + +/// 系统更新与重启配置；所有字段只由 Host 解释 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(default)] +pub struct SystemUpdateConfig { + pub version: String, + pub git_sha: String, + pub build_time: String, + pub deployment_mode: String, + pub build_type: String, + pub update_repository: Option, + pub github_api_base: String, + pub executable_path: Option, + /// 未显式指定时，由组合根传入 API 实际使用的静态资源目录 + pub web_dist_dir: Option, + pub update_state_file: PathBuf, + pub update_lock_file: PathBuf, + pub update_temp_dir: PathBuf, + pub self_restart_enabled: bool, +} + +impl Default for SystemUpdateConfig { + fn default() -> Self { + let deployment_mode = + environment_value("CPR_DEPLOYMENT_MODE").unwrap_or_else(|| "source".to_owned()); + let executable_path = environment_value("CPR_UPDATE_EXE_PATH") + .map(PathBuf::from) + .or_else(|| { + (deployment_mode == "docker") + .then(|| PathBuf::from("/app/bin").join(APP_BINARY_NAME)) + }); + let update_state_file = environment_value("CPR_UPDATE_STATE_FILE") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("update-state.json")); + let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE") + .map(PathBuf::from) + .unwrap_or_else(|| update_state_file.with_file_name("update.lock")); + let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| default_temp_dir(&update_state_file)); + Self { + version: option_env!("CPR_VERSION") + .unwrap_or(env!("CARGO_PKG_VERSION")) + .to_owned(), + git_sha: option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_owned(), + build_time: option_env!("CPR_BUILD_TIME") + .unwrap_or("unknown") + .to_owned(), + deployment_mode, + build_type: option_env!("CPR_BUILD_TYPE").unwrap_or("source").to_owned(), + update_repository: Some( + environment_value("CPR_UPDATE_REPOSITORY") + .unwrap_or_else(|| DEFAULT_UPDATE_REPOSITORY.to_owned()), + ), + github_api_base: environment_value("CPR_GITHUB_API_BASE") + .unwrap_or_else(|| DEFAULT_GITHUB_API_BASE.to_owned()), + executable_path, + web_dist_dir: None, + update_state_file, + update_lock_file, + update_temp_dir, + self_restart_enabled: environment_value("CPR_ENABLE_SELF_RESTART").as_deref() + == Some("true"), + } + } +} + +impl SystemUpdateConfig { + pub(crate) fn resolve_and_validate( + &mut self, + source_dir: &Path, + runtime_data_dir: &Path, + asset_directory: &Path, + ) -> Result<(), ConfigError> { + let web_dist_dir = self + .web_dist_dir + .get_or_insert_with(|| asset_directory.to_path_buf()); + if web_dist_dir.as_os_str().is_empty() { + return Err(ConfigError::InvalidField("host.system_update.web_dist_dir")); + } + for path in [self.executable_path.as_mut(), self.web_dist_dir.as_mut()] + .into_iter() + .flatten() + { + if path.is_relative() { + *path = source_dir.join(&*path); + } + } + for path in [ + &mut self.update_state_file, + &mut self.update_lock_file, + &mut self.update_temp_dir, + ] { + if path.is_relative() { + *path = runtime_data_dir.join(&*path); + } + } + if self.version.trim().is_empty() { + return Err(ConfigError::InvalidField("host.system_update.version")); + } + if !matches!( + self.deployment_mode.as_str(), + "source" | "binary" | "docker" + ) { + return Err(ConfigError::InvalidField( + "host.system_update.deployment_mode", + )); + } + Ok(()) + } + + pub(crate) fn executable_path(&self) -> Result { + if let Some(path) = &self.executable_path { + return Ok(path.clone()); + } + env::current_exe() + .and_then(fs::canonicalize) + .map_err(|error| internal(format!("failed to resolve executable: {error}"))) + } + + pub(crate) fn official_plugins_dir(&self) -> Result { + let executable = self.executable_path()?; + let parent = executable + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .ok_or_else(|| internal("failed to resolve official plugin release directory"))?; + Ok(parent.join("plugins").join("official")) + } + + pub(crate) fn web_dist_dir(&self) -> Result<&Path, OperationError> { + self.web_dist_dir + .as_deref() + .filter(|path| !path.as_os_str().is_empty()) + .ok_or_else(|| conflict("web assets directory is not configured")) + } +} + +pub(super) fn environment_value(key: &str) -> Option { + env::var(key) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn default_temp_dir(state_file: &Path) -> PathBuf { + state_file + .parent() + .map(|parent| parent.join("update-tmp")) + .unwrap_or_else(|| std::env::temp_dir().join("codex-proxy-rs-update")) +} diff --git a/backend/crates/gateway-host/src/system_update/download.rs b/backend/crates/gateway-host/src/system_update/download.rs index c99ed258d..bd7685304 100644 --- a/backend/crates/gateway-host/src/system_update/download.rs +++ b/backend/crates/gateway-host/src/system_update/download.rs @@ -8,8 +8,9 @@ use std::time::Duration; use futures::StreamExt as _; use sha2::Digest as _; +use super::events::UpdateEvents; use super::release::{download_client, validate_download_url}; -use super::{OperationError, UpdateEvents, invalid, upstream}; +use super::{OperationError, invalid, upstream}; pub(crate) const MAX_DOWNLOAD_SIZE: u64 = 500 * 1024 * 1024; pub(crate) const MAX_CHECKSUM_SIZE: u64 = 1024 * 1024; diff --git a/backend/crates/gateway-host/src/system_update/error.rs b/backend/crates/gateway-host/src/system_update/error.rs new file mode 100644 index 000000000..bdf7d85bb --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/error.rs @@ -0,0 +1,21 @@ +//! 系统更新各基础设施步骤共用的管理端错误分类 + +use gateway_admin::ports::system::{SystemOperationError, SystemOperationErrorKind}; + +pub(super) type OperationError = SystemOperationError; + +pub(super) fn invalid(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Invalid, message) +} + +pub(super) fn conflict(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Conflict, message) +} + +pub(super) fn upstream(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Upstream, message) +} + +pub(super) fn internal(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Internal, message) +} diff --git a/backend/crates/gateway-host/src/system_update/events.rs b/backend/crates/gateway-host/src/system_update/events.rs new file mode 100644 index 000000000..f58fcbd8c --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/events.rs @@ -0,0 +1,130 @@ +//! 更新日志事件的有界广播、递增标识与终态标记 + +use chrono::Utc; +use gateway_admin::model::system::{SystemUpdateEvent, SystemUpdateEventLevel}; +use std::sync::atomic::{AtomicU64, Ordering}; +use tokio::sync::broadcast; + +pub(super) struct UpdateEvents { + sender: broadcast::Sender<(SystemUpdateEvent, bool)>, + sequence: AtomicU64, +} + +impl Default for UpdateEvents { + fn default() -> Self { + let (sender, _) = broadcast::channel(256); + Self { + sender, + sequence: AtomicU64::new(0), + } + } +} + +impl UpdateEvents { + pub(super) fn subscribe(&self) -> broadcast::Receiver<(SystemUpdateEvent, bool)> { + self.sender.subscribe() + } + + pub(super) fn info( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Info, + operation_id, + step, + message, + None, + false, + ); + } + + pub(super) fn success( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Success, + operation_id, + step, + message, + None, + false, + ); + } + + pub(super) fn success_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Success, + operation_id, + step, + message, + Some(100), + true, + ); + } + + pub(super) fn error_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Error, + operation_id, + step, + message, + None, + true, + ); + } + + pub(super) fn emit_progress(&self, operation_id: &str, message: &str, progress: u8) { + self.emit( + SystemUpdateEventLevel::Info, + Some(operation_id), + Some("download"), + message, + Some(progress), + false, + ); + } + + fn emit( + &self, + level: SystemUpdateEventLevel, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + progress_percent: Option, + terminal: bool, + ) { + let now = Utc::now(); + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); + let event = SystemUpdateEvent { + id: format!( + "update-log-{}-{sequence}", + now.timestamp_nanos_opt() + .unwrap_or_else(|| now.timestamp_millis()) + ), + operation_id: operation_id.map(str::to_owned), + level, + step: step.map(str::to_owned), + message: message.into(), + terminal, + progress_percent, + occurred_at: now, + }; + let _ = self.sender.send((event, terminal)); + } +} diff --git a/backend/crates/gateway-host/src/system_update/installation.rs b/backend/crates/gateway-host/src/system_update/installation.rs index d4d827b81..8b107548c 100644 --- a/backend/crates/gateway-host/src/system_update/installation.rs +++ b/backend/crates/gateway-host/src/system_update/installation.rs @@ -10,6 +10,8 @@ use sha2::{Digest as _, Sha256}; use super::swap::backup_path_for; use super::{OperationError, SystemUpdateConfig, conflict, internal}; +pub(super) const OFFICIAL_PLUGIN_MANIFEST: &str = "plugin-release-manifest.json"; + #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] pub(crate) struct ReleaseFiles { @@ -36,9 +38,7 @@ impl ReleaseFiles { } fn read(binary: &Path, web: &Path, plugins: &Path) -> Result { - if !web.join("index.html").is_file() - || !plugins.join(super::OFFICIAL_PLUGIN_MANIFEST).is_file() - { + if !web.join("index.html").is_file() || !plugins.join(OFFICIAL_PLUGIN_MANIFEST).is_file() { return Err(conflict("安装文件不完整，请核对部署后重试")); } let mut binary_hash = Sha256::new(); diff --git a/backend/crates/gateway-host/src/system_update/mod.rs b/backend/crates/gateway-host/src/system_update/mod.rs index d26757702..38b9e17e4 100644 --- a/backend/crates/gateway-host/src/system_update/mod.rs +++ b/backend/crates/gateway-host/src/system_update/mod.rs @@ -1,962 +1,19 @@ -//! 版本、安全自更新、回滚与进程重启的 Host-owned 实现 +//! 版本、安全自更新、回滚与进程重启的 Host-owned 实现入口 mod archive; +mod config; mod download; +mod error; +mod events; mod installation; +mod operations; mod process; mod release; mod state; mod swap; -use std::env; -use std::fs; -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +pub use config::SystemUpdateConfig; +pub use operations::ProcessSystemOperations; +pub use release::validate_download_url; -use async_trait::async_trait; -use chrono::Utc; -use gateway_admin::model::system::{ - SystemOperationAccepted, SystemOperationKind, SystemUpdateChannel as UpdateChannel, - SystemUpdateDetail, SystemUpdateEvent, SystemUpdateEventLevel, SystemUpdatePolicy, - SystemUpdateStatus, SystemVersion, -}; -use gateway_admin::ports::system::{ - SystemOperationError, SystemOperationErrorKind, SystemOperations, SystemRestartPreflight, - SystemUpdateCandidate, SystemUpdateEventStream, SystemUpdatePreflight, -}; -use gateway_core::lifecycle::CancellationToken; -use serde::Deserialize; -use tokio::sync::{Mutex as AsyncMutex, broadcast}; - -use crate::config::ConfigError; - -use self::archive::extract_release; -use self::download::{MAX_CHECKSUM_SIZE, MAX_DOWNLOAD_SIZE, download_file, verify_checksum}; -use self::installation::ReleaseFiles; -use self::process::{environment_value, spawn_replacement}; -use self::release::{ - ReleaseCache, confirmed_target, detail_from_release, fetch_latest, select_archive, - validate_update_target, version_channel, -}; -use self::state::{ - OperationFileLock, UpdateOperation, UpdateTempDir, finish, operation_id, read_status, - reconcile_installation, recover_interrupted, set_running, -}; -use self::swap::{replace_release_files, rollback_official_plugins_dir, rollback_release}; - -pub use self::release::validate_download_url; - -const APP_BINARY_NAME: &str = "codex-proxy-rs"; -const DEFAULT_GITHUB_API_BASE: &str = "https://api.github.com/repos"; -const DEFAULT_UPDATE_REPOSITORY: &str = "zyycn/codex-proxy-rs"; -const OFFICIAL_PLUGIN_MANIFEST: &str = "plugin-release-manifest.json"; -const MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE: u64 = 256 * 1024; - -type OperationError = SystemOperationError; - -/// 系统更新与重启配置；所有字段只由 Host 解释 -#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] -#[serde(default)] -pub struct SystemUpdateConfig { - pub version: String, - pub git_sha: String, - pub build_time: String, - pub deployment_mode: String, - pub build_type: String, - pub update_repository: Option, - pub github_api_base: String, - pub executable_path: Option, - /// 未显式指定时，由组合根传入 API 实际使用的静态资源目录 - pub web_dist_dir: Option, - pub update_state_file: PathBuf, - pub update_lock_file: PathBuf, - pub update_temp_dir: PathBuf, - pub self_restart_enabled: bool, -} - -impl Default for SystemUpdateConfig { - fn default() -> Self { - let deployment_mode = - environment_value("CPR_DEPLOYMENT_MODE").unwrap_or_else(|| "source".to_owned()); - let executable_path = environment_value("CPR_UPDATE_EXE_PATH") - .map(PathBuf::from) - .or_else(|| { - (deployment_mode == "docker") - .then(|| PathBuf::from("/app/bin").join(APP_BINARY_NAME)) - }); - let update_state_file = environment_value("CPR_UPDATE_STATE_FILE") - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from("update-state.json")); - let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE") - .map(PathBuf::from) - .unwrap_or_else(|| update_state_file.with_file_name("update.lock")); - let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| state::default_temp_dir(&update_state_file)); - Self { - version: option_env!("CPR_VERSION") - .unwrap_or(env!("CARGO_PKG_VERSION")) - .to_owned(), - git_sha: option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_owned(), - build_time: option_env!("CPR_BUILD_TIME") - .unwrap_or("unknown") - .to_owned(), - deployment_mode, - build_type: option_env!("CPR_BUILD_TYPE").unwrap_or("source").to_owned(), - update_repository: Some( - environment_value("CPR_UPDATE_REPOSITORY") - .unwrap_or_else(|| DEFAULT_UPDATE_REPOSITORY.to_owned()), - ), - github_api_base: environment_value("CPR_GITHUB_API_BASE") - .unwrap_or_else(|| DEFAULT_GITHUB_API_BASE.to_owned()), - executable_path, - web_dist_dir: None, - update_state_file, - update_lock_file, - update_temp_dir, - self_restart_enabled: environment_value("CPR_ENABLE_SELF_RESTART").as_deref() - == Some("true"), - } - } -} - -impl SystemUpdateConfig { - pub(crate) fn resolve_and_validate( - &mut self, - source_dir: &Path, - runtime_data_dir: &Path, - asset_directory: &Path, - ) -> Result<(), ConfigError> { - let web_dist_dir = self - .web_dist_dir - .get_or_insert_with(|| asset_directory.to_path_buf()); - if web_dist_dir.as_os_str().is_empty() { - return Err(ConfigError::InvalidField("host.system_update.web_dist_dir")); - } - for path in [self.executable_path.as_mut(), self.web_dist_dir.as_mut()] - .into_iter() - .flatten() - { - if path.is_relative() { - *path = source_dir.join(&*path); - } - } - for path in [ - &mut self.update_state_file, - &mut self.update_lock_file, - &mut self.update_temp_dir, - ] { - if path.is_relative() { - *path = runtime_data_dir.join(&*path); - } - } - if self.version.trim().is_empty() { - return Err(ConfigError::InvalidField("host.system_update.version")); - } - if !matches!( - self.deployment_mode.as_str(), - "source" | "binary" | "docker" - ) { - return Err(ConfigError::InvalidField( - "host.system_update.deployment_mode", - )); - } - Ok(()) - } - - fn update_support_error(&self) -> Option { - if !matches!(self.build_type.as_str(), "release" | "experimental") { - return Some("在线更新需要官方发布构建".to_owned()); - } - let Some(channel) = version_channel(&self.version) else { - return Some("当前版本不符合发行命名规范，无法确定更新通道".to_owned()); - }; - if self.build_type == "experimental" && channel != UpdateChannel::Experimental { - return Some("实验构建必须使用 exp.N 版本，无法在线更新".to_owned()); - } - let Some(repository) = self.update_repository.as_deref() else { - return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned()); - }; - if let Err(error) = release::validate_repository(repository) { - return Some(error.to_string()); - } - if let Err(error) = release::validate_api_base(&self.github_api_base) { - return Some(error); - } - if let Err(error) = self.web_dist_dir() { - return Some(error.to_string()); - } - None - } - - fn release_cache_key(&self) -> String { - format!( - "{}|{}|{}|{}|{}", - self.update_repository.as_deref().unwrap_or_default(), - self.github_api_base, - self.version, - self.deployment_mode, - self.build_type, - ) - } - - pub(crate) fn executable_path(&self) -> Result { - if let Some(path) = &self.executable_path { - return Ok(path.clone()); - } - env::current_exe() - .and_then(fs::canonicalize) - .map_err(|error| internal(format!("failed to resolve executable: {error}"))) - } - - pub(crate) fn official_plugins_dir(&self) -> Result { - let executable = self.executable_path()?; - let parent = executable - .parent() - .filter(|path| !path.as_os_str().is_empty()) - .ok_or_else(|| internal("failed to resolve official plugin release directory"))?; - Ok(parent.join("plugins").join("official")) - } - - pub(crate) fn web_dist_dir(&self) -> Result<&Path, OperationError> { - self.web_dist_dir - .as_deref() - .filter(|path| !path.as_os_str().is_empty()) - .ok_or_else(|| conflict("web assets directory is not configured")) - } -} - -/// gateway-admin 消费的真实进程/文件系统操作实现 -#[derive(Clone)] -pub struct ProcessSystemOperations { - cancellation: CancellationToken, - events: Arc, - config: Arc, - operation_lock: Arc>, - release_cache: Arc, - running_files: Arc>, -} - -impl ProcessSystemOperations { - #[must_use] - pub fn new(cancellation: CancellationToken, mut config: SystemUpdateConfig) -> Self { - // 组合根在接收请求前固定安装路径和发行文件；旧程序重命名后， - // current_exe 可能指向备份，后续更新、回滚和重启不能再反查运行文件 - let running_files = Arc::new(config.executable_path().and_then(|executable| { - config.executable_path = Some(executable); - ReleaseFiles::installed(&config) - })); - Self { - cancellation, - events: Arc::new(UpdateEvents::default()), - config: Arc::new(config), - operation_lock: Arc::new(AsyncMutex::new(())), - release_cache: Arc::new(ReleaseCache::default()), - running_files, - } - } - - fn start_update( - &self, - target_version: Option, - channel: Option, - preflight: Arc, - ) -> Result { - let operation_lock = Arc::clone(&self.operation_lock) - .try_lock_owned() - .map_err(|_| conflict("system operation is already running"))?; - if self.cancellation.is_cancelled() { - return Err(conflict("服务正在关闭")); - } - if let Some(reason) = self.config.update_support_error() { - self.events - .error_terminal(None, Some("preflight"), reason.clone()); - return Err(conflict(reason)); - } - let target = confirmed_target(target_version)?; - let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; - let selected = self.resolve_channel(channel)?; - if let Err(error) = validate_update_target(&self.config.version, &target, selected) { - self.events - .error_terminal(None, Some("preflight"), error.to_string()); - return Err(error); - } - let status = self.reconcile_installation()?; - if status.need_restart { - return Err(conflict("更新已完成，请先重启服务")); - } - let operation_id = operation_id("update"); - let mut operation = UpdateOperation::start( - &self.config.update_state_file, - &operation_id, - &target, - &self.config.version, - file_lock, - )?; - let accepted = SystemOperationAccepted::Update { - operation_id: operation_id.clone(), - deployment_mode: self.config.deployment_mode.clone(), - message: "更新已开始".to_owned(), - target_version: target.clone(), - }; - let service = self.clone(); - // 任务持有互斥锁和落盘守卫，HTTP 断开不会取消更新；Host 关闭仍会收敛终态 - drop(tokio::spawn(async move { - let _operation_lock = operation_lock; - let result = tokio::select! { - biased; - () = service.cancellation.cancelled() => Err(conflict("服务关闭，更新已中断")), - result = service.perform_update_inner(&target, selected, &operation_id, preflight.as_ref()) => result, - }; - if let Err(error) = operation.complete(&result) { - tracing::warn!(error = %error, "系统更新终态落盘失败"); - return; - } - match result { - Ok(_) => service.events.success_terminal( - Some(&operation_id), - Some("done"), - "更新文件已替换，等待服务重启生效", - ), - Err(error) => service.events.error_terminal( - Some(&operation_id), - Some("failed"), - error.to_string(), - ), - } - })); - Ok(accepted) - } - - async fn perform_update_inner( - &self, - target: &str, - channel: UpdateChannel, - operation_id: &str, - preflight: &dyn SystemUpdatePreflight, - ) -> Result { - let repository = self - .config - .update_repository - .as_deref() - .ok_or_else(|| conflict("update repository is not configured"))?; - self.events.info( - Some(operation_id), - Some("release"), - "正在获取最新 Release 信息", - ); - let release = fetch_latest( - &self.config.github_api_base, - repository, - &self.config.version, - channel, - ) - .await? - .ok_or_else(|| conflict("当前发行通道没有可用更新"))?; - let detail = detail_from_release(&self.config, &release, channel); - if detail.latest_version != target { - return Err(conflict("远端最新版本已变化，请重新确认")); - } - if !detail.has_update { - return Err(conflict("当前没有可用更新")); - } - self.events.info( - Some(operation_id), - Some("prepare"), - format!("准备更新到 v{target}"), - ); - self.install_release(&release, target, operation_id, preflight) - .await - } - - async fn install_release( - &self, - release: &release::GitHubRelease, - version: &str, - operation_id: &str, - preflight: &dyn SystemUpdatePreflight, - ) -> Result { - self.events.info( - Some(operation_id), - Some("asset"), - "正在选择匹配当前平台的更新包", - ); - let archive = select_archive(release, version)?; - self.events.info( - Some(operation_id), - Some("asset"), - format!( - "已选择更新包 {} ({})", - archive.name, - download::format_bytes(archive.size) - ), - ); - self.events - .info(Some(operation_id), Some("verify"), "正在校验更新资源"); - if archive.size == 0 || archive.size > MAX_DOWNLOAD_SIZE { - return Err(invalid("release archive size is invalid")); - } - let checksum = release - .assets - .iter() - .find(|asset| asset.name == "checksums.txt") - .ok_or_else(|| upstream("release checksums.txt is required"))?; - if checksum.size == 0 || checksum.size > MAX_CHECKSUM_SIZE { - return Err(invalid("release checksum size is invalid")); - } - self.events - .info(Some(operation_id), Some("prepare"), "正在创建临时更新目录"); - fs::create_dir_all(&self.config.update_temp_dir) - .map_err(|error| internal(format!("failed to prepare update temp dir: {error}")))?; - let temp_root = fs::canonicalize(&self.config.update_temp_dir) - .map_err(|error| internal(format!("failed to resolve update temp dir: {error}")))?; - let temp = UpdateTempDir::create(&temp_root)?; - let archive_path = temp.path().join(&archive.name); - self.events - .info(Some(operation_id), Some("download"), "开始下载更新包"); - download_file( - &archive.browser_download_url, - &archive_path, - archive.size, - &self.config.github_api_base, - operation_id, - &self.events, - ) - .await?; - self.events - .success(Some(operation_id), Some("download"), "更新包下载完成"); - self.events - .info(Some(operation_id), Some("checksum"), "正在校验 checksum"); - verify_checksum( - &archive_path, - &archive.name, - &checksum.browser_download_url, - checksum.size, - &self.config.github_api_base, - ) - .await?; - self.events - .success(Some(operation_id), Some("checksum"), "checksum 校验通过"); - self.events - .info(Some(operation_id), Some("extract"), "正在解压更新包"); - let extracted = extract_release(&archive_path, temp.path())?; - self.events - .success(Some(operation_id), Some("extract"), "更新包解压完成"); - let release_manifest = read_release_manifest(&extracted.official_plugins_dir)?; - let candidate = SystemUpdateCandidate { - target_version: version.to_owned(), - release_manifest: release_manifest.clone(), - }; - self.events.info( - Some(operation_id), - Some("preflight"), - "正在校验目标发行信息", - ); - let plugin_revision = preflight.validate(candidate).await?; - preflight.confirm_revision(plugin_revision).await?; - self.events.success( - Some(operation_id), - Some("preflight"), - "目标发行信息校验通过，插件兼容性将在重启前检查", - ); - self.events - .info(Some(operation_id), Some("replace"), "正在替换应用文件"); - replace_release_files( - &self.config.executable_path()?, - self.config.web_dist_dir()?, - extracted, - )?; - let mut applied = AppliedReleaseGuard::new(&self.config); - let applied_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; - if applied_manifest.as_ref() != release_manifest.as_ref() { - return Err(applied.rollback(conflict("已应用的更新候选与预检候选不一致"))); - } - if let Err(error) = preflight.confirm_revision(plugin_revision).await { - return Err(applied.rollback(error)); - } - let files = ReleaseFiles::installed(&self.config)?; - applied.commit(); - self.events - .success(Some(operation_id), Some("replace"), "应用文件替换完成"); - Ok(files) - } - - fn resolve_channel( - &self, - requested: Option, - ) -> Result { - let channel = requested.unwrap_or_else(|| { - version_channel(&self.config.version).unwrap_or(UpdateChannel::Stable) - }); - if !update_policy(&self.config, channel) - .available_channels - .contains(&channel) - { - return Err(conflict("当前发行版本不支持此更新通道")); - } - Ok(channel) - } - - fn installed_restart_candidate(&self) -> Result, OperationError> { - if self.config.build_type == "source" { - return Ok(None); - } - let release_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; - #[derive(Deserialize)] - struct ReleaseIdentity { - gateway_version: String, - } - let identity: ReleaseIdentity = serde_json::from_slice(&release_manifest) - .map_err(|_| invalid("已安装发行清单缺少目标版本"))?; - Ok(Some(SystemUpdateCandidate { - target_version: identity.gateway_version, - release_manifest, - })) - } - - fn reconcile_installation(&self) -> Result { - let running_files = self.running_files.as_ref().as_ref().map_err(Clone::clone)?; - reconcile_installation(&self.config, running_files) - } -} - -#[async_trait] -impl SystemOperations for ProcessSystemOperations { - async fn version(&self) -> Result { - let detail = self.update_detail(false, None).await?; - Ok(SystemVersion { - version: self.config.version.clone(), - git_sha: self.config.git_sha.clone(), - build_time: self.config.build_time.clone(), - deployment_mode: self.config.deployment_mode.clone(), - update_channel: version_channel(&self.config.version) - .map(|channel| channel.as_str().to_owned()) - .unwrap_or_else(|| "unknown".to_owned()), - latest_version: detail.latest_version, - has_update: detail.has_update, - update_cached: detail.cached, - update_warning: detail.warning, - }) - } - - async fn update_detail( - &self, - refresh: bool, - channel: Option, - ) -> Result { - let channel = self.resolve_channel(channel)?; - match self - .release_cache - .detail(&self.config, refresh, channel) - .await - { - Ok(detail) => Ok(detail), - Err(error) => Ok(base_update_detail( - &self.config, - channel, - self.config.update_support_error(), - Some(error.to_string()), - )), - } - } - - fn update_events(&self) -> SystemUpdateEventStream { - let receiver = self.events.subscribe(); - Box::pin(futures::stream::unfold( - (receiver, false), - |(mut receiver, close)| async move { - if close { - return None; - } - loop { - match receiver.recv().await { - Ok((event, terminal)) => return Some((event, (receiver, terminal))), - Err(broadcast::error::RecvError::Lagged(_)) => continue, - Err(broadcast::error::RecvError::Closed) => return None, - } - } - }, - )) - } - - async fn perform_update( - &self, - target_version: Option, - channel: Option, - preflight: Arc, - ) -> Result { - self.start_update(target_version, channel, preflight) - } - - async fn update_status(&self) -> Result { - if let Ok(_operation) = self.operation_lock.try_lock() { - recover_interrupted( - &self.config.update_state_file, - &self.config.update_lock_file, - )?; - if self.config.update_support_error().is_none() { - match OperationFileLock::acquire(&self.config.update_lock_file) { - Ok(_lock) => return self.reconcile_installation(), - Err(error) if error.kind() == SystemOperationErrorKind::Conflict => {} - Err(error) => return Err(error), - } - } - } - read_status(&self.config.update_state_file, false) - } - - async fn rollback( - &self, - preflight: Arc, - ) -> Result { - let _operation = self - .operation_lock - .try_lock() - .map_err(|_| conflict("system operation is already running"))?; - if let Some(reason) = self.config.update_support_error() { - return Err(conflict(reason)); - } - let operation_id = operation_id("rollback"); - let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; - let target_version = self - .reconcile_installation()? - .previous_version - .ok_or_else(|| conflict("没有可用于回滚的上一版本"))?; - let release_manifest = - read_release_manifest(&rollback_official_plugins_dir(&self.config)?)?; - let plugin_revision = preflight - .validate_rollback(SystemUpdateCandidate { - target_version, - release_manifest: Arc::clone(&release_manifest), - }) - .await?; - preflight.confirm_revision(plugin_revision).await?; - set_running( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Rollback, - None, - &self.config.version, - )?; - let result = match rollback_release(&self.config) { - Ok(()) => { - let mut applied = AppliedReleaseGuard::new(&self.config); - let result = match read_release_manifest(&self.config.official_plugins_dir()?) { - Ok(manifest) if manifest.as_ref() == release_manifest.as_ref() => { - preflight.confirm_revision(plugin_revision).await - } - Ok(_) => Err(conflict("已应用的回滚候选与预检候选不一致")), - Err(error) => Err(error), - }; - match result { - Ok(()) => { - applied.commit(); - Ok(()) - } - Err(error) => Err(applied.rollback(error)), - } - } - Err(error) => Err(error), - }; - finish( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Rollback, - None, - None, - result.as_ref().err().map(ToString::to_string), - )?; - match &result { - Ok(()) => self.events.success_terminal( - Some(&operation_id), - Some("done"), - "previous release restored", - ), - Err(error) => { - self.events - .error_terminal(Some(&operation_id), Some("failed"), error.to_string()) - } - } - drop(file_lock); - result?; - Ok(SystemOperationAccepted::Rollback { - operation_id, - message: "回滚完成，请重启服务。".to_owned(), - need_restart: true, - }) - } - - async fn restart_candidate(&self) -> Result, OperationError> { - let _operation = self - .operation_lock - .try_lock() - .map_err(|_| conflict("system operation is already running"))?; - self.installed_restart_candidate() - } - - async fn restart( - &self, - preflight: Arc, - ) -> Result { - // 与 update/rollback 互斥：更新替换文件期间触发自重启会让新进程 - // 载入半成品产物 - let _operation = self - .operation_lock - .try_lock() - .map_err(|_| conflict("system operation is already running"))?; - if !self.config.self_restart_enabled { - return Err(conflict("self restart is disabled")); - } - let _file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; - let candidate = self.installed_restart_candidate()?; - let files = candidate - .as_ref() - .map(|_| ReleaseFiles::installed(&self.config)) - .transpose()?; - preflight.prepare(candidate).await?; - if let Some(files) = files - && files != ReleaseFiles::installed(&self.config)? - { - return Err(conflict("重启检查期间安装文件已变化，请重新确认")); - } - let message = if self.config.deployment_mode == "docker" { - "已安排进程内重启" - } else { - spawn_replacement(&self.config)?; - "已安排自重启" - }; - let operation_id = operation_id("restart"); - let cancellation = self.cancellation.clone(); - drop(tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(500)).await; - cancellation.cancel(); - })); - Ok(SystemOperationAccepted::Restart { - operation_id, - message: message.to_owned(), - }) - } -} - -fn read_release_manifest(directory: &Path) -> Result, OperationError> { - let path = directory.join(OFFICIAL_PLUGIN_MANIFEST); - let metadata = fs::symlink_metadata(&path) - .map_err(|error| invalid(format!("official plugin manifest is unavailable: {error}")))?; - if !metadata.file_type().is_file() || metadata.file_type().is_symlink() { - return Err(invalid("official plugin manifest is not a regular file")); - } - if metadata.len() == 0 || metadata.len() > MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE { - return Err(invalid("official plugin manifest size is invalid")); - } - fs::read(path) - .map(Arc::<[u8]>::from) - .map_err(|error| invalid(format!("official plugin manifest cannot be read: {error}"))) -} - -fn restore_applied_release(config: &SystemUpdateConfig, cause: OperationError) -> OperationError { - match rollback_release(config) { - Ok(()) => cause, - Err(error) => internal(format!( - "release verification failed and restoring the previous files failed: {error}" - )), - } -} - -struct AppliedReleaseGuard<'a> { - config: &'a SystemUpdateConfig, - armed: bool, -} - -impl<'a> AppliedReleaseGuard<'a> { - const fn new(config: &'a SystemUpdateConfig) -> Self { - Self { - config, - armed: true, - } - } - - const fn commit(&mut self) { - self.armed = false; - } - - fn rollback(mut self, cause: OperationError) -> OperationError { - self.armed = false; - restore_applied_release(self.config, cause) - } -} - -impl Drop for AppliedReleaseGuard<'_> { - fn drop(&mut self) { - if self.armed - && let Err(error) = rollback_release(self.config) - { - tracing::error!(error = %error, "系统发行物临界区取消后的恢复失败"); - } - } -} - -fn update_policy(config: &SystemUpdateConfig, channel: UpdateChannel) -> SystemUpdatePolicy { - use UpdateChannel::{Alpha, Beta, Experimental, Rc, Stable}; - SystemUpdatePolicy { - channel, - available_channels: if version_channel(&config.version) == Some(Experimental) { - vec![Experimental] - } else { - vec![Stable, Rc, Beta, Alpha] - }, - } -} - -fn base_update_detail( - config: &SystemUpdateConfig, - channel: UpdateChannel, - unsupported_reason: Option, - warning: Option, -) -> SystemUpdateDetail { - SystemUpdateDetail { - policy: update_policy(config, channel), - current_version: config.version.clone(), - latest_version: config.version.clone(), - has_update: false, - deployment_mode: config.deployment_mode.clone(), - build_type: config.build_type.clone(), - release_url: None, - notes: None, - cached: false, - update_supported: unsupported_reason.is_none() && warning.is_none(), - unsupported_reason, - warning, - } -} - -pub(crate) struct UpdateEvents { - sender: broadcast::Sender<(SystemUpdateEvent, bool)>, - sequence: AtomicU64, -} - -impl Default for UpdateEvents { - fn default() -> Self { - let (sender, _) = broadcast::channel(256); - Self { - sender, - sequence: AtomicU64::new(0), - } - } -} - -impl UpdateEvents { - fn subscribe(&self) -> broadcast::Receiver<(SystemUpdateEvent, bool)> { - self.sender.subscribe() - } - - fn info(&self, operation_id: Option<&str>, step: Option<&str>, message: impl Into) { - self.emit( - SystemUpdateEventLevel::Info, - operation_id, - step, - message, - None, - false, - ); - } - - fn success(&self, operation_id: Option<&str>, step: Option<&str>, message: impl Into) { - self.emit( - SystemUpdateEventLevel::Success, - operation_id, - step, - message, - None, - false, - ); - } - - fn success_terminal( - &self, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - ) { - self.emit( - SystemUpdateEventLevel::Success, - operation_id, - step, - message, - Some(100), - true, - ); - } - - fn error_terminal( - &self, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - ) { - self.emit( - SystemUpdateEventLevel::Error, - operation_id, - step, - message, - None, - true, - ); - } - - pub(crate) fn emit_progress(&self, operation_id: &str, message: &str, progress: u8) { - self.emit( - SystemUpdateEventLevel::Info, - Some(operation_id), - Some("download"), - message, - Some(progress), - false, - ); - } - - fn emit( - &self, - level: SystemUpdateEventLevel, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - progress_percent: Option, - terminal: bool, - ) { - let now = Utc::now(); - let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); - let event = SystemUpdateEvent { - id: format!( - "update-log-{}-{sequence}", - now.timestamp_nanos_opt() - .unwrap_or_else(|| now.timestamp_millis()) - ), - operation_id: operation_id.map(str::to_owned), - level, - step: step.map(str::to_owned), - message: message.into(), - terminal, - progress_percent, - occurred_at: now, - }; - let _ = self.sender.send((event, terminal)); - } -} - -fn invalid(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Invalid, message) -} - -fn conflict(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Conflict, message) -} - -fn upstream(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Upstream, message) -} - -fn internal(message: impl Into) -> OperationError { - SystemOperationError::new(SystemOperationErrorKind::Internal, message) -} +use error::{OperationError, conflict, internal, invalid, upstream}; diff --git a/backend/crates/gateway-host/src/system_update/operations.rs b/backend/crates/gateway-host/src/system_update/operations.rs new file mode 100644 index 000000000..1e485419c --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/operations.rs @@ -0,0 +1,605 @@ +//! 系统更新、回滚与重启的受理、互斥和文件交换生命周期 + +use std::{fs, path::Path, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use gateway_admin::model::system::{ + SystemOperationAccepted, SystemOperationKind, SystemUpdateChannel as UpdateChannel, + SystemUpdateDetail, SystemUpdateStatus, SystemVersion, +}; +use gateway_admin::ports::system::{ + SystemOperationErrorKind, SystemOperations, SystemRestartPreflight, SystemUpdateCandidate, + SystemUpdateEventStream, SystemUpdatePreflight, +}; +use gateway_core::lifecycle::CancellationToken; +use serde::Deserialize; +use tokio::sync::{Mutex as AsyncMutex, broadcast}; + +use super::archive::extract_release; +use super::download::{ + MAX_CHECKSUM_SIZE, MAX_DOWNLOAD_SIZE, download_file, format_bytes, verify_checksum, +}; +use super::events::UpdateEvents; +use super::installation::{OFFICIAL_PLUGIN_MANIFEST, ReleaseFiles}; +use super::process::spawn_replacement; +use super::release::{ + GitHubRelease, ReleaseCache, base_update_detail, confirmed_target, detail_from_release, + fetch_latest, select_archive, update_policy, validate_update_target, version_channel, +}; +use super::state::{ + OperationFileLock, UpdateOperation, UpdateTempDir, finish, operation_id, read_status, + reconcile_installation, recover_interrupted, set_running, +}; +use super::swap::{replace_release_files, rollback_official_plugins_dir, rollback_release}; +use super::{OperationError, SystemUpdateConfig, conflict, internal, invalid, upstream}; + +const MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE: u64 = 256 * 1024; + +/// gateway-admin 消费的真实进程/文件系统操作实现 +#[derive(Clone)] +pub struct ProcessSystemOperations { + cancellation: CancellationToken, + events: Arc, + config: Arc, + operation_lock: Arc>, + release_cache: Arc, + running_files: Arc>, +} + +impl ProcessSystemOperations { + #[must_use] + pub fn new(cancellation: CancellationToken, mut config: SystemUpdateConfig) -> Self { + // 组合根在接收请求前固定安装路径和发行文件；旧程序重命名后， + // current_exe 可能指向备份，后续更新、回滚和重启不能再反查运行文件 + let running_files = Arc::new(config.executable_path().and_then(|executable| { + config.executable_path = Some(executable); + ReleaseFiles::installed(&config) + })); + Self { + cancellation, + events: Arc::new(UpdateEvents::default()), + config: Arc::new(config), + operation_lock: Arc::new(AsyncMutex::new(())), + release_cache: Arc::new(ReleaseCache::default()), + running_files, + } + } + + fn start_update( + &self, + target_version: Option, + channel: Option, + preflight: Arc, + ) -> Result { + let operation_lock = Arc::clone(&self.operation_lock) + .try_lock_owned() + .map_err(|_| conflict("system operation is already running"))?; + if self.cancellation.is_cancelled() { + return Err(conflict("服务正在关闭")); + } + if let Some(reason) = self.config.update_support_error() { + self.events + .error_terminal(None, Some("preflight"), reason.clone()); + return Err(conflict(reason)); + } + let target = confirmed_target(target_version)?; + let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + let selected = self.resolve_channel(channel)?; + if let Err(error) = validate_update_target(&self.config.version, &target, selected) { + self.events + .error_terminal(None, Some("preflight"), error.to_string()); + return Err(error); + } + let status = self.reconcile_installation()?; + if status.need_restart { + return Err(conflict("更新已完成，请先重启服务")); + } + let operation_id = operation_id("update"); + let mut operation = UpdateOperation::start( + &self.config.update_state_file, + &operation_id, + &target, + &self.config.version, + file_lock, + )?; + let accepted = SystemOperationAccepted::Update { + operation_id: operation_id.clone(), + deployment_mode: self.config.deployment_mode.clone(), + message: "更新已开始".to_owned(), + target_version: target.clone(), + }; + let service = self.clone(); + // 任务持有互斥锁和落盘守卫，HTTP 断开不会取消更新；Host 关闭仍会收敛终态 + drop(tokio::spawn(async move { + let _operation_lock = operation_lock; + let result = tokio::select! { + biased; + () = service.cancellation.cancelled() => Err(conflict("服务关闭，更新已中断")), + result = service.perform_update_inner(&target, selected, &operation_id, preflight.as_ref()) => result, + }; + if let Err(error) = operation.complete(&result) { + tracing::warn!(error = %error, "系统更新终态落盘失败"); + return; + } + match result { + Ok(_) => service.events.success_terminal( + Some(&operation_id), + Some("done"), + "更新文件已替换，等待服务重启生效", + ), + Err(error) => service.events.error_terminal( + Some(&operation_id), + Some("failed"), + error.to_string(), + ), + } + })); + Ok(accepted) + } + + async fn perform_update_inner( + &self, + target: &str, + channel: UpdateChannel, + operation_id: &str, + preflight: &dyn SystemUpdatePreflight, + ) -> Result { + let repository = self + .config + .update_repository + .as_deref() + .ok_or_else(|| conflict("update repository is not configured"))?; + self.events.info( + Some(operation_id), + Some("release"), + "正在获取最新 Release 信息", + ); + let release = fetch_latest( + &self.config.github_api_base, + repository, + &self.config.version, + channel, + ) + .await? + .ok_or_else(|| conflict("当前发行通道没有可用更新"))?; + let detail = detail_from_release(&self.config, &release, channel); + if detail.latest_version != target { + return Err(conflict("远端最新版本已变化，请重新确认")); + } + if !detail.has_update { + return Err(conflict("当前没有可用更新")); + } + self.events.info( + Some(operation_id), + Some("prepare"), + format!("准备更新到 v{target}"), + ); + self.install_release(&release, target, operation_id, preflight) + .await + } + + async fn install_release( + &self, + release: &GitHubRelease, + version: &str, + operation_id: &str, + preflight: &dyn SystemUpdatePreflight, + ) -> Result { + self.events.info( + Some(operation_id), + Some("asset"), + "正在选择匹配当前平台的更新包", + ); + let archive = select_archive(release, version)?; + self.events.info( + Some(operation_id), + Some("asset"), + format!( + "已选择更新包 {} ({})", + archive.name, + format_bytes(archive.size) + ), + ); + self.events + .info(Some(operation_id), Some("verify"), "正在校验更新资源"); + if archive.size == 0 || archive.size > MAX_DOWNLOAD_SIZE { + return Err(invalid("release archive size is invalid")); + } + let checksum = release + .assets + .iter() + .find(|asset| asset.name == "checksums.txt") + .ok_or_else(|| upstream("release checksums.txt is required"))?; + if checksum.size == 0 || checksum.size > MAX_CHECKSUM_SIZE { + return Err(invalid("release checksum size is invalid")); + } + self.events + .info(Some(operation_id), Some("prepare"), "正在创建临时更新目录"); + fs::create_dir_all(&self.config.update_temp_dir) + .map_err(|error| internal(format!("failed to prepare update temp dir: {error}")))?; + let temp_root = fs::canonicalize(&self.config.update_temp_dir) + .map_err(|error| internal(format!("failed to resolve update temp dir: {error}")))?; + let temp = UpdateTempDir::create(&temp_root)?; + let archive_path = temp.path().join(&archive.name); + self.events + .info(Some(operation_id), Some("download"), "开始下载更新包"); + download_file( + &archive.browser_download_url, + &archive_path, + archive.size, + &self.config.github_api_base, + operation_id, + &self.events, + ) + .await?; + self.events + .success(Some(operation_id), Some("download"), "更新包下载完成"); + self.events + .info(Some(operation_id), Some("checksum"), "正在校验 checksum"); + verify_checksum( + &archive_path, + &archive.name, + &checksum.browser_download_url, + checksum.size, + &self.config.github_api_base, + ) + .await?; + self.events + .success(Some(operation_id), Some("checksum"), "checksum 校验通过"); + self.events + .info(Some(operation_id), Some("extract"), "正在解压更新包"); + let extracted = extract_release(&archive_path, temp.path())?; + self.events + .success(Some(operation_id), Some("extract"), "更新包解压完成"); + let release_manifest = read_release_manifest(&extracted.official_plugins_dir)?; + let candidate = SystemUpdateCandidate { + target_version: version.to_owned(), + release_manifest: release_manifest.clone(), + }; + self.events.info( + Some(operation_id), + Some("preflight"), + "正在校验目标发行信息", + ); + let plugin_revision = preflight.validate(candidate).await?; + preflight.confirm_revision(plugin_revision).await?; + self.events.success( + Some(operation_id), + Some("preflight"), + "目标发行信息校验通过，插件兼容性将在重启前检查", + ); + self.events + .info(Some(operation_id), Some("replace"), "正在替换应用文件"); + replace_release_files( + &self.config.executable_path()?, + self.config.web_dist_dir()?, + extracted, + )?; + let mut applied = AppliedReleaseGuard::new(&self.config); + let applied_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; + if applied_manifest.as_ref() != release_manifest.as_ref() { + return Err(applied.rollback(conflict("已应用的更新候选与预检候选不一致"))); + } + if let Err(error) = preflight.confirm_revision(plugin_revision).await { + return Err(applied.rollback(error)); + } + let files = ReleaseFiles::installed(&self.config)?; + applied.commit(); + self.events + .success(Some(operation_id), Some("replace"), "应用文件替换完成"); + Ok(files) + } + + fn resolve_channel( + &self, + requested: Option, + ) -> Result { + let channel = requested.unwrap_or_else(|| { + version_channel(&self.config.version).unwrap_or(UpdateChannel::Stable) + }); + if !update_policy(&self.config, channel) + .available_channels + .contains(&channel) + { + return Err(conflict("当前发行版本不支持此更新通道")); + } + Ok(channel) + } + + fn installed_restart_candidate(&self) -> Result, OperationError> { + if self.config.build_type == "source" { + return Ok(None); + } + let release_manifest = read_release_manifest(&self.config.official_plugins_dir()?)?; + #[derive(Deserialize)] + struct ReleaseIdentity { + gateway_version: String, + } + let identity: ReleaseIdentity = serde_json::from_slice(&release_manifest) + .map_err(|_| invalid("已安装发行清单缺少目标版本"))?; + Ok(Some(SystemUpdateCandidate { + target_version: identity.gateway_version, + release_manifest, + })) + } + + fn reconcile_installation(&self) -> Result { + let running_files = self.running_files.as_ref().as_ref().map_err(Clone::clone)?; + reconcile_installation(&self.config, running_files) + } +} + +#[async_trait] +impl SystemOperations for ProcessSystemOperations { + async fn version(&self) -> Result { + let detail = self.update_detail(false, None).await?; + Ok(SystemVersion { + version: self.config.version.clone(), + git_sha: self.config.git_sha.clone(), + build_time: self.config.build_time.clone(), + deployment_mode: self.config.deployment_mode.clone(), + update_channel: version_channel(&self.config.version) + .map(|channel| channel.as_str().to_owned()) + .unwrap_or_else(|| "unknown".to_owned()), + latest_version: detail.latest_version, + has_update: detail.has_update, + update_cached: detail.cached, + update_warning: detail.warning, + }) + } + + async fn update_detail( + &self, + refresh: bool, + channel: Option, + ) -> Result { + let channel = self.resolve_channel(channel)?; + match self + .release_cache + .detail(&self.config, refresh, channel) + .await + { + Ok(detail) => Ok(detail), + Err(error) => Ok(base_update_detail( + &self.config, + channel, + self.config.update_support_error(), + Some(error.to_string()), + )), + } + } + + fn update_events(&self) -> SystemUpdateEventStream { + let receiver = self.events.subscribe(); + Box::pin(futures::stream::unfold( + (receiver, false), + |(mut receiver, close)| async move { + if close { + return None; + } + loop { + match receiver.recv().await { + Ok((event, terminal)) => return Some((event, (receiver, terminal))), + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => return None, + } + } + }, + )) + } + + async fn perform_update( + &self, + target_version: Option, + channel: Option, + preflight: Arc, + ) -> Result { + self.start_update(target_version, channel, preflight) + } + + async fn update_status(&self) -> Result { + if let Ok(_operation) = self.operation_lock.try_lock() { + recover_interrupted( + &self.config.update_state_file, + &self.config.update_lock_file, + )?; + if self.config.update_support_error().is_none() { + match OperationFileLock::acquire(&self.config.update_lock_file) { + Ok(_lock) => return self.reconcile_installation(), + Err(error) if error.kind() == SystemOperationErrorKind::Conflict => {} + Err(error) => return Err(error), + } + } + } + read_status(&self.config.update_state_file, false) + } + + async fn rollback( + &self, + preflight: Arc, + ) -> Result { + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + if let Some(reason) = self.config.update_support_error() { + return Err(conflict(reason)); + } + let operation_id = operation_id("rollback"); + let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + let target_version = self + .reconcile_installation()? + .previous_version + .ok_or_else(|| conflict("没有可用于回滚的上一版本"))?; + let release_manifest = + read_release_manifest(&rollback_official_plugins_dir(&self.config)?)?; + let plugin_revision = preflight + .validate_rollback(SystemUpdateCandidate { + target_version, + release_manifest: Arc::clone(&release_manifest), + }) + .await?; + preflight.confirm_revision(plugin_revision).await?; + set_running( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Rollback, + None, + &self.config.version, + )?; + let result = match rollback_release(&self.config) { + Ok(()) => { + let mut applied = AppliedReleaseGuard::new(&self.config); + let result = match read_release_manifest(&self.config.official_plugins_dir()?) { + Ok(manifest) if manifest.as_ref() == release_manifest.as_ref() => { + preflight.confirm_revision(plugin_revision).await + } + Ok(_) => Err(conflict("已应用的回滚候选与预检候选不一致")), + Err(error) => Err(error), + }; + match result { + Ok(()) => { + applied.commit(); + Ok(()) + } + Err(error) => Err(applied.rollback(error)), + } + } + Err(error) => Err(error), + }; + finish( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Rollback, + None, + None, + result.as_ref().err().map(ToString::to_string), + )?; + match &result { + Ok(()) => self.events.success_terminal( + Some(&operation_id), + Some("done"), + "previous release restored", + ), + Err(error) => { + self.events + .error_terminal(Some(&operation_id), Some("failed"), error.to_string()) + } + } + drop(file_lock); + result?; + Ok(SystemOperationAccepted::Rollback { + operation_id, + message: "回滚完成，请重启服务。".to_owned(), + need_restart: true, + }) + } + + async fn restart_candidate(&self) -> Result, OperationError> { + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + self.installed_restart_candidate() + } + + async fn restart( + &self, + preflight: Arc, + ) -> Result { + // 与 update/rollback 互斥：更新替换文件期间触发自重启会让新进程 + // 载入半成品产物 + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + if !self.config.self_restart_enabled { + return Err(conflict("self restart is disabled")); + } + let _file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + let candidate = self.installed_restart_candidate()?; + let files = candidate + .as_ref() + .map(|_| ReleaseFiles::installed(&self.config)) + .transpose()?; + preflight.prepare(candidate).await?; + if let Some(files) = files + && files != ReleaseFiles::installed(&self.config)? + { + return Err(conflict("重启检查期间安装文件已变化，请重新确认")); + } + let message = if self.config.deployment_mode == "docker" { + "已安排进程内重启" + } else { + spawn_replacement(&self.config)?; + "已安排自重启" + }; + let operation_id = operation_id("restart"); + let cancellation = self.cancellation.clone(); + drop(tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(500)).await; + cancellation.cancel(); + })); + Ok(SystemOperationAccepted::Restart { + operation_id, + message: message.to_owned(), + }) + } +} + +fn read_release_manifest(directory: &Path) -> Result, OperationError> { + let path = directory.join(OFFICIAL_PLUGIN_MANIFEST); + let metadata = fs::symlink_metadata(&path) + .map_err(|error| invalid(format!("official plugin manifest is unavailable: {error}")))?; + if !metadata.file_type().is_file() || metadata.file_type().is_symlink() { + return Err(invalid("official plugin manifest is not a regular file")); + } + if metadata.len() == 0 || metadata.len() > MAX_OFFICIAL_PLUGIN_MANIFEST_SIZE { + return Err(invalid("official plugin manifest size is invalid")); + } + fs::read(path) + .map(Arc::<[u8]>::from) + .map_err(|error| invalid(format!("official plugin manifest cannot be read: {error}"))) +} + +fn restore_applied_release(config: &SystemUpdateConfig, cause: OperationError) -> OperationError { + match rollback_release(config) { + Ok(()) => cause, + Err(error) => internal(format!( + "release verification failed and restoring the previous files failed: {error}" + )), + } +} + +struct AppliedReleaseGuard<'a> { + config: &'a SystemUpdateConfig, + armed: bool, +} + +impl<'a> AppliedReleaseGuard<'a> { + const fn new(config: &'a SystemUpdateConfig) -> Self { + Self { + config, + armed: true, + } + } + + const fn commit(&mut self) { + self.armed = false; + } + + fn rollback(mut self, cause: OperationError) -> OperationError { + self.armed = false; + restore_applied_release(self.config, cause) + } +} + +impl Drop for AppliedReleaseGuard<'_> { + fn drop(&mut self) { + if self.armed + && let Err(error) = rollback_release(self.config) + { + tracing::error!(error = %error, "系统发行物临界区取消后的恢复失败"); + } + } +} diff --git a/backend/crates/gateway-host/src/system_update/process.rs b/backend/crates/gateway-host/src/system_update/process.rs index 4f6d8f6c5..d95fe6fd9 100644 --- a/backend/crates/gateway-host/src/system_update/process.rs +++ b/backend/crates/gateway-host/src/system_update/process.rs @@ -4,6 +4,7 @@ use std::env; use std::fs; use std::process::{Command, Stdio}; +use super::config::environment_value; use super::{OperationError, SystemUpdateConfig, internal}; const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS"; @@ -59,10 +60,3 @@ pub(crate) fn spawn_replacement(config: &SystemUpdateConfig) -> Result<(), Opera .map(|_| ()) .map_err(|error| internal(format!("failed to schedule replacement process: {error}"))) } - -pub(crate) fn environment_value(key: &str) -> Option { - env::var(key) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} diff --git a/backend/crates/gateway-host/src/system_update/release.rs b/backend/crates/gateway-host/src/system_update/release.rs index 69f9114e5..0544d4aef 100644 --- a/backend/crates/gateway-host/src/system_update/release.rs +++ b/backend/crates/gateway-host/src/system_update/release.rs @@ -4,7 +4,9 @@ use std::env; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; -use gateway_admin::model::system::{SystemUpdateChannel as UpdateChannel, SystemUpdateDetail}; +use gateway_admin::model::system::{ + SystemUpdateChannel as UpdateChannel, SystemUpdateDetail, SystemUpdatePolicy, +}; use serde::Deserialize; use tokio::sync::Mutex; @@ -53,12 +55,7 @@ impl ReleaseCache { channel: UpdateChannel, ) -> Result { if let Some(reason) = config.update_support_error() { - return Ok(super::base_update_detail( - config, - channel, - Some(reason), - None, - )); + return Ok(base_update_detail(config, channel, Some(reason), None)); } let repository = config .update_repository @@ -81,14 +78,7 @@ impl ReleaseCache { { Ok(release) => { let detail = release.as_ref().map_or_else( - || { - super::base_update_detail( - config, - channel, - config.update_support_error(), - None, - ) - }, + || base_update_detail(config, channel, config.update_support_error(), None), |release| detail_from_release(config, release, channel), ); let mut entry = self.entry.lock().await; @@ -218,10 +208,10 @@ pub(crate) fn detail_from_release( if unsupported_reason.is_some() || (!has_update && latest_version != normalize_version(&config.version)) { - return super::base_update_detail(config, channel, unsupported_reason, None); + return base_update_detail(config, channel, unsupported_reason, None); } SystemUpdateDetail { - policy: super::update_policy(config, channel), + policy: update_policy(config, channel), current_version: config.version.clone(), latest_version, has_update, @@ -488,3 +478,78 @@ fn url_host_is_loopback(url: &reqwest::Url) -> bool { .is_ok_and(|address| address.is_loopback()) }) } + +impl SystemUpdateConfig { + pub(super) fn update_support_error(&self) -> Option { + if !matches!(self.build_type.as_str(), "release" | "experimental") { + return Some("在线更新需要官方发布构建".to_owned()); + } + let Some(channel) = version_channel(&self.version) else { + return Some("当前版本不符合发行命名规范，无法确定更新通道".to_owned()); + }; + if self.build_type == "experimental" && channel != UpdateChannel::Experimental { + return Some("实验构建必须使用 exp.N 版本，无法在线更新".to_owned()); + } + let Some(repository) = self.update_repository.as_deref() else { + return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned()); + }; + if let Err(error) = validate_repository(repository) { + return Some(error.to_string()); + } + if let Err(error) = validate_api_base(&self.github_api_base) { + return Some(error); + } + if let Err(error) = self.web_dist_dir() { + return Some(error.to_string()); + } + None + } + + fn release_cache_key(&self) -> String { + format!( + "{}|{}|{}|{}|{}", + self.update_repository.as_deref().unwrap_or_default(), + self.github_api_base, + self.version, + self.deployment_mode, + self.build_type, + ) + } +} + +pub(super) fn update_policy( + config: &SystemUpdateConfig, + channel: UpdateChannel, +) -> SystemUpdatePolicy { + use UpdateChannel::{Alpha, Beta, Experimental, Rc, Stable}; + SystemUpdatePolicy { + channel, + available_channels: if version_channel(&config.version) == Some(Experimental) { + vec![Experimental] + } else { + vec![Stable, Rc, Beta, Alpha] + }, + } +} + +pub(super) fn base_update_detail( + config: &SystemUpdateConfig, + channel: UpdateChannel, + unsupported_reason: Option, + warning: Option, +) -> SystemUpdateDetail { + SystemUpdateDetail { + policy: update_policy(config, channel), + current_version: config.version.clone(), + latest_version: config.version.clone(), + has_update: false, + deployment_mode: config.deployment_mode.clone(), + build_type: config.build_type.clone(), + release_url: None, + notes: None, + cached: false, + update_supported: unsupported_reason.is_none() && warning.is_none(), + unsupported_reason, + warning, + } +} diff --git a/backend/crates/gateway-host/src/system_update/state.rs b/backend/crates/gateway-host/src/system_update/state.rs index 2daa64721..274208a78 100644 --- a/backend/crates/gateway-host/src/system_update/state.rs +++ b/backend/crates/gateway-host/src/system_update/state.rs @@ -10,6 +10,7 @@ use chrono::{DateTime, Utc}; use gateway_admin::model::system::{ SystemOperationKind, SystemOperationState, SystemOperationStatus, SystemUpdateStatus, }; +use gateway_admin::ports::system::SystemOperationErrorKind; use serde::{Deserialize, Serialize}; use super::installation::ReleaseFiles; @@ -235,7 +236,7 @@ pub(crate) fn recover_interrupted(path: &Path, lock_path: &Path) -> Result<(), O } let _lock = match OperationFileLock::acquire(lock_path) { Ok(lock) => lock, - Err(error) if error.kind() == super::SystemOperationErrorKind::Conflict => return Ok(()), + Err(error) if error.kind() == SystemOperationErrorKind::Conflict => return Ok(()), Err(error) => return Err(error), }; if let (Some(operation_id), Some(kind)) = (state.operation.operation_id, state.operation.kind) { @@ -394,13 +395,6 @@ pub(crate) fn reconcile_installation( read_status(path, need_restart) } -pub(crate) fn default_temp_dir(state_file: &Path) -> PathBuf { - state_file - .parent() - .map(|parent| parent.join("update-tmp")) - .unwrap_or_else(|| std::env::temp_dir().join("codex-proxy-rs-update")) -} - fn stale_lock(path: &Path) -> Result { let modified = fs::metadata(path) .and_then(|metadata| metadata.modified()) diff --git a/backend/crates/gateway-host/src/workers.rs b/backend/crates/gateway-host/src/workers.rs index 989f5f785..889bb5045 100644 --- a/backend/crates/gateway-host/src/workers.rs +++ b/backend/crates/gateway-host/src/workers.rs @@ -117,10 +117,13 @@ impl WorkerSupervisor { /// 不会被二次 poll，逾期的在 abort 后等待其真正终止 pub async fn shutdown(&self, timeout: Duration) { self.cancellation.cancel(); - let handles = std::mem::take(&mut *lock_unpoisoned(&self.handles)); + let mut shutdown = WorkerShutdown { + handles: std::mem::take(&mut *lock_unpoisoned(&self.handles)), + health: &self.health, + }; let deadline = tokio::time::Instant::now() + timeout; - for (_, mut handle) in handles { - if tokio::time::timeout_at(deadline, &mut handle) + for (_, handle) in &mut shutdown.handles { + if tokio::time::timeout_at(deadline, &mut *handle) .await .is_err() { @@ -128,6 +131,21 @@ impl WorkerSupervisor { let _ = handle.await; } } + } +} + +struct WorkerShutdown<'a> { + handles: Vec<(WorkerId, JoinHandle<()>)>, + health: &'a WorkerHealthRegistry, +} + +impl Drop for WorkerShutdown<'_> { + fn drop(&mut self) { + // shutdown future 可在任一 join 点被丢弃；句柄始终留在同一 owner 内， + // 防止当前等待或尚未轮到的任务因 JoinHandle 析构而脱离监督 + for (_, handle) in &self.handles { + handle.abort(); + } self.health.stop_all(); } } diff --git a/backend/crates/gateway-host/tests/command_line.rs b/backend/crates/gateway-host/tests/command_line.rs new file mode 100644 index 000000000..ce776141e --- /dev/null +++ b/backend/crates/gateway-host/tests/command_line.rs @@ -0,0 +1,85 @@ +//! 验证 CLI 不进入 HTTP serve 时的信号取消与 Host 析构 + +use crate::support::host::{configuration, run_in_child}; + +#[test] +fn dropping_command_line_host_cancels_issued_capabilities() { + run_in_child( + "command_line::dropping_command_line_host_cancels_issued_capabilities", + async { + let directory = tempfile::tempdir().unwrap(); + let host = gateway_host::initialize_command_line(configuration(directory.path())) + .await + .unwrap(); + let command = host.cancellation().child_token(); + let connection = host.connection_lifecycle().cancellation(); + + drop(host); + + assert!(command.is_cancelled()); + assert!(connection.is_cancelled()); + }, + ); +} + +#[cfg(unix)] +#[test] +fn sigterm_cancels_command_line_without_http_serve() { + use std::io::{BufRead as _, Write as _}; + use std::process::{Command, Stdio}; + + use crate::support::host::{CHILD_PROCESS_ENV, child_command}; + + const NAME: &str = "command_line::sigterm_cancels_command_line_without_http_serve"; + const READY: &str = "host command signal listener ready"; + + if std::env::var(CHILD_PROCESS_ENV).as_deref() == Ok(NAME) { + run_in_child(NAME, async { + let directory = tempfile::tempdir().unwrap(); + let host = gateway_host::initialize_command_line(configuration(directory.path())) + .await + .unwrap(); + let command = host.cancellation().child_token(); + // SignalGuard 必须先 poll，父进程才可发送信号，避免测试命中默认 SIGTERM 行为 + tokio::task::yield_now().await; + println!("{READY}"); + std::io::stdout().flush().unwrap(); + command.cancelled().await; + assert!(host.connection_lifecycle().cancellation().is_cancelled()); + assert!(host.worker_health().snapshot().is_empty()); + }); + return; + } + + let mut child = child_command(NAME) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let mut output = std::io::BufReader::new(child.stdout.take().unwrap()); + let mut line = String::new(); + loop { + assert_ne!( + output.read_line(&mut line).unwrap(), + 0, + "child was not ready" + ); + if line.contains(READY) { + break; + } + line.clear(); + } + assert!( + Command::new("kill") + .args(["-TERM", &child.id().to_string()]) + .status() + .unwrap() + .success() + ); + let output = child.wait_with_output().unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); +} diff --git a/backend/crates/gateway-host/tests/main.rs b/backend/crates/gateway-host/tests/main.rs index 712529fd7..c71ba8584 100644 --- a/backend/crates/gateway-host/tests/main.rs +++ b/backend/crates/gateway-host/tests/main.rs @@ -1,6 +1,7 @@ //! 宿主配置、网络、进程、日志与后台任务的测试入口 mod client_distribution; +mod command_line; mod config; mod logging; mod official_plugins; diff --git a/backend/crates/gateway-host/tests/serve.rs b/backend/crates/gateway-host/tests/serve.rs index 170d63b30..e6fda0b2a 100644 --- a/backend/crates/gateway-host/tests/serve.rs +++ b/backend/crates/gateway-host/tests/serve.rs @@ -1,12 +1,25 @@ //! 验证宿主服务关闭时的连接排空、超时与并发唤醒 use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; +use axum::{Router, routing::get}; +use futures::future::BoxFuture; +use gateway_core::diagnostics::{OperationalDiagnostics, OperationalFailure}; use gateway_core::lifecycle::CancellationToken; use gateway_core::lifecycle::ConnectionLifecycle as _; +use gateway_core::task::{ + DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerId, WorkerKind, + WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest, + WorkerRegistration, WorkerRunnable, WorkerTaskError, +}; use gateway_host::HostBundle; use gateway_host::serve::{ConnectionTracker, bind_listener}; +use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; +use tokio::sync::{Semaphore, mpsc}; + +use crate::support::host::{configuration, run_in_child}; #[test] fn host_bundle_serve_is_a_consuming_process_entrypoint() { @@ -95,3 +108,354 @@ async fn bind_listener_should_retry_until_previous_listener_releases_port() { address ); } + +#[test] +fn http_drain_keeps_writer_running_until_the_request_finishes() { + run_in_child( + "serve::http_drain_keeps_writer_running_until_the_request_finishes", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let connections = host.connection_lifecycle(); + let probe = start_workers(&host, false); + probe.started.cancelled().await; + let (router, entered, finish_request) = request_router(&host, &probe); + let server = tokio::spawn(host.serve(router)); + let mut client = connect_request(address).await; + entered.cancelled().await; + + cancellation.cancel(); + wait_for_draining(connections.as_ref()).await; + assert!(connections.try_register().is_err()); + assert!( + tokio::time::timeout(Duration::from_millis(50), probe.cancelled.cancelled()) + .await + .is_err(), + "writer cancelled before the in-flight request finished" + ); + finish_request.add_permits(1); + let mut response = String::new(); + client.read_to_string(&mut response).await.unwrap(); + server.await.unwrap().unwrap(); + + assert!(response.starts_with("HTTP/1.1 200 OK")); + assert_eq!(*probe.writes.lock().unwrap(), [1]); + assert!(probe.finished.is_cancelled()); + }, + ); +} + +#[test] +fn drain_timeout_starts_a_separate_worker_shutdown_phase() { + run_in_child( + "serve::drain_timeout_starts_a_separate_worker_shutdown_phase", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let connections = host.connection_lifecycle(); + let probe = start_workers(&host, true); + probe.started.cancelled().await; + let (router, entered, finish_request) = request_router(&host, &probe); + let server = tokio::spawn(host.serve(router)); + let client = connect_request(address).await; + entered.cancelled().await; + + let started = tokio::time::Instant::now(); + cancellation.cancel(); + wait_for_draining(connections.as_ref()).await; + probe.cancelled.cancelled().await; + assert!(started.elapsed() >= Duration::from_secs(1)); + assert!(!server.is_finished(), "worker cleanup was not joined"); + assert!(!probe.finished.is_cancelled()); + probe.finish_shutdown.add_permits(1); + server.await.unwrap().unwrap(); + assert!(probe.finished.is_cancelled()); + + // 超时后真实连接可以仍在 axum 任务中；释放测试端，避免遗留挂起 handler + drop(client); + finish_request.add_permits(1); + }, + ); +} + +#[test] +fn bind_failure_cancels_and_joins_workers_before_returning() { + run_in_child( + "serve::bind_failure_cancels_and_joins_workers_before_returning", + async { + let directory = tempfile::tempdir().unwrap(); + let mut config = configuration(directory.path()); + config.listen.host = "invalid address".to_owned(); + let host = gateway_host::initialize(config).await.unwrap(); + let cancellation = host.cancellation(); + let probe = start_workers(&host, false); + probe.started.cancelled().await; + + let result = host.serve(Router::new()).await; + + assert!(matches!( + result, + Err(gateway_host::HostError::Serve( + gateway_host::serve::ServeError::Bind(_) + )) + )); + assert!(cancellation.is_cancelled()); + assert!(probe.finished.is_cancelled()); + }, + ); +} + +#[test] +fn dropping_serve_cancels_public_operations_and_aborts_workers() { + run_in_child( + "serve::dropping_serve_cancels_public_operations_and_aborts_workers", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let probe = start_workers(&host, true); + probe.started.cancelled().await; + let server = tokio::spawn(host.serve(Router::new())); + let client = connect_request(address).await; + + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + probe.dropped.cancelled().await; + + assert!(cancellation.is_cancelled()); + assert!(!probe.finished.is_cancelled()); + drop(client); + }, + ); +} + +#[test] +fn dropping_serve_during_worker_shutdown_aborts_the_joined_worker() { + run_in_child( + "serve::dropping_serve_during_worker_shutdown_aborts_the_joined_worker", + async { + let directory = tempfile::tempdir().unwrap(); + let (host, address) = listening_host(directory.path()).await; + let cancellation = host.cancellation(); + let probe = start_workers(&host, true); + probe.started.cancelled().await; + let server = tokio::spawn(host.serve(Router::new())); + let client = connect_request(address).await; + + cancellation.cancel(); + probe.cancelled.cancelled().await; + assert!(!probe.finished.is_cancelled()); + assert!( + !server.is_finished(), + "worker shutdown must still be waiting" + ); + // Worker 已观察独立的取消信号，serve 正持有 shutdown 的 join future + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(1), probe.dropped.cancelled()) + .await + .expect("aborting shutdown must not detach the worker"); + + assert!(!probe.finished.is_cancelled()); + drop(client); + }, + ); +} + +async fn listening_host(directory: &std::path::Path) -> (HostBundle, std::net::SocketAddr) { + let reservation = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let address = reservation.local_addr().unwrap(); + let mut config = configuration(directory); + config.listen.port = address.port(); + let host = gateway_host::initialize(config).await.unwrap(); + drop(reservation); + (host, address) +} + +async fn connect_request(address: std::net::SocketAddr) -> tokio::net::TcpStream { + let mut client = loop { + match tokio::net::TcpStream::connect(address).await { + Ok(client) => break client, + Err(error) if error.kind() == std::io::ErrorKind::ConnectionRefused => { + tokio::time::sleep(Duration::from_millis(5)).await; + } + Err(error) => panic!("connect host: {error}"), + } + }; + client + .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n") + .await + .unwrap(); + client +} + +async fn wait_for_draining(connections: &dyn gateway_core::lifecycle::ConnectionLifecycle) { + while !connections.is_draining() { + tokio::task::yield_now().await; + } +} + +fn request_router( + host: &HostBundle, + probe: &WriterProbe, +) -> (Router, CancellationToken, Arc) { + let entered = CancellationToken::new(); + let finish = Arc::new(Semaphore::new(0)); + let connections = host.connection_lifecycle(); + let sender = probe.sender.clone(); + let router = Router::new().route( + "/", + get({ + let entered = entered.clone(); + let finish = Arc::clone(&finish); + move || { + let entered = entered.clone(); + let finish = Arc::clone(&finish); + let connections = Arc::clone(&connections); + let sender = sender.clone(); + async move { + let _connection = connections.try_register().unwrap(); + entered.cancel(); + finish.acquire().await.unwrap().forget(); + // 对应请求结束后写入终态的时机；drain 超时后的写入允许被拒绝 + if sender.send(1).is_ok() { + "recorded" + } else { + "writer stopped" + } + } + } + }), + ); + (router, entered, finish) +} + +struct WriterProbe { + sender: mpsc::UnboundedSender, + writes: Arc>>, + started: CancellationToken, + cancelled: CancellationToken, + finish_shutdown: Arc, + finished: CancellationToken, + dropped: CancellationToken, +} + +struct WriterTask { + received: tokio::sync::Mutex>, + probe: Arc, +} + +struct TaskLifetime(CancellationToken); + +impl Drop for TaskLifetime { + fn drop(&mut self) { + self.0.cancel(); + } +} + +impl DaemonTask for WriterTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let _lifetime = TaskLifetime(self.probe.dropped.clone()); + let mut received = self.received.lock().await; + self.probe.started.cancel(); + loop { + tokio::select! { + biased; + () = cancellation.cancelled() => break, + value = received.recv() => { + self.probe.writes.lock().unwrap().push(value.unwrap()); + } + } + } + self.probe.cancelled.cancel(); + received.close(); + while let Some(value) = received.recv().await { + self.probe.writes.lock().unwrap().push(value); + } + self.probe.finish_shutdown.acquire().await.unwrap().forget(); + self.probe.finished.cancel(); + Ok(()) + }) + } +} + +struct IdleTask; + +impl DaemonTask for IdleTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + cancellation.cancelled().await; + Ok(()) + }) + } +} + +struct UnusedLease; + +impl WorkerLeaderLeasePort for UnusedLease { + fn try_acquire( + &self, + _: WorkerLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Err(WorkerLeaseError::safe("daemon must not acquire a lease")) }) + } +} + +struct Diagnostics; + +#[async_trait::async_trait] +impl OperationalDiagnostics for Diagnostics { + async fn record_failure( + &self, + _: OperationalFailure, + ) -> Result<(), gateway_core::error::StoreError> { + Ok(()) + } +} + +fn start_workers(host: &HostBundle, hold_shutdown: bool) -> Arc { + let (sender, received) = mpsc::unbounded_channel(); + let probe = Arc::new(WriterProbe { + sender, + writes: Arc::new(Mutex::new(Vec::new())), + started: CancellationToken::new(), + cancelled: CancellationToken::new(), + finish_shutdown: Arc::new(Semaphore::new(usize::from(!hold_shutdown))), + finished: CancellationToken::new(), + dropped: CancellationToken::new(), + }); + let mut writer = Some(WriterTask { + received: tokio::sync::Mutex::new(received), + probe: Arc::clone(&probe), + }); + let plan = WorkerKind::ALL + .into_iter() + .map(|kind| { + let task: Box = if kind == WorkerKind::OpsFlush { + Box::new(writer.take().unwrap()) + } else { + Box::new(IdleTask) + }; + WorkerContribution::Registration( + WorkerRegistration::try_new( + WorkerId::try_new(kind, "host_test").unwrap(), + WorkerRunnable::Daemon { + restart: DaemonRestartPolicy::try_new( + Duration::from_secs(1), + Duration::from_secs(2), + ) + .unwrap(), + task, + }, + ) + .unwrap(), + ) + }) + .collect(); + host.start_workers(plan, Arc::new(UnusedLease), Arc::new(Diagnostics)) + .unwrap(); + probe +} diff --git a/backend/crates/gateway-host/tests/support/host.rs b/backend/crates/gateway-host/tests/support/host.rs new file mode 100644 index 000000000..7bf4ae22c --- /dev/null +++ b/backend/crates/gateway-host/tests/support/host.rs @@ -0,0 +1,65 @@ +//! 真实 Host 生命周期测试的配置与独立进程，隔离全局日志订阅器和信号监听 + +use std::{future::Future, path::Path, process::Command, time::Duration}; + +use gateway_host::config::{FileLoggingConfig, HostConfig, ListenConfig, LoggingConfig}; + +pub const CHILD_PROCESS_ENV: &str = "CPR_HOST_LIFECYCLE_TEST_CHILD"; + +pub fn run_in_child(name: &str, scenario: impl Future) { + if std::env::var(CHILD_PROCESS_ENV).as_deref() == Ok(name) { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + tokio::time::timeout(Duration::from_secs(10), scenario) + .await + .expect("host lifecycle scenario timed out"); + }); + return; + } + let output = child_command(name).output().unwrap(); + assert!( + output.status.success(), + "child failed: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} + +pub fn child_command(name: &str) -> Command { + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", name, "--nocapture"]) + .env(CHILD_PROCESS_ENV, name) + .env("RUST_LOG", "off"); + command +} + +pub fn configuration(directory: &Path) -> HostConfig { + HostConfig { + timezone: Default::default(), + listen: ListenConfig { + host: "127.0.0.1".to_owned(), + port: 0, + }, + runtime_data_dir: directory.to_path_buf(), + logging: LoggingConfig { + level: "off".to_owned(), + stdout: false, + file: FileLoggingConfig { + enabled: true, + directory: directory.join("logs"), + retention_days: 1, + max_file_size_mb: 1, + }, + oauth_recovery: false, + request_dump: false, + request_dump_retention_days: 1, + }, + system_update: Default::default(), + drain_timeout_seconds: 1, + worker_shutdown_timeout_seconds: 1, + } +} diff --git a/backend/crates/gateway-host/tests/support/mod.rs b/backend/crates/gateway-host/tests/support/mod.rs index 205a7c466..dd1f2b0b0 100644 --- a/backend/crates/gateway-host/tests/support/mod.rs +++ b/backend/crates/gateway-host/tests/support/mod.rs @@ -1,3 +1,4 @@ //! 宿主测试共用的网络与服务辅助入口 +pub mod host; pub mod network; diff --git a/backend/crates/gateway-host/tests/system_update/download.rs b/backend/crates/gateway-host/tests/system_update/download.rs new file mode 100644 index 000000000..16c74c30b --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/download.rs @@ -0,0 +1,14 @@ +//! 更新下载来源信任边界测试 + +use super::*; + +#[test] +fn update_should_trust_github_release_asset_redirect_host() { + assert!( + validate_download_url( + "https://release-assets.githubusercontent.com/github-production-release-asset/archive", + "https://api.github.com/repos", + ) + .is_ok() + ); +} diff --git a/backend/crates/gateway-host/tests/system_update/events.rs b/backend/crates/gateway-host/tests/system_update/events.rs new file mode 100644 index 000000000..f232a2b4f --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/events.rs @@ -0,0 +1,154 @@ +//! 系统更新事件的订阅、阶段进度与终态可见性测试 + +use super::*; + +#[tokio::test] +async fn update_events_should_close_after_terminal_update_log() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let mut events = service.update_events(); + let _ = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await; + let first = events.next().await.expect("terminal event"); + + assert_eq!(first.level, SystemUpdateEventLevel::Error); + assert!(events.next().await.is_none()); +} + +#[tokio::test] +async fn update_events_should_open_authenticated_sse_stream() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let mut stream = service.update_events(); + let _ = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await; + + assert!(stream.next().await.is_some()); +} + +#[tokio::test] +async fn update_events_should_preserve_complete_release_stage_sequence() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + let events = service.update_events(); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + let steps = events + .map(|event| event.step.expect("release update step")) + .collect::>() + .await; + + assert_eq!( + steps, + [ + "release", + "prepare", + "asset", + "asset", + "verify", + "prepare", + "download", + "download", + "download", + "checksum", + "checksum", + "extract", + "extract", + "preflight", + "preflight", + "replace", + "replace", + "done", + ] + ); +} + +#[tokio::test] +async fn update_events_should_report_download_bytes_and_percent() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + let events = service.update_events(); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + let progress = events + .filter_map( + |event| async move { event.progress_percent.map(|value| (value, event.message)) }, + ) + .collect::>() + .await; + + assert!(progress.iter().any(|(percent, message)| { + *percent == 100 && message.starts_with("已下载 ") && message.ends_with(" (100%)") + })); +} + +#[tokio::test] +async fn terminal_event_should_only_be_visible_after_status_is_persisted() { + let upstream = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &upstream, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Mismatch, + ) + .await; + let service = fixture.service(&upstream); + let mut events = service.update_events(); + let SystemOperationAccepted::Update { operation_id, .. } = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("accepted") + else { + panic!("update") + }; + tokio::time::timeout(Duration::from_secs(5), async { + while let Some(event) = events.next().await { + assert_eq!(event.operation_id.as_deref(), Some(operation_id.as_str())); + if event.terminal { + let status = service.update_status().await.expect("terminal status"); + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert!(status.operation.finished_at.is_some()); + return; + } + } + panic!("missing terminal event"); + }) + .await + .expect("terminal event"); +} diff --git a/backend/crates/gateway-host/tests/system_update/mod.rs b/backend/crates/gateway-host/tests/system_update/mod.rs index b3cd7142c..197117bb1 100644 --- a/backend/crates/gateway-host/tests/system_update/mod.rs +++ b/backend/crates/gateway-host/tests/system_update/mod.rs @@ -1,7 +1,12 @@ -//! 系统更新测试入口，以及重启准备和操作互斥测试 +//! 系统更新测试入口与发行包、前置校验和本地安装夹具 mod channels; +mod download; +mod events; mod installation; +mod operations; +mod release; +mod state; use std::fs; use std::path::PathBuf; @@ -164,1789 +169,6 @@ impl SystemUpdatePreflight for BlockingRollbackPreflight { } } -#[tokio::test] -async fn restart_should_not_shutdown_when_replacement_spawn_fails() { - let fixture = Fixture::new(); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - config.executable_path = Some(fixture.root.path().join("missing")); - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - - assert!( - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .is_err() - ); - assert!(!shutdown.is_cancelled()); -} - -#[tokio::test] -async fn restart_should_request_process_restart_inside_docker() { - let fixture = Fixture::new(); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - config.deployment_mode = "docker".to_owned(); - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("restart accepted"); - - tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) - .await - .expect("shutdown requested"); -} - -#[tokio::test] -async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { - let fixture = Fixture::new(); - let replacement_marker = fixture.root.path().join("replacement-ran"); - fixture.write_executable("#!/bin/sh\n: > \"${0%/*}/replacement-ran\"\n"); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - - assert_eq!( - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("replacement scheduled") - .kind(), - SystemOperationKind::Restart - ); - tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) - .await - .expect("shutdown requested after spawn"); - tokio::time::timeout(Duration::from_secs(2), async { - while !replacement_marker.is_file() { - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("replacement executed before fixture cleanup"); -} - -#[cfg(target_os = "linux")] -#[tokio::test] -async fn restart_should_use_startup_path_after_running_executable_is_renamed() { - const CHILD_ENV: &str = "CPR_TEST_RESTART_RENAMED_EXECUTABLE_CHILD"; - let executable = std::env::current_exe().expect("test executable"); - if std::env::var_os(CHILD_ENV).is_none() { - // 用例只移动目录项，不修改旧程序内容；同盘硬链接避免复制整个调试二进制 - let directory = - tempfile::tempdir_in(executable.parent().unwrap()).expect("isolated test directory"); - let child = directory.path().join("restart-test"); - fs::hard_link(&executable, &child).expect("link isolated test executable"); - let output = std::process::Command::new(child) - .args([ - "--exact", - "system_update::restart_should_use_startup_path_after_running_executable_is_renamed", - ]) - .env(CHILD_ENV, "1") - .env("CPR_RESTART_DELAY_MS", "1200") - .output() - .expect("isolated restart test"); - assert!( - output.status.success(), - "stdout: {}\nstderr: {}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - return; - } - if executable.file_name() == Some(std::ffi::OsStr::new("codex-proxy-rs.backup")) { - // 误启动旧程序时留下可观察结果并退出，避免旧副本再次发起重启 - fs::write( - executable - .parent() - .expect("installation directory") - .join("replacement-version"), - "old", - ) - .expect("old replacement marker"); - return; - } - - let fixture = Fixture::new_in(executable.parent().unwrap()); - fs::rename(executable, fixture.executable()).expect("move isolated running executable"); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.executable_path = None; - config.self_restart_enabled = true; - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let backup = fixture.root.path().join("codex-proxy-rs.backup"); - fs::rename(fixture.executable(), &backup).expect("back up running executable"); - fixture.write_executable("#!/bin/sh\nprintf new > \"${0%/*}/replacement-version\"\n"); - assert_eq!(std::env::current_exe().expect("renamed executable"), backup); - - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("restart accepted"); - let marker = fixture.root.path().join("replacement-version"); - tokio::time::timeout(Duration::from_secs(5), async { - while !marker.is_file() { - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("replacement executed"); - assert_eq!( - fs::read_to_string(marker).expect("replacement version"), - "new" - ); -} - -#[tokio::test] -async fn restart_should_conflict_while_another_system_operation_is_running() { - let fixture = Fixture::new(); - // 裸 TCP 监听不回包：update 持有操作锁后停在 release 拉取阶段，accept - // 信号保证断言时锁一定已被占用 - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("release listener"); - let api_base = format!("http://{}/repos", listener.local_addr().expect("addr")); - let (connected_tx, connected_rx) = tokio::sync::oneshot::channel(); - let upstream = tokio::spawn(async move { - let connection = listener.accept().await; - let _ = connected_tx.send(()); - tokio::time::sleep(Duration::from_secs(30)).await; - drop(connection); - }); - let mut config = fixture.config(&api_base); - config.self_restart_enabled = true; - config.deployment_mode = "docker".to_owned(); - let shutdown = CancellationToken::new(); - let service = Arc::new(ProcessSystemOperations::new(shutdown.clone(), config)); - let update = tokio::spawn({ - let service = Arc::clone(&service); - async move { - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - } - }); - tokio::time::timeout(Duration::from_secs(5), connected_rx) - .await - .expect("update reaches release fetch while holding the operation lock") - .expect("connected signal"); - - let error = service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect_err("restart during update"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert!(!shutdown.is_cancelled()); - - update.await.expect("request").expect("accepted"); - assert!(fixture.lock().exists()); - upstream.abort(); - assert_eq!( - wait_for_update(&service).await.operation.status, - SystemOperationStatus::Failed - ); - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .expect("restart after lock release"); -} - -#[tokio::test] -async fn rollback_should_restore_binary_web_and_version_state() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - let updated = complete_update(&service, TARGET_VERSION).await; - assert_eq!( - updated.operation.status, - SystemOperationStatus::Succeeded, - "{updated:?}" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); - assert!(fixture.official().join("new-plugin.tar.gz").is_file()); - assert!(!fixture.official().join("old-plugin.tar.gz").exists()); - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt as _; - - assert_eq!( - fs::metadata(fixture.official()) - .expect("official directory") - .permissions() - .mode() - & 0o777, - 0o555 - ); - } - service - .rollback(Arc::new(AllowingUpdatePreflight)) - .await - .expect("rollback"); - - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - OLD_RELEASE_MANIFEST.as_bytes() - ); - assert_eq!( - fs::read(fixture.official().join("old-plugin.tar.gz")).expect("old plugin"), - b"old-plugin" - ); - assert!(!fixture.official().join("new-plugin.tar.gz").exists()); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Succeeded - ); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded, - "第二次更新必须能清理只读的官方目录备份" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); -} - -#[tokio::test] -async fn update_should_restore_all_files_when_plugin_revision_changes_during_swap() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - SystemOperations::perform_update( - &service, - Some(TARGET_VERSION.to_owned()), - None, - Arc::new(ChangingRevisionPreflight { - confirmations: AtomicUsize::new(0), - }), - ) - .await - .expect("update accepted"); - - let status = wait_for_update(&service).await; - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - OLD_RELEASE_MANIFEST.as_bytes() - ); - assert!(fixture.official().join("old-plugin.tar.gz").is_file()); - assert!(!fixture.official().join("new-plugin.tar.gz").exists()); -} - -#[tokio::test] -async fn rollback_should_restore_current_release_when_plugin_revision_changes_during_swap() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - - let error = service - .rollback(Arc::new(ChangingRollbackRevisionPreflight { - confirmations: AtomicUsize::new(0), - })) - .await - .expect_err("revision changed after rollback swap"); - - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Failed - ); -} - -#[tokio::test] -async fn cancelled_rollback_should_restore_current_release_after_the_swap() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - let reached_after_swap = Arc::new(tokio::sync::Notify::new()); - let task_service = service.clone(); - let task_reached = Arc::clone(&reached_after_swap); - let task = tokio::spawn(async move { - task_service - .rollback(Arc::new(BlockingRollbackPreflight { - confirmations: AtomicUsize::new(0), - reached_after_swap: task_reached, - })) - .await - }); - tokio::time::timeout(Duration::from_secs(5), reached_after_swap.notified()) - .await - .expect("post-swap confirmation reached"); - task.abort(); - assert!( - task.await - .expect_err("rollback task cancelled") - .is_cancelled() - ); - - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - b"new-manifest" - ); - assert_eq!( - service - .update_status() - .await - .expect("recovered status") - .operation - .status, - SystemOperationStatus::Failed - ); -} - -#[tokio::test] -async fn update_detail_should_report_refresh_failure_and_clear_previous_update() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, TARGET_VERSION).await; - let service = fixture.service(&server); - service - .update_detail(true, None) - .await - .expect("prime cache"); - - let detail = service - .update_detail(true, None) - .await - .expect("failure detail"); - assert!(!detail.has_update); - assert!(detail.warning.is_some()); - assert!(!detail.cached); - assert_eq!(detail.latest_version, "1.0.0"); - assert!(detail.notes.is_none()); - let version = service - .version() - .await - .expect("version after failed refresh"); - assert!(!version.has_update); - assert!(version.update_warning.is_some()); -} - -#[tokio::test] -async fn update_detail_should_offer_latest_non_draft_release_in_current_channel() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { - "tag_name": "v1.1.0-beta.3", - "name": "Draft 1.1.0-beta.3", - "prerelease": true, - "draft": true, - "assets": [], - }, - { - "tag_name": "v1.1.0-beta.2", - "name": "Release 1.1.0-beta.2", - "body": "notes", - "html_url": "https://github.com/owner/repository/releases", - "prerelease": true, - "assets": [], - }, - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "1.1.0-beta.1".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "1.1.0-beta.2"); - assert!(detail.has_update); -} - -#[tokio::test] -async fn update_detail_should_reject_untrusted_github_api_base() { - let fixture = Fixture::new(); - let service = ProcessSystemOperations::new( - CancellationToken::new(), - fixture.config("https://api.github.example/repos"), - ); - - let detail = service - .update_detail(true, None) - .await - .expect("safe rejection"); - assert!(!detail.update_supported); - assert!(detail.warning.is_some() || detail.unsupported_reason.is_some()); -} - -#[tokio::test] -async fn update_detail_should_withhold_updates_for_source_builds() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.build_type = "source".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "1.0.0"); - assert!(!detail.update_supported); - assert!(detail.unsupported_reason.is_some()); - assert!(!detail.has_update); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - assert!(!service.version().await.expect("version").has_update); - assert!( - server - .received_requests() - .await - .expect("requests") - .is_empty() - ); -} - -#[tokio::test] -async fn experimental_build_should_ignore_newer_stable_and_other_experimental_releases() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": "v4.0.0", "prerelease": false }, - { "tag_name": "v3.12.0", "prerelease": false }, - { "tag_name": "v3.10.0", "prerelease": false }, - { "tag_name": "v3.11.0-exp.3", "prerelease": true }, - { "tag_name": "v3.10.1-exp.3", "prerelease": true }, - { "tag_name": "v3.11.0-beta.1", "prerelease": true }, - { "tag_name": "v3.10.0-exp.other.3", "prerelease": true }, - { "tag_name": "v3.10.0-exp.3", "prerelease": true, "draft": true }, - { "tag_name": "v3.10.0-exp.4", "prerelease": false }, - { "tag_name": "invalid", "prerelease": true } - ]))) - .expect(1) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.10.0-exp.2".to_owned(); - config.build_type = "experimental".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "3.10.0-exp.2"); - assert!(!detail.has_update); - assert!(detail.update_supported); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - assert!(detail.warning.is_none()); - let version = service.version().await.expect("version"); - assert!(!version.has_update); - assert_eq!(version.update_channel, "exp"); - let cached = service - .update_detail(false, None) - .await - .expect("cached detail"); - assert!(cached.cached); - assert!(!cached.has_update); -} - -#[tokio::test] -async fn experimental_update_should_find_highest_same_channel_version_across_pages() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut first_page = vec![ - serde_json::json!({ - "tag_name": "v3.12.0", "prerelease": false - }); - 100 - ]; - first_page[1] = serde_json::json!({ "tag_name": "v3.10.0-exp.3", "prerelease": true }); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "1")) - .and(query_param("per_page", "100")) - .respond_with(ResponseTemplate::new(200).set_body_json(first_page)) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "2")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": "v3.10.0-exp.10", "prerelease": true, "body": "exp.10 notes" }, - { "tag_name": "v3.10.0-exp.2", "prerelease": true } - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.10.0-exp.2".to_owned(); - config.build_type = "experimental".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.latest_version, "3.10.0-exp.10"); - assert_eq!(detail.notes.as_deref(), Some("exp.10 notes")); - assert!(detail.has_update); - assert!(detail.update_supported); - assert!(service.version().await.expect("version").has_update); -} - -#[tokio::test] -async fn experimental_update_should_install_same_channel_release() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - "3.10.0-exp.3", - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.10.0-exp.2".to_owned(); - config.build_type = "experimental".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - assert_eq!( - complete_update(&service, "3.10.0-exp.3") - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Succeeded - ); -} - -#[tokio::test] -async fn update_should_reject_cross_channel_targets_before_fetching_or_replacing_files() { - for (build_type, current, target) in [ - ("experimental", "3.10.0-exp.2", "3.12.0"), - ("experimental", "3.10.0-exp.2", "3.10.0"), - ("experimental", "3.10.0-exp.2", "3.10.1-exp.3"), - ("experimental", "3.10.0-exp.2", "3.10.0-exp.1"), - ("release", "3.10.0-beta.2", "3.10.0-alpha.3"), - ("release", "3.10.0", "3.10.0+new"), - ("release", "3.10.0", "3.9.0"), - ("experimental", "3.10.0-exp.2", "3.11.0-beta.1"), - ("experimental", "3.10.0-exp.2", "3.10.0-exp.other.3"), - ("release", "3.10.0", "3.11.0-exp.1"), - ("release", "3.10.0", "3.11.0-beta.1"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - config.build_type = build_type.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let error = service - .perform_test_update(Some(target.to_owned())) - .await - .expect_err("cross-channel update"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert!(error.to_string().contains("所选通道")); - assert!( - server - .received_requests() - .await - .expect("requests") - .is_empty() - ); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert!(!fixture.state().exists()); - } -} - -#[tokio::test] -async fn update_checks_should_follow_stage_promotion_and_experimental_isolation() { - let cases = [ - ("3.12.0-alpha.1", "alpha", [true, true, true, true, false]), - ("3.12.0-beta.1", "beta", [false, true, true, true, false]), - ("3.12.0-rc.1", "rc", [false, false, true, true, false]), - ("3.12.0-exp.1", "exp", [false, false, false, false, true]), - ("3.11.0", "stable", [false, false, false, true, false]), - ]; - for (current, channel, expected) in cases { - for (target, allowed) in [ - "3.12.0-alpha.2", - "3.12.0-beta.2", - "3.12.0-rc.2", - "3.12.0", - "3.12.0-exp.2", - ] - .into_iter() - .zip(expected) - { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, target).await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert_eq!(detail.has_update, allowed, "{current} -> {target}"); - assert_eq!( - detail.latest_version, - if allowed { target } else { current } - ); - assert_eq!(detail.notes.is_some(), allowed); - assert_eq!(detail.release_url.is_some(), allowed); - assert!(detail.warning.is_none()); - let version = service.version().await.expect("version"); - assert_eq!(version.has_update, allowed); - assert_eq!(version.update_channel, channel); - } - } -} - -#[tokio::test] -async fn update_detail_should_keep_current_release_notes_without_allowing_reinstallation() { - for current in [ - "3.12.1", - "3.12.1-alpha.1", - "3.12.1-beta.1", - "3.12.1-rc.1", - "3.12.1-exp.1", - "3.12.1+build.1", - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let release_url = format!("https://github.com/owner/repository/releases/tag/v{current}"); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { - "tag_name": format!("v{current}"), - "prerelease": !semver::Version::parse(current).expect("version").pre.is_empty(), - "body": "## 当前版本\n\n- 修复更新状态", - "html_url": release_url, - } - ]))) - .expect(1) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - if current.contains("-exp.") { - config.build_type = "experimental".to_owned(); - } - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service - .update_detail(true, None) - .await - .expect("current release"); - assert_eq!(detail.latest_version, current); - assert!(!detail.has_update); - assert!(detail.update_supported); - assert_eq!( - detail.notes.as_deref(), - Some("## 当前版本\n\n- 修复更新状态") - ); - assert_eq!(detail.release_url.as_deref(), Some(release_url.as_str())); - let cached = service - .update_detail(false, None) - .await - .expect("cached release"); - assert!(cached.cached); - assert_eq!(cached.notes, detail.notes); - assert_eq!(cached.release_url, detail.release_url); - assert!(!service.version().await.expect("version").has_update); - - let error = service - .perform_test_update(Some(current.to_owned())) - .await - .expect_err("current release must not be installed again"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert!(!fixture.state().exists()); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - } -} - -#[tokio::test] -async fn update_detail_should_prefer_available_update_notes_over_current_release() { - for current_first in [true, false] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut releases = vec![ - serde_json::json!({ - "tag_name": "v1.0.0", "prerelease": false, "body": "current notes" - }), - serde_json::json!({ - "tag_name": "v1.9.9", "prerelease": false, "body": "update notes", - "html_url": "https://github.com/owner/repository/releases/tag/v1.9.9" - }), - ]; - if !current_first { - releases.reverse(); - } - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(releases)) - .mount(&server) - .await; - - let detail = fixture - .service(&server) - .update_detail(true, None) - .await - .expect("detail"); - assert!(detail.has_update); - assert_eq!(detail.latest_version, TARGET_VERSION); - assert_eq!(detail.notes.as_deref(), Some("update notes")); - assert_eq!( - detail.release_url.as_deref(), - Some("https://github.com/owner/repository/releases/tag/v1.9.9") - ); - } -} - -#[tokio::test] -async fn update_detail_should_find_current_release_notes_across_pages() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "1")) - .respond_with(ResponseTemplate::new(200).set_body_json(vec![ - serde_json::json!({ "tag_name": "v2.0.0", "prerelease": false }); - 100 - ])) - .expect(1) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .and(query_param("page", "2")) - .respond_with(release_response("1.0.0", Vec::new())) - .expect(1) - .mount(&server) - .await; - - let detail = fixture - .service(&server) - .update_detail(true, None) - .await - .expect("detail"); - assert!(!detail.has_update); - assert_eq!(detail.latest_version, "1.0.0"); - assert_eq!(detail.notes.as_deref(), Some("notes")); - assert!(detail.release_url.is_some()); -} - -#[tokio::test] -async fn update_detail_should_not_use_draft_or_mislabeled_current_release_notes() { - for (current, prerelease) in [("1.0.0", false), ("1.0.0-beta.1", true)] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": format!("v{current}"), "prerelease": prerelease, "draft": true, "body": "draft notes" }, - { "tag_name": format!("v{current}"), "prerelease": !prerelease, "body": "mislabeled notes" }, - { "tag_name": "v0.9.0", "prerelease": false, "body": "older notes" }, - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(!detail.has_update); - assert_eq!(detail.latest_version, current); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - } -} - -#[tokio::test] -async fn update_checks_should_ignore_downgrades_unknown_stages_and_build_metadata() { - for (current, target) in [ - ("3.12.0-beta.2", "3.12.0-beta.1"), - ("3.12.0", "3.12.0+new-build"), - ("3.12.0-beta.1+aaa", "3.12.0-beta.1+zzz"), - ("3.12.0", "3.11.0"), - ("3.12.0-alpha.1", "3.12.0-preview.2"), - ("3.12.0-alpha.1", "3.12.0-beta"), - ("3.12.0-alpha.1", "3.12.0-beta.0"), - ("3.12.0-alpha.1", "3.12.0-beta.2.extra"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, target).await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(!detail.has_update, "{current} -> {target}"); - assert_eq!(detail.latest_version, current); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); - } -} - -#[tokio::test] -async fn update_checks_should_fail_closed_for_unknown_current_channels() { - for (current, build_type) in [ - ("3.12.0-preview.1", "release"), - ("3.12.0-alpha", "release"), - ("3.12.0-exp.0", "experimental"), - ("invalid", "release"), - ("3.12.0", "experimental"), - ("3.12.0-beta.1", "experimental"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - config.build_type = build_type.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(!detail.has_update); - assert!(!detail.update_supported); - assert!(detail.unsupported_reason.is_some()); - assert!( - server - .received_requests() - .await - .expect("requests") - .is_empty() - ); - } -} - -#[tokio::test] -async fn stable_update_should_select_allowed_version_below_newer_major_releases() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - Mock::given(method("GET")) - .and(path("/repos/owner/repository/releases")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ - { "tag_name": "v4.0.0", "prerelease": false }, - { "tag_name": "v3.13.0-beta.1", "prerelease": true }, - { "tag_name": "v3.12.1", "prerelease": true }, - { "tag_name": "v3.12.0", "prerelease": false, "draft": true }, - { "tag_name": "v3.11.2", "prerelease": false }, - { "tag_name": "v3.11.10", "prerelease": false, "body": "stable notes" } - ]))) - .mount(&server) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = "3.11.0".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let detail = service.update_detail(true, None).await.expect("detail"); - assert!(detail.has_update); - assert_eq!(detail.latest_version, "3.11.10"); - assert_eq!(detail.notes.as_deref(), Some("stable notes")); -} - -#[tokio::test] -async fn prerelease_updates_should_install_later_stages_and_stable_release() { - for (current, target) in [ - ("3.12.0-alpha.1", "3.12.0-beta.1"), - ("3.12.0-beta.1", "3.12.0-rc.1"), - ("3.12.0-rc.1", "3.12.0"), - ("3.12.0-alpha.1", "3.12.0"), - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release(&server, target, ArchiveKind::Safe, ChecksumKind::Valid) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - config.version = current.to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - assert!( - service - .update_detail(true, None) - .await - .expect("detail") - .has_update - ); - assert_eq!( - complete_update(&service, target).await.operation.status, - SystemOperationStatus::Succeeded - ); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - assert_eq!( - service - .update_status() - .await - .expect("status") - .operation - .status, - SystemOperationStatus::Succeeded - ); - } -} - -#[tokio::test] -async fn update_detail_should_use_cached_release_when_not_refreshed() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture.mount_release_once(&server, TARGET_VERSION).await; - let service = fixture.service(&server); - service - .update_detail(true, None) - .await - .expect("prime cache"); - - assert!(service.update_detail(false, None).await.is_ok()); -} - -#[tokio::test] -async fn update_detail_should_withhold_cross_major_release() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release_once(&server, CROSS_MAJOR_VERSION) - .await; - - let detail = fixture - .service(&server) - .update_detail(true, None) - .await - .expect("detail"); - assert_eq!(detail.latest_version, "1.0.0"); - assert!(!detail.has_update); - assert!(detail.warning.is_none()); - assert!(detail.notes.is_none()); - assert!(detail.release_url.is_none()); -} - -#[tokio::test] -async fn update_should_reject_cross_major_target_before_fetching_release() { - let fixture = Fixture::new(); - let service = ProcessSystemOperations::new( - CancellationToken::new(), - fixture.config("http://127.0.0.1:1/repos"), - ); - - let error = service - .perform_test_update(Some(CROSS_MAJOR_VERSION.to_owned())) - .await - .expect_err("cross-major target must be rejected"); - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); -} - -#[tokio::test] -async fn update_events_should_close_after_terminal_update_log() { - let fixture = Fixture::new(); - let mut config = fixture.config("https://api.github.com/repos"); - config.build_type = "source".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let mut events = service.update_events(); - let _ = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await; - let first = events.next().await.expect("terminal event"); - - assert_eq!(first.level, SystemUpdateEventLevel::Error); - assert!(events.next().await.is_none()); -} - -#[tokio::test] -async fn update_events_should_open_authenticated_sse_stream() { - let fixture = Fixture::new(); - let mut config = fixture.config("https://api.github.com/repos"); - config.build_type = "source".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let mut stream = service.update_events(); - let _ = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await; - - assert!(stream.next().await.is_some()); -} - -#[tokio::test] -async fn update_events_should_preserve_complete_release_stage_sequence() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - let events = service.update_events(); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - let steps = events - .map(|event| event.step.expect("release update step")) - .collect::>() - .await; - - assert_eq!( - steps, - [ - "release", - "prepare", - "asset", - "asset", - "verify", - "prepare", - "download", - "download", - "download", - "checksum", - "checksum", - "extract", - "extract", - "preflight", - "preflight", - "replace", - "replace", - "done", - ] - ); -} - -#[tokio::test] -async fn update_events_should_report_download_bytes_and_percent() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let service = fixture.service(&server); - let events = service.update_events(); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - let progress = events - .filter_map( - |event| async move { event.progress_percent.map(|value| (value, event.message)) }, - ) - .collect::>() - .await; - - assert!(progress.iter().any(|(percent, message)| { - *percent == 100 && message.starts_with("已下载 ") && message.ends_with(" (100%)") - })); -} - -#[tokio::test] -async fn update_should_fail_when_release_checksum_is_missing() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Missing, - ) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_fail_when_release_checksum_mismatches() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Mismatch, - ) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_reject_insecure_release_archive() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_custom_release(&server, "http://github.com/archive", None) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_reject_release_archive_from_untrusted_host() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_custom_release(&server, "https://evil.example/archive", None) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[test] -fn update_should_trust_github_release_asset_redirect_host() { - assert!( - validate_download_url( - "https://release-assets.githubusercontent.com/github-production-release-asset/archive", - "https://api.github.com/repos", - ) - .is_ok() - ); -} - -#[tokio::test] -async fn update_should_reject_release_archive_with_unsafe_path() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::UnsafePath, - ChecksumKind::Valid, - ) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_require_a_flat_regular_official_plugin_bundle() { - for archive_kind in [ - ArchiveKind::MissingOfficialManifest, - ArchiveKind::NestedOfficialAsset, - ArchiveKind::OfficialSymlink, - ] { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release(&server, TARGET_VERSION, archive_kind, ChecksumKind::Valid) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); - } -} - -#[tokio::test] -async fn update_should_reject_untrusted_github_api_base() { - let fixture = Fixture::new(); - let service = ProcessSystemOperations::new( - CancellationToken::new(), - fixture.config("https://api.github.example/repos"), - ); - let error = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect_err("untrusted API rejected"); - - assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); -} - -#[tokio::test] -async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release(&server, "1.9.8", ArchiveKind::Safe, ChecksumKind::Valid) - .await; - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); -} - -#[tokio::test] -async fn update_should_remove_stale_file_lock_and_continue() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - fs::write(fixture.lock(), "stale").expect("lock"); - filetime::set_file_mtime(fixture.lock(), FileTime::from_unix_time(1, 0)).expect("old mtime"); - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded, - ); -} - -#[tokio::test] -async fn update_and_rollback_should_use_the_default_api_asset_directory() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let mut host: gateway_host::HostConfig = serde_json::from_value(serde_json::json!({ - "listen": { "host": "127.0.0.1", "port": 8080 }, - "runtime_data_dir": "../.runtime/data", - "logging": { - "level": "info", "stdout": true, - "file": { "enabled": false, "directory": "../.runtime/logs", "max_file_size_mb": 20 } - } - })) - .expect("binary host config"); - host.system_update = fixture.config(&format!("{}/repos", server.uri())); - host.system_update.web_dist_dir = None; - host.resolve_and_validate(&fixture.root.path().join("deploy"), &fixture.web()) - .expect("inherit API asset directory"); - let service = ProcessSystemOperations::new(CancellationToken::new(), host.system_update); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"new-web" - ); - - service - .rollback(Arc::new(AllowingUpdatePreflight)) - .await - .expect("rollback"); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); -} - -#[tokio::test] -async fn update_should_replace_web_assets_across_filesystems() { - let Ok(external) = tempfile::tempdir_in("/dev/shm") else { - return; - }; - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let mut config = fixture.config(&format!("{}/repos", server.uri())); - let web_dist = external.path().join("dist"); - config.web_dist_dir = Some(web_dist.clone()); - fs::create_dir_all(&web_dist).expect("web dir"); - fs::write(web_dist.join("index.html"), "old-web").expect("web"); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - assert_eq!( - complete_update(&service, TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Succeeded - ); - - assert_eq!( - fs::read(web_dist.join("index.html")).expect("web"), - b"new-web" - ); -} - -#[tokio::test] -async fn update_should_restore_web_assets_when_binary_backup_fails() { - let server = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &server, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - fs::create_dir(fixture.executable().with_extension("backup")).expect("blocking backup dir"); - - assert_eq!( - complete_update(&fixture.service(&server), TARGET_VERSION) - .await - .operation - .status, - SystemOperationStatus::Failed, - ); - assert_eq!( - fs::read(fixture.web().join("index.html")).expect("web"), - b"old-web" - ); - assert_eq!( - fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), - OLD_RELEASE_MANIFEST.as_bytes() - ); - assert!(fixture.official().join("old-plugin.tar.gz").is_file()); - assert!(!fixture.official().join("new-plugin.tar.gz").exists()); -} - -#[tokio::test] -async fn update_status_should_reconcile_manual_deployment_without_reusing_legacy_success() { - let fixture = Fixture::new(); - fs::write(fixture.state(), r#"{"previousVersion":"3.14.0","currentVersion":"3.14.1","operation":{"operationId":"legacy","kind":"update","status":"succeeded","targetVersion":"3.14.1"}}"#).expect("legacy state"); - let mut config = fixture.config("https://api.github.com/repos"); - config.version = "3.15.0".to_owned(); - config.deployment_mode = "docker".to_owned(); - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let status = service.update_status().await.expect("status"); - assert!(!status.need_restart); - assert_eq!(status.current_version.as_deref(), Some("3.15.0")); - assert!(status.previous_version.is_none()); - assert_eq!(status.operation.target_version.as_deref(), Some("3.14.1")); - assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); -} - -#[tokio::test] -async fn version_should_return_backend_build_metadata() { - let fixture = Fixture::new(); - let mut config = fixture.config("https://api.github.com/repos"); - config.version = "2.3.4".to_owned(); - config.git_sha = "abc123".to_owned(); - config.update_repository = None; - let service = ProcessSystemOperations::new(CancellationToken::new(), config); - let version = service.version().await.expect("version"); - - assert_eq!( - (version.version.as_str(), version.git_sha.as_str()), - ("2.3.4", "abc123") - ); -} - -#[tokio::test] -async fn accepted_update_should_survive_a_lost_http_response() { - let upstream = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &upstream, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - Mock::given(method("GET")) - .and(path("/archive")) - .respond_with( - ResponseTemplate::new(200) - .set_body_bytes(release_archive(ArchiveKind::Safe)) - .set_delay(Duration::from_secs(2)), - ) - .with_priority(1) - .mount(&upstream) - .await; - let service = Arc::new(fixture.service(&upstream)); - let handler_service = Arc::clone(&service); - let accepted = Arc::new(tokio::sync::Notify::new()); - let handler_accepted = Arc::clone(&accepted); - let router = axum::Router::new().route( - "/update", - axum::routing::post(move || { - let service = Arc::clone(&handler_service); - let accepted = Arc::clone(&handler_accepted); - async move { - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect("accepted"); - accepted.notify_one(); - // 模拟受理响应丢失，客户端断开会取消此 HTTP handler - std::future::pending::().await - } - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("listener"); - let address = listener.local_addr().expect("address"); - let server = tokio::spawn(async move { - axum::serve(listener, router).await.expect("serve"); - }); - let client = tokio::spawn(async move { - reqwest::Client::new() - .post(format!("http://{address}/update")) - .send() - .await - }); - tokio::time::timeout(Duration::from_secs(1), accepted.notified()) - .await - .expect("accepted before slow download finishes"); - client.abort(); - let _ = client.await; - assert_eq!( - service - .update_status() - .await - .expect("running status") - .operation - .status, - SystemOperationStatus::Running - ); - assert_eq!( - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect_err("duplicate update") - .kind(), - SystemOperationErrorKind::Conflict - ); - assert_eq!( - service - .rollback(Arc::new(AllowingUpdatePreflight)) - .await - .expect_err("rollback during update") - .kind(), - SystemOperationErrorKind::Conflict - ); - let status = wait_for_update(&service).await; - assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); - assert!(status.need_restart); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"new-binary" - ); - assert!(!fixture.lock().exists()); - assert_eq!( - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect_err("restart required") - .kind(), - SystemOperationErrorKind::Conflict - ); - let mut restarted_config = fixture.config(&format!("{}/repos", upstream.uri())); - restarted_config.version = TARGET_VERSION.to_owned(); - let restarted = ProcessSystemOperations::new(CancellationToken::new(), restarted_config); - assert!( - !restarted - .update_status() - .await - .expect("restarted status") - .need_restart - ); - server.abort(); -} - -#[tokio::test] -async fn host_shutdown_should_finish_an_accepted_update_and_release_its_lock() { - let upstream = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &upstream, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Valid, - ) - .await; - let cancellation = CancellationToken::new(); - let service = ProcessSystemOperations::new( - cancellation.clone(), - fixture.config(&format!("{}/repos", upstream.uri())), - ); - // 在后台任务第一次 poll 之前触发关闭，覆盖受理后的生命周期交接 - service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect("accepted"); - cancellation.cancel(); - let status = wait_for_update(&service).await; - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert!(status.operation.error.expect("reason").contains("中断")); - assert!(!fixture.lock().exists()); - assert_eq!( - fs::read(fixture.executable()).expect("binary"), - b"old-binary" - ); -} - -#[test] -fn dropped_update_task_should_persist_failure_even_before_its_first_poll() { - let fixture = Fixture::new(); - let service = fixture.service_for_url("http://127.0.0.1:1/repos"); - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("runtime"); - runtime - .block_on(service.perform_test_update(Some(TARGET_VERSION.to_owned()))) - .expect("accepted"); - drop(runtime); - let status: serde_json::Value = - serde_json::from_slice(&fs::read(fixture.state()).expect("state")).expect("JSON"); - assert_eq!(status["operation"]["status"], "failed"); - assert!(status["operation"]["finishedAt"].is_string()); - assert!(!fixture.lock().exists()); -} - -#[tokio::test] -async fn status_should_recover_abandoned_running_state_without_stealing_a_live_lock() { - let fixture = Fixture::new(); - fs::write(fixture.state(), r#"{"currentVersion":"1.0.0","operation":{"operationId":"abandoned","kind":"update","status":"running","targetVersion":"1.9.9"}}"#).expect("state"); - fs::write(fixture.lock(), "another process").expect("live lock"); - let service = fixture.service_for_url("http://127.0.0.1:1/repos"); - assert_eq!( - service - .update_status() - .await - .expect("locked status") - .operation - .status, - SystemOperationStatus::Running - ); - fs::remove_file(fixture.lock()).expect("release abandoned lock"); - let status = service.update_status().await.expect("recovered status"); - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert!(status.operation.finished_at.is_some()); - assert_eq!(status.operation.operation_id.as_deref(), Some("abandoned")); -} - -#[tokio::test] -async fn terminal_event_should_only_be_visible_after_status_is_persisted() { - let upstream = MockServer::start().await; - let fixture = Fixture::new(); - fixture - .mount_release( - &upstream, - TARGET_VERSION, - ArchiveKind::Safe, - ChecksumKind::Mismatch, - ) - .await; - let service = fixture.service(&upstream); - let mut events = service.update_events(); - let SystemOperationAccepted::Update { operation_id, .. } = service - .perform_test_update(Some(TARGET_VERSION.to_owned())) - .await - .expect("accepted") - else { - panic!("update") - }; - tokio::time::timeout(Duration::from_secs(5), async { - while let Some(event) = events.next().await { - assert_eq!(event.operation_id.as_deref(), Some(operation_id.as_str())); - if event.terminal { - let status = service.update_status().await.expect("terminal status"); - assert_eq!(status.operation.status, SystemOperationStatus::Failed); - assert!(status.operation.finished_at.is_some()); - return; - } - } - panic!("missing terminal event"); - }) - .await - .expect("terminal event"); -} - async fn wait_for_update(service: &ProcessSystemOperations) -> SystemUpdateStatus { tokio::time::timeout(Duration::from_secs(5), async { loop { @@ -2278,28 +500,3 @@ impl gateway_admin::ports::system::SystemRestartPreflight for RejectingRestartPr )) } } - -#[tokio::test] -async fn restart_checks_plugins_before_shutdown_and_releases_locks_on_rejection() { - let fixture = Fixture::new(); - let mut config = fixture.config("http://127.0.0.1:1/repos"); - config.self_restart_enabled = true; - config.deployment_mode = "docker".into(); - let shutdown = CancellationToken::new(); - let service = ProcessSystemOperations::new(shutdown.clone(), config); - assert!( - service - .restart(Arc::new(RejectingRestartPreflight)) - .await - .is_err() - ); - assert!(!shutdown.is_cancelled()); - assert!(service.restart_candidate().await.unwrap().is_some()); - service - .restart(Arc::new(AllowingUpdatePreflight)) - .await - .unwrap(); - tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) - .await - .unwrap(); -} diff --git a/backend/crates/gateway-host/tests/system_update/operations.rs b/backend/crates/gateway-host/tests/system_update/operations.rs new file mode 100644 index 000000000..5387b7ddd --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/operations.rs @@ -0,0 +1,1083 @@ +//! 系统更新、回滚与重启的受理、互斥和文件交换测试 + +use super::*; + +#[tokio::test] +async fn restart_should_not_shutdown_when_replacement_spawn_fails() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.executable_path = Some(fixture.root.path().join("missing")); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + + assert!( + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .is_err() + ); + assert!(!shutdown.is_cancelled()); +} + +#[tokio::test] +async fn restart_should_request_process_restart_inside_docker() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.deployment_mode = "docker".to_owned(); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("restart accepted"); + + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .expect("shutdown requested"); +} + +#[tokio::test] +async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { + let fixture = Fixture::new(); + let replacement_marker = fixture.root.path().join("replacement-ran"); + fixture.write_executable("#!/bin/sh\n: > \"${0%/*}/replacement-ran\"\n"); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + + assert_eq!( + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("replacement scheduled") + .kind(), + SystemOperationKind::Restart + ); + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .expect("shutdown requested after spawn"); + tokio::time::timeout(Duration::from_secs(2), async { + while !replacement_marker.is_file() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("replacement executed before fixture cleanup"); +} + +#[cfg(target_os = "linux")] +#[tokio::test] +async fn restart_should_use_startup_path_after_running_executable_is_renamed() { + const CHILD_ENV: &str = "CPR_TEST_RESTART_RENAMED_EXECUTABLE_CHILD"; + let executable = std::env::current_exe().expect("test executable"); + if std::env::var_os(CHILD_ENV).is_none() { + // 用例只移动目录项，不修改旧程序内容；同盘硬链接避免复制整个调试二进制 + let directory = + tempfile::tempdir_in(executable.parent().unwrap()).expect("isolated test directory"); + let child = directory.path().join("restart-test"); + fs::hard_link(&executable, &child).expect("link isolated test executable"); + let output = std::process::Command::new(child) + .args([ + "--exact", + "system_update::restart_should_use_startup_path_after_running_executable_is_renamed", + ]) + .env(CHILD_ENV, "1") + .env("CPR_RESTART_DELAY_MS", "1200") + .output() + .expect("isolated restart test"); + assert!( + output.status.success(), + "stdout: {}\nstderr: {}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + if executable.file_name() == Some(std::ffi::OsStr::new("codex-proxy-rs.backup")) { + // 误启动旧程序时留下可观察结果并退出，避免旧副本再次发起重启 + fs::write( + executable + .parent() + .expect("installation directory") + .join("replacement-version"), + "old", + ) + .expect("old replacement marker"); + return; + } + + let fixture = Fixture::new_in(executable.parent().unwrap()); + fs::rename(executable, fixture.executable()).expect("move isolated running executable"); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.executable_path = None; + config.self_restart_enabled = true; + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let backup = fixture.root.path().join("codex-proxy-rs.backup"); + fs::rename(fixture.executable(), &backup).expect("back up running executable"); + fixture.write_executable("#!/bin/sh\nprintf new > \"${0%/*}/replacement-version\"\n"); + assert_eq!(std::env::current_exe().expect("renamed executable"), backup); + + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("restart accepted"); + let marker = fixture.root.path().join("replacement-version"); + tokio::time::timeout(Duration::from_secs(5), async { + while !marker.is_file() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("replacement executed"); + assert_eq!( + fs::read_to_string(marker).expect("replacement version"), + "new" + ); +} + +#[tokio::test] +async fn restart_should_conflict_while_another_system_operation_is_running() { + let fixture = Fixture::new(); + // 裸 TCP 监听不回包：update 持有操作锁后停在 release 拉取阶段，accept + // 信号保证断言时锁一定已被占用 + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("release listener"); + let api_base = format!("http://{}/repos", listener.local_addr().expect("addr")); + let (connected_tx, connected_rx) = tokio::sync::oneshot::channel(); + let upstream = tokio::spawn(async move { + let connection = listener.accept().await; + let _ = connected_tx.send(()); + tokio::time::sleep(Duration::from_secs(30)).await; + drop(connection); + }); + let mut config = fixture.config(&api_base); + config.self_restart_enabled = true; + config.deployment_mode = "docker".to_owned(); + let shutdown = CancellationToken::new(); + let service = Arc::new(ProcessSystemOperations::new(shutdown.clone(), config)); + let update = tokio::spawn({ + let service = Arc::clone(&service); + async move { + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + } + }); + tokio::time::timeout(Duration::from_secs(5), connected_rx) + .await + .expect("update reaches release fetch while holding the operation lock") + .expect("connected signal"); + + let error = service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect_err("restart during update"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert!(!shutdown.is_cancelled()); + + update.await.expect("request").expect("accepted"); + assert!(fixture.lock().exists()); + upstream.abort(); + assert_eq!( + wait_for_update(&service).await.operation.status, + SystemOperationStatus::Failed + ); + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .expect("restart after lock release"); +} + +#[tokio::test] +async fn rollback_should_restore_binary_web_and_version_state() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + let updated = complete_update(&service, TARGET_VERSION).await; + assert_eq!( + updated.operation.status, + SystemOperationStatus::Succeeded, + "{updated:?}" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); + assert!(fixture.official().join("new-plugin.tar.gz").is_file()); + assert!(!fixture.official().join("old-plugin.tar.gz").exists()); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + + assert_eq!( + fs::metadata(fixture.official()) + .expect("official directory") + .permissions() + .mode() + & 0o777, + 0o555 + ); + } + service + .rollback(Arc::new(AllowingUpdatePreflight)) + .await + .expect("rollback"); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + OLD_RELEASE_MANIFEST.as_bytes() + ); + assert_eq!( + fs::read(fixture.official().join("old-plugin.tar.gz")).expect("old plugin"), + b"old-plugin" + ); + assert!(!fixture.official().join("new-plugin.tar.gz").exists()); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded, + "第二次更新必须能清理只读的官方目录备份" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); +} + +#[tokio::test] +async fn update_should_restore_all_files_when_plugin_revision_changes_during_swap() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + SystemOperations::perform_update( + &service, + Some(TARGET_VERSION.to_owned()), + None, + Arc::new(ChangingRevisionPreflight { + confirmations: AtomicUsize::new(0), + }), + ) + .await + .expect("update accepted"); + + let status = wait_for_update(&service).await; + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + OLD_RELEASE_MANIFEST.as_bytes() + ); + assert!(fixture.official().join("old-plugin.tar.gz").is_file()); + assert!(!fixture.official().join("new-plugin.tar.gz").exists()); +} + +#[tokio::test] +async fn rollback_should_restore_current_release_when_plugin_revision_changes_during_swap() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + + let error = service + .rollback(Arc::new(ChangingRollbackRevisionPreflight { + confirmations: AtomicUsize::new(0), + })) + .await + .expect_err("revision changed after rollback swap"); + + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Failed + ); +} + +#[tokio::test] +async fn cancelled_rollback_should_restore_current_release_after_the_swap() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + let reached_after_swap = Arc::new(tokio::sync::Notify::new()); + let task_service = service.clone(); + let task_reached = Arc::clone(&reached_after_swap); + let task = tokio::spawn(async move { + task_service + .rollback(Arc::new(BlockingRollbackPreflight { + confirmations: AtomicUsize::new(0), + reached_after_swap: task_reached, + })) + .await + }); + tokio::time::timeout(Duration::from_secs(5), reached_after_swap.notified()) + .await + .expect("post-swap confirmation reached"); + task.abort(); + assert!( + task.await + .expect_err("rollback task cancelled") + .is_cancelled() + ); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + b"new-manifest" + ); + assert_eq!( + service + .update_status() + .await + .expect("recovered status") + .operation + .status, + SystemOperationStatus::Failed + ); +} + +#[tokio::test] +async fn experimental_update_should_install_same_channel_release() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + "3.10.0-exp.3", + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.10.0-exp.2".to_owned(); + config.build_type = "experimental".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + assert_eq!( + complete_update(&service, "3.10.0-exp.3") + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); +} + +#[tokio::test] +async fn update_should_reject_cross_channel_targets_before_fetching_or_replacing_files() { + for (build_type, current, target) in [ + ("experimental", "3.10.0-exp.2", "3.12.0"), + ("experimental", "3.10.0-exp.2", "3.10.0"), + ("experimental", "3.10.0-exp.2", "3.10.1-exp.3"), + ("experimental", "3.10.0-exp.2", "3.10.0-exp.1"), + ("release", "3.10.0-beta.2", "3.10.0-alpha.3"), + ("release", "3.10.0", "3.10.0+new"), + ("release", "3.10.0", "3.9.0"), + ("experimental", "3.10.0-exp.2", "3.11.0-beta.1"), + ("experimental", "3.10.0-exp.2", "3.10.0-exp.other.3"), + ("release", "3.10.0", "3.11.0-exp.1"), + ("release", "3.10.0", "3.11.0-beta.1"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + config.build_type = build_type.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let error = service + .perform_test_update(Some(target.to_owned())) + .await + .expect_err("cross-channel update"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert!(error.to_string().contains("所选通道")); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert!(!fixture.state().exists()); + } +} + +#[tokio::test] +async fn prerelease_updates_should_install_later_stages_and_stable_release() { + for (current, target) in [ + ("3.12.0-alpha.1", "3.12.0-beta.1"), + ("3.12.0-beta.1", "3.12.0-rc.1"), + ("3.12.0-rc.1", "3.12.0"), + ("3.12.0-alpha.1", "3.12.0"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, target, ArchiveKind::Safe, ChecksumKind::Valid) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + assert!( + service + .update_detail(true, None) + .await + .expect("detail") + .has_update + ); + assert_eq!( + complete_update(&service, target).await.operation.status, + SystemOperationStatus::Succeeded + ); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); + } +} + +#[tokio::test] +async fn update_should_reject_cross_major_target_before_fetching_release() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("http://127.0.0.1:1/repos"), + ); + + let error = service + .perform_test_update(Some(CROSS_MAJOR_VERSION.to_owned())) + .await + .expect_err("cross-major target must be rejected"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); +} + +#[tokio::test] +async fn update_should_fail_when_release_checksum_is_missing() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Missing, + ) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_fail_when_release_checksum_mismatches() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Mismatch, + ) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_reject_insecure_release_archive() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_custom_release(&server, "http://github.com/archive", None) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_reject_release_archive_from_untrusted_host() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_custom_release(&server, "https://evil.example/archive", None) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_reject_release_archive_with_unsafe_path() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::UnsafePath, + ChecksumKind::Valid, + ) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_require_a_flat_regular_official_plugin_bundle() { + for archive_kind in [ + ArchiveKind::MissingOfficialManifest, + ArchiveKind::NestedOfficialAsset, + ArchiveKind::OfficialSymlink, + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, TARGET_VERSION, archive_kind, ChecksumKind::Valid) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); + } +} + +#[tokio::test] +async fn update_should_reject_untrusted_github_api_base() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.example/repos"), + ); + let error = service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("untrusted API rejected"); + + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); +} + +#[tokio::test] +async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, "1.9.8", ArchiveKind::Safe, ChecksumKind::Valid) + .await; + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); +} + +#[tokio::test] +async fn update_should_remove_stale_file_lock_and_continue() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fs::write(fixture.lock(), "stale").expect("lock"); + filetime::set_file_mtime(fixture.lock(), FileTime::from_unix_time(1, 0)).expect("old mtime"); + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded, + ); +} + +#[tokio::test] +async fn update_and_rollback_should_use_the_default_api_asset_directory() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut host: gateway_host::HostConfig = serde_json::from_value(serde_json::json!({ + "listen": { "host": "127.0.0.1", "port": 8080 }, + "runtime_data_dir": "../.runtime/data", + "logging": { + "level": "info", "stdout": true, + "file": { "enabled": false, "directory": "../.runtime/logs", "max_file_size_mb": 20 } + } + })) + .expect("binary host config"); + host.system_update = fixture.config(&format!("{}/repos", server.uri())); + host.system_update.web_dist_dir = None; + host.resolve_and_validate(&fixture.root.path().join("deploy"), &fixture.web()) + .expect("inherit API asset directory"); + let service = ProcessSystemOperations::new(CancellationToken::new(), host.system_update); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); + + service + .rollback(Arc::new(AllowingUpdatePreflight)) + .await + .expect("rollback"); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); +} + +#[tokio::test] +async fn update_should_replace_web_assets_across_filesystems() { + let Ok(external) = tempfile::tempdir_in("/dev/shm") else { + return; + }; + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + let web_dist = external.path().join("dist"); + config.web_dist_dir = Some(web_dist.clone()); + fs::create_dir_all(&web_dist).expect("web dir"); + fs::write(web_dist.join("index.html"), "old-web").expect("web"); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + assert_eq!( + complete_update(&service, TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Succeeded + ); + + assert_eq!( + fs::read(web_dist.join("index.html")).expect("web"), + b"new-web" + ); +} + +#[tokio::test] +async fn update_should_restore_web_assets_when_binary_backup_fails() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fs::create_dir(fixture.executable().with_extension("backup")).expect("blocking backup dir"); + + assert_eq!( + complete_update(&fixture.service(&server), TARGET_VERSION) + .await + .operation + .status, + SystemOperationStatus::Failed, + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + fs::read(fixture.official().join("plugin-release-manifest.json")).expect("manifest"), + OLD_RELEASE_MANIFEST.as_bytes() + ); + assert!(fixture.official().join("old-plugin.tar.gz").is_file()); + assert!(!fixture.official().join("new-plugin.tar.gz").exists()); +} + +#[tokio::test] +async fn version_should_return_backend_build_metadata() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.version = "2.3.4".to_owned(); + config.git_sha = "abc123".to_owned(); + config.update_repository = None; + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let version = service.version().await.expect("version"); + + assert_eq!( + (version.version.as_str(), version.git_sha.as_str()), + ("2.3.4", "abc123") + ); +} + +#[tokio::test] +async fn accepted_update_should_survive_a_lost_http_response() { + let upstream = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &upstream, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + Mock::given(method("GET")) + .and(path("/archive")) + .respond_with( + ResponseTemplate::new(200) + .set_body_bytes(release_archive(ArchiveKind::Safe)) + .set_delay(Duration::from_secs(2)), + ) + .with_priority(1) + .mount(&upstream) + .await; + let service = Arc::new(fixture.service(&upstream)); + let handler_service = Arc::clone(&service); + let accepted = Arc::new(tokio::sync::Notify::new()); + let handler_accepted = Arc::clone(&accepted); + let router = axum::Router::new().route( + "/update", + axum::routing::post(move || { + let service = Arc::clone(&handler_service); + let accepted = Arc::clone(&handler_accepted); + async move { + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("accepted"); + accepted.notify_one(); + // 模拟受理响应丢失，客户端断开会取消此 HTTP handler + std::future::pending::().await + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener"); + let address = listener.local_addr().expect("address"); + let server = tokio::spawn(async move { + axum::serve(listener, router).await.expect("serve"); + }); + let client = tokio::spawn(async move { + reqwest::Client::new() + .post(format!("http://{address}/update")) + .send() + .await + }); + tokio::time::timeout(Duration::from_secs(1), accepted.notified()) + .await + .expect("accepted before slow download finishes"); + client.abort(); + let _ = client.await; + assert_eq!( + service + .update_status() + .await + .expect("running status") + .operation + .status, + SystemOperationStatus::Running + ); + assert_eq!( + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("duplicate update") + .kind(), + SystemOperationErrorKind::Conflict + ); + assert_eq!( + service + .rollback(Arc::new(AllowingUpdatePreflight)) + .await + .expect_err("rollback during update") + .kind(), + SystemOperationErrorKind::Conflict + ); + let status = wait_for_update(&service).await; + assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); + assert!(status.need_restart); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert!(!fixture.lock().exists()); + assert_eq!( + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("restart required") + .kind(), + SystemOperationErrorKind::Conflict + ); + let mut restarted_config = fixture.config(&format!("{}/repos", upstream.uri())); + restarted_config.version = TARGET_VERSION.to_owned(); + let restarted = ProcessSystemOperations::new(CancellationToken::new(), restarted_config); + assert!( + !restarted + .update_status() + .await + .expect("restarted status") + .need_restart + ); + server.abort(); +} + +#[tokio::test] +async fn host_shutdown_should_finish_an_accepted_update_and_release_its_lock() { + let upstream = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &upstream, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let cancellation = CancellationToken::new(); + let service = ProcessSystemOperations::new( + cancellation.clone(), + fixture.config(&format!("{}/repos", upstream.uri())), + ); + // 在后台任务第一次 poll 之前触发关闭，覆盖受理后的生命周期交接 + service + .perform_test_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("accepted"); + cancellation.cancel(); + let status = wait_for_update(&service).await; + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert!(status.operation.error.expect("reason").contains("中断")); + assert!(!fixture.lock().exists()); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); +} + +#[tokio::test] +async fn restart_checks_plugins_before_shutdown_and_releases_locks_on_rejection() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.deployment_mode = "docker".into(); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + assert!( + service + .restart(Arc::new(RejectingRestartPreflight)) + .await + .is_err() + ); + assert!(!shutdown.is_cancelled()); + assert!(service.restart_candidate().await.unwrap().is_some()); + service + .restart(Arc::new(AllowingUpdatePreflight)) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .unwrap(); +} diff --git a/backend/crates/gateway-host/tests/system_update/release.rs b/backend/crates/gateway-host/tests/system_update/release.rs new file mode 100644 index 000000000..b54675d06 --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/release.rs @@ -0,0 +1,512 @@ +//! 发行查询、缓存、版本通道与目标策略测试 + +use super::*; + +#[tokio::test] +async fn update_detail_should_report_refresh_failure_and_clear_previous_update() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, TARGET_VERSION).await; + let service = fixture.service(&server); + service + .update_detail(true, None) + .await + .expect("prime cache"); + + let detail = service + .update_detail(true, None) + .await + .expect("failure detail"); + assert!(!detail.has_update); + assert!(detail.warning.is_some()); + assert!(!detail.cached); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(detail.notes.is_none()); + let version = service + .version() + .await + .expect("version after failed refresh"); + assert!(!version.has_update); + assert!(version.update_warning.is_some()); +} + +#[tokio::test] +async fn update_detail_should_offer_latest_non_draft_release_in_current_channel() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { + "tag_name": "v1.1.0-beta.3", + "name": "Draft 1.1.0-beta.3", + "prerelease": true, + "draft": true, + "assets": [], + }, + { + "tag_name": "v1.1.0-beta.2", + "name": "Release 1.1.0-beta.2", + "body": "notes", + "html_url": "https://github.com/owner/repository/releases", + "prerelease": true, + "assets": [], + }, + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "1.1.0-beta.1".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "1.1.0-beta.2"); + assert!(detail.has_update); +} + +#[tokio::test] +async fn update_detail_should_reject_untrusted_github_api_base() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.example/repos"), + ); + + let detail = service + .update_detail(true, None) + .await + .expect("safe rejection"); + assert!(!detail.update_supported); + assert!(detail.warning.is_some() || detail.unsupported_reason.is_some()); +} + +#[tokio::test] +async fn update_detail_should_withhold_updates_for_source_builds() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(!detail.update_supported); + assert!(detail.unsupported_reason.is_some()); + assert!(!detail.has_update); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + assert!(!service.version().await.expect("version").has_update); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); +} + +#[tokio::test] +async fn experimental_build_should_ignore_newer_stable_and_other_experimental_releases() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": "v4.0.0", "prerelease": false }, + { "tag_name": "v3.12.0", "prerelease": false }, + { "tag_name": "v3.10.0", "prerelease": false }, + { "tag_name": "v3.11.0-exp.3", "prerelease": true }, + { "tag_name": "v3.10.1-exp.3", "prerelease": true }, + { "tag_name": "v3.11.0-beta.1", "prerelease": true }, + { "tag_name": "v3.10.0-exp.other.3", "prerelease": true }, + { "tag_name": "v3.10.0-exp.3", "prerelease": true, "draft": true }, + { "tag_name": "v3.10.0-exp.4", "prerelease": false }, + { "tag_name": "invalid", "prerelease": true } + ]))) + .expect(1) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.10.0-exp.2".to_owned(); + config.build_type = "experimental".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "3.10.0-exp.2"); + assert!(!detail.has_update); + assert!(detail.update_supported); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + assert!(detail.warning.is_none()); + let version = service.version().await.expect("version"); + assert!(!version.has_update); + assert_eq!(version.update_channel, "exp"); + let cached = service + .update_detail(false, None) + .await + .expect("cached detail"); + assert!(cached.cached); + assert!(!cached.has_update); +} + +#[tokio::test] +async fn experimental_update_should_find_highest_same_channel_version_across_pages() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut first_page = vec![ + serde_json::json!({ + "tag_name": "v3.12.0", "prerelease": false + }); + 100 + ]; + first_page[1] = serde_json::json!({ "tag_name": "v3.10.0-exp.3", "prerelease": true }); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "1")) + .and(query_param("per_page", "100")) + .respond_with(ResponseTemplate::new(200).set_body_json(first_page)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "2")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": "v3.10.0-exp.10", "prerelease": true, "body": "exp.10 notes" }, + { "tag_name": "v3.10.0-exp.2", "prerelease": true } + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.10.0-exp.2".to_owned(); + config.build_type = "experimental".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.latest_version, "3.10.0-exp.10"); + assert_eq!(detail.notes.as_deref(), Some("exp.10 notes")); + assert!(detail.has_update); + assert!(detail.update_supported); + assert!(service.version().await.expect("version").has_update); +} + +#[tokio::test] +async fn update_checks_should_follow_stage_promotion_and_experimental_isolation() { + let cases = [ + ("3.12.0-alpha.1", "alpha", [true, true, true, true, false]), + ("3.12.0-beta.1", "beta", [false, true, true, true, false]), + ("3.12.0-rc.1", "rc", [false, false, true, true, false]), + ("3.12.0-exp.1", "exp", [false, false, false, false, true]), + ("3.11.0", "stable", [false, false, false, true, false]), + ]; + for (current, channel, expected) in cases { + for (target, allowed) in [ + "3.12.0-alpha.2", + "3.12.0-beta.2", + "3.12.0-rc.2", + "3.12.0", + "3.12.0-exp.2", + ] + .into_iter() + .zip(expected) + { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, target).await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert_eq!(detail.has_update, allowed, "{current} -> {target}"); + assert_eq!( + detail.latest_version, + if allowed { target } else { current } + ); + assert_eq!(detail.notes.is_some(), allowed); + assert_eq!(detail.release_url.is_some(), allowed); + assert!(detail.warning.is_none()); + let version = service.version().await.expect("version"); + assert_eq!(version.has_update, allowed); + assert_eq!(version.update_channel, channel); + } + } +} + +#[tokio::test] +async fn update_detail_should_keep_current_release_notes_without_allowing_reinstallation() { + for current in [ + "3.12.1", + "3.12.1-alpha.1", + "3.12.1-beta.1", + "3.12.1-rc.1", + "3.12.1-exp.1", + "3.12.1+build.1", + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let release_url = format!("https://github.com/owner/repository/releases/tag/v{current}"); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { + "tag_name": format!("v{current}"), + "prerelease": !semver::Version::parse(current).expect("version").pre.is_empty(), + "body": "## 当前版本\n\n- 修复更新状态", + "html_url": release_url, + } + ]))) + .expect(1) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + if current.contains("-exp.") { + config.build_type = "experimental".to_owned(); + } + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service + .update_detail(true, None) + .await + .expect("current release"); + assert_eq!(detail.latest_version, current); + assert!(!detail.has_update); + assert!(detail.update_supported); + assert_eq!( + detail.notes.as_deref(), + Some("## 当前版本\n\n- 修复更新状态") + ); + assert_eq!(detail.release_url.as_deref(), Some(release_url.as_str())); + let cached = service + .update_detail(false, None) + .await + .expect("cached release"); + assert!(cached.cached); + assert_eq!(cached.notes, detail.notes); + assert_eq!(cached.release_url, detail.release_url); + assert!(!service.version().await.expect("version").has_update); + + let error = service + .perform_test_update(Some(current.to_owned())) + .await + .expect_err("current release must not be installed again"); + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); + assert!(!fixture.state().exists()); + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + } +} + +#[tokio::test] +async fn update_detail_should_prefer_available_update_notes_over_current_release() { + for current_first in [true, false] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut releases = vec![ + serde_json::json!({ + "tag_name": "v1.0.0", "prerelease": false, "body": "current notes" + }), + serde_json::json!({ + "tag_name": "v1.9.9", "prerelease": false, "body": "update notes", + "html_url": "https://github.com/owner/repository/releases/tag/v1.9.9" + }), + ]; + if !current_first { + releases.reverse(); + } + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(releases)) + .mount(&server) + .await; + + let detail = fixture + .service(&server) + .update_detail(true, None) + .await + .expect("detail"); + assert!(detail.has_update); + assert_eq!(detail.latest_version, TARGET_VERSION); + assert_eq!(detail.notes.as_deref(), Some("update notes")); + assert_eq!( + detail.release_url.as_deref(), + Some("https://github.com/owner/repository/releases/tag/v1.9.9") + ); + } +} + +#[tokio::test] +async fn update_detail_should_find_current_release_notes_across_pages() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "1")) + .respond_with(ResponseTemplate::new(200).set_body_json(vec![ + serde_json::json!({ "tag_name": "v2.0.0", "prerelease": false }); + 100 + ])) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .and(query_param("page", "2")) + .respond_with(release_response("1.0.0", Vec::new())) + .expect(1) + .mount(&server) + .await; + + let detail = fixture + .service(&server) + .update_detail(true, None) + .await + .expect("detail"); + assert!(!detail.has_update); + assert_eq!(detail.latest_version, "1.0.0"); + assert_eq!(detail.notes.as_deref(), Some("notes")); + assert!(detail.release_url.is_some()); +} + +#[tokio::test] +async fn update_detail_should_not_use_draft_or_mislabeled_current_release_notes() { + for (current, prerelease) in [("1.0.0", false), ("1.0.0-beta.1", true)] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": format!("v{current}"), "prerelease": prerelease, "draft": true, "body": "draft notes" }, + { "tag_name": format!("v{current}"), "prerelease": !prerelease, "body": "mislabeled notes" }, + { "tag_name": "v0.9.0", "prerelease": false, "body": "older notes" }, + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(!detail.has_update); + assert_eq!(detail.latest_version, current); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + } +} + +#[tokio::test] +async fn update_checks_should_ignore_downgrades_unknown_stages_and_build_metadata() { + for (current, target) in [ + ("3.12.0-beta.2", "3.12.0-beta.1"), + ("3.12.0", "3.12.0+new-build"), + ("3.12.0-beta.1+aaa", "3.12.0-beta.1+zzz"), + ("3.12.0", "3.11.0"), + ("3.12.0-alpha.1", "3.12.0-preview.2"), + ("3.12.0-alpha.1", "3.12.0-beta"), + ("3.12.0-alpha.1", "3.12.0-beta.0"), + ("3.12.0-alpha.1", "3.12.0-beta.2.extra"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, target).await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(!detail.has_update, "{current} -> {target}"); + assert_eq!(detail.latest_version, current); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); + } +} + +#[tokio::test] +async fn update_checks_should_fail_closed_for_unknown_current_channels() { + for (current, build_type) in [ + ("3.12.0-preview.1", "release"), + ("3.12.0-alpha", "release"), + ("3.12.0-exp.0", "experimental"), + ("invalid", "release"), + ("3.12.0", "experimental"), + ("3.12.0-beta.1", "experimental"), + ] { + let server = MockServer::start().await; + let fixture = Fixture::new(); + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = current.to_owned(); + config.build_type = build_type.to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(!detail.has_update); + assert!(!detail.update_supported); + assert!(detail.unsupported_reason.is_some()); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); + } +} + +#[tokio::test] +async fn stable_update_should_select_allowed_version_below_newer_major_releases() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([ + { "tag_name": "v4.0.0", "prerelease": false }, + { "tag_name": "v3.13.0-beta.1", "prerelease": true }, + { "tag_name": "v3.12.1", "prerelease": true }, + { "tag_name": "v3.12.0", "prerelease": false, "draft": true }, + { "tag_name": "v3.11.2", "prerelease": false }, + { "tag_name": "v3.11.10", "prerelease": false, "body": "stable notes" } + ]))) + .mount(&server) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.version = "3.11.0".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let detail = service.update_detail(true, None).await.expect("detail"); + assert!(detail.has_update); + assert_eq!(detail.latest_version, "3.11.10"); + assert_eq!(detail.notes.as_deref(), Some("stable notes")); +} + +#[tokio::test] +async fn update_detail_should_use_cached_release_when_not_refreshed() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, TARGET_VERSION).await; + let service = fixture.service(&server); + service + .update_detail(true, None) + .await + .expect("prime cache"); + + assert!(service.update_detail(false, None).await.is_ok()); +} + +#[tokio::test] +async fn update_detail_should_withhold_cross_major_release() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release_once(&server, CROSS_MAJOR_VERSION) + .await; + + let detail = fixture + .service(&server) + .update_detail(true, None) + .await + .expect("detail"); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(!detail.has_update); + assert!(detail.warning.is_none()); + assert!(detail.notes.is_none()); + assert!(detail.release_url.is_none()); +} diff --git a/backend/crates/gateway-host/tests/system_update/state.rs b/backend/crates/gateway-host/tests/system_update/state.rs new file mode 100644 index 000000000..04cf95e83 --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/state.rs @@ -0,0 +1,60 @@ +//! 系统更新状态恢复与任务中止持久化测试 + +use super::*; + +#[tokio::test] +async fn update_status_should_reconcile_manual_deployment_without_reusing_legacy_success() { + let fixture = Fixture::new(); + fs::write(fixture.state(), r#"{"previousVersion":"3.14.0","currentVersion":"3.14.1","operation":{"operationId":"legacy","kind":"update","status":"succeeded","targetVersion":"3.14.1"}}"#).expect("legacy state"); + let mut config = fixture.config("https://api.github.com/repos"); + config.version = "3.15.0".to_owned(); + config.deployment_mode = "docker".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let status = service.update_status().await.expect("status"); + assert!(!status.need_restart); + assert_eq!(status.current_version.as_deref(), Some("3.15.0")); + assert!(status.previous_version.is_none()); + assert_eq!(status.operation.target_version.as_deref(), Some("3.14.1")); + assert_eq!(status.operation.status, SystemOperationStatus::Succeeded); +} + +#[test] +fn dropped_update_task_should_persist_failure_even_before_its_first_poll() { + let fixture = Fixture::new(); + let service = fixture.service_for_url("http://127.0.0.1:1/repos"); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + runtime + .block_on(service.perform_test_update(Some(TARGET_VERSION.to_owned()))) + .expect("accepted"); + drop(runtime); + let status: serde_json::Value = + serde_json::from_slice(&fs::read(fixture.state()).expect("state")).expect("JSON"); + assert_eq!(status["operation"]["status"], "failed"); + assert!(status["operation"]["finishedAt"].is_string()); + assert!(!fixture.lock().exists()); +} + +#[tokio::test] +async fn status_should_recover_abandoned_running_state_without_stealing_a_live_lock() { + let fixture = Fixture::new(); + fs::write(fixture.state(), r#"{"currentVersion":"1.0.0","operation":{"operationId":"abandoned","kind":"update","status":"running","targetVersion":"1.9.9"}}"#).expect("state"); + fs::write(fixture.lock(), "another process").expect("live lock"); + let service = fixture.service_for_url("http://127.0.0.1:1/repos"); + assert_eq!( + service + .update_status() + .await + .expect("locked status") + .operation + .status, + SystemOperationStatus::Running + ); + fs::remove_file(fixture.lock()).expect("release abandoned lock"); + let status = service.update_status().await.expect("recovered status"); + assert_eq!(status.operation.status, SystemOperationStatus::Failed); + assert!(status.operation.finished_at.is_some()); + assert_eq!(status.operation.operation_id.as_deref(), Some("abandoned")); +} diff --git a/backend/crates/gateway-host/tests/workers.rs b/backend/crates/gateway-host/tests/workers.rs index 9ff6ecb3b..f86ff63f1 100644 --- a/backend/crates/gateway-host/tests/workers.rs +++ b/backend/crates/gateway-host/tests/workers.rs @@ -668,6 +668,40 @@ async fn shutdown_timeout_aborts_uncooperative_task_and_drops_guard() { ); } +#[tokio::test] +async fn dropping_shutdown_aborts_current_and_unjoined_workers() { + let task = HungTask { + entered: Arc::new(AtomicBool::new(false)), + notification: Arc::new(Notify::new()), + }; + let leases = Arc::new(FakeLeasePort::default()); + let supervisor = supervisor(); + supervisor + .start( + all_active_plan(task), + leases.clone(), + Arc::new(RecordingDiagnostics::default()), + ) + .expect("complete plan"); + wait_until(|| lock_unpoisoned(&leases.active_guards).len() == ACTIVE_KINDS.len()).await; + + let mut shutdown = Box::pin(supervisor.shutdown(Duration::from_secs(60))); + assert!(futures::poll!(&mut shutdown).is_pending()); + drop(shutdown); + // Supervisor 仍存活，shutdown 自己必须回收当前等待和尚未轮到的所有任务 + wait_until(|| lock_unpoisoned(&leases.active_guards).is_empty()).await; + assert!( + supervisor + .health_source() + .snapshot() + .into_iter() + .all(|entry| matches!( + entry.state, + WorkerRuntimeState::Stopped | WorkerRuntimeState::Disabled + )) + ); +} + #[tokio::test] async fn shutdown_timeout_must_not_rejoin_workers_that_already_exited() { // 挂死任务放在 join 顺序最末的 kind（Retention），保证超时触发时 diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs index aae868186..a3b5aadaa 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/mod.rs @@ -13,7 +13,7 @@ use gateway_admin::{ model::{AdminError, Revision}, ports::plugins::PluginStore, }; -use gateway_core::runtime::extensions::ExtensionSetReference; +use gateway_core::routing::extensions::ExtensionSetReference; use gateway_plugin_sdk::call::management::CommandDescriptor; use crate::{RpcLimits, RpcSession, callback::PluginCallbacks}; diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs b/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs index 03cbd504d..782006158 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs @@ -16,7 +16,7 @@ use gateway_core::{ WebSocketResponseObservation, }, metering::{CostEstimateStatus, CostSource}, - runtime::extensions::ExtensionSetReference, + routing::extensions::ExtensionSetReference, upstream::UpstreamSendState, }; use gateway_plugin_sdk::{ diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/execution.rs similarity index 96% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/execution.rs index 121f16d97..30a11da26 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/execution.rs @@ -28,7 +28,7 @@ use crate::{ callback::{MiddlewareBodyAuthority, MiddlewareCompletionBody, MiddlewareInvocation}, }; -use super::{MiddlewareEntry, PluginRequestPolicyPlan}; +use super::super::{MiddlewareEntry, PluginRequestPolicyPlan}; const MAXIMUM_MAPPED_FRAMES_PER_SOURCE: usize = 64; const MAXIMUM_MAPPED_BYTES_PER_SOURCE: usize = 8 * 1024 * 1024; @@ -80,12 +80,15 @@ impl MiddlewarePlan for PluginRequestPolicyPlan { } fn handle_service( &self, - context: gateway_core::middleware::service::Context, - input: gateway_core::middleware::service::Value, - next: gateway_core::middleware::service::Next, + context: gateway_core::engine::middleware::service::Context, + input: gateway_core::engine::middleware::service::Value, + next: gateway_core::engine::middleware::service::Next, ) -> BoxFuture< 'static, - Result, + Result< + gateway_core::engine::middleware::service::Value, + gateway_core::engine::middleware::service::Error, + >, > { self.chain( MiddlewareMount::Service, @@ -104,12 +107,12 @@ impl MiddlewarePlan for PluginRequestPolicyPlan { } fn handle_websocket( &self, - context: gateway_core::middleware::websocket::Context, - message: gateway_core::middleware::websocket::Message, - next: gateway_core::middleware::websocket::Next, + context: gateway_core::engine::middleware::websocket::Context, + message: gateway_core::engine::middleware::websocket::Message, + next: gateway_core::engine::middleware::websocket::Next, ) -> BoxFuture< 'static, - Result, MiddlewareError>, + Result, MiddlewareError>, > { self.chain( MiddlewareMount::WebSocket, @@ -130,10 +133,11 @@ impl MiddlewarePlan for PluginRequestPolicyPlan { fn handle_http( &self, - context: gateway_core::middleware::http::Context, - request: gateway_core::middleware::http::Request, - next: gateway_core::middleware::http::Next, - ) -> BoxFuture<'static, Result> { + context: gateway_core::engine::middleware::http::Context, + request: gateway_core::engine::middleware::http::Request, + next: gateway_core::engine::middleware::http::Next, + ) -> BoxFuture<'static, Result> + { self.chain( MiddlewareMount::Http, context, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/http.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/http.rs similarity index 97% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/http.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/http.rs index 99586684e..57766a92f 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/http.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/http.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use futures::stream; use gateway_admin::model::plugins::instances::PluginFailurePolicy; -use gateway_core::{engine::middleware::MiddlewareError, middleware::http as core}; +use gateway_core::{engine::middleware::MiddlewareError, engine::middleware::http as core}; use gateway_plugin_sdk::{ Stage, call::middleware::{HANDLE_METHOD, http as wire}, @@ -12,7 +12,7 @@ use gateway_plugin_sdk::{ use http_body::Frame; use http_body_util::{BodyExt as _, StreamBody}; -use super::MiddlewareEntry; +use super::super::MiddlewareEntry; use crate::callback::http_middleware::{Invocation, headers}; pub(super) async fn invoke( diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs new file mode 100644 index 000000000..6b4e3cdd6 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/mod.rs @@ -0,0 +1,6 @@ +//! HTTP、WebSocket、逻辑请求与服务中间件的协议适配入口 + +mod execution; +mod http; +mod service; +mod websocket; diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/service.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/service.rs similarity index 96% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/service.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/service.rs index ef930e6c9..b836a7fd3 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/service.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/service.rs @@ -2,10 +2,10 @@ use std::sync::Arc; -use super::MiddlewareEntry; +use super::super::MiddlewareEntry; use crate::callback::services::Invocation; use gateway_admin::model::plugins::instances::PluginFailurePolicy; -use gateway_core::middleware::service as core; +use gateway_core::engine::middleware::service as core; use gateway_plugin_sdk::{ Stage, call::{middleware::HANDLE_METHOD, services as wire}, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/websocket.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/websocket.rs similarity index 95% rename from backend/crates/gateway-plugin/runtime/src/adapter/policy/websocket.rs rename to backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/websocket.rs index 2daa33bcc..749420dfc 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/websocket.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/middleware/websocket.rs @@ -2,10 +2,10 @@ use std::sync::Arc; -use super::MiddlewareEntry; +use super::super::MiddlewareEntry; use crate::callback::websocket_middleware::Invocation; use gateway_admin::model::plugins::instances::PluginFailurePolicy; -use gateway_core::{engine::middleware::MiddlewareError, middleware::websocket as core}; +use gateway_core::{engine::middleware::MiddlewareError, engine::middleware::websocket as core}; use gateway_plugin_sdk::{ Stage, call::middleware::{HANDLE_METHOD, MiddlewareHeader, websocket as wire}, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs index 3a6ba63cc..b212343f0 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs @@ -1,11 +1,8 @@ //! 编译插件请求策略与中间件绑定，组合为请求级执行计划 -mod http; mod middleware; mod retry; mod route_schedule; -mod service; -mod websocket; use std::{collections::BTreeMap, fmt, sync::Arc, time::Duration}; diff --git a/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs b/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs new file mode 100644 index 000000000..308023a36 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs @@ -0,0 +1,128 @@ +//! 单次宿主回调的资源状态、流式预算与统一终结 + +use super::{CallbackScope, MiddlewareCallback, denied, http_middleware, model}; +use gateway_core::lifecycle::CancellationToken; +use gateway_host::outbound::HttpBody; +use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault}; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::Duration, +}; +use tokio::time::Instant; + +pub(super) struct CallResources { + pub(super) deadline: Instant, + operation_timeout: Duration, + pub(super) cancellation: CancellationToken, + pub(super) scope: Arc, + pub(super) http_resources: Arc, + pub(super) model_bindings: tokio::sync::Mutex< + BTreeMap, + >, + pub(super) state: Mutex, +} + +impl CallResources { + pub(super) fn new( + context: &CallContext, + scope: Arc, + middleware: Option>, + ) -> Self { + Self { + deadline: Instant::now() + Duration::from_millis(context.timeout_ms), + operation_timeout: Duration::from_millis(context.timeout_ms), + cancellation: CancellationToken::new(), + scope, + http_resources: middleware + .as_ref() + .and_then(|middleware| middleware.http_resources()) + .unwrap_or_default(), + model_bindings: tokio::sync::Mutex::new(BTreeMap::new()), + state: Mutex::new(CallState { + middleware, + ..CallState::default() + }), + } + } + + pub(super) fn set_streaming(&self, streaming: bool) { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .resource_stream = streaming; + } + + pub(super) fn close(&self) { + self.cancellation.cancel(); + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.closed = true; + // RPC End 可以先于消费者取完已入队正文到达，这里只撤销新 callback + // 正文包装器持有 invocation，最后一个 owner 释放时再关闭，保留计量和终态信封 + state.middleware.take(); + for stream in state.streams.values() { + stream.close(); + } + state.streams.clear(); + for stream in state.model_streams.values() { + stream.close(); + } + state.model_streams.clear(); + } + + pub(super) fn request_settings( + &self, + ) -> Option { + let middleware = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .middleware + .clone(); + middleware.and_then(|middleware| middleware.request_settings()) + } + + /// 建立受管流后，网络操作各自计时；连接空闲时间不消耗后续操作的预算 + pub(super) fn timeout(&self) -> Result { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.closed { + return Err(denied()); + } + if state.resource_stream { + return Ok(self.operation_timeout); + } + self.deadline + .checked_duration_since(Instant::now()) + .filter(|timeout| !timeout.is_zero()) + .ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed")) + } +} + +#[derive(Default)] +pub(super) struct CallState { + pub(super) closed: bool, + pub(super) resource_stream: bool, + pub(super) middleware: Option>, + pub(super) streams: BTreeMap>, + pub(super) model_streams: BTreeMap>, +} + +pub(super) struct HttpStream { + pub(super) body: tokio::sync::Mutex>, + pub(super) closed: tokio::sync::watch::Sender, +} + +impl HttpStream { + pub(super) fn close(&self) { + self.closed.send_replace(true); + if let Ok(mut body) = self.body.try_lock() { + body.take(); + } + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs index 83cb05274..1c88e9870 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs @@ -1,7 +1,7 @@ //! 主动 HTTP 子调用只进入 API 分派端口，正文复用当前 RPC 的资源池 use super::{CallResources, http_middleware, invalid, services::ServicePorts}; use crate::RpcReply; -use gateway_core::middleware::http as core; +use gateway_core::engine::middleware::http as core; use gateway_plugin_sdk::{CallContext, PluginFault, call::middleware::http as wire}; pub(super) async fn call( diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs index 30d2f5ba6..e06fc05dd 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/mod.rs @@ -9,7 +9,7 @@ use std::{ }; use futures::future::BoxFuture; -use gateway_core::{engine::middleware::MiddlewareError, middleware::http as core}; +use gateway_core::{engine::middleware::MiddlewareError, engine::middleware::http as core}; use gateway_plugin_sdk::{ ErrorCode, PluginFault, call::middleware::{MiddlewareHeader, NEXT_METHOD, http as wire}, @@ -27,13 +27,13 @@ pub(crate) struct Invocation { } struct State { - settings: Option, + settings: Option, request: Option, next: Option, } pub(super) fn resolve_settings( - settings: Option, + settings: Option, value: serde_json::Value, timeout_ms: Option, instance_id: &str, @@ -171,7 +171,7 @@ impl Invocation { } impl MiddlewareCallback for Invocation { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs index d908bfbb5..85af02e17 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_middleware/websocket.rs @@ -2,7 +2,7 @@ use super::*; use crate::callback::websocket_middleware::Invocation as Messages; -use gateway_core::middleware::{http::upgrade, websocket as ws}; +use gateway_core::engine::middleware::{http::upgrade, websocket as ws}; use gateway_plugin_sdk::call::middleware::websocket as message; use gateway_plugin_sdk::call::middleware::{BODY_CLOSE_METHOD, BODY_READ_METHOD}; diff --git a/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs index cbb3a675a..23fae55dd 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/middleware/mod.rs @@ -7,7 +7,7 @@ use futures::future::BoxFuture; use std::sync::{Arc, Mutex}; pub(crate) trait MiddlewareCallback: Send + Sync { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { None } @@ -25,7 +25,7 @@ pub(crate) trait MiddlewareCallback: Send + Sync { } impl MiddlewareCallback for MiddlewareInvocation { - fn request_settings(&self) -> Option { + fn request_settings(&self) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) diff --git a/backend/crates/gateway-plugin/runtime/src/callback/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/mod.rs index 9923bf480..82a654ba1 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/mod.rs @@ -3,6 +3,7 @@ mod accounts; mod admin; mod affinity; +mod call_resources; mod data; pub(crate) mod error; mod facts; @@ -23,19 +24,19 @@ pub(crate) mod websocket_middleware; use std::{ collections::BTreeMap, sync::{Arc, Mutex, Weak}, - time::Duration, }; use futures::future::BoxFuture; use gateway_admin::model::AdminError; use gateway_core::{account::OutboundProxy, lifecycle::CancellationToken}; -use gateway_host::outbound::{HttpBody, HttpClient, NetworkPolicy}; +use gateway_host::outbound::{HttpClient, NetworkPolicy}; use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault}; use tokio::time::Instant; use crate::{CallbackHandler, RpcReply}; pub(crate) use accounts::PluginAccountPortSlot; pub(crate) use affinity::PluginAffinityPortSlot; +use call_resources::{CallResources, HttpStream}; pub(crate) use keys::PluginClientKeyPortSlot; pub(crate) use middleware::{ MiddlewareBinding, MiddlewareBodyAuthority, MiddlewareCallback, MiddlewareCompletionBody, @@ -97,71 +98,6 @@ pub(crate) struct PluginCallbacks { maximum_payload: usize, } -struct CallResources { - deadline: Instant, - operation_timeout: Duration, - cancellation: CancellationToken, - scope: Arc, - http_resources: Arc, - model_bindings: tokio::sync::Mutex< - BTreeMap, - >, - state: Mutex, -} - -impl CallResources { - fn request_settings(&self) -> Option { - let middleware = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .middleware - .clone(); - middleware.and_then(|middleware| middleware.request_settings()) - } - - /// 建立受管流后，网络操作各自计时；连接空闲时间不消耗后续操作的预算 - fn timeout(&self) -> Result { - let state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state.closed { - return Err(denied()); - } - if state.resource_stream { - return Ok(self.operation_timeout); - } - self.deadline - .checked_duration_since(Instant::now()) - .filter(|timeout| !timeout.is_zero()) - .ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed")) - } -} - -#[derive(Default)] -struct CallState { - closed: bool, - resource_stream: bool, - middleware: Option>, - streams: BTreeMap>, - model_streams: BTreeMap>, -} - -struct HttpStream { - body: tokio::sync::Mutex>, - closed: tokio::sync::watch::Sender, -} - -impl HttpStream { - fn close(&self) { - self.closed.send_replace(true); - if let Ok(mut body) = self.body.try_lock() { - body.take(); - } - } -} - impl PluginCallbacks { pub(crate) fn new( instance: &gateway_admin::model::plugins::instances::PluginInstance, @@ -411,21 +347,7 @@ impl CallbackHandler for PluginCallbacks { .unwrap_or_else(std::sync::PoisonError::into_inner) .insert( context.call_id, - Arc::new(CallResources { - deadline: Instant::now() + Duration::from_millis(context.timeout_ms), - operation_timeout: Duration::from_millis(context.timeout_ms), - cancellation: CancellationToken::new(), - scope, - http_resources: middleware - .as_ref() - .and_then(|middleware| middleware.http_resources()) - .unwrap_or_default(), - model_bindings: tokio::sync::Mutex::new(BTreeMap::new()), - state: Mutex::new(CallState { - middleware, - ..CallState::default() - }), - }), + Arc::new(CallResources::new(context, scope, middleware)), ); } @@ -436,10 +358,7 @@ impl CallbackHandler for PluginCallbacks { .unwrap_or_else(std::sync::PoisonError::into_inner) .get(&context.call_id) { - call.state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .resource_stream = context.resource_stream; + call.set_streaming(context.resource_stream); } } @@ -450,25 +369,7 @@ impl CallbackHandler for PluginCallbacks { .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(&context.call_id) { - call.cancellation.cancel(); - let mut state = call - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state.closed = true; - // RPC End 可以先于消费者取完已入队的中间件正文到达 - // 这里只撤销新的 - // callback 入口；正文包装器持有 invocation，最后一个 owner 释放时再关闭 - // 下游，避免提前丢失尚未搬运的计量/终态信封 - state.middleware.take(); - for stream in state.streams.values() { - stream.close(); - } - state.streams.clear(); - for stream in state.model_streams.values() { - stream.close(); - } - state.model_streams.clear(); + call.close(); } } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/services.rs b/backend/crates/gateway-plugin/runtime/src/callback/services.rs index 9496c779d..01814dd15 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/services.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/services.rs @@ -3,8 +3,8 @@ use super::{CallbackScope, MiddlewareCallback, invalid}; use crate::RpcReply; use futures::future::BoxFuture; -use gateway_admin::{model::AdminError, service::Registry}; -use gateway_core::middleware::service as core; +use gateway_admin::{model::AdminError, public_service::Registry}; +use gateway_core::engine::middleware::service as core; use gateway_plugin_sdk::{ CallContext, PluginFault, call::{middleware::NEXT_METHOD, services as wire}, @@ -14,13 +14,13 @@ use std::sync::{Arc, Mutex, OnceLock, Weak}; #[derive(Default)] pub(crate) struct ServicePorts { registry: OnceLock>, - http: OnceLock>, + http: OnceLock>, } impl ServicePorts { pub(crate) fn bind_http( &self, - dispatcher: &Arc, + dispatcher: &Arc, ) -> Result<(), AdminError> { self.http .set(Arc::downgrade(dispatcher)) @@ -29,7 +29,7 @@ impl ServicePorts { pub(super) fn http( &self, - ) -> Result, PluginFault> { + ) -> Result, PluginFault> { self.http.get().and_then(Weak::upgrade).ok_or_else(|| { PluginFault::new( gateway_plugin_sdk::ErrorCode::Fault, @@ -64,7 +64,7 @@ impl ServicePorts { })?; let call = registry.call(&request.operation, request.input); let origin = scope.origin.as_ref(); - let parent = gateway_admin::service::Origin { + let parent = gateway_admin::public_service::Origin { request_id: origin.map_or_else( || { context @@ -80,11 +80,11 @@ impl ServicePorts { ), cancellation, extensions: scope.child_extensions(&context.instance_id)?, - plan: origin.map_or(gateway_admin::service::Plan::Current, |origin| { - gateway_admin::service::Plan::Frozen(Some(origin.plan.clone())) + plan: origin.map_or(gateway_admin::public_service::Plan::Current, |origin| { + gateway_admin::public_service::Plan::Frozen(Some(origin.plan.clone())) }), }; - let result = gateway_admin::service::scope(parent, call).await; + let result = gateway_admin::public_service::scope(parent, call).await; encode(result) } } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs b/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs index 4fa7da028..190a42474 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/websocket_middleware.rs @@ -3,7 +3,7 @@ use super::MiddlewareCallback; use crate::RpcReply; use bytes::Bytes; use futures::future::BoxFuture; -use gateway_core::middleware::websocket as core; +use gateway_core::engine::middleware::websocket as core; use gateway_plugin_sdk::{ ErrorCode, PluginFault, call::middleware::{ diff --git a/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs b/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs index c9d3d187a..c030b4c89 100644 --- a/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs +++ b/backend/crates/gateway-plugin/runtime/src/compatibility/fast.rs @@ -2,7 +2,7 @@ use gateway_core::{ account::FastMode, - settings::{ExecutionSettings, RequestSettings}, + routing::request_settings::{ExecutionSettings, RequestSettings}, }; use gateway_plugin_sdk::{ Capability, ErrorCode, PluginFault, diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs new file mode 100644 index 000000000..27a1a63ed --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/diagnostics.rs @@ -0,0 +1,255 @@ +//! 将已发布与准备中的实例事实投影为运行诊断，不参与候选提交 + +use super::super::restart_circuit::RestartIdentity; +use super::{ + PluginRuntime, + set::{PreparedSet, instance_fingerprint}, +}; +use async_trait::async_trait; +use gateway_admin::model::{ + AdminError, AdminErrorKind, Revision, + plugins::instances::{ + PluginInstance, PluginInstanceRuntime, PluginInstanceRuntimeFailure, + PluginInstanceRuntimeStatus, PluginInstanceSnapshot, + }, +}; +use gateway_core::routing::extensions::ExtensionSetReference; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{Arc, Weak}, +}; + +const MAXIMUM_PREPARATION_DIAGNOSTICS: usize = 8; +#[derive(Clone)] +pub(super) enum PreparationDiagnostic { + Preparing, + Prepared { set_id: String }, + Failed(PluginInstanceRuntimeFailure), +} +impl PluginRuntime { + pub(super) fn record_preparation(&self, revision: u64, diagnostic: PreparationDiagnostic) { + let mut diagnostics = self + .preparation_diagnostics + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + diagnostics.insert(revision, diagnostic); + while diagnostics.len() > MAXIMUM_PREPARATION_DIAGNOSTICS { + diagnostics.pop_first(); + } + } + + fn preparation_diagnostic(&self, revision: u64) -> Option { + self.preparation_diagnostics + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&revision) + .cloned() + } +} +#[async_trait] +impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for PluginRuntime { + async fn runtime_diagnostics( + &self, + snapshot: &PluginInstanceSnapshot, + published_revision: Option, + published: Option<&ExtensionSetReference>, + ) -> Option> { + let sets = { + let mut prepared = self.prepared.lock().await; + prepared.retain(|_, set| set.strong_count() > 0); + prepared + .values() + .filter_map(Weak::upgrade) + .collect::>() + }; + let open_restart_circuits = snapshot + .instances + .iter() + .filter(|instance| instance.enabled) + .filter_map(|instance| { + let identity = + RestartIdentity::new(instance.id.clone(), instance_fingerprint(instance).ok()?); + self.restart_circuit_is_open(&identity) + .then(|| instance.id.clone()) + }) + .collect::>(); + let published_set = published + .and_then(|reference| sets.iter().find(|set| set.id == *reference.id()).cloned()); + let published_ready = published.is_some_and(ExtensionSetReference::can_serve); + let preparation = self.preparation_diagnostic(snapshot.config_revision.get()); + let projection = InstanceRuntimeProjection { + target_revision: snapshot.config_revision, + published_revision, + published_ready, + published_set: published_set.as_deref(), + sets: &sets, + preparation: preparation.as_ref(), + open_restart_circuits: &open_restart_circuits, + }; + Some( + snapshot + .instances + .iter() + .map(|instance| { + let runtime = projection.project(instance); + (instance.id.clone(), runtime) + }) + .collect(), + ) + } +} +struct InstanceRuntimeProjection<'a> { + target_revision: Revision, + published_revision: Option, + published_ready: bool, + published_set: Option<&'a PreparedSet>, + sets: &'a [Arc], + preparation: Option<&'a PreparationDiagnostic>, + open_restart_circuits: &'a BTreeSet, +} + +impl InstanceRuntimeProjection<'_> { + fn project(&self, expected: &PluginInstance) -> PluginInstanceRuntime { + let active = self.published_set.and_then(|set| { + set.sessions + .iter() + .find(|instance| instance.instance_id == expected.id) + }); + let actual_matches = active.is_some_and(|instance| { + instance.revision == expected.revision + && instance.artifact_sha256 == expected.artifact_sha256 + }); + let target_is_published = self.published_revision == Some(self.target_revision.get()); + let candidate_set = match self.preparation { + Some(PreparationDiagnostic::Prepared { set_id }) => Some(set_id.as_str()), + _ => None, + }; + let published_set_id = self.published_set.map(|set| set.id.as_str()); + let mut draining_revisions = BTreeSet::new(); + for set in self.sets { + if Some(set.id.as_str()) == published_set_id || Some(set.id.as_str()) == candidate_set { + continue; + } + if let Some(instance) = set + .sessions + .iter() + .find(|instance| instance.instance_id == expected.id) + && instance.revision.get() < expected.revision.get() + { + draining_revisions.insert(instance.revision.get()); + } + } + if !expected.enabled + && let Some(instance) = active + && instance.revision != expected.revision + { + draining_revisions.insert(instance.revision.get()); + } + + let diagnostic = active.map(|instance| instance.session.diagnostic()); + let mut failure = None; + let status = if !expected.enabled { + if active.is_some() || !draining_revisions.is_empty() { + PluginInstanceRuntimeStatus::Draining + } else { + PluginInstanceRuntimeStatus::Disabled + } + } else if target_is_published + && let Some(reason) = self + .published_set + .and_then(|set| set.failures.get(&expected.id)) + { + failure = Some(reason.clone()); + PluginInstanceRuntimeStatus::PreparationFailed + } else if self.open_restart_circuits.contains(&expected.id) { + let message = match diagnostic { + Some(crate::rpc::RpcSessionDiagnostic::Failed { message, .. }) => { + format!("{message}，已暂停自动重启") + } + _ => "插件实例连续异常退出，已暂停自动重启".to_owned(), + }; + failure = Some(runtime_failure("restart_circuit_open", &message)); + PluginInstanceRuntimeStatus::Faulted + } else if actual_matches + && let Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) = diagnostic + { + failure = Some(runtime_failure(code, message)); + PluginInstanceRuntimeStatus::Faulted + } else if !target_is_published { + match self.preparation { + Some(PreparationDiagnostic::Preparing) => PluginInstanceRuntimeStatus::Preparing, + Some(PreparationDiagnostic::Failed(reason)) => { + failure = Some(reason.clone()); + PluginInstanceRuntimeStatus::PreparationFailed + } + Some(PreparationDiagnostic::Prepared { .. }) | None => { + PluginInstanceRuntimeStatus::AwaitingPublication + } + } + } else if !actual_matches { + failure = Some(runtime_failure( + "published_instance_mismatch", + "已发布插件实例与期望版本不一致", + )); + PluginInstanceRuntimeStatus::Blocked + } else { + match diagnostic { + Some(crate::rpc::RpcSessionDiagnostic::Ready) if self.published_ready => { + PluginInstanceRuntimeStatus::Running + } + Some(crate::rpc::RpcSessionDiagnostic::Ready) => { + failure = Some(runtime_failure( + "published_set_unready", + "同一发布集合中存在未就绪插件实例", + )); + PluginInstanceRuntimeStatus::Blocked + } + Some(crate::rpc::RpcSessionDiagnostic::Quiescing) => { + PluginInstanceRuntimeStatus::Draining + } + Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) => { + failure = Some(runtime_failure(code, message)); + PluginInstanceRuntimeStatus::Faulted + } + None => { + failure = Some(runtime_failure( + "instance_not_published", + "已发布集合缺少该插件实例", + )); + PluginInstanceRuntimeStatus::Blocked + } + } + }; + + PluginInstanceRuntime { + status, + actual_revision: active.map(|instance| instance.revision.get()), + actual_artifact_sha256: active.map(|instance| instance.artifact_sha256.clone()), + failure, + draining_revisions: draining_revisions.into_iter().collect(), + } + } +} + +fn runtime_failure(code: &str, message: &str) -> PluginInstanceRuntimeFailure { + PluginInstanceRuntimeFailure { + code: code.to_owned(), + message: message.to_owned(), + } +} + +pub(super) fn runtime_failure_from_admin(error: &AdminError) -> PluginInstanceRuntimeFailure { + let code = match error.kind() { + AdminErrorKind::Invalid => "invalid", + AdminErrorKind::Unauthorized => "unauthorized", + AdminErrorKind::Forbidden => "forbidden", + AdminErrorKind::NotFound => "not_found", + AdminErrorKind::Conflict => "conflict", + AdminErrorKind::RateLimited => "rate_limited", + AdminErrorKind::BadGateway => "bad_gateway", + AdminErrorKind::UpstreamResultUnknown => "upstream_result_unknown", + AdminErrorKind::Unavailable => "unavailable", + AdminErrorKind::Internal => "internal", + }; + runtime_failure(code, error.message()) +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs index 1dd82987d..5331e1369 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs @@ -1,11 +1,20 @@ //! 准备单个插件实例的进程、RPC 会话与能力贡献 +use super::super::restart_circuit::RestartIdentity; use super::{ - AdminError, Arc, Duration, Handshake, PluginCallbackPorts, PluginCallbacks, PluginInstance, - PluginPrivateState, PluginRuntime, PluginStateStoreErrorKind, PreparedContributions, - PreparedInstance, Registration, RestartIdentity, Revision, RpcSession, Stage, ValidatedPackage, - instance_fingerprint, + PluginRuntime, + set::{PreparedContributions, PreparedInstance, instance_fingerprint}, }; +use crate::{ + RpcSession, ValidatedPackage, + callback::{PluginCallbackPorts, PluginCallbacks, private_state::PluginPrivateState}, +}; +use gateway_admin::{ + model::{AdminError, Revision, plugins::instances::PluginInstance}, + ports::plugins::PluginStateStoreErrorKind, +}; +use gateway_plugin_sdk::{Handshake, Stage, call::registration::Registration}; +use std::{sync::Arc, time::Duration}; impl PluginRuntime { pub(super) async fn prepare_instance( diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs index a91e1012c..edb32f731 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs @@ -1,6 +1,7 @@ //! 对账只从已发布集合开始；通知合并，实例串行，失败和停用不阻塞其他实例 -use super::{PluginRuntime, PreparedSet, RpcSession}; +use super::{PluginRuntime, set::PreparedSet}; +use crate::RpcSession; use futures::{StreamExt as _, future::BoxFuture}; use gateway_core::{ lifecycle::CancellationToken, diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs new file mode 100644 index 000000000..2a18fddec --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/management.rs @@ -0,0 +1,137 @@ +//! 按发布集合适配插件管理接口与命令行会话 + +use super::PluginRuntime; +use async_trait::async_trait; +use gateway_admin::model::AdminError; +use gateway_core::routing::extensions::ExtensionSetReference; + +impl PluginRuntime { + /// CLI 复用同一候选准备与弱索引；读取帮助不会发布 Core 快照或执行命令 + pub async fn prepare_command_line(&self) -> Result { + let snapshot = self + .store + .load_instances() + .await + .map_err(|_| AdminError::unavailable("无法读取已安装插件"))?; + let reference = self.prepare_snapshot(snapshot, None).await?; + let set = self.prepared_set(&reference).await?; + Ok(crate::PluginCommandSession::new( + reference, + set.commands.clone(), + set.sessions + .iter() + .map(|instance| instance.session.clone()) + .collect(), + self.store.clone(), + self.config.rpc_limits, + )) + } +} +#[async_trait] +impl gateway_admin::ports::plugin_management::PluginManagement for PluginRuntime { + async fn validate_target( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + ) -> Result<(), AdminError> { + let set = self.prepared_set(published).await?; + set.management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))?; + + Ok(()) + } + + async fn start_callback( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + command: gateway_admin::model::plugins::management::StartPluginManagementCallback, + context: &gateway_admin::model::auth::AdminRequestContext, + ) -> Result + { + let store = self + .oauth_pending + .as_ref() + .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .start_callback(store.as_ref(), command, context) + .await + } + + async fn callback( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + state: &str, + request: gateway_admin::model::plugins::management::PluginManagementRequest, + ) -> Result + { + let store = self + .oauth_pending + .as_ref() + .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .callback(store.as_ref(), state, request, self.config.rpc_limits) + .await + } + + async fn views( + &self, + published: &ExtensionSetReference, + ) -> Result, AdminError> + { + Ok(self + .prepared_set(published) + .await? + .management + .iter() + .map(|entry| entry.view.clone()) + .collect()) + } + + async fn resource( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + path: &str, + public: bool, + ) -> Result + { + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .resource(path, public) + } + + async fn handle( + &self, + published: &ExtensionSetReference, + target: &gateway_admin::model::plugins::management::PluginManagementTarget, + request: gateway_admin::model::plugins::management::PluginManagementRequest, + ) -> Result + { + self.prepared_set(published) + .await? + .management + .iter() + .find(|entry| &entry.view.target == target) + .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? + .handle(request, self.config.rpc_limits) + .await + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs index f6cf6cfdc..c22a17afd 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/mod.rs @@ -1,1267 +1,12 @@ -//! 组装插件运行时端口，准备并持有可发布的实例与能力集合 +//! 插件运行时端口与代次准备的公共入口 +mod diagnostics; mod instance; mod maintenance; +mod management; +mod runtime; +mod set; +mod snapshot; +mod state; -use std::{ - collections::{BTreeMap, BTreeSet}, - path::PathBuf, - sync::{ - Arc, Mutex as SyncMutex, Weak, - atomic::{AtomicBool, Ordering}, - }, - time::Duration, -}; - -use async_trait::async_trait; -use futures::future::BoxFuture; -use gateway_admin::{ - model::{ - AdminError, AdminErrorKind, Revision, - plugins::{ - instances::{ - PluginInstance, PluginInstanceRuntime, PluginInstanceRuntimeFailure, - PluginInstanceRuntimeStatus, PluginInstanceSnapshot, - }, - state::{ - ApplyPluginStateMigration, PluginStateConfiguration, PluginStateMigrationAction, - PluginStateMigrationChange, PluginStateTransition, - }, - }, - }, - ports::{ - plugin_accounts::PluginAccountAccess, - plugin_client_keys::PluginClientKeyAccess, - plugins::{PluginPreparation, PluginStateStore, PluginStateStoreErrorKind, PluginStore}, - }, -}; -use gateway_core::{ - engine::{ - authentication::FrontendAuthenticationExtensionIndex, - middleware::{MiddlewareExtensionIndex, MiddlewarePlan}, - nested::{AffinityLookupPort, NestedModelExecutionPort}, - observation::{RequestObserverExtensionIndex, RequestObserverPlan}, - policy::{RequestPolicyExtensionIndex, RequestPolicyPlan}, - }, - provider_ports::OAuthPendingFlowPort, - routing::ConfigRevision, - runtime::extensions::{ - ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetLease, - ExtensionSetReference, - }, -}; -use gateway_plugin_sdk::call::registration::Registration; -use gateway_plugin_sdk::{ - Handshake, Stage, - call::host::{ - StateMigrationChange, StateMigrationRecord, StateMigrationRequest, StateMigrationResult, - }, -}; -use secrecy::ExposeSecret as _; -use sha2::{Digest as _, Sha256}; -use tokio::sync::{Mutex, Semaphore}; - -use crate::{ - PackageLimits, RpcLimits, RpcSession, ValidatedPackage, - callback::{ - PluginAccountPortSlot, PluginAffinityPortSlot, PluginCallbackPorts, PluginCallbacks, - PluginClientKeyPortSlot, PluginModelPortSlot, private_state::PluginPrivateState, - }, -}; - -use super::restart_circuit::{PluginRestartCircuitConfig, RestartCircuits, RestartIdentity}; - -const RUNTIME_SHUTDOWN_GRACE: Duration = Duration::from_secs(2); -const MAXIMUM_PREPARATION_DIAGNOSTICS: usize = 8; - -pub struct PluginRuntimeConfig { - pub cache_directory: PathBuf, - pub host_version: semver::Version, - pub package_limits: PackageLimits, - pub rpc_limits: RpcLimits, - pub restart_circuit: PluginRestartCircuitConfig, -} - -/// 只索引有主人的候选；当前代次仅由 Core 的发布视图持有 -pub struct PluginRuntime { - service_ports: Arc, - pub(super) store: Arc, - state: Arc, - config: PluginRuntimeConfig, - prepare_lock: Mutex<()>, - prepared: Mutex>>, - preparation_diagnostics: SyncMutex>, - restart_circuits: Arc, - shutting_down: Arc, - shutdown_lock: Mutex<()>, - processes: Arc, - validators: Arc, - log_slots: Arc, - pub(super) account_ports: Arc, - client_key_ports: Arc, - resource_ports: Arc, - model_ports: Arc, - affinity_ports: Arc, - observers: RequestObserverExtensionIndex, - policies: RequestPolicyExtensionIndex, - middleware: MiddlewareExtensionIndex, - authentications: FrontendAuthenticationExtensionIndex, - oauth_pending: Option>, - http: Arc, -} - -#[derive(Clone)] -enum PreparationDiagnostic { - Preparing, - Prepared { set_id: String }, - Failed(PluginInstanceRuntimeFailure), -} - -pub(super) struct PreparedSet { - id: ExtensionSetId, - _observers: Option>, - _policies: Option>, - _middleware: Option>, - upstream_adapters: Option>, - _authentication: - Option>, - sessions: Vec, - failures: BTreeMap, - shutting_down: Arc, - commands: Vec>, - management: Vec, - model_aliases: Vec, -} - -#[derive(Default)] -struct PreparedContributions { - sessions: Vec, - observer_entries: Vec, - commands: Vec>, - management: Vec, - policy_entries: Vec, - upstream_entries: Vec, - authentication_entries: - Vec, - model_aliases: Vec, -} - -impl PreparedContributions { - fn append(&mut self, other: Self) { - self.sessions.extend(other.sessions); - self.observer_entries.extend(other.observer_entries); - self.commands.extend(other.commands); - self.management.extend(other.management); - self.model_aliases.extend(other.model_aliases); - self.policy_entries.extend(other.policy_entries); - self.upstream_entries.extend(other.upstream_entries); - self.authentication_entries - .extend(other.authentication_entries); - } -} - -struct PreparedInstance { - instance_id: String, - artifact_sha256: String, - revision: Revision, - session: Arc, - private_state: Arc, - maintenance: bool, -} - -impl PreparedSet { - fn sessions_ready(&self) -> bool { - self.sessions - .iter() - .all(|instance| instance.session.is_ready()) - } -} - -impl ExtensionSetLease for PreparedSet { - fn upstream_adapters( - &self, - ) -> Option> { - self.upstream_adapters.clone() - } - - fn is_ready(&self) -> bool { - self.sessions_ready() && self.can_serve() - } - - fn model_aliases(&self) -> &[gateway_core::routing::ContributedModelAlias] { - &self.model_aliases - } - - fn can_serve(&self) -> bool { - // 发布的是可用能力和故障绑定组成的完整计划；单个进程退出不能使原生转发失效 - !self.shutting_down.load(Ordering::Acquire) - } -} - -impl Drop for PreparedSet { - fn drop(&mut self) { - for instance in self.sessions.drain(..) { - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - let session = instance.session; - session.quiesce(); - session.shutdown(Duration::from_secs(1)).await; - }); - } - } - } -} - -impl PluginRuntime { - #[must_use] - pub fn new( - store: Arc, - state: Arc, - config: PluginRuntimeConfig, - http: Arc, - processes: Arc, - ) -> Self { - let restart_circuit = config.restart_circuit; - Self { - store, - state, - log_slots: Arc::new(Semaphore::new( - config - .rpc_limits - .maximum_callbacks - .min(Semaphore::MAX_PERMITS), - )), - config, - prepare_lock: Mutex::new(()), - prepared: Mutex::new(BTreeMap::new()), - preparation_diagnostics: SyncMutex::new(BTreeMap::new()), - restart_circuits: RestartCircuits::new(restart_circuit), - shutting_down: Arc::new(AtomicBool::new(false)), - shutdown_lock: Mutex::new(()), - processes, - validators: Arc::new(Semaphore::new(2)), - service_ports: Arc::new(crate::callback::services::ServicePorts::default()), - account_ports: Arc::new(PluginAccountPortSlot::new()), - client_key_ports: Arc::new(PluginClientKeyPortSlot::new()), - resource_ports: Arc::new(crate::callback::PluginResourcePorts::new()), - model_ports: Arc::new(PluginModelPortSlot::new()), - affinity_ports: Arc::new(PluginAffinityPortSlot::new()), - observers: RequestObserverExtensionIndex::default(), - policies: RequestPolicyExtensionIndex::default(), - middleware: MiddlewareExtensionIndex::default(), - authentications: FrontendAuthenticationExtensionIndex::default(), - oauth_pending: None, - http, - } - } - - #[must_use] - pub fn with_oauth_pending(mut self, port: Arc) -> Self { - self.oauth_pending = Some(port); - self - } - - /// API 完成组装后绑定唯一内部 HTTP 分派端口；不持有服务端强引用 - pub fn bind_http( - &self, - dispatcher: &Arc, - ) -> Result<(), AdminError> { - self.service_ports.bind_http(dispatcher) - } - - pub fn bind_services( - &self, - registry: &Arc, - ) -> Result<(), AdminError> { - self.service_ports.bind(registry) - } - - pub fn bind_account_ports( - &self, - access: &Arc, - ) -> Result<(), AdminError> { - self.account_ports.bind(access) - } - - /// Key 目录与预算端口由 Admin 组合并保活，完整字段由对应公开合同返回 - pub fn bind_client_key_ports( - &self, - access: &Arc, - ) -> Result<(), AdminError> { - self.client_key_ports.bind(access) - } - - /// Core 激活后一次性绑定嵌套执行与亲和查询端口；Runtime 仅保存 Weak，避免组合根强环 - pub fn bind_model_ports( - &self, - models: &Arc, - affinity: &Arc, - ) -> Result<(), AdminError> { - self.model_ports.bind(models)?; - self.affinity_ports.bind(affinity) - } - - /// 调用方必须先停止接收 HTTP 与 Worker；这里拒绝新候选并等待全部插件 I/O 退出 - pub async fn shutdown(&self) { - self.shutting_down.store(true, Ordering::Release); - let _shutdown = self.shutdown_lock.lock().await; - let _prepare = self.prepare_lock.lock().await; - let sets = { - let mut prepared = self.prepared.lock().await; - prepared.retain(|_, set| set.strong_count() > 0); - let sets = prepared - .values() - .filter_map(Weak::upgrade) - .collect::>(); - prepared.clear(); - sets - }; - let mut sessions = Vec::new(); - for set in &sets { - for instance in &set.sessions { - if !sessions - .iter() - .any(|session| Arc::ptr_eq(session, &instance.session)) - { - sessions.push(Arc::clone(&instance.session)); - } - } - } - shutdown_sessions(sessions, RUNTIME_SHUTDOWN_GRACE).await; - } - - #[must_use] - pub fn observer_registry(&self) -> RequestObserverExtensionIndex { - self.observers.clone() - } - - #[must_use] - pub fn policy_registry(&self) -> RequestPolicyExtensionIndex { - self.policies.clone() - } - - #[must_use] - pub fn middleware_registry(&self) -> MiddlewareExtensionIndex { - self.middleware.clone() - } - - #[must_use] - pub fn frontend_authentication_registry(&self) -> FrontendAuthenticationExtensionIndex { - self.authentications.clone() - } - - fn record_preparation(&self, revision: u64, diagnostic: PreparationDiagnostic) { - let mut diagnostics = self - .preparation_diagnostics - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - diagnostics.insert(revision, diagnostic); - while diagnostics.len() > MAXIMUM_PREPARATION_DIAGNOSTICS { - diagnostics.pop_first(); - } - } - - fn preparation_diagnostic(&self, revision: u64) -> Option { - self.preparation_diagnostics - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .get(&revision) - .cloned() - } - - fn restart_circuit_is_open(&self, identity: &RestartIdentity) -> bool { - self.restart_circuits.is_open(identity) - } - - async fn prepare_snapshot( - &self, - snapshot: PluginInstanceSnapshot, - required_revision: Option, - ) -> Result { - let target_revision = snapshot.config_revision.get(); - self.record_preparation(target_revision, PreparationDiagnostic::Preparing); - let result = self.prepare_inner(snapshot, required_revision).await; - match &result { - Ok(reference) => self.record_preparation( - target_revision, - PreparationDiagnostic::Prepared { - set_id: reference.id().as_str().to_owned(), - }, - ), - Err(error) => self.record_preparation( - target_revision, - PreparationDiagnostic::Failed(runtime_failure_from_admin(error)), - ), - } - result - } - - async fn prepare_inner( - &self, - snapshot: PluginInstanceSnapshot, - required_revision: Option, - ) -> Result { - if self.shutting_down.load(Ordering::Acquire) { - return Err(AdminError::unavailable("插件运行时正在关闭")); - } - if snapshot.instances.len() > 64 { - return Err(AdminError::invalid("最多配置 64 个插件实例")); - } - let fingerprint = fingerprint(&snapshot)?; - // 准备串行化与弱索引分锁；诊断和已发布请求不会等待外部 I/O - let _prepare = self.prepare_lock.lock().await; - if self.shutting_down.load(Ordering::Acquire) { - return Err(AdminError::unavailable("插件运行时正在关闭")); - } - let existing = { - let mut prepared = self.prepared.lock().await; - prepared.retain(|_, set| set.strong_count() > 0); - prepared - .get(&fingerprint) - .and_then(Weak::upgrade) - .filter(|set| set.sessions_ready()) - }; - if let Some(set) = existing { - if snapshot.instances.iter().any(|instance| { - required_revision == Some(instance.revision) - && set.failures.contains_key(&instance.id) - }) { - return Err(AdminError::unavailable("插件启动失败，请重新启用后重试")); - } - return Ok(ExtensionSetReference::new(set.id.clone(), set)); - } - let mut contributions = PreparedContributions::default(); - let mut failures = BTreeMap::new(); - let mut identities = BTreeSet::new(); - for instance in snapshot - .instances - .iter() - .filter(|instance| instance.enabled) - { - if !identities.insert(instance.id.clone()) { - return Err(AdminError::invalid("插件实例 ID 重复")); - } - let result = tokio::time::timeout( - // 沿用准备阶段原有的 30 秒预算；打包校验与能力恢复也在预算内 - Duration::from_secs(30), - self.prepare_instance(instance.clone(), snapshot.config_revision), - ) - .await - .unwrap_or_else(|_| Err(AdminError::unavailable("插件启动超时"))); - match result { - Ok(candidate) => contributions.append(candidate), - Err(error) if required_revision == Some(instance.revision) => return Err(error), - Err(error) => { - // 目录没有拒绝请求用的 binding，不能在恢复失败时悄悄撤销仍启用的别名 - // 不复制其他候选的注册结果；让发布事务保留调用方持有的旧有效快照 - let has_retained_catalog = self - .prepared - .lock() - .await - .values() - .filter_map(Weak::upgrade) - .any(|set| { - set.model_aliases - .iter() - .any(|alias| alias.owner == instance.id) - }); - if has_retained_catalog { - return Err(AdminError::unavailable( - "已启用的模型目录插件未就绪，请修复或停用后重试", - )); - } - // 恢复失败只隔离该实例；绑定的拒绝策略仍保留，不能绕过认证或必需处理 - contributions - .policy_entries - .extend(crate::adapter::policy::unavailable_entries(instance)?); - contributions.upstream_entries.extend( - crate::adapter::upstream_adapter::unavailable_entries(instance)?, - ); - if let Some(entry) = - crate::adapter::frontend_authentication::unavailable_entry(instance) - { - contributions.authentication_entries.push(entry); - } - failures.insert(instance.id.clone(), runtime_failure_from_admin(&error)); - } - } - } - let PreparedContributions { - sessions, - observer_entries, - commands, - management, - policy_entries, - upstream_entries, - authentication_entries, - model_aliases, - } = contributions; - let mut alias_owners = BTreeMap::new(); - for alias in &model_aliases { - if let Some(previous) = alias_owners.insert(&alias.id, &alias.owner) { - return Err(AdminError::invalid(format!( - "模型别名 {} 在插件实例 {} 与 {} 之间冲突", - alias.id, previous, alias.owner - ))); - } - } - let id = ExtensionSetId::new(uuid::Uuid::new_v4().to_string()) - .map_err(|_| AdminError::internal("插件集合 ID 无效"))?; - let observers = crate::adapter::observer::PluginObserverPlan::compile( - observer_entries, - self.config.rpc_limits.maximum_calls, - self.config.rpc_limits.maximum_call_timeout, - self.config.rpc_limits.maximum_buffered_body_bytes, - ) - .map(|plan| { - self.observers - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件观察计划注册冲突")) - }) - .transpose()?; - let policy_plan = crate::adapter::policy::PluginRequestPolicyPlan::compile( - policy_entries, - self.config.rpc_limits.maximum_call_timeout, - )?; - let policies = policy_plan - .as_ref() - .filter(|plan| plan.has_request_policy()) - .cloned() - .map(|plan| { - self.policies - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件请求策略计划注册冲突")) - }) - .transpose()?; - let middleware = policy_plan - .filter(|plan| plan.has_middleware()) - .map(|plan| { - self.middleware - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件中间件计划注册冲突")) - }) - .transpose()?; - let authentication = - crate::adapter::frontend_authentication::PluginFrontendAuthenticationPlan::compile( - authentication_entries, - self.config.rpc_limits.maximum_call_timeout, - )? - .map(|plan| { - self.authentications - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件入口认证计划注册冲突")) - }) - .transpose()?; - let set = Arc::new(PreparedSet { - id, - _observers: observers, - _policies: policies, - _middleware: middleware, - upstream_adapters: - crate::adapter::upstream_adapter::PluginUpstreamAdapterPlan::compile( - upstream_entries, - )?, - _authentication: authentication, - sessions, - commands, - management, - failures, - model_aliases, - shutting_down: self.shutting_down.clone(), - }); - self.prepared - .lock() - .await - .insert(fingerprint, Arc::downgrade(&set)); - Ok(ExtensionSetReference::new(set.id.clone(), set)) - } - - /// CLI 复用同一候选准备与弱索引；读取帮助不会发布 Core 快照或执行命令 - pub async fn prepare_command_line(&self) -> Result { - let snapshot = self - .store - .load_instances() - .await - .map_err(|_| AdminError::unavailable("无法读取已安装插件"))?; - let reference = self.prepare_snapshot(snapshot, None).await?; - let set = self.prepared_set(&reference).await?; - Ok(crate::PluginCommandSession::new( - reference, - set.commands.clone(), - set.sessions - .iter() - .map(|instance| instance.session.clone()) - .collect(), - self.store.clone(), - self.config.rpc_limits, - )) - } - - pub(super) async fn prepared_set( - &self, - reference: &ExtensionSetReference, - ) -> Result, AdminError> { - self.prepared - .lock() - .await - .values() - .filter_map(Weak::upgrade) - .find(|set| set.id == *reference.id()) - .ok_or_else(|| AdminError::conflict("插件候选已过期，请重试")) - } - - async fn run_state_migration( - &self, - prepared: &ExtensionSetReference, - transition: PluginStateTransition, - ) -> Result<(), AdminError> { - let set = self.prepared_set(prepared).await?; - let instance = set - .sessions - .iter() - .find(|instance| { - instance.instance_id == transition.instance_id - && instance.artifact_sha256 == transition.artifact_sha256 - }) - .ok_or_else(|| AdminError::conflict("状态迁移候选插件已失效"))?; - for namespace in &transition.namespaces { - let mut batches = 0usize; - loop { - batches += 1; - if batches > 200 { - return Err(AdminError::invalid("插件状态迁移超过有界批次数")); - } - let batch = self - .state - .migration_batch(&transition.id, &namespace.namespace, 64) - .await - .map_err(map_state_admin_error)?; - let expected_keys = batch - .records - .iter() - .map(|record| record.key.clone()) - .collect::>(); - if batch.records.is_empty() { - self.state - .apply_migration_batch(ApplyPluginStateMigration { - transition_id: transition.id.clone(), - namespace: namespace.namespace.clone(), - cursor: batch.cursor, - expected_keys, - changes: Vec::new(), - }) - .await - .map_err(map_state_admin_error)?; - break; - } - let request = StateMigrationRequest { - namespace: namespace.namespace.clone(), - from_schema_version: namespace.from_schema_version, - to_schema_version: namespace.to_schema_version, - records: batch - .records - .iter() - .map(|record| StateMigrationRecord { - key: record.key.clone(), - value: record.value.clone(), - version: record.version, - }) - .collect(), - }; - let reply = instance - .session - .call( - "plugin.state.migrate", - instance - .session - .context(Stage::Configuration, Duration::from_secs(30)), - serde_json::json!({}), - serde_json::to_vec(&request) - .map_err(|_| AdminError::internal("状态迁移批次无法编码"))?, - ) - .await - .map_err(|_| AdminError::invalid("插件状态迁移调用失败"))?; - if reply.result != serde_json::json!({}) { - return Err(AdminError::invalid("插件状态迁移结果信封无效")); - } - let result: StateMigrationResult = serde_json::from_slice(&reply.payload) - .map_err(|_| AdminError::invalid("插件状态迁移结果无效"))?; - if result.changes.len() != batch.records.len() - || result - .changes - .iter() - .zip(&expected_keys) - .any(|(change, expected)| change.key() != expected) - { - return Err(AdminError::invalid("插件状态迁移结果未逐项对应源批次")); - } - let mut changes = Vec::with_capacity(result.changes.len()); - for (change, source) in result.changes.into_iter().zip(&batch.records) { - let (key, action) = match change { - StateMigrationChange::Keep { key } => { - instance - .private_state - .validate_value(&namespace.namespace, &source.value)?; - (key, PluginStateMigrationAction::Keep) - } - StateMigrationChange::Replace { key, value } => { - instance - .private_state - .validate_value(&namespace.namespace, &value)?; - (key, PluginStateMigrationAction::Replace(value)) - } - StateMigrationChange::Delete { key } => { - (key, PluginStateMigrationAction::Delete) - } - }; - changes.push(PluginStateMigrationChange { key, action }); - } - self.state - .apply_migration_batch(ApplyPluginStateMigration { - transition_id: transition.id.clone(), - namespace: namespace.namespace.clone(), - cursor: batch.cursor, - expected_keys, - changes, - }) - .await - .map_err(map_state_admin_error)?; - } - } - Ok(()) - } -} - -#[async_trait] -impl gateway_admin::ports::plugin_management::PluginManagement for PluginRuntime { - async fn validate_target( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - ) -> Result<(), AdminError> { - let set = self.prepared_set(published).await?; - set.management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))?; - - Ok(()) - } - - async fn start_callback( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - command: gateway_admin::model::plugins::management::StartPluginManagementCallback, - context: &gateway_admin::model::auth::AdminRequestContext, - ) -> Result - { - let store = self - .oauth_pending - .as_ref() - .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .start_callback(store.as_ref(), command, context) - .await - } - - async fn callback( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - state: &str, - request: gateway_admin::model::plugins::management::PluginManagementRequest, - ) -> Result - { - let store = self - .oauth_pending - .as_ref() - .ok_or_else(|| AdminError::unavailable("插件登录状态端口不可用"))?; - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .callback(store.as_ref(), state, request, self.config.rpc_limits) - .await - } - - async fn views( - &self, - published: &ExtensionSetReference, - ) -> Result, AdminError> - { - Ok(self - .prepared_set(published) - .await? - .management - .iter() - .map(|entry| entry.view.clone()) - .collect()) - } - - async fn resource( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - path: &str, - public: bool, - ) -> Result - { - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .resource(path, public) - } - - async fn handle( - &self, - published: &ExtensionSetReference, - target: &gateway_admin::model::plugins::management::PluginManagementTarget, - request: gateway_admin::model::plugins::management::PluginManagementRequest, - ) -> Result - { - self.prepared_set(published) - .await? - .management - .iter() - .find(|entry| &entry.view.target == target) - .ok_or_else(|| AdminError::conflict("插件页面版本已变化，请刷新页面"))? - .handle(request, self.config.rpc_limits) - .await - } -} - -#[async_trait] -impl PluginPreparation for PluginRuntime { - async fn configuration_ready( - &self, - instance: PluginInstance, - metadata: &gateway_admin::model::plugins::PluginArtifactMetadata, - ) -> Result { - if instance.artifact_sha256 != metadata.sha256 { - return Err(AdminError::invalid("插件配置与制品不匹配")); - } - // 安装时已验证声明；只读列表不重新解包，也不占用安装与启用的校验配额 - let schema = metadata.configuration_schema.clone(); - let secret_fields = metadata.secret_fields.iter().cloned().collect(); - tokio::task::spawn_blocking(move || { - super::configuration::configuration_ready(&instance, &schema, &secret_fields) - }) - .await - .map_err(|_| AdminError::internal("插件校验任务失败"))? - } - - async fn validate( - &self, - instance: PluginInstance, - ) -> Result { - let artifact = self - .store - .load_artifact(&instance.artifact_sha256) - .await - .map_err(|_| AdminError::not_found("插件制品不存在"))?; - let limits = self.config.package_limits; - let slot = self - .validators - .clone() - .try_acquire_owned() - .map_err(|_| AdminError::unavailable("插件校验繁忙"))?; - tokio::task::spawn_blocking(move || { - let _slot = slot; - let package = - ValidatedPackage::read(artifact.archive, Some(&instance.artifact_sha256), limits) - .map_err(|_| AdminError::invalid("插件制品校验失败"))?; - let (_, state) = super::configuration::validate(&instance, package.manifest())?; - crate::adapter::observer::validate_bindings(package.manifest(), &instance.bindings)?; - crate::adapter::policy::validate_bindings(package.manifest(), &instance.bindings)?; - crate::adapter::catalog::validate_bindings(package.manifest(), &instance.bindings)?; - crate::adapter::upstream_adapter::validate_bindings( - package.manifest(), - &instance.bindings, - )?; - crate::adapter::frontend_authentication::validate_bindings( - package.manifest(), - &instance.bindings, - )?; - Ok(state) - }) - .await - .map_err(|_| AdminError::internal("插件校验任务失败"))? - } - async fn prepare( - &self, - snapshot: PluginInstanceSnapshot, - ) -> Result { - let required_revision = snapshot.config_revision; - self.prepare_snapshot(snapshot, Some(required_revision)) - .await - } -} - -#[async_trait] -impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for PluginRuntime { - async fn runtime_diagnostics( - &self, - snapshot: &PluginInstanceSnapshot, - published_revision: Option, - published: Option<&ExtensionSetReference>, - ) -> Option> { - let sets = { - let mut prepared = self.prepared.lock().await; - prepared.retain(|_, set| set.strong_count() > 0); - prepared - .values() - .filter_map(Weak::upgrade) - .collect::>() - }; - let open_restart_circuits = snapshot - .instances - .iter() - .filter(|instance| instance.enabled) - .filter_map(|instance| { - let identity = - RestartIdentity::new(instance.id.clone(), instance_fingerprint(instance).ok()?); - self.restart_circuit_is_open(&identity) - .then(|| instance.id.clone()) - }) - .collect::>(); - let published_set = published - .and_then(|reference| sets.iter().find(|set| set.id == *reference.id()).cloned()); - let published_ready = published.is_some_and(ExtensionSetReference::can_serve); - let preparation = self.preparation_diagnostic(snapshot.config_revision.get()); - let projection = InstanceRuntimeProjection { - target_revision: snapshot.config_revision, - published_revision, - published_ready, - published_set: published_set.as_deref(), - sets: &sets, - preparation: preparation.as_ref(), - open_restart_circuits: &open_restart_circuits, - }; - Some( - snapshot - .instances - .iter() - .map(|instance| { - let runtime = projection.project(instance); - (instance.id.clone(), runtime) - }) - .collect(), - ) - } -} - -#[async_trait] -impl gateway_admin::ports::plugins::PluginStateLifecycle for PluginRuntime { - async fn activate_state( - &self, - prepared: &ExtensionSetReference, - instance: &PluginInstance, - ) -> Result<(), AdminError> { - if !instance.enabled { - return Ok(()); - } - let set = self.prepared_set(prepared).await?; - let prepared = set - .sessions - .iter() - .find(|prepared| { - prepared.instance_id == instance.id - && prepared.artifact_sha256 == instance.artifact_sha256 - }) - .ok_or_else(|| AdminError::conflict("插件候选实例已失效"))?; - let owner = self - .state - .load_owner(prepared.private_state.owner_request(instance)) - .await - .map_err(map_state_admin_error)? - .ok_or_else(|| AdminError::conflict("插件状态 fence 尚未提交"))?; - prepared.private_state.activate(owner); - Ok(()) - } - - async fn quiesce_instance(&self, instance_id: &str, artifact_sha256: &str, revision: Revision) { - let mut sessions = Vec::new(); - for set in self - .prepared - .lock() - .await - .values() - .filter_map(Weak::upgrade) - { - for prepared in &set.sessions { - if prepared.instance_id == instance_id - && prepared.artifact_sha256 == artifact_sha256 - && prepared.revision == revision - && !sessions - .iter() - .any(|session| Arc::ptr_eq(session, &prepared.session)) - { - sessions.push(prepared.session.clone()); - } - } - } - shutdown_sessions(sessions, Duration::from_secs(2)).await; - } - - async fn migrate_state( - &self, - prepared: &ExtensionSetReference, - transition: PluginStateTransition, - ) -> Result<(), AdminError> { - self.run_state_migration(prepared, transition).await - } -} - -struct InstanceRuntimeProjection<'a> { - target_revision: Revision, - published_revision: Option, - published_ready: bool, - published_set: Option<&'a PreparedSet>, - sets: &'a [Arc], - preparation: Option<&'a PreparationDiagnostic>, - open_restart_circuits: &'a BTreeSet, -} - -impl InstanceRuntimeProjection<'_> { - fn project(&self, expected: &PluginInstance) -> PluginInstanceRuntime { - let active = self.published_set.and_then(|set| { - set.sessions - .iter() - .find(|instance| instance.instance_id == expected.id) - }); - let actual_matches = active.is_some_and(|instance| { - instance.revision == expected.revision - && instance.artifact_sha256 == expected.artifact_sha256 - }); - let target_is_published = self.published_revision == Some(self.target_revision.get()); - let candidate_set = match self.preparation { - Some(PreparationDiagnostic::Prepared { set_id }) => Some(set_id.as_str()), - _ => None, - }; - let published_set_id = self.published_set.map(|set| set.id.as_str()); - let mut draining_revisions = BTreeSet::new(); - for set in self.sets { - if Some(set.id.as_str()) == published_set_id || Some(set.id.as_str()) == candidate_set { - continue; - } - if let Some(instance) = set - .sessions - .iter() - .find(|instance| instance.instance_id == expected.id) - && instance.revision.get() < expected.revision.get() - { - draining_revisions.insert(instance.revision.get()); - } - } - if !expected.enabled - && let Some(instance) = active - && instance.revision != expected.revision - { - draining_revisions.insert(instance.revision.get()); - } - - let diagnostic = active.map(|instance| instance.session.diagnostic()); - let mut failure = None; - let status = if !expected.enabled { - if active.is_some() || !draining_revisions.is_empty() { - PluginInstanceRuntimeStatus::Draining - } else { - PluginInstanceRuntimeStatus::Disabled - } - } else if target_is_published - && let Some(reason) = self - .published_set - .and_then(|set| set.failures.get(&expected.id)) - { - failure = Some(reason.clone()); - PluginInstanceRuntimeStatus::PreparationFailed - } else if self.open_restart_circuits.contains(&expected.id) { - let message = match diagnostic { - Some(crate::rpc::RpcSessionDiagnostic::Failed { message, .. }) => { - format!("{message}，已暂停自动重启") - } - _ => "插件实例连续异常退出，已暂停自动重启".to_owned(), - }; - failure = Some(runtime_failure("restart_circuit_open", &message)); - PluginInstanceRuntimeStatus::Faulted - } else if actual_matches - && let Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) = diagnostic - { - failure = Some(runtime_failure(code, message)); - PluginInstanceRuntimeStatus::Faulted - } else if !target_is_published { - match self.preparation { - Some(PreparationDiagnostic::Preparing) => PluginInstanceRuntimeStatus::Preparing, - Some(PreparationDiagnostic::Failed(reason)) => { - failure = Some(reason.clone()); - PluginInstanceRuntimeStatus::PreparationFailed - } - Some(PreparationDiagnostic::Prepared { .. }) | None => { - PluginInstanceRuntimeStatus::AwaitingPublication - } - } - } else if !actual_matches { - failure = Some(runtime_failure( - "published_instance_mismatch", - "已发布插件实例与期望版本不一致", - )); - PluginInstanceRuntimeStatus::Blocked - } else { - match diagnostic { - Some(crate::rpc::RpcSessionDiagnostic::Ready) if self.published_ready => { - PluginInstanceRuntimeStatus::Running - } - Some(crate::rpc::RpcSessionDiagnostic::Ready) => { - failure = Some(runtime_failure( - "published_set_unready", - "同一发布集合中存在未就绪插件实例", - )); - PluginInstanceRuntimeStatus::Blocked - } - Some(crate::rpc::RpcSessionDiagnostic::Quiescing) => { - PluginInstanceRuntimeStatus::Draining - } - Some(crate::rpc::RpcSessionDiagnostic::Failed { code, message }) => { - failure = Some(runtime_failure(code, message)); - PluginInstanceRuntimeStatus::Faulted - } - None => { - failure = Some(runtime_failure( - "instance_not_published", - "已发布集合缺少该插件实例", - )); - PluginInstanceRuntimeStatus::Blocked - } - } - }; - - PluginInstanceRuntime { - status, - actual_revision: active.map(|instance| instance.revision.get()), - actual_artifact_sha256: active.map(|instance| instance.artifact_sha256.clone()), - failure, - draining_revisions: draining_revisions.into_iter().collect(), - } - } -} - -fn runtime_failure(code: &str, message: &str) -> PluginInstanceRuntimeFailure { - PluginInstanceRuntimeFailure { - code: code.to_owned(), - message: message.to_owned(), - } -} - -fn runtime_failure_from_admin(error: &AdminError) -> PluginInstanceRuntimeFailure { - let code = match error.kind() { - AdminErrorKind::Invalid => "invalid", - AdminErrorKind::Unauthorized => "unauthorized", - AdminErrorKind::Forbidden => "forbidden", - AdminErrorKind::NotFound => "not_found", - AdminErrorKind::Conflict => "conflict", - AdminErrorKind::RateLimited => "rate_limited", - AdminErrorKind::BadGateway => "bad_gateway", - AdminErrorKind::UpstreamResultUnknown => "upstream_result_unknown", - AdminErrorKind::Unavailable => "unavailable", - AdminErrorKind::Internal => "internal", - }; - runtime_failure(code, error.message()) -} - -async fn shutdown_sessions(sessions: Vec>, grace: Duration) { - for session in &sessions { - session.quiesce(); - } - let deadline = tokio::time::Instant::now() + grace; - futures::future::join_all( - sessions - .iter() - .map(|session| session.shutdown_until(deadline)), - ) - .await; -} - -fn map_state_admin_error( - error: gateway_admin::ports::plugins::PluginStateStoreError, -) -> AdminError { - match error.kind() { - PluginStateStoreErrorKind::Invalid => AdminError::invalid("插件状态请求不合法"), - PluginStateStoreErrorKind::NotFound => AdminError::not_found("插件状态迁移不存在"), - PluginStateStoreErrorKind::Conflict => AdminError::conflict("插件状态版本或迁移游标冲突"), - PluginStateStoreErrorKind::Quota => AdminError::conflict("插件状态超过声明配额"), - PluginStateStoreErrorKind::PermissionDenied => { - AdminError::conflict("插件状态 fence 已失效") - } - PluginStateStoreErrorKind::Unavailable => AdminError::unavailable("插件状态存储暂不可用"), - } -} - -impl ExtensionPreparationPort for PluginRuntime { - fn prepare( - &self, - revision: ConfigRevision, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { - let snapshot = self - .store - .load_instances() - .await - .map_err(|_| ExtensionPreparationError)?; - if snapshot.config_revision.get() != revision.get() { - return Err(ExtensionPreparationError); - } - self.prepare_snapshot(snapshot, None) - .await - .map_err(|_| ExtensionPreparationError) - }) - } -} - -fn fingerprint(snapshot: &PluginInstanceSnapshot) -> Result { - let mut instances: Vec<_> = snapshot - .instances - .iter() - .filter(|instance| instance.enabled) - .collect(); - instances.sort_by(|left, right| left.id.cmp(&right.id)); - let mut hash = Sha256::new(); - for instance in instances { - hash.update(instance_fingerprint(instance)?.as_bytes()); - } - Ok(hex::encode(hash.finalize())) -} - -fn instance_fingerprint( - instance: &gateway_admin::model::plugins::instances::PluginInstance, -) -> Result { - let secrets: BTreeMap<_, _> = instance - .secrets - .iter() - .map(|(name, value)| (name, value.expose_secret())) - .collect(); - let mut value = serde_json::json!({"id":instance.id,"artifact":instance.artifact_sha256,"revision":instance.revision.get(),"trusted":instance.trusted_process,"configuration":instance.configuration,"secrets":secrets,"bindings":instance.bindings}); - // JSONB 恢复可以改变对象键顺序，不能因此替换已经准备好的同一配置 - value.sort_all_objects(); - let bytes = serde_json::to_vec(&value).map_err(|_| AdminError::invalid("插件配置无法编码"))?; - Ok(hex::encode(Sha256::digest(bytes))) -} - -impl PluginRuntime { - pub fn bind_resource_ports( - &self, - access: &Arc, - ) -> Result<(), AdminError> { - self.resource_ports.bind(access) - } -} +pub use runtime::{PluginRuntime, PluginRuntimeConfig}; diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs new file mode 100644 index 000000000..be6244605 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs @@ -0,0 +1,240 @@ +//! 插件运行时配置、端口装配与整体关闭；候选和调用均复用同一状态 + +use super::super::restart_circuit::{PluginRestartCircuitConfig, RestartCircuits, RestartIdentity}; +use super::{diagnostics::PreparationDiagnostic, set::PreparedSet}; +use crate::{ + PackageLimits, RpcLimits, RpcSession, + callback::{ + PluginAccountPortSlot, PluginAffinityPortSlot, PluginClientKeyPortSlot, PluginModelPortSlot, + }, +}; +use gateway_admin::{ + model::AdminError, + ports::{ + plugin_accounts::PluginAccountAccess, + plugin_client_keys::PluginClientKeyAccess, + plugins::{PluginStateStore, PluginStore}, + }, +}; +use gateway_core::{ + engine::{ + authentication::FrontendAuthenticationExtensionIndex, + extensions::ExecutionExtensionIndex, + nested::{AffinityLookupPort, NestedModelExecutionPort}, + observation::RequestObserverExtensionIndex, + policy::RequestPolicyExtensionIndex, + }, + provider_ports::OAuthPendingFlowPort, +}; +use std::{ + collections::BTreeMap, + path::PathBuf, + sync::{ + Arc, Mutex as SyncMutex, Weak, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; +use tokio::sync::{Mutex, Semaphore}; + +const RUNTIME_SHUTDOWN_GRACE: Duration = Duration::from_secs(2); +pub struct PluginRuntimeConfig { + pub cache_directory: PathBuf, + pub host_version: semver::Version, + pub package_limits: PackageLimits, + pub rpc_limits: RpcLimits, + pub restart_circuit: PluginRestartCircuitConfig, +} + +/// 只索引有主人的候选；当前代次仅由 Core 的发布视图持有 +pub struct PluginRuntime { + pub(super) service_ports: Arc, + pub(super) store: Arc, + pub(super) state: Arc, + pub(super) config: PluginRuntimeConfig, + pub(super) prepare_lock: Mutex<()>, + pub(super) prepared: Mutex>>, + pub(super) preparation_diagnostics: SyncMutex>, + pub(super) restart_circuits: Arc, + pub(super) shutting_down: Arc, + shutdown_lock: Mutex<()>, + pub(super) processes: Arc, + pub(super) validators: Arc, + pub(super) log_slots: Arc, + pub(super) account_ports: Arc, + pub(super) client_key_ports: Arc, + pub(super) resource_ports: Arc, + pub(super) model_ports: Arc, + pub(super) affinity_ports: Arc, + pub(super) observers: RequestObserverExtensionIndex, + pub(super) policies: RequestPolicyExtensionIndex, + pub(super) execution: ExecutionExtensionIndex, + pub(super) authentications: FrontendAuthenticationExtensionIndex, + pub(super) oauth_pending: Option>, + pub(super) http: Arc, +} +impl PluginRuntime { + #[must_use] + pub fn new( + store: Arc, + state: Arc, + config: PluginRuntimeConfig, + http: Arc, + processes: Arc, + ) -> Self { + let restart_circuit = config.restart_circuit; + Self { + store, + state, + log_slots: Arc::new(Semaphore::new( + config + .rpc_limits + .maximum_callbacks + .min(Semaphore::MAX_PERMITS), + )), + config, + prepare_lock: Mutex::new(()), + prepared: Mutex::new(BTreeMap::new()), + preparation_diagnostics: SyncMutex::new(BTreeMap::new()), + restart_circuits: RestartCircuits::new(restart_circuit), + shutting_down: Arc::new(AtomicBool::new(false)), + shutdown_lock: Mutex::new(()), + processes, + validators: Arc::new(Semaphore::new(2)), + service_ports: Arc::new(crate::callback::services::ServicePorts::default()), + account_ports: Arc::new(PluginAccountPortSlot::new()), + client_key_ports: Arc::new(PluginClientKeyPortSlot::new()), + resource_ports: Arc::new(crate::callback::PluginResourcePorts::new()), + model_ports: Arc::new(PluginModelPortSlot::new()), + affinity_ports: Arc::new(PluginAffinityPortSlot::new()), + observers: RequestObserverExtensionIndex::default(), + policies: RequestPolicyExtensionIndex::default(), + execution: ExecutionExtensionIndex::default(), + authentications: FrontendAuthenticationExtensionIndex::default(), + oauth_pending: None, + http, + } + } + + #[must_use] + pub fn with_oauth_pending(mut self, port: Arc) -> Self { + self.oauth_pending = Some(port); + self + } + + /// API 完成组装后绑定唯一内部 HTTP 分派端口；不持有服务端强引用 + pub fn bind_http( + &self, + dispatcher: &Arc, + ) -> Result<(), AdminError> { + self.service_ports.bind_http(dispatcher) + } + + pub fn bind_services( + &self, + registry: &Arc, + ) -> Result<(), AdminError> { + self.service_ports.bind(registry) + } + + pub fn bind_account_ports( + &self, + access: &Arc, + ) -> Result<(), AdminError> { + self.account_ports.bind(access) + } + + /// Key 目录与预算端口由 Admin 组合并保活，完整字段由对应公开合同返回 + pub fn bind_client_key_ports( + &self, + access: &Arc, + ) -> Result<(), AdminError> { + self.client_key_ports.bind(access) + } + + /// Core 激活后一次性绑定嵌套执行与亲和查询端口；Runtime 仅保存 Weak，避免组合根强环 + pub fn bind_model_ports( + &self, + models: &Arc, + affinity: &Arc, + ) -> Result<(), AdminError> { + self.model_ports.bind(models)?; + self.affinity_ports.bind(affinity) + } + + /// 调用方必须先停止接收 HTTP 与 Worker；这里拒绝新候选并等待全部插件 I/O 退出 + pub async fn shutdown(&self) { + self.shutting_down.store(true, Ordering::Release); + let _shutdown = self.shutdown_lock.lock().await; + let _prepare = self.prepare_lock.lock().await; + let sets = { + let mut prepared = self.prepared.lock().await; + prepared.retain(|_, set| set.strong_count() > 0); + let sets = prepared + .values() + .filter_map(Weak::upgrade) + .collect::>(); + prepared.clear(); + sets + }; + let mut sessions = Vec::new(); + for set in &sets { + for instance in &set.sessions { + if !sessions + .iter() + .any(|session| Arc::ptr_eq(session, &instance.session)) + { + sessions.push(Arc::clone(&instance.session)); + } + } + } + shutdown_sessions(sessions, RUNTIME_SHUTDOWN_GRACE).await; + } + + #[must_use] + pub fn observer_registry(&self) -> RequestObserverExtensionIndex { + self.observers.clone() + } + + #[must_use] + pub fn policy_registry(&self) -> RequestPolicyExtensionIndex { + self.policies.clone() + } + + /// 执行计划按同一发布身份解析，索引不延长代次寿命 + #[must_use] + pub fn execution_registry(&self) -> ExecutionExtensionIndex { + self.execution.clone() + } + + #[must_use] + pub fn frontend_authentication_registry(&self) -> FrontendAuthenticationExtensionIndex { + self.authentications.clone() + } +} +pub(super) async fn shutdown_sessions(sessions: Vec>, grace: Duration) { + for session in &sessions { + session.quiesce(); + } + let deadline = tokio::time::Instant::now() + grace; + futures::future::join_all( + sessions + .iter() + .map(|session| session.shutdown_until(deadline)), + ) + .await; +} +impl PluginRuntime { + pub fn bind_resource_ports( + &self, + access: &Arc, + ) -> Result<(), AdminError> { + self.resource_ports.bind(access) + } +} + +impl PluginRuntime { + pub(super) fn restart_circuit_is_open(&self, identity: &RestartIdentity) -> bool { + self.restart_circuits.is_open(identity) + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs new file mode 100644 index 000000000..aae6eb632 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs @@ -0,0 +1,140 @@ +//! 发布集合及其保活、候选贡献与配置身份指纹 + +use crate::{RpcSession, callback::private_state::PluginPrivateState}; +use gateway_admin::model::{ + AdminError, Revision, + plugins::instances::{PluginInstanceRuntimeFailure, PluginInstanceSnapshot}, +}; +use gateway_core::{ + engine::{ + extensions::ExecutionExtensionPlans, observation::RequestObserverPlan, + policy::RequestPolicyPlan, + }, + routing::extensions::{ExtensionSetId, ExtensionSetLease}, +}; +use secrecy::ExposeSecret as _; +use sha2::{Digest as _, Sha256}; +use std::{ + collections::BTreeMap, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +pub(super) struct PreparedSet { + pub(super) id: ExtensionSetId, + pub(super) _observers: Option>, + pub(super) _policies: Option>, + pub(super) _execution: Option>, + pub(super) _authentication: + Option>, + pub(super) sessions: Vec, + pub(super) failures: BTreeMap, + pub(super) shutting_down: Arc, + pub(super) commands: Vec>, + pub(super) management: Vec, + pub(super) model_aliases: Vec, +} + +#[derive(Default)] +pub(super) struct PreparedContributions { + pub(super) sessions: Vec, + pub(super) observer_entries: Vec, + pub(super) commands: Vec>, + pub(super) management: Vec, + pub(super) policy_entries: Vec, + pub(super) upstream_entries: Vec, + pub(super) authentication_entries: + Vec, + pub(super) model_aliases: Vec, +} + +impl PreparedContributions { + pub(super) fn append(&mut self, other: Self) { + self.sessions.extend(other.sessions); + self.observer_entries.extend(other.observer_entries); + self.commands.extend(other.commands); + self.management.extend(other.management); + self.model_aliases.extend(other.model_aliases); + self.policy_entries.extend(other.policy_entries); + self.upstream_entries.extend(other.upstream_entries); + self.authentication_entries + .extend(other.authentication_entries); + } +} + +pub(super) struct PreparedInstance { + pub(super) instance_id: String, + pub(super) artifact_sha256: String, + pub(super) revision: Revision, + pub(super) session: Arc, + pub(super) private_state: Arc, + pub(super) maintenance: bool, +} + +impl PreparedSet { + pub(super) fn sessions_ready(&self) -> bool { + self.sessions + .iter() + .all(|instance| instance.session.is_ready()) + } +} + +impl ExtensionSetLease for PreparedSet { + fn is_ready(&self) -> bool { + self.sessions_ready() && self.can_serve() + } + + fn model_aliases(&self) -> &[gateway_core::routing::ContributedModelAlias] { + &self.model_aliases + } + + fn can_serve(&self) -> bool { + // 发布的是可用能力和故障绑定组成的完整计划；单个进程退出不能使原生转发失效 + !self.shutting_down.load(Ordering::Acquire) + } +} + +impl Drop for PreparedSet { + fn drop(&mut self) { + for instance in self.sessions.drain(..) { + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let session = instance.session; + session.quiesce(); + session.shutdown(Duration::from_secs(1)).await; + }); + } + } + } +} +pub(super) fn fingerprint(snapshot: &PluginInstanceSnapshot) -> Result { + let mut instances: Vec<_> = snapshot + .instances + .iter() + .filter(|instance| instance.enabled) + .collect(); + instances.sort_by(|left, right| left.id.cmp(&right.id)); + let mut hash = Sha256::new(); + for instance in instances { + hash.update(instance_fingerprint(instance)?.as_bytes()); + } + Ok(hex::encode(hash.finalize())) +} + +pub(super) fn instance_fingerprint( + instance: &gateway_admin::model::plugins::instances::PluginInstance, +) -> Result { + let secrets: BTreeMap<_, _> = instance + .secrets + .iter() + .map(|(name, value)| (name, value.expose_secret())) + .collect(); + let mut value = serde_json::json!({"id":instance.id,"artifact":instance.artifact_sha256,"revision":instance.revision.get(),"trusted":instance.trusted_process,"configuration":instance.configuration,"secrets":secrets,"bindings":instance.bindings}); + // JSONB 恢复可以改变对象键顺序，不能因此替换已经准备好的同一配置 + value.sort_all_objects(); + let bytes = serde_json::to_vec(&value).map_err(|_| AdminError::invalid("插件配置无法编码"))?; + Ok(hex::encode(Sha256::digest(bytes))) +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs new file mode 100644 index 000000000..e55ee2e13 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs @@ -0,0 +1,343 @@ +//! 校验实例配置并准备可发布的完整能力集合 + +use super::{ + PluginRuntime, + diagnostics::{PreparationDiagnostic, runtime_failure_from_admin}, + set::{PreparedContributions, PreparedSet, fingerprint}, +}; +use crate::ValidatedPackage; +use async_trait::async_trait; +use futures::future::BoxFuture; +use gateway_admin::{ + model::{ + AdminError, Revision, + plugins::{ + instances::{PluginInstance, PluginInstanceSnapshot}, + state::PluginStateConfiguration, + }, + }, + ports::plugins::PluginPreparation, +}; +use gateway_core::{ + routing::ConfigRevision, + routing::extensions::{ + ExtensionPreparationError, ExtensionPreparationPort, ExtensionSetId, ExtensionSetReference, + }, +}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{Arc, Weak, atomic::Ordering}, + time::Duration, +}; + +impl PluginRuntime { + pub(super) async fn prepare_snapshot( + &self, + snapshot: PluginInstanceSnapshot, + required_revision: Option, + ) -> Result { + let target_revision = snapshot.config_revision.get(); + self.record_preparation(target_revision, PreparationDiagnostic::Preparing); + let result = self.prepare_inner(snapshot, required_revision).await; + match &result { + Ok(reference) => self.record_preparation( + target_revision, + PreparationDiagnostic::Prepared { + set_id: reference.id().as_str().to_owned(), + }, + ), + Err(error) => self.record_preparation( + target_revision, + PreparationDiagnostic::Failed(runtime_failure_from_admin(error)), + ), + } + result + } + + async fn prepare_inner( + &self, + snapshot: PluginInstanceSnapshot, + required_revision: Option, + ) -> Result { + if self.shutting_down.load(Ordering::Acquire) { + return Err(AdminError::unavailable("插件运行时正在关闭")); + } + if snapshot.instances.len() > 64 { + return Err(AdminError::invalid("最多配置 64 个插件实例")); + } + let fingerprint = fingerprint(&snapshot)?; + // 准备串行化与弱索引分锁；诊断和已发布请求不会等待外部 I/O + let _prepare = self.prepare_lock.lock().await; + if self.shutting_down.load(Ordering::Acquire) { + return Err(AdminError::unavailable("插件运行时正在关闭")); + } + let existing = { + let mut prepared = self.prepared.lock().await; + prepared.retain(|_, set| set.strong_count() > 0); + prepared + .get(&fingerprint) + .and_then(Weak::upgrade) + .filter(|set| set.sessions_ready()) + }; + if let Some(set) = existing { + if snapshot.instances.iter().any(|instance| { + required_revision == Some(instance.revision) + && set.failures.contains_key(&instance.id) + }) { + return Err(AdminError::unavailable("插件启动失败，请重新启用后重试")); + } + return Ok(ExtensionSetReference::new(set.id.clone(), set)); + } + let mut contributions = PreparedContributions::default(); + let mut failures = BTreeMap::new(); + let mut identities = BTreeSet::new(); + for instance in snapshot + .instances + .iter() + .filter(|instance| instance.enabled) + { + if !identities.insert(instance.id.clone()) { + return Err(AdminError::invalid("插件实例 ID 重复")); + } + let result = tokio::time::timeout( + // 沿用准备阶段原有的 30 秒预算；打包校验与能力恢复也在预算内 + Duration::from_secs(30), + self.prepare_instance(instance.clone(), snapshot.config_revision), + ) + .await + .unwrap_or_else(|_| Err(AdminError::unavailable("插件启动超时"))); + match result { + Ok(candidate) => contributions.append(candidate), + Err(error) if required_revision == Some(instance.revision) => return Err(error), + Err(error) => { + // 目录没有拒绝请求用的 binding，不能在恢复失败时悄悄撤销仍启用的别名 + // 不复制其他候选的注册结果；让发布事务保留调用方持有的旧有效快照 + let has_retained_catalog = self + .prepared + .lock() + .await + .values() + .filter_map(Weak::upgrade) + .any(|set| { + set.model_aliases + .iter() + .any(|alias| alias.owner == instance.id) + }); + if has_retained_catalog { + return Err(AdminError::unavailable( + "已启用的模型目录插件未就绪，请修复或停用后重试", + )); + } + // 恢复失败只隔离该实例；绑定的拒绝策略仍保留，不能绕过认证或必需处理 + contributions + .policy_entries + .extend(crate::adapter::policy::unavailable_entries(instance)?); + contributions.upstream_entries.extend( + crate::adapter::upstream_adapter::unavailable_entries(instance)?, + ); + if let Some(entry) = + crate::adapter::frontend_authentication::unavailable_entry(instance) + { + contributions.authentication_entries.push(entry); + } + failures.insert(instance.id.clone(), runtime_failure_from_admin(&error)); + } + } + } + let PreparedContributions { + sessions, + observer_entries, + commands, + management, + policy_entries, + upstream_entries, + authentication_entries, + model_aliases, + } = contributions; + let mut alias_owners = BTreeMap::new(); + for alias in &model_aliases { + if let Some(previous) = alias_owners.insert(&alias.id, &alias.owner) { + return Err(AdminError::invalid(format!( + "模型别名 {} 在插件实例 {} 与 {} 之间冲突", + alias.id, previous, alias.owner + ))); + } + } + let id = ExtensionSetId::new(uuid::Uuid::new_v4().to_string()) + .map_err(|_| AdminError::internal("插件集合 ID 无效"))?; + let observers = crate::adapter::observer::PluginObserverPlan::compile( + observer_entries, + self.config.rpc_limits.maximum_calls, + self.config.rpc_limits.maximum_call_timeout, + self.config.rpc_limits.maximum_buffered_body_bytes, + ) + .map(|plan| { + self.observers + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件观察计划注册冲突")) + }) + .transpose()?; + let policy_plan = crate::adapter::policy::PluginRequestPolicyPlan::compile( + policy_entries, + self.config.rpc_limits.maximum_call_timeout, + )?; + let policies = policy_plan + .as_ref() + .filter(|plan| plan.has_request_policy()) + .cloned() + .map(|plan| { + self.policies + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件请求策略计划注册冲突")) + }) + .transpose()?; + let middleware = policy_plan + .filter(|plan| plan.has_middleware()) + .map(|plan| plan as Arc); + let upstream_adapters = + crate::adapter::upstream_adapter::PluginUpstreamAdapterPlan::compile(upstream_entries)?; + let execution = if middleware.is_some() || upstream_adapters.is_some() { + Some( + self.execution + .register( + id.clone(), + Arc::new(gateway_core::engine::extensions::ExecutionExtensionPlans { + middleware, + upstream_adapters, + }), + ) + .map_err(|_| AdminError::invalid("插件执行计划注册冲突"))?, + ) + } else { + None + }; + let authentication = + crate::adapter::frontend_authentication::PluginFrontendAuthenticationPlan::compile( + authentication_entries, + self.config.rpc_limits.maximum_call_timeout, + )? + .map(|plan| { + self.authentications + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件入口认证计划注册冲突")) + }) + .transpose()?; + let set = Arc::new(PreparedSet { + id, + _observers: observers, + _policies: policies, + _execution: execution, + _authentication: authentication, + sessions, + commands, + management, + failures, + model_aliases, + shutting_down: self.shutting_down.clone(), + }); + self.prepared + .lock() + .await + .insert(fingerprint, Arc::downgrade(&set)); + Ok(ExtensionSetReference::new(set.id.clone(), set)) + } + pub(super) async fn prepared_set( + &self, + reference: &ExtensionSetReference, + ) -> Result, AdminError> { + self.prepared + .lock() + .await + .values() + .filter_map(Weak::upgrade) + .find(|set| set.id == *reference.id()) + .ok_or_else(|| AdminError::conflict("插件候选已过期，请重试")) + } +} +#[async_trait] +impl PluginPreparation for PluginRuntime { + async fn configuration_ready( + &self, + instance: PluginInstance, + metadata: &gateway_admin::model::plugins::PluginArtifactMetadata, + ) -> Result { + if instance.artifact_sha256 != metadata.sha256 { + return Err(AdminError::invalid("插件配置与制品不匹配")); + } + // 安装时已验证声明；只读列表不重新解包，也不占用安装与启用的校验配额 + let schema = metadata.configuration_schema.clone(); + let secret_fields = metadata.secret_fields.iter().cloned().collect(); + tokio::task::spawn_blocking(move || { + super::super::configuration::configuration_ready(&instance, &schema, &secret_fields) + }) + .await + .map_err(|_| AdminError::internal("插件校验任务失败"))? + } + + async fn validate( + &self, + instance: PluginInstance, + ) -> Result { + let artifact = self + .store + .load_artifact(&instance.artifact_sha256) + .await + .map_err(|_| AdminError::not_found("插件制品不存在"))?; + let limits = self.config.package_limits; + let slot = self + .validators + .clone() + .try_acquire_owned() + .map_err(|_| AdminError::unavailable("插件校验繁忙"))?; + tokio::task::spawn_blocking(move || { + let _slot = slot; + let package = + ValidatedPackage::read(artifact.archive, Some(&instance.artifact_sha256), limits) + .map_err(|_| AdminError::invalid("插件制品校验失败"))?; + let (_, state) = super::super::configuration::validate(&instance, package.manifest())?; + crate::adapter::observer::validate_bindings(package.manifest(), &instance.bindings)?; + crate::adapter::policy::validate_bindings(package.manifest(), &instance.bindings)?; + crate::adapter::catalog::validate_bindings(package.manifest(), &instance.bindings)?; + crate::adapter::upstream_adapter::validate_bindings( + package.manifest(), + &instance.bindings, + )?; + crate::adapter::frontend_authentication::validate_bindings( + package.manifest(), + &instance.bindings, + )?; + Ok(state) + }) + .await + .map_err(|_| AdminError::internal("插件校验任务失败"))? + } + async fn prepare( + &self, + snapshot: PluginInstanceSnapshot, + ) -> Result { + let required_revision = snapshot.config_revision; + self.prepare_snapshot(snapshot, Some(required_revision)) + .await + } +} + +impl ExtensionPreparationPort for PluginRuntime { + fn prepare( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let snapshot = self + .store + .load_instances() + .await + .map_err(|_| ExtensionPreparationError)?; + if snapshot.config_revision.get() != revision.get() { + return Err(ExtensionPreparationError); + } + self.prepare_snapshot(snapshot, None) + .await + .map_err(|_| ExtensionPreparationError) + }) + } +} diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs new file mode 100644 index 000000000..2b83fa96a --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs @@ -0,0 +1,226 @@ +//! 插件状态的激活、迁移与实例排空，事务事实由状态端口提交 + +use super::{PluginRuntime, runtime::shutdown_sessions}; +use async_trait::async_trait; +use gateway_admin::{ + model::{ + AdminError, Revision, + plugins::{ + instances::PluginInstance, + state::{ + ApplyPluginStateMigration, PluginStateMigrationAction, PluginStateMigrationChange, + PluginStateTransition, + }, + }, + }, + ports::plugins::PluginStateStoreErrorKind, +}; +use gateway_core::routing::extensions::ExtensionSetReference; +use gateway_plugin_sdk::{ + Stage, + call::host::{ + StateMigrationChange, StateMigrationRecord, StateMigrationRequest, StateMigrationResult, + }, +}; +use std::{ + sync::{Arc, Weak}, + time::Duration, +}; + +impl PluginRuntime { + async fn run_state_migration( + &self, + prepared: &ExtensionSetReference, + transition: PluginStateTransition, + ) -> Result<(), AdminError> { + let set = self.prepared_set(prepared).await?; + let instance = set + .sessions + .iter() + .find(|instance| { + instance.instance_id == transition.instance_id + && instance.artifact_sha256 == transition.artifact_sha256 + }) + .ok_or_else(|| AdminError::conflict("状态迁移候选插件已失效"))?; + for namespace in &transition.namespaces { + let mut batches = 0usize; + loop { + batches += 1; + if batches > 200 { + return Err(AdminError::invalid("插件状态迁移超过有界批次数")); + } + let batch = self + .state + .migration_batch(&transition.id, &namespace.namespace, 64) + .await + .map_err(map_state_admin_error)?; + let expected_keys = batch + .records + .iter() + .map(|record| record.key.clone()) + .collect::>(); + if batch.records.is_empty() { + self.state + .apply_migration_batch(ApplyPluginStateMigration { + transition_id: transition.id.clone(), + namespace: namespace.namespace.clone(), + cursor: batch.cursor, + expected_keys, + changes: Vec::new(), + }) + .await + .map_err(map_state_admin_error)?; + break; + } + let request = StateMigrationRequest { + namespace: namespace.namespace.clone(), + from_schema_version: namespace.from_schema_version, + to_schema_version: namespace.to_schema_version, + records: batch + .records + .iter() + .map(|record| StateMigrationRecord { + key: record.key.clone(), + value: record.value.clone(), + version: record.version, + }) + .collect(), + }; + let reply = instance + .session + .call( + "plugin.state.migrate", + instance + .session + .context(Stage::Configuration, Duration::from_secs(30)), + serde_json::json!({}), + serde_json::to_vec(&request) + .map_err(|_| AdminError::internal("状态迁移批次无法编码"))?, + ) + .await + .map_err(|_| AdminError::invalid("插件状态迁移调用失败"))?; + if reply.result != serde_json::json!({}) { + return Err(AdminError::invalid("插件状态迁移结果信封无效")); + } + let result: StateMigrationResult = serde_json::from_slice(&reply.payload) + .map_err(|_| AdminError::invalid("插件状态迁移结果无效"))?; + if result.changes.len() != batch.records.len() + || result + .changes + .iter() + .zip(&expected_keys) + .any(|(change, expected)| change.key() != expected) + { + return Err(AdminError::invalid("插件状态迁移结果未逐项对应源批次")); + } + let mut changes = Vec::with_capacity(result.changes.len()); + for (change, source) in result.changes.into_iter().zip(&batch.records) { + let (key, action) = match change { + StateMigrationChange::Keep { key } => { + instance + .private_state + .validate_value(&namespace.namespace, &source.value)?; + (key, PluginStateMigrationAction::Keep) + } + StateMigrationChange::Replace { key, value } => { + instance + .private_state + .validate_value(&namespace.namespace, &value)?; + (key, PluginStateMigrationAction::Replace(value)) + } + StateMigrationChange::Delete { key } => { + (key, PluginStateMigrationAction::Delete) + } + }; + changes.push(PluginStateMigrationChange { key, action }); + } + self.state + .apply_migration_batch(ApplyPluginStateMigration { + transition_id: transition.id.clone(), + namespace: namespace.namespace.clone(), + cursor: batch.cursor, + expected_keys, + changes, + }) + .await + .map_err(map_state_admin_error)?; + } + } + Ok(()) + } +} +#[async_trait] +impl gateway_admin::ports::plugins::PluginStateLifecycle for PluginRuntime { + async fn activate_state( + &self, + prepared: &ExtensionSetReference, + instance: &PluginInstance, + ) -> Result<(), AdminError> { + if !instance.enabled { + return Ok(()); + } + let set = self.prepared_set(prepared).await?; + let prepared = set + .sessions + .iter() + .find(|prepared| { + prepared.instance_id == instance.id + && prepared.artifact_sha256 == instance.artifact_sha256 + }) + .ok_or_else(|| AdminError::conflict("插件候选实例已失效"))?; + let owner = self + .state + .load_owner(prepared.private_state.owner_request(instance)) + .await + .map_err(map_state_admin_error)? + .ok_or_else(|| AdminError::conflict("插件状态 fence 尚未提交"))?; + prepared.private_state.activate(owner); + Ok(()) + } + + async fn quiesce_instance(&self, instance_id: &str, artifact_sha256: &str, revision: Revision) { + let mut sessions = Vec::new(); + for set in self + .prepared + .lock() + .await + .values() + .filter_map(Weak::upgrade) + { + for prepared in &set.sessions { + if prepared.instance_id == instance_id + && prepared.artifact_sha256 == artifact_sha256 + && prepared.revision == revision + && !sessions + .iter() + .any(|session| Arc::ptr_eq(session, &prepared.session)) + { + sessions.push(prepared.session.clone()); + } + } + } + shutdown_sessions(sessions, Duration::from_secs(2)).await; + } + + async fn migrate_state( + &self, + prepared: &ExtensionSetReference, + transition: PluginStateTransition, + ) -> Result<(), AdminError> { + self.run_state_migration(prepared, transition).await + } +} +pub(super) fn map_state_admin_error( + error: gateway_admin::ports::plugins::PluginStateStoreError, +) -> AdminError { + match error.kind() { + PluginStateStoreErrorKind::Invalid => AdminError::invalid("插件状态请求不合法"), + PluginStateStoreErrorKind::NotFound => AdminError::not_found("插件状态迁移不存在"), + PluginStateStoreErrorKind::Conflict => AdminError::conflict("插件状态版本或迁移游标冲突"), + PluginStateStoreErrorKind::Quota => AdminError::conflict("插件状态超过声明配额"), + PluginStateStoreErrorKind::PermissionDenied => { + AdminError::conflict("插件状态 fence 已失效") + } + PluginStateStoreErrorKind::Unavailable => AdminError::unavailable("插件状态存储暂不可用"), + } +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/data.rs b/backend/crates/gateway-plugin/runtime/tests/generation/data.rs index a3dfa1baa..95fc020e3 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/data.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/data.rs @@ -15,8 +15,8 @@ use gateway_admin::{ }; use gateway_core::{ account::{AccountWeight, CredentialState, ProviderAccountId, QuotaState}, + routing::extensions::ExtensionPreparationPort, routing::{ConfigRevision, ProviderKind}, - runtime::extensions::ExtensionPreparationPort, }; use gateway_plugin_sdk::{Capability, Contributions, Stage}; use serde_json::{Value, json}; @@ -180,7 +180,12 @@ async fn management_facts_validate_queries_without_permission_declarations() { ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); assert!(runtime.policy_registry().resolve(&generation).is_none()); let view = runtime.views(&generation).await.unwrap().remove(0); let response = runtime @@ -266,7 +271,12 @@ async fn quota_refresh_is_available_without_permission_declarations() { ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); assert!(runtime.policy_registry().resolve(&generation).is_none()); let view = runtime.views(&generation).await.unwrap().remove(0); let response = runtime diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs b/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs index d1bcc37b2..1537fa9f4 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/observer.rs @@ -251,7 +251,7 @@ async fn frozen_plan_applies_scope_order_and_instance_deduplication() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -322,7 +322,7 @@ async fn websocket_observer_preserves_order_generation_and_complete_payload() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -416,7 +416,7 @@ async fn websocket_observer_bounds_event_count_and_payload_bytes_without_blockin limits, ) .await; - let count_generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let count_generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &count_runtime, ConfigRevision::new(1).unwrap(), ) @@ -487,7 +487,7 @@ async fn websocket_observer_bounds_event_count_and_payload_bytes_without_blockin RpcLimits::default(), ) .await; - let byte_generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let byte_generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &byte_runtime, ConfigRevision::new(1).unwrap(), ) @@ -547,7 +547,7 @@ async fn malformed_observer_response_stops_only_the_observer() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -598,7 +598,7 @@ async fn observer_failure_continues_the_plan_and_timeout_is_bounded() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -690,7 +690,7 @@ async fn client_scope_filters_observations_and_exposes_only_the_key_identifier() RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -767,7 +767,7 @@ async fn completed_observation_contains_terminal_cost_timings_and_failure() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -874,7 +874,7 @@ async fn observer_backpressure_drops_excess_and_inflight_dispatch_keeps_generati limits, ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) @@ -924,7 +924,7 @@ async fn one_observer_matches_completed_and_websocket_scopes_independently() { RpcLimits::default(), ) .await; - let generation = gateway_core::runtime::extensions::ExtensionPreparationPort::prepare( + let generation = gateway_core::routing::extensions::ExtensionPreparationPort::prepare( &runtime, ConfigRevision::new(1).unwrap(), ) diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs index a9334a57a..157ee2e2b 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/compatibility.rs @@ -1,7 +1,7 @@ //! 验证旧中间件的 Fast 设置适配与不兼容插件的实例故障隔离 use super::*; -use gateway_core::{policy::ClientPolicy, settings::RequestSettings}; +use gateway_core::{policy::ClientPolicy, routing::request_settings::RequestSettings}; use serde_json::{Value, json}; #[tokio::test] @@ -75,7 +75,10 @@ async fn legacy_middleware_preserves_three_state_fast_and_projects_sources() { bindings: vec![binding(MIDDLEWARE_CONTRIBUTION,"request",0,PluginFailurePolicy::Reject)], }], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let next = (Downstream { calls: Arc::default(), closes: Arc::default(), @@ -158,7 +161,10 @@ async fn unsupported_contract_starts_and_runtime_faults_only_stop_the_failed_ins } let generation = prepare(&runtime).await; assert!(generation.is_ready(), "版本警告允许随宿主启动"); - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); for _ in 0..2 { let response = plan diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs index 6d60184cf..31d74a8da 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/facts/mod.rs @@ -41,7 +41,10 @@ async fn facts_are_readable_without_replacing_or_duplicating_the_original_execut bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, "attempt", 0, PluginFailurePolicy::Reject)], }], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let metadata = ProviderCallMetadata::new( ProviderKind::new("openai").unwrap(), UpstreamModelId::new("native-model").unwrap(), diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs index 93055beca..561776160 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/dispatch.rs @@ -116,7 +116,10 @@ async fn setup( ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); (cache, runtime, generation, plan) } @@ -378,7 +381,7 @@ async fn cancelling_a_parent_drops_its_active_http_child() { #[tokio::test] async fn active_http_children_inherit_settings_and_preserve_override_sources() { - use gateway_core::{routing::RuntimeSnapshot, settings::RequestSettings}; + use gateway_core::{routing::RuntimeSnapshot, routing::request_settings::RequestSettings}; let (_cache, runtime, generation, plan) = setup("settings").await; let settings = Arc::new(Mutex::new(Vec::new())); let dispatcher: Arc = Arc::new(Dispatcher { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs index c093e6af1..161b69492 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs @@ -4,7 +4,7 @@ mod dispatch; mod websocket; use super::*; -use gateway_core::middleware::{compose, http as core}; +use gateway_core::{engine::middleware::http as core, middleware::compose}; use http_body::Frame; use http_body_util::{BodyExt as _, StreamBody}; @@ -54,7 +54,10 @@ async fn http_plugin_with_limits( ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); assert!(plan.has_http()); (cache, runtime, generation, plan) } @@ -521,7 +524,7 @@ async fn http_onion_performance_baseline() { .collect(); let (_cache, runtime) = setup_package(instances, package.clone()).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation); + let plan = runtime.execution_registry().middleware(&generation); let snapshot = gateway_core::routing::RuntimeSnapshot::new( ConfigRevision::new(1).unwrap(), gateway_core::settings::SettingsValues::new( @@ -537,7 +540,8 @@ async fn http_onion_performance_baseline() { vec![], ) .unwrap(); - let settings = gateway_core::settings::RequestSettings::new(Arc::new(snapshot)); + let settings = + gateway_core::routing::request_settings::RequestSettings::new(Arc::new(snapshot)); for size in [64 * 1024, 2 * 1024 * 1024] { let payload = Bytes::from(vec![b'a'; size]); for _ in 0..5 { @@ -585,7 +589,7 @@ async fn http_onion_performance_baseline() { #[cfg(not(debug_assertions))] async fn benchmark_http_request( plan: Option<&gateway_core::engine::middleware::FrozenMiddlewarePlan>, - settings: &gateway_core::settings::RequestSettings, + settings: &gateway_core::routing::request_settings::RequestSettings, payload: Bytes, uppercase: bool, ) -> (Duration, Duration) { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs index b724bd3e8..b7fa01445 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs @@ -2,7 +2,7 @@ use super::*; use futures::FutureExt as _; -use gateway_core::middleware::{http::upgrade, websocket as ws}; +use gateway_core::engine::middleware::{http::upgrade, websocket as ws}; use std::sync::atomic::AtomicBool; struct Transport { @@ -75,8 +75,8 @@ async fn session_model_and_network_callbacks_outlive_handshake_and_each_executio environment.store.start_command_line_writes().unwrap(); let snapshot = bundle.snapshots().snapshot_for_diagnostics().unwrap(); let plan = runtime - .middleware_registry() - .resolve(snapshot.extensions().unwrap()) + .execution_registry() + .middleware(snapshot.extensions().unwrap()) .unwrap(); let (send, incoming) = tokio::sync::mpsc::channel(1); let (outgoing, mut receive) = tokio::sync::mpsc::channel(1); diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs index cbd7ea501..4d30c7702 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs @@ -51,8 +51,8 @@ use gateway_core::{ lifecycle::CancellationToken, operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload}, policy::ClientApiKeyId, + routing::extensions::{ExtensionPreparationPort, ExtensionSetReference}, routing::{AccountGroupId, ConfigRevision, PublicModelId}, - runtime::extensions::{ExtensionPreparationPort, ExtensionSetReference}, }; use gateway_plugin_runtime::{ PackageInspector, PackageLimits, PluginRuntime, PluginRuntimeConfig, RpcLimits, @@ -766,8 +766,8 @@ async fn attempt_middleware_can_delegate_after_a_managed_http_side_effect() { .await; let generation = prepare(&runtime).await; let plan = runtime - .middleware_registry() - .resolve(&generation) + .execution_registry() + .middleware(&generation) .expect("middleware plan"); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); @@ -821,8 +821,8 @@ async fn middleware_next_and_lazy_body_mapping_use_one_real_rpc_call() { .await; let generation = prepare(&runtime).await; let plan = runtime - .middleware_registry() - .resolve(&generation) + .execution_registry() + .middleware(&generation) .expect("middleware plan"); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); @@ -974,7 +974,10 @@ async fn middleware_classifies_real_rpc_output_after_a_terminal_frame() { }]) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); let response = plan @@ -1035,7 +1038,10 @@ async fn outer_plugin_receives_inner_rejection_details_and_preserves_rejection_s InstanceFixture {id:"inner-error", configuration:serde_json::json!({"mode":"rejected"}), bindings:vec![binding(MIDDLEWARE_CONTRIBUTION,"request",1,PluginFailurePolicy::Reject)]}, ],package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let result = plan .handle( middleware_context(ClientTransport::HttpJson), @@ -1090,7 +1096,10 @@ async fn real_request_plugin_receives_downstream_error_details_without_consuming bindings:vec![binding(MIDDLEWARE_CONTRIBUTION,"request",0,PluginFailurePolicy::Reject)], }],package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let result = plan .handle( middleware_context(ClientTransport::HttpJson), @@ -1140,7 +1149,9 @@ async fn real_request_plugin_receives_downstream_error_details_without_consuming #[tokio::test] async fn real_request_plugins_rewrite_settings_in_onion_order() { - use gateway_core::{policy::RateLimits, settings::ExecutionSettings, settings::SettingsValues}; + use gateway_core::{ + policy::RateLimits, routing::request_settings::ExecutionSettings, settings::SettingsValues, + }; let profiles = |identity| { BTreeMap::from([( gateway_core::identity::ProviderKind::new("openai").unwrap(), @@ -1178,26 +1189,27 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() { .with_request_profiles(profiles("host")), ) .unwrap(); - let configuration = gateway_core::settings::RequestSettings::new(Arc::new(host_snapshot)) - .with_execution( - &gateway_core::policy::ClientPolicy::new( - middleware_context(ClientTransport::HttpJson) - .client_key_id() - .clone(), - gateway_core::policy::PlaintextClientApiKey::new("fixture-key").unwrap(), - Arc::new( - gateway_core::routing::FrozenAccountScope::new( - Arc::default(), - gateway_core::routing::ClientRoutingScope::all_accounts(), - ) - .with_request_profiles(profiles("key-default")) - .with_fast_mode(baseline.fast_mode), + let configuration = + gateway_core::routing::request_settings::RequestSettings::new(Arc::new(host_snapshot)) + .with_execution( + &gateway_core::policy::ClientPolicy::new( + middleware_context(ClientTransport::HttpJson) + .client_key_id() + .clone(), + gateway_core::policy::PlaintextClientApiKey::new("fixture-key").unwrap(), + Arc::new( + gateway_core::routing::FrozenAccountScope::new( + Arc::default(), + gateway_core::routing::ClientRoutingScope::all_accounts(), + ) + .with_request_profiles(profiles("key-default")) + .with_fast_mode(baseline.fast_mode), + ), + true, + baseline.client_limits, ), - true, - baseline.client_limits, - ), - baseline.timeout_ms, - ); + baseline.timeout_ms, + ); let mut first = baseline_json.clone(); first["fast_mode"] = serde_json::json!("default"); first["runtime"]["request_interval_ms"] = serde_json::json!(0); @@ -1223,7 +1235,10 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() { InstanceFixture { id: "settings-pass", configuration: serde_json::json!({"mode":"settings","expected_settings":last}), bindings: vec![binding(MIDDLEWARE_CONTRIBUTION,"request",2,PluginFailurePolicy::Reject)] }, ], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); // 连续调用使用同一宿主基线，前一次插件改写不能泄漏到后一次 for _ in 0..2 { let expected: ExecutionSettings = serde_json::from_value(last.clone()).unwrap(); @@ -1310,7 +1325,10 @@ async fn sdk_middleware_entry_registers_and_maps_a_real_process_response() { ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let closes = Arc::new(AtomicUsize::new(0)); let reads = Arc::new(AtomicUsize::new(0)); @@ -1448,7 +1466,12 @@ async fn instances_without_data_plane_bindings_publish_neither_plan() { .await; let generation = prepare(&runtime).await; assert!(runtime.policy_registry().resolve(&generation).is_none()); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); drop(generation); super::wait_until_empty(cache.path()).await; } @@ -1522,7 +1545,10 @@ async fn sdk_capability_declarations_require_request_stage() { bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, if attempt { "attempt" } else { "request" }, 0, PluginFailurePolicy::Reject)], }], package).await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let context = if attempt { attempt_middleware_context(Arc::default()) @@ -1604,7 +1630,10 @@ async fn assert_sdk_transparency(mode: &'static str) { ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let mut headers = middleware_headers(); headers.extend([ MiddlewareHeader::new("X-Repeated", Bytes::from_static(b"first")), @@ -1777,7 +1806,10 @@ async fn a_failed_middleware_preserves_scope_and_the_configured_failure_policy() }]) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let response = plan .handle( @@ -1853,7 +1885,10 @@ async fn middleware_preserves_large_body_and_following_requests() { }]) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); for length in [4 * 1024 * 1024 + 17, 32] { let payload = Bytes::from( (0..length) diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs index d44f76fea..aeefa995a 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mounts.rs @@ -1,7 +1,10 @@ //! 验证混合中间件挂载保持独立顺序与各次调用的续接状态 use super::*; -use gateway_core::middleware::{compose, http, service, websocket}; +use gateway_core::{ + engine::middleware::{http, service, websocket}, + middleware::compose, +}; use http_body_util::BodyExt as _; struct Sender; @@ -71,7 +74,10 @@ async fn mixed_mounts_keep_independent_order_and_fresh_continuations() { ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); assert!(plan.has_http() && plan.has_websocket() && plan.has_service()); let mut expected = Vec::new(); for direction in [ diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs index a39bb461f..f46f0337c 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/service/mod.rs @@ -294,10 +294,10 @@ async fn cli_service_callback_uses_the_same_public_settings_and_plugin_chain() { crate::support::native::admin_registry(), Arc::new(Pricing), ); - let mut registry = gateway_admin::service::Registry::new({ + let mut registry = gateway_admin::public_service::Registry::new({ let snapshots = core.snapshots(); - let middleware = runtime.middleware_registry(); - Arc::new(move || middleware.resolve(snapshots.snapshot_for_diagnostics()?.extensions()?)) + let middleware = runtime.execution_registry(); + Arc::new(move || middleware.middleware(snapshots.snapshot_for_diagnostics()?.extensions()?)) }); registry.register_settings(&settings).unwrap(); let registry = Arc::new(registry); @@ -345,7 +345,7 @@ async fn cli_service_callback_uses_the_same_public_settings_and_plugin_chain() { #[tokio::test] async fn service_chain_preserves_onion_order_parent_scope_and_native_results_without_a_database() { use super::*; - use gateway_core::middleware::{compose, service as core}; + use gateway_core::{engine::middleware::service as core, middleware::compose}; let directory = tempfile::tempdir().unwrap(); let records = directory.path().join("onion.jsonl"); @@ -377,7 +377,10 @@ async fn service_chain_preserves_onion_order_parent_scope_and_native_results_wit ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let input = json!({"opaque":[null,false,0,""]}); for reject in [false, true] { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs index d6df33241..db14cf597 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/websocket/mod.rs @@ -1,7 +1,7 @@ //! 验证真实插件 WebSocket 消息转发、改写与控制消息处理 use super::*; -use gateway_core::middleware::{compose, websocket as core}; +use gateway_core::{engine::middleware::websocket as core, middleware::compose}; #[derive(Default)] struct Sender(std::sync::Mutex>); @@ -46,7 +46,10 @@ async fn websocket_plugin( ) .await; let generation = prepare(&runtime).await; - let plan = runtime.middleware_registry().resolve(&generation).unwrap(); + let plan = runtime + .execution_registry() + .middleware(&generation) + .unwrap(); assert!(plan.has_websocket()); (cache, runtime, generation, plan) } diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs new file mode 100644 index 000000000..fcf3bbfa6 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/diagnostics.rs @@ -0,0 +1,191 @@ +//! 验证准备、发布进程故障与排空状态的诊断投影 + +use super::wait_until_unready; +use gateway_admin::model::Revision; +use gateway_admin::model::plugins::instances::PluginInstanceRuntimeStatus; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; + +#[tokio::test] +async fn diagnostics_distinguish_running_failed_preparation_and_draining_generations() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = store.snapshot.lock().unwrap().clone(); + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .unwrap(); + + let running = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("runtime diagnostics"); + let running = &running["instance-one"]; + assert_eq!(running.status, PluginInstanceRuntimeStatus::Running); + assert_eq!(running.actual_revision, Some(1)); + assert_eq!( + running.actual_artifact_sha256.as_deref(), + Some(snapshot.instances[0].artifact_sha256.as_str()) + ); + assert!(running.failure.is_none()); + + let mut failed = snapshot.clone(); + failed.config_revision = Revision::new(2).unwrap(); + failed.instances[0].revision = Revision::new(2).unwrap(); + failed.instances[0].configuration = serde_json::json!({"startup":"fail"}); + assert!( + PluginPreparation::prepare(&runtime, failed.clone()) + .await + .is_err() + ); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &failed, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("failed preparation diagnostics"); + let failed_runtime = &diagnostics["instance-one"]; + assert_eq!( + failed_runtime.status, + PluginInstanceRuntimeStatus::PreparationFailed + ); + assert!(failed_runtime.failure.is_some()); + assert_eq!(failed_runtime.actual_revision, Some(1)); + + let mut disabled = failed; + disabled.config_revision = Revision::new(3).unwrap(); + disabled.instances[0].revision = Revision::new(3).unwrap(); + disabled.instances[0].enabled = false; + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &disabled, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("draining diagnostics"); + let draining = &diagnostics["instance-one"]; + assert_eq!(draining.status, PluginInstanceRuntimeStatus::Draining); + assert_eq!(draining.draining_revisions, [1]); + + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn preparing_diagnostics_never_wait_for_the_serialized_prepare_io() { + let (cache, store, runtime) = super::super::setup().await; + let source = store.snapshot.lock().unwrap().clone(); + let mut candidate = source.clone(); + candidate.config_revision = Revision::new(2).unwrap(); + candidate.instances[0].revision = Revision::new(2).unwrap(); + candidate.instances[0].configuration = serde_json::json!({"startup_delay_ms":400}); + let diagnostic_snapshot = candidate.clone(); + let mut preparation = Box::pin(PluginPreparation::prepare(&runtime, candidate)); + assert!(futures::poll!(preparation.as_mut()).is_pending()); + + let diagnostics = tokio::time::timeout( + std::time::Duration::from_millis(100), + gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &diagnostic_snapshot, + None, + None, + ), + ) + .await + .expect("diagnostics must not wait for plugin startup") + .expect("runtime diagnostics"); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Preparing + ); + + let candidate = preparation.await.expect("delayed candidate"); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn diagnostics_report_a_published_process_fault_without_plugin_details() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({"exit_after_ready_ms":20}); + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("candidate starts before the controlled exit"); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while active.is_ready() { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("worker exit observed"); + + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("runtime diagnostics"); + let fault = &diagnostics["instance-one"]; + assert_eq!(fault.status, PluginInstanceRuntimeStatus::Faulted); + assert_eq!( + fault.failure.as_ref().map(|failure| failure.code.as_str()), + Some("closed") + ); + assert_eq!( + fault + .failure + .as_ref() + .map(|failure| failure.message.as_str()), + Some("插件进程或传输已停止") + ); + + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn a_published_process_crash_preserves_the_serving_snapshot_and_other_plugin_status() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = { + let mut snapshot = store.snapshot.lock().unwrap(); + let mut failed = snapshot.instances[0].clone(); + failed.id = "crashing-plugin".into(); + failed.configuration = serde_json::json!({"exit_after_ready_ms":100}); + snapshot.instances.push(failed); + snapshot.clone() + }; + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + wait_until_unready(&generation).await; + assert!(generation.can_serve()); + assert!(!generation.is_ready()); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&generation), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Running + ); + assert_eq!( + diagnostics["crashing-plugin"].status, + PluginInstanceRuntimeStatus::Faulted + ); + drop(generation); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs new file mode 100644 index 000000000..d711f0584 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs @@ -0,0 +1,511 @@ +//! 验证实例握手、声明与重启熔断的生命周期 + +use super::wait_until_unready; +use gateway_admin::model::Revision; +use gateway_admin::model::plugins::instances::PluginInstanceRuntimeStatus; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; +use gateway_plugin_sdk::Capability; +use gateway_plugin_sdk::Contributions; +use gateway_plugin_sdk::Stage; +use std::num::NonZeroUsize; +use std::path::Path; +use std::time::Duration; + +fn startup_count(marker: &Path) -> usize { + std::fs::read_to_string(marker) + .map(|content| content.lines().count()) + .unwrap_or_default() +} + +#[tokio::test] +async fn registration_with_a_changed_contribution_id_is_rejected() { + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::Observer, + vec![Stage::Observation], + Vec::new(), + Vec::new(), + ), + ])) + .await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("the declared contribution prepares before its registration id changes"); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; + + snapshot.instances[0].configuration = serde_json::json!({"registration_mismatch":true}); + let error = PluginPreparation::prepare(&runtime, snapshot) + .await + .expect_err("a changed registration contribution id must be rejected"); + assert!( + error.to_string().contains("插件注册结果与清单不符"), + "unexpected preparation failure: {error}" + ); +} + +#[tokio::test] +async fn restart_circuit_caps_short_lived_crashes_and_a_new_revision_recovers() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_secs(1), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("starts.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker": marker, + "exit_after_ready_ms": 20, + }); + let mut active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("first incarnation"); + for expected_start in 2..=3 { + wait_until_unready(&active).await; + let replacement = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("failure budget still permits a restart"); + drop(active); + active = replacement; + assert_eq!(startup_count(&marker), expected_start); + } + wait_until_unready(&active).await; + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err(), + "the fourth incarnation must not start" + ); + assert_eq!(startup_count(&marker), 3); + + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(snapshot.config_revision.get()), + Some(&active), + ) + .await + .expect("runtime diagnostics"); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("restart_circuit_open") + ); + + let mut disabled = snapshot; + disabled.config_revision = Revision::new(2).unwrap(); + disabled.instances[0].revision = Revision::new(2).unwrap(); + disabled.instances[0].enabled = false; + let disabled_generation = PluginPreparation::prepare(&runtime, disabled.clone()) + .await + .expect("disabling the instance does not reuse its open circuit"); + assert_eq!(startup_count(&marker), 3); + + let mut repaired = disabled; + repaired.config_revision = Revision::new(3).unwrap(); + repaired.instances[0].revision = Revision::new(3).unwrap(); + repaired.instances[0].enabled = true; + repaired.instances[0].configuration = serde_json::json!({"startup_marker":marker}); + let replacement = PluginPreparation::prepare(&runtime, repaired) + .await + .expect("a new instance revision has an independent failure budget"); + assert!(replacement.is_ready()); + assert_eq!(startup_count(&marker), 4); + + drop(active); + drop(disabled_generation); + drop(replacement); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_resets_after_a_stable_incarnation() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_millis(80), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("stable-reset-starts.jsonl"); + let exit = controls.path().join("stable-exit"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "startup_failures":[true,true,false,true,true], + "exit_after_ready_signals":[null,null,exit], + }); + // 握手前失败的有效运行时间固定为零；不能用短 sleep 假定宿主一定及时观察到退出 + for _ in 0..2 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + } + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("the third incarnation starts before the circuit opens"); + assert!(active.is_ready()); + // 明确等到稳定窗口之后才允许第三个进程退出；调度变慢不会把短命进程误变为稳定进程 + tokio::time::sleep(restart_circuit.stability_window).await; + tokio::fs::write(exit, b"exit").await.unwrap(); + wait_until_unready(&active).await; + for expected_start in 4..=5 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + assert_eq!(startup_count(&marker), expected_start); + } + assert!( + PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err(), + "three new failures after the stable incarnation reopen the circuit" + ); + assert_eq!(startup_count(&marker), 5); + + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_ignores_planned_generation_shutdown() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::MIN, + stability_window: Duration::from_secs(60), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("planned-shutdown-starts.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({"startup_marker":marker}); + let first = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("first generation"); + let instance = &snapshot.instances[0]; + gateway_admin::ports::plugins::PluginStateLifecycle::quiesce_instance( + &runtime, + &instance.id, + &instance.artifact_sha256, + instance.revision, + ) + .await; + assert!(!first.is_ready()); + + let second = PluginPreparation::prepare(&runtime, snapshot) + .await + .expect("planned shutdown must not open the one-failure circuit"); + assert!(second.is_ready()); + assert_eq!(startup_count(&marker), 2); + + drop(first); + drop(second); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_counts_exit_before_the_candidate_is_indexed() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_secs(1), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("registration-exits.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "exit_during_registration":true, + }); + for _ in 0..3 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + } + assert_eq!(startup_count(&marker), 3); + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + assert_eq!( + startup_count(&marker), + 3, + "open circuit must gate process spawn" + ); + + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, &snapshot, None, None, + ) + .await + .expect("runtime diagnostics"); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("restart_circuit_open") + ); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_counts_a_started_process_that_exits_during_handshake() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(2).unwrap(), + stability_window: Duration::from_millis(10), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("handshake-exits.jsonl"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "startup":"fail", + "startup_fail_delay_ms":50, + }); + for _ in 0..2 { + assert!( + PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .is_err() + ); + } + assert_eq!(startup_count(&marker), 2); + assert!( + PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err() + ); + assert_eq!( + startup_count(&marker), + 2, + "host-side gating must happen before another process spawn" + ); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_does_not_let_an_older_generation_clear_newer_failures() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::new(3).unwrap(), + stability_window: Duration::from_millis(50), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("overlapping-generations.jsonl"); + let old_exit = controls.path().join("old-generation-exit"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "startup_failures":[false,true,true,true], + "exit_after_ready_signals":[old_exit], + }); + let mut other = snapshot.instances[0].clone(); + other.id = "other-instance".into(); + other.name = "Other".into(); + other.configuration = serde_json::json!({}); + snapshot.instances.push(other); + let published = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("published generation"); + assert!(published.is_ready()); + // 旧代次明确越过稳定窗口；新代次在握手前失败，运行时间固定为零 + // 不依赖 15 ms 退出与注册完成的竞速，也不让调度延迟重置新失败预算 + tokio::time::sleep(restart_circuit.stability_window).await; + + let mut failed_candidate = None; + for revision in 2..=4 { + snapshot.config_revision = Revision::new(revision).unwrap(); + snapshot.instances[1].revision = Revision::new(revision).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":revision}); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("the unrelated instance starts while the newer incarnation fails"); + assert_eq!(startup_count(&marker), revision as usize); + drop(failed_candidate.replace(candidate)); + assert!(published.is_ready(), "the older generation remains healthy"); + } + drop(failed_candidate.take()); + std::fs::write(&old_exit, b"exit").unwrap(); + wait_until_unready(&published).await; + snapshot.config_revision = Revision::new(5).unwrap(); + snapshot.instances[1].revision = Revision::new(5).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":5}); + let circuit = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, &snapshot, None, None, + ) + .await + .unwrap(); + assert_eq!( + circuit["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("restart_circuit_open"), + ); + let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("an older failed instance is isolated during unrelated recovery"); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(5), + Some(&recovered), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("unavailable"), + "the older healthy session must not erase newer failures for the same identity" + ); + assert_eq!(startup_count(&marker), 4); + + drop(recovered); + drop(published); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn restart_circuit_restores_an_older_failure_after_candidate_shutdown() { + let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { + maximum_failures: NonZeroUsize::MIN, + stability_window: Duration::from_secs(1), + }; + let (cache, store, runtime) = super::super::setup_with_restart_circuit(restart_circuit).await; + let controls = tempfile::tempdir().unwrap(); + let marker = controls.path().join("abandoned-candidate.jsonl"); + let old_exit = controls.path().join("abandoned-candidate-old-exit"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "startup_marker":marker, + "exit_after_ready_delays_ms":[null,null], + "exit_after_ready_signals":[old_exit,null], + }); + let mut other = snapshot.instances[0].clone(); + other.id = "other-instance".into(); + other.name = "Other".into(); + other.configuration = serde_json::json!({}); + snapshot.instances.push(other); + let published = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("published generation"); + + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[1].revision = Revision::new(2).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":2}); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("healthy replacement candidate"); + std::fs::write(&old_exit, b"exit").unwrap(); + wait_until_unready(&published).await; + assert!(candidate.is_ready()); + drop(candidate); + + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let diagnostics = + gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&published), + ) + .await + .expect("runtime diagnostics"); + if diagnostics["instance-one"] + .failure + .as_ref() + .is_some_and(|failure| failure.code == "restart_circuit_open") + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("older failure is restored after candidate shutdown"); + + snapshot.config_revision = Revision::new(3).unwrap(); + snapshot.instances[1].revision = Revision::new(3).unwrap(); + snapshot.instances[1].configuration = serde_json::json!({"revision":3}); + let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .expect("an older failed instance is isolated during unrelated recovery"); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(3), + Some(&recovered), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["instance-one"] + .failure + .as_ref() + .map(|failure| failure.code.as_str()), + Some("unavailable"), + ); + assert_eq!(startup_count(&marker), 2); + + drop(recovered); + drop(published); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn a_host_version_mismatch_still_restores_the_enabled_plugin() { + use std::sync::Arc; + let (cache, store, _) = super::super::setup().await; + let runtime = gateway_plugin_runtime::PluginRuntime::new( + store.clone(), + store.clone(), + gateway_plugin_runtime::PluginRuntimeConfig { + cache_directory: cache.path().to_owned(), + host_version: "3.0.0".parse().unwrap(), + package_limits: Default::default(), + rpc_limits: Default::default(), + restart_circuit: Default::default(), + }, + Arc::new(gateway_host::outbound::HttpClient::new().unwrap()), + Arc::new(gateway_host::process::ProcessSupervisor::default()), + ); + let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert!(restored.can_serve()); + let snapshot = store.snapshot.lock().unwrap().clone(); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&restored), + ) + .await + .unwrap(); + assert!(restored.is_ready()); + assert!(snapshot.instances[0].enabled); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Running + ); + assert!(diagnostics["instance-one"].failure.is_none()); + drop(restored); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs new file mode 100644 index 000000000..e674a5fd1 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/management.rs @@ -0,0 +1,50 @@ +//! 验证管理调用绑定、执行身份与候选配置 + +use gateway_admin::model::plugins::instances::PluginCapabilityBinding; +use gateway_admin::model::plugins::instances::PluginFailurePolicy; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_plugin_sdk::Capability; +use gateway_plugin_sdk::Contributions; +use gateway_plugin_sdk::Stage; + +#[tokio::test] +async fn management_needs_no_binding_and_rejects_stale_execution_identity_configuration() { + let (_cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::Management, + vec![Stage::Management], + Vec::new(), + Vec::new(), + ), + ])) + .await; + let mut instance = store.snapshot.lock().unwrap().instances[0].clone(); + assert!( + PluginPreparation::validate(&runtime, instance.clone()) + .await + .is_ok() + ); + instance.bindings = vec![PluginCapabilityBinding { + contribution: "test.example.management".into(), + stage: "management".into(), + order: 0, + failure_policy: PluginFailurePolicy::Reject, + client_key_ids: vec!["key-one".into(), "key-two".into()], + account_group_ids: Vec::new(), + provider_ids: Vec::new(), + models: Vec::new(), + event: None, + identity_bindings: Vec::new(), + }]; + assert!( + PluginPreparation::validate(&runtime, instance.clone()) + .await + .is_err() + ); + instance.enabled = false; + assert!( + PluginPreparation::validate(&runtime, instance) + .await + .is_err() + ); +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs index 101d49c44..e7531c5cc 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs @@ -1,19 +1,15 @@ -//! 验证插件发布准备的目录冻结、声明校验与运行时关闭 +//! 插件发布准备各职责的集成测试入口 -use gateway_admin::{ - model::{ - Revision, - plugins::instances::{ - PluginCapabilityBinding, PluginFailurePolicy, PluginInstanceRuntimeStatus, - }, - }, - ports::plugins::PluginPreparation, -}; -use gateway_core::{routing::ConfigRevision, runtime::extensions::ExtensionPreparationPort}; -use gateway_plugin_sdk::{Capability, Contributions, Stage}; -use std::{num::NonZeroUsize, path::Path, time::Duration}; +mod diagnostics; +mod instance; +mod management; +mod runtime; +mod set; +mod snapshot; -async fn wait_until_unready(reference: &gateway_core::runtime::extensions::ExtensionSetReference) { +use std::time::Duration; + +async fn wait_until_unready(reference: &gateway_core::routing::extensions::ExtensionSetReference) { tokio::time::timeout(Duration::from_secs(2), async { while reference.is_ready() { tokio::time::sleep(Duration::from_millis(10)).await; @@ -22,1003 +18,3 @@ async fn wait_until_unready(reference: &gateway_core::runtime::extensions::Exten .await .expect("worker exit observed"); } - -fn startup_count(marker: &Path) -> usize { - std::fs::read_to_string(marker) - .map(|content| content.lines().count()) - .unwrap_or_default() -} - -#[tokio::test] -async fn model_catalog_is_frozen_without_entering_the_request_chain() { - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::ModelCatalog, - vec![Stage::Registration], - vec![], - vec![], - )])) - .await; - store.snapshot.lock().unwrap().instances[0].configuration = serde_json::json!({ - "model_catalog": {"models": [{"id": "my-model", "provider": "openai", "model": "upstream-model"}]} - }); - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert_eq!(generation.model_aliases().len(), 1); - assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); - assert_eq!( - generation.model_aliases()[0].target.as_str(), - "upstream-model" - ); - assert!(runtime.policy_registry().resolve(&generation).is_none()); - assert!(runtime.middleware_registry().resolve(&generation).is_none()); - let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert_eq!(reused.id(), generation.id()); - assert_eq!(reused.model_aliases(), generation.model_aliases()); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[0].revision = Revision::new(2).unwrap(); - snapshot.instances[0].configuration = serde_json::json!({"startup":"fail"}); - } - assert!( - ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .is_err(), - "failed recovery must not publish a catalog with silently removed aliases" - ); - assert!(generation.is_ready()); - assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[0].enabled = false; - } - let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .unwrap(); - assert!(disabled.model_aliases().is_empty()); - assert_eq!(generation.model_aliases().len(), 1); - drop(disabled); - drop(reused); - drop(generation); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn model_catalog_rejects_invalid_registrations_and_cross_instance_conflicts() { - for models in [ - serde_json::json!([{"id":"alias","provider":"custom","model":"target"}]), - serde_json::json!([{"id":"alias","provider":"openai","model":"alias"}]), - serde_json::json!([ - {"id":"alias","provider":"openai","model":"target"}, - {"id":"alias","provider":"xai","model":"target"} - ]), - serde_json::json!([{"id":"alias","provider":"openai","model":"target","capabilities":[]}]), - ] { - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::ModelCatalog, - vec![Stage::Registration], - vec![], - vec![], - )])) - .await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = - serde_json::json!({"model_catalog":{"models":models}}); - assert!( - PluginPreparation::prepare(&runtime, snapshot) - .await - .is_err() - ); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; - } - - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::ModelCatalog, - vec![Stage::Registration], - vec![], - vec![], - )])) - .await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "model_catalog":{"models":[{"id":"alias","provider":"openai","model":"target"}]} - }); - let mut other = snapshot.instances[0].clone(); - other.id = "instance-two".into(); - snapshot.instances.push(other); - let error = PluginPreparation::prepare(&runtime, snapshot) - .await - .unwrap_err(); - assert!(error.to_string().contains("instance-one")); - assert!(error.to_string().contains("instance-two")); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn runtime_shutdown_reaps_live_generation_and_rejects_late_preparation() { - let (cache, store, runtime) = super::setup().await; - let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert!(published.is_ready()); - assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); - - runtime.shutdown().await; - assert!(!published.is_ready()); - super::wait_until_empty(cache.path()).await; - assert!( - ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .is_err() - ); - runtime.shutdown().await; - - drop(published); - drop(store); -} - -#[tokio::test] -async fn management_needs_no_binding_and_rejects_stale_execution_identity_configuration() { - let (_cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::Management, - vec![Stage::Management], - Vec::new(), - Vec::new(), - )])) - .await; - let mut instance = store.snapshot.lock().unwrap().instances[0].clone(); - assert!( - PluginPreparation::validate(&runtime, instance.clone()) - .await - .is_ok() - ); - instance.bindings = vec![PluginCapabilityBinding { - contribution: "test.example.management".into(), - stage: "management".into(), - order: 0, - failure_policy: PluginFailurePolicy::Reject, - client_key_ids: vec!["key-one".into(), "key-two".into()], - account_group_ids: Vec::new(), - provider_ids: Vec::new(), - models: Vec::new(), - event: None, - identity_bindings: Vec::new(), - }]; - assert!( - PluginPreparation::validate(&runtime, instance.clone()) - .await - .is_err() - ); - instance.enabled = false; - assert!( - PluginPreparation::validate(&runtime, instance) - .await - .is_err() - ); -} - -#[tokio::test] -async fn registration_with_a_changed_contribution_id_is_rejected() { - let (cache, store, runtime) = - super::setup_with_contributions(Contributions::from([crate::support::contribution( - Capability::Observer, - vec![Stage::Observation], - Vec::new(), - Vec::new(), - )])) - .await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("the declared contribution prepares before its registration id changes"); - drop(candidate); - super::wait_until_empty(cache.path()).await; - - snapshot.instances[0].configuration = serde_json::json!({"registration_mismatch":true}); - let error = PluginPreparation::prepare(&runtime, snapshot) - .await - .expect_err("a changed registration contribution id must be rejected"); - assert!( - error.to_string().contains("插件注册结果与清单不符"), - "unexpected preparation failure: {error}" - ); -} - -#[tokio::test] -async fn candidate_is_reused_after_commit_and_released_when_no_generation_owner_remains() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = - serde_json::from_str(r#"{"z":1,"a":{"z":2,"a":3}}"#).unwrap(); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .unwrap(); - assert!( - runtime.observer_registry().resolve(&candidate).is_none(), - "没有观察绑定时不应编译派发计划" - ); - snapshot.instances[0].configuration = - serde_json::from_str(r#"{"a":{"a":3,"z":2},"z":1}"#).unwrap(); - snapshot.config_revision = Revision::new(2).unwrap(); - *store.snapshot.lock().unwrap() = snapshot; - let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .unwrap(); - assert_eq!(candidate.id(), published.id()); - assert!(published.is_ready()); - drop(candidate); - assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); - drop(published); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn diagnostics_distinguish_running_failed_preparation_and_draining_generations() { - let (cache, store, runtime) = super::setup().await; - let snapshot = store.snapshot.lock().unwrap().clone(); - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .unwrap(); - - let running = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("runtime diagnostics"); - let running = &running["instance-one"]; - assert_eq!(running.status, PluginInstanceRuntimeStatus::Running); - assert_eq!(running.actual_revision, Some(1)); - assert_eq!( - running.actual_artifact_sha256.as_deref(), - Some(snapshot.instances[0].artifact_sha256.as_str()) - ); - assert!(running.failure.is_none()); - - let mut failed = snapshot.clone(); - failed.config_revision = Revision::new(2).unwrap(); - failed.instances[0].revision = Revision::new(2).unwrap(); - failed.instances[0].configuration = serde_json::json!({"startup":"fail"}); - assert!( - PluginPreparation::prepare(&runtime, failed.clone()) - .await - .is_err() - ); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &failed, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("failed preparation diagnostics"); - let failed_runtime = &diagnostics["instance-one"]; - assert_eq!( - failed_runtime.status, - PluginInstanceRuntimeStatus::PreparationFailed - ); - assert!(failed_runtime.failure.is_some()); - assert_eq!(failed_runtime.actual_revision, Some(1)); - - let mut disabled = failed; - disabled.config_revision = Revision::new(3).unwrap(); - disabled.instances[0].revision = Revision::new(3).unwrap(); - disabled.instances[0].enabled = false; - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &disabled, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("draining diagnostics"); - let draining = &diagnostics["instance-one"]; - assert_eq!(draining.status, PluginInstanceRuntimeStatus::Draining); - assert_eq!(draining.draining_revisions, [1]); - - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn preparing_diagnostics_never_wait_for_the_serialized_prepare_io() { - let (cache, store, runtime) = super::setup().await; - let source = store.snapshot.lock().unwrap().clone(); - let mut candidate = source.clone(); - candidate.config_revision = Revision::new(2).unwrap(); - candidate.instances[0].revision = Revision::new(2).unwrap(); - candidate.instances[0].configuration = serde_json::json!({"startup_delay_ms":400}); - let diagnostic_snapshot = candidate.clone(); - let mut preparation = Box::pin(PluginPreparation::prepare(&runtime, candidate)); - assert!(futures::poll!(preparation.as_mut()).is_pending()); - - let diagnostics = tokio::time::timeout( - std::time::Duration::from_millis(100), - gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &diagnostic_snapshot, - None, - None, - ), - ) - .await - .expect("diagnostics must not wait for plugin startup") - .expect("runtime diagnostics"); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Preparing - ); - - let candidate = preparation.await.expect("delayed candidate"); - drop(candidate); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn diagnostics_report_a_published_process_fault_without_plugin_details() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({"exit_after_ready_ms":20}); - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("candidate starts before the controlled exit"); - tokio::time::timeout(std::time::Duration::from_secs(2), async { - while active.is_ready() { - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - } - }) - .await - .expect("worker exit observed"); - - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("runtime diagnostics"); - let fault = &diagnostics["instance-one"]; - assert_eq!(fault.status, PluginInstanceRuntimeStatus::Faulted); - assert_eq!( - fault.failure.as_ref().map(|failure| failure.code.as_str()), - Some("closed") - ); - assert_eq!( - fault - .failure - .as_ref() - .map(|failure| failure.message.as_str()), - Some("插件进程或传输已停止") - ); - - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_caps_short_lived_crashes_and_a_new_revision_recovers() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_secs(1), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("starts.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker": marker, - "exit_after_ready_ms": 20, - }); - let mut active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("first incarnation"); - for expected_start in 2..=3 { - wait_until_unready(&active).await; - let replacement = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("failure budget still permits a restart"); - drop(active); - active = replacement; - assert_eq!(startup_count(&marker), expected_start); - } - wait_until_unready(&active).await; - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err(), - "the fourth incarnation must not start" - ); - assert_eq!(startup_count(&marker), 3); - - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(snapshot.config_revision.get()), - Some(&active), - ) - .await - .expect("runtime diagnostics"); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("restart_circuit_open") - ); - - let mut disabled = snapshot; - disabled.config_revision = Revision::new(2).unwrap(); - disabled.instances[0].revision = Revision::new(2).unwrap(); - disabled.instances[0].enabled = false; - let disabled_generation = PluginPreparation::prepare(&runtime, disabled.clone()) - .await - .expect("disabling the instance does not reuse its open circuit"); - assert_eq!(startup_count(&marker), 3); - - let mut repaired = disabled; - repaired.config_revision = Revision::new(3).unwrap(); - repaired.instances[0].revision = Revision::new(3).unwrap(); - repaired.instances[0].enabled = true; - repaired.instances[0].configuration = serde_json::json!({"startup_marker":marker}); - let replacement = PluginPreparation::prepare(&runtime, repaired) - .await - .expect("a new instance revision has an independent failure budget"); - assert!(replacement.is_ready()); - assert_eq!(startup_count(&marker), 4); - - drop(active); - drop(disabled_generation); - drop(replacement); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_resets_after_a_stable_incarnation() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_millis(80), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("stable-reset-starts.jsonl"); - let exit = controls.path().join("stable-exit"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "startup_failures":[true,true,false,true,true], - "exit_after_ready_signals":[null,null,exit], - }); - // 握手前失败的有效运行时间固定为零；不能用短 sleep 假定宿主一定及时观察到退出 - for _ in 0..2 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - } - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("the third incarnation starts before the circuit opens"); - assert!(active.is_ready()); - // 明确等到稳定窗口之后才允许第三个进程退出；调度变慢不会把短命进程误变为稳定进程 - tokio::time::sleep(restart_circuit.stability_window).await; - tokio::fs::write(exit, b"exit").await.unwrap(); - wait_until_unready(&active).await; - for expected_start in 4..=5 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - assert_eq!(startup_count(&marker), expected_start); - } - assert!( - PluginPreparation::prepare(&runtime, snapshot) - .await - .is_err(), - "three new failures after the stable incarnation reopen the circuit" - ); - assert_eq!(startup_count(&marker), 5); - - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_ignores_planned_generation_shutdown() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::MIN, - stability_window: Duration::from_secs(60), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("planned-shutdown-starts.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({"startup_marker":marker}); - let first = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("first generation"); - let instance = &snapshot.instances[0]; - gateway_admin::ports::plugins::PluginStateLifecycle::quiesce_instance( - &runtime, - &instance.id, - &instance.artifact_sha256, - instance.revision, - ) - .await; - assert!(!first.is_ready()); - - let second = PluginPreparation::prepare(&runtime, snapshot) - .await - .expect("planned shutdown must not open the one-failure circuit"); - assert!(second.is_ready()); - assert_eq!(startup_count(&marker), 2); - - drop(first); - drop(second); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_counts_exit_before_the_candidate_is_indexed() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_secs(1), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("registration-exits.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "exit_during_registration":true, - }); - for _ in 0..3 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - } - assert_eq!(startup_count(&marker), 3); - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - assert_eq!( - startup_count(&marker), - 3, - "open circuit must gate process spawn" - ); - - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, &snapshot, None, None, - ) - .await - .expect("runtime diagnostics"); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("restart_circuit_open") - ); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_counts_a_started_process_that_exits_during_handshake() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(2).unwrap(), - stability_window: Duration::from_millis(10), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("handshake-exits.jsonl"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "startup":"fail", - "startup_fail_delay_ms":50, - }); - for _ in 0..2 { - assert!( - PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .is_err() - ); - } - assert_eq!(startup_count(&marker), 2); - assert!( - PluginPreparation::prepare(&runtime, snapshot) - .await - .is_err() - ); - assert_eq!( - startup_count(&marker), - 2, - "host-side gating must happen before another process spawn" - ); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_does_not_let_an_older_generation_clear_newer_failures() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::new(3).unwrap(), - stability_window: Duration::from_millis(50), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("overlapping-generations.jsonl"); - let old_exit = controls.path().join("old-generation-exit"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "startup_failures":[false,true,true,true], - "exit_after_ready_signals":[old_exit], - }); - let mut other = snapshot.instances[0].clone(); - other.id = "other-instance".into(); - other.name = "Other".into(); - other.configuration = serde_json::json!({}); - snapshot.instances.push(other); - let published = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("published generation"); - assert!(published.is_ready()); - // 旧代次明确越过稳定窗口；新代次在握手前失败，运行时间固定为零 - // 不依赖 15 ms 退出与注册完成的竞速，也不让调度延迟重置新失败预算 - tokio::time::sleep(restart_circuit.stability_window).await; - - let mut failed_candidate = None; - for revision in 2..=4 { - snapshot.config_revision = Revision::new(revision).unwrap(); - snapshot.instances[1].revision = Revision::new(revision).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":revision}); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("the unrelated instance starts while the newer incarnation fails"); - assert_eq!(startup_count(&marker), revision as usize); - drop(failed_candidate.replace(candidate)); - assert!(published.is_ready(), "the older generation remains healthy"); - } - drop(failed_candidate.take()); - std::fs::write(&old_exit, b"exit").unwrap(); - wait_until_unready(&published).await; - snapshot.config_revision = Revision::new(5).unwrap(); - snapshot.instances[1].revision = Revision::new(5).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":5}); - let circuit = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, &snapshot, None, None, - ) - .await - .unwrap(); - assert_eq!( - circuit["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("restart_circuit_open"), - ); - let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("an older failed instance is isolated during unrelated recovery"); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(5), - Some(&recovered), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("unavailable"), - "the older healthy session must not erase newer failures for the same identity" - ); - assert_eq!(startup_count(&marker), 4); - - drop(recovered); - drop(published); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn restart_circuit_restores_an_older_failure_after_candidate_shutdown() { - let restart_circuit = gateway_plugin_runtime::PluginRestartCircuitConfig { - maximum_failures: NonZeroUsize::MIN, - stability_window: Duration::from_secs(1), - }; - let (cache, store, runtime) = super::setup_with_restart_circuit(restart_circuit).await; - let controls = tempfile::tempdir().unwrap(); - let marker = controls.path().join("abandoned-candidate.jsonl"); - let old_exit = controls.path().join("abandoned-candidate-old-exit"); - let mut snapshot = store.snapshot.lock().unwrap().clone(); - snapshot.instances[0].configuration = serde_json::json!({ - "startup_marker":marker, - "exit_after_ready_delays_ms":[null,null], - "exit_after_ready_signals":[old_exit,null], - }); - let mut other = snapshot.instances[0].clone(); - other.id = "other-instance".into(); - other.name = "Other".into(); - other.configuration = serde_json::json!({}); - snapshot.instances.push(other); - let published = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("published generation"); - - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[1].revision = Revision::new(2).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":2}); - let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("healthy replacement candidate"); - std::fs::write(&old_exit, b"exit").unwrap(); - wait_until_unready(&published).await; - assert!(candidate.is_ready()); - drop(candidate); - - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let diagnostics = - gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&published), - ) - .await - .expect("runtime diagnostics"); - if diagnostics["instance-one"] - .failure - .as_ref() - .is_some_and(|failure| failure.code == "restart_circuit_open") - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("older failure is restored after candidate shutdown"); - - snapshot.config_revision = Revision::new(3).unwrap(); - snapshot.instances[1].revision = Revision::new(3).unwrap(); - snapshot.instances[1].configuration = serde_json::json!({"revision":3}); - let recovered = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .expect("an older failed instance is isolated during unrelated recovery"); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(3), - Some(&recovered), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["instance-one"] - .failure - .as_ref() - .map(|failure| failure.code.as_str()), - Some("unavailable"), - ); - assert_eq!(startup_count(&marker), 2); - - drop(recovered); - drop(published); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn failed_candidate_leaves_the_previous_generation_ready() { - let (cache, store, runtime) = super::setup().await; - let snapshot = store.snapshot.lock().unwrap().clone(); - let active = PluginPreparation::prepare(&runtime, snapshot.clone()) - .await - .unwrap(); - let mut candidate = snapshot; - candidate.instances[0].configuration = serde_json::json!({"startup":"fail"}); - assert!( - PluginPreparation::prepare(&runtime, candidate) - .await - .is_err() - ); - assert!(active.is_ready()); - drop(active); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn disabled_missing_artifacts_do_not_block_the_remaining_generation_or_offline_recovery() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - let mut missing = snapshot.instances[0].clone(); - missing.id = "missing".into(); - missing.enabled = false; - missing.artifact_sha256 = "0".repeat(64); - snapshot.instances.push(missing); - let active = PluginPreparation::prepare(&runtime, snapshot) - .await - .unwrap(); - assert!(active.is_ready()); - drop(active); - super::wait_until_empty(cache.path()).await; - let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert!(restored.is_ready()); - drop(restored); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn preparation_cannot_mix_a_newer_persisted_revision_with_an_older_snapshot() { - let (_cache, _store, runtime) = super::setup().await; - assert!( - ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .is_err() - ); -} - -#[tokio::test] -async fn restoring_a_failed_plugin_keeps_other_instances_ready_and_reports_only_its_failure() { - let (cache, store, runtime) = super::setup().await; - let snapshot = { - let mut snapshot = store.snapshot.lock().unwrap(); - let mut failed = snapshot.instances[0].clone(); - failed.id = "failed-plugin".into(); - failed.configuration = serde_json::json!({"startup":"fail"}); - snapshot.instances.insert(0, failed); - snapshot.clone() - }; - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .expect("one failed plugin must not stop restoration"); - assert!(generation.can_serve()); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&generation), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["failed-plugin"].status, - PluginInstanceRuntimeStatus::PreparationFailed - ); - assert!(diagnostics["failed-plugin"].failure.is_some()); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Running - ); - let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert_eq!( - generation.id(), - reused.id(), - "quarantined failures do not restart on every request" - ); - drop(reused); - drop(generation); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn a_published_process_crash_preserves_the_serving_snapshot_and_other_plugin_status() { - let (cache, store, runtime) = super::setup().await; - let snapshot = { - let mut snapshot = store.snapshot.lock().unwrap(); - let mut failed = snapshot.instances[0].clone(); - failed.id = "crashing-plugin".into(); - failed.configuration = serde_json::json!({"exit_after_ready_ms":100}); - snapshot.instances.push(failed); - snapshot.clone() - }; - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - wait_until_unready(&generation).await; - assert!(generation.can_serve()); - assert!(!generation.is_ready()); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&generation), - ) - .await - .unwrap(); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Running - ); - assert_eq!( - diagnostics["crashing-plugin"].status, - PluginInstanceRuntimeStatus::Faulted - ); - drop(generation); - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn a_host_version_mismatch_still_restores_the_enabled_plugin() { - use std::sync::Arc; - let (cache, store, _) = super::setup().await; - let runtime = gateway_plugin_runtime::PluginRuntime::new( - store.clone(), - store.clone(), - gateway_plugin_runtime::PluginRuntimeConfig { - cache_directory: cache.path().to_owned(), - host_version: "3.0.0".parse().unwrap(), - package_limits: Default::default(), - rpc_limits: Default::default(), - restart_circuit: Default::default(), - }, - Arc::new(gateway_host::outbound::HttpClient::new().unwrap()), - Arc::new(gateway_host::process::ProcessSupervisor::default()), - ); - let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - assert!(restored.can_serve()); - let snapshot = store.snapshot.lock().unwrap().clone(); - let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( - &runtime, - &snapshot, - Some(1), - Some(&restored), - ) - .await - .unwrap(); - assert!(restored.is_ready()); - assert!(snapshot.instances[0].enabled); - assert_eq!( - diagnostics["instance-one"].status, - PluginInstanceRuntimeStatus::Running - ); - assert!(diagnostics["instance-one"].failure.is_none()); - drop(restored); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; -} - -#[tokio::test] -async fn an_existing_failed_plugin_does_not_prevent_editing_an_unrelated_instance() { - let (cache, store, runtime) = super::setup().await; - let mut snapshot = store.snapshot.lock().unwrap().clone(); - let mut failed = snapshot.instances[0].clone(); - failed.id = "failed-plugin".into(); - failed.configuration = serde_json::json!({"startup":"fail"}); - snapshot.instances.push(failed); - snapshot.config_revision = Revision::new(2).unwrap(); - snapshot.instances[0].revision = snapshot.config_revision; - let candidate = PluginPreparation::prepare(&runtime, snapshot) - .await - .expect("only the instance being changed must successfully prepare"); - assert!(candidate.can_serve()); - drop(candidate); - super::wait_until_empty(cache.path()).await; -} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs new file mode 100644 index 000000000..ee3e633dd --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/runtime.rs @@ -0,0 +1,27 @@ +//! 验证运行时整体关闭回收当前代次并拒绝新候选 + +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; + +#[tokio::test] +async fn runtime_shutdown_reaps_live_generation_and_rejects_late_preparation() { + let (cache, store, runtime) = super::super::setup().await; + let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert!(published.is_ready()); + assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); + + runtime.shutdown().await; + assert!(!published.is_ready()); + super::super::wait_until_empty(cache.path()).await; + assert!( + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .is_err() + ); + runtime.shutdown().await; + + drop(published); + drop(store); +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs new file mode 100644 index 000000000..0fd200820 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/set.rs @@ -0,0 +1,34 @@ +//! 验证候选复用与最后一个代次持有者释放后的回收 + +use gateway_admin::model::Revision; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; + +#[tokio::test] +async fn candidate_is_reused_after_commit_and_released_when_no_generation_owner_remains() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = + serde_json::from_str(r#"{"z":1,"a":{"z":2,"a":3}}"#).unwrap(); + let candidate = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .unwrap(); + assert!( + runtime.observer_registry().resolve(&candidate).is_none(), + "没有观察绑定时不应编译派发计划" + ); + snapshot.instances[0].configuration = + serde_json::from_str(r#"{"a":{"a":3,"z":2},"z":1}"#).unwrap(); + snapshot.config_revision = Revision::new(2).unwrap(); + *store.snapshot.lock().unwrap() = snapshot; + let published = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .unwrap(); + assert_eq!(candidate.id(), published.id()); + assert!(published.is_ready()); + drop(candidate); + assert_eq!(std::fs::read_dir(cache.path()).unwrap().count(), 1); + drop(published); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs new file mode 100644 index 000000000..fe01dcb11 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/snapshot.rs @@ -0,0 +1,248 @@ +//! 验证完整快照准备、能力目录与失败实例恢复 + +use gateway_admin::model::Revision; +use gateway_admin::model::plugins::instances::PluginInstanceRuntimeStatus; +use gateway_admin::ports::plugins::PluginPreparation; +use gateway_core::routing::ConfigRevision; +use gateway_core::routing::extensions::ExtensionPreparationPort; +use gateway_plugin_sdk::Capability; +use gateway_plugin_sdk::Contributions; +use gateway_plugin_sdk::Stage; + +#[tokio::test] +async fn model_catalog_is_frozen_without_entering_the_request_chain() { + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::ModelCatalog, + vec![Stage::Registration], + vec![], + vec![], + ), + ])) + .await; + store.snapshot.lock().unwrap().instances[0].configuration = serde_json::json!({ + "model_catalog": {"models": [{"id": "my-model", "provider": "openai", "model": "upstream-model"}]} + }); + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert_eq!(generation.model_aliases().len(), 1); + assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); + assert_eq!( + generation.model_aliases()[0].target.as_str(), + "upstream-model" + ); + assert!(runtime.policy_registry().resolve(&generation).is_none()); + assert!( + runtime + .execution_registry() + .middleware(&generation) + .is_none() + ); + let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert_eq!(reused.id(), generation.id()); + assert_eq!(reused.model_aliases(), generation.model_aliases()); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[0].revision = Revision::new(2).unwrap(); + snapshot.instances[0].configuration = serde_json::json!({"startup":"fail"}); + } + assert!( + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .is_err(), + "failed recovery must not publish a catalog with silently removed aliases" + ); + assert!(generation.is_ready()); + assert_eq!(generation.model_aliases()[0].id.as_str(), "my-model"); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[0].enabled = false; + } + let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .unwrap(); + assert!(disabled.model_aliases().is_empty()); + assert_eq!(generation.model_aliases().len(), 1); + drop(disabled); + drop(reused); + drop(generation); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn model_catalog_rejects_invalid_registrations_and_cross_instance_conflicts() { + for models in [ + serde_json::json!([{"id":"alias","provider":"custom","model":"target"}]), + serde_json::json!([{"id":"alias","provider":"openai","model":"alias"}]), + serde_json::json!([ + {"id":"alias","provider":"openai","model":"target"}, + {"id":"alias","provider":"xai","model":"target"} + ]), + serde_json::json!([{"id":"alias","provider":"openai","model":"target","capabilities":[]}]), + ] { + let (cache, store, runtime) = + super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::ModelCatalog, + vec![Stage::Registration], + vec![], + vec![], + ), + ])) + .await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = + serde_json::json!({"model_catalog":{"models":models}}); + assert!( + PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err() + ); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; + } + + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::ModelCatalog, + vec![Stage::Registration], + vec![], + vec![], + ), + ])) + .await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "model_catalog":{"models":[{"id":"alias","provider":"openai","model":"target"}]} + }); + let mut other = snapshot.instances[0].clone(); + other.id = "instance-two".into(); + snapshot.instances.push(other); + let error = PluginPreparation::prepare(&runtime, snapshot) + .await + .unwrap_err(); + assert!(error.to_string().contains("instance-one")); + assert!(error.to_string().contains("instance-two")); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn failed_candidate_leaves_the_previous_generation_ready() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = store.snapshot.lock().unwrap().clone(); + let active = PluginPreparation::prepare(&runtime, snapshot.clone()) + .await + .unwrap(); + let mut candidate = snapshot; + candidate.instances[0].configuration = serde_json::json!({"startup":"fail"}); + assert!( + PluginPreparation::prepare(&runtime, candidate) + .await + .is_err() + ); + assert!(active.is_ready()); + drop(active); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn disabled_missing_artifacts_do_not_block_the_remaining_generation_or_offline_recovery() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + let mut missing = snapshot.instances[0].clone(); + missing.id = "missing".into(); + missing.enabled = false; + missing.artifact_sha256 = "0".repeat(64); + snapshot.instances.push(missing); + let active = PluginPreparation::prepare(&runtime, snapshot) + .await + .unwrap(); + assert!(active.is_ready()); + drop(active); + super::super::wait_until_empty(cache.path()).await; + let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert!(restored.is_ready()); + drop(restored); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn preparation_cannot_mix_a_newer_persisted_revision_with_an_older_snapshot() { + let (_cache, _store, runtime) = super::super::setup().await; + assert!( + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .is_err() + ); +} + +#[tokio::test] +async fn restoring_a_failed_plugin_keeps_other_instances_ready_and_reports_only_its_failure() { + let (cache, store, runtime) = super::super::setup().await; + let snapshot = { + let mut snapshot = store.snapshot.lock().unwrap(); + let mut failed = snapshot.instances[0].clone(); + failed.id = "failed-plugin".into(); + failed.configuration = serde_json::json!({"startup":"fail"}); + snapshot.instances.insert(0, failed); + snapshot.clone() + }; + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .expect("one failed plugin must not stop restoration"); + assert!(generation.can_serve()); + let diagnostics = gateway_admin::ports::plugins::PluginRuntimeDiagnostics::runtime_diagnostics( + &runtime, + &snapshot, + Some(1), + Some(&generation), + ) + .await + .unwrap(); + assert_eq!( + diagnostics["failed-plugin"].status, + PluginInstanceRuntimeStatus::PreparationFailed + ); + assert!(diagnostics["failed-plugin"].failure.is_some()); + assert_eq!( + diagnostics["instance-one"].status, + PluginInstanceRuntimeStatus::Running + ); + let reused = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + assert_eq!( + generation.id(), + reused.id(), + "quarantined failures do not restart on every request" + ); + drop(reused); + drop(generation); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn an_existing_failed_plugin_does_not_prevent_editing_an_unrelated_instance() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + let mut failed = snapshot.instances[0].clone(); + failed.id = "failed-plugin".into(); + failed.configuration = serde_json::json!({"startup":"fail"}); + snapshot.instances.push(failed); + snapshot.config_revision = Revision::new(2).unwrap(); + snapshot.instances[0].revision = snapshot.config_revision; + let candidate = PluginPreparation::prepare(&runtime, snapshot) + .await + .expect("only the instance being changed must successfully prepare"); + assert!(candidate.can_serve()); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs b/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs index ccf90a50b..a5421e9e1 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/private_state.rs @@ -38,8 +38,8 @@ use gateway_core::{ }, identity::ProviderKind, operation::OperationKind, + routing::extensions::ExtensionPreparationPort, routing::{ConfigRevision, PublicModelId}, - runtime::extensions::ExtensionPreparationPort, upstream::UpstreamSendState, }; use gateway_plugin_runtime::{ diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs index d287a56aa..7589dd7a8 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs @@ -27,8 +27,8 @@ use gateway_core::{ metering::{CalculatedCost, Usage}, operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState}, policy::ClientApiKeyId, + routing::extensions::{ExtensionPreparationPort, ExtensionSetReference}, routing::{ConfigRevision, UpstreamModelId}, - runtime::extensions::{ExtensionPreparationPort, ExtensionSetReference}, upstream::{UpstreamSendState, UpstreamTransport}, }; use gateway_plugin_sdk::{Capability, Contributions, Stage}; @@ -136,11 +136,21 @@ async fn setup( (cache, store, runtime) } -fn context(generation: &ExtensionSetReference, key: &str) -> AttemptContext { - context_with_fast_mode(generation, key, gateway_core::account::FastMode::Default) +fn context( + runtime: &gateway_plugin_runtime::PluginRuntime, + generation: &ExtensionSetReference, + key: &str, +) -> AttemptContext { + context_with_fast_mode( + runtime, + generation, + key, + gateway_core::account::FastMode::Default, + ) } fn context_with_fast_mode( + runtime: &gateway_plugin_runtime::PluginRuntime, generation: &ExtensionSetReference, key: &str, mode: gateway_core::account::FastMode, @@ -150,7 +160,7 @@ fn context_with_fast_mode( ModelRequestId::new("req_one").unwrap(), ClientApiKeyId::new(key).unwrap(), ) - .with_upstream_adapters(generation.upstream_adapters()) + .with_upstream_adapters(runtime.execution_registry().upstream_adapters(generation)) .with_fast_mode(mode), NonZeroU32::new(1).unwrap(), SystemTime::now() + Duration::from_secs(5), @@ -166,12 +176,13 @@ fn context_with_fast_mode( } fn execute( + runtime: &gateway_plugin_runtime::PluginRuntime, generation: &ExtensionSetReference, account: Arc, key: &str, previous: Option, ) -> EventStream { - let context = context(generation, key); + let context = context(runtime, generation, key); execute_with_context(account, context, previous) } @@ -235,7 +246,7 @@ async fn legacy_adapter_decodes_exact_v1_metadata_for_all_fast_modes() { .unwrap(); let mut stream = execute_with_context( Account::new(), - context_with_fast_mode(&generation, "key-one", mode), + context_with_fast_mode(&runtime, &generation, "key-one", mode), None, ); let mut completed = 0; @@ -277,7 +288,7 @@ async fn adapter_http_and_sse_are_cold_authenticated_and_settled_once() { .await .unwrap(); let account = Account::new(); - let mut stream = execute(&generation, Arc::clone(&account), "key-one", None); + let mut stream = execute(&runtime, &generation, Arc::clone(&account), "key-one", None); assert!( server.received_requests().await.unwrap().is_empty(), "构造冷流不能先发送业务请求" @@ -321,7 +332,7 @@ async fn adapter_accepts_service_tier_resolved_by_the_terminal_response() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut tiers = Vec::new(); let mut completed = false; while let Some(event) = stream.next().await { @@ -368,7 +379,7 @@ async fn adapter_failure_preserves_upstream_diagnostics_and_native_feedback() { .await .unwrap(); let account = Account::new(); - let mut stream = execute(&generation, Arc::clone(&account), "key-one", None); + let mut stream = execute(&runtime, &generation, Arc::clone(&account), "key-one", None); let error = stream.next().await.unwrap().unwrap_err(); assert_eq!(error.kind(), ProviderErrorKind::RateLimited); assert_eq!(error.send_state(), UpstreamSendState::Sent); @@ -409,7 +420,7 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut state = None; while let Some(event) = stream.next().await { state = event.unwrap().take_session_update().or(state); @@ -420,7 +431,7 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() { ("key-one", Account::with_identity("acct_other", 1)), ("key-one", Account::with_identity("acct_one", 2)), ] { - let error = execute(&generation, account, key, state.clone()) + let error = execute(&runtime, &generation, account, key, state.clone()) .next() .await .unwrap() @@ -438,7 +449,7 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); while let Some(event) = stream.next().await { event.unwrap(); } @@ -467,7 +478,7 @@ async fn adapter_does_not_publish_completed_when_rpc_fails_after_terminal_frame( let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); stream.next().await.unwrap().unwrap(); let error = stream.next().await.unwrap().err().unwrap(); assert_eq!(error.kind(), ProviderErrorKind::Unavailable); @@ -576,7 +587,7 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut first = execute(&generation, Account::new(), "key-one", None); + let mut first = execute(&runtime, &generation, Account::new(), "key-one", None); let mut state = None; while let Some(event) = first.next().await { state = event.unwrap().take_session_update().or(state); @@ -584,24 +595,42 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle drop(first); let state = state.unwrap(); assert!(state.extension_owner().unwrap().connection_local); - let error = execute(&generation, Account::new(), "key-two", Some(state.clone())) - .next() - .await - .unwrap() - .err() - .unwrap(); + let error = execute( + &runtime, + &generation, + Account::new(), + "key-two", + Some(state.clone()), + ) + .next() + .await + .unwrap() + .err() + .unwrap(); assert_eq!(error.send_state(), UpstreamSendState::NotSent); - let mut second = execute(&generation, Account::new(), "key-one", Some(state.clone())); + let mut second = execute( + &runtime, + &generation, + Account::new(), + "key-one", + Some(state.clone()), + ); while let Some(event) = second.next().await { event.unwrap(); } drop(second); - let error = execute(&generation, Account::new(), "key-one", Some(state)) - .next() - .await - .unwrap() - .err() - .unwrap(); + let error = execute( + &runtime, + &generation, + Account::new(), + "key-one", + Some(state), + ) + .next() + .await + .unwrap() + .err() + .unwrap(); assert_eq!(error.send_state(), UpstreamSendState::NotSent); drop(generation); runtime.shutdown().await; @@ -642,7 +671,7 @@ async fn websocket_adapter_can_send_while_a_read_callback_is_pending() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut continuation = None; tokio::time::timeout(Duration::from_secs(3), async { while let Some(event) = stream.next().await { @@ -696,7 +725,7 @@ async fn cancelling_adapter_body_closes_pending_websocket_and_releases_selected_ .await .unwrap(); let account = Account::new(); - let mut stream = execute(&generation, account.clone(), "key-one", None); + let mut stream = execute(&runtime, &generation, account.clone(), "key-one", None); let mut reading = Box::pin(stream.next()); tokio::select! { result = &mut reading => panic!("上游等待期间不应提前产生结果：{}", result.is_some()), @@ -734,7 +763,7 @@ async fn usage_without_response_model_still_uses_native_account_pricing() { let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let mut costs = 0; while let Some(event) = stream.next().await { costs += event @@ -834,7 +863,7 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); let account = Account::new(); - let mut old_stream = execute(&old, account.clone(), "key-one", None); + let mut old_stream = execute(&runtime, &old, account.clone(), "key-one", None); old_stream.next().await.unwrap().unwrap(); { let mut snapshot = store.snapshot.lock().unwrap(); @@ -846,7 +875,7 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); assert_ne!(old.id(), new.id()); - let mut new_stream = execute(&new, account.clone(), "key-one", None); + let mut new_stream = execute(&runtime, &new, account.clone(), "key-one", None); new_stream.next().await.unwrap().unwrap(); drop(old); let old_state = old_stream @@ -868,7 +897,7 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); assert!( - context(&disabled, "key-one") + context(&runtime, &disabled, "key-one") .upstream_adapter( &ProviderKind::new("openai").unwrap(), &UpstreamModelId::new("native-model").unwrap() @@ -903,13 +932,19 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a .await .unwrap(); // 恢复同一配置仍是新代次，不能复活旧进程中的续接身份 - let error = execute(&restored, account.clone(), "key-one", Some(old_state)) - .next() - .await - .unwrap() - .unwrap_err(); + let error = execute( + &runtime, + &restored, + account.clone(), + "key-one", + Some(old_state), + ) + .next() + .await + .unwrap() + .unwrap_err(); assert_eq!(error.send_state(), UpstreamSendState::NotSent); - let mut fresh = execute(&restored, account.clone(), "key-one", None); + let mut fresh = execute(&runtime, &restored, account.clone(), "key-one", None); while let Some(event) = fresh.next().await { event.unwrap(); } @@ -923,15 +958,15 @@ async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_a #[tokio::test] async fn upstream_adapter_can_dispatch_http_after_its_instance_entered_the_scope() { struct Http(Arc); - impl gateway_core::middleware::http::Dispatcher for Http { + impl gateway_core::engine::middleware::http::Dispatcher for Http { fn dispatch( &self, - context: gateway_core::middleware::http::Context, - request: gateway_core::middleware::http::Request, + context: gateway_core::engine::middleware::http::Context, + request: gateway_core::engine::middleware::http::Request, ) -> BoxFuture< 'static, Result< - gateway_core::middleware::http::Response, + gateway_core::engine::middleware::http::Response, gateway_core::engine::middleware::MiddlewareError, >, > { @@ -944,8 +979,8 @@ async fn upstream_adapter_can_dispatch_http_after_its_instance_entered_the_scope assert_eq!(request.uri(), "/api/admin/settings"); self.0.fetch_add(1, Ordering::SeqCst); Box::pin(async { - Ok(gateway_core::middleware::http::Response::new( - gateway_core::middleware::http::empty_body(), + Ok(gateway_core::engine::middleware::http::Response::new( + gateway_core::engine::middleware::http::empty_body(), )) }) } @@ -961,13 +996,13 @@ async fn upstream_adapter_can_dispatch_http_after_its_instance_entered_the_scope })); let (_cache, _, runtime) = setup(config).await; let calls = Arc::new(AtomicUsize::new(0)); - let dispatcher: Arc = + let dispatcher: Arc = Arc::new(Http(calls.clone())); runtime.bind_http(&dispatcher).unwrap(); let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); while let Some(event) = stream.next().await { event.unwrap(); } @@ -1004,7 +1039,7 @@ async fn malformed_adapter_events_stop_only_the_plugin_session() { ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await .unwrap(); - let mut stream = execute(&generation, Account::new(), "key-one", None); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); let error = stream.next().await.unwrap().unwrap_err(); assert_eq!(error.kind(), ProviderErrorKind::Protocol); if has_sequence_error { diff --git a/backend/crates/gateway-plugin/runtime/tests/support/environment.rs b/backend/crates/gateway-plugin/runtime/tests/support/environment.rs index 5b6ce314b..c61b1b789 100644 --- a/backend/crates/gateway-plugin/runtime/tests/support/environment.rs +++ b/backend/crates/gateway-plugin/runtime/tests/support/environment.rs @@ -439,7 +439,7 @@ impl Environment { Some(runtime.clone()), Some(runtime.observer_registry()), Some(runtime.policy_registry()), - Some(runtime.middleware_registry()), + Some(runtime.execution_registry()), Some(runtime.frontend_authentication_registry()), ); let access = gateway_admin::initialize_plugin_accounts( @@ -472,7 +472,7 @@ impl Environment { Some(runtime.clone()), Some(runtime.observer_registry()), Some(runtime.policy_registry()), - Some(runtime.middleware_registry()), + Some(runtime.execution_registry()), ) .activate() .await @@ -536,10 +536,10 @@ impl Environment { timezone: Default::default(), service_middleware: { let snapshots = core.snapshots(); - let middleware = runtime.middleware_registry(); + let middleware = runtime.execution_registry(); std::sync::Arc::new(move || { let snapshot = snapshots.snapshot_for_diagnostics()?; - middleware.resolve(snapshot.extensions()?) + middleware.middleware(snapshot.extensions()?) }) }, plugin_preparation: runtime.clone(), diff --git a/backend/crates/gateway-plugin/sdk/README.md b/backend/crates/gateway-plugin/sdk/README.md index 485af3fcb..160964c82 100644 --- a/backend/crates/gateway-plugin/sdk/README.md +++ b/backend/crates/gateway-plugin/sdk/README.md @@ -101,7 +101,7 @@ sequenceDiagram ## 本地验证 -设置服务的 SDK 类型和操作由宿主领域类型及[操作声明](../../gateway-admin/src/service/settings.rs)生成，SDK 构建不依赖宿主源码。 +设置服务的 SDK 类型和操作由宿主领域类型及[操作声明](../../gateway-admin/src/public_service/settings.rs)生成，SDK 构建不依赖宿主源码。 修改宿主合同后，在仓库根目录更新生成文件；Admin 合同测试会检查生成结果是否同步： ```bash diff --git a/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs b/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs index 42f123220..bab9cdc92 100644 --- a/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs +++ b/backend/crates/gateway-plugin/sdk/src/call/services/settings.rs @@ -1,6 +1,6 @@ //! 宿主设置服务的操作标识与请求、响应数据合同 //! -//! 从宿主设置类型与 service/settings.rs 生成；更新命令见 SDK 维护说明 +//! 从宿主设置类型与 public_service/settings.rs 生成；更新命令见 SDK 维护说明 use super::Operation; use serde::{Deserialize, Serialize}; use std::{ diff --git a/backend/crates/gateway-plugin/sdk/src/client/session.rs b/backend/crates/gateway-plugin/sdk/src/client/session.rs deleted file mode 100644 index 6f7e22b44..000000000 --- a/backend/crates/gateway-plugin/sdk/src/client/session.rs +++ /dev/null @@ -1,1479 +0,0 @@ -//! 插件侧双向 RPC 会话的握手、调用分派、流控与取消生命周期 - -use std::{ - collections::{BTreeMap, BTreeSet, VecDeque}, - future::Future, - num::NonZeroUsize, - pin::Pin, - sync::{Arc, Mutex}, - time::Duration, -}; - -use serde_json::Value; -use tokio::{ - io::{AsyncRead, AsyncWrite}, - sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch}, - task::JoinHandle, - time::Instant, -}; - -use crate::{ - CallContext, ErrorCode, Frame, FrameError, Handshake, Message, PROTOCOL_VERSION, PluginFault, -}; - -use super::frame::{read_frame, validate_frame, write_frame}; - -const MAXIMUM_STREAM_CHUNK_BYTES: usize = 16 * 1024 * 1024; -const MAXIMUM_CONCURRENCY: usize = 256; -const MAXIMUM_BUFFERED_STREAM_CHUNKS: usize = 65_536; - -/// 插件侧会话的内存、并发与期限边界 -#[derive(Debug, Clone, Copy)] -pub struct SessionConfig { - /// 单个业务流分块的预算，不限制普通调用正文的总长度 - pub maximum_stream_chunk_bytes: usize, - pub maximum_calls: usize, - pub maximum_callbacks: usize, - pub maximum_buffered_stream_chunks: usize, - pub handshake_timeout: Duration, - pub maximum_call_timeout: Duration, -} - -impl Default for SessionConfig { - fn default() -> Self { - Self { - // 流分块还需匹配宿主授予的信用窗口 - maximum_stream_chunk_bytes: 1024 * 1024, - maximum_calls: 32, - maximum_callbacks: 32, - maximum_buffered_stream_chunks: 1_024, - handshake_timeout: Duration::from_secs(5), - maximum_call_timeout: Duration::from_secs(120), - } - } -} - -impl SessionConfig { - fn validate(self) -> Result { - if self.maximum_stream_chunk_bytes < 1_024 - || self.maximum_stream_chunk_bytes > MAXIMUM_STREAM_CHUNK_BYTES - || self.maximum_calls == 0 - || self.maximum_calls > MAXIMUM_CONCURRENCY - || self.maximum_callbacks == 0 - || self.maximum_callbacks > MAXIMUM_CONCURRENCY - || self.maximum_buffered_stream_chunks == 0 - || self.maximum_buffered_stream_chunks > MAXIMUM_BUFFERED_STREAM_CHUNKS - || self.handshake_timeout.is_zero() - || self.maximum_call_timeout.is_zero() - { - return Err(SessionError::Configuration); - } - Ok(self) - } -} - -/// 插件侧传输和会话错误 -#[derive(Debug, thiserror::Error)] -pub enum SessionError { - #[error("plugin session configuration is invalid")] - Configuration, - #[error("plugin handshake is invalid")] - Handshake, - #[error("plugin protocol is invalid")] - Protocol, - #[error("plugin transport is closed")] - Closed, - #[error("plugin operation exceeded its deadline")] - Timeout, - #[error("plugin call was cancelled")] - Cancelled, - #[error("plugin session capacity is exhausted")] - Capacity, - #[error("plugin handler stopped unexpectedly")] - HandlerStopped, - #[error("host callback returned an error: {0:?}")] - Remote(PluginFault), - #[error(transparent)] - Frame(#[from] FrameError), -} - -impl SessionError { - /// 将本地会话失败收敛为可返回宿主的稳定插件错误 - #[must_use] - pub fn into_plugin_fault(self) -> PluginFault { - match self { - Self::Remote(fault) => fault, - Self::Timeout => PluginFault::new(ErrorCode::Timeout, "host callback timed out"), - Self::Cancelled => PluginFault::new(ErrorCode::Cancelled, "parent call was cancelled"), - Self::Capacity => PluginFault::new(ErrorCode::Capacity, "plugin capacity is exhausted"), - Self::Configuration | Self::Protocol | Self::Handshake | Self::Frame(_) => { - PluginFault::new(ErrorCode::InvalidInput, "plugin session input is invalid") - } - Self::Closed | Self::HandlerStopped => { - PluginFault::new(ErrorCode::Fault, "plugin session is unavailable") - } - } - } -} - -/// 一次宿主调用的可克隆取消信号 -#[derive(Clone)] -pub struct CallCancellation { - receiver: watch::Receiver, -} - -impl CallCancellation { - #[must_use] - pub fn is_cancelled(&self) -> bool { - *self.receiver.borrow() - } - - pub async fn cancelled(&self) { - let mut receiver = self.receiver.clone(); - if *receiver.borrow() { - return; - } - while receiver.changed().await.is_ok() { - if *receiver.borrow() { - return; - } - } - } -} - -struct CancellationSource { - sender: watch::Sender, -} - -impl CancellationSource { - fn new() -> (Self, CallCancellation) { - let (sender, receiver) = watch::channel(false); - (Self { sender }, CallCancellation { receiver }) - } - - fn cancel(&self) { - self.sender.send_replace(true); - } -} - -/// 宿主回调的成功结果；不实现 `Debug`，避免载荷进入普通诊断 -pub struct HostReply { - pub result: Value, - pub payload: Vec, -} - -/// 绑定父调用、期限和取消信号的宿主回调客户端 -#[derive(Clone)] -pub struct HostClient { - parent_id: u64, - deadline: tokio::sync::watch::Receiver>, - cancellation: CallCancellation, - callbacks: Arc, - output: Output, - maximum_stream_chunk_bytes: usize, - credits: Arc, -} - -impl HostClient { - pub(crate) fn maximum_stream_chunk_bytes(&self) -> usize { - let window = lock_unpoisoned(&self.credits.state).maximum_bytes; - if window == 0 { - self.maximum_stream_chunk_bytes - } else { - self.maximum_stream_chunk_bytes - .min(usize::try_from(window).unwrap_or(usize::MAX)) - } - } - - /// 发起一次与父调用关联的宿主回调 - /// - /// # Errors - /// - /// 方法或帧无效、容量耗尽、父调用取消、期限到达、连接关闭或宿主返回错误时失败 - pub async fn call( - &self, - method: impl Into, - params: Value, - payload: Vec, - ) -> Result { - let method = method.into(); - if method.is_empty() || method.len() > 128 { - return Err(SessionError::Protocol); - } - let ordered = self.callbacks.order.lock().await; - let (id, received) = self.callbacks.begin(self.parent_id)?; - let frame = Frame { - message: Message::Callback { - id, - parent_id: self.parent_id, - method, - params, - }, - payload, - }; - if validate_frame(&frame).is_err() { - self.callbacks.retire(id); - return Err(SessionError::Protocol); - } - let sent = tokio::select! { - biased; - () = self.cancellation.cancelled() => Err(SessionError::Cancelled), - () = callback_deadline(self.deadline.clone()) => Err(SessionError::Timeout), - result = self.output.send_data(frame) => result, - }; - if let Err(error) = sent { - self.callbacks.retire(id); - return Err(error); - } - drop(ordered); - tokio::select! { - biased; - () = self.cancellation.cancelled() => { - self.callbacks.retire(id); - Err(SessionError::Cancelled) - } - () = callback_deadline(self.deadline.clone()) => { - self.callbacks.retire(id); - Err(SessionError::Timeout) - } - result = received => result.map_err(|_| SessionError::Closed)?, - } - } -} - -/// 交给插件业务处理器的一次宿主调用；不实现 `Debug`，避免敏感载荷泄露 -pub struct PluginCall { - pub method: String, - pub context: CallContext, - pub params: Value, - pub payload: Vec, - pub host: HostClient, - pub cancellation: CallCancellation, -} - -/// 插件业务调用 future -pub type CallFuture<'a> = Pin> + Send + 'a>>; - -/// 插件只实现业务调用；会话负责握手、关联、流控、取消和关闭 -pub trait PluginHandler: Send + Sync + 'static { - fn call(&self, call: PluginCall) -> CallFuture<'_>; - - /// 取消时释放只属于该调用的临时业务状态 - /// 实现必须幂等、快速且不得执行阻塞 I/O - fn cancel(&self, _context: &CallContext) {} - - /// 会话停止接收新调用时触发 - /// 实现必须幂等、快速且不得执行阻塞 I/O - fn quiesce(&self) {} - - /// 会话关闭时释放插件自有临时状态 - /// 实现必须幂等、快速且不得执行阻塞 I/O - fn shutdown(&self) {} -} - -/// 一次业务成功结果；`stream` 存在时 SDK 自动产生单一流终态 -pub struct CallReply { - result: Value, - payload: Vec, - stream: Option, -} - -impl CallReply { - #[must_use] - pub fn unary(result: Value, payload: Vec) -> Self { - Self { - result, - payload, - stream: None, - } - } - - #[must_use] - pub fn stream(result: Value, payload: Vec, stream: ResponseStream) -> Self { - Self { - result, - payload, - stream: Some(stream), - } - } -} - -enum StreamSource { - Buffered(VecDeque>), - Channel(mpsc::Receiver, PluginFault>>), - Pull(Box), -} - -pub type PullResponseFuture<'a> = - Pin, PluginFault>>> + Send + 'a>>; - -pub trait PullResponseStream: Send { - fn next(&mut self) -> PullResponseFuture<'_>; -} - -/// SDK 管理 sequence、Credit 和终态的响应流 -pub struct ResponseStream { - source: StreamSource, - declared_capacity: usize, -} - -impl ResponseStream { - /// 从已完成业务校验的有限分块构造流；SDK 会在发送 `Result` 前预检全部分块 - #[must_use] - pub fn from_chunks(chunks: Vec>) -> Self { - let declared_capacity = chunks.len().max(1); - Self { - source: StreamSource::Buffered(chunks.into()), - declared_capacity, - } - } - - /// 创建由生产者驱动的有界流队列 - #[must_use] - pub fn channel(capacity: NonZeroUsize) -> (StreamSender, Self) { - let (sender, receiver) = mpsc::channel(capacity.get()); - ( - StreamSender { sender }, - Self { - source: StreamSource::Channel(receiver), - declared_capacity: capacity.get(), - }, - ) - } - - /// 按消费进度拉取下一帧；返回 None 表示完成，丢弃流会丢弃生产者 - pub fn pull(source: Box) -> Self { - Self { - source: StreamSource::Pull(source), - declared_capacity: 1, - } - } - - fn validate_buffered( - &self, - maximum_chunks: usize, - maximum_stream_chunk_bytes: usize, - window_bytes: u64, - ) -> Result<(), PluginFault> { - if self.declared_capacity > maximum_chunks { - return Err(capacity_fault( - "response stream exceeds the local queue limit", - )); - } - if let StreamSource::Buffered(chunks) = &self.source { - for chunk in chunks { - validate_stream_chunk(chunk, maximum_stream_chunk_bytes, window_bytes)?; - } - } - Ok(()) - } - - pub(crate) async fn next(&mut self) -> Option, PluginFault>> { - match &mut self.source { - StreamSource::Buffered(chunks) => chunks.pop_front().map(Ok), - StreamSource::Channel(receiver) => receiver.recv().await, - StreamSource::Pull(source) => source.next().await, - } - } -} - -/// 动态响应流的有界生产端；最后一个 sender 被释放表示成功终态 -#[derive(Clone)] -pub struct StreamSender { - sender: mpsc::Sender, PluginFault>>, -} - -impl StreamSender { - /// 排队一个业务分块；实际线协议流控由会话完成 - pub async fn send(&self, payload: Vec) -> Result<(), SessionError> { - self.sender - .send(Ok(payload)) - .await - .map_err(|_| SessionError::Closed) - } - - /// 以业务错误结束动态流 - pub async fn fail(&self, fault: PluginFault) -> Result<(), SessionError> { - self.sender - .send(Err(fault)) - .await - .map_err(|_| SessionError::Closed) - } -} - -/// 已读取并校验 Hello、尚未向宿主发送 Ready 的插件会话 -pub struct PluginSession { - reader: R, - writer: W, - handshake: Handshake, - config: SessionConfig, -} - -impl PluginSession -where - R: AsyncRead + Unpin + Send + 'static, - W: AsyncWrite + Unpin + Send + 'static, -{ - /// 读取并校验宿主握手 - /// 业务处理器可以在 `run` 前读取不可变 Handshake 构造注册描述 - /// - /// # Errors - /// - /// 配置无效、握手超时、传输关闭或 Hello 不符合协议时失败 - pub async fn accept( - mut reader: R, - writer: W, - config: SessionConfig, - ) -> Result { - let config = config.validate()?; - let frame = tokio::time::timeout(config.handshake_timeout, read_frame(&mut reader)) - .await - .map_err(|_| SessionError::Timeout)??; - let Message::Hello { handshake } = frame.message else { - return Err(SessionError::Handshake); - }; - if !frame.payload.is_empty() - || handshake.protocol_version != PROTOCOL_VERSION - || handshake.plugin_id.is_empty() - || handshake.instance_id.is_empty() - || handshake.generation == 0 - || handshake.incarnation.is_empty() - { - return Err(SessionError::Handshake); - } - Ok(Self { - reader, - writer, - handshake, - config, - }) - } - - #[must_use] - pub const fn handshake(&self) -> &Handshake { - &self.handshake - } - - /// 发送 Ready 并持续处理调用、回调、流控、取消与关闭 - /// - /// # Errors - /// - /// 帧损坏、协议违规、处理器异常停止或传输关闭时失败 - pub async fn run(mut self, handler: H) -> Result<(), SessionError> - where - H: PluginHandler, - { - let data_capacity = self - .config - .maximum_calls - .checked_add(self.config.maximum_callbacks) - .ok_or(SessionError::Configuration)?; - let control_reserve = self - .config - .maximum_calls - .checked_add(8) - .ok_or(SessionError::Configuration)?; - let queue_capacity = data_capacity - .checked_add(control_reserve) - .ok_or(SessionError::Configuration)?; - let (outbound, received) = mpsc::channel(queue_capacity); - let output = Output { - sender: outbound, - data_slots: Arc::new(Semaphore::new(data_capacity)), - }; - let callbacks = Arc::new(CallbackRegistry::new(self.config.maximum_callbacks)); - let handler: Arc = Arc::new(handler); - let lifecycle = LifecycleHooks::new(Arc::clone(&handler), control_reserve)?; - write_frame( - &mut self.writer, - &Frame::control(Message::Ready { - protocol_version: PROTOCOL_VERSION, - incarnation: self.handshake.incarnation.clone(), - }), - ) - .await?; - let writer = tokio::spawn(writer_loop(self.writer, received)); - let (completed, completions) = mpsc::channel(self.config.maximum_calls); - let mut driver = SessionDriver { - config: self.config, - instance_id: self.handshake.instance_id, - generation: self.handshake.generation, - incarnation: self.handshake.incarnation, - output, - callbacks, - handler, - lifecycle, - writer, - completed, - completions, - active: BTreeMap::new(), - last_call: 0, - quiescing: false, - }; - let result = driver.drive(self.reader).await; - driver.close().await; - result - } -} - -struct Output { - sender: mpsc::Sender, - data_slots: Arc, -} - -impl Clone for Output { - fn clone(&self) -> Self { - Self { - sender: self.sender.clone(), - data_slots: Arc::clone(&self.data_slots), - } - } -} - -impl Output { - async fn send_data(&self, frame: Frame) -> Result<(), SessionError> { - validate_frame(&frame)?; - let permit = Arc::clone(&self.data_slots) - .acquire_owned() - .await - .map_err(|_| SessionError::Closed)?; - self.sender - .send(Outbound { - frame, - _data_slot: Some(permit), - }) - .await - .map_err(|_| SessionError::Closed) - } - - async fn send_data_until( - &self, - frame: Frame, - deadline: Option, - cancellation: &CallCancellation, - ) -> Result<(), SessionError> { - tokio::select! { - biased; - () = cancellation.cancelled() => Err(SessionError::Cancelled), - () = wait_deadline(deadline) => Err(SessionError::Timeout), - result = self.send_data(frame) => result, - } - } - - fn send_control(&self, frame: Frame) -> Result<(), SessionError> { - validate_frame(&frame)?; - self.sender - .try_send(Outbound { - frame, - _data_slot: None, - }) - .map_err(|_| SessionError::Capacity) - } -} - -struct Outbound { - frame: Frame, - _data_slot: Option, -} - -async fn writer_loop( - mut writer: W, - mut received: mpsc::Receiver, -) -> Result<(), SessionError> { - while let Some(outbound) = received.recv().await { - write_frame(&mut writer, &outbound.frame).await?; - } - Ok(()) -} - -struct CallbackRegistry { - state: Mutex, - order: tokio::sync::Mutex<()>, - maximum: usize, - maximum_retired: usize, -} - -struct CallbackState { - next_id: u64, - pending: BTreeMap, - retired: BTreeSet, - retired_order: VecDeque, -} - -struct PendingCallback { - parent_id: u64, - result: oneshot::Sender>, -} - -impl CallbackRegistry { - fn new(maximum: usize) -> Self { - Self { - state: Mutex::new(CallbackState { - next_id: 2, - pending: BTreeMap::new(), - retired: BTreeSet::new(), - retired_order: VecDeque::new(), - }), - order: tokio::sync::Mutex::new(()), - maximum, - maximum_retired: maximum.saturating_mul(2), - } - } - - fn begin( - &self, - parent_id: u64, - ) -> Result<(u64, oneshot::Receiver>), SessionError> { - let mut state = lock_unpoisoned(&self.state); - if state.pending.len() >= self.maximum { - return Err(SessionError::Capacity); - } - let id = state.next_id; - state.next_id = state.next_id.checked_add(2).ok_or(SessionError::Capacity)?; - let (result, received) = oneshot::channel(); - state - .pending - .insert(id, PendingCallback { parent_id, result }); - Ok((id, received)) - } - - fn resolve(&self, frame: Frame) -> Result<(), SessionError> { - let id = match &frame.message { - Message::Result { id, .. } | Message::Error { id, .. } => *id, - _ => return Err(SessionError::Protocol), - }; - let pending = { - let mut state = lock_unpoisoned(&self.state); - if let Some(pending) = state.pending.remove(&id) { - Some(pending) - } else if state.retired.remove(&id) { - state.retired_order.retain(|retired| *retired != id); - None - } else { - return Err(SessionError::Protocol); - } - }; - let Some(pending) = pending else { - return Ok(()); - }; - let result = match frame.message { - Message::Result { result, .. } => Ok(HostReply { - result, - payload: frame.payload, - }), - Message::Error { error, .. } if frame.payload.is_empty() => { - Err(SessionError::Remote(error)) - } - _ => return Err(SessionError::Protocol), - }; - let _ = pending.result.send(result); - Ok(()) - } - - fn retire(&self, id: u64) { - let mut state = lock_unpoisoned(&self.state); - if state.pending.remove(&id).is_some() { - retire_callback(&mut state, id, self.maximum_retired); - } - } - - fn finish_parent(&self, parent_id: u64) { - let pending = { - let mut state = lock_unpoisoned(&self.state); - let mut remaining = BTreeMap::new(); - let mut finished = Vec::new(); - for (id, pending) in std::mem::take(&mut state.pending) { - if pending.parent_id == parent_id { - retire_callback(&mut state, id, self.maximum_retired); - finished.push(pending); - } else { - remaining.insert(id, pending); - } - } - state.pending = remaining; - finished - }; - for pending in pending { - let _ = pending.result.send(Err(SessionError::Cancelled)); - } - } - - fn close(&self) { - let pending = { - let mut state = lock_unpoisoned(&self.state); - std::mem::take(&mut state.pending) - }; - for pending in pending.into_values() { - let _ = pending.result.send(Err(SessionError::Closed)); - } - } -} - -fn retire_callback(state: &mut CallbackState, id: u64, maximum: usize) { - if state.retired.insert(id) { - state.retired_order.push_back(id); - } - while state.retired_order.len() > maximum { - if let Some(expired) = state.retired_order.pop_front() { - state.retired.remove(&expired); - } - } -} - -enum LifecycleEvent { - Cancel(CallContext), - Quiesce, - Shutdown, -} - -#[derive(Clone)] -struct LifecycleHooks { - events: std::sync::mpsc::SyncSender, -} - -impl LifecycleHooks { - fn new(handler: Arc, capacity: usize) -> Result { - let (events, received) = std::sync::mpsc::sync_channel(capacity); - std::thread::Builder::new() - .name("gateway-plugin-sdk-lifecycle".into()) - .spawn(move || { - while let Ok(event) = received.recv() { - match event { - LifecycleEvent::Cancel(context) => handler.cancel(&context), - LifecycleEvent::Quiesce => handler.quiesce(), - LifecycleEvent::Shutdown => { - handler.shutdown(); - break; - } - } - } - }) - .map_err(|_| SessionError::Closed)?; - Ok(Self { events }) - } - - fn cancel(&self, context: CallContext) { - let _ = self.events.try_send(LifecycleEvent::Cancel(context)); - } - - fn quiesce(&self) { - let _ = self.events.try_send(LifecycleEvent::Quiesce); - } - - fn shutdown(&self) { - let _ = self.events.try_send(LifecycleEvent::Shutdown); - } -} - -struct ActiveCall { - context: CallContext, - credits: Arc, - cancellation: CancellationSource, - task: JoinHandle<()>, -} - -#[derive(Clone, Copy)] -enum CallTaskOutcome { - Complete, - TransportFailed, - HandlerStopped, -} - -struct CallFinished { - id: u64, - outcome: CallTaskOutcome, -} - -struct CompletionGuard { - id: u64, - completed: mpsc::Sender, - armed: bool, -} - -impl CompletionGuard { - fn new(id: u64, completed: mpsc::Sender) -> Self { - Self { - id, - completed, - armed: true, - } - } - - fn finish(mut self, outcome: CallTaskOutcome) { - self.armed = false; - let _ = self.completed.try_send(CallFinished { - id: self.id, - outcome, - }); - } -} - -impl Drop for CompletionGuard { - fn drop(&mut self) { - if self.armed { - let _ = self.completed.try_send(CallFinished { - id: self.id, - outcome: CallTaskOutcome::HandlerStopped, - }); - } - } -} - -struct SessionDriver { - config: SessionConfig, - instance_id: String, - generation: u64, - incarnation: String, - output: Output, - callbacks: Arc, - handler: Arc, - lifecycle: LifecycleHooks, - writer: JoinHandle>, - completed: mpsc::Sender, - completions: mpsc::Receiver, - active: BTreeMap, - last_call: u64, - quiescing: bool, -} - -impl SessionDriver { - async fn drive(&mut self, mut reader: R) -> Result<(), SessionError> { - loop { - // read_exact 不是取消安全的；完成通知可以穿插处理，但必须保留同一个半帧读取 future - let incoming = read_frame(&mut reader); - tokio::pin!(incoming); - let frame = loop { - tokio::select! { - biased; - writer = &mut self.writer => { - return writer.map_err(|_| SessionError::Closed)?; - } - Some(finished) = self.completions.recv() => { - self.finish_call(finished)?; - } - frame = &mut incoming => { - break frame.map_err(|error| match error { - FrameError::Io(io_error) - if io_error.kind() == std::io::ErrorKind::UnexpectedEof => - { - SessionError::Closed - } - error => SessionError::Frame(error), - })?; - } - } - }; - if self.handle_frame(frame)? { - return Ok(()); - } - } - } - - fn handle_frame(&mut self, frame: Frame) -> Result { - match frame.message { - Message::Call { - id, - method, - context, - params, - } => { - self.start_call(id, method, context, params, frame.payload)?; - Ok(false) - } - Message::Credit { id, bytes, frames } if frame.payload.is_empty() => { - if id == 0 - || id.is_multiple_of(2) - || id > self.last_call - || bytes == 0 - || frames == 0 - { - return Err(SessionError::Protocol); - } - // End/Cancelled 已进入同一 FIFO 后，宿主先前为已消费分块排队的 - // Credit 仍可能稍晚抵达 - // 此时调用资源已经回收，忽略迟到信用即可； - // 未来或回调 ID 仍按协议错误处理 - if let Some(active) = self.active.get(&id) { - active.credits.grant(bytes, frames)?; - } - Ok(false) - } - Message::Cancel { id } if frame.payload.is_empty() => { - self.cancel_call(id)?; - Ok(false) - } - Message::Result { id, result } => { - self.callbacks.resolve(Frame { - message: Message::Result { id, result }, - payload: frame.payload, - })?; - Ok(false) - } - Message::Error { id, error } => { - self.callbacks.resolve(Frame { - message: Message::Error { id, error }, - payload: frame.payload, - })?; - Ok(false) - } - Message::Quiesce if frame.payload.is_empty() => { - if !self.quiescing { - self.quiescing = true; - self.lifecycle.quiesce(); - } - Ok(false) - } - Message::Shutdown if frame.payload.is_empty() => Ok(true), - _ => Err(SessionError::Protocol), - } - } - - fn start_call( - &mut self, - id: u64, - method: String, - context: CallContext, - params: Value, - payload: Vec, - ) -> Result<(), SessionError> { - if id == 0 - || id.is_multiple_of(2) - || id <= self.last_call - || context.call_id != id - || context.instance_id != self.instance_id - || context.generation != self.generation - || context.incarnation != self.incarnation - || context.timeout_ms == 0 - || Duration::from_millis(context.timeout_ms) > self.config.maximum_call_timeout - || method.is_empty() - || method.len() > 128 - { - return Err(SessionError::Protocol); - } - self.last_call = id; - if self.quiescing { - return self.output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Capacity, "plugin is quiescing"), - )); - } - if self.active.len() >= self.config.maximum_calls { - return self.output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Capacity, "plugin call capacity is exhausted"), - )); - } - let deadline = Instant::now() - .checked_add(Duration::from_millis(context.timeout_ms)) - .ok_or(SessionError::Protocol)?; - let deadline = Some(deadline); - let (deadline_source, callback_deadline) = tokio::sync::watch::channel(deadline); - let credits = Arc::new(CreditWindow::default()); - let (cancellation, cancellation_signal) = CancellationSource::new(); - let host = HostClient { - parent_id: id, - deadline: callback_deadline, - cancellation: cancellation_signal.clone(), - callbacks: Arc::clone(&self.callbacks), - output: self.output.clone(), - maximum_stream_chunk_bytes: self.config.maximum_stream_chunk_bytes, - credits: credits.clone(), - }; - let call = PluginCall { - method, - context: context.clone(), - params, - payload, - host, - cancellation: cancellation_signal, - }; - let handler = Arc::clone(&self.handler); - let output = self.output.clone(); - let completed = self.completed.clone(); - let task_credits = Arc::clone(&credits); - let maximum_stream_chunk_bytes = self.config.maximum_stream_chunk_bytes; - let maximum_chunks = self.config.maximum_buffered_stream_chunks; - let task_context = context.clone(); - let lifecycle = self.lifecycle.clone(); - let task = tokio::spawn(async move { - let guard = CompletionGuard::new(id, completed); - let result = run_call( - handler, - call, - task_context, - task_credits, - output, - lifecycle, - deadline, - deadline_source, - maximum_stream_chunk_bytes, - maximum_chunks, - ) - .await; - guard.finish(if result.is_ok() { - CallTaskOutcome::Complete - } else { - CallTaskOutcome::TransportFailed - }); - }); - self.active.insert( - id, - ActiveCall { - context, - credits, - cancellation, - task, - }, - ); - Ok(()) - } - - fn cancel_call(&mut self, id: u64) -> Result<(), SessionError> { - if id == 0 || id.is_multiple_of(2) || id > self.last_call { - return Err(SessionError::Protocol); - } - if let Some(active) = self.active.remove(&id) { - active.cancellation.cancel(); - self.callbacks.finish_parent(id); - self.lifecycle.cancel(active.context); - active.task.abort(); - } - self.output - .send_control(Frame::control(Message::Cancelled { id })) - } - - fn finish_call(&mut self, finished: CallFinished) -> Result<(), SessionError> { - let Some(_active) = self.active.remove(&finished.id) else { - return Ok(()); - }; - self.callbacks.finish_parent(finished.id); - match finished.outcome { - CallTaskOutcome::Complete => Ok(()), - CallTaskOutcome::TransportFailed => Err(SessionError::Closed), - CallTaskOutcome::HandlerStopped => Err(SessionError::HandlerStopped), - } - } - - async fn close(&mut self) { - for (id, active) in std::mem::take(&mut self.active) { - active.cancellation.cancel(); - self.callbacks.finish_parent(id); - self.lifecycle.cancel(active.context); - active.task.abort(); - } - self.callbacks.close(); - self.lifecycle.shutdown(); - self.writer.abort(); - let _ = (&mut self.writer).await; - } -} - -#[expect( - clippy::too_many_arguments, - reason = "调用任务显式接收其完整资源边界，避免隐藏第二份会话状态" -)] -async fn run_call( - handler: Arc, - call: PluginCall, - context: CallContext, - credits: Arc, - output: Output, - lifecycle: LifecycleHooks, - deadline: Option, - deadline_source: tokio::sync::watch::Sender>, - maximum_stream_chunk_bytes: usize, - maximum_chunks: usize, -) -> Result<(), SessionError> { - let id = context.call_id; - let cancellation = call.cancellation.clone(); - let response = tokio::select! { - biased; - () = cancellation.cancelled() => return Ok(()), - () = wait_deadline(deadline) => { - lifecycle.cancel(context.clone()); - output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), - ))?; - return Ok(()); - } - response = handler.call(call) => response, - }; - let mut reply = match response { - Ok(reply) => reply, - Err(error) => { - let _ = send_initial( - &output, - bounded_fault_frame(id, error), - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await?; - return Ok(()); - } - }; - let initial = Frame { - message: Message::Result { - id, - result: reply.result, - }, - payload: reply.payload, - }; - if validate_frame(&initial).is_err() { - output.send_control(fault_frame( - id, - PluginFault::new( - ErrorCode::InvalidInput, - "plugin response metadata exceeds its limit", - ), - ))?; - return Ok(()); - } - let Some(mut stream) = reply.stream.take() else { - let _ = send_initial( - &output, - initial, - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await?; - return Ok(()); - }; - let window_bytes = match credits.initial_window(deadline, &cancellation).await { - Ok(window) => window, - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(fault_frame( - id, - PluginFault::new(ErrorCode::Timeout, "stream credit timed out"), - ))?; - return Ok(()); - } - Err(SessionError::Cancelled) => return Ok(()), - Err(error) => return Err(error), - }; - if let Err(fault) = - stream.validate_buffered(maximum_chunks, maximum_stream_chunk_bytes, window_bytes) - { - output.send_control(fault_frame(id, fault))?; - return Ok(()); - } - if !send_initial( - &output, - initial, - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await? - { - return Ok(()); - } - let deadline = if context.resource_stream { - None - } else { - deadline - }; - deadline_source.send_replace(deadline); - let mut sequence = 0_u64; - loop { - let item = tokio::select! { - biased; - () = cancellation.cancelled() => return Ok(()), - () = wait_deadline(deadline) => { - lifecycle.cancel(context.clone()); - output.send_control(end_frame( - id, - Some(PluginFault::new(ErrorCode::Timeout, "plugin stream timed out")), - ))?; - return Ok(()); - } - item = stream.next() => item, - }; - let Some(item) = item else { - output.send_control(end_frame(id, None))?; - return Ok(()); - }; - let payload = match item { - Ok(payload) => payload, - Err(fault) => { - output.send_control(bounded_end_frame(id, fault))?; - return Ok(()); - } - }; - if let Err(fault) = - validate_stream_chunk(&payload, maximum_stream_chunk_bytes, window_bytes) - { - output.send_control(end_frame(id, Some(fault)))?; - return Ok(()); - } - match credits - .take(payload.len() as u64, deadline, &cancellation) - .await - { - Ok(()) => {} - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(end_frame( - id, - Some(PluginFault::new( - ErrorCode::Timeout, - "stream credit timed out", - )), - ))?; - return Ok(()); - } - Err(SessionError::Cancelled) => return Ok(()), - Err(error) => return Err(error), - } - let frame = Frame { - message: Message::Stream { id, sequence }, - payload, - }; - if !send_stream_data( - &output, - frame, - id, - &context, - &lifecycle, - deadline, - &cancellation, - ) - .await? - { - return Ok(()); - } - sequence = sequence.checked_add(1).ok_or(SessionError::Protocol)?; - } -} - -async fn send_initial( - output: &Output, - frame: Frame, - context: &CallContext, - lifecycle: &LifecycleHooks, - deadline: Option, - cancellation: &CallCancellation, -) -> Result { - match output.send_data_until(frame, deadline, cancellation).await { - Ok(()) => Ok(true), - Err(SessionError::Cancelled) => Ok(false), - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(fault_frame( - context.call_id, - PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), - ))?; - Ok(false) - } - Err(error) => Err(error), - } -} - -async fn send_stream_data( - output: &Output, - frame: Frame, - id: u64, - context: &CallContext, - lifecycle: &LifecycleHooks, - deadline: Option, - cancellation: &CallCancellation, -) -> Result { - match output.send_data_until(frame, deadline, cancellation).await { - Ok(()) => Ok(true), - Err(SessionError::Cancelled) => Ok(false), - Err(SessionError::Timeout) => { - lifecycle.cancel(context.clone()); - output.send_control(end_frame( - id, - Some(PluginFault::new( - ErrorCode::Timeout, - "plugin stream timed out", - )), - ))?; - Ok(false) - } - Err(error) => Err(error), - } -} - -#[derive(Default)] -struct CreditWindow { - state: Mutex, - changed: tokio::sync::Notify, -} - -#[derive(Default)] -struct CreditState { - bytes: u64, - frames: u64, - maximum_bytes: u64, -} - -impl CreditWindow { - fn grant(&self, bytes: u32, frames: u32) -> Result<(), SessionError> { - if bytes == 0 || frames == 0 { - return Err(SessionError::Protocol); - } - let mut state = lock_unpoisoned(&self.state); - state.bytes = state - .bytes - .checked_add(u64::from(bytes)) - .ok_or(SessionError::Protocol)?; - state.frames = state - .frames - .checked_add(u64::from(frames)) - .ok_or(SessionError::Protocol)?; - state.maximum_bytes = state.maximum_bytes.max(state.bytes); - drop(state); - // 每个调用只有一个顺序消费 Credit 的任务；notify_one 会保留许可，避免 - // grant 恰好发生在状态检查和等待注册之间时丢失唤醒 - self.changed.notify_one(); - Ok(()) - } - - async fn initial_window( - &self, - deadline: Option, - cancellation: &CallCancellation, - ) -> Result { - loop { - let changed = self.changed.notified(); - { - let state = lock_unpoisoned(&self.state); - if state.bytes > 0 && state.frames > 0 { - return Ok(state.maximum_bytes); - } - } - tokio::select! { - biased; - () = cancellation.cancelled() => return Err(SessionError::Cancelled), - () = wait_deadline(deadline) => return Err(SessionError::Timeout), - () = changed => {} - } - } - } - - async fn take( - &self, - bytes: u64, - deadline: Option, - cancellation: &CallCancellation, - ) -> Result<(), SessionError> { - loop { - let changed = self.changed.notified(); - { - let mut state = lock_unpoisoned(&self.state); - if bytes > state.maximum_bytes { - return Err(SessionError::Protocol); - } - if state.bytes >= bytes && state.frames > 0 { - state.bytes -= bytes; - state.frames -= 1; - return Ok(()); - } - } - tokio::select! { - biased; - () = cancellation.cancelled() => return Err(SessionError::Cancelled), - () = wait_deadline(deadline) => return Err(SessionError::Timeout), - () = changed => {} - } - } - } -} - -fn validate_stream_chunk( - payload: &[u8], - maximum_stream_chunk_bytes: usize, - window_bytes: u64, -) -> Result<(), PluginFault> { - if payload.is_empty() { - return Err(PluginFault::new( - ErrorCode::InvalidInput, - "stream chunks must not be empty", - )); - } - if payload.len() as u64 > window_bytes { - return Err(PluginFault::new( - ErrorCode::InvalidInput, - "response event exceeds the host stream credit window", - )); - } - if payload.len() > maximum_stream_chunk_bytes { - return Err(PluginFault::new( - ErrorCode::InvalidInput, - "response event exceeds the stream chunk budget", - )); - } - Ok(()) -} - -fn capacity_fault(message: &'static str) -> PluginFault { - PluginFault::new(ErrorCode::Capacity, message) -} - -fn fault_frame(id: u64, error: PluginFault) -> Frame { - Frame::control(Message::Error { id, error }) -} - -fn end_frame(id: u64, error: Option) -> Frame { - Frame::control(Message::End { id, error }) -} - -fn bounded_fault_frame(id: u64, error: PluginFault) -> Frame { - let frame = fault_frame(id, error); - if validate_frame(&frame).is_ok() { - frame - } else { - fault_frame( - id, - PluginFault::new(ErrorCode::Fault, "plugin error metadata exceeds its limit"), - ) - } -} - -fn bounded_end_frame(id: u64, error: PluginFault) -> Frame { - let frame = end_frame(id, Some(error)); - if validate_frame(&frame).is_ok() { - frame - } else { - end_frame( - id, - Some(PluginFault::new( - ErrorCode::Fault, - "plugin stream error metadata exceeds its limit", - )), - ) - } -} - -fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { - mutex - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -async fn wait_deadline(deadline: Option) { - match deadline { - Some(deadline) => tokio::time::sleep_until(deadline).await, - None => std::future::pending().await, - } -} - -async fn callback_deadline(mut deadline: tokio::sync::watch::Receiver>) { - loop { - let current = *deadline.borrow_and_update(); - tokio::select! { - () = wait_deadline(current) => return, - changed = deadline.changed() => { - if changed.is_err() { return wait_deadline(current).await; } - } - } - } -} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/call.rs b/backend/crates/gateway-plugin/sdk/src/client/session/call.rs new file mode 100644 index 000000000..ebc21bd69 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/call.rs @@ -0,0 +1,66 @@ +//! 插件业务调用、处理器与响应合同 + +use std::{future::Future, pin::Pin}; + +use serde_json::Value; + +use crate::{CallContext, PluginFault}; + +use super::{CallCancellation, HostClient, ResponseStream}; + +/// 交给插件业务处理器的一次宿主调用；不实现 `Debug`，避免敏感载荷泄露 +pub struct PluginCall { + pub method: String, + pub context: CallContext, + pub params: Value, + pub payload: Vec, + pub host: HostClient, + pub cancellation: CallCancellation, +} + +/// 插件业务调用 future +pub type CallFuture<'a> = Pin> + Send + 'a>>; + +/// 插件只实现业务调用；会话负责握手、关联、流控、取消和关闭 +pub trait PluginHandler: Send + Sync + 'static { + fn call(&self, call: PluginCall) -> CallFuture<'_>; + + /// 取消时释放只属于该调用的临时业务状态 + /// 实现必须幂等、快速且不得执行阻塞 I/O + fn cancel(&self, _context: &CallContext) {} + + /// 会话停止接收新调用时触发 + /// 实现必须幂等、快速且不得执行阻塞 I/O + fn quiesce(&self) {} + + /// 会话关闭时释放插件自有临时状态 + /// 实现必须幂等、快速且不得执行阻塞 I/O + fn shutdown(&self) {} +} + +/// 一次业务成功结果；`stream` 存在时 SDK 自动产生单一流终态 +pub struct CallReply { + pub(super) result: Value, + pub(super) payload: Vec, + pub(super) stream: Option, +} + +impl CallReply { + #[must_use] + pub fn unary(result: Value, payload: Vec) -> Self { + Self { + result, + payload, + stream: None, + } + } + + #[must_use] + pub fn stream(result: Value, payload: Vec, stream: ResponseStream) -> Self { + Self { + result, + payload, + stream: Some(stream), + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/config.rs b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs new file mode 100644 index 000000000..6c880ad38 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs @@ -0,0 +1,54 @@ +//! 插件侧会话的内存、并发与期限配置 + +use std::time::Duration; + +use super::SessionError; + +const MAXIMUM_STREAM_CHUNK_BYTES: usize = 16 * 1024 * 1024; +const MAXIMUM_CONCURRENCY: usize = 256; +const MAXIMUM_BUFFERED_STREAM_CHUNKS: usize = 65_536; + +/// 插件侧会话的内存、并发与期限边界 +#[derive(Debug, Clone, Copy)] +pub struct SessionConfig { + /// 单个业务流分块的预算，不限制普通调用正文的总长度 + pub maximum_stream_chunk_bytes: usize, + pub maximum_calls: usize, + pub maximum_callbacks: usize, + pub maximum_buffered_stream_chunks: usize, + pub handshake_timeout: Duration, + pub maximum_call_timeout: Duration, +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + // 流分块还需匹配宿主授予的信用窗口 + maximum_stream_chunk_bytes: 1024 * 1024, + maximum_calls: 32, + maximum_callbacks: 32, + maximum_buffered_stream_chunks: 1_024, + handshake_timeout: Duration::from_secs(5), + maximum_call_timeout: Duration::from_secs(120), + } + } +} + +impl SessionConfig { + pub(super) fn validate(self) -> Result { + if self.maximum_stream_chunk_bytes < 1_024 + || self.maximum_stream_chunk_bytes > MAXIMUM_STREAM_CHUNK_BYTES + || self.maximum_calls == 0 + || self.maximum_calls > MAXIMUM_CONCURRENCY + || self.maximum_callbacks == 0 + || self.maximum_callbacks > MAXIMUM_CONCURRENCY + || self.maximum_buffered_stream_chunks == 0 + || self.maximum_buffered_stream_chunks > MAXIMUM_BUFFERED_STREAM_CHUNKS + || self.handshake_timeout.is_zero() + || self.maximum_call_timeout.is_zero() + { + return Err(SessionError::Configuration); + } + Ok(self) + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/credit.rs b/backend/crates/gateway-plugin/sdk/src/client/session/credit.rs new file mode 100644 index 000000000..8e6810364 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/credit.rs @@ -0,0 +1,109 @@ +//! 响应流字节与帧信用的唯一状态和等待唤醒 + +use std::sync::Mutex; + +use tokio::time::Instant; + +use super::{CallCancellation, SessionError, lifecycle::wait_deadline}; + +#[derive(Default)] +pub(super) struct CreditWindow { + state: Mutex, + changed: tokio::sync::Notify, +} + +#[derive(Default)] +struct CreditState { + bytes: u64, + frames: u64, + maximum_bytes: u64, +} + +impl CreditWindow { + pub(super) fn maximum_bytes(&self) -> u64 { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .maximum_bytes + } + + pub(super) fn grant(&self, bytes: u32, frames: u32) -> Result<(), SessionError> { + if bytes == 0 || frames == 0 { + return Err(SessionError::Protocol); + } + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.bytes = state + .bytes + .checked_add(u64::from(bytes)) + .ok_or(SessionError::Protocol)?; + state.frames = state + .frames + .checked_add(u64::from(frames)) + .ok_or(SessionError::Protocol)?; + state.maximum_bytes = state.maximum_bytes.max(state.bytes); + drop(state); + // 每个调用只有一个顺序消费 Credit 的任务；notify_one 会保留许可，避免 + // grant 恰好发生在状态检查和等待注册之间时丢失唤醒 + self.changed.notify_one(); + Ok(()) + } + + pub(super) async fn initial_window( + &self, + deadline: Option, + cancellation: &CallCancellation, + ) -> Result { + loop { + let changed = self.changed.notified(); + { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.bytes > 0 && state.frames > 0 { + return Ok(state.maximum_bytes); + } + } + tokio::select! { + biased; + () = cancellation.cancelled() => return Err(SessionError::Cancelled), + () = wait_deadline(deadline) => return Err(SessionError::Timeout), + () = changed => {} + } + } + } + + pub(super) async fn take( + &self, + bytes: u64, + deadline: Option, + cancellation: &CallCancellation, + ) -> Result<(), SessionError> { + loop { + let changed = self.changed.notified(); + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if bytes > state.maximum_bytes { + return Err(SessionError::Protocol); + } + if state.bytes >= bytes && state.frames > 0 { + state.bytes -= bytes; + state.frames -= 1; + return Ok(()); + } + } + tokio::select! { + biased; + () = cancellation.cancelled() => return Err(SessionError::Cancelled), + () = wait_deadline(deadline) => return Err(SessionError::Timeout), + () = changed => {} + } + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs new file mode 100644 index 000000000..0c850465b --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs @@ -0,0 +1,712 @@ +//! 插件会话握手、在途调用分派与单一终结责任 + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use serde_json::Value; +use tokio::{ + io::{AsyncRead, AsyncWrite}, + sync::{Semaphore, mpsc}, + task::JoinHandle, + time::Instant, +}; + +use crate::{ + CallContext, ErrorCode, Frame, FrameError, Handshake, Message, PROTOCOL_VERSION, PluginFault, +}; + +use super::super::frame::{read_frame, validate_frame, write_frame}; +use super::{ + CallCancellation, HostClient, PluginCall, PluginHandler, SessionConfig, SessionError, + credit::CreditWindow, + lifecycle::{CancellationSource, LifecycleHooks, wait_deadline}, + output::{Output, writer_loop}, + registry::CallbackRegistry, + stream::validate_stream_chunk, +}; + +/// 已读取并校验 Hello、尚未向宿主发送 Ready 的插件会话 +pub struct PluginSession { + reader: R, + writer: W, + handshake: Handshake, + config: SessionConfig, +} + +impl PluginSession +where + R: AsyncRead + Unpin + Send + 'static, + W: AsyncWrite + Unpin + Send + 'static, +{ + /// 读取并校验宿主握手 + /// 业务处理器可以在 `run` 前读取不可变 Handshake 构造注册描述 + /// + /// # Errors + /// + /// 配置无效、握手超时、传输关闭或 Hello 不符合协议时失败 + pub async fn accept( + mut reader: R, + writer: W, + config: SessionConfig, + ) -> Result { + let config = config.validate()?; + let frame = tokio::time::timeout(config.handshake_timeout, read_frame(&mut reader)) + .await + .map_err(|_| SessionError::Timeout)??; + let Message::Hello { handshake } = frame.message else { + return Err(SessionError::Handshake); + }; + if !frame.payload.is_empty() + || handshake.protocol_version != PROTOCOL_VERSION + || handshake.plugin_id.is_empty() + || handshake.instance_id.is_empty() + || handshake.generation == 0 + || handshake.incarnation.is_empty() + { + return Err(SessionError::Handshake); + } + Ok(Self { + reader, + writer, + handshake, + config, + }) + } + + #[must_use] + pub const fn handshake(&self) -> &Handshake { + &self.handshake + } + + /// 发送 Ready 并持续处理调用、回调、流控、取消与关闭 + /// + /// # Errors + /// + /// 帧损坏、协议违规、处理器异常停止或传输关闭时失败 + pub async fn run(mut self, handler: H) -> Result<(), SessionError> + where + H: PluginHandler, + { + let data_capacity = self + .config + .maximum_calls + .checked_add(self.config.maximum_callbacks) + .ok_or(SessionError::Configuration)?; + let control_reserve = self + .config + .maximum_calls + .checked_add(8) + .ok_or(SessionError::Configuration)?; + let queue_capacity = data_capacity + .checked_add(control_reserve) + .ok_or(SessionError::Configuration)?; + let (outbound, received) = mpsc::channel(queue_capacity); + let output = Output { + sender: outbound, + data_slots: Arc::new(Semaphore::new(data_capacity)), + }; + let callbacks = Arc::new(CallbackRegistry::new(self.config.maximum_callbacks)); + let handler: Arc = Arc::new(handler); + let lifecycle = LifecycleHooks::new(Arc::clone(&handler), control_reserve)?; + write_frame( + &mut self.writer, + &Frame::control(Message::Ready { + protocol_version: PROTOCOL_VERSION, + incarnation: self.handshake.incarnation.clone(), + }), + ) + .await?; + let writer = tokio::spawn(writer_loop(self.writer, received)); + let (completed, completions) = mpsc::channel(self.config.maximum_calls); + let mut driver = SessionDriver { + config: self.config, + instance_id: self.handshake.instance_id, + generation: self.handshake.generation, + incarnation: self.handshake.incarnation, + output, + callbacks, + handler, + lifecycle, + writer, + completed, + completions, + active: BTreeMap::new(), + last_call: 0, + quiescing: false, + }; + let result = driver.drive(self.reader).await; + driver.close().await; + result + } +} + +struct ActiveCall { + context: CallContext, + credits: Arc, + cancellation: CancellationSource, + task: JoinHandle<()>, +} + +#[derive(Clone, Copy)] +enum CallTaskOutcome { + Complete, + TransportFailed, + HandlerStopped, +} + +struct CallFinished { + id: u64, + outcome: CallTaskOutcome, +} + +struct CompletionGuard { + id: u64, + completed: mpsc::Sender, + armed: bool, +} + +impl CompletionGuard { + fn new(id: u64, completed: mpsc::Sender) -> Self { + Self { + id, + completed, + armed: true, + } + } + + fn finish(mut self, outcome: CallTaskOutcome) { + self.armed = false; + let _ = self.completed.try_send(CallFinished { + id: self.id, + outcome, + }); + } +} + +impl Drop for CompletionGuard { + fn drop(&mut self) { + if self.armed { + let _ = self.completed.try_send(CallFinished { + id: self.id, + outcome: CallTaskOutcome::HandlerStopped, + }); + } + } +} + +struct SessionDriver { + config: SessionConfig, + instance_id: String, + generation: u64, + incarnation: String, + output: Output, + callbacks: Arc, + handler: Arc, + lifecycle: LifecycleHooks, + writer: JoinHandle>, + completed: mpsc::Sender, + completions: mpsc::Receiver, + active: BTreeMap, + last_call: u64, + quiescing: bool, +} + +impl SessionDriver { + async fn drive(&mut self, mut reader: R) -> Result<(), SessionError> { + loop { + // read_exact 不是取消安全的；完成通知可以穿插处理，但必须保留同一个半帧读取 future + let incoming = read_frame(&mut reader); + tokio::pin!(incoming); + let frame = loop { + tokio::select! { + biased; + writer = &mut self.writer => { + return writer.map_err(|_| SessionError::Closed)?; + } + Some(finished) = self.completions.recv() => { + self.finish_call(finished)?; + } + frame = &mut incoming => { + break frame.map_err(|error| match error { + FrameError::Io(io_error) + if io_error.kind() == std::io::ErrorKind::UnexpectedEof => + { + SessionError::Closed + } + error => SessionError::Frame(error), + })?; + } + } + }; + if self.handle_frame(frame)? { + return Ok(()); + } + } + } + + fn handle_frame(&mut self, frame: Frame) -> Result { + match frame.message { + Message::Call { + id, + method, + context, + params, + } => { + self.start_call(id, method, context, params, frame.payload)?; + Ok(false) + } + Message::Credit { id, bytes, frames } if frame.payload.is_empty() => { + if id == 0 + || id.is_multiple_of(2) + || id > self.last_call + || bytes == 0 + || frames == 0 + { + return Err(SessionError::Protocol); + } + // End/Cancelled 已进入同一 FIFO 后，宿主先前为已消费分块排队的 + // Credit 仍可能稍晚抵达 + // 此时调用资源已经回收，忽略迟到信用即可； + // 未来或回调 ID 仍按协议错误处理 + if let Some(active) = self.active.get(&id) { + active.credits.grant(bytes, frames)?; + } + Ok(false) + } + Message::Cancel { id } if frame.payload.is_empty() => { + self.cancel_call(id)?; + Ok(false) + } + Message::Result { id, result } => { + self.callbacks.resolve(Frame { + message: Message::Result { id, result }, + payload: frame.payload, + })?; + Ok(false) + } + Message::Error { id, error } => { + self.callbacks.resolve(Frame { + message: Message::Error { id, error }, + payload: frame.payload, + })?; + Ok(false) + } + Message::Quiesce if frame.payload.is_empty() => { + if !self.quiescing { + self.quiescing = true; + self.lifecycle.quiesce(); + } + Ok(false) + } + Message::Shutdown if frame.payload.is_empty() => Ok(true), + _ => Err(SessionError::Protocol), + } + } + + fn start_call( + &mut self, + id: u64, + method: String, + context: CallContext, + params: Value, + payload: Vec, + ) -> Result<(), SessionError> { + if id == 0 + || id.is_multiple_of(2) + || id <= self.last_call + || context.call_id != id + || context.instance_id != self.instance_id + || context.generation != self.generation + || context.incarnation != self.incarnation + || context.timeout_ms == 0 + || Duration::from_millis(context.timeout_ms) > self.config.maximum_call_timeout + || method.is_empty() + || method.len() > 128 + { + return Err(SessionError::Protocol); + } + self.last_call = id; + if self.quiescing { + return self.output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Capacity, "plugin is quiescing"), + )); + } + if self.active.len() >= self.config.maximum_calls { + return self.output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Capacity, "plugin call capacity is exhausted"), + )); + } + let deadline = Instant::now() + .checked_add(Duration::from_millis(context.timeout_ms)) + .ok_or(SessionError::Protocol)?; + let deadline = Some(deadline); + let (deadline_source, callback_deadline) = tokio::sync::watch::channel(deadline); + let credits = Arc::new(CreditWindow::default()); + let (cancellation, cancellation_signal) = CancellationSource::new(); + let host = HostClient { + parent_id: id, + deadline: callback_deadline, + cancellation: cancellation_signal.clone(), + callbacks: Arc::clone(&self.callbacks), + output: self.output.clone(), + maximum_stream_chunk_bytes: self.config.maximum_stream_chunk_bytes, + credits: credits.clone(), + }; + let call = PluginCall { + method, + context: context.clone(), + params, + payload, + host, + cancellation: cancellation_signal, + }; + let handler = Arc::clone(&self.handler); + let output = self.output.clone(); + let completed = self.completed.clone(); + let task_credits = Arc::clone(&credits); + let maximum_stream_chunk_bytes = self.config.maximum_stream_chunk_bytes; + let maximum_chunks = self.config.maximum_buffered_stream_chunks; + let task_context = context.clone(); + let lifecycle = self.lifecycle.clone(); + let task = tokio::spawn(async move { + let guard = CompletionGuard::new(id, completed); + let result = run_call( + handler, + call, + task_context, + task_credits, + output, + lifecycle, + deadline, + deadline_source, + maximum_stream_chunk_bytes, + maximum_chunks, + ) + .await; + guard.finish(if result.is_ok() { + CallTaskOutcome::Complete + } else { + CallTaskOutcome::TransportFailed + }); + }); + self.active.insert( + id, + ActiveCall { + context, + credits, + cancellation, + task, + }, + ); + Ok(()) + } + + fn cancel_call(&mut self, id: u64) -> Result<(), SessionError> { + if id == 0 || id.is_multiple_of(2) || id > self.last_call { + return Err(SessionError::Protocol); + } + if let Some(active) = self.active.remove(&id) { + active.cancellation.cancel(); + self.callbacks.finish_parent(id); + self.lifecycle.cancel(active.context); + active.task.abort(); + } + self.output + .send_control(Frame::control(Message::Cancelled { id })) + } + + fn finish_call(&mut self, finished: CallFinished) -> Result<(), SessionError> { + let Some(_active) = self.active.remove(&finished.id) else { + return Ok(()); + }; + self.callbacks.finish_parent(finished.id); + match finished.outcome { + CallTaskOutcome::Complete => Ok(()), + CallTaskOutcome::TransportFailed => Err(SessionError::Closed), + CallTaskOutcome::HandlerStopped => Err(SessionError::HandlerStopped), + } + } + + async fn close(&mut self) { + for (id, active) in std::mem::take(&mut self.active) { + active.cancellation.cancel(); + self.callbacks.finish_parent(id); + self.lifecycle.cancel(active.context); + active.task.abort(); + } + self.callbacks.close(); + self.lifecycle.shutdown(); + self.writer.abort(); + let _ = (&mut self.writer).await; + } +} + +#[expect( + clippy::too_many_arguments, + reason = "调用任务显式接收其完整资源边界，避免隐藏第二份会话状态" +)] +async fn run_call( + handler: Arc, + call: PluginCall, + context: CallContext, + credits: Arc, + output: Output, + lifecycle: LifecycleHooks, + deadline: Option, + deadline_source: tokio::sync::watch::Sender>, + maximum_stream_chunk_bytes: usize, + maximum_chunks: usize, +) -> Result<(), SessionError> { + let id = context.call_id; + let cancellation = call.cancellation.clone(); + let response = tokio::select! { + biased; + () = cancellation.cancelled() => return Ok(()), + () = wait_deadline(deadline) => { + lifecycle.cancel(context.clone()); + output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), + ))?; + return Ok(()); + } + response = handler.call(call) => response, + }; + let mut reply = match response { + Ok(reply) => reply, + Err(error) => { + let _ = send_initial( + &output, + bounded_fault_frame(id, error), + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await?; + return Ok(()); + } + }; + let initial = Frame { + message: Message::Result { + id, + result: reply.result, + }, + payload: reply.payload, + }; + if validate_frame(&initial).is_err() { + output.send_control(fault_frame( + id, + PluginFault::new( + ErrorCode::InvalidInput, + "plugin response metadata exceeds its limit", + ), + ))?; + return Ok(()); + } + let Some(mut stream) = reply.stream.take() else { + let _ = send_initial( + &output, + initial, + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await?; + return Ok(()); + }; + let window_bytes = match credits.initial_window(deadline, &cancellation).await { + Ok(window) => window, + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(fault_frame( + id, + PluginFault::new(ErrorCode::Timeout, "stream credit timed out"), + ))?; + return Ok(()); + } + Err(SessionError::Cancelled) => return Ok(()), + Err(error) => return Err(error), + }; + if let Err(fault) = + stream.validate_buffered(maximum_chunks, maximum_stream_chunk_bytes, window_bytes) + { + output.send_control(fault_frame(id, fault))?; + return Ok(()); + } + if !send_initial( + &output, + initial, + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await? + { + return Ok(()); + } + let deadline = if context.resource_stream { + None + } else { + deadline + }; + deadline_source.send_replace(deadline); + let mut sequence = 0_u64; + loop { + let item = tokio::select! { + biased; + () = cancellation.cancelled() => return Ok(()), + () = wait_deadline(deadline) => { + lifecycle.cancel(context.clone()); + output.send_control(end_frame( + id, + Some(PluginFault::new(ErrorCode::Timeout, "plugin stream timed out")), + ))?; + return Ok(()); + } + item = stream.next() => item, + }; + let Some(item) = item else { + output.send_control(end_frame(id, None))?; + return Ok(()); + }; + let payload = match item { + Ok(payload) => payload, + Err(fault) => { + output.send_control(bounded_end_frame(id, fault))?; + return Ok(()); + } + }; + if let Err(fault) = + validate_stream_chunk(&payload, maximum_stream_chunk_bytes, window_bytes) + { + output.send_control(end_frame(id, Some(fault)))?; + return Ok(()); + } + match credits + .take(payload.len() as u64, deadline, &cancellation) + .await + { + Ok(()) => {} + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(end_frame( + id, + Some(PluginFault::new( + ErrorCode::Timeout, + "stream credit timed out", + )), + ))?; + return Ok(()); + } + Err(SessionError::Cancelled) => return Ok(()), + Err(error) => return Err(error), + } + let frame = Frame { + message: Message::Stream { id, sequence }, + payload, + }; + if !send_stream_data( + &output, + frame, + id, + &context, + &lifecycle, + deadline, + &cancellation, + ) + .await? + { + return Ok(()); + } + sequence = sequence.checked_add(1).ok_or(SessionError::Protocol)?; + } +} + +async fn send_initial( + output: &Output, + frame: Frame, + context: &CallContext, + lifecycle: &LifecycleHooks, + deadline: Option, + cancellation: &CallCancellation, +) -> Result { + match output.send_data_until(frame, deadline, cancellation).await { + Ok(()) => Ok(true), + Err(SessionError::Cancelled) => Ok(false), + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(fault_frame( + context.call_id, + PluginFault::new(ErrorCode::Timeout, "plugin call timed out"), + ))?; + Ok(false) + } + Err(error) => Err(error), + } +} + +async fn send_stream_data( + output: &Output, + frame: Frame, + id: u64, + context: &CallContext, + lifecycle: &LifecycleHooks, + deadline: Option, + cancellation: &CallCancellation, +) -> Result { + match output.send_data_until(frame, deadline, cancellation).await { + Ok(()) => Ok(true), + Err(SessionError::Cancelled) => Ok(false), + Err(SessionError::Timeout) => { + lifecycle.cancel(context.clone()); + output.send_control(end_frame( + id, + Some(PluginFault::new( + ErrorCode::Timeout, + "plugin stream timed out", + )), + ))?; + Ok(false) + } + Err(error) => Err(error), + } +} + +fn fault_frame(id: u64, error: PluginFault) -> Frame { + Frame::control(Message::Error { id, error }) +} + +fn end_frame(id: u64, error: Option) -> Frame { + Frame::control(Message::End { id, error }) +} + +fn bounded_fault_frame(id: u64, error: PluginFault) -> Frame { + let frame = fault_frame(id, error); + if validate_frame(&frame).is_ok() { + frame + } else { + fault_frame( + id, + PluginFault::new(ErrorCode::Fault, "plugin error metadata exceeds its limit"), + ) + } +} + +fn bounded_end_frame(id: u64, error: PluginFault) -> Frame { + let frame = end_frame(id, Some(error)); + if validate_frame(&frame).is_ok() { + frame + } else { + end_frame( + id, + Some(PluginFault::new( + ErrorCode::Fault, + "plugin stream error metadata exceeds its limit", + )), + ) + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/error.rs b/backend/crates/gateway-plugin/sdk/src/client/session/error.rs new file mode 100644 index 000000000..834027fc8 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/error.rs @@ -0,0 +1,47 @@ +//! 插件侧传输和会话错误及稳定业务错误投影 + +use crate::{ErrorCode, FrameError, PluginFault}; + +/// 插件侧传输和会话错误 +#[derive(Debug, thiserror::Error)] +pub enum SessionError { + #[error("plugin session configuration is invalid")] + Configuration, + #[error("plugin handshake is invalid")] + Handshake, + #[error("plugin protocol is invalid")] + Protocol, + #[error("plugin transport is closed")] + Closed, + #[error("plugin operation exceeded its deadline")] + Timeout, + #[error("plugin call was cancelled")] + Cancelled, + #[error("plugin session capacity is exhausted")] + Capacity, + #[error("plugin handler stopped unexpectedly")] + HandlerStopped, + #[error("host callback returned an error: {0:?}")] + Remote(PluginFault), + #[error(transparent)] + Frame(#[from] FrameError), +} + +impl SessionError { + /// 将本地会话失败收敛为可返回宿主的稳定插件错误 + #[must_use] + pub fn into_plugin_fault(self) -> PluginFault { + match self { + Self::Remote(fault) => fault, + Self::Timeout => PluginFault::new(ErrorCode::Timeout, "host callback timed out"), + Self::Cancelled => PluginFault::new(ErrorCode::Cancelled, "parent call was cancelled"), + Self::Capacity => PluginFault::new(ErrorCode::Capacity, "plugin capacity is exhausted"), + Self::Configuration | Self::Protocol | Self::Handshake | Self::Frame(_) => { + PluginFault::new(ErrorCode::InvalidInput, "plugin session input is invalid") + } + Self::Closed | Self::HandlerStopped => { + PluginFault::new(ErrorCode::Fault, "plugin session is unavailable") + } + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs b/backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs new file mode 100644 index 000000000..fefc65cdb --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/lifecycle.rs @@ -0,0 +1,116 @@ +//! 调用取消、期限等待与插件生命周期钩子的分派 + +use std::sync::Arc; + +use tokio::{sync::watch, time::Instant}; + +use crate::CallContext; + +use super::{PluginHandler, SessionError}; + +/// 一次宿主调用的可克隆取消信号 +#[derive(Clone)] +pub struct CallCancellation { + receiver: watch::Receiver, +} + +impl CallCancellation { + #[must_use] + pub fn is_cancelled(&self) -> bool { + *self.receiver.borrow() + } + + pub async fn cancelled(&self) { + let mut receiver = self.receiver.clone(); + if *receiver.borrow() { + return; + } + while receiver.changed().await.is_ok() { + if *receiver.borrow() { + return; + } + } + } +} + +pub(super) struct CancellationSource { + sender: watch::Sender, +} + +impl CancellationSource { + pub(super) fn new() -> (Self, CallCancellation) { + let (sender, receiver) = watch::channel(false); + (Self { sender }, CallCancellation { receiver }) + } + + pub(super) fn cancel(&self) { + self.sender.send_replace(true); + } +} + +enum LifecycleEvent { + Cancel(CallContext), + Quiesce, + Shutdown, +} + +#[derive(Clone)] +pub(super) struct LifecycleHooks { + events: std::sync::mpsc::SyncSender, +} + +impl LifecycleHooks { + pub(super) fn new( + handler: Arc, + capacity: usize, + ) -> Result { + let (events, received) = std::sync::mpsc::sync_channel(capacity); + std::thread::Builder::new() + .name("gateway-plugin-sdk-lifecycle".into()) + .spawn(move || { + while let Ok(event) = received.recv() { + match event { + LifecycleEvent::Cancel(context) => handler.cancel(&context), + LifecycleEvent::Quiesce => handler.quiesce(), + LifecycleEvent::Shutdown => { + handler.shutdown(); + break; + } + } + } + }) + .map_err(|_| SessionError::Closed)?; + Ok(Self { events }) + } + + pub(super) fn cancel(&self, context: CallContext) { + let _ = self.events.try_send(LifecycleEvent::Cancel(context)); + } + + pub(super) fn quiesce(&self) { + let _ = self.events.try_send(LifecycleEvent::Quiesce); + } + + pub(super) fn shutdown(&self) { + let _ = self.events.try_send(LifecycleEvent::Shutdown); + } +} + +pub(super) async fn wait_deadline(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } +} + +pub(super) async fn callback_deadline(mut deadline: tokio::sync::watch::Receiver>) { + loop { + let current = *deadline.borrow_and_update(); + tokio::select! { + () = wait_deadline(current) => return, + changed = deadline.changed() => { + if changed.is_err() { return wait_deadline(current).await; } + } + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs b/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs new file mode 100644 index 000000000..7e7b5f62b --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs @@ -0,0 +1,19 @@ +//! 插件侧双向 RPC 会话的公开合同与内部职责入口 + +mod call; +mod config; +mod credit; +mod driver; +mod error; +mod lifecycle; +mod output; +mod registry; +mod stream; + +pub use call::{CallFuture, CallReply, PluginCall, PluginHandler}; +pub use config::SessionConfig; +pub use driver::PluginSession; +pub use error::SessionError; +pub use lifecycle::CallCancellation; +pub use registry::{HostClient, HostReply}; +pub use stream::{PullResponseFuture, PullResponseStream, ResponseStream, StreamSender}; diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/output.rs b/backend/crates/gateway-plugin/sdk/src/client/session/output.rs new file mode 100644 index 000000000..aa2a7cffb --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/output.rs @@ -0,0 +1,84 @@ +//! 会话输出队列、控制帧预留与唯一串行写泵 + +use std::sync::Arc; + +use tokio::{ + io::AsyncWrite, + sync::{OwnedSemaphorePermit, Semaphore, mpsc}, + time::Instant, +}; + +use crate::Frame; + +use super::super::frame::{validate_frame, write_frame}; +use super::{CallCancellation, SessionError, lifecycle::wait_deadline}; + +pub(super) struct Output { + pub(super) sender: mpsc::Sender, + pub(super) data_slots: Arc, +} + +impl Clone for Output { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + data_slots: Arc::clone(&self.data_slots), + } + } +} + +impl Output { + pub(super) async fn send_data(&self, frame: Frame) -> Result<(), SessionError> { + validate_frame(&frame)?; + let permit = Arc::clone(&self.data_slots) + .acquire_owned() + .await + .map_err(|_| SessionError::Closed)?; + self.sender + .send(Outbound { + frame, + _data_slot: Some(permit), + }) + .await + .map_err(|_| SessionError::Closed) + } + + pub(super) async fn send_data_until( + &self, + frame: Frame, + deadline: Option, + cancellation: &CallCancellation, + ) -> Result<(), SessionError> { + tokio::select! { + biased; + () = cancellation.cancelled() => Err(SessionError::Cancelled), + () = wait_deadline(deadline) => Err(SessionError::Timeout), + result = self.send_data(frame) => result, + } + } + + pub(super) fn send_control(&self, frame: Frame) -> Result<(), SessionError> { + validate_frame(&frame)?; + self.sender + .try_send(Outbound { + frame, + _data_slot: None, + }) + .map_err(|_| SessionError::Capacity) + } +} + +pub(super) struct Outbound { + frame: Frame, + _data_slot: Option, +} + +pub(super) async fn writer_loop( + mut writer: W, + mut received: mpsc::Receiver, +) -> Result<(), SessionError> { + while let Some(outbound) = received.recv().await { + write_frame(&mut writer, &outbound.frame).await?; + } + Ok(()) +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs new file mode 100644 index 000000000..cfe11e86e --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs @@ -0,0 +1,251 @@ +//! 宿主回调客户端、关联编号与有界退休回调注册表 + +use std::{ + collections::{BTreeMap, BTreeSet, VecDeque}, + sync::{Arc, Mutex}, +}; + +use serde_json::Value; +use tokio::{sync::oneshot, time::Instant}; + +use crate::{Frame, Message}; + +use super::super::frame::validate_frame; +use super::{ + CallCancellation, SessionError, credit::CreditWindow, lifecycle::callback_deadline, + output::Output, +}; + +/// 宿主回调的成功结果；不实现 `Debug`，避免载荷进入普通诊断 +pub struct HostReply { + pub result: Value, + pub payload: Vec, +} + +/// 绑定父调用、期限和取消信号的宿主回调客户端 +#[derive(Clone)] +pub struct HostClient { + pub(super) parent_id: u64, + pub(super) deadline: tokio::sync::watch::Receiver>, + pub(super) cancellation: CallCancellation, + pub(super) callbacks: Arc, + pub(super) output: Output, + pub(super) maximum_stream_chunk_bytes: usize, + pub(super) credits: Arc, +} + +impl HostClient { + pub(crate) fn maximum_stream_chunk_bytes(&self) -> usize { + let window = self.credits.maximum_bytes(); + if window == 0 { + self.maximum_stream_chunk_bytes + } else { + self.maximum_stream_chunk_bytes + .min(usize::try_from(window).unwrap_or(usize::MAX)) + } + } + + /// 发起一次与父调用关联的宿主回调 + /// + /// # Errors + /// + /// 方法或帧无效、容量耗尽、父调用取消、期限到达、连接关闭或宿主返回错误时失败 + pub async fn call( + &self, + method: impl Into, + params: Value, + payload: Vec, + ) -> Result { + let method = method.into(); + if method.is_empty() || method.len() > 128 { + return Err(SessionError::Protocol); + } + let ordered = self.callbacks.order.lock().await; + let (id, received) = self.callbacks.begin(self.parent_id)?; + let frame = Frame { + message: Message::Callback { + id, + parent_id: self.parent_id, + method, + params, + }, + payload, + }; + if validate_frame(&frame).is_err() { + self.callbacks.retire(id); + return Err(SessionError::Protocol); + } + let sent = tokio::select! { + biased; + () = self.cancellation.cancelled() => Err(SessionError::Cancelled), + () = callback_deadline(self.deadline.clone()) => Err(SessionError::Timeout), + result = self.output.send_data(frame) => result, + }; + if let Err(error) = sent { + self.callbacks.retire(id); + return Err(error); + } + drop(ordered); + tokio::select! { + biased; + () = self.cancellation.cancelled() => { + self.callbacks.retire(id); + Err(SessionError::Cancelled) + } + () = callback_deadline(self.deadline.clone()) => { + self.callbacks.retire(id); + Err(SessionError::Timeout) + } + result = received => result.map_err(|_| SessionError::Closed)?, + } + } +} + +pub(super) struct CallbackRegistry { + state: Mutex, + order: tokio::sync::Mutex<()>, + maximum: usize, + maximum_retired: usize, +} + +struct CallbackState { + next_id: u64, + pending: BTreeMap, + retired: BTreeSet, + retired_order: VecDeque, +} + +struct PendingCallback { + parent_id: u64, + result: oneshot::Sender>, +} + +impl CallbackRegistry { + pub(super) fn new(maximum: usize) -> Self { + Self { + state: Mutex::new(CallbackState { + next_id: 2, + pending: BTreeMap::new(), + retired: BTreeSet::new(), + retired_order: VecDeque::new(), + }), + order: tokio::sync::Mutex::new(()), + maximum, + maximum_retired: maximum.saturating_mul(2), + } + } + + fn begin( + &self, + parent_id: u64, + ) -> Result<(u64, oneshot::Receiver>), SessionError> { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.pending.len() >= self.maximum { + return Err(SessionError::Capacity); + } + let id = state.next_id; + state.next_id = state.next_id.checked_add(2).ok_or(SessionError::Capacity)?; + let (result, received) = oneshot::channel(); + state + .pending + .insert(id, PendingCallback { parent_id, result }); + Ok((id, received)) + } + + pub(super) fn resolve(&self, frame: Frame) -> Result<(), SessionError> { + let id = match &frame.message { + Message::Result { id, .. } | Message::Error { id, .. } => *id, + _ => return Err(SessionError::Protocol), + }; + let pending = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(pending) = state.pending.remove(&id) { + Some(pending) + } else if state.retired.remove(&id) { + state.retired_order.retain(|retired| *retired != id); + None + } else { + return Err(SessionError::Protocol); + } + }; + let Some(pending) = pending else { + return Ok(()); + }; + let result = match frame.message { + Message::Result { result, .. } => Ok(HostReply { + result, + payload: frame.payload, + }), + Message::Error { error, .. } if frame.payload.is_empty() => { + Err(SessionError::Remote(error)) + } + _ => return Err(SessionError::Protocol), + }; + let _ = pending.result.send(result); + Ok(()) + } + + fn retire(&self, id: u64) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.pending.remove(&id).is_some() { + retire_callback(&mut state, id, self.maximum_retired); + } + } + + pub(super) fn finish_parent(&self, parent_id: u64) { + let pending = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut remaining = BTreeMap::new(); + let mut finished = Vec::new(); + for (id, pending) in std::mem::take(&mut state.pending) { + if pending.parent_id == parent_id { + retire_callback(&mut state, id, self.maximum_retired); + finished.push(pending); + } else { + remaining.insert(id, pending); + } + } + state.pending = remaining; + finished + }; + for pending in pending { + let _ = pending.result.send(Err(SessionError::Cancelled)); + } + } + + pub(super) fn close(&self) { + let pending = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + std::mem::take(&mut state.pending) + }; + for pending in pending.into_values() { + let _ = pending.result.send(Err(SessionError::Closed)); + } + } +} + +fn retire_callback(state: &mut CallbackState, id: u64, maximum: usize) { + if state.retired.insert(id) { + state.retired_order.push_back(id); + } + while state.retired_order.len() > maximum { + if let Some(expired) = state.retired_order.pop_front() { + state.retired.remove(&expired); + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/stream.rs b/backend/crates/gateway-plugin/sdk/src/client/session/stream.rs new file mode 100644 index 000000000..c10690122 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/stream.rs @@ -0,0 +1,142 @@ +//! 插件响应流的有界生产、拉取与分块校验 + +use std::{collections::VecDeque, future::Future, num::NonZeroUsize, pin::Pin}; + +use tokio::sync::mpsc; + +use crate::{ErrorCode, PluginFault}; + +use super::SessionError; + +enum StreamSource { + Buffered(VecDeque>), + Channel(mpsc::Receiver, PluginFault>>), + Pull(Box), +} + +pub type PullResponseFuture<'a> = + Pin, PluginFault>>> + Send + 'a>>; + +pub trait PullResponseStream: Send { + fn next(&mut self) -> PullResponseFuture<'_>; +} + +/// SDK 管理 sequence、Credit 和终态的响应流 +pub struct ResponseStream { + source: StreamSource, + declared_capacity: usize, +} + +impl ResponseStream { + /// 从已完成业务校验的有限分块构造流；SDK 会在发送 `Result` 前预检全部分块 + #[must_use] + pub fn from_chunks(chunks: Vec>) -> Self { + let declared_capacity = chunks.len().max(1); + Self { + source: StreamSource::Buffered(chunks.into()), + declared_capacity, + } + } + + /// 创建由生产者驱动的有界流队列 + #[must_use] + pub fn channel(capacity: NonZeroUsize) -> (StreamSender, Self) { + let (sender, receiver) = mpsc::channel(capacity.get()); + ( + StreamSender { sender }, + Self { + source: StreamSource::Channel(receiver), + declared_capacity: capacity.get(), + }, + ) + } + + /// 按消费进度拉取下一帧；返回 None 表示完成，丢弃流会丢弃生产者 + pub fn pull(source: Box) -> Self { + Self { + source: StreamSource::Pull(source), + declared_capacity: 1, + } + } + + pub(super) fn validate_buffered( + &self, + maximum_chunks: usize, + maximum_stream_chunk_bytes: usize, + window_bytes: u64, + ) -> Result<(), PluginFault> { + if self.declared_capacity > maximum_chunks { + return Err(capacity_fault( + "response stream exceeds the local queue limit", + )); + } + if let StreamSource::Buffered(chunks) = &self.source { + for chunk in chunks { + validate_stream_chunk(chunk, maximum_stream_chunk_bytes, window_bytes)?; + } + } + Ok(()) + } + + pub(crate) async fn next(&mut self) -> Option, PluginFault>> { + match &mut self.source { + StreamSource::Buffered(chunks) => chunks.pop_front().map(Ok), + StreamSource::Channel(receiver) => receiver.recv().await, + StreamSource::Pull(source) => source.next().await, + } + } +} + +/// 动态响应流的有界生产端；最后一个 sender 被释放表示成功终态 +#[derive(Clone)] +pub struct StreamSender { + sender: mpsc::Sender, PluginFault>>, +} + +impl StreamSender { + /// 排队一个业务分块；实际线协议流控由会话完成 + pub async fn send(&self, payload: Vec) -> Result<(), SessionError> { + self.sender + .send(Ok(payload)) + .await + .map_err(|_| SessionError::Closed) + } + + /// 以业务错误结束动态流 + pub async fn fail(&self, fault: PluginFault) -> Result<(), SessionError> { + self.sender + .send(Err(fault)) + .await + .map_err(|_| SessionError::Closed) + } +} + +pub(super) fn validate_stream_chunk( + payload: &[u8], + maximum_stream_chunk_bytes: usize, + window_bytes: u64, +) -> Result<(), PluginFault> { + if payload.is_empty() { + return Err(PluginFault::new( + ErrorCode::InvalidInput, + "stream chunks must not be empty", + )); + } + if payload.len() as u64 > window_bytes { + return Err(PluginFault::new( + ErrorCode::InvalidInput, + "response event exceeds the host stream credit window", + )); + } + if payload.len() > maximum_stream_chunk_bytes { + return Err(PluginFault::new( + ErrorCode::InvalidInput, + "response event exceeds the stream chunk budget", + )); + } + Ok(()) +} + +fn capacity_fault(message: &'static str) -> PluginFault { + PluginFault::new(ErrorCode::Capacity, message) +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/frame.rs b/backend/crates/gateway-plugin/sdk/tests/client/frame.rs new file mode 100644 index 000000000..a7ed59402 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/frame.rs @@ -0,0 +1,87 @@ +//! 二进制帧往返、长度预算与不完整帧拒绝测试 + +use gateway_plugin_sdk::{ + Frame, FrameError, Message, + client::{read_frame, write_frame}, +}; +use serde_json::json; + +#[tokio::test] +async fn binary_payload_round_trips_without_json_encoding() { + let frame = Frame { + message: Message::Stream { + id: 17, + sequence: 2, + }, + payload: vec![0, 255, 13, 10, 128], + }; + let mut bytes = Vec::new(); + write_frame(&mut bytes, &frame).await.unwrap(); + assert!(bytes.ends_with(&frame.payload)); + assert_eq!(read_frame(&mut bytes.as_slice()).await.unwrap(), frame); +} + +#[tokio::test] +async fn malicious_lengths_fail_before_reading_or_allocating_payload() { + let bytes = [0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0]; + assert!(matches!( + read_frame(&mut bytes.as_slice()).await, + Err(FrameError::Length) + )); +} + +#[tokio::test] +async fn large_payload_is_exact_across_multiple_io_buffers_and_following_frames() { + let frame = Frame { + message: Message::Result { + id: 1, + result: json!({}), + }, + payload: (0..101 * 1024 * 1024 + 13) + .map(|index| (index % 251) as u8) + .collect(), + }; + let next = Frame::control(Message::Cancel { id: 3 }); + let (mut writer, mut reader) = tokio::io::duplex(16 * 1024); + let writing = tokio::spawn(async move { + write_frame(&mut writer, &frame).await.unwrap(); + write_frame(&mut writer, &next).await.unwrap(); + (frame, next) + }); + let received = read_frame(&mut reader).await.unwrap(); + let following = read_frame(&mut reader).await.unwrap(); + let (original, next) = writing.await.unwrap(); + assert_eq!(received, original); + assert_eq!(following, next); +} + +#[tokio::test] +async fn truncated_large_payload_never_becomes_a_successful_message() { + let message = serde_json::to_vec(&Message::Result { + id: 1, + result: json!({}), + }) + .unwrap(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(message.len() as u32).to_be_bytes()); + bytes.extend_from_slice(&(101_u64 * 1024 * 1024).to_be_bytes()); + bytes.extend_from_slice(&message); + bytes.extend_from_slice(&[1, 2, 3]); + assert!(matches!( + read_frame(&mut bytes.as_slice()).await, + Err(FrameError::Io(_)) + )); +} + +#[tokio::test] +async fn truncated_frame_never_becomes_a_successful_message() { + let mut bytes = Vec::new(); + write_frame(&mut bytes, &Frame::control(Message::Cancel { id: 3 })) + .await + .unwrap(); + bytes.pop(); + assert!(matches!( + read_frame(&mut bytes.as_slice()).await, + Err(FrameError::Io(_)) + )); +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/mod.rs b/backend/crates/gateway-plugin/sdk/tests/client/mod.rs index e239f0e45..551836fae 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/mod.rs @@ -1,5 +1,6 @@ //! 插件客户端、处理器注册与会话的测试入口 +mod frame; mod http; mod plugin; mod read; diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session.rs b/backend/crates/gateway-plugin/sdk/tests/client/session.rs deleted file mode 100644 index e3ec1c119..000000000 --- a/backend/crates/gateway-plugin/sdk/tests/client/session.rs +++ /dev/null @@ -1,1760 +0,0 @@ -//! 验证插件会话的二进制传输、帧长度限制与不完整消息拒绝 - -use std::{ - future::pending, - num::NonZeroUsize, - sync::{ - Arc, - atomic::{AtomicBool, AtomicUsize, Ordering}, - }, - time::Duration, -}; - -use gateway_plugin_sdk::{ - CallContext, Capability, ContributionDeclaration, Contributions, ErrorCode, Frame, FrameError, - Handshake, Message, PROTOCOL_VERSION, PluginFault, Stage, - call::middleware::{ - BODY_READ_METHOD, HANDLE_METHOD, MiddlewareBodyDisposition, MiddlewareBodyFrame, - MiddlewareBodyFraming, MiddlewareBodyHandle, MiddlewareBodyReadResult, MiddlewareHeader, - MiddlewareHeaderMutation, MiddlewareMount, MiddlewareNextRequest, MiddlewareNextResponse, - MiddlewareRequestBody, MiddlewareRequestHead, MiddlewareResponseBody, - MiddlewareResponseHead, MiddlewareTransport, NEXT_METHOD, - }, - client::{ - CallFuture, CallReply, MiddlewarePlugin, PluginCall, PluginHandler, PluginSession, - RequestCall, ResponseStream, SessionConfig, SessionError, read_frame, write_frame, - }, -}; -use serde_json::{Value, json}; -use tokio::{ - io::{AsyncWriteExt, DuplexStream, ReadHalf, WriteHalf}, - task::JoinHandle, -}; - -const MAXIMUM_STREAM_CHUNK_BYTES: usize = 64 * 1024; - -#[tokio::test] -async fn binary_payload_round_trips_without_json_encoding() { - let frame = Frame { - message: Message::Stream { - id: 17, - sequence: 2, - }, - payload: vec![0, 255, 13, 10, 128], - }; - let mut bytes = Vec::new(); - write_frame(&mut bytes, &frame).await.unwrap(); - assert!(bytes.ends_with(&frame.payload)); - assert_eq!(read_frame(&mut bytes.as_slice()).await.unwrap(), frame); -} - -#[tokio::test] -async fn malicious_lengths_fail_before_reading_or_allocating_payload() { - let bytes = [0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0]; - assert!(matches!( - read_frame(&mut bytes.as_slice()).await, - Err(FrameError::Length) - )); -} - -#[tokio::test] -async fn large_payload_is_exact_across_multiple_io_buffers_and_following_frames() { - let frame = Frame { - message: Message::Result { - id: 1, - result: json!({}), - }, - payload: (0..101 * 1024 * 1024 + 13) - .map(|index| (index % 251) as u8) - .collect(), - }; - let next = Frame::control(Message::Cancel { id: 3 }); - let (mut writer, mut reader) = tokio::io::duplex(16 * 1024); - let writing = tokio::spawn(async move { - write_frame(&mut writer, &frame).await.unwrap(); - write_frame(&mut writer, &next).await.unwrap(); - (frame, next) - }); - let received = read_frame(&mut reader).await.unwrap(); - let following = read_frame(&mut reader).await.unwrap(); - let (original, next) = writing.await.unwrap(); - assert_eq!(received, original); - assert_eq!(following, next); -} - -#[tokio::test] -async fn truncated_large_payload_never_becomes_a_successful_message() { - let message = serde_json::to_vec(&Message::Result { - id: 1, - result: json!({}), - }) - .unwrap(); - let mut bytes = Vec::new(); - bytes.extend_from_slice(&(message.len() as u32).to_be_bytes()); - bytes.extend_from_slice(&(101_u64 * 1024 * 1024).to_be_bytes()); - bytes.extend_from_slice(&message); - bytes.extend_from_slice(&[1, 2, 3]); - assert!(matches!( - read_frame(&mut bytes.as_slice()).await, - Err(FrameError::Io(_)) - )); -} - -#[tokio::test] -async fn truncated_frame_never_becomes_a_successful_message() { - let mut bytes = Vec::new(); - write_frame(&mut bytes, &Frame::control(Message::Cancel { id: 3 })) - .await - .unwrap(); - bytes.pop(); - assert!(matches!( - read_frame(&mut bytes.as_slice()).await, - Err(FrameError::Io(_)) - )); -} - -#[derive(Clone, Default)] -struct TestHandler { - lifecycle: Arc, - cancel_delay: Option, -} - -#[derive(Default)] -struct LifecycleState { - cancel_started: AtomicBool, - cancelled: AtomicUsize, - quiesced: AtomicBool, - shutdown: AtomicBool, -} - -impl PluginHandler for TestHandler { - fn call(&self, call: PluginCall) -> CallFuture<'_> { - Box::pin(async move { - match call.method.as_str() { - "echo" => Ok(CallReply::unary(call.params, call.payload)), - "key_facts" => { - let result = call - .host - .key_facts(gateway_plugin_sdk::call::data::ClientKeyFactsQuery { - client_key_id: "key_1".into(), - }) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "get_budget" => { - let result = call - .host - .get_key_budget( - gateway_plugin_sdk::call::key_budgets::GetKeyBudgetRequest { - client_key_id: "key_1".into(), - }, - ) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "update_budget_limits" => { - let result = call - .host - .update_key_budget_limits( - gateway_plugin_sdk::call::key_budgets::UpdateKeyBudgetLimitsRequest { - client_key_id: "key_1".into(), - daily_limit_usd: None, - weekly_limit_usd: Some("12.5".into()), - }, - ) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "refresh_quota" => { - let result = call - .host - .refresh_account_quota(gateway_plugin_sdk::call::data::QuotaFactsQuery { - account_id: "acct_1".into(), - }) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "reset_budget" => { - use gateway_plugin_sdk::call::{ - host::KeyListRequest, - key_budgets::{BudgetPeriod, ResetKeyBudgetRequest}, - }; - let keys = call - .host - .list_keys(KeyListRequest { - cursor: None, - limit: 10, - }) - .await?; - let result = call - .host - .reset_key_budget(ResetKeyBudgetRequest { - client_key_id: keys.keys[0].id.clone(), - period: BudgetPeriod::Weekly, - }) - .await?; - Ok(CallReply::unary( - serde_json::to_value(result).unwrap(), - vec![], - )) - } - "slow" => { - tokio::time::sleep(Duration::from_millis(80)).await; - Ok(CallReply::unary(call.params, call.payload)) - } - "callback" => { - let reply = call - .host - .call("host.log", call.params, call.payload) - .await - .map_err(SessionError::into_plugin_fault)?; - Ok(CallReply::unary(reply.result, reply.payload)) - } - "resource_stream" => Ok(CallReply::stream( - json!({}), - Vec::new(), - ResponseStream::pull(Box::new(ResourceStream(Some(call.host)))), - )), - "stream" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"abcd".to_vec(), b"efgh".to_vec()]), - )), - "long_stream" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"x".to_vec(); 16]), - )), - "oversized_stream" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"abcde".to_vec()]), - )), - "too_many_chunks" => Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - ResponseStream::from_chunks(vec![b"x".to_vec(); 17]), - )), - "dynamic_oversized_stream" => { - let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); - tokio::spawn(async move { - let _ = sender.send(b"abcde".to_vec()).await; - }); - Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - stream, - )) - } - "oversized_fault" => Err(PluginFault::new( - ErrorCode::Upstream, - "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), - )), - "oversized_result" => Ok(CallReply::unary( - Value::String("x".repeat(MAXIMUM_STREAM_CHUNK_BYTES)), - Vec::new(), - )), - "dynamic_oversized_fault" => { - let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); - tokio::spawn(async move { - let _ = sender - .fail(PluginFault::new( - ErrorCode::Upstream, - "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), - )) - .await; - }); - Ok(CallReply::stream( - json!({"stream": true}), - Vec::new(), - stream, - )) - } - "pending" => pending::>().await, - _ => Err(PluginFault::new( - ErrorCode::Unsupported, - "unsupported test method", - )), - } - }) - } - - fn cancel(&self, _context: &CallContext) { - self.lifecycle.cancel_started.store(true, Ordering::Release); - if let Some(delay) = self.cancel_delay { - std::thread::sleep(delay); - } - self.lifecycle.cancelled.fetch_add(1, Ordering::Relaxed); - } - - fn quiesce(&self) { - self.lifecycle.quiesced.store(true, Ordering::Release); - } - - fn shutdown(&self) { - self.lifecycle.shutdown.store(true, Ordering::Release); - } -} - -pub(super) struct HostPeer { - pub(super) reader: ReadHalf, - pub(super) writer: WriteHalf, -} - -pub(super) async fn start_session( - handler: H, -) -> (HostPeer, JoinHandle>) { - start_session_with_capacity(handler, MAXIMUM_STREAM_CHUNK_BYTES * 2).await -} - -async fn start_session_with_capacity( - handler: H, - capacity: usize, -) -> (HostPeer, JoinHandle>) { - let (host, plugin) = tokio::io::duplex(capacity); - let (plugin_reader, plugin_writer) = tokio::io::split(plugin); - let task = tokio::spawn(async move { - let session = PluginSession::accept( - plugin_reader, - plugin_writer, - SessionConfig { - maximum_stream_chunk_bytes: MAXIMUM_STREAM_CHUNK_BYTES, - maximum_calls: 4, - maximum_callbacks: 4, - maximum_buffered_stream_chunks: 16, - handshake_timeout: Duration::from_secs(1), - maximum_call_timeout: Duration::from_secs(2), - }, - ) - .await?; - session.run(handler).await - }); - let (mut reader, mut writer) = tokio::io::split(host); - write_frame( - &mut writer, - &Frame::control(Message::Hello { - handshake: handshake(), - }), - ) - .await - .unwrap(); - let ready = read_frame(&mut reader).await.unwrap(); - assert!(matches!( - ready, - Frame { - message: Message::Ready { - protocol_version: PROTOCOL_VERSION, - ref incarnation, - }, - ref payload, - } if incarnation == "test-incarnation" && payload.is_empty() - )); - (HostPeer { reader, writer }, task) -} - -fn middleware_contributions() -> Contributions { - Contributions::from([( - Capability::Middleware, - ContributionDeclaration { - id: "test.example.middleware".into(), - version: 4, - stages: vec![Stage::Request], - input_formats: vec!["openai".into()], - output_formats: vec!["openai".into()], - }, - )]) -} - -fn middleware_plugin(map_response: bool) -> impl PluginHandler { - MiddlewarePlugin::new( - &middleware_contributions(), - move |call: RequestCall| async move { - let RequestCall { - mut request, next, .. - } = call; - if map_response { - request.body.push(b'!'); - request.remove_header("x-old"); - request.head.headers.push(MiddlewareHeader { - name: "x-direct".into(), - value: b"request".to_vec(), - }); - } - let mut response = next.run(request).await?; - if map_response { - response.remove_header("x-upstream"); - response.headers.push(MiddlewareHeader { - name: "x-direct-response".into(), - value: b"response".to_vec(), - }); - response.body = response.body.map_frames(|mut frame| { - frame.payload.make_ascii_uppercase(); - Ok(vec![frame]) - })?; - } - Ok(response) - }, - ) - .unwrap() -} - -#[test] -fn middleware_plugin_rejects_declarations_without_matching_handlers() { - let declaration = middleware_contributions() - .remove(&Capability::Middleware) - .unwrap(); - let mut invalid = vec![Contributions::new()]; - invalid.push(Contributions::from([( - Capability::Scheduler, - declaration.clone(), - )])); - invalid.push(Contributions::from([ - (Capability::Middleware, declaration.clone()), - (Capability::Scheduler, declaration.clone()), - ])); - let mut unsupported = declaration.clone(); - unsupported.version = 2; - invalid.push(Contributions::from([(Capability::Middleware, unsupported)])); - for stages in [ - vec![], - vec![Stage::Registration], - vec![Stage::Request, Stage::Request], - vec![Stage::Request, Stage::Attempt, Stage::Request], - ] { - let mut invalid_stage = declaration.clone(); - invalid_stage.stages = stages; - invalid.push(Contributions::from([( - Capability::Middleware, - invalid_stage, - )])); - } - for contributes in invalid { - let result = MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { - call.next.run(call.request).await - }); - assert!(matches!(result, Err(SessionError::Configuration))); - } - let mut both_mounts = declaration; - both_mounts.stages.push(Stage::Attempt); - let contributes = Contributions::from([(Capability::Middleware, both_mounts)]); - assert!( - MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { - call.next.run(call.request).await - }) - .is_ok() - ); -} - -#[tokio::test] -async fn middleware_plugin_registers_without_invoking_business_handler() { - let contributes = middleware_contributions(); - let plugin = MiddlewarePlugin::new(&contributes, |_: RequestCall| async { - panic!("registration must not invoke middleware") - }) - .unwrap(); - let (mut host, task) = start_session(plugin).await; - let mut registration_context = context(1, Duration::from_secs(1)); - registration_context.stage = Stage::Registration; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id: 1, - method: "plugin.register".into(), - context: registration_context, - params: json!({}), - }, - payload: vec![], - }, - ) - .await - .unwrap(); - let reply = receive(&mut host).await; - let Message::Result { id: 1, result } = reply.message else { - panic!("expected registration result") - }; - let registration: gateway_plugin_sdk::call::registration::Registration = - serde_json::from_value(result).unwrap(); - assert_eq!(registration.contributes, contributes); - assert!(reply.payload.is_empty()); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn middleware_plugin_rejects_invalid_calls_before_business_dispatch() { - let plugin = MiddlewarePlugin::new(&middleware_contributions(), |_: RequestCall| async { - panic!("invalid call must not invoke middleware") - }) - .unwrap(); - let (mut host, task) = start_session(plugin).await; - let cases = [ - ("plugin.register", Stage::Request, json!({}), vec![]), - ("plugin.register", Stage::Registration, json!(null), vec![]), - ( - "plugin.register", - Stage::Registration, - json!({"extra":true}), - vec![], - ), - ("plugin.register", Stage::Registration, json!({}), vec![1]), - (HANDLE_METHOD, Stage::Attempt, json!({}), vec![]), - (HANDLE_METHOD, Stage::Request, json!({}), vec![]), - ("unknown.execute", Stage::Request, json!({}), vec![]), - ]; - for (index, (method, stage, params, payload)) in cases.into_iter().enumerate() { - let id = u64::try_from(index).unwrap() * 2 + 1; - let mut call_context = middleware_context(id); - call_context.stage = stage; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id, - method: method.into(), - context: call_context, - params, - }, - payload, - }, - ) - .await - .unwrap(); - let reply = receive(&mut host).await; - let Message::Error { - id: reply_id, - error, - } = reply.message - else { - panic!("expected invalid call error") - }; - assert_eq!(reply_id, id); - assert_eq!( - error.code, - if method == "unknown.execute" { - ErrorCode::Unsupported - } else { - ErrorCode::InvalidInput - } - ); - } - shutdown(&mut host, task).await; -} - -fn handshake() -> Handshake { - Handshake { - protocol_version: PROTOCOL_VERSION, - artifact_sha256: "a".repeat(64), - plugin_id: "test.example".into(), - instance_id: "test-instance".into(), - generation: 7, - incarnation: "test-incarnation".into(), - configuration: json!({}), - - contributes: Contributions::new(), - } -} - -pub(super) fn context(id: u64, timeout: Duration) -> CallContext { - CallContext { - call_id: id, - instance_id: "test-instance".into(), - generation: 7, - incarnation: "test-incarnation".into(), - stage: Stage::Request, - timeout_ms: timeout.as_millis() as u64, - resource_stream: false, - resource_scope_id: format!("scope-{id}"), - request_id: None, - attempt_id: None, - account_id: None, - credential_revision: None, - } -} - -fn middleware_context(id: u64) -> CallContext { - CallContext { - stage: Stage::Request, - request_id: Some("request-middleware".into()), - ..context(id, Duration::from_secs(1)) - } -} - -pub(super) fn middleware_request() -> MiddlewareRequestHead { - MiddlewareRequestHead { - settings_sources: serde_json::Value::Null, - settings: serde_json::Value::Null, - client_key_id: "fixture-key".into(), - account_group_ids: vec!["fixture-group".into()], - request_id: "request-middleware".into(), - mount: MiddlewareMount::Request, - attempt_index: None, - operation: "generate".into(), - protocol: "openai".into(), - endpoint: "responses".into(), - transport: MiddlewareTransport::HttpSse, - provider: None, - model: Some("public-model".into()), - account_id: None, - headers: vec![MiddlewareHeader { - name: "x-old".into(), - value: b"old".to_vec(), - }], - } -} - -pub(super) async fn send_call( - host: &mut HostPeer, - id: u64, - method: &str, - params: Value, - payload: Vec, -) { - send_call_with_timeout(host, id, method, params, payload, Duration::from_secs(1)).await; -} - -async fn send_call_with_timeout( - host: &mut HostPeer, - id: u64, - method: &str, - params: Value, - payload: Vec, - timeout: Duration, -) { - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id, - method: method.into(), - context: context(id, timeout), - params, - }, - payload, - }, - ) - .await - .unwrap(); -} - -async fn send_control(host: &mut HostPeer, message: Message) { - write_frame(&mut host.writer, &Frame::control(message)) - .await - .unwrap(); -} - -pub(super) async fn receive(host: &mut HostPeer) -> Frame { - tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) - .await - .expect("plugin response timed out") - .expect("plugin response frame must be valid") -} - -pub(super) async fn shutdown(host: &mut HostPeer, task: JoinHandle>) { - send_control(host, Message::Shutdown).await; - tokio::time::timeout(Duration::from_secs(1), task) - .await - .expect("plugin session did not shut down") - .expect("plugin session task panicked") - .expect("plugin session shutdown failed"); -} - -#[tokio::test] -async fn typed_budget_helpers_preserve_wire_contract_and_propagate_failure_without_retry() { - for failed in [false, true] { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "reset_budget", json!({}), vec![]).await; - let list = receive(&mut host).await; - let Message::Callback { - id, - parent_id: 1, - method, - params, - } = list.message - else { - panic!("expected Key list callback") - }; - assert_eq!(method, "host.keys.list"); - assert_eq!(params, json!({"cursor":null,"limit":10})); - assert!(list.payload.is_empty()); - send_control(&mut host, Message::Result { id, result: json!({"keys":[{"id":"key_1","name":"test","enabled":true}],"next_cursor":null}) }).await; - let reset = receive(&mut host).await; - let Message::Callback { - id, - parent_id: 1, - method, - params, - } = reset.message - else { - panic!("expected budget reset callback") - }; - assert_eq!(method, "host.keys.reset_budget"); - assert_eq!(params, json!({})); - assert_eq!( - serde_json::from_slice::(&reset.payload).unwrap(), - json!({"client_key_id":"key_1","period":"weekly"}) - ); - if failed { - send_control( - &mut host, - Message::Error { - id, - error: PluginFault::new(ErrorCode::Conflict, "stale instance"), - }, - ) - .await; - let Message::Error { id: 1, error } = receive(&mut host).await.message else { - panic!("expected error propagated without another callback") - }; - assert_eq!(error.code, ErrorCode::Conflict); - } else { - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id, - result: json!({}), - }, - payload: serde_json::to_vec(&json!({"client_key_id":"key_1"})).unwrap(), - }, - ) - .await - .unwrap(); - let Message::Result { id: 1, result } = receive(&mut host).await.message else { - panic!("expected successful reset") - }; - assert_eq!(result, json!({"client_key_id":"key_1"})); - } - shutdown(&mut host, task).await; - } -} - -#[tokio::test] -async fn independent_calls_complete_concurrently() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "slow", json!("slow"), vec![1]).await; - send_call(&mut host, 3, "echo", json!("fast"), vec![3]).await; - - let fast = receive(&mut host).await; - let slow = receive(&mut host).await; - assert!(matches!( - fast, - Frame { - message: Message::Result { id: 3, result }, - payload, - } if result == json!("fast") && payload == [3] - )); - assert!(matches!( - slow, - Frame { - message: Message::Result { id: 1, result }, - payload, - } if result == json!("slow") && payload == [1] - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn completed_call_does_not_discard_a_partially_read_next_frame() { - let next = Frame { - message: Message::Call { - id: 3, - method: "echo".into(), - context: context(3, Duration::from_secs(1)), - params: json!({"fragmented": true}), - }, - payload: vec![0, 255, 13, 10, 128], - }; - let mut encoded = Vec::new(); - write_frame(&mut encoded, &next).await.unwrap(); - let payload_start = encoded.len() - next.payload.len(); - for split in [2, 6, 10, payload_start + 2] { - tokio::time::timeout(Duration::from_secs(2), async { - // 单字节缓冲保证前缀已经进入读取 future，而不是仍滞留在传输队列 - let (mut host, task) = start_session_with_capacity(TestHandler::default(), 1).await; - send_call(&mut host, 1, "slow", json!("first"), Vec::new()).await; - host.writer.write_all(&encoded[..split]).await.unwrap(); - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - host.writer.write_all(&encoded[split..]).await.unwrap(); - assert_eq!( - receive(&mut host).await, - Frame { - message: Message::Result { - id: 3, - result: json!({"fragmented": true}) - }, - payload: next.payload.clone(), - } - ); - shutdown(&mut host, task).await; - }) - .await - .expect("fragmented frame must survive concurrent call completion"); - } -} - -#[tokio::test] -async fn host_callback_round_trips_while_parent_call_is_waiting() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "callback", json!({"event": "ready"}), vec![7]).await; - - let callback = receive(&mut host).await; - let callback_id = match callback.message { - Message::Callback { - id, - parent_id: 1, - ref method, - ref params, - } if id.is_multiple_of(2) - && method == "host.log" - && params == &json!({"event": "ready"}) => - { - id - } - message => panic!("unexpected callback frame: {message:?}"), - }; - assert_eq!(callback.payload, [7]); - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: callback_id, - result: json!({"recorded": true}), - }, - payload: vec![9], - }, - ) - .await - .unwrap(); - - let result = receive(&mut host).await; - assert!(matches!( - result, - Frame { - message: Message::Result { id: 1, result }, - payload, - } if result == json!({"recorded": true}) && payload == [9] - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn callbacks_from_concurrent_calls_remain_correlated() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "callback", json!({"parent": 1}), vec![1]).await; - send_call(&mut host, 3, "callback", json!({"parent": 3}), vec![3]).await; - - let mut callbacks = Vec::new(); - for _ in 0..2 { - let frame = receive(&mut host).await; - let Message::Callback { - id, - parent_id, - method, - params, - } = frame.message - else { - panic!("expected callback") - }; - assert_eq!(method, "host.log"); - assert_eq!(params, json!({"parent": parent_id})); - assert_eq!(frame.payload, [parent_id as u8]); - callbacks.push((id, parent_id)); - } - assert!(callbacks[0].0 < callbacks[1].0); - assert!(callbacks.iter().all(|(id, _)| id.is_multiple_of(2))); - - for (callback_id, parent_id) in callbacks.into_iter().rev() { - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: callback_id, - result: json!({"parent": parent_id}), - }, - payload: vec![parent_id as u8], - }, - ) - .await - .unwrap(); - } - let mut completed = Vec::new(); - for _ in 0..2 { - let frame = receive(&mut host).await; - let Message::Result { id, result } = frame.message else { - panic!("expected parent result") - }; - assert_eq!(result, json!({"parent": id})); - assert_eq!(frame.payload, [id as u8]); - completed.push(id); - } - completed.sort_unstable(); - assert_eq!(completed, [1, 3]); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn middleware_next_and_body_mapping_reuse_callback_and_credit_flow() { - let (mut host, task) = start_session(middleware_plugin(true)).await; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id: 1, - method: HANDLE_METHOD.into(), - context: middleware_context(1), - params: serde_json::to_value(middleware_request()).unwrap(), - }, - payload: b"request".to_vec(), - }, - ) - .await - .unwrap(); - - let next = receive(&mut host).await; - let next_id = match next.message { - Message::Callback { - id, - parent_id: 1, - ref method, - ref params, - } if method == NEXT_METHOD => { - let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); - assert_eq!(request.body, MiddlewareRequestBody::Replace); - assert!(request.protocol.is_none()); - assert_eq!(request.header_mutations.len(), 2); - assert!(matches!( - &request.header_mutations[0], - MiddlewareHeaderMutation::Remove { name } if name == "x-old" - )); - assert!(matches!( - &request.header_mutations[1], - MiddlewareHeaderMutation::Append { name, value } - if name == "x-direct" && value == b"request" - )); - id - } - message => panic!("unexpected middleware next frame: {message:?}"), - }; - assert_eq!(next.payload, b"request!"); - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: next_id, - result: serde_json::to_value(MiddlewareNextResponse { - metadata: None, - response: "response-1".into(), - protocol: "openai".into(), - status: 200, - headers: vec![MiddlewareHeader { - name: "x-upstream".into(), - value: b"old".to_vec(), - }], - body: Some(MiddlewareBodyHandle { - handle: "body-1".into(), - framing: MiddlewareBodyFraming::SseEvent, - }), - }) - .unwrap(), - }, - payload: Vec::new(), - }, - ) - .await - .unwrap(); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1_024, - frames: 2, - }, - ) - .await; - - let initial = receive(&mut host).await; - let Message::Result { id: 1, result } = initial.message else { - panic!("expected middleware result") - }; - let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); - assert_eq!(head.header_mutations.len(), 2); - assert!(matches!( - &head.header_mutations[0], - MiddlewareHeaderMutation::Remove { name } if name == "x-upstream" - )); - assert!(matches!( - &head.header_mutations[1], - MiddlewareHeaderMutation::Append { name, value } - if name == "x-direct-response" && value == b"response" - )); - assert!(matches!( - head.body, - MiddlewareResponseBody::Stream { - framing: MiddlewareBodyFraming::SseEvent - } - )); - - let read = receive(&mut host).await; - let read_id = match read.message { - Message::Callback { - id, - parent_id: 1, - ref method, - .. - } if method == BODY_READ_METHOD => id, - message => panic!("unexpected middleware body read: {message:?}"), - }; - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: read_id, - result: serde_json::to_value(MiddlewareBodyReadResult { - framing: MiddlewareBodyFraming::SseEvent, - source_id: 1, - eof: false, - terminal: true, - }) - .unwrap(), - }, - payload: b"data: done\n\n".to_vec(), - }, - ) - .await - .unwrap(); - let mapped = receive(&mut host).await; - let Frame { - message: Message::Stream { id: 1, sequence: 0 }, - payload, - } = mapped - else { - panic!("expected mapped middleware frame") - }; - let mapped = MiddlewareBodyFrame::decode(&payload).unwrap(); - assert_eq!(mapped.payload, b"DATA: DONE\n\n"); - assert!(!mapped.terminal); - assert_eq!(mapped.source_id(), 1); - assert_eq!(mapped.disposition(), MiddlewareBodyDisposition::Only); - - let eof = receive(&mut host).await; - let eof_id = match eof.message { - Message::Callback { - id, - parent_id: 1, - ref method, - .. - } if method == BODY_READ_METHOD => id, - message => panic!("unexpected middleware eof read: {message:?}"), - }; - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: eof_id, - result: serde_json::to_value(MiddlewareBodyReadResult { - framing: MiddlewareBodyFraming::SseEvent, - source_id: 0, - eof: true, - terminal: false, - }) - .unwrap(), - }, - payload: Vec::new(), - }, - ) - .await - .unwrap(); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn untouched_middleware_response_transfers_opaque_body_without_reading_it() { - let (mut host, task) = start_session(middleware_plugin(false)).await; - write_frame( - &mut host.writer, - &Frame { - message: Message::Call { - id: 1, - method: HANDLE_METHOD.into(), - context: middleware_context(1), - params: serde_json::to_value(middleware_request()).unwrap(), - }, - payload: b"hidden-by-preserve".to_vec(), - }, - ) - .await - .unwrap(); - - let next = receive(&mut host).await; - let next_id = match next.message { - Message::Callback { - id, - parent_id: 1, - ref method, - ref params, - } if method == NEXT_METHOD => { - let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); - assert_eq!(request.body, MiddlewareRequestBody::Preserve); - id - } - message => panic!("unexpected middleware next frame: {message:?}"), - }; - assert!(next.payload.is_empty()); - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id: next_id, - result: serde_json::to_value(MiddlewareNextResponse { - metadata: None, - response: "response-1".into(), - protocol: "openai".into(), - status: 200, - headers: Vec::new(), - body: Some(MiddlewareBodyHandle { - handle: "body-pass-through".into(), - framing: MiddlewareBodyFraming::RawBytes, - }), - }) - .unwrap(), - }, - payload: Vec::new(), - }, - ) - .await - .unwrap(); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - - let initial = receive(&mut host).await; - let Message::Result { id: 1, result } = initial.message else { - panic!("expected middleware result") - }; - let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); - assert!(matches!( - head.body, - MiddlewareResponseBody::PassThrough { - body: MiddlewareBodyHandle { ref handle, .. } - } if handle == "body-pass-through" - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn stream_waits_for_credit_and_rejects_uncreditable_buffered_chunk_before_result() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "stream", json!({}), Vec::new()).await; - assert!( - tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) - .await - .is_err() - ); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 4, - frames: 1, - }, - ) - .await; - - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - assert!(matches!( - receive(&mut host).await, - Frame { - message: Message::Stream { id: 1, sequence: 0 }, - payload, - } if payload == b"abcd" - )); - assert!( - tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) - .await - .is_err() - ); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 4, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await, - Frame { - message: Message::Stream { id: 1, sequence: 1 }, - payload, - } if payload == b"efgh" - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - - send_call(&mut host, 3, "oversized_stream", json!({}), Vec::new()).await; - send_control( - &mut host, - Message::Credit { - id: 3, - bytes: 4, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 3, - error: PluginFault { - code: ErrorCode::InvalidInput, - .. - }, - } - )); - - send_call(&mut host, 5, "too_many_chunks", json!({}), Vec::new()).await; - send_control( - &mut host, - Message::Credit { - id: 5, - bytes: 4, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 5, - error: PluginFault { - code: ErrorCode::Capacity, - .. - }, - } - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn late_credit_after_a_long_stream_terminal_does_not_close_the_session() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "long_stream", json!({}), Vec::new()).await; - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - for sequence in 0..16 { - assert!(matches!( - receive(&mut host).await, - Frame { - message: Message::Stream { id: 1, sequence: observed }, - payload, - } if observed == sequence && payload == b"x" - )); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - } - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - - send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 3, result } if result == json!("still-ready") - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn malformed_late_credit_remains_a_protocol_error() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 0, - frames: 1, - }, - ) - .await; - assert!(matches!( - tokio::time::timeout(Duration::from_secs(1), task) - .await - .expect("protocol error did not close the session") - .expect("plugin session task panicked"), - Err(SessionError::Protocol) - )); -} - -#[tokio::test] -async fn dynamic_stream_ends_with_one_error_when_a_later_chunk_cannot_fit_credit() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call( - &mut host, - 1, - "dynamic_oversized_stream", - json!({}), - Vec::new(), - ) - .await; - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 4, - frames: 1, - }, - ) - .await; - - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { - id: 1, - error: Some(PluginFault { - code: ErrorCode::InvalidInput, - .. - }), - } - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn oversized_handler_faults_are_replaced_without_closing_the_session() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "oversized_result", json!({}), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 1, - error: PluginFault { - code: ErrorCode::InvalidInput, - .. - }, - } - )); - - send_call(&mut host, 3, "oversized_fault", json!({}), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 3, - error: PluginFault { - code: ErrorCode::Fault, - ref message, - .. - }, - } if message == "plugin error metadata exceeds its limit" - )); - - send_call( - &mut host, - 5, - "dynamic_oversized_fault", - json!({}), - Vec::new(), - ) - .await; - send_control( - &mut host, - Message::Credit { - id: 5, - bytes: 1_024, - frames: 1, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 5, .. } - )); - assert!(matches!( - receive(&mut host).await.message, - Message::End { - id: 5, - error: Some(PluginFault { - code: ErrorCode::Fault, - ref message, - .. - }), - } if message == "plugin stream error metadata exceeds its limit" - )); - - send_call(&mut host, 7, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 7, result } if result == json!("still-ready") - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn call_deadline_cancels_handler_and_keeps_the_session_available() { - let handler = TestHandler::default(); - let state = Arc::clone(&handler.lifecycle); - let (mut host, task) = start_session(handler).await; - send_call_with_timeout( - &mut host, - 1, - "pending", - json!({}), - Vec::new(), - Duration::from_millis(20), - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 1, - error: PluginFault { - code: ErrorCode::Timeout, - .. - }, - } - )); - tokio::time::timeout(Duration::from_millis(100), async { - while !state.cancel_started.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("deadline did not notify the lifecycle hook"); - - send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 3, result } if result == json!("still-ready") - )); - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn cancelling_parent_retires_callback_and_late_reply_does_not_close_session() { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, "callback", json!({}), Vec::new()).await; - let callback = receive(&mut host).await; - let Message::Callback { - id: callback_id, .. - } = callback.message - else { - panic!("expected callback") - }; - - send_control(&mut host, Message::Cancel { id: 1 }).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Cancelled { id: 1 } - )); - send_control( - &mut host, - Message::Result { - id: callback_id, - result: json!({}), - }, - ) - .await; - send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 3, result } if result == json!("still-ready") - )); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn blocking_cancel_hook_cannot_delay_cancelled_or_eof_settlement() { - let handler = TestHandler { - cancel_delay: Some(Duration::from_millis(200)), - ..TestHandler::default() - }; - let state = Arc::clone(&handler.lifecycle); - let (mut host, task) = start_session(handler).await; - send_call(&mut host, 1, "pending", json!({}), Vec::new()).await; - send_control(&mut host, Message::Cancel { id: 1 }).await; - - let cancelled = tokio::time::timeout(Duration::from_millis(50), receive(&mut host)) - .await - .expect("blocking lifecycle hook delayed Cancelled"); - assert!(matches!(cancelled.message, Message::Cancelled { id: 1 })); - tokio::time::timeout(Duration::from_millis(50), async { - while !state.cancel_started.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("cancel hook never started"); - - host.writer.shutdown().await.unwrap(); - let result = tokio::time::timeout(Duration::from_millis(50), task) - .await - .expect("blocking lifecycle hook delayed EOF settlement") - .expect("plugin session task panicked"); - assert!(matches!(result, Err(SessionError::Closed))); -} - -#[tokio::test] -async fn quiesce_rejects_new_calls_and_shutdown_closes_the_session() { - let handler = TestHandler::default(); - let state = Arc::clone(&handler.lifecycle); - let (mut host, task) = start_session(handler).await; - send_control(&mut host, Message::Quiesce).await; - send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; - - assert!(matches!( - receive(&mut host).await.message, - Message::Error { - id: 1, - error: PluginFault { - code: ErrorCode::Capacity, - .. - }, - } - )); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 1, - frames: 1, - }, - ) - .await; - tokio::time::timeout(Duration::from_millis(50), async { - while !state.quiesced.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("quiesce hook was not signalled"); - - shutdown(&mut host, task).await; -} - -#[tokio::test] -async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() { - for (entry, method, request, response) in [ - ( - "key_facts", - "host.data.keys.get", - json!({"client_key_id":"key_1"}), - json!({"schema_version":1,"client_key_id":"key_1","enabled":false,"group_ids":["grp_1"]}), - ), - ( - "get_budget", - "host.keys.get_budget", - json!({"client_key_id":"key_1"}), - json!({ - "client_key_id":"key_1", "daily_limit_usd":"10", "weekly_limit_usd":"20", - "daily_used_usd":"1.25", "weekly_used_usd":"2.5", "daily_resets_at_ms":null, "weekly_resets_at_ms":null, - }), - ), - ( - "update_budget_limits", - "host.keys.update_budget_limits", - json!({"client_key_id":"key_1","weekly_limit_usd":"12.5"}), - json!({"client_key_id":"key_1"}), - ), - ( - "refresh_quota", - "host.quota_observations.refresh", - json!({"account_id":"acct_1"}), - json!({"schema_version":1,"account_id":"acct_1","observed_at_ms":null,"windows":[]}), - ), - ] { - for failed in [false, true] { - let (mut host, task) = start_session(TestHandler::default()).await; - send_call(&mut host, 1, entry, json!({}), vec![]).await; - let callback = receive(&mut host).await; - let Message::Callback { - id, - parent_id: 1, - method: actual, - params, - } = callback.message - else { - panic!("expected callback"); - }; - assert_eq!(actual, method); - assert_eq!(params, json!({})); - assert_eq!( - serde_json::from_slice::(&callback.payload).unwrap(), - request - ); - if failed { - send_control( - &mut host, - Message::Error { - id, - error: PluginFault::new(ErrorCode::Conflict, "changed"), - }, - ) - .await; - let Message::Error { id: 1, error } = receive(&mut host).await.message else { - panic!("expected propagated error without retry"); - }; - assert_eq!(error.code, ErrorCode::Conflict); - } else { - let mut response_payload = response.clone(); - if matches!(entry, "key_facts" | "refresh_quota") { - response_payload["future_fact"] = json!({"value":1}); - } - write_frame( - &mut host.writer, - &Frame { - message: Message::Result { - id, - result: json!({}), - }, - payload: serde_json::to_vec(&response_payload).unwrap(), - }, - ) - .await - .unwrap(); - let Message::Result { id: 1, result } = receive(&mut host).await.message else { - panic!("expected projected response"); - }; - assert_eq!(result, response); - } - shutdown(&mut host, task).await; - } - } -} - -struct ResourceStream(Option); -impl gateway_plugin_sdk::client::PullResponseStream for ResourceStream { - fn next(&mut self) -> gateway_plugin_sdk::client::PullResponseFuture<'_> { - Box::pin(async move { - let host = self.0.take()?; - tokio::time::sleep(Duration::from_millis(700)).await; - Some( - host.call("host.log", json!({}), Vec::new()) - .await - .map(|_| b"live".to_vec()) - .map_err(SessionError::into_plugin_fault), - ) - }) - } -} - -#[tokio::test] -async fn resource_stream_keeps_callbacks_and_credit_alive_after_the_initial_deadline() { - let (mut host, task) = start_session(TestHandler::default()).await; - let mut context = context(1, Duration::from_millis(500)); - context.resource_stream = true; - write_frame( - &mut host.writer, - &Frame::control(Message::Call { - id: 1, - method: "resource_stream".into(), - context, - params: json!({}), - }), - ) - .await - .unwrap(); - send_control( - &mut host, - Message::Credit { - id: 1, - bytes: 16, - frames: 2, - }, - ) - .await; - assert!(matches!( - receive(&mut host).await.message, - Message::Result { id: 1, .. } - )); - let Message::Callback { - id, parent_id: 1, .. - } = receive(&mut host).await.message - else { - panic!("stream callback must survive initial timeout"); - }; - send_control( - &mut host, - Message::Result { - id, - result: json!({}), - }, - ) - .await; - assert!( - matches!(receive(&mut host).await, Frame { message: Message::Stream { id: 1, sequence: 0 }, payload } if payload == b"live") - ); - assert!(matches!( - receive(&mut host).await.message, - Message::End { id: 1, error: None } - )); - shutdown(&mut host, task).await; -} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs new file mode 100644 index 000000000..72af4f63e --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/credit.rs @@ -0,0 +1,271 @@ +//! 流信用等待、迟到信用、分块上限与资源流期限测试 + +use super::*; + +#[tokio::test] +async fn stream_waits_for_credit_and_rejects_uncreditable_buffered_chunk_before_result() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "stream", json!({}), Vec::new()).await; + assert!( + tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) + .await + .is_err() + ); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 4, + frames: 1, + }, + ) + .await; + + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + assert!(matches!( + receive(&mut host).await, + Frame { + message: Message::Stream { id: 1, sequence: 0 }, + payload, + } if payload == b"abcd" + )); + assert!( + tokio::time::timeout(Duration::from_millis(25), read_frame(&mut host.reader)) + .await + .is_err() + ); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 4, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await, + Frame { + message: Message::Stream { id: 1, sequence: 1 }, + payload, + } if payload == b"efgh" + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + + send_call(&mut host, 3, "oversized_stream", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 3, + bytes: 4, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 3, + error: PluginFault { + code: ErrorCode::InvalidInput, + .. + }, + } + )); + + send_call(&mut host, 5, "too_many_chunks", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 5, + bytes: 4, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 5, + error: PluginFault { + code: ErrorCode::Capacity, + .. + }, + } + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn late_credit_after_a_long_stream_terminal_does_not_close_the_session() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "long_stream", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + for sequence in 0..16 { + assert!(matches!( + receive(&mut host).await, + Frame { + message: Message::Stream { id: 1, sequence: observed }, + payload, + } if observed == sequence && payload == b"x" + )); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + } + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + + send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 3, result } if result == json!("still-ready") + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn malformed_late_credit_remains_a_protocol_error() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 0, + frames: 1, + }, + ) + .await; + assert!(matches!( + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("protocol error did not close the session") + .expect("plugin session task panicked"), + Err(SessionError::Protocol) + )); +} + +#[tokio::test] +async fn dynamic_stream_ends_with_one_error_when_a_later_chunk_cannot_fit_credit() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call( + &mut host, + 1, + "dynamic_oversized_stream", + json!({}), + Vec::new(), + ) + .await; + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 4, + frames: 1, + }, + ) + .await; + + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { + id: 1, + error: Some(PluginFault { + code: ErrorCode::InvalidInput, + .. + }), + } + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn resource_stream_keeps_callbacks_and_credit_alive_after_the_initial_deadline() { + let (mut host, task) = start_session(TestHandler::default()).await; + let mut context = context(1, Duration::from_millis(500)); + context.resource_stream = true; + write_frame( + &mut host.writer, + &Frame::control(Message::Call { + id: 1, + method: "resource_stream".into(), + context, + params: json!({}), + }), + ) + .await + .unwrap(); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 16, + frames: 2, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + let Message::Callback { + id, parent_id: 1, .. + } = receive(&mut host).await.message + else { + panic!("stream callback must survive initial timeout"); + }; + send_control( + &mut host, + Message::Result { + id, + result: json!({}), + }, + ) + .await; + assert!( + matches!(receive(&mut host).await, Frame { message: Message::Stream { id: 1, sequence: 0 }, payload } if payload == b"live") + ); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + shutdown(&mut host, task).await; +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs new file mode 100644 index 000000000..06a3acff6 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs @@ -0,0 +1,545 @@ +//! 会话调用分派、并发帧读取与处理器响应边界测试 + +use super::*; + +#[test] +fn middleware_plugin_rejects_declarations_without_matching_handlers() { + let declaration = middleware_contributions() + .remove(&Capability::Middleware) + .unwrap(); + let mut invalid = vec![Contributions::new()]; + invalid.push(Contributions::from([( + Capability::Scheduler, + declaration.clone(), + )])); + invalid.push(Contributions::from([ + (Capability::Middleware, declaration.clone()), + (Capability::Scheduler, declaration.clone()), + ])); + let mut unsupported = declaration.clone(); + unsupported.version = 2; + invalid.push(Contributions::from([(Capability::Middleware, unsupported)])); + for stages in [ + vec![], + vec![Stage::Registration], + vec![Stage::Request, Stage::Request], + vec![Stage::Request, Stage::Attempt, Stage::Request], + ] { + let mut invalid_stage = declaration.clone(); + invalid_stage.stages = stages; + invalid.push(Contributions::from([( + Capability::Middleware, + invalid_stage, + )])); + } + for contributes in invalid { + let result = MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { + call.next.run(call.request).await + }); + assert!(matches!(result, Err(SessionError::Configuration))); + } + let mut both_mounts = declaration; + both_mounts.stages.push(Stage::Attempt); + let contributes = Contributions::from([(Capability::Middleware, both_mounts)]); + assert!( + MiddlewarePlugin::new(&contributes, |call: RequestCall| async move { + call.next.run(call.request).await + }) + .is_ok() + ); +} + +#[tokio::test] +async fn middleware_plugin_registers_without_invoking_business_handler() { + let contributes = middleware_contributions(); + let plugin = MiddlewarePlugin::new(&contributes, |_: RequestCall| async { + panic!("registration must not invoke middleware") + }) + .unwrap(); + let (mut host, task) = start_session(plugin).await; + let mut registration_context = context(1, Duration::from_secs(1)); + registration_context.stage = Stage::Registration; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id: 1, + method: "plugin.register".into(), + context: registration_context, + params: json!({}), + }, + payload: vec![], + }, + ) + .await + .unwrap(); + let reply = receive(&mut host).await; + let Message::Result { id: 1, result } = reply.message else { + panic!("expected registration result") + }; + let registration: gateway_plugin_sdk::call::registration::Registration = + serde_json::from_value(result).unwrap(); + assert_eq!(registration.contributes, contributes); + assert!(reply.payload.is_empty()); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn middleware_plugin_rejects_invalid_calls_before_business_dispatch() { + let plugin = MiddlewarePlugin::new(&middleware_contributions(), |_: RequestCall| async { + panic!("invalid call must not invoke middleware") + }) + .unwrap(); + let (mut host, task) = start_session(plugin).await; + let cases = [ + ("plugin.register", Stage::Request, json!({}), vec![]), + ("plugin.register", Stage::Registration, json!(null), vec![]), + ( + "plugin.register", + Stage::Registration, + json!({"extra":true}), + vec![], + ), + ("plugin.register", Stage::Registration, json!({}), vec![1]), + (HANDLE_METHOD, Stage::Attempt, json!({}), vec![]), + (HANDLE_METHOD, Stage::Request, json!({}), vec![]), + ("unknown.execute", Stage::Request, json!({}), vec![]), + ]; + for (index, (method, stage, params, payload)) in cases.into_iter().enumerate() { + let id = u64::try_from(index).unwrap() * 2 + 1; + let mut call_context = middleware_context(id); + call_context.stage = stage; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id, + method: method.into(), + context: call_context, + params, + }, + payload, + }, + ) + .await + .unwrap(); + let reply = receive(&mut host).await; + let Message::Error { + id: reply_id, + error, + } = reply.message + else { + panic!("expected invalid call error") + }; + assert_eq!(reply_id, id); + assert_eq!( + error.code, + if method == "unknown.execute" { + ErrorCode::Unsupported + } else { + ErrorCode::InvalidInput + } + ); + } + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn independent_calls_complete_concurrently() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "slow", json!("slow"), vec![1]).await; + send_call(&mut host, 3, "echo", json!("fast"), vec![3]).await; + + let fast = receive(&mut host).await; + let slow = receive(&mut host).await; + assert!(matches!( + fast, + Frame { + message: Message::Result { id: 3, result }, + payload, + } if result == json!("fast") && payload == [3] + )); + assert!(matches!( + slow, + Frame { + message: Message::Result { id: 1, result }, + payload, + } if result == json!("slow") && payload == [1] + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn completed_call_does_not_discard_a_partially_read_next_frame() { + let next = Frame { + message: Message::Call { + id: 3, + method: "echo".into(), + context: context(3, Duration::from_secs(1)), + params: json!({"fragmented": true}), + }, + payload: vec![0, 255, 13, 10, 128], + }; + let mut encoded = Vec::new(); + write_frame(&mut encoded, &next).await.unwrap(); + let payload_start = encoded.len() - next.payload.len(); + for split in [2, 6, 10, payload_start + 2] { + tokio::time::timeout(Duration::from_secs(2), async { + // 单字节缓冲保证前缀已经进入读取 future，而不是仍滞留在传输队列 + let (mut host, task) = start_session_with_capacity(TestHandler::default(), 1).await; + send_call(&mut host, 1, "slow", json!("first"), Vec::new()).await; + host.writer.write_all(&encoded[..split]).await.unwrap(); + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 1, .. } + )); + host.writer.write_all(&encoded[split..]).await.unwrap(); + assert_eq!( + receive(&mut host).await, + Frame { + message: Message::Result { + id: 3, + result: json!({"fragmented": true}) + }, + payload: next.payload.clone(), + } + ); + shutdown(&mut host, task).await; + }) + .await + .expect("fragmented frame must survive concurrent call completion"); + } +} + +#[tokio::test] +async fn middleware_next_and_body_mapping_reuse_callback_and_credit_flow() { + let (mut host, task) = start_session(middleware_plugin(true)).await; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id: 1, + method: HANDLE_METHOD.into(), + context: middleware_context(1), + params: serde_json::to_value(middleware_request()).unwrap(), + }, + payload: b"request".to_vec(), + }, + ) + .await + .unwrap(); + + let next = receive(&mut host).await; + let next_id = match next.message { + Message::Callback { + id, + parent_id: 1, + ref method, + ref params, + } if method == NEXT_METHOD => { + let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); + assert_eq!(request.body, MiddlewareRequestBody::Replace); + assert!(request.protocol.is_none()); + assert_eq!(request.header_mutations.len(), 2); + assert!(matches!( + &request.header_mutations[0], + MiddlewareHeaderMutation::Remove { name } if name == "x-old" + )); + assert!(matches!( + &request.header_mutations[1], + MiddlewareHeaderMutation::Append { name, value } + if name == "x-direct" && value == b"request" + )); + id + } + message => panic!("unexpected middleware next frame: {message:?}"), + }; + assert_eq!(next.payload, b"request!"); + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: next_id, + result: serde_json::to_value(MiddlewareNextResponse { + metadata: None, + response: "response-1".into(), + protocol: "openai".into(), + status: 200, + headers: vec![MiddlewareHeader { + name: "x-upstream".into(), + value: b"old".to_vec(), + }], + body: Some(MiddlewareBodyHandle { + handle: "body-1".into(), + framing: MiddlewareBodyFraming::SseEvent, + }), + }) + .unwrap(), + }, + payload: Vec::new(), + }, + ) + .await + .unwrap(); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1_024, + frames: 2, + }, + ) + .await; + + let initial = receive(&mut host).await; + let Message::Result { id: 1, result } = initial.message else { + panic!("expected middleware result") + }; + let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); + assert_eq!(head.header_mutations.len(), 2); + assert!(matches!( + &head.header_mutations[0], + MiddlewareHeaderMutation::Remove { name } if name == "x-upstream" + )); + assert!(matches!( + &head.header_mutations[1], + MiddlewareHeaderMutation::Append { name, value } + if name == "x-direct-response" && value == b"response" + )); + assert!(matches!( + head.body, + MiddlewareResponseBody::Stream { + framing: MiddlewareBodyFraming::SseEvent + } + )); + + let read = receive(&mut host).await; + let read_id = match read.message { + Message::Callback { + id, + parent_id: 1, + ref method, + .. + } if method == BODY_READ_METHOD => id, + message => panic!("unexpected middleware body read: {message:?}"), + }; + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: read_id, + result: serde_json::to_value(MiddlewareBodyReadResult { + framing: MiddlewareBodyFraming::SseEvent, + source_id: 1, + eof: false, + terminal: true, + }) + .unwrap(), + }, + payload: b"data: done\n\n".to_vec(), + }, + ) + .await + .unwrap(); + let mapped = receive(&mut host).await; + let Frame { + message: Message::Stream { id: 1, sequence: 0 }, + payload, + } = mapped + else { + panic!("expected mapped middleware frame") + }; + let mapped = MiddlewareBodyFrame::decode(&payload).unwrap(); + assert_eq!(mapped.payload, b"DATA: DONE\n\n"); + assert!(!mapped.terminal); + assert_eq!(mapped.source_id(), 1); + assert_eq!(mapped.disposition(), MiddlewareBodyDisposition::Only); + + let eof = receive(&mut host).await; + let eof_id = match eof.message { + Message::Callback { + id, + parent_id: 1, + ref method, + .. + } if method == BODY_READ_METHOD => id, + message => panic!("unexpected middleware eof read: {message:?}"), + }; + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: eof_id, + result: serde_json::to_value(MiddlewareBodyReadResult { + framing: MiddlewareBodyFraming::SseEvent, + source_id: 0, + eof: true, + terminal: false, + }) + .unwrap(), + }, + payload: Vec::new(), + }, + ) + .await + .unwrap(); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn untouched_middleware_response_transfers_opaque_body_without_reading_it() { + let (mut host, task) = start_session(middleware_plugin(false)).await; + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id: 1, + method: HANDLE_METHOD.into(), + context: middleware_context(1), + params: serde_json::to_value(middleware_request()).unwrap(), + }, + payload: b"hidden-by-preserve".to_vec(), + }, + ) + .await + .unwrap(); + + let next = receive(&mut host).await; + let next_id = match next.message { + Message::Callback { + id, + parent_id: 1, + ref method, + ref params, + } if method == NEXT_METHOD => { + let request: MiddlewareNextRequest = serde_json::from_value(params.clone()).unwrap(); + assert_eq!(request.body, MiddlewareRequestBody::Preserve); + id + } + message => panic!("unexpected middleware next frame: {message:?}"), + }; + assert!(next.payload.is_empty()); + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: next_id, + result: serde_json::to_value(MiddlewareNextResponse { + metadata: None, + response: "response-1".into(), + protocol: "openai".into(), + status: 200, + headers: Vec::new(), + body: Some(MiddlewareBodyHandle { + handle: "body-pass-through".into(), + framing: MiddlewareBodyFraming::RawBytes, + }), + }) + .unwrap(), + }, + payload: Vec::new(), + }, + ) + .await + .unwrap(); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + + let initial = receive(&mut host).await; + let Message::Result { id: 1, result } = initial.message else { + panic!("expected middleware result") + }; + let head: MiddlewareResponseHead = serde_json::from_value(result).unwrap(); + assert!(matches!( + head.body, + MiddlewareResponseBody::PassThrough { + body: MiddlewareBodyHandle { ref handle, .. } + } if handle == "body-pass-through" + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn oversized_handler_faults_are_replaced_without_closing_the_session() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "oversized_result", json!({}), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 1, + error: PluginFault { + code: ErrorCode::InvalidInput, + .. + }, + } + )); + + send_call(&mut host, 3, "oversized_fault", json!({}), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 3, + error: PluginFault { + code: ErrorCode::Fault, + ref message, + .. + }, + } if message == "plugin error metadata exceeds its limit" + )); + + send_call( + &mut host, + 5, + "dynamic_oversized_fault", + json!({}), + Vec::new(), + ) + .await; + send_control( + &mut host, + Message::Credit { + id: 5, + bytes: 1_024, + frames: 1, + }, + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 5, .. } + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { + id: 5, + error: Some(PluginFault { + code: ErrorCode::Fault, + ref message, + .. + }), + } if message == "plugin stream error metadata exceeds its limit" + )); + + send_call(&mut host, 7, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 7, result } if result == json!("still-ready") + )); + shutdown(&mut host, task).await; +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs new file mode 100644 index 000000000..edfd2e860 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/lifecycle.rs @@ -0,0 +1,112 @@ +//! 会话期限、取消钩子、排空与关闭测试 + +use super::*; + +#[tokio::test] +async fn call_deadline_cancels_handler_and_keeps_the_session_available() { + let handler = TestHandler::default(); + let state = Arc::clone(&handler.lifecycle); + let (mut host, task) = start_session(handler).await; + send_call_with_timeout( + &mut host, + 1, + "pending", + json!({}), + Vec::new(), + Duration::from_millis(20), + ) + .await; + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 1, + error: PluginFault { + code: ErrorCode::Timeout, + .. + }, + } + )); + tokio::time::timeout(Duration::from_millis(100), async { + while !state.cancel_started.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("deadline did not notify the lifecycle hook"); + + send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 3, result } if result == json!("still-ready") + )); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn blocking_cancel_hook_cannot_delay_cancelled_or_eof_settlement() { + let handler = TestHandler { + cancel_delay: Some(Duration::from_millis(200)), + ..TestHandler::default() + }; + let state = Arc::clone(&handler.lifecycle); + let (mut host, task) = start_session(handler).await; + send_call(&mut host, 1, "pending", json!({}), Vec::new()).await; + send_control(&mut host, Message::Cancel { id: 1 }).await; + + let cancelled = tokio::time::timeout(Duration::from_millis(50), receive(&mut host)) + .await + .expect("blocking lifecycle hook delayed Cancelled"); + assert!(matches!(cancelled.message, Message::Cancelled { id: 1 })); + tokio::time::timeout(Duration::from_millis(50), async { + while !state.cancel_started.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancel hook never started"); + + host.writer.shutdown().await.unwrap(); + let result = tokio::time::timeout(Duration::from_millis(50), task) + .await + .expect("blocking lifecycle hook delayed EOF settlement") + .expect("plugin session task panicked"); + assert!(matches!(result, Err(SessionError::Closed))); +} + +#[tokio::test] +async fn quiesce_rejects_new_calls_and_shutdown_closes_the_session() { + let handler = TestHandler::default(); + let state = Arc::clone(&handler.lifecycle); + let (mut host, task) = start_session(handler).await; + send_control(&mut host, Message::Quiesce).await; + send_call(&mut host, 1, "echo", json!({}), Vec::new()).await; + + assert!(matches!( + receive(&mut host).await.message, + Message::Error { + id: 1, + error: PluginFault { + code: ErrorCode::Capacity, + .. + }, + } + )); + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1, + frames: 1, + }, + ) + .await; + tokio::time::timeout(Duration::from_millis(50), async { + while !state.quiesced.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("quiesce hook was not signalled"); + + shutdown(&mut host, task).await; +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs new file mode 100644 index 000000000..5b86d9bec --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs @@ -0,0 +1,470 @@ +//! 插件会话测试入口与双工传输、处理器和宿主回调夹具 + +mod credit; +mod driver; +mod lifecycle; +mod registry; + +use std::{ + future::pending, + num::NonZeroUsize, + sync::{ + Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use gateway_plugin_sdk::{ + CallContext, Capability, ContributionDeclaration, Contributions, ErrorCode, Frame, Handshake, + Message, PROTOCOL_VERSION, PluginFault, Stage, + call::middleware::{ + BODY_READ_METHOD, HANDLE_METHOD, MiddlewareBodyDisposition, MiddlewareBodyFrame, + MiddlewareBodyFraming, MiddlewareBodyHandle, MiddlewareBodyReadResult, MiddlewareHeader, + MiddlewareHeaderMutation, MiddlewareMount, MiddlewareNextRequest, MiddlewareNextResponse, + MiddlewareRequestBody, MiddlewareRequestHead, MiddlewareResponseBody, + MiddlewareResponseHead, MiddlewareTransport, NEXT_METHOD, + }, + client::{ + CallFuture, CallReply, MiddlewarePlugin, PluginCall, PluginHandler, PluginSession, + RequestCall, ResponseStream, SessionConfig, SessionError, read_frame, write_frame, + }, +}; +use serde_json::{Value, json}; +use tokio::{ + io::{AsyncWriteExt, DuplexStream, ReadHalf, WriteHalf}, + task::JoinHandle, +}; + +const MAXIMUM_STREAM_CHUNK_BYTES: usize = 64 * 1024; + +#[derive(Clone, Default)] +struct TestHandler { + lifecycle: Arc, + cancel_delay: Option, +} + +#[derive(Default)] +struct LifecycleState { + cancel_started: AtomicBool, + cancelled: AtomicUsize, + quiesced: AtomicBool, + shutdown: AtomicBool, +} + +impl PluginHandler for TestHandler { + fn call(&self, call: PluginCall) -> CallFuture<'_> { + Box::pin(async move { + match call.method.as_str() { + "echo" => Ok(CallReply::unary(call.params, call.payload)), + "key_facts" => { + let result = call + .host + .key_facts(gateway_plugin_sdk::call::data::ClientKeyFactsQuery { + client_key_id: "key_1".into(), + }) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "get_budget" => { + let result = call + .host + .get_key_budget( + gateway_plugin_sdk::call::key_budgets::GetKeyBudgetRequest { + client_key_id: "key_1".into(), + }, + ) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "update_budget_limits" => { + let result = call + .host + .update_key_budget_limits( + gateway_plugin_sdk::call::key_budgets::UpdateKeyBudgetLimitsRequest { + client_key_id: "key_1".into(), + daily_limit_usd: None, + weekly_limit_usd: Some("12.5".into()), + }, + ) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "refresh_quota" => { + let result = call + .host + .refresh_account_quota(gateway_plugin_sdk::call::data::QuotaFactsQuery { + account_id: "acct_1".into(), + }) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "reset_budget" => { + use gateway_plugin_sdk::call::{ + host::KeyListRequest, + key_budgets::{BudgetPeriod, ResetKeyBudgetRequest}, + }; + let keys = call + .host + .list_keys(KeyListRequest { + cursor: None, + limit: 10, + }) + .await?; + let result = call + .host + .reset_key_budget(ResetKeyBudgetRequest { + client_key_id: keys.keys[0].id.clone(), + period: BudgetPeriod::Weekly, + }) + .await?; + Ok(CallReply::unary( + serde_json::to_value(result).unwrap(), + vec![], + )) + } + "slow" => { + tokio::time::sleep(Duration::from_millis(80)).await; + Ok(CallReply::unary(call.params, call.payload)) + } + "callback" => { + let reply = call + .host + .call("host.log", call.params, call.payload) + .await + .map_err(SessionError::into_plugin_fault)?; + Ok(CallReply::unary(reply.result, reply.payload)) + } + "resource_stream" => Ok(CallReply::stream( + json!({}), + Vec::new(), + ResponseStream::pull(Box::new(ResourceStream(Some(call.host)))), + )), + "stream" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"abcd".to_vec(), b"efgh".to_vec()]), + )), + "long_stream" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"x".to_vec(); 16]), + )), + "oversized_stream" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"abcde".to_vec()]), + )), + "too_many_chunks" => Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + ResponseStream::from_chunks(vec![b"x".to_vec(); 17]), + )), + "dynamic_oversized_stream" => { + let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); + tokio::spawn(async move { + let _ = sender.send(b"abcde".to_vec()).await; + }); + Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + stream, + )) + } + "oversized_fault" => Err(PluginFault::new( + ErrorCode::Upstream, + "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), + )), + "oversized_result" => Ok(CallReply::unary( + Value::String("x".repeat(MAXIMUM_STREAM_CHUNK_BYTES)), + Vec::new(), + )), + "dynamic_oversized_fault" => { + let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); + tokio::spawn(async move { + let _ = sender + .fail(PluginFault::new( + ErrorCode::Upstream, + "x".repeat(MAXIMUM_STREAM_CHUNK_BYTES), + )) + .await; + }); + Ok(CallReply::stream( + json!({"stream": true}), + Vec::new(), + stream, + )) + } + "pending" => pending::>().await, + _ => Err(PluginFault::new( + ErrorCode::Unsupported, + "unsupported test method", + )), + } + }) + } + + fn cancel(&self, _context: &CallContext) { + self.lifecycle.cancel_started.store(true, Ordering::Release); + if let Some(delay) = self.cancel_delay { + std::thread::sleep(delay); + } + self.lifecycle.cancelled.fetch_add(1, Ordering::Relaxed); + } + + fn quiesce(&self) { + self.lifecycle.quiesced.store(true, Ordering::Release); + } + + fn shutdown(&self) { + self.lifecycle.shutdown.store(true, Ordering::Release); + } +} + +pub(super) struct HostPeer { + pub(super) reader: ReadHalf, + pub(super) writer: WriteHalf, +} + +pub(super) async fn start_session( + handler: H, +) -> (HostPeer, JoinHandle>) { + start_session_with_capacity(handler, MAXIMUM_STREAM_CHUNK_BYTES * 2).await +} + +async fn start_session_with_capacity( + handler: H, + capacity: usize, +) -> (HostPeer, JoinHandle>) { + let (host, plugin) = tokio::io::duplex(capacity); + let (plugin_reader, plugin_writer) = tokio::io::split(plugin); + let task = tokio::spawn(async move { + let session = PluginSession::accept( + plugin_reader, + plugin_writer, + SessionConfig { + maximum_stream_chunk_bytes: MAXIMUM_STREAM_CHUNK_BYTES, + maximum_calls: 4, + maximum_callbacks: 4, + maximum_buffered_stream_chunks: 16, + handshake_timeout: Duration::from_secs(1), + maximum_call_timeout: Duration::from_secs(2), + }, + ) + .await?; + session.run(handler).await + }); + let (mut reader, mut writer) = tokio::io::split(host); + write_frame( + &mut writer, + &Frame::control(Message::Hello { + handshake: handshake(), + }), + ) + .await + .unwrap(); + let ready = read_frame(&mut reader).await.unwrap(); + assert!(matches!( + ready, + Frame { + message: Message::Ready { + protocol_version: PROTOCOL_VERSION, + ref incarnation, + }, + ref payload, + } if incarnation == "test-incarnation" && payload.is_empty() + )); + (HostPeer { reader, writer }, task) +} + +fn middleware_contributions() -> Contributions { + Contributions::from([( + Capability::Middleware, + ContributionDeclaration { + id: "test.example.middleware".into(), + version: 4, + stages: vec![Stage::Request], + input_formats: vec!["openai".into()], + output_formats: vec!["openai".into()], + }, + )]) +} + +fn middleware_plugin(map_response: bool) -> impl PluginHandler { + MiddlewarePlugin::new( + &middleware_contributions(), + move |call: RequestCall| async move { + let RequestCall { + mut request, next, .. + } = call; + if map_response { + request.body.push(b'!'); + request.remove_header("x-old"); + request.head.headers.push(MiddlewareHeader { + name: "x-direct".into(), + value: b"request".to_vec(), + }); + } + let mut response = next.run(request).await?; + if map_response { + response.remove_header("x-upstream"); + response.headers.push(MiddlewareHeader { + name: "x-direct-response".into(), + value: b"response".to_vec(), + }); + response.body = response.body.map_frames(|mut frame| { + frame.payload.make_ascii_uppercase(); + Ok(vec![frame]) + })?; + } + Ok(response) + }, + ) + .unwrap() +} + +fn handshake() -> Handshake { + Handshake { + protocol_version: PROTOCOL_VERSION, + artifact_sha256: "a".repeat(64), + plugin_id: "test.example".into(), + instance_id: "test-instance".into(), + generation: 7, + incarnation: "test-incarnation".into(), + configuration: json!({}), + + contributes: Contributions::new(), + } +} + +pub(super) fn context(id: u64, timeout: Duration) -> CallContext { + CallContext { + call_id: id, + instance_id: "test-instance".into(), + generation: 7, + incarnation: "test-incarnation".into(), + stage: Stage::Request, + timeout_ms: timeout.as_millis() as u64, + resource_stream: false, + resource_scope_id: format!("scope-{id}"), + request_id: None, + attempt_id: None, + account_id: None, + credential_revision: None, + } +} + +fn middleware_context(id: u64) -> CallContext { + CallContext { + stage: Stage::Request, + request_id: Some("request-middleware".into()), + ..context(id, Duration::from_secs(1)) + } +} + +pub(super) fn middleware_request() -> MiddlewareRequestHead { + MiddlewareRequestHead { + settings_sources: serde_json::Value::Null, + settings: serde_json::Value::Null, + client_key_id: "fixture-key".into(), + account_group_ids: vec!["fixture-group".into()], + request_id: "request-middleware".into(), + mount: MiddlewareMount::Request, + attempt_index: None, + operation: "generate".into(), + protocol: "openai".into(), + endpoint: "responses".into(), + transport: MiddlewareTransport::HttpSse, + provider: None, + model: Some("public-model".into()), + account_id: None, + headers: vec![MiddlewareHeader { + name: "x-old".into(), + value: b"old".to_vec(), + }], + } +} + +pub(super) async fn send_call( + host: &mut HostPeer, + id: u64, + method: &str, + params: Value, + payload: Vec, +) { + send_call_with_timeout(host, id, method, params, payload, Duration::from_secs(1)).await; +} + +async fn send_call_with_timeout( + host: &mut HostPeer, + id: u64, + method: &str, + params: Value, + payload: Vec, + timeout: Duration, +) { + write_frame( + &mut host.writer, + &Frame { + message: Message::Call { + id, + method: method.into(), + context: context(id, timeout), + params, + }, + payload, + }, + ) + .await + .unwrap(); +} + +async fn send_control(host: &mut HostPeer, message: Message) { + write_frame(&mut host.writer, &Frame::control(message)) + .await + .unwrap(); +} + +pub(super) async fn receive(host: &mut HostPeer) -> Frame { + tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) + .await + .expect("plugin response timed out") + .expect("plugin response frame must be valid") +} + +pub(super) async fn shutdown(host: &mut HostPeer, task: JoinHandle>) { + send_control(host, Message::Shutdown).await; + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("plugin session did not shut down") + .expect("plugin session task panicked") + .expect("plugin session shutdown failed"); +} + +struct ResourceStream(Option); +impl gateway_plugin_sdk::client::PullResponseStream for ResourceStream { + fn next(&mut self) -> gateway_plugin_sdk::client::PullResponseFuture<'_> { + Box::pin(async move { + let host = self.0.take()?; + tokio::time::sleep(Duration::from_millis(700)).await; + Some( + host.call("host.log", json!({}), Vec::new()) + .await + .map(|_| b"live".to_vec()) + .map_err(SessionError::into_plugin_fault), + ) + }) + } +} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs new file mode 100644 index 000000000..cda12ef26 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs @@ -0,0 +1,299 @@ +//! 宿主回调关联、退休回调与强类型调用合同测试 + +use super::*; + +#[tokio::test] +async fn typed_budget_helpers_preserve_wire_contract_and_propagate_failure_without_retry() { + for failed in [false, true] { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "reset_budget", json!({}), vec![]).await; + let list = receive(&mut host).await; + let Message::Callback { + id, + parent_id: 1, + method, + params, + } = list.message + else { + panic!("expected Key list callback") + }; + assert_eq!(method, "host.keys.list"); + assert_eq!(params, json!({"cursor":null,"limit":10})); + assert!(list.payload.is_empty()); + send_control(&mut host, Message::Result { id, result: json!({"keys":[{"id":"key_1","name":"test","enabled":true}],"next_cursor":null}) }).await; + let reset = receive(&mut host).await; + let Message::Callback { + id, + parent_id: 1, + method, + params, + } = reset.message + else { + panic!("expected budget reset callback") + }; + assert_eq!(method, "host.keys.reset_budget"); + assert_eq!(params, json!({})); + assert_eq!( + serde_json::from_slice::(&reset.payload).unwrap(), + json!({"client_key_id":"key_1","period":"weekly"}) + ); + if failed { + send_control( + &mut host, + Message::Error { + id, + error: PluginFault::new(ErrorCode::Conflict, "stale instance"), + }, + ) + .await; + let Message::Error { id: 1, error } = receive(&mut host).await.message else { + panic!("expected error propagated without another callback") + }; + assert_eq!(error.code, ErrorCode::Conflict); + } else { + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id, + result: json!({}), + }, + payload: serde_json::to_vec(&json!({"client_key_id":"key_1"})).unwrap(), + }, + ) + .await + .unwrap(); + let Message::Result { id: 1, result } = receive(&mut host).await.message else { + panic!("expected successful reset") + }; + assert_eq!(result, json!({"client_key_id":"key_1"})); + } + shutdown(&mut host, task).await; + } +} + +#[tokio::test] +async fn host_callback_round_trips_while_parent_call_is_waiting() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({"event": "ready"}), vec![7]).await; + + let callback = receive(&mut host).await; + let callback_id = match callback.message { + Message::Callback { + id, + parent_id: 1, + ref method, + ref params, + } if id.is_multiple_of(2) + && method == "host.log" + && params == &json!({"event": "ready"}) => + { + id + } + message => panic!("unexpected callback frame: {message:?}"), + }; + assert_eq!(callback.payload, [7]); + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: callback_id, + result: json!({"recorded": true}), + }, + payload: vec![9], + }, + ) + .await + .unwrap(); + + let result = receive(&mut host).await; + assert!(matches!( + result, + Frame { + message: Message::Result { id: 1, result }, + payload, + } if result == json!({"recorded": true}) && payload == [9] + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn callbacks_from_concurrent_calls_remain_correlated() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({"parent": 1}), vec![1]).await; + send_call(&mut host, 3, "callback", json!({"parent": 3}), vec![3]).await; + + let mut callbacks = Vec::new(); + for _ in 0..2 { + let frame = receive(&mut host).await; + let Message::Callback { + id, + parent_id, + method, + params, + } = frame.message + else { + panic!("expected callback") + }; + assert_eq!(method, "host.log"); + assert_eq!(params, json!({"parent": parent_id})); + assert_eq!(frame.payload, [parent_id as u8]); + callbacks.push((id, parent_id)); + } + assert!(callbacks[0].0 < callbacks[1].0); + assert!(callbacks.iter().all(|(id, _)| id.is_multiple_of(2))); + + for (callback_id, parent_id) in callbacks.into_iter().rev() { + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id: callback_id, + result: json!({"parent": parent_id}), + }, + payload: vec![parent_id as u8], + }, + ) + .await + .unwrap(); + } + let mut completed = Vec::new(); + for _ in 0..2 { + let frame = receive(&mut host).await; + let Message::Result { id, result } = frame.message else { + panic!("expected parent result") + }; + assert_eq!(result, json!({"parent": id})); + assert_eq!(frame.payload, [id as u8]); + completed.push(id); + } + completed.sort_unstable(); + assert_eq!(completed, [1, 3]); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn cancelling_parent_retires_callback_and_late_reply_does_not_close_session() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({}), Vec::new()).await; + let callback = receive(&mut host).await; + let Message::Callback { + id: callback_id, .. + } = callback.message + else { + panic!("expected callback") + }; + + send_control(&mut host, Message::Cancel { id: 1 }).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Cancelled { id: 1 } + )); + send_control( + &mut host, + Message::Result { + id: callback_id, + result: json!({}), + }, + ) + .await; + send_call(&mut host, 3, "echo", json!("still-ready"), Vec::new()).await; + assert!(matches!( + receive(&mut host).await.message, + Message::Result { id: 3, result } if result == json!("still-ready") + )); + + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() { + for (entry, method, request, response) in [ + ( + "key_facts", + "host.data.keys.get", + json!({"client_key_id":"key_1"}), + json!({"schema_version":1,"client_key_id":"key_1","enabled":false,"group_ids":["grp_1"]}), + ), + ( + "get_budget", + "host.keys.get_budget", + json!({"client_key_id":"key_1"}), + json!({ + "client_key_id":"key_1", "daily_limit_usd":"10", "weekly_limit_usd":"20", + "daily_used_usd":"1.25", "weekly_used_usd":"2.5", "daily_resets_at_ms":null, "weekly_resets_at_ms":null, + }), + ), + ( + "update_budget_limits", + "host.keys.update_budget_limits", + json!({"client_key_id":"key_1","weekly_limit_usd":"12.5"}), + json!({"client_key_id":"key_1"}), + ), + ( + "refresh_quota", + "host.quota_observations.refresh", + json!({"account_id":"acct_1"}), + json!({"schema_version":1,"account_id":"acct_1","observed_at_ms":null,"windows":[]}), + ), + ] { + for failed in [false, true] { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, entry, json!({}), vec![]).await; + let callback = receive(&mut host).await; + let Message::Callback { + id, + parent_id: 1, + method: actual, + params, + } = callback.message + else { + panic!("expected callback"); + }; + assert_eq!(actual, method); + assert_eq!(params, json!({})); + assert_eq!( + serde_json::from_slice::(&callback.payload).unwrap(), + request + ); + if failed { + send_control( + &mut host, + Message::Error { + id, + error: PluginFault::new(ErrorCode::Conflict, "changed"), + }, + ) + .await; + let Message::Error { id: 1, error } = receive(&mut host).await.message else { + panic!("expected propagated error without retry"); + }; + assert_eq!(error.code, ErrorCode::Conflict); + } else { + let mut response_payload = response.clone(); + if matches!(entry, "key_facts" | "refresh_quota") { + response_payload["future_fact"] = json!({"value":1}); + } + write_frame( + &mut host.writer, + &Frame { + message: Message::Result { + id, + result: json!({}), + }, + payload: serde_json::to_vec(&response_payload).unwrap(), + }, + ) + .await + .unwrap(); + let Message::Result { id: 1, result } = receive(&mut host).await.message else { + panic!("expected projected response"); + }; + assert_eq!(result, response); + } + shutdown(&mut host, task).await; + } + } +} diff --git a/backend/crates/gateway-store/src/backup/s3.rs b/backend/crates/gateway-store/src/backup/s3.rs index 0cbba7772..2e78fa03a 100644 --- a/backend/crates/gateway-store/src/backup/s3.rs +++ b/backend/crates/gateway-store/src/backup/s3.rs @@ -7,10 +7,7 @@ use std::collections::HashMap; use std::path::Path; -use std::sync::{ - Arc, Mutex, - atomic::{AtomicU64, Ordering}, -}; +use std::sync::Mutex; use std::time::Duration; use aws_credential_types::Credentials; @@ -24,9 +21,9 @@ use aws_smithy_types::error::metadata::ProvideErrorMetadata; use aws_smithy_types::retry::RetryConfig; use aws_smithy_types::timeout::TimeoutConfig; use chrono::{DateTime, Utc}; +use futures::{StreamExt as _, TryStreamExt as _, stream}; use secrecy::ExposeSecret as _; use tokio::io::{AsyncReadExt as _, AsyncSeekExt as _, SeekFrom}; -use tokio::sync::Semaphore; use uuid::Uuid; use gateway_admin::model::backup::{ @@ -242,7 +239,7 @@ impl BackupObjectStorePort for S3ObjectStoreAdapter { })? .to_owned(); - let upload = upload_parts(&client, config, &request, upload_id.clone(), file_size); + let upload = upload_parts(&client, config, &request, &upload_id, file_size); let outcome = tokio::select! { result = upload => result, _ = request.cancellation.cancelled() => { @@ -358,78 +355,43 @@ async fn upload_parts( client: &S3Client, config: &BackupStorageConfig, request: &UploadObjectRequest, - upload_id: String, + upload_id: &str, file_size: u64, ) -> Result<(), BackupError> { let total_parts = file_size.div_ceil(u64::try_from(PART_SIZE).unwrap_or(u64::MAX)); - let next_part = Arc::new(AtomicU64::new(1)); - let semaphore = Arc::new(Semaphore::new(CONCURRENCY)); - let (sender, mut receiver) = tokio::sync::mpsc::channel::<(i32, String)>(CONCURRENCY * 2); - - let mut workers = Vec::with_capacity(CONCURRENCY); - for _ in 0..CONCURRENCY { - let client = client.clone(); - let semaphore = Arc::clone(&semaphore); - let next_part = Arc::clone(&next_part); - let sender = sender.clone(); - let source = request.source.clone(); - let bucket = config.bucket.clone(); - let key = request.object_key.clone(); - let upload_id = upload_id.clone(); - workers.push(tokio::spawn(async move { - loop { - let part = next_part.fetch_add(1, Ordering::Relaxed); - if part > total_parts { - break; - } - let _permit = semaphore.acquire().await.map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "上传并发闸门关闭".to_owned()) - })?; - let offset = (part - 1) * (PART_SIZE as u64); - let remaining = file_size.checked_sub(offset).ok_or_else(|| { - BackupError::new(code::S3_UPLOAD_FAILED, "分片偏移超出归档大小".to_owned()) - })?; - let part_len = usize::try_from(remaining.min(PART_SIZE as u64)).map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "分片大小溢出".to_owned()) - })?; - let bytes = read_part(&source, offset, part_len).await?; - let part_number = i32::try_from(part).map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "分片号溢出".to_owned()) - })?; - let response = client - .upload_part() - .bucket(&bucket) - .key(&key) - .upload_id(&upload_id) - .part_number(part_number) - .body(ByteStream::from(bytes)) - .send() - .await - .map_err(map_s3_error)?; - let etag = response - .e_tag() - .ok_or_else(|| { - BackupError::new(code::S3_UPLOAD_FAILED, "S3 未返回分片 ETag".to_owned()) - })? - .to_owned(); - if sender.send((part_number, etag)).await.is_err() { - break; - } - } - Ok::<(), BackupError>(()) - })); - } - drop(sender); - - let mut parts = Vec::new(); - while let Some((part, etag)) = receiver.recv().await { - parts.push((part, etag)); - } - for worker in workers { - worker.await.map_err(|_| { - BackupError::new(code::S3_UPLOAD_FAILED, "上传分片任务崩溃".to_owned()) - })??; - } + // 分片 future 由父上传持有；取消或首次失败时一并释放，SDK 重试不能逃逸到清理之后 + let mut parts: Vec<_> = stream::iter(1..=total_parts) + .map(|part| async move { + let offset = (part - 1) * (PART_SIZE as u64); + let remaining = file_size.checked_sub(offset).ok_or_else(|| { + BackupError::new(code::S3_UPLOAD_FAILED, "分片偏移超出归档大小".to_owned()) + })?; + let part_len = usize::try_from(remaining.min(PART_SIZE as u64)) + .map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "分片大小溢出".to_owned()))?; + let bytes = read_part(&request.source, offset, part_len).await?; + let part_number = i32::try_from(part) + .map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "分片号溢出".to_owned()))?; + let response = client + .upload_part() + .bucket(&config.bucket) + .key(&request.object_key) + .upload_id(upload_id) + .part_number(part_number) + .body(ByteStream::from(bytes)) + .send() + .await + .map_err(map_s3_error)?; + let etag = response + .e_tag() + .ok_or_else(|| { + BackupError::new(code::S3_UPLOAD_FAILED, "S3 未返回分片 ETag".to_owned()) + })? + .to_owned(); + Ok::<_, BackupError>((part_number, etag)) + }) + .buffer_unordered(CONCURRENCY) + .try_collect() + .await?; parts.sort_by_key(|(part, _)| *part); let completed = CompletedMultipartUpload::builder() @@ -449,7 +411,7 @@ async fn upload_parts( .complete_multipart_upload() .bucket(&config.bucket) .key(&request.object_key) - .upload_id(&upload_id) + .upload_id(upload_id) .multipart_upload(completed) .send() .await diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs index 190880430..377eb1daa 100644 --- a/backend/crates/gateway-store/src/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -371,12 +371,13 @@ pub(super) async fn load_client_key_budgets( .iter() .map(|record| record.id.as_str()) .collect::>(); + // 起止相等是重置后的未开启窗口，不能因应用时钟领先数据库而投影为活跃预算 let rows = sqlx::query( "select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text, - (case when w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used, - (case when w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used, - case when w.daily_end > now() then w.daily_end end as daily_end, - case when w.weekly_end > now() then w.weekly_end end as weekly_end + (case when w.daily_end > w.daily_start and w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used, + (case when w.weekly_end > w.weekly_start and w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used, + case when w.daily_end > w.daily_start and w.daily_end > now() then w.daily_end end as daily_end, + case when w.weekly_end > w.weekly_start and w.weekly_end > now() then w.weekly_end end as weekly_end from client_api_keys k left join client_key_budget_windows w on w.client_api_key_id = k.id where k.id = any($1)", ) diff --git a/backend/crates/gateway-store/src/postgres/connection.rs b/backend/crates/gateway-store/src/postgres/connection.rs new file mode 100644 index 000000000..6cb104040 --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/connection.rs @@ -0,0 +1,104 @@ +//! PostgreSQL 迁移屏障、连接池与会话级资源预算 + +use sqlx::{ + PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; + +use crate::{ + POSTGRES_IDLE_TRANSACTION_TIMEOUT, POSTGRES_LOCK_TIMEOUT, POSTGRES_STATEMENT_TIMEOUT, + StoreError, StorePoolConfig, StoreResult, postgres_unavailable, +}; + +static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations"); + +/// 建立 PostgreSQL pool 并只执行冻结的 migration 集 +pub async fn connect_and_migrate( + database_url: &str, + pool_config: StorePoolConfig, +) -> StoreResult { + if database_url.trim().is_empty() { + return Err(StoreError::InvalidData { + source: None, + entity: "PostgreSQL configuration", + message: "database URL is empty".to_owned(), + }); + } + pool_config.validate()?; + let connect_options = database_url + .parse::() + .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; + let migration_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + connect_options + .clone() + .application_name("codex-proxy-rs:migration"), + ) + .await + .map_err(|source| postgres_unavailable("connect PostgreSQL for migrations", source))?; + if let Err(error) = MIGRATOR.run(&migration_pool).await { + migration_pool.close().await; + return Err(postgres_unavailable("apply PostgreSQL migrations", error)); + } + migration_pool.close().await; + + connect_pool(connect_options, pool_config, false).await +} + +/// 帮助查询不执行迁移，并用连接默认只读事务阻止意外业务写入 +pub(crate) async fn connect_read_only( + database_url: &str, + pool_config: StorePoolConfig, +) -> StoreResult { + pool_config.validate()?; + let options = database_url + .parse::() + .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; + connect_pool(options, pool_config, true).await +} + +async fn connect_pool( + connect_options: PgConnectOptions, + pool_config: StorePoolConfig, + read_only: bool, +) -> StoreResult { + let statement_timeout = postgres_duration_setting(POSTGRES_STATEMENT_TIMEOUT); + let lock_timeout = postgres_duration_setting(POSTGRES_LOCK_TIMEOUT); + let idle_in_transaction_session_timeout = + postgres_duration_setting(POSTGRES_IDLE_TRANSACTION_TIMEOUT); + let pool = PgPoolOptions::new() + .max_connections(pool_config.max_connections) + .acquire_timeout(std::time::Duration::from_secs( + pool_config.acquire_timeout_seconds, + )) + .after_connect(move |connection, _metadata| { + let statement_timeout = statement_timeout.clone(); + let lock_timeout = lock_timeout.clone(); + let idle_in_transaction_session_timeout = idle_in_transaction_session_timeout.clone(); + Box::pin(async move { + sqlx::query( + "select set_config('statement_timeout', $1, false), + set_config('lock_timeout', $2, false), + set_config('idle_in_transaction_session_timeout', $3, false), + set_config('default_transaction_read_only', $4, false), + set_config('TimeZone', 'UTC', false)", + ) + .bind(statement_timeout) + .bind(lock_timeout) + .bind(idle_in_transaction_session_timeout) + .bind(if read_only { "on" } else { "off" }) + .execute(connection) + .await?; + Ok(()) + }) + }) + .connect_with(connect_options.application_name("codex-proxy-rs")) + .await + .map_err(|source| postgres_unavailable("connect PostgreSQL", source))?; + Ok(pool) +} + +fn postgres_duration_setting(duration: std::time::Duration) -> String { + format!("{}ms", duration.as_millis()) +} diff --git a/backend/crates/gateway-store/src/postgres/control_plane.rs b/backend/crates/gateway-store/src/postgres/control_plane.rs new file mode 100644 index 000000000..1dcf96173 --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/control_plane.rs @@ -0,0 +1,286 @@ +//! 控制面设置与 Key 修改的 revision 锁、条件写入和同事务审计 + +use async_trait::async_trait; +use sqlx::PgPool; + +use crate::{Revision, StoreError, StoreResult, postgres_unavailable}; + +use super::{ + AdminAuditEvent, NewClientApiKey, RuntimeSettings, RuntimeSettingsUpdate, + UpdateClientApiKeyDetails, append_admin_audit_event_in_transaction, + bump_config_revision_in_transaction, delete_client_api_key_in_transaction, + insert_client_api_key_in_transaction, load_runtime_settings_from_pool, + load_runtime_settings_in_transaction, set_client_api_key_enabled_in_transaction, + update_admin_api_key_in_transaction, update_client_api_key_in_transaction, + update_runtime_settings_in_transaction, +}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ControlPlaneSnapshot { + pub settings: RuntimeSettings, +} + +#[derive(Debug, Clone)] +pub struct ControlPlaneReplacement { + pub expected_revision: Revision, + pub settings: RuntimeSettingsUpdate, + pub audit: AdminAuditEvent, +} + +#[async_trait] +pub trait ControlPlaneRepository: Send + Sync { + async fn load_control_plane(&self) -> StoreResult; + + async fn replace_control_plane( + &self, + replacement: ControlPlaneReplacement, + ) -> StoreResult; + + /// 更新 admin_api_key 字段并推进 config revision + async fn replace_admin_api_key( + &self, + admin_api_key: Option, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn create_client_api_key( + &self, + key: NewClientApiKey, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn update_client_api_key( + &self, + key: UpdateClientApiKeyDetails, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn set_client_api_key_enabled( + &self, + id: &str, + enabled: bool, + audit: AdminAuditEvent, + ) -> StoreResult; + + async fn delete_client_api_key( + &self, + id: &str, + audit: AdminAuditEvent, + ) -> StoreResult; +} + +#[derive(Clone)] +pub struct PgControlPlaneRepository { + pub(super) pool: PgPool, +} + +impl PgControlPlaneRepository { + #[must_use] + pub const fn new(pool: PgPool) -> Self { + Self { pool } + } +} + +#[async_trait] +impl ControlPlaneRepository for PgControlPlaneRepository { + async fn load_control_plane(&self) -> StoreResult { + let settings = load_runtime_settings_from_pool(&self.pool).await?; + Ok(ControlPlaneSnapshot { settings }) + } + + async fn replace_control_plane( + &self, + replacement: ControlPlaneReplacement, + ) -> StoreResult { + replacement.settings.validate()?; + let mut transaction = + self.pool.begin().await.map_err(|source| { + postgres_unavailable("begin control plane replacement", source) + })?; + let result: StoreResult<_> = async { + // 锁住所有配置写入共同使用的行，版本检查与修改不能被其他事务穿插 + let current = sqlx::query_scalar::<_, i64>( + "select config_revision from runtime_settings where id = 1 for update", + ) + .fetch_one(&mut *transaction) + .await + .map_err(|source| postgres_unavailable("lock control plane revision", source))?; + if u64::try_from(current).ok() != Some(replacement.expected_revision.get()) { + return Err(StoreError::Conflict { + source: None, + entity: "runtime settings", + id: "1".to_owned(), + kind: crate::ConflictKind::StaleRevision, + }); + } + let revision = + update_runtime_settings_in_transaction(&mut transaction, &replacement.settings) + .await?; + append_admin_audit_event_in_transaction(&mut transaction, replacement.audit, revision) + .await?; + load_control_plane_in_transaction(&mut transaction).await + } + .await; + match result { + Ok(snapshot) => { + transaction.commit().await.map_err(|source| { + postgres_unavailable("commit control plane replacement", source) + })?; + Ok(snapshot) + } + Err(error) => { + let error = match transaction.rollback().await { + Ok(()) => error, + Err(cleanup) => error.with_cleanup(cleanup), + }; + Err(error) + } + } + } + + async fn replace_admin_api_key( + &self, + admin_api_key: Option, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation(ControlPlaneMutation::SetAdminApiKey(admin_api_key), audit) + .await + } + + async fn create_client_api_key( + &self, + key: NewClientApiKey, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation(ControlPlaneMutation::CreateClientApiKey(key), audit) + .await + } + + async fn update_client_api_key( + &self, + key: UpdateClientApiKeyDetails, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation(ControlPlaneMutation::UpdateClientApiKey(key), audit) + .await + } + + async fn set_client_api_key_enabled( + &self, + id: &str, + enabled: bool, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation( + ControlPlaneMutation::SetClientApiKeyEnabled { + id: id.to_owned(), + enabled, + }, + audit, + ) + .await + } + + async fn delete_client_api_key( + &self, + id: &str, + audit: AdminAuditEvent, + ) -> StoreResult { + self.apply_targeted_mutation( + ControlPlaneMutation::DeleteClientApiKey(id.to_owned()), + audit, + ) + .await + } +} + +enum ControlPlaneMutation { + CreateClientApiKey(NewClientApiKey), + UpdateClientApiKey(UpdateClientApiKeyDetails), + SetClientApiKeyEnabled { id: String, enabled: bool }, + DeleteClientApiKey(String), + SetAdminApiKey(Option), +} + +impl PgControlPlaneRepository { + async fn apply_targeted_mutation( + &self, + mutation: ControlPlaneMutation, + mut audit: AdminAuditEvent, + ) -> StoreResult { + let mut transaction = self.pool.begin().await.map_err(|source| { + postgres_unavailable("begin targeted control plane mutation", source) + })?; + let result: StoreResult<_> = async { + let revision = bump_config_revision_in_transaction(&mut transaction).await?; + match mutation { + ControlPlaneMutation::CreateClientApiKey(key) => { + audit.changed_fields.push(if key.group_ids.is_empty() { + "routing_scope:all".to_owned() + } else { + "routing_scope:groups".to_owned() + }); + insert_client_api_key_in_transaction(&mut transaction, &key).await?; + } + ControlPlaneMutation::UpdateClientApiKey(key) => { + let previously_restricted = sqlx::query_scalar::<_, bool>( + "select exists( + select 1 from client_api_key_groups where client_api_key_id = $1 + )", + ) + .bind(&key.id) + .fetch_one(&mut *transaction) + .await + .map_err(|source| { + postgres_unavailable("load client API key routing scope for audit", source) + })?; + if previously_restricted && key.group_ids.is_empty() { + audit + .changed_fields + .push("routing_scope:groups->all".to_owned()); + } else if !previously_restricted && !key.group_ids.is_empty() { + audit + .changed_fields + .push("routing_scope:all->groups".to_owned()); + } + update_client_api_key_in_transaction(&mut transaction, &key).await?; + } + ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => { + set_client_api_key_enabled_in_transaction(&mut transaction, &id, enabled) + .await?; + } + ControlPlaneMutation::DeleteClientApiKey(id) => { + delete_client_api_key_in_transaction(&mut transaction, &id).await?; + } + ControlPlaneMutation::SetAdminApiKey(key) => { + update_admin_api_key_in_transaction(&mut transaction, key).await?; + } + } + append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?; + Ok(revision) + } + .await; + match result { + Ok(revision) => { + transaction.commit().await.map_err(|source| { + postgres_unavailable("commit targeted control plane mutation", source) + })?; + Ok(revision) + } + Err(error) => { + let error = match transaction.rollback().await { + Ok(()) => error, + Err(cleanup) => error.with_cleanup(cleanup), + }; + Err(error) + } + } + } +} + +async fn load_control_plane_in_transaction( + transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>, +) -> StoreResult { + let settings = load_runtime_settings_in_transaction(transaction).await?; + Ok(ControlPlaneSnapshot { settings }) +} diff --git a/backend/crates/gateway-store/src/postgres/mod.rs b/backend/crates/gateway-store/src/postgres/mod.rs index 72c086924..aaff5b330 100644 --- a/backend/crates/gateway-store/src/postgres/mod.rs +++ b/backend/crates/gateway-store/src/postgres/mod.rs @@ -1,15 +1,4 @@ -//! PostgreSQL 业务表的 adapters - -use async_trait::async_trait; -use sqlx::{ - PgPool, - postgres::{PgConnectOptions, PgPoolOptions}, -}; - -use crate::{ - POSTGRES_IDLE_TRANSACTION_TIMEOUT, POSTGRES_LOCK_TIMEOUT, POSTGRES_STATEMENT_TIMEOUT, Revision, - StoreError, StorePoolConfig, StoreResult, postgres_unavailable, -}; +//! PostgreSQL 业务表、连接与控制面事务的适配入口 mod account_groups; mod admin_security_audit; @@ -17,6 +6,8 @@ mod admission_recovery; mod backup; mod client_budgets; mod client_keys; +mod connection; +mod control_plane; mod execution; mod execution_buffer; mod observability; @@ -25,6 +16,7 @@ mod plugins; mod pricing; mod provider_accounts; mod proxies; +mod query_pattern; mod retention; mod runtime_settings; mod snapshot; @@ -36,6 +28,11 @@ pub use admission_recovery::*; pub use backup::*; pub use client_budgets::PgClientBudgetStore; pub use client_keys::*; +pub use connection::connect_and_migrate; +pub(crate) use connection::connect_read_only; +pub use control_plane::{ + ControlPlaneReplacement, ControlPlaneRepository, ControlPlaneSnapshot, PgControlPlaneRepository, +}; pub use execution::*; pub use execution_buffer::*; pub use observability::*; @@ -43,6 +40,7 @@ pub use ops_events::*; pub use plugins::PgPluginStore; pub use provider_accounts::*; pub use proxies::PgProxyRepository; +use query_pattern::literal_prefix_pattern; pub use retention::*; pub use runtime_settings::*; pub use snapshot::*; @@ -50,378 +48,3 @@ pub(crate) use usage_facts::{ completed_usage_fact_predicate, push_completed_usage_fact_filter, push_unrecovered_request_filter, }; - -static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations"); - -/// 建立 PostgreSQL pool 并只执行冻结的 migration 集 -pub async fn connect_and_migrate( - database_url: &str, - pool_config: StorePoolConfig, -) -> StoreResult { - if database_url.trim().is_empty() { - return Err(StoreError::InvalidData { - source: None, - entity: "PostgreSQL configuration", - message: "database URL is empty".to_owned(), - }); - } - pool_config.validate()?; - let connect_options = database_url - .parse::() - .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; - let migration_pool = PgPoolOptions::new() - .max_connections(1) - .connect_with( - connect_options - .clone() - .application_name("codex-proxy-rs:migration"), - ) - .await - .map_err(|source| postgres_unavailable("connect PostgreSQL for migrations", source))?; - if let Err(error) = MIGRATOR.run(&migration_pool).await { - migration_pool.close().await; - return Err(postgres_unavailable("apply PostgreSQL migrations", error)); - } - migration_pool.close().await; - - connect_pool(connect_options, pool_config, false).await -} - -/// 帮助查询不执行迁移，并用连接默认只读事务阻止意外业务写入 -pub(crate) async fn connect_read_only( - database_url: &str, - pool_config: StorePoolConfig, -) -> StoreResult { - pool_config.validate()?; - let options = database_url - .parse::() - .map_err(|source| postgres_unavailable("parse PostgreSQL connection options", source))?; - connect_pool(options, pool_config, true).await -} - -async fn connect_pool( - connect_options: PgConnectOptions, - pool_config: StorePoolConfig, - read_only: bool, -) -> StoreResult { - let statement_timeout = postgres_duration_setting(POSTGRES_STATEMENT_TIMEOUT); - let lock_timeout = postgres_duration_setting(POSTGRES_LOCK_TIMEOUT); - let idle_in_transaction_session_timeout = - postgres_duration_setting(POSTGRES_IDLE_TRANSACTION_TIMEOUT); - let pool = PgPoolOptions::new() - .max_connections(pool_config.max_connections) - .acquire_timeout(std::time::Duration::from_secs( - pool_config.acquire_timeout_seconds, - )) - .after_connect(move |connection, _metadata| { - let statement_timeout = statement_timeout.clone(); - let lock_timeout = lock_timeout.clone(); - let idle_in_transaction_session_timeout = idle_in_transaction_session_timeout.clone(); - Box::pin(async move { - sqlx::query( - "select set_config('statement_timeout', $1, false), - set_config('lock_timeout', $2, false), - set_config('idle_in_transaction_session_timeout', $3, false), - set_config('default_transaction_read_only', $4, false), - set_config('TimeZone', 'UTC', false)", - ) - .bind(statement_timeout) - .bind(lock_timeout) - .bind(idle_in_transaction_session_timeout) - .bind(if read_only { "on" } else { "off" }) - .execute(connection) - .await?; - Ok(()) - }) - }) - .connect_with(connect_options.application_name("codex-proxy-rs")) - .await - .map_err(|source| postgres_unavailable("connect PostgreSQL", source))?; - Ok(pool) -} - -fn postgres_duration_setting(duration: std::time::Duration) -> String { - format!("{}ms", duration.as_millis()) -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ControlPlaneSnapshot { - pub settings: RuntimeSettings, -} - -#[derive(Debug, Clone)] -pub struct ControlPlaneReplacement { - pub expected_revision: Revision, - pub settings: RuntimeSettingsUpdate, - pub audit: AdminAuditEvent, -} - -#[async_trait] -pub trait ControlPlaneRepository: Send + Sync { - async fn load_control_plane(&self) -> StoreResult; - - async fn replace_control_plane( - &self, - replacement: ControlPlaneReplacement, - ) -> StoreResult; - - /// 更新 admin_api_key 字段并推进 config revision - async fn replace_admin_api_key( - &self, - admin_api_key: Option, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn create_client_api_key( - &self, - key: NewClientApiKey, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn update_client_api_key( - &self, - key: UpdateClientApiKeyDetails, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn set_client_api_key_enabled( - &self, - id: &str, - enabled: bool, - audit: AdminAuditEvent, - ) -> StoreResult; - - async fn delete_client_api_key( - &self, - id: &str, - audit: AdminAuditEvent, - ) -> StoreResult; -} - -#[derive(Clone)] -pub struct PgControlPlaneRepository { - pool: PgPool, -} - -impl PgControlPlaneRepository { - #[must_use] - pub const fn new(pool: PgPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl ControlPlaneRepository for PgControlPlaneRepository { - async fn load_control_plane(&self) -> StoreResult { - let settings = load_runtime_settings_from_pool(&self.pool).await?; - Ok(ControlPlaneSnapshot { settings }) - } - - async fn replace_control_plane( - &self, - replacement: ControlPlaneReplacement, - ) -> StoreResult { - replacement.settings.validate()?; - let mut transaction = - self.pool.begin().await.map_err(|source| { - postgres_unavailable("begin control plane replacement", source) - })?; - let result: StoreResult<_> = async { - // 锁住所有配置写入共同使用的行，版本检查与修改不能被其他事务穿插 - let current = sqlx::query_scalar::<_, i64>( - "select config_revision from runtime_settings where id = 1 for update", - ) - .fetch_one(&mut *transaction) - .await - .map_err(|source| postgres_unavailable("lock control plane revision", source))?; - if u64::try_from(current).ok() != Some(replacement.expected_revision.get()) { - return Err(StoreError::Conflict { - source: None, - entity: "runtime settings", - id: "1".to_owned(), - kind: crate::ConflictKind::StaleRevision, - }); - } - let revision = - update_runtime_settings_in_transaction(&mut transaction, &replacement.settings) - .await?; - append_admin_audit_event_in_transaction(&mut transaction, replacement.audit, revision) - .await?; - load_control_plane_in_transaction(&mut transaction).await - } - .await; - match result { - Ok(snapshot) => { - transaction.commit().await.map_err(|source| { - postgres_unavailable("commit control plane replacement", source) - })?; - Ok(snapshot) - } - Err(error) => { - let error = match transaction.rollback().await { - Ok(()) => error, - Err(cleanup) => error.with_cleanup(cleanup), - }; - Err(error) - } - } - } - - async fn replace_admin_api_key( - &self, - admin_api_key: Option, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation(ControlPlaneMutation::SetAdminApiKey(admin_api_key), audit) - .await - } - - async fn create_client_api_key( - &self, - key: NewClientApiKey, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation(ControlPlaneMutation::CreateClientApiKey(key), audit) - .await - } - - async fn update_client_api_key( - &self, - key: UpdateClientApiKeyDetails, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation(ControlPlaneMutation::UpdateClientApiKey(key), audit) - .await - } - - async fn set_client_api_key_enabled( - &self, - id: &str, - enabled: bool, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation( - ControlPlaneMutation::SetClientApiKeyEnabled { - id: id.to_owned(), - enabled, - }, - audit, - ) - .await - } - - async fn delete_client_api_key( - &self, - id: &str, - audit: AdminAuditEvent, - ) -> StoreResult { - self.apply_targeted_mutation( - ControlPlaneMutation::DeleteClientApiKey(id.to_owned()), - audit, - ) - .await - } -} - -enum ControlPlaneMutation { - CreateClientApiKey(NewClientApiKey), - UpdateClientApiKey(UpdateClientApiKeyDetails), - SetClientApiKeyEnabled { id: String, enabled: bool }, - DeleteClientApiKey(String), - SetAdminApiKey(Option), -} - -impl PgControlPlaneRepository { - async fn apply_targeted_mutation( - &self, - mutation: ControlPlaneMutation, - mut audit: AdminAuditEvent, - ) -> StoreResult { - let mut transaction = self.pool.begin().await.map_err(|source| { - postgres_unavailable("begin targeted control plane mutation", source) - })?; - let result: StoreResult<_> = async { - let revision = bump_config_revision_in_transaction(&mut transaction).await?; - match mutation { - ControlPlaneMutation::CreateClientApiKey(key) => { - audit.changed_fields.push(if key.group_ids.is_empty() { - "routing_scope:all".to_owned() - } else { - "routing_scope:groups".to_owned() - }); - insert_client_api_key_in_transaction(&mut transaction, &key).await?; - } - ControlPlaneMutation::UpdateClientApiKey(key) => { - let previously_restricted = sqlx::query_scalar::<_, bool>( - "select exists( - select 1 from client_api_key_groups where client_api_key_id = $1 - )", - ) - .bind(&key.id) - .fetch_one(&mut *transaction) - .await - .map_err(|source| { - postgres_unavailable("load client API key routing scope for audit", source) - })?; - if previously_restricted && key.group_ids.is_empty() { - audit - .changed_fields - .push("routing_scope:groups->all".to_owned()); - } else if !previously_restricted && !key.group_ids.is_empty() { - audit - .changed_fields - .push("routing_scope:all->groups".to_owned()); - } - update_client_api_key_in_transaction(&mut transaction, &key).await?; - } - ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => { - set_client_api_key_enabled_in_transaction(&mut transaction, &id, enabled) - .await?; - } - ControlPlaneMutation::DeleteClientApiKey(id) => { - delete_client_api_key_in_transaction(&mut transaction, &id).await?; - } - ControlPlaneMutation::SetAdminApiKey(key) => { - update_admin_api_key_in_transaction(&mut transaction, key).await?; - } - } - append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?; - Ok(revision) - } - .await; - match result { - Ok(revision) => { - transaction.commit().await.map_err(|source| { - postgres_unavailable("commit targeted control plane mutation", source) - })?; - Ok(revision) - } - Err(error) => { - let error = match transaction.rollback().await { - Ok(()) => error, - Err(cleanup) => error.with_cleanup(cleanup), - }; - Err(error) - } - } - } -} - -async fn load_control_plane_in_transaction( - transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>, -) -> StoreResult { - let settings = load_runtime_settings_in_transaction(transaction).await?; - Ok(ControlPlaneSnapshot { settings }) -} - -fn literal_prefix_pattern(value: &str) -> String { - let mut escaped = String::with_capacity(value.len().saturating_add(1)); - for character in value.to_lowercase().chars() { - if matches!(character, '\\' | '%' | '_') { - escaped.push('\\'); - } - escaped.push(character); - } - escaped.push('%'); - escaped -} diff --git a/backend/crates/gateway-store/src/postgres/query_pattern.rs b/backend/crates/gateway-store/src/postgres/query_pattern.rs new file mode 100644 index 000000000..e44a664be --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/query_pattern.rs @@ -0,0 +1,13 @@ +//! 管理实体搜索共用的 SQL 字面前缀转义 + +pub(super) fn literal_prefix_pattern(value: &str) -> String { + let mut escaped = String::with_capacity(value.len().saturating_add(1)); + for character in value.to_lowercase().chars() { + if matches!(character, '\\' | '%' | '_') { + escaped.push('\\'); + } + escaped.push(character); + } + escaped.push('%'); + escaped +} diff --git a/backend/crates/gateway-store/tests/backup/s3.rs b/backend/crates/gateway-store/tests/backup/s3.rs index cf7efee07..d7a4c2461 100644 --- a/backend/crates/gateway-store/tests/backup/s3.rs +++ b/backend/crates/gateway-store/tests/backup/s3.rs @@ -1,9 +1,9 @@ -//! 验证 S3 上传完整发送超过异步文件默认缓冲区的分块 +//! 验证 S3 分片完整读取、并发上限与取消和失败时的请求回收 -use std::io; +use std::{io, time::Duration}; use chrono::Utc; -use gateway_admin::model::backup::{BackupObjectMetadata, BackupStorageConfig}; +use gateway_admin::model::backup::{BackupError, BackupObjectMetadata, BackupStorageConfig, code}; use gateway_admin::ports::backup::{BackupObjectStorePort, UploadObjectRequest}; use gateway_core::lifecycle::CancellationToken; use gateway_store::backup::s3::S3ObjectStoreAdapter; @@ -12,6 +12,7 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; use tokio::net::{TcpListener, TcpStream}; const PART_SIZE_ABOVE_TOKIO_DEFAULT: usize = 2 * 1024 * 1024 + 1; +const MULTIPART_SIZE: usize = 16 * 1024 * 1024; #[tokio::test] async fn upload_file_sends_full_part_beyond_tokio_default_buffer_limit() { @@ -64,6 +65,223 @@ async fn upload_file_sends_full_part_beyond_tokio_default_buffer_limit() { assert_eq!(uploaded_len, PART_SIZE_ABOVE_TOKIO_DEFAULT); } +#[tokio::test] +async fn cancelled_upload_does_not_retry_a_pending_part_after_returning() { + let probe = start_upload(MULTIPART_SIZE).await; + let (mut part, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + + probe.cancellation.cancel(); + let (mut abort, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("DELETE ") && request.contains("uploadId=")); + respond(&mut abort, &xml_response("")).await; + let error = tokio::time::timeout(Duration::from_secs(3), probe.upload) + .await + .expect("cancelled upload returns without waiting for the part response") + .expect("upload task did not panic") + .expect_err("upload was cancelled"); + assert_eq!(error.code(), code::CANCELLED); + + assert_no_part_retry(&probe.listener, &mut part).await; +} + +#[tokio::test] +async fn a_failed_part_stops_other_pending_parts_before_upload_returns() { + let probe = start_upload(MULTIPART_SIZE * 2).await; + let (mut failed, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + let (mut pending, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + respond( + &mut failed, + &error_response("403 Forbidden", "AccessDenied"), + ) + .await; + + // 另一分片仍未收到响应，首次失败必须立即撤销它并进入 multipart 清理 + let (mut abort, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("DELETE ") && request.contains("uploadId=")); + respond(&mut abort, &xml_response("")).await; + let (mut delete, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("DELETE ") && !request.contains("uploadId=")); + respond(&mut delete, &xml_response("")).await; + let error = tokio::time::timeout(Duration::from_secs(3), probe.upload) + .await + .expect("failed upload returns without waiting for the other part") + .expect("upload task did not panic") + .expect_err("part upload failed"); + assert_eq!(error.code(), code::S3_PERMISSION_DENIED); + + assert_no_part_retry(&probe.listener, &mut pending).await; +} + +#[tokio::test] +async fn dropping_the_upload_future_stops_pending_part_retries() { + let probe = start_upload(MULTIPART_SIZE).await; + let (mut part, request, _) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + + probe.upload.abort(); + assert!( + probe + .upload + .await + .expect_err("parent task was aborted") + .is_cancelled() + ); + + assert_no_part_retry(&probe.listener, &mut part).await; +} + +#[tokio::test] +async fn multipart_upload_bounds_concurrency_and_completes_parts_in_order() { + let probe = start_upload(MULTIPART_SIZE * 4 + 1).await; + let mut pending = Vec::new(); + for _ in 0..4 { + let (part, request, body) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ")); + assert_eq!(body.len(), MULTIPART_SIZE); + pending.push(part); + } + assert!( + tokio::time::timeout(Duration::from_millis(150), probe.listener.accept()) + .await + .is_err(), + "a fifth part must wait for one of the four active uploads" + ); + respond(&mut pending.pop().unwrap(), part_response()).await; + let (mut last, request, body) = accept_request(&probe.listener).await; + assert!(request.starts_with("PUT ") && request.contains("partNumber=5")); + assert_eq!(body.len(), 1); + respond(&mut last, part_response()).await; + for mut part in pending.into_iter().rev() { + respond(&mut part, part_response()).await; + } + let (mut complete, request, body) = accept_request(&probe.listener).await; + assert!(request.starts_with("POST ") && request.contains("uploadId=")); + let body = String::from_utf8(body).unwrap(); + let positions: Vec<_> = (1..=5) + .map(|part| { + body.find(&format!("{part}")) + .unwrap() + }) + .collect(); + assert!(positions.windows(2).all(|pair| pair[0] < pair[1])); + respond( + &mut complete, + &xml_response( + "\"probe-part\"", + ), + ) + .await; + tokio::time::timeout(Duration::from_secs(3), probe.upload) + .await + .expect("multipart upload completes") + .expect("upload task did not panic") + .expect("all parts uploaded"); +} + +struct UploadProbe { + listener: TcpListener, + cancellation: CancellationToken, + upload: tokio::task::JoinHandle>, + _directory: tempfile::TempDir, +} + +async fn start_upload(size: usize) -> UploadProbe { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("archive.dump"); + let file = tokio::fs::File::create(&source).await.unwrap(); + file.set_len(size as u64).await.unwrap(); + drop(file); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let storage = BackupStorageConfig { + storage_revision: 1, + endpoint: format!("http://{}", listener.local_addr().unwrap()), + region: "auto".into(), + bucket: "backup-bucket".into(), + access_key_id: "test-access-key".into(), + secret_access_key: SecretString::from("test-secret-key"), + prefix: "backups".into(), + force_path_style: true, + }; + let cancellation = CancellationToken::new(); + let request = UploadObjectRequest { + object_key: "backups/archive.dump".into(), + source, + metadata: BackupObjectMetadata::new( + "backup_test".into(), + "a".repeat(64), + Utc::now(), + size as u64, + ), + cancellation: cancellation.clone(), + }; + let upload = tokio::spawn(async move { + S3ObjectStoreAdapter::new() + .upload_file(&storage, request) + .await + }); + let (mut create, request, _) = accept_request(&listener).await; + assert!(request.starts_with("POST ") && request.contains("uploads")); + respond( + &mut create, + &xml_response( + "probe-upload", + ), + ) + .await; + UploadProbe { + listener, + cancellation, + upload, + _directory: directory, + } +} + +async fn accept_request(listener: &TcpListener) -> (TcpStream, String, Vec) { + tokio::time::timeout(Duration::from_secs(3), async { + let (mut stream, _) = listener.accept().await.expect("accept S3 request"); + let (request, body) = read_http_request(&mut stream) + .await + .expect("read S3 request"); + (stream, request, body) + }) + .await + .expect("expected S3 request arrives without waiting for another part") +} + +async fn respond(stream: &mut TcpStream, response: &str) { + stream.write_all(response.as_bytes()).await.unwrap(); + stream.shutdown().await.unwrap(); +} + +async fn assert_no_part_retry(listener: &TcpListener, pending: &mut TcpStream) { + // 旧请求可能已被关闭，写入失败是正常撤销结果；仍活动的 SDK 调用收到 503 后会重试 + let _ = pending + .write_all(error_response("503 Service Unavailable", "SlowDown").as_bytes()) + .await; + let _ = pending.shutdown().await; + assert!( + tokio::time::timeout(Duration::from_secs(2), listener.accept()) + .await + .is_err(), + "a retired part must not issue another UploadPart request" + ); +} + +fn part_response() -> &'static str { + "HTTP/1.1 200 OK\r\ncontent-length: 0\r\nconnection: close\r\netag: \"probe-part\"\r\n\r\n" +} + +fn error_response(status: &str, code: &str) -> String { + let body = format!("{code}probe error"); + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/xml\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ) +} + async fn serve_multipart_probe(listener: TcpListener) -> io::Result { let mut uploaded_len = None; for _ in 0..3 { diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs index 2704c998a..6fcbba9e2 100644 --- a/backend/crates/gateway-store/tests/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -76,6 +76,41 @@ fn context() -> MutationContext { } } +#[tokio::test] +async fn unopened_future_budget_windows_do_not_report_usage_or_reset_times() { + let Some(database) = TestDatabase::create("budget_unopened_future").await else { + return; + }; + seed(&database, "key", "10", "20").await; + PgClientBudgetStore::new(database.pool.clone()) + .settle(charge("key", "before", "3")) + .await + .unwrap(); + // 起止相等表示窗口未开启；未来边界确定性覆盖应用时钟领先数据库的场景 + sqlx::query( + "update client_key_budget_windows set + daily_start = now() + interval '1 hour', daily_end = now() + interval '1 hour', + weekly_start = now() + interval '1 hour', weekly_end = now() + interval '1 hour' + where client_api_key_id = $1", + ) + .bind("key") + .execute(&database.pool) + .await + .unwrap(); + + let budget = status(&database, "key").await; + assert_eq!( + ( + budget.daily_used_usd.canonical(), + budget.weekly_used_usd.canonical(), + budget.daily_resets_at, + budget.weekly_resets_at, + ), + ("0".to_owned(), "0".to_owned(), None, None) + ); + database.close().await; +} + #[tokio::test] async fn manual_reset_clears_selected_windows_and_preserves_policy_and_history() { let Some(database) = TestDatabase::create("budget_manual_reset").await else { diff --git a/backend/crates/gateway-store/tests/postgres/connection.rs b/backend/crates/gateway-store/tests/postgres/connection.rs index 9fd564814..3ef4a780c 100644 --- a/backend/crates/gateway-store/tests/postgres/connection.rs +++ b/backend/crates/gateway-store/tests/postgres/connection.rs @@ -1,7 +1,9 @@ -//! 验证实际基础设施错误在 Store 边界保留来源而不进入普通格式化 +//! PostgreSQL 连接错误、迁移应用与事务归属测试 use std::error::Error as _; +use super::*; + #[tokio::test] async fn invalid_postgres_options_preserve_the_sqlx_source() { let error = gateway_store::postgres::connect_and_migrate( @@ -21,3 +23,188 @@ async fn invalid_postgres_options_preserve_the_sqlx_source() { ); assert!(!format!("{error:?} {error}").contains("invalid-port")); } + +#[tokio::test] +async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly() { + let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else { + return; + }; + let database = format!("cpr_store_migrator_{}", Uuid::new_v4().simple()); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("connect migration test PostgreSQL"); + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "create database \"{database}\"" + ))) + .execute(&admin) + .await + .expect("create migration test database"); + + let isolated_url = PgConnectOptions::from_str(&database_url) + .expect("parse migration test PostgreSQL URL") + .database(&database) + .to_url_lossy() + .to_string(); + let pool_config = gateway_store::StorePoolConfig::default(); + let first = connect_and_migrate(&isolated_url, pool_config) + .await + .expect("apply migrations through production migrator"); + let session_settings = sqlx::query_as::<_, (String, i64, i64, i64)>( + "select current_setting('application_name'), + extract(epoch from current_setting('statement_timeout')::interval)::bigint, + extract(epoch from current_setting('lock_timeout')::interval)::bigint, + extract(epoch from current_setting( + 'idle_in_transaction_session_timeout' + )::interval)::bigint", + ) + .fetch_one(&first) + .await + .expect("load runtime PostgreSQL session settings"); + let timezone: String = sqlx::query_scalar("show time zone") + .fetch_one(&first) + .await + .unwrap(); + assert_eq!( + timezone, "UTC", + "runtime sessions use an explicit technical time basis" + ); + let first_tables = sqlx::query_scalar::<_, String>( + "select table_name + from information_schema.tables + where table_schema = 'public' + order by table_name", + ) + .fetch_all(&first) + .await + .expect("load migrated tables"); + first.close().await; + + let second = connect_and_migrate(&isolated_url, pool_config) + .await + .expect("reopen database through production migrator"); + let migration_count = + sqlx::query_scalar::<_, i64>("select count(*) from _sqlx_migrations where success") + .fetch_one(&second) + .await + .expect("count successful migrations"); + let response_id_types = sqlx::query_scalar::<_, String>( + "select data_type + from information_schema.columns + where table_schema = 'public' + and table_name = 'model_requests' + and column_name in ('client_response_id', 'upstream_response_id') + order by column_name", + ) + .fetch_all(&second) + .await + .expect("load opaque response ID column types"); + let raw_response_id_index_exists = sqlx::query_scalar::<_, bool>( + "select exists ( + select 1 + from pg_indexes + where schemaname = 'public' + and indexname = 'model_requests_client_response_uq' + )", + ) + .fetch_one(&second) + .await + .expect("check removed raw response ID index"); + let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>( + "select exists ( + select 1 from information_schema.columns + where table_schema = 'public' + and table_name = 'client_api_keys' + and column_name = 'provider_kind' + )", + ) + .fetch_one(&second) + .await + .expect("check removed client key provider column"); + let routing_history_columns = sqlx::query_scalar::<_, String>( + "select column_name from information_schema.columns + where table_schema = 'public' + and table_name = 'model_requests' + and column_name in ( + 'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot' + ) + order by column_name", + ) + .fetch_all(&second) + .await + .expect("load routing history columns"); + second.close().await; + + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "drop database \"{database}\" with (force)" + ))) + .execute(&admin) + .await + .expect("drop migration test database"); + admin.close().await; + + assert_eq!( + first_tables, + [ + "_sqlx_migrations", + "account_group_accounts", + "account_groups", + "admin_audit_events", + "admin_users", + "authorization_receipts", + "backup_records", + "backup_settings", + "client_api_key_groups", + "client_api_keys", + "client_key_budget_windows", + "client_key_charge_events", + "model_requests", + "ops_events", + "outbound_proxies", + "plugin_artifact_credentials", + "plugin_artifact_platforms", + "plugin_artifacts", + "plugin_group_resources", + "plugin_instance_secrets", + "plugin_instances", + "plugin_key_resources", + "plugin_source_credentials", + "plugin_state_generations", + "plugin_state_records", + "plugin_update_sources", + "plugin_version_configurations", + "provider_accounts", + "runtime_settings", + ] + ); + assert_eq!(session_settings, ("codex-proxy-rs".to_owned(), 30, 5, 30)); + assert_eq!( + migration_count, + i64::try_from(TEST_MIGRATOR.iter().count()) + .expect("migration count fits PostgreSQL bigint") + ); + assert_eq!(response_id_types, ["bytea", "bytea"]); + assert!(!raw_response_id_index_exists); + assert!(!legacy_key_provider_column_exists); + assert_eq!( + routing_history_columns, + [ + "routing_group_names_snapshot", + "routing_group_refs", + "routing_scope", + ] + ); +} + +#[test] +fn migrations_should_leave_transaction_ownership_to_sqlx() { + let transaction_statements = TEST_MIGRATOR + .iter() + .flat_map(|migration| migration.sql.as_str().lines()) + .map(str::trim) + .filter(|line| matches!(*line, "begin;" | "commit;")) + .count(); + + assert_eq!(transaction_statements, 0); +} diff --git a/backend/crates/gateway-store/tests/postgres/control_plane.rs b/backend/crates/gateway-store/tests/postgres/control_plane.rs new file mode 100644 index 000000000..53dd10cd4 --- /dev/null +++ b/backend/crates/gateway-store/tests/postgres/control_plane.rs @@ -0,0 +1,104 @@ +//! 控制面替换事务的修订冲突与设置保留测试 + +use chrono::Utc; + +use super::{TestDatabase, runtime_settings::settings_with_margin}; + +#[tokio::test] +async fn control_plane_replacement_commits_one_writer_per_revision_and_preserves_newer_settings() { + use gateway_store::postgres::{ + AdminAuditActorKind, AdminAuditEvent, ControlPlaneReplacement, ControlPlaneRepository, + PgControlPlaneRepository, + }; + use gateway_store::{ConflictKind, StoreError}; + let Some(database) = TestDatabase::create("settings_compare_replace").await else { + return; + }; + let repository = PgControlPlaneRepository::new(database.pool.clone()); + let revision = repository + .load_control_plane() + .await + .unwrap() + .settings + .config_revision; + let replacement = |id: &str, margin| ControlPlaneReplacement { + expected_revision: revision, + settings: settings_with_margin(margin), + audit: AdminAuditEvent { + id: id.into(), + actor_kind: AdminAuditActorKind::System, + actor_admin_user_id: None, + actor_ref: "system".into(), + admin_request_id: Some(id.into()), + action: "settings.replace".into(), + entity_kind: "runtime_settings".into(), + entity_ref: "1".into(), + config_revision: None, + changed_fields: vec!["refresh_margin_seconds".into()], + created_at: Utc::now(), + }, + }; + let (first, second) = tokio::join!( + repository.replace_control_plane(replacement("first", 1800)), + repository.replace_control_plane(replacement("second", 7200)), + ); + let (saved, conflict) = match (first, second) { + (Ok(saved), Err(error)) | (Err(error), Ok(saved)) => (saved.settings, error), + other => panic!("expected exactly one successful transaction: {other:?}"), + }; + assert!(matches!( + conflict, + StoreError::Conflict { + kind: ConflictKind::StaleRevision, + .. + } + )); + let current = repository.load_control_plane().await.unwrap().settings; + assert_eq!(current.config_revision.get(), revision.get() + 1); + assert_eq!( + current.values.refresh_margin_seconds, + saved.values.refresh_margin_seconds + ); + let audit_count: i64 = sqlx::query_scalar( + "select count(*) from admin_audit_events where action = 'settings.replace'", + ) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(audit_count, 1); + // API Key 更新也推进相同版本，旧设置快照不能复活已经替换的 Key + let mut key_audit = replacement("key", 3600).audit; + key_audit.action = "settings.admin_key".into(); + repository + .replace_admin_api_key(Some("new-test-key".into()), key_audit) + .await + .unwrap(); + let mut stale = replacement("stale", 3600); + stale.expected_revision = saved.config_revision; + assert!(matches!( + repository.replace_control_plane(stale).await, + Err(StoreError::Conflict { + kind: ConflictKind::StaleRevision, + .. + }) + )); + let mut fresh = replacement("fresh", 3600); + fresh.expected_revision = repository + .load_control_plane() + .await + .unwrap() + .settings + .config_revision; + repository.replace_control_plane(fresh).await.unwrap(); + assert_eq!( + repository + .load_control_plane() + .await + .unwrap() + .settings + .admin_api_key + .as_deref(), + Some("new-test-key") + ); + database.close().await; +} diff --git a/backend/crates/gateway-store/tests/postgres/mod.rs b/backend/crates/gateway-store/tests/postgres/mod.rs index 1c26a49ef..2bb898ece 100644 --- a/backend/crates/gateway-store/tests/postgres/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/mod.rs @@ -1,4 +1,4 @@ -//! PostgreSQL 适配测试入口，以及迁移应用和事务归属测试 +//! PostgreSQL 适配测试入口与隔离数据库夹具 use std::{str::FromStr, time::Duration}; @@ -19,6 +19,7 @@ mod backup; mod client_budgets; mod client_keys; mod connection; +mod control_plane; mod execution; mod execution_buffer; mod health; @@ -110,188 +111,3 @@ impl TestDatabase { self.admin.close().await; } } - -#[tokio::test] -async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly() { - let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else { - return; - }; - let database = format!("cpr_store_migrator_{}", Uuid::new_v4().simple()); - let admin = PgPoolOptions::new() - .max_connections(1) - .connect(&database_url) - .await - .expect("connect migration test PostgreSQL"); - sqlx::raw_sql(sqlx::AssertSqlSafe(format!( - "create database \"{database}\"" - ))) - .execute(&admin) - .await - .expect("create migration test database"); - - let isolated_url = PgConnectOptions::from_str(&database_url) - .expect("parse migration test PostgreSQL URL") - .database(&database) - .to_url_lossy() - .to_string(); - let pool_config = gateway_store::StorePoolConfig::default(); - let first = connect_and_migrate(&isolated_url, pool_config) - .await - .expect("apply migrations through production migrator"); - let session_settings = sqlx::query_as::<_, (String, i64, i64, i64)>( - "select current_setting('application_name'), - extract(epoch from current_setting('statement_timeout')::interval)::bigint, - extract(epoch from current_setting('lock_timeout')::interval)::bigint, - extract(epoch from current_setting( - 'idle_in_transaction_session_timeout' - )::interval)::bigint", - ) - .fetch_one(&first) - .await - .expect("load runtime PostgreSQL session settings"); - let timezone: String = sqlx::query_scalar("show time zone") - .fetch_one(&first) - .await - .unwrap(); - assert_eq!( - timezone, "UTC", - "runtime sessions use an explicit technical time basis" - ); - let first_tables = sqlx::query_scalar::<_, String>( - "select table_name - from information_schema.tables - where table_schema = 'public' - order by table_name", - ) - .fetch_all(&first) - .await - .expect("load migrated tables"); - first.close().await; - - let second = connect_and_migrate(&isolated_url, pool_config) - .await - .expect("reopen database through production migrator"); - let migration_count = - sqlx::query_scalar::<_, i64>("select count(*) from _sqlx_migrations where success") - .fetch_one(&second) - .await - .expect("count successful migrations"); - let response_id_types = sqlx::query_scalar::<_, String>( - "select data_type - from information_schema.columns - where table_schema = 'public' - and table_name = 'model_requests' - and column_name in ('client_response_id', 'upstream_response_id') - order by column_name", - ) - .fetch_all(&second) - .await - .expect("load opaque response ID column types"); - let raw_response_id_index_exists = sqlx::query_scalar::<_, bool>( - "select exists ( - select 1 - from pg_indexes - where schemaname = 'public' - and indexname = 'model_requests_client_response_uq' - )", - ) - .fetch_one(&second) - .await - .expect("check removed raw response ID index"); - let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>( - "select exists ( - select 1 from information_schema.columns - where table_schema = 'public' - and table_name = 'client_api_keys' - and column_name = 'provider_kind' - )", - ) - .fetch_one(&second) - .await - .expect("check removed client key provider column"); - let routing_history_columns = sqlx::query_scalar::<_, String>( - "select column_name from information_schema.columns - where table_schema = 'public' - and table_name = 'model_requests' - and column_name in ( - 'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot' - ) - order by column_name", - ) - .fetch_all(&second) - .await - .expect("load routing history columns"); - second.close().await; - - sqlx::raw_sql(sqlx::AssertSqlSafe(format!( - "drop database \"{database}\" with (force)" - ))) - .execute(&admin) - .await - .expect("drop migration test database"); - admin.close().await; - - assert_eq!( - first_tables, - [ - "_sqlx_migrations", - "account_group_accounts", - "account_groups", - "admin_audit_events", - "admin_users", - "authorization_receipts", - "backup_records", - "backup_settings", - "client_api_key_groups", - "client_api_keys", - "client_key_budget_windows", - "client_key_charge_events", - "model_requests", - "ops_events", - "outbound_proxies", - "plugin_artifact_credentials", - "plugin_artifact_platforms", - "plugin_artifacts", - "plugin_group_resources", - "plugin_instance_secrets", - "plugin_instances", - "plugin_key_resources", - "plugin_source_credentials", - "plugin_state_generations", - "plugin_state_records", - "plugin_update_sources", - "plugin_version_configurations", - "provider_accounts", - "runtime_settings", - ] - ); - assert_eq!(session_settings, ("codex-proxy-rs".to_owned(), 30, 5, 30)); - assert_eq!( - migration_count, - i64::try_from(TEST_MIGRATOR.iter().count()) - .expect("migration count fits PostgreSQL bigint") - ); - assert_eq!(response_id_types, ["bytea", "bytea"]); - assert!(!raw_response_id_index_exists); - assert!(!legacy_key_provider_column_exists); - assert_eq!( - routing_history_columns, - [ - "routing_group_names_snapshot", - "routing_group_refs", - "routing_scope", - ] - ); -} - -#[test] -fn migrations_should_leave_transaction_ownership_to_sqlx() { - let transaction_statements = TEST_MIGRATOR - .iter() - .flat_map(|migration| migration.sql.as_str().lines()) - .map(str::trim) - .filter(|line| matches!(*line, "begin;" | "commit;")) - .count(); - - assert_eq!(transaction_statements, 0); -} diff --git a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs index 2c331d64f..8f2adee0b 100644 --- a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs +++ b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs @@ -9,7 +9,7 @@ use gateway_store::postgres::{ use super::TestDatabase; -fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate { +pub(super) fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate { RuntimeSettingsUpdate { request_profile_updates: BTreeMap::new(), rotation_strategy: "smart".to_owned(), @@ -942,105 +942,6 @@ async fn request_profile_projection_is_revision_consistent_and_includes_key_over database.close().await; } -#[tokio::test] -async fn control_plane_replacement_commits_one_writer_per_revision_and_preserves_newer_settings() { - use gateway_store::postgres::{ - AdminAuditActorKind, AdminAuditEvent, ControlPlaneReplacement, ControlPlaneRepository, - PgControlPlaneRepository, - }; - use gateway_store::{ConflictKind, StoreError}; - let Some(database) = TestDatabase::create("settings_compare_replace").await else { - return; - }; - let repository = PgControlPlaneRepository::new(database.pool.clone()); - let revision = repository - .load_control_plane() - .await - .unwrap() - .settings - .config_revision; - let replacement = |id: &str, margin| ControlPlaneReplacement { - expected_revision: revision, - settings: settings_with_margin(margin), - audit: AdminAuditEvent { - id: id.into(), - actor_kind: AdminAuditActorKind::System, - actor_admin_user_id: None, - actor_ref: "system".into(), - admin_request_id: Some(id.into()), - action: "settings.replace".into(), - entity_kind: "runtime_settings".into(), - entity_ref: "1".into(), - config_revision: None, - changed_fields: vec!["refresh_margin_seconds".into()], - created_at: Utc::now(), - }, - }; - let (first, second) = tokio::join!( - repository.replace_control_plane(replacement("first", 1800)), - repository.replace_control_plane(replacement("second", 7200)), - ); - let (saved, conflict) = match (first, second) { - (Ok(saved), Err(error)) | (Err(error), Ok(saved)) => (saved.settings, error), - other => panic!("expected exactly one successful transaction: {other:?}"), - }; - assert!(matches!( - conflict, - StoreError::Conflict { - kind: ConflictKind::StaleRevision, - .. - } - )); - let current = repository.load_control_plane().await.unwrap().settings; - assert_eq!(current.config_revision.get(), revision.get() + 1); - assert_eq!( - current.values.refresh_margin_seconds, - saved.values.refresh_margin_seconds - ); - let audit_count: i64 = sqlx::query_scalar( - "select count(*) from admin_audit_events where action = 'settings.replace'", - ) - .fetch_one(&database.pool) - .await - .unwrap(); - assert_eq!(audit_count, 1); - // API Key 更新也推进相同版本，旧设置快照不能复活已经替换的 Key - let mut key_audit = replacement("key", 3600).audit; - key_audit.action = "settings.admin_key".into(); - repository - .replace_admin_api_key(Some("new-test-key".into()), key_audit) - .await - .unwrap(); - let mut stale = replacement("stale", 3600); - stale.expected_revision = saved.config_revision; - assert!(matches!( - repository.replace_control_plane(stale).await, - Err(StoreError::Conflict { - kind: ConflictKind::StaleRevision, - .. - }) - )); - let mut fresh = replacement("fresh", 3600); - fresh.expected_revision = repository - .load_control_plane() - .await - .unwrap() - .settings - .config_revision; - repository.replace_control_plane(fresh).await.unwrap(); - assert_eq!( - repository - .load_control_plane() - .await - .unwrap() - .settings - .admin_api_key - .as_deref(), - Some("new-test-key") - ); - database.close().await; -} - #[tokio::test] async fn warmup_cursor_survives_restart_and_settings_updates_without_rewinding() { use gateway_core::{provider_ports::ProviderRuntimePolicyPort, time::DeploymentTimeZone}; diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index 739ec6077..4095487e7 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -6,6 +6,7 @@ use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAl use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; +use crate::request_identity::{account_session_with_headers, follows_session_with_headers}; use crate::transport::protocol::responses::CodexResponsesRequest; use crate::transport::request::derive_conversation_anchor; @@ -171,80 +172,6 @@ pub(crate) fn derive_endpoint_affinity_with_headers( }) } -/// 只有能识别出的后代线程禁止迁移会话；没有身份的普通客户端沿用现有调度 -pub(crate) fn follows_session_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> bool { - let metadata = headers - .iter() - .find(|h| h.name().eq_ignore_ascii_case("x-codex-turn-metadata")) - .and_then(|h| std::str::from_utf8(h.value()).ok()); - let thread = account_thread_with_headers(body, context, headers); - if let Some((session, thread)) = - account_session_with_headers(body, context, headers).zip(thread) - { - return session != thread; - } - gateway_protocol::openai::codex_responses_request_semantics_with_turn_metadata( - body, - metadata.or_else(|| context.get("turn_metadata").and_then(Value::as_str)), - ) - .subagent_kind - .is_some() -} - -fn account_thread_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> Option { - let header = |name: &str| { - headers - .iter() - .find(|h| h.name().eq_ignore_ascii_case(name)) - .and_then(|h| std::str::from_utf8(h.value()).ok()) - }; - header("x-codex-turn-metadata") - .and_then(gateway_protocol::openai::turn_metadata_thread_id) - .or_else(|| { - header("thread-id") - .and_then(|value| non_empty(Some(value))) - .map(str::to_owned) - }) - .or_else(|| gateway_protocol::openai::codex_account_thread_id(body, context)) -} - -/// 中间件最终请求头按实际传输的覆盖语义参与身份复验 -pub(crate) fn account_session_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> Option { - let header = |name: &str| { - headers - .iter() - .find(|h| h.name().eq_ignore_ascii_case(name)) - .and_then(|h| std::str::from_utf8(h.value()).ok()) - }; - if let Some(session) = - header("x-codex-turn-metadata").and_then(gateway_protocol::openai::turn_metadata_session_id) - { - return Some(session); - } - let mut context = context.clone(); - for (name, field) in [ - ("session-id", "session_id"), - ("x-codex-turn-metadata", "turn_metadata"), - ] { - if let Some(value) = header(name) { - context.insert(field.to_owned(), Value::String(value.to_owned())); - } - } - gateway_protocol::openai::codex_account_session_id(body, &context) -} - pub(crate) fn derive_live_session_affinity( request: &gateway_core::operation::ProviderHttpRequest, headers: &[gateway_core::engine::middleware::MiddlewareHeader], @@ -376,24 +303,3 @@ pub(crate) fn derive_turn_alias( let turn_id = non_empty(Some(turn_id))?; opaque_affinity_key("client-turn", &format!("{}\0{turn_id}", client.as_str())) } - -pub(crate) fn turn_id_with_headers( - body: &Map, - context: &Map, - headers: &[gateway_core::engine::middleware::MiddlewareHeader], -) -> Option { - headers - .iter() - .find(|header| header.name().eq_ignore_ascii_case("x-codex-turn-metadata")) - .and_then(|header| std::str::from_utf8(header.value()).ok()) - .and_then(gateway_protocol::openai::turn_metadata_turn_id) - .or_else(|| { - headers - .iter() - .find(|header| header.name().eq_ignore_ascii_case("x-client-turn-id")) - .and_then(|header| std::str::from_utf8(header.value()).ok()) - .and_then(|value| non_empty(Some(value))) - .map(str::to_owned) - }) - .or_else(|| gateway_protocol::openai::codex_turn_id(body, context)) -} diff --git a/backend/crates/providers/openai/src/credential/api_key.rs b/backend/crates/providers/openai/src/credential/api_key.rs index e2e4e3eca..7678cc2d5 100644 --- a/backend/crates/providers/openai/src/credential/api_key.rs +++ b/backend/crates/providers/openai/src/credential/api_key.rs @@ -5,7 +5,7 @@ use std::fmt; use secrecy::SecretString; use serde::{Deserialize, Serialize}; -pub const CODEX_AUTHENTICATION_KIND_API_KEY: &str = "api_key"; +pub(super) use crate::CODEX_AUTHENTICATION_KIND_API_KEY; use super::types::ResponsesTransport; diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index eb9f03a27..aad85e27d 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -745,7 +745,7 @@ impl CodexCredentialCatalogService { .map_err(|error| CodexCredentialCatalogError::Upstream { detail: error.to_string(), })? - .with_authentication(&credential.authentication) + .with_responses_api_base_url(credential.authentication.responses_api_base_url()) .fetch_models_with_context( CodexRequestContext::auxiliary( authorization.expose_secret(), diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index 435044416..5f85bd933 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -17,19 +17,17 @@ pub mod token_client; mod types; pub(crate) use affinity::{ - CodexSessionAffinity, account_session_with_headers, derive_codex_cyber_policy_session_key, + CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_endpoint_session_affinity, derive_codex_session_affinity, derive_codex_transport_key, derive_endpoint_affinity_with_headers, derive_live_session_affinity, derive_previous_response_id_hash, derive_turn_alias, - follows_session_with_headers, turn_id_with_headers, }; pub(crate) use oauth::oauth_owner_ref; pub(crate) use types::parse_access_token_expiration; -pub use api_key::{ - ApiKeyAuthentication, ApiKeyConfiguration, ApiKeyCredentialData, - CODEX_AUTHENTICATION_KIND_API_KEY, -}; +pub(crate) use crate::CODEX_AUTHENTICATION_KIND_API_KEY; + +pub use api_key::{ApiKeyAuthentication, ApiKeyConfiguration, ApiKeyCredentialData}; pub use admin::{ CodexCprExportDocument, CodexCredentialAdmin, CodexCredentialAdminError, diff --git a/backend/crates/providers/openai/src/credential/quota/document.rs b/backend/crates/providers/openai/src/credential/quota/document.rs index 8301a07e8..1ab272832 100644 --- a/backend/crates/providers/openai/src/credential/quota/document.rs +++ b/backend/crates/providers/openai/src/credential/quota/document.rs @@ -172,3 +172,75 @@ fn normalize_limit_id(value: &str) -> Option { (!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)) .then(|| value.to_ascii_lowercase().replace('-', "_")) } + +/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位 +/// +/// 生产响应头可能携带 `used_percent=0`、零时长和空 reset 的占位，`/usage` 也会 +/// 返回 `secondary_window: null` +/// 正常被动同步保留该响应事实，Admin 展示会将其 +/// 隐藏；但主动 `/usage` 刷新或 402 确认投影会把一个存在但无事实的字段写成 +/// 100%，从而显示并不存在的“次级额度” +/// 因此这些非被动写入口在落库前移除 +/// 无事实值 +/// 带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按 +/// 原有额度逻辑保留 +pub(super) fn normalize_quota_window_placeholders( + mut quota: Map, +) -> Map { + quota = canonicalize_rate_limit_document(quota); + if let Some(rate_limit) = quota + .get_mut(RATE_LIMITS_BY_LIMIT_ID) + .and_then(Value::as_object_mut) + .and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID)) + .and_then(Value::as_object_mut) + { + drop_secondary_window_placeholder(rate_limit); + } + quota +} + +fn drop_secondary_window_placeholder(rate_limit: &mut Map) { + let placeholder = rate_limit.get("secondary_window").is_some_and(|window| { + window.is_null() + || window + .as_object() + .is_some_and(secondary_window_is_placeholder) + }); + if placeholder { + rate_limit.remove("secondary_window"); + } +} + +fn secondary_window_is_placeholder(window: &Map) -> bool { + window.iter().all(|(field, value)| match field.as_str() { + "used_percent" => value + .as_f64() + .is_some_and(|used_percent| used_percent.is_finite() && used_percent == 0.0), + "limit_reached" => value.as_bool() == Some(false), + _ => false, + }) +} + +/// 上游额度可确认套餐变更；同族泛化值不能丢弃 JWT 已给出的具体 SKU +pub(super) fn observed_account_plan( + current: Option<&str>, + observed: Option<&str>, +) -> Option { + let plan = observed?.trim().to_ascii_lowercase(); + if plan.is_empty() || plan == "unknown" { + return None; + } + // 套餐族沿用官方 codex_protocol::account::PlanType 的分类 + let current = current.unwrap_or_default().trim().to_ascii_lowercase(); + let generalized = matches!( + (plan.as_str(), current.as_str()), + ( + "team", + "self_serve_business_prolite" | "self_serve_business_usage_based" + ) | ( + "business", + "ent26" | "enterprise_cbp_automation" | "enterprise_cbp_usage_based" + ) | ("edu" | "education", "edu_plus" | "edu_pro") + ); + (!generalized).then_some(plan) +} diff --git a/backend/crates/providers/openai/src/credential/quota/mod.rs b/backend/crates/providers/openai/src/credential/quota/mod.rs index a0bd5ad76..387c7119f 100644 --- a/backend/crates/providers/openai/src/credential/quota/mod.rs +++ b/backend/crates/providers/openai/src/credential/quota/mod.rs @@ -1,4 +1,4 @@ -//! Codex quota 服务编排：权威额度事实、展示投影、主动/被动同步与 429 冷却 +//! Codex quota 门面与唯一服务状态；实现按读取、排程、预热和重置消费组织 //! //! - [`document`]：上游 `/usage` 输入到多桶 map 的单向规范化 //! - [`snapshot`]：快照/窗口解析、聚合、O(1) 滚动与调度信号 @@ -7,14 +7,23 @@ mod document; pub(crate) mod evidence; +mod passive; mod recovery; +mod reset_credits; +mod scheduling; +mod service; pub(crate) mod snapshot; +mod warmup; + +pub use reset_credits::CodexResetCreditsError; +use reset_credits::ResetCreditLocks; +use scheduling::CodexQuotaSchedulingProjection; pub use snapshot::{ CodexAccountQuotaSnapshot, CodexQuotaFact, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole, parse_codex_quota_usage, }; -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet}; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant, SystemTime}; @@ -50,8 +59,8 @@ use crate::transport::{ use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; use super::security::CodexRuntimeCredential; use document::{ - DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId, - canonicalize_rate_limit_document, + DEFAULT_CODEX_LIMIT_ID, RateLimitSnapshotsByLimitId, normalize_quota_window_placeholders, + observed_account_plan, }; use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure}; use recovery::{RECOVERY_FIELD, reconcile_refresh}; @@ -123,74 +132,6 @@ pub struct CodexWarmupSummary { pub failed: u64, } -#[derive(Clone, Copy, PartialEq, Eq)] -enum WarmupTerminal { - Missing, - Completed, - Failed, -} - -fn observe_warmup_events(events: Vec, terminal: &mut WarmupTerminal) { - for event in events { - let parsed = serde_json::from_str::(&event.data).ok(); - let kind = parsed - .as_ref() - .and_then(|data| data.get("type")) - .and_then(Value::as_str) - .or(event.event.as_deref()); - match kind { - Some("response.completed") => { - // 只有明确完成的响应才算预热成功；失败事件即使随后出现完成帧也优先 - if *terminal != WarmupTerminal::Failed { - *terminal = if parsed.is_some() - && parsed - .as_ref() - .and_then(|data| data.pointer("/response/status")) - .and_then(Value::as_str) - .is_none_or(|status| status == "completed") - { - WarmupTerminal::Completed - } else { - WarmupTerminal::Failed - }; - } - } - Some("response.failed" | "response.incomplete" | "error") => { - *terminal = WarmupTerminal::Failed; - } - _ => {} - } - } -} - -async fn consume_warmup_sse( - response: &mut CodexBackendStreamingResponse, -) -> Result<(), &'static str> { - let mut decoder = SseEventDecoder::default(); - let mut terminal = WarmupTerminal::Missing; - let mut received_bytes = 0usize; - while let Some(chunk) = response.body.next().await { - let chunk = chunk.map_err(|_| "stream_read_failed")?; - received_bytes = received_bytes.saturating_add(chunk.len()); - if received_bytes > WARMUP_STREAM_MAX_BYTES { - return Err("stream_size_limit_exceeded"); - } - let events = decoder.push(&chunk).map_err(|_| "invalid_sse")?; - observe_warmup_events(events, &mut terminal); - if terminal != WarmupTerminal::Missing { - break; - } - } - if terminal == WarmupTerminal::Missing { - observe_warmup_events(decoder.finish().map_err(|_| "invalid_sse")?, &mut terminal); - } - match terminal { - WarmupTerminal::Completed => Ok(()), - WarmupTerminal::Missing => Err("completion_event_missing"), - WarmupTerminal::Failed => Err("response_failed"), - } -} - #[derive(Debug, Error)] pub enum CodexCredentialQuotaError { #[error("Codex quota response is invalid")] @@ -215,54 +156,6 @@ pub enum CodexCredentialQuotaError { }, } -/// 主动额度重置卡查询/消费失败 -#[derive(Error)] -pub enum CodexResetCreditsError { - #[error("Codex reset-credit credential data is invalid")] - InvalidCredentialData, - #[error("Codex OAuth access token must be refreshed before using reset credits")] - CredentialRefreshRequired { upstream_body: Option }, - #[error("Codex reset-credit account was not found")] - NotFound, - #[error("provider account store is unavailable: {detail}")] - Store { detail: String }, - #[error("Codex reset-credit upstream returned HTTP {status}")] - Upstream { - status: u16, - body: String, - retry_after_seconds: Option, - }, - #[error("Codex reset-credit query transport is unavailable")] - TransportUnavailable, - #[error("Codex reset-credit consume result is unknown")] - ConsumeResultUnknown, -} - -impl std::fmt::Debug for CodexResetCreditsError { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"), - Self::CredentialRefreshRequired { .. } => { - formatter.write_str("CredentialRefreshRequired { upstream_body: }") - } - Self::NotFound => formatter.write_str("NotFound"), - Self::Store { .. } => formatter.write_str("Store { detail: }"), - Self::Upstream { - status, - retry_after_seconds, - .. - } => formatter - .debug_struct("Upstream") - .field("status", status) - .field("body", &"") - .field("retry_after_seconds", retry_after_seconds) - .finish(), - Self::TransportUnavailable => formatter.write_str("TransportUnavailable"), - Self::ConsumeResultUnknown => formatter.write_str("ConsumeResultUnknown"), - } - } -} - impl From for CodexCredentialQuotaError { fn from(error: gateway_core::error::StoreError) -> Self { Self::Store { @@ -271,38 +164,6 @@ impl From for CodexCredentialQuotaError { } } -/// 客户端错误携带的 HTTP 状态码；传输失败等为 `None` -fn upstream_error_status(error: &CodexClientError) -> Option { - match error { - CodexClientError::Upstream { status, .. } => Some(status.as_u16()), - _ => None, - } -} - -/// 从上游错误体提取稳定错误码：优先 `/error/code`，其次 `/code` -/// -/// 只接受有界的 ASCII 标识，避免把自由文本作为错误码写入诊断 -fn upstream_error_code(error: &CodexClientError) -> Option { - let CodexClientError::Upstream { body, .. } = error else { - return None; - }; - let value = serde_json::from_str::(body).ok()?; - let code = value - .pointer("/error/code") - .or_else(|| value.pointer("/code")) - .and_then(Value::as_str)? - .trim(); - if code.is_empty() - || code.len() > 64 - || !code - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.')) - { - return None; - } - Some(code.to_owned()) -} - pub struct CodexCredentialQuotaService { repository: CodexCredentialRepository, store: Arc, @@ -315,7 +176,7 @@ pub struct CodexCredentialQuotaService { /// 冻结策略的短 TTL 缓存：失败路径热读，避免每个容量错误都查询设置 freeze_policy_cache: Mutex>, scheduling: CodexQuotaSchedulingProjection, - reset_consume_locks: Mutex>>>, + reset_consume_locks: ResetCreditLocks, /// 首查随机延迟采样器；测试注入固定值获得确定性节奏。 initial_sync_delays: CodexInitialSyncDelays, } @@ -327,1833 +188,13 @@ pub type CodexInitialSyncDelays = Arc Duration + Send + Sync>; /// 冻结策略缓存活跃期；过期后下一次容量错误重新读取运行时设置 const FREEZE_POLICY_CACHE_TTL: Duration = Duration::from_secs(30); -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum QuotaRefreshAuthority { - ObserveAccess, - PreserveAccess, -} - -#[derive(Clone, Default)] -struct CodexQuotaSchedulingProjection { - state: Arc>, - hydration: Arc>, -} - -#[derive(Default)] -struct CodexQuotaProjectionState { - next_version: u64, - entries: BTreeMap, - last_periodic_refresh_at: BTreeMap, - initial_refresh_not_before: BTreeMap, -} - -struct CodexQuotaRefreshAttempt { - monotonic_at: Instant, - wall_at: SystemTime, -} - -#[derive(Debug, Clone, Copy)] -struct CodexQuotaSchedulingEntry { - version: u64, - revision: CredentialRevision, - expires_at: Instant, - signals: Option, -} - -#[derive(Clone)] -struct CodexQuotaHydrationTarget { - account: ProviderAccount, - expected_version: Option, -} - -struct FetchedCodexQuota { - account: ProviderAccount, - value: Value, -} - struct PreparedCodexRuntimeCredential { account: ProviderAccount, credential: CodexRuntimeCredential, } -enum CodexQuotaFetchAttemptError { - InvalidCredential, - Upstream(CodexClientError), -} - -enum CodexQuotaFetchError { - InvalidCredential, - Upstream { - account: Box, - error: CodexClientError, - }, -} - -enum ResetCreditAttemptError { - InvalidCredential, - Upstream(CodexClientError), -} - -impl CodexQuotaSchedulingProjection { - fn invalidate(&self, account_ids: &[ProviderAccountId]) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - for account_id in account_ids { - state.entries.remove(account_id); - state.last_periodic_refresh_at.remove(account_id); - state.initial_refresh_not_before.remove(account_id); - } - } - - fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let now = Instant::now(); - accounts - .iter() - .filter_map(|account| { - let current = state.entries.get(account.id()); - current - .is_none_or(|entry| { - entry.revision != account.revision() || now >= entry.expires_at - }) - .then(|| CodexQuotaHydrationTarget { - account: account.clone(), - expected_version: current.map(|entry| entry.version), - }) - }) - .collect() - } - - fn observe(&self, snapshot: &CodexAccountQuotaSnapshot) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { - return false; - }; - self.replace( - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - scheduling_signals_from_snapshot(snapshot), - ); - true - } - - fn mark_unknown_if_unchanged(&self, target: &CodexQuotaHydrationTarget, ttl: Duration) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return; - } - insert_projection_entry( - &mut state, - target.account.id().clone(), - target.account.revision(), - ttl, - None, - ); - } - - fn observe_if_unchanged( - &self, - target: &CodexQuotaHydrationTarget, - snapshot: &CodexAccountQuotaSnapshot, - ) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { - return false; - }; - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return true; - } - insert_projection_entry( - &mut state, - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - scheduling_signals_from_snapshot(snapshot), - ); - true - } - - fn replace( - &self, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, - ) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - insert_projection_entry(&mut state, account_id, revision, ttl, signals); - } - - fn signals(&self, account: &ProviderAccount) -> Option { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state - .entries - .get(account.id()) - .filter(|entry| { - entry.revision == account.revision() && Instant::now() < entry.expires_at - }) - .and_then(|entry| entry.signals) - } - - fn reserve_periodic_refreshes( - &self, - accounts: Vec, - observed_snapshots: &BTreeMap, - now: SystemTime, - ) -> Vec { - let candidates = accounts - .into_iter() - .filter_map(|account| { - let snapshot = observed_snapshots.get(account.id()); - quota_refresh_candidate(account, snapshot, now) - }) - .collect::>(); - let candidate_ids = candidates - .iter() - .map(|(account, _)| account.id().clone()) - .collect::>(); - let refreshed_at = Instant::now(); - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state - .last_periodic_refresh_at - .retain(|account_id, _| candidate_ids.contains(account_id)); - - // 已耗尽账号首次立即复核，之后每 30 分钟复核，以发现官方提前重置 - // reset + 2 分钟额外触发一次复核，给上游重置留出传播时间 - // 正常账号仅在非零用量窗口经过宽限期后参与，未更新时复用周期节流 - let mut reserved = Vec::new(); - for (account, target_reset) in candidates { - if !periodic_quota_refresh_due(&state, account.id(), target_reset, now, refreshed_at) { - continue; - } - state.last_periodic_refresh_at.insert( - account.id().clone(), - CodexQuotaRefreshAttempt { - monotonic_at: refreshed_at, - wall_at: now, - }, - ); - reserved.push(account); - } - reserved - } - - /// 为尚无 quota 快照的账号分配随机首查时点，并返回本轮已到期的候选。 - /// - /// `quota_observed_at` 为空代表首次异步观察尚未成功；不另建同步状态表。 - /// 首次进入候选的账号抽一次随机启动延迟记为到期时刻，避免新导入账号池 - /// 在 worker 唤醒边界形成批量 `/usage` 突发；到期时刻跨轮保留，零延迟 - /// （含随机源退化）时保持首轮即查的旧行为。单轮数量仍受 - /// `INITIAL_QUOTA_SYNC_BATCH` 上限约束。 - fn reserve_initial_refreshes( - &self, - accounts: &[ProviderAccount], - observed_ids: &BTreeSet, - now: SystemTime, - draw_start_delay: &dyn Fn() -> Duration, - ) -> Vec { - let candidates = accounts - .iter() - .filter(|account| { - !observed_ids.contains(account.id()) && eligible_initial_quota_sync(account, now) - }) - .collect::>(); - let candidate_ids = candidates - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let monotonic_now = Instant::now(); - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - // 已产生快照或离开候选集的账号不再保留首查状态。 - state - .initial_refresh_not_before - .retain(|account_id, _| candidate_ids.contains(account_id)); - candidates - .into_iter() - .filter(|account| { - let due_at = state - .initial_refresh_not_before - .entry(account.id().clone()) - .or_insert_with(|| monotonic_now + draw_start_delay()); - *due_at <= monotonic_now - }) - .take(INITIAL_QUOTA_SYNC_BATCH) - .cloned() - .collect() - } - - /// 首查尝试后仍无快照的账号重抽随机延迟，避免失败账号回到固定 30s 重试节奏。 - /// - /// 重排同时随机化每轮的候选集合，避免持续失败时按账号顺序只重试同一批。 - fn defer_failed_initial_refreshes( - &self, - account_ids: &BTreeSet, - draw_start_delay: &dyn Fn() -> Duration, - ) { - if account_ids.is_empty() { - return; - } - let now = Instant::now(); - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - for account_id in account_ids { - // 账号是否仍在首查路径由下一轮 reserve 按候选集剪枝，这里无需区分。 - state - .initial_refresh_not_before - .insert(account_id.clone(), now + draw_start_delay()); - } - } -} - -fn quota_refresh_candidate( - account: ProviderAccount, - snapshot: Option<&CodexAccountQuotaSnapshot>, - now: SystemTime, -) -> Option<(ProviderAccount, Option)> { - if !eligible_periodic_quota_refresh(&account, now) { - return None; - } - if account.quota().is_exhausted() { - let reset_at = account.quota().reset_at(); - return Some((account, reset_at)); - } - // 正常账号首次进入周期候选也要等待宽限期，不能由“无刷新历史”绕过 - let snapshot = snapshot?; - let expired_window_reset = snapshot - .windows() - .iter() - .filter(|window| { - window.used_percent().is_some_and(|used| used > 0.0) || window.limit_reached() - }) - .filter_map(CodexQuotaWindow::reset_at) - .map(SystemTime::from) - .filter(|reset| { - reset - .checked_add(QUOTA_RESET_GRACE) - .is_some_and(|due_at| due_at <= now) - }) - .min(); - expired_window_reset.map(|reset_at| (account, Some(reset_at))) -} - -fn periodic_quota_refresh_due( - state: &CodexQuotaProjectionState, - account_id: &ProviderAccountId, - reset_at: Option, - now: SystemTime, - monotonic_now: Instant, -) -> bool { - state - .last_periodic_refresh_at - .get(account_id) - .is_none_or(|last| { - monotonic_now.saturating_duration_since(last.monotonic_at) - >= PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL - || reset_at - .and_then(|reset| reset.checked_add(QUOTA_RESET_GRACE)) - // 已在该边界之后复核过时回到周期重试，避免过期 reset 每轮触发 - .is_some_and(|due_at| last.wall_at < due_at && due_at <= now) - }) -} - -fn insert_projection_entry( - state: &mut CodexQuotaProjectionState, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, -) { - state.next_version = state.next_version.saturating_add(1); - state.entries.insert( - account_id, - CodexQuotaSchedulingEntry { - version: state.next_version, - revision, - expires_at: Instant::now() + ttl, - signals, - }, - ); -} - -impl CodexCredentialQuotaService { - pub fn new( - repository: CodexCredentialRepository, - profile: CodexWireProfileState, - http: Client, - base_url: String, - cooldowns: Arc, - leases: Arc, - runtime_policy: Arc, - ) -> Self { - Self { - store: Arc::clone(repository.store()), - repository, - profile, - http, - base_url, - cooldowns, - leases, - runtime_policy, - freeze_policy_cache: Mutex::new(None), - scheduling: CodexQuotaSchedulingProjection::default(), - reset_consume_locks: Mutex::new(HashMap::new()), - initial_sync_delays: Arc::new(|| { - crate::jitter::uniform_delay( - crate::jitter::random_u64(), - INITIAL_QUOTA_SYNC_MAX_START_DELAY, - ) - }), - } - } - - /// 注入首查随机延迟采样器（测试用确定性节奏）；生产默认均匀 - /// `[0, INITIAL_QUOTA_SYNC_MAX_START_DELAY)`，零延迟等价旧行为 - #[doc(hidden)] - pub fn with_initial_sync_delays(mut self, delays: CodexInitialSyncDelays) -> Self { - self.initial_sync_delays = delays; - self - } - - /// 读取容量熔断策略（带短 TTL 缓存）；读取失败退化为关闭，熔断不得放大失败 - async fn freeze_policy(&self) -> ProviderFreezePolicy { - { - let cache = self.freeze_policy_cache.lock().await; - if let Some((policy, loaded_at)) = cache.as_ref() - && loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL - { - return policy.clone(); - } - } - let loaded = self.runtime_policy.load_freeze_policy().await.ok(); - let mut cache = self.freeze_policy_cache.lock().await; - if let Some(policy) = loaded { - *cache = Some((policy.clone(), Instant::now())); - return policy; - } - cache - .as_ref() - .filter(|(_, loaded_at)| loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL) - .map_or_else(ProviderFreezePolicy::disabled, |(policy, _)| policy.clone()) - } - - /// 容量熔断入口：滑动窗口内累计容量类失败，达到阈值即写入带 - /// `CapacityFreeze` 类别的账号级冷却 - /// 调度侧立即屏蔽该账号，恢复由 - /// freeze-recovery worker 处理；本路径只依赖 Redis 可丢失事实 - pub async fn apply_capacity_failure(&self, account: &ProviderAccount, observed_at: SystemTime) { - let policy = self.freeze_policy().await; - if !policy.enabled() { - return; - } - let in_flight = self.current_in_flight(account.id()).await; - let Ok(count) = self - .cooldowns - .record_capacity_failure(account.id(), policy.window(), in_flight) - .await - else { - return; - }; - if count < policy.threshold() { - return; - } - let Some(until) = observed_at.checked_add(policy.freeze_duration()) else { - return; - }; - let cooldown = ProviderCooldown::new_with_kind( - account.id().clone(), - account.revision(), - until, - if policy.probe_enabled() { - ProviderCooldownKind::CapacityFreezeProbe - } else { - ProviderCooldownKind::CapacityFreeze - }, - ); - if self.cooldowns.put_if_later(cooldown).await.is_ok() { - tracing::warn!( - account_id = account.id().as_str(), - threshold = policy.threshold(), - window_seconds = policy.window().as_secs(), - freeze_seconds = policy.freeze_duration().as_secs(), - peak_in_flight = in_flight, - "账号容量熔断触发：冻结该账号一段时间", - ); - } - } - - async fn current_in_flight(&self, account_id: &ProviderAccountId) -> u32 { - self.leases - .account_in_flight(std::slice::from_ref(account_id)) - .await - .ok() - .and_then(|signals| signals.get(account_id).copied()) - .unwrap_or(0) - } - - /// 查询当前账号由 Codex Desktop 暴露的主动额度重置卡 - pub async fn list_reset_credits( - &self, - account_id: &ProviderAccountId, - ) -> Result { - let account = self.reset_credit_account(account_id).await?; - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - let credential = self - .repository - .load_runtime_credential(&account) - .await - .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; - let prepared = PreparedCodexRuntimeCredential { - account, - credential, - }; - let request_id = format!("reset_credits_{}", Uuid::now_v7().simple()); - list_reset_credits_once(&client, &prepared, &request_id) - .await - .map_err(|error| map_reset_credit_attempt_error(error, false)) - } - - /// 消费一张主动额度重置卡 - /// 相同账号在本进程内串行，且不做传输重试 - pub async fn consume_reset_credit( - &self, - account_id: &ProviderAccountId, - credit_id: Option<&str>, - redeem_request_id: Uuid, - ) -> Result { - let lock = { - let mut locks = self.reset_consume_locks.lock().await; - Arc::clone( - locks - .entry(account_id.clone()) - .or_insert_with(|| Arc::new(Mutex::new(()))), - ) - }; - let _guard = lock.lock().await; - let account = self.reset_credit_account(account_id).await?; - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - let credential = self - .repository - .load_runtime_credential(&account) - .await - .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; - let prepared = PreparedCodexRuntimeCredential { - account, - credential, - }; - let request_id = format!("reset_credit_consume_{}", Uuid::now_v7().simple()); - consume_reset_credit_once( - &client, - &prepared, - &request_id, - credit_id, - redeem_request_id, - ) - .await - .map_err(|error| map_reset_credit_attempt_error(error, true)) - } - - async fn reset_credit_account( - &self, - account_id: &ProviderAccountId, - ) -> Result { - let account = self - .store - .get_account(account_id) - .await - .map_err(|error| CodexResetCreditsError::Store { - detail: error.to_string(), - })? - .filter(|account| account.provider().as_str() == "openai") - .ok_or(CodexResetCreditsError::NotFound)?; - if !access_token_is_current(&account, SystemTime::now()) { - return Err(CodexResetCreditsError::CredentialRefreshRequired { - upstream_body: None, - }); - } - Ok(account) - } - - /// 真实推理错误只更新额度访问事实，不伪造 Provider JSON 或展示百分比 - pub(crate) async fn record_confirmed_exhaustion( - &self, - account: &ProviderAccount, - evidence: QuotaEvidence, - reset_at: Option, - observed_at: SystemTime, - ) -> Result<(), CodexCredentialQuotaError> { - let outcome = self - .store - .apply_quota_access(QuotaAccessChange { - account_id: account.id().clone(), - expected_revision: account.revision(), - state: QuotaState::exhausted(evidence, observed_at, reset_at), - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - Ok(()) - } - - /// 成功推理是额度可访问的权威证据，同时解除账号级 429 冷却 - pub async fn record_successful_inference( - &self, - account: &ProviderAccount, - observed_at: SystemTime, - ) -> Result<(), CodexCredentialQuotaError> { - // 已经发出的并发请求可能晚于耗尽事实成功返回；真实成功仍是独立的 - // Allowed 权威证据，但调度不会为了试探恢复而放行耗尽账号 - if account.quota().access() == QuotaAccessState::Exhausted { - let outcome = self - .store - .apply_quota_access(QuotaAccessChange { - account_id: account.id().clone(), - expected_revision: account.revision(), - state: QuotaState::allowed(observed_at), - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - } - self.cooldowns - .clear_after_success(account.id(), account.revision()) - .await - .map_err(|error| CodexCredentialQuotaError::Store { - detail: error.to_string(), - })?; - Ok(()) - } - - /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 - pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { - if self.scheduling.hydration_targets(accounts).is_empty() { - return; - } - let _hydration = self.scheduling.hydration.lock().await; - let pending = self.scheduling.hydration_targets(accounts); - if pending.is_empty() { - return; - } - let account_ids = pending - .iter() - .map(|target| target.account.id().clone()) - .collect::>(); - let Ok(observations) = self.store.get_quotas(&account_ids).await else { - for target in &pending { - self.scheduling - .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); - } - return; - }; - let observations = observations - .into_iter() - .map(|observation| (observation.account_id.clone(), observation)) - .collect::>(); - for target in pending { - let snapshot = observations - .get(target.account.id()) - .filter(|observation| observation.expected_revision == target.account.revision()) - .and_then(quota_snapshot_from_observation); - if !snapshot - .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) - { - self.scheduling - .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); - } - } - } - - #[must_use] - pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { - self.scheduling.signals(account) - } - - pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { - self.scheduling.invalidate(account_ids); - } - - pub async fn synchronize(&self) -> Result { - self.synchronize_at(SystemTime::now()).await - } - - /// 按本轮调度时刻选择到期账号；实际 HTTP 观察与落库仍使用发生时刻 - pub async fn synchronize_at( - &self, - now: SystemTime, - ) -> Result { - let mut accounts = self.repository.list_for_provider().await?; - accounts.retain(|account| { - account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH - }); - let mut summary = CodexQuotaSyncSummary::default(); - let account_ids = accounts - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let observed = self.store.get_quotas(&account_ids).await?; - let observed_snapshots = observed - .iter() - .filter_map(|obs| { - quota_snapshot_from_observation(obs) - .map(|snapshot| (obs.account_id.clone(), snapshot)) - }) - .collect::>(); - let observed_ids = observed - .into_iter() - .map(|observation| observation.account_id) - .collect::>(); - let initial = self.scheduling.reserve_initial_refreshes( - &accounts, - &observed_ids, - now, - self.initial_sync_delays.as_ref(), - ); - // 本轮首查后仍无快照的账号需要重排随机延迟。 - let mut pending_initial: BTreeSet = - initial.iter().map(|account| account.id().clone()).collect(); - let periodic = - self.scheduling - .reserve_periodic_refreshes(accounts, &observed_snapshots, now); - let accounts = initial - .into_iter() - .chain(periodic) - .fold(BTreeMap::new(), |mut unique, account| { - unique.insert(account.id().clone(), account); - unique - }) - .into_values() - .collect::>(); - if accounts.is_empty() { - return Ok(summary); - } - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - for account in accounts { - let observed_at = SystemTime::now(); - match self.fetch_usage(&client, &account).await { - Ok(FetchedCodexQuota { account, value }) => { - // 单账号解析或落库失败只影响该账号；其余账号继续同步 - match self - .apply_fetched_quota(&account, &value, observed_at, &mut summary) - .await - { - Ok(()) => { - pending_initial.remove(account.id()); - } - Err(error) => { - summary.transient += 1; - tracing::warn!( - account_id = %account.id(), - error = %error, - "OpenAI quota synchronization skipped one account" - ); - } - } - } - Err(CodexQuotaFetchError::InvalidCredential) => { - summary.stale += 1; - } - Err(CodexQuotaFetchError::Upstream { account, error }) => { - match classify_quota_endpoint_failure(&error) { - Some(QuotaEndpointFailure::Exhausted(evidence)) => { - summary.exhausted += 1; - if let Err(write_error) = self - .record_confirmed_exhaustion(&account, evidence, None, observed_at) - .await - { - summary.transient += 1; - tracing::warn!( - account_id = %account.id(), - error = %write_error, - "OpenAI quota exhaustion fact write failed" - ); - } - } - Some(QuotaEndpointFailure::Credential { state, reason }) => { - summary.banned += 1; - self.persist_credential_failure(&account, state, reason, observed_at) - .await; - } - None => { - summary.transient += 1; - tracing::warn!( - account_id = %account.id(), - error = %error, - upstream_status = upstream_error_status(&error), - upstream_code = ?upstream_error_code(&error), - "OpenAI quota upstream rejection; refresh cycle will retry later" - ); - } - } - } - } - } - // 首查失败或未能落库的账号重排随机延迟；已成功落库的账号由快照事实自然退出首查路径。 - self.scheduling - .defer_failed_initial_refreshes(&pending_initial, self.initial_sync_delays.as_ref()); - Ok(summary) - } - - #[must_use] - pub fn runtime_policy(&self) -> &Arc { - &self.runtime_policy - } - - /// 批量预激活 OAuth 账号的 5h 配额滑动窗口 - pub async fn execute_warmup( - &self, - model: &str, - ) -> Result { - let mut accounts = self.repository.list_for_provider().await?; - accounts.retain(|account| { - account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH - && account.enabled() - && matches!( - account.credential_state(), - CredentialState::Unknown | CredentialState::Ready - ) - }); - let mut summary = CodexWarmupSummary::default(); - if accounts.is_empty() { - return Ok(summary); - } - let account_ids = accounts - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let observed = self.store.get_quotas(&account_ids).await?; - let observed_snapshots = observed - .iter() - .filter_map(|obs| { - quota_snapshot_from_observation(obs) - .map(|snapshot| (obs.account_id.clone(), snapshot)) - }) - .collect::>(); - - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - - let now_utc = chrono::Utc::now(); - for account in accounts { - if let Some(snapshot) = observed_snapshots.get(account.id()) { - // 1. 周线已触顶或耗尽跳过 - let weekly_exhausted = snapshot - .windows() - .iter() - .any(|w| w.kind() == CodexQuotaWindowKind::Weekly && w.limit_reached()); - if weekly_exhausted { - summary.skipped_exhausted += 1; - continue; - } - // 2. 5h 窗口当前活跃且距重置时间 > 30 分钟跳过 - let has_active_5h = snapshot.windows().iter().any(|w| { - w.kind() == CodexQuotaWindowKind::ShortTerm - && w.reset_at().is_some_and(|reset_at| { - reset_at > now_utc + chrono::Duration::minutes(30) - }) - }); - if has_active_5h { - summary.skipped_active += 1; - continue; - } - } - - let credential = match self.repository.load_runtime_credential(&account).await { - Ok(cred) => cred, - Err(_) => { - summary.failed += 1; - continue; - } - }; - let authorization = match credential.authentication.authorization_header() { - Ok(auth) => auth, - Err(_) => { - summary.failed += 1; - continue; - } - }; - - let mut body = Map::new(); - body.insert("model".to_owned(), Value::String(model.to_owned())); - body.insert( - "input".to_owned(), - serde_json::json!([{ - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "hello"}] - }]), - ); - body.insert("stream".to_owned(), Value::Bool(true)); - body.insert("store".to_owned(), Value::Bool(false)); - body.insert( - "service_tier".to_owned(), - Value::String("default".to_owned()), - ); - body.insert( - "reasoning".to_owned(), - serde_json::json!({"effort": "none"}), - ); - body.insert("text".to_owned(), serde_json::json!({"verbosity": "low"})); - - let upstream_request = CodexResponsesRequest::from_body(body); - let request_id = format!("warmup_{}", Uuid::now_v7().simple()); - let client_for_account = match client.for_account(&account) { - Ok(c) => c, - Err(error) => { - tracing::warn!(account_id = %account.id(), error = %error, "warmup client for account failed"); - summary.failed += 1; - continue; - } - }; - let context = crate::transport::CodexRequestContext::auxiliary( - authorization.expose_secret(), - account.upstream_account_id(), - &request_id, - None, - ); - - // HTTP 200 只确认响应头；必须消费 SSE 直到终态才能确认预热结果和额度事件 - let attempt = tokio::time::timeout(WARMUP_REQUEST_TIMEOUT, async { - let mut response = client_for_account - .create_response_stream_http_sse(&upstream_request, context) - .await - .map_err(|_| "request_rejected")?; - let terminal = consume_warmup_sse(&mut response).await; - let rate_limit_updates = match response.rate_limit_updates.as_ref() { - Some(updates) => std::mem::take(&mut *updates.lock().await), - None => Vec::new(), - }; - Ok::<_, &'static str>((terminal, response.rate_limit_headers, rate_limit_updates)) - }) - .await; - match attempt { - Ok(Ok((terminal, headers, rate_limit_updates))) => { - match terminal { - Ok(()) => { - // 被动额度同步会把成功请求的窗口事实标为可用，失败流不能复用该结论 - if !headers.is_empty() - && let Err(error) = - self.synchronize_passive_headers(&account, &headers).await - { - tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit header sync failed"); - } - if !rate_limit_updates.is_empty() - && let Err(error) = self - .synchronize_passive_rate_limits(&account, &rate_limit_updates) - .await - { - tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit event sync failed"); - } - summary.warmed_up += 1; - tracing::info!( - account_id = %account.id(), - model, - "OpenAI account warmed up successfully" - ); - } - Err(reason) => { - summary.failed += 1; - tracing::warn!( - account_id = %account.id(), - reason, - "OpenAI account warmup stream failed" - ); - } - } - } - Ok(Err(reason)) => { - summary.failed += 1; - tracing::warn!( - account_id = %account.id(), - reason, - "OpenAI account warmup request rejected by upstream" - ); - } - Err(_) => { - summary.failed += 1; - tracing::warn!( - account_id = %account.id(), - "OpenAI account warmup request timed out" - ); - } - } - tokio::time::sleep(Duration::from_secs(2)).await; - } - - Ok(summary) - } - - /// 解析并 revision-fenced 落库单账号的 Provider quota JSON - async fn apply_fetched_quota( - &self, - account: &ProviderAccount, - value: &Value, - observed_at: SystemTime, - summary: &mut CodexQuotaSyncSummary, - ) -> Result<(), CodexCredentialQuotaError> { - let mut object = normalize_quota_window_placeholders( - value - .as_object() - .cloned() - .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, - ); - object.remove(RECOVERY_FIELD); - let mut snapshot = parse_account_quota_snapshot( - account.id().clone(), - account.revision(), - observed_at, - &Value::Object(object.clone()), - )?; - let previous = if account.quota().is_exhausted() { - self.read_snapshot_for(account).await? - } else { - None - }; - let state = reconcile_refresh( - account.quota(), - &mut snapshot, - previous.as_ref(), - &mut object, - )?; - let snapshot = snapshot.with_quota_state(state); - let outcome = self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(object), - observed_at, - state, - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - summary.stale += 1; - return Ok(()); - } - self.scheduling.observe(&snapshot); - if snapshot.quota().is_exhausted() { - summary.exhausted += 1; - } else { - summary.updated += 1; - } - Ok(()) - } - - /// 把正常推理响应携带的限流事实合并进 Provider 原始 quota JSON - pub async fn synchronize_passive_headers( - &self, - account: &ProviderAccount, - headers: &[(String, String)], - ) -> Result { - let Some(rate_limits) = parse_rate_limit_headers(headers) else { - return Ok(false); - }; - self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits)) - .await - } - - /// 把一次推理响应中采集的结构化限流观察合并后单次落库 - pub async fn synchronize_passive_rate_limits( - &self, - account: &ProviderAccount, - rate_limits: &[ParsedRateLimits], - ) -> Result { - if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH - || rate_limits.is_empty() - { - return Ok(false); - } - let has_quota_facts = rate_limits.iter().any(|observation| { - observation - .limits - .values() - .any(|details| passive_rate_limit_snapshot(details).is_some()) - }); - let existing = self - .store - .get_quotas(std::slice::from_ref(account.id())) - .await? - .into_iter() - .find(|observation| { - observation.account_id == *account.id() - && observation.expected_revision == account.revision() - }); - let existing_state = existing.as_ref().map(|observation| observation.state); - let existing = existing - .map(|observation| observation.quota.into_inner()) - .unwrap_or_default(); - // 只用本次响应明确携带的套餐更新账号，不能把合并前的旧快照重新当作新证据 - let observed_plan = rate_limits - .iter() - .rev() - .filter_map(|observation| { - observation.plan_type.as_deref().filter(|plan| { - !plan.trim().is_empty() && !plan.trim().eq_ignore_ascii_case("unknown") - }) - }) - .next(); - let plan_type = observed_account_plan(account.plan_type(), observed_plan); - // 套餐、credits 等元数据可以更新，但没有额度窗口事实时必须保留旧观察时刻， - // 也不能借旧快照重新推导 quota state - if !has_quota_facts { - let Some(state) = existing_state else { - return Ok(false); - }; - let outcome = self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type, - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)), - observed_at: SystemTime::now(), - state, - }) - .await?; - return Ok(outcome != QuotaWriteOutcome::Conflict); - } - let quota = merge_passive_quota(existing, rate_limits); - let observed_at = SystemTime::now(); - let snapshot = parse_account_quota_snapshot( - account.id().clone(), - account.revision(), - observed_at, - &Value::Object(quota.clone()), - )?; - // 这些 headers 来自一次成功推理，访问结论优先于可能滞后的百分比 - let state = QuotaState::allowed(observed_at); - let outcome = self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type, - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(quota), - observed_at, - state, - }) - .await?; - if outcome == QuotaWriteOutcome::Conflict { - return Ok(false); - } - self.scheduling.observe(&snapshot); - Ok(true) - } - - /// 真实 429 的单一事实入口：写入 Redis 临时限流冷却（`until = now + retry_after`） - /// 凭据与额度主窗口（额度重置时间）都不改变——临时限流是独立维度， - /// 到期由 Redis key 过期自动解除，不污染配额耗尽状态 - /// 已有更晚的冷却不会被缩短（put_if_later） - pub async fn apply_rate_limit_429( - &self, - account: &ProviderAccount, - retry_after: Option, - observed_at: SystemTime, - ) -> Result<(), CodexCredentialQuotaError> { - let until = observed_at - .checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)) - .unwrap_or(observed_at); - self.cooldowns - .put_if_later(ProviderCooldown::new( - account.id().clone(), - account.revision(), - until, - )) - .await - .map_err(|error| CodexCredentialQuotaError::Store { - detail: error.to_string(), - })?; - Ok(()) - } - - /// 读取有效的账号冷却事实；等待恢复探测的冻结到期后仍有效 - pub async fn cooldown( - &self, - account_id: &ProviderAccountId, - ) -> Result, CodexCredentialQuotaError> { - let Some(cooldown) = self.cooldowns.read(account_id).await.map_err(|error| { - CodexCredentialQuotaError::Store { - detail: error.to_string(), - } - })? - else { - return Ok(None); - }; - let state = cooldown.scheduling_state(); - Ok(state.is_active(SystemTime::now()).then_some(state)) - } - - /// 读取单账号最后一次落库的 Provider quota，并由 Codex 域解析展示窗口 - pub async fn read_account( - &self, - account_id: &ProviderAccountId, - ) -> Result, CodexCredentialQuotaError> { - let account = self - .store - .get_account(account_id) - .await? - .filter(|account| account.provider().as_str() == "openai") - .ok_or(CodexCredentialQuotaError::NotFound)?; - self.read_snapshot_for(&account).await - } - - async fn read_snapshot_for( - &self, - account: &ProviderAccount, - ) -> Result, CodexCredentialQuotaError> { - if account.provider().as_str() != "openai" { - return Err(CodexCredentialQuotaError::NotFound); - } - let account_id = account.id(); - let Some(observation) = self - .store - .get_quotas(std::slice::from_ref(account_id)) - .await? - .into_iter() - .next() - else { - return Ok(None); - }; - if observation.account_id != *account_id - || observation.expected_revision != account.revision() - { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - if observation.quota.expose_to_provider().is_empty() { - return Ok(None); - } - let observed_at = observation.observed_at; - let snapshot = parse_account_quota_snapshot( - account_id.clone(), - account.revision(), - observed_at, - &Value::Object(observation.quota.expose_to_provider().clone()), - )? - .with_quota_state(observation.state); - self.scheduling.observe(&snapshot); - Ok(Some(snapshot)) - } - - /// 只刷新指定账号，revision-fenced 写入动态 Provider JSON 后返回解析快照 - pub async fn refresh_account( - &self, - account_id: &ProviderAccountId, - ) -> Result { - self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::ObserveAccess) - .await - } - - /// 真实限额失败后的异步刷新只补齐展示快照，不允许 usage 快照撤销已确认的失败状态 - pub(crate) async fn refresh_account_after_failure( - &self, - account_id: &ProviderAccountId, - ) -> Result { - self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::PreserveAccess) - .await - } - - async fn refresh_account_snapshot( - &self, - account_id: &ProviderAccountId, - authority: QuotaRefreshAuthority, - ) -> Result { - let account = self - .store - .get_account(account_id) - .await? - .filter(|account| account.provider().as_str() == "openai") - .ok_or(CodexCredentialQuotaError::NotFound)?; - if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH { - return Err(CodexCredentialQuotaError::NotFound); - } - let observed_at = SystemTime::now(); - if !access_token_is_current(&account, observed_at) { - return Err(CodexCredentialQuotaError::CredentialRefreshRequired); - } - let client = CodexBackendClient::new( - self.http.clone(), - self.base_url.clone(), - self.profile.clone(), - ); - let FetchedCodexQuota { account, value } = match self.fetch_usage(&client, &account).await { - Ok(fetched) => fetched, - Err(CodexQuotaFetchError::InvalidCredential) => { - return Err(CodexCredentialQuotaError::InvalidCredentialData); - } - Err(CodexQuotaFetchError::Upstream { account, error }) => { - match classify_quota_endpoint_failure(&error) { - Some(QuotaEndpointFailure::Exhausted(evidence)) => { - self.record_confirmed_exhaustion(&account, evidence, None, observed_at) - .await?; - } - Some(QuotaEndpointFailure::Credential { state, reason }) => { - self.persist_credential_failure(&account, state, reason, observed_at) - .await; - } - None => {} - } - return Err(CodexCredentialQuotaError::Upstream { - detail: error.to_string(), - status: upstream_error_status(&error), - code: upstream_error_code(&error), - }); - } - }; - let mut object = normalize_quota_window_placeholders( - value - .as_object() - .cloned() - .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, - ); - object.remove(RECOVERY_FIELD); - let mut snapshot = parse_account_quota_snapshot( - account.id().clone(), - account.revision(), - observed_at, - &Value::Object(object.clone()), - )?; - let previous = if account.quota().is_exhausted() { - self.read_snapshot_for(&account).await? - } else { - None - }; - if account.quota().is_exhausted() - && matches!(authority, QuotaRefreshAuthority::PreserveAccess) - { - if self - .store - .touch_quota_observation(QuotaObservationTouch { - account_id: account.id().clone(), - expected_revision: account.revision(), - observed_at, - }) - .await? - == QuotaWriteOutcome::Conflict - { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - return Ok(match previous { - Some(previous) => { - let previous = previous.with_observed_at(observed_at); - self.scheduling.observe(&previous); - previous - } - None => snapshot - .with_quota_state(account.quota()) - .with_observed_at(observed_at), - }); - } - let state = match authority { - QuotaRefreshAuthority::ObserveAccess => reconcile_refresh( - account.quota(), - &mut snapshot, - previous.as_ref(), - &mut object, - )?, - QuotaRefreshAuthority::PreserveAccess => account.quota(), - }; - let snapshot = snapshot.with_quota_state(state); - if self - .store - .compare_and_swap_quota(QuotaObservation { - plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: OpaqueProviderData::new(object), - observed_at, - state, - }) - .await? - == QuotaWriteOutcome::Conflict - { - return Err(CodexCredentialQuotaError::RevisionConflict); - } - self.scheduling.observe(&snapshot); - Ok(snapshot) - } - - async fn persist_credential_failure( - &self, - account: &ProviderAccount, - credential_state: CredentialState, - reason: AccountErrorReason, - observed_at: SystemTime, - ) { - if let Err(error) = self - .repository - .apply_state_with_reason(account, credential_state, observed_at, Some(reason), None) - .await - { - tracing::warn!( - account_id = %account.id(), - ?credential_state, - reason = reason.as_str(), - error = %error, - "OpenAI quota credential fact write failed" - ); - } - } - - async fn fetch_usage( - &self, - client: &CodexBackendClient, - account: &ProviderAccount, - ) -> Result { - let credential = self - .repository - .load_runtime_credential(account) - .await - .map_err(|_| CodexQuotaFetchError::InvalidCredential)?; - let prepared = PreparedCodexRuntimeCredential { - account: account.clone(), - credential, - }; - let result = fetch_usage_with_5xx_retry(client, &prepared).await; - match result { - Ok(value) => Ok(FetchedCodexQuota { - account: prepared.account, - value, - }), - Err(CodexQuotaFetchAttemptError::InvalidCredential) => { - Err(CodexQuotaFetchError::InvalidCredential) - } - Err(CodexQuotaFetchAttemptError::Upstream(error)) => { - Err(CodexQuotaFetchError::Upstream { - account: Box::new(prepared.account), - error, - }) - } - } - } -} - -async fn list_reset_credits_once( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, - request_id: &str, -) -> Result { - let authorization = prepared - .credential - .authentication - .authorization_header() - .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; - client - .for_account(&prepared.account) - .map_err(ResetCreditAttemptError::Upstream)? - .list_rate_limit_reset_credits(CodexRequestContext::auxiliary( - authorization.expose_secret(), - prepared.account.upstream_account_id(), - request_id, - None, - )) - .await - .map_err(ResetCreditAttemptError::Upstream) -} - -async fn consume_reset_credit_once( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, - request_id: &str, - credit_id: Option<&str>, - redeem_request_id: Uuid, -) -> Result { - let authorization = prepared - .credential - .authentication - .authorization_header() - .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; - client - .for_account(&prepared.account) - .map_err(ResetCreditAttemptError::Upstream)? - .consume_rate_limit_reset_credit( - CodexRequestContext::auxiliary( - authorization.expose_secret(), - prepared.account.upstream_account_id(), - request_id, - None, - ), - credit_id, - redeem_request_id, - ) - .await - .map_err(ResetCreditAttemptError::Upstream) -} - -fn map_reset_credit_attempt_error( - error: ResetCreditAttemptError, - consume: bool, -) -> CodexResetCreditsError { - match error { - ResetCreditAttemptError::InvalidCredential => CodexResetCreditsError::InvalidCredentialData, - ResetCreditAttemptError::Upstream(error) => map_reset_credit_client_error(error, consume), - } -} - -fn map_reset_credit_client_error(error: CodexClientError, consume: bool) -> CodexResetCreditsError { - match error { - CodexClientError::Upstream { - status, - body, - diagnostics, - .. - } if status == reqwest::StatusCode::UNAUTHORIZED - && reset_credit_response_was_explicit_rejection(status, &diagnostics) => - { - CodexResetCreditsError::CredentialRefreshRequired { - upstream_body: Some(body), - } - } - CodexClientError::Upstream { - status, - body, - retry_after_seconds, - diagnostics, - .. - } => { - // 2xx 后发生的解码/响应体上限错误使用 synthetic 502 表示，但额度卡 - // 可能已经消费 - // 只有 transport 记录的真实非成功状态与错误状态一致时， - // 才能把它当作确定的上游拒绝并允许前端清除 pending 幂等键 - if consume && !reset_credit_response_was_explicit_rejection(status, &diagnostics) { - CodexResetCreditsError::ConsumeResultUnknown - } else { - CodexResetCreditsError::Upstream { - status: status.as_u16(), - body, - retry_after_seconds, - } - } - } - _ if consume => CodexResetCreditsError::ConsumeResultUnknown, - _ => CodexResetCreditsError::TransportUnavailable, - } -} - -fn reset_credit_response_was_explicit_rejection( - status: reqwest::StatusCode, - diagnostics: &crate::transport::CodexUpstreamDiagnostics, -) -> bool { - !status.is_success() && diagnostics.status_code == Some(status.as_u16()) -} - -async fn fetch_usage_once( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, -) -> Result { - let authorization = prepared - .credential - .authentication - .authorization_header() - .map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?; - let request_id = format!("quota_{}", Uuid::now_v7().simple()); - client - .for_account(&prepared.account) - .map_err(CodexQuotaFetchAttemptError::Upstream)? - .fetch_usage(CodexRequestContext::auxiliary( - authorization.expose_secret(), - prepared.account.upstream_account_id(), - &request_id, - None, - )) - .await - .map_err(CodexQuotaFetchAttemptError::Upstream) -} - -/// 对 5xx 上游拒绝做有限次指数退避重试（1s/2s），吞掉瞬时抖动 -/// -/// 4xx（含 402/429）不重试：它们已经走额度状态转换，重试只会放大上游负载 -async fn fetch_usage_with_5xx_retry( - client: &CodexBackendClient, - prepared: &PreparedCodexRuntimeCredential, -) -> Result { - let mut attempt = 0_u32; - loop { - let result = fetch_usage_once(client, prepared).await; - let retryable = match &result { - Ok(_) => false, - Err(CodexQuotaFetchAttemptError::Upstream(CodexClientError::Upstream { - status, - .. - })) => status.is_server_error(), - Err(_) => false, - }; - if !retryable || attempt >= QUOTA_FETCH_5XX_MAX_RETRIES { - return result; - } - attempt += 1; - let delay = QUOTA_FETCH_5XX_BASE_DELAY.saturating_mul(attempt); - tracing::warn!( - account_id = %prepared.account.id(), - retry_attempt = attempt, - retry_delay_ms = delay.as_millis(), - "OpenAI quota usage 5xx upstream rejection; retrying with backoff" - ); - tokio::time::sleep(delay).await; - } -} - -/// 上游额度可确认套餐变更；同族泛化值不能丢弃 JWT 已给出的具体 SKU -fn observed_account_plan(current: Option<&str>, observed: Option<&str>) -> Option { - let plan = observed?.trim().to_ascii_lowercase(); - if plan.is_empty() || plan == "unknown" { - return None; - } - // 套餐族沿用官方 codex_protocol::account::PlanType 的分类 - let current = current.unwrap_or_default().trim().to_ascii_lowercase(); - let generalized = matches!( - (plan.as_str(), current.as_str()), - ( - "team", - "self_serve_business_prolite" | "self_serve_business_usage_based" - ) | ( - "business", - "ent26" | "enterprise_cbp_automation" | "enterprise_cbp_usage_based" - ) | ("edu" | "education", "edu_plus" | "edu_pro") - ); - (!generalized).then_some(plan) -} - -fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool { - account.enabled() && access_token_is_current(account, now) -} - -fn eligible_initial_quota_sync(account: &ProviderAccount, now: SystemTime) -> bool { - // 首次观察只兜底刚入库、尚无 quota 快照的账号 - // 已耗尽等运行时状态必须由 - // periodic 路径处理，才能保留同一账号的最小复核间隔 - account.credential_state() == CredentialState::Ready - && eligible_periodic_quota_refresh(account, now) -} - fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool { account .access_token_expires_at() .is_none_or(|expires_at| expires_at > now) } - -fn merge_passive_quota( - mut quota: Map, - observations: &[ParsedRateLimits], -) -> Map { - let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota); - for rate_limits in observations { - let default_is_named_alias = default_limit_is_named_alias(rate_limits); - let mut resolved_limit_ids = BTreeMap::new(); - for (wire_limit_id, details) in &rate_limits.limits { - // HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口 - // 同一 wire 观察里存在相同具名事实时丢弃镜像，不触碰 core 桶 - if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias { - continue; - } - let Some(limit_id) = snapshots.resolve_limit_id(details) else { - continue; - }; - let Some(rate_limit) = passive_rate_limit_snapshot(details) else { - continue; - }; - snapshots.upsert(&limit_id, details, rate_limit); - resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id); - } - - let active_limit = rate_limits - .active_limit - .as_deref() - .and_then(|wire_limit_id| { - resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| { - (!rate_limits.limits.contains_key(wire_limit_id)) - .then(|| wire_limit_id.to_owned()) - }) - }) - .or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned()); - if let Some(active_limit) = active_limit { - quota.insert("active_limit".to_owned(), Value::String(active_limit)); - } - merge_passive_metadata(&mut quota, rate_limits); - } - snapshots.write_to_document(&mut quota); - quota -} - -fn merge_passive_metadata(quota: &mut Map, rate_limits: &ParsedRateLimits) { - if let Some(plan_type) = rate_limits.plan_type.as_ref() { - quota.insert("plan_type".to_owned(), Value::String(plan_type.clone())); - } - if let Some(promo_message) = rate_limits.promo_message.as_ref() { - quota.insert( - "promo_message".to_owned(), - Value::String(promo_message.clone()), - ); - } - if let Some(reached_type) = rate_limits.rate_limit_reached_type.as_ref() { - quota.insert( - "rate_limit_reached_type".to_owned(), - Value::String(reached_type.clone()), - ); - } - if let Some(credits) = rate_limits.credits.as_ref() { - let mut value = quota - .remove("credits") - .and_then(|value| value.as_object().cloned()) - .unwrap_or_default(); - value.insert("has_credits".to_owned(), Value::Bool(credits.has_credits)); - value.insert("unlimited".to_owned(), Value::Bool(credits.unlimited)); - if let Some(balance) = credits.balance.as_ref() { - value.insert("balance".to_owned(), Value::String(balance.clone())); - } else { - // 点数对象明确更新但未提供余额时，不能把旧余额继续作为当前余额展示 - value.remove("balance"); - } - quota.insert("credits".to_owned(), Value::Object(value)); - } -} - -fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool { - let Some(default_primary) = rate_limits - .limits - .get(DEFAULT_CODEX_LIMIT_ID) - .and_then(|details| details.primary) - else { - return false; - }; - rate_limits.limits.iter().any(|(limit_id, details)| { - limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary) - }) -} - -/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位 -/// -/// 生产响应头可能携带 `used_percent=0`、零时长和空 reset 的占位，`/usage` 也会 -/// 返回 `secondary_window: null` -/// 正常被动同步保留该响应事实，Admin 展示会将其 -/// 隐藏；但主动 `/usage` 刷新或 402 确认投影会把一个存在但无事实的字段写成 -/// 100%，从而显示并不存在的“次级额度” -/// 因此这些非被动写入口在落库前移除 -/// 无事实值 -/// 带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按 -/// 原有额度逻辑保留 -fn normalize_quota_window_placeholders(mut quota: Map) -> Map { - quota = canonicalize_rate_limit_document(quota); - if let Some(rate_limit) = quota - .get_mut(RATE_LIMITS_BY_LIMIT_ID) - .and_then(Value::as_object_mut) - .and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID)) - .and_then(Value::as_object_mut) - { - drop_secondary_window_placeholder(rate_limit); - } - quota -} - -fn drop_secondary_window_placeholder(rate_limit: &mut Map) { - let placeholder = rate_limit.get("secondary_window").is_some_and(|window| { - window.is_null() - || window - .as_object() - .is_some_and(secondary_window_is_placeholder) - }); - if placeholder { - rate_limit.remove("secondary_window"); - } -} - -fn secondary_window_is_placeholder(window: &Map) -> bool { - window.iter().all(|(field, value)| match field.as_str() { - "used_percent" => value - .as_f64() - .is_some_and(|used_percent| used_percent.is_finite() && used_percent == 0.0), - "limit_reached" => value.as_bool() == Some(false), - _ => false, - }) -} - -fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option> { - // 响应头的窗口属于同一次上游观测；跨响应合并会制造不存在的额度窗口 - let mut snapshot = Map::new(); - if let Some(allowed) = details.allowed { - snapshot.insert("allowed".to_owned(), Value::Bool(allowed)); - } - if let Some(limit_reached) = details.limit_reached { - snapshot.insert("limit_reached".to_owned(), Value::Bool(limit_reached)); - } - for (field, window) in [ - ("primary_window", details.primary), - ("secondary_window", details.secondary), - ] { - let Some(window) = window else { - continue; - }; - snapshot.insert( - field.to_owned(), - Value::Object(passive_rate_limit_window(window)), - ); - } - (!snapshot.is_empty()).then_some(snapshot) -} - -fn passive_rate_limit_window(window: RateLimitWindow) -> Map { - let mut snapshot = Map::new(); - if let Some(number) = serde_json::Number::from_f64(window.used_percent) { - snapshot.insert("used_percent".to_owned(), Value::Number(number)); - } - if let Some(seconds) = window - .window_minutes - .and_then(|minutes| minutes.checked_mul(60)) - { - snapshot.insert("limit_window_seconds".to_owned(), Value::from(seconds)); - } - if let Some(reset_at) = window.reset_at { - snapshot.insert("reset_at".to_owned(), Value::from(reset_at)); - } - snapshot -} diff --git a/backend/crates/providers/openai/src/credential/quota/passive.rs b/backend/crates/providers/openai/src/credential/quota/passive.rs new file mode 100644 index 000000000..b91553adf --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/passive.rs @@ -0,0 +1,236 @@ +//! 推理响应携带的额度观察合并与 Provider 快照规范化 + +use super::*; + +impl CodexCredentialQuotaService { + /// 把正常推理响应携带的限流事实合并进 Provider 原始 quota JSON + pub async fn synchronize_passive_headers( + &self, + account: &ProviderAccount, + headers: &[(String, String)], + ) -> Result { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return Ok(false); + }; + self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits)) + .await + } + + /// 把一次推理响应中采集的结构化限流观察合并后单次落库 + pub async fn synchronize_passive_rate_limits( + &self, + account: &ProviderAccount, + rate_limits: &[ParsedRateLimits], + ) -> Result { + if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH + || rate_limits.is_empty() + { + return Ok(false); + } + let has_quota_facts = rate_limits.iter().any(|observation| { + observation + .limits + .values() + .any(|details| passive_rate_limit_snapshot(details).is_some()) + }); + let existing = self + .store + .get_quotas(std::slice::from_ref(account.id())) + .await? + .into_iter() + .find(|observation| { + observation.account_id == *account.id() + && observation.expected_revision == account.revision() + }); + let existing_state = existing.as_ref().map(|observation| observation.state); + let existing = existing + .map(|observation| observation.quota.into_inner()) + .unwrap_or_default(); + // 只用本次响应明确携带的套餐更新账号，不能把合并前的旧快照重新当作新证据 + let observed_plan = rate_limits + .iter() + .rev() + .filter_map(|observation| { + observation.plan_type.as_deref().filter(|plan| { + !plan.trim().is_empty() && !plan.trim().eq_ignore_ascii_case("unknown") + }) + }) + .next(); + let plan_type = observed_account_plan(account.plan_type(), observed_plan); + // 套餐、credits 等元数据可以更新，但没有额度窗口事实时必须保留旧观察时刻， + // 也不能借旧快照重新推导 quota state + if !has_quota_facts { + let Some(state) = existing_state else { + return Ok(false); + }; + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type, + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)), + observed_at: SystemTime::now(), + state, + }) + .await?; + return Ok(outcome != QuotaWriteOutcome::Conflict); + } + let quota = merge_passive_quota(existing, rate_limits); + let observed_at = SystemTime::now(); + let snapshot = parse_account_quota_snapshot( + account.id().clone(), + account.revision(), + observed_at, + &Value::Object(quota.clone()), + )?; + // 这些 headers 来自一次成功推理，访问结论优先于可能滞后的百分比 + let state = QuotaState::allowed(observed_at); + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type, + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(quota), + observed_at, + state, + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Ok(false); + } + self.scheduling.observe(&snapshot); + Ok(true) + } +} + +fn merge_passive_quota( + mut quota: Map, + observations: &[ParsedRateLimits], +) -> Map { + let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota); + for rate_limits in observations { + let default_is_named_alias = default_limit_is_named_alias(rate_limits); + let mut resolved_limit_ids = BTreeMap::new(); + for (wire_limit_id, details) in &rate_limits.limits { + // HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口 + // 同一 wire 观察里存在相同具名事实时丢弃镜像，不触碰 core 桶 + if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias { + continue; + } + let Some(limit_id) = snapshots.resolve_limit_id(details) else { + continue; + }; + let Some(rate_limit) = passive_rate_limit_snapshot(details) else { + continue; + }; + snapshots.upsert(&limit_id, details, rate_limit); + resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id); + } + + let active_limit = rate_limits + .active_limit + .as_deref() + .and_then(|wire_limit_id| { + resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| { + (!rate_limits.limits.contains_key(wire_limit_id)) + .then(|| wire_limit_id.to_owned()) + }) + }) + .or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned()); + if let Some(active_limit) = active_limit { + quota.insert("active_limit".to_owned(), Value::String(active_limit)); + } + merge_passive_metadata(&mut quota, rate_limits); + } + snapshots.write_to_document(&mut quota); + quota +} + +fn merge_passive_metadata(quota: &mut Map, rate_limits: &ParsedRateLimits) { + if let Some(plan_type) = rate_limits.plan_type.as_ref() { + quota.insert("plan_type".to_owned(), Value::String(plan_type.clone())); + } + if let Some(promo_message) = rate_limits.promo_message.as_ref() { + quota.insert( + "promo_message".to_owned(), + Value::String(promo_message.clone()), + ); + } + if let Some(reached_type) = rate_limits.rate_limit_reached_type.as_ref() { + quota.insert( + "rate_limit_reached_type".to_owned(), + Value::String(reached_type.clone()), + ); + } + if let Some(credits) = rate_limits.credits.as_ref() { + let mut value = quota + .remove("credits") + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + value.insert("has_credits".to_owned(), Value::Bool(credits.has_credits)); + value.insert("unlimited".to_owned(), Value::Bool(credits.unlimited)); + if let Some(balance) = credits.balance.as_ref() { + value.insert("balance".to_owned(), Value::String(balance.clone())); + } else { + // 点数对象明确更新但未提供余额时，不能把旧余额继续作为当前余额展示 + value.remove("balance"); + } + quota.insert("credits".to_owned(), Value::Object(value)); + } +} + +fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool { + let Some(default_primary) = rate_limits + .limits + .get(DEFAULT_CODEX_LIMIT_ID) + .and_then(|details| details.primary) + else { + return false; + }; + rate_limits.limits.iter().any(|(limit_id, details)| { + limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary) + }) +} + +fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option> { + // 响应头的窗口属于同一次上游观测；跨响应合并会制造不存在的额度窗口 + let mut snapshot = Map::new(); + if let Some(allowed) = details.allowed { + snapshot.insert("allowed".to_owned(), Value::Bool(allowed)); + } + if let Some(limit_reached) = details.limit_reached { + snapshot.insert("limit_reached".to_owned(), Value::Bool(limit_reached)); + } + for (field, window) in [ + ("primary_window", details.primary), + ("secondary_window", details.secondary), + ] { + let Some(window) = window else { + continue; + }; + snapshot.insert( + field.to_owned(), + Value::Object(passive_rate_limit_window(window)), + ); + } + (!snapshot.is_empty()).then_some(snapshot) +} + +fn passive_rate_limit_window(window: RateLimitWindow) -> Map { + let mut snapshot = Map::new(); + if let Some(number) = serde_json::Number::from_f64(window.used_percent) { + snapshot.insert("used_percent".to_owned(), Value::Number(number)); + } + if let Some(seconds) = window + .window_minutes + .and_then(|minutes| minutes.checked_mul(60)) + { + snapshot.insert("limit_window_seconds".to_owned(), Value::from(seconds)); + } + if let Some(reset_at) = window.reset_at { + snapshot.insert("reset_at".to_owned(), Value::from(reset_at)); + } + snapshot +} diff --git a/backend/crates/providers/openai/src/credential/quota/reset_credits.rs b/backend/crates/providers/openai/src/credential/quota/reset_credits.rs new file mode 100644 index 000000000..c8d0f49b6 --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/reset_credits.rs @@ -0,0 +1,315 @@ +//! 主动额度重置卡查询、按账号串行消费与不确定结果分类 + +use super::*; + +/// 仅索引活跃的消费或排队者，锁资源由请求持有 +#[derive(Default)] +pub(super) struct ResetCreditLocks { + entries: + std::sync::Mutex>>>, +} + +impl ResetCreditLocks { + fn entry(&self, account_id: &ProviderAccountId) -> ResetCreditEntry<'_> { + let mut entries = self + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let lock = entries + .get(account_id) + .and_then(std::sync::Weak::upgrade) + .unwrap_or_else(|| { + let lock = Arc::new(Mutex::new(())); + entries.insert(account_id.clone(), Arc::downgrade(&lock)); + lock + }); + ResetCreditEntry { + owner: self, + account_id: account_id.clone(), + lock: Some(lock), + } + } +} + +struct ResetCreditEntry<'a> { + owner: &'a ResetCreditLocks, + account_id: ProviderAccountId, + lock: Option>>, +} + +impl ResetCreditEntry<'_> { + async fn lock(&self) -> tokio::sync::MutexGuard<'_, ()> { + self.lock + .as_ref() + .expect("账号锁只在 entry 销毁时释放") + .lock() + .await + } +} + +impl Drop for ResetCreditEntry<'_> { + fn drop(&mut self) { + let mut entries = self + .owner + .entries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // 引用释放也须在表锁内完成，避免并发 Drop 都看见彼此而留下死 Weak + // 有持锁或排队者时仍保留同一锁身份；等待 future 取消也走此清理路径 + drop(self.lock.take()); + if entries + .get(&self.account_id) + .is_some_and(|lock| lock.strong_count() == 0) + { + entries.remove(&self.account_id); + } + } +} + +/// 主动额度重置卡查询/消费失败 +#[derive(Error)] +pub enum CodexResetCreditsError { + #[error("Codex reset-credit credential data is invalid")] + InvalidCredentialData, + #[error("Codex OAuth access token must be refreshed before using reset credits")] + CredentialRefreshRequired { upstream_body: Option }, + #[error("Codex reset-credit account was not found")] + NotFound, + #[error("provider account store is unavailable: {detail}")] + Store { detail: String }, + #[error("Codex reset-credit upstream returned HTTP {status}")] + Upstream { + status: u16, + body: String, + retry_after_seconds: Option, + }, + #[error("Codex reset-credit query transport is unavailable")] + TransportUnavailable, + #[error("Codex reset-credit consume result is unknown")] + ConsumeResultUnknown, +} + +impl std::fmt::Debug for CodexResetCreditsError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"), + Self::CredentialRefreshRequired { .. } => { + formatter.write_str("CredentialRefreshRequired { upstream_body: }") + } + Self::NotFound => formatter.write_str("NotFound"), + Self::Store { .. } => formatter.write_str("Store { detail: }"), + Self::Upstream { + status, + retry_after_seconds, + .. + } => formatter + .debug_struct("Upstream") + .field("status", status) + .field("body", &"") + .field("retry_after_seconds", retry_after_seconds) + .finish(), + Self::TransportUnavailable => formatter.write_str("TransportUnavailable"), + Self::ConsumeResultUnknown => formatter.write_str("ConsumeResultUnknown"), + } + } +} + +enum ResetCreditAttemptError { + InvalidCredential, + Upstream(CodexClientError), +} + +impl CodexCredentialQuotaService { + /// 查询当前账号由 Codex Desktop 暴露的主动额度重置卡 + pub async fn list_reset_credits( + &self, + account_id: &ProviderAccountId, + ) -> Result { + let account = self.reset_credit_account(account_id).await?; + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + let credential = self + .repository + .load_runtime_credential(&account) + .await + .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; + let prepared = PreparedCodexRuntimeCredential { + account, + credential, + }; + let request_id = format!("reset_credits_{}", Uuid::now_v7().simple()); + list_reset_credits_once(&client, &prepared, &request_id) + .await + .map_err(|error| map_reset_credit_attempt_error(error, false)) + } + + /// 消费一张主动额度重置卡 + /// 相同账号在本进程内串行，且不做传输重试 + pub async fn consume_reset_credit( + &self, + account_id: &ProviderAccountId, + credit_id: Option<&str>, + redeem_request_id: Uuid, + ) -> Result { + let entry = self.reset_consume_locks.entry(account_id); + let _guard = entry.lock().await; + let account = self.reset_credit_account(account_id).await?; + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + let credential = self + .repository + .load_runtime_credential(&account) + .await + .map_err(|_| CodexResetCreditsError::InvalidCredentialData)?; + let prepared = PreparedCodexRuntimeCredential { + account, + credential, + }; + let request_id = format!("reset_credit_consume_{}", Uuid::now_v7().simple()); + consume_reset_credit_once( + &client, + &prepared, + &request_id, + credit_id, + redeem_request_id, + ) + .await + .map_err(|error| map_reset_credit_attempt_error(error, true)) + } + + async fn reset_credit_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + let account = self + .store + .get_account(account_id) + .await + .map_err(|error| CodexResetCreditsError::Store { + detail: error.to_string(), + })? + .filter(|account| account.provider().as_str() == "openai") + .ok_or(CodexResetCreditsError::NotFound)?; + if !access_token_is_current(&account, SystemTime::now()) { + return Err(CodexResetCreditsError::CredentialRefreshRequired { + upstream_body: None, + }); + } + Ok(account) + } +} + +async fn list_reset_credits_once( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, + request_id: &str, +) -> Result { + let authorization = prepared + .credential + .authentication + .authorization_header() + .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; + client + .for_account(&prepared.account) + .map_err(ResetCreditAttemptError::Upstream)? + .list_rate_limit_reset_credits(CodexRequestContext::auxiliary( + authorization.expose_secret(), + prepared.account.upstream_account_id(), + request_id, + None, + )) + .await + .map_err(ResetCreditAttemptError::Upstream) +} + +async fn consume_reset_credit_once( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, + request_id: &str, + credit_id: Option<&str>, + redeem_request_id: Uuid, +) -> Result { + let authorization = prepared + .credential + .authentication + .authorization_header() + .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; + client + .for_account(&prepared.account) + .map_err(ResetCreditAttemptError::Upstream)? + .consume_rate_limit_reset_credit( + CodexRequestContext::auxiliary( + authorization.expose_secret(), + prepared.account.upstream_account_id(), + request_id, + None, + ), + credit_id, + redeem_request_id, + ) + .await + .map_err(ResetCreditAttemptError::Upstream) +} + +fn map_reset_credit_attempt_error( + error: ResetCreditAttemptError, + consume: bool, +) -> CodexResetCreditsError { + match error { + ResetCreditAttemptError::InvalidCredential => CodexResetCreditsError::InvalidCredentialData, + ResetCreditAttemptError::Upstream(error) => map_reset_credit_client_error(error, consume), + } +} + +fn map_reset_credit_client_error(error: CodexClientError, consume: bool) -> CodexResetCreditsError { + match error { + CodexClientError::Upstream { + status, + body, + diagnostics, + .. + } if status == reqwest::StatusCode::UNAUTHORIZED + && reset_credit_response_was_explicit_rejection(status, &diagnostics) => + { + CodexResetCreditsError::CredentialRefreshRequired { + upstream_body: Some(body), + } + } + CodexClientError::Upstream { + status, + body, + retry_after_seconds, + diagnostics, + .. + } => { + // 2xx 后发生的解码/响应体上限错误使用 synthetic 502 表示，但额度卡 + // 可能已经消费 + // 只有 transport 记录的真实非成功状态与错误状态一致时， + // 才能把它当作确定的上游拒绝并允许前端清除 pending 幂等键 + if consume && !reset_credit_response_was_explicit_rejection(status, &diagnostics) { + CodexResetCreditsError::ConsumeResultUnknown + } else { + CodexResetCreditsError::Upstream { + status: status.as_u16(), + body, + retry_after_seconds, + } + } + } + _ if consume => CodexResetCreditsError::ConsumeResultUnknown, + _ => CodexResetCreditsError::TransportUnavailable, + } +} + +fn reset_credit_response_was_explicit_rejection( + status: reqwest::StatusCode, + diagnostics: &crate::transport::CodexUpstreamDiagnostics, +) -> bool { + !status.is_success() && diagnostics.status_code == Some(status.as_u16()) +} diff --git a/backend/crates/providers/openai/src/credential/quota/scheduling.rs b/backend/crates/providers/openai/src/credential/quota/scheduling.rs new file mode 100644 index 000000000..4b3226ea8 --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/scheduling.rs @@ -0,0 +1,414 @@ +//! 请求级额度投影、首查与周期刷新排程 + +use super::*; + +impl CodexCredentialQuotaService { + /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 + pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { + if self.scheduling.hydration_targets(accounts).is_empty() { + return; + } + let _hydration = self.scheduling.hydration.lock().await; + let pending = self.scheduling.hydration_targets(accounts); + if pending.is_empty() { + return; + } + let account_ids = pending + .iter() + .map(|target| target.account.id().clone()) + .collect::>(); + let Ok(observations) = self.store.get_quotas(&account_ids).await else { + for target in &pending { + self.scheduling + .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); + } + return; + }; + let observations = observations + .into_iter() + .map(|observation| (observation.account_id.clone(), observation)) + .collect::>(); + for target in pending { + let snapshot = observations + .get(target.account.id()) + .filter(|observation| observation.expected_revision == target.account.revision()) + .and_then(quota_snapshot_from_observation); + if !snapshot + .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) + { + self.scheduling + .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); + } + } + } + + #[must_use] + pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { + self.scheduling.signals(account) + } + + pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { + self.scheduling.invalidate(account_ids); + } +} + +#[derive(Clone, Default)] +pub(super) struct CodexQuotaSchedulingProjection { + state: Arc>, + hydration: Arc>, +} + +#[derive(Default)] +struct CodexQuotaProjectionState { + next_version: u64, + entries: BTreeMap, + last_periodic_refresh_at: BTreeMap, + initial_refresh_not_before: BTreeMap, +} + +struct CodexQuotaRefreshAttempt { + monotonic_at: Instant, + wall_at: SystemTime, +} + +#[derive(Debug, Clone, Copy)] +struct CodexQuotaSchedulingEntry { + version: u64, + revision: CredentialRevision, + expires_at: Instant, + signals: Option, +} + +#[derive(Clone)] +struct CodexQuotaHydrationTarget { + account: ProviderAccount, + expected_version: Option, +} + +impl CodexQuotaSchedulingProjection { + fn invalidate(&self, account_ids: &[ProviderAccountId]) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for account_id in account_ids { + state.entries.remove(account_id); + state.last_periodic_refresh_at.remove(account_id); + state.initial_refresh_not_before.remove(account_id); + } + } + + fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + accounts + .iter() + .filter_map(|account| { + let current = state.entries.get(account.id()); + current + .is_none_or(|entry| { + entry.revision != account.revision() || now >= entry.expires_at + }) + .then(|| CodexQuotaHydrationTarget { + account: account.clone(), + expected_version: current.map(|entry| entry.version), + }) + }) + .collect() + } + + pub(super) fn observe(&self, snapshot: &CodexAccountQuotaSnapshot) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { + return false; + }; + self.replace( + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + scheduling_signals_from_snapshot(snapshot), + ); + true + } + + fn mark_unknown_if_unchanged(&self, target: &CodexQuotaHydrationTarget, ttl: Duration) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return; + } + insert_projection_entry( + &mut state, + target.account.id().clone(), + target.account.revision(), + ttl, + None, + ); + } + + fn observe_if_unchanged( + &self, + target: &CodexQuotaHydrationTarget, + snapshot: &CodexAccountQuotaSnapshot, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot) else { + return false; + }; + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return true; + } + insert_projection_entry( + &mut state, + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + scheduling_signals_from_snapshot(snapshot), + ); + true + } + + fn replace( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, + ) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + insert_projection_entry(&mut state, account_id, revision, ttl, signals); + } + + fn signals(&self, account: &ProviderAccount) -> Option { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .entries + .get(account.id()) + .filter(|entry| { + entry.revision == account.revision() && Instant::now() < entry.expires_at + }) + .and_then(|entry| entry.signals) + } + + pub(super) fn reserve_periodic_refreshes( + &self, + accounts: Vec, + observed_snapshots: &BTreeMap, + now: SystemTime, + ) -> Vec { + let candidates = accounts + .into_iter() + .filter_map(|account| { + let snapshot = observed_snapshots.get(account.id()); + quota_refresh_candidate(account, snapshot, now) + }) + .collect::>(); + let candidate_ids = candidates + .iter() + .map(|(account, _)| account.id().clone()) + .collect::>(); + let refreshed_at = Instant::now(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .last_periodic_refresh_at + .retain(|account_id, _| candidate_ids.contains(account_id)); + + // 已耗尽账号首次立即复核，之后每 30 分钟复核，以发现官方提前重置 + // reset + 2 分钟额外触发一次复核，给上游重置留出传播时间 + // 正常账号仅在非零用量窗口经过宽限期后参与，未更新时复用周期节流 + let mut reserved = Vec::new(); + for (account, target_reset) in candidates { + if !periodic_quota_refresh_due(&state, account.id(), target_reset, now, refreshed_at) { + continue; + } + state.last_periodic_refresh_at.insert( + account.id().clone(), + CodexQuotaRefreshAttempt { + monotonic_at: refreshed_at, + wall_at: now, + }, + ); + reserved.push(account); + } + reserved + } + + /// 为尚无 quota 快照的账号分配随机首查时点，并返回本轮已到期的候选。 + /// + /// `quota_observed_at` 为空代表首次异步观察尚未成功；不另建同步状态表。 + /// 首次进入候选的账号抽一次随机启动延迟记为到期时刻，避免新导入账号池 + /// 在 worker 唤醒边界形成批量 `/usage` 突发；到期时刻跨轮保留，零延迟 + /// （含随机源退化）时保持首轮即查的旧行为。单轮数量仍受 + /// `INITIAL_QUOTA_SYNC_BATCH` 上限约束。 + pub(super) fn reserve_initial_refreshes( + &self, + accounts: &[ProviderAccount], + observed_ids: &BTreeSet, + now: SystemTime, + draw_start_delay: &dyn Fn() -> Duration, + ) -> Vec { + let candidates = accounts + .iter() + .filter(|account| { + !observed_ids.contains(account.id()) && eligible_initial_quota_sync(account, now) + }) + .collect::>(); + let candidate_ids = candidates + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let monotonic_now = Instant::now(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // 已产生快照或离开候选集的账号不再保留首查状态。 + state + .initial_refresh_not_before + .retain(|account_id, _| candidate_ids.contains(account_id)); + candidates + .into_iter() + .filter(|account| { + let due_at = state + .initial_refresh_not_before + .entry(account.id().clone()) + .or_insert_with(|| monotonic_now + draw_start_delay()); + *due_at <= monotonic_now + }) + .take(INITIAL_QUOTA_SYNC_BATCH) + .cloned() + .collect() + } + + /// 首查尝试后仍无快照的账号重抽随机延迟，避免失败账号回到固定 30s 重试节奏。 + /// + /// 重排同时随机化每轮的候选集合，避免持续失败时按账号顺序只重试同一批。 + pub(super) fn defer_failed_initial_refreshes( + &self, + account_ids: &BTreeSet, + draw_start_delay: &dyn Fn() -> Duration, + ) { + if account_ids.is_empty() { + return; + } + let now = Instant::now(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for account_id in account_ids { + // 账号是否仍在首查路径由下一轮 reserve 按候选集剪枝，这里无需区分。 + state + .initial_refresh_not_before + .insert(account_id.clone(), now + draw_start_delay()); + } + } +} + +fn quota_refresh_candidate( + account: ProviderAccount, + snapshot: Option<&CodexAccountQuotaSnapshot>, + now: SystemTime, +) -> Option<(ProviderAccount, Option)> { + if !eligible_periodic_quota_refresh(&account, now) { + return None; + } + if account.quota().is_exhausted() { + let reset_at = account.quota().reset_at(); + return Some((account, reset_at)); + } + // 正常账号首次进入周期候选也要等待宽限期，不能由“无刷新历史”绕过 + let snapshot = snapshot?; + let expired_window_reset = snapshot + .windows() + .iter() + .filter(|window| { + window.used_percent().is_some_and(|used| used > 0.0) || window.limit_reached() + }) + .filter_map(CodexQuotaWindow::reset_at) + .map(SystemTime::from) + .filter(|reset| { + reset + .checked_add(QUOTA_RESET_GRACE) + .is_some_and(|due_at| due_at <= now) + }) + .min(); + expired_window_reset.map(|reset_at| (account, Some(reset_at))) +} + +fn periodic_quota_refresh_due( + state: &CodexQuotaProjectionState, + account_id: &ProviderAccountId, + reset_at: Option, + now: SystemTime, + monotonic_now: Instant, +) -> bool { + state + .last_periodic_refresh_at + .get(account_id) + .is_none_or(|last| { + monotonic_now.saturating_duration_since(last.monotonic_at) + >= PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL + || reset_at + .and_then(|reset| reset.checked_add(QUOTA_RESET_GRACE)) + // 已在该边界之后复核过时回到周期重试，避免过期 reset 每轮触发 + .is_some_and(|due_at| last.wall_at < due_at && due_at <= now) + }) +} + +fn insert_projection_entry( + state: &mut CodexQuotaProjectionState, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, +) { + state.next_version = state.next_version.saturating_add(1); + state.entries.insert( + account_id, + CodexQuotaSchedulingEntry { + version: state.next_version, + revision, + expires_at: Instant::now() + ttl, + signals, + }, + ); +} + +fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool { + account.enabled() && access_token_is_current(account, now) +} + +fn eligible_initial_quota_sync(account: &ProviderAccount, now: SystemTime) -> bool { + // 首次观察只兜底刚入库、尚无 quota 快照的账号 + // 已耗尽等运行时状态必须由 + // periodic 路径处理，才能保留同一账号的最小复核间隔 + account.credential_state() == CredentialState::Ready + && eligible_periodic_quota_refresh(account, now) +} diff --git a/backend/crates/providers/openai/src/credential/quota/service.rs b/backend/crates/providers/openai/src/credential/quota/service.rs new file mode 100644 index 000000000..b5b824f21 --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/service.rs @@ -0,0 +1,756 @@ +//! 额度权威读取、刷新、错误反馈与服务组装 + +use super::*; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum QuotaRefreshAuthority { + ObserveAccess, + PreserveAccess, +} + +struct FetchedCodexQuota { + account: ProviderAccount, + value: Value, +} + +enum CodexQuotaFetchAttemptError { + InvalidCredential, + Upstream(CodexClientError), +} + +enum CodexQuotaFetchError { + InvalidCredential, + Upstream { + account: Box, + error: CodexClientError, + }, +} + +/// 客户端错误携带的 HTTP 状态码；传输失败等为 `None` +fn upstream_error_status(error: &CodexClientError) -> Option { + match error { + CodexClientError::Upstream { status, .. } => Some(status.as_u16()), + _ => None, + } +} + +/// 从上游错误体提取稳定错误码：优先 `/error/code`，其次 `/code` +/// +/// 只接受有界的 ASCII 标识，避免把自由文本作为错误码写入诊断 +fn upstream_error_code(error: &CodexClientError) -> Option { + let CodexClientError::Upstream { body, .. } = error else { + return None; + }; + let value = serde_json::from_str::(body).ok()?; + let code = value + .pointer("/error/code") + .or_else(|| value.pointer("/code")) + .and_then(Value::as_str)? + .trim(); + if code.is_empty() + || code.len() > 64 + || !code + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.')) + { + return None; + } + Some(code.to_owned()) +} + +impl CodexCredentialQuotaService { + pub fn new( + repository: CodexCredentialRepository, + profile: CodexWireProfileState, + http: Client, + base_url: String, + cooldowns: Arc, + leases: Arc, + runtime_policy: Arc, + ) -> Self { + Self { + store: Arc::clone(repository.store()), + repository, + profile, + http, + base_url, + cooldowns, + leases, + runtime_policy, + freeze_policy_cache: Mutex::new(None), + scheduling: CodexQuotaSchedulingProjection::default(), + reset_consume_locks: ResetCreditLocks::default(), + initial_sync_delays: Arc::new(|| { + crate::jitter::uniform_delay( + crate::jitter::random_u64(), + INITIAL_QUOTA_SYNC_MAX_START_DELAY, + ) + }), + } + } + + /// 注入首查随机延迟采样器（测试用确定性节奏）；生产默认均匀 + /// `[0, INITIAL_QUOTA_SYNC_MAX_START_DELAY)`，零延迟等价旧行为 + #[doc(hidden)] + pub fn with_initial_sync_delays(mut self, delays: CodexInitialSyncDelays) -> Self { + self.initial_sync_delays = delays; + self + } + + /// 读取容量熔断策略（带短 TTL 缓存）；读取失败退化为关闭，熔断不得放大失败 + async fn freeze_policy(&self) -> ProviderFreezePolicy { + { + let cache = self.freeze_policy_cache.lock().await; + if let Some((policy, loaded_at)) = cache.as_ref() + && loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL + { + return policy.clone(); + } + } + let loaded = self.runtime_policy.load_freeze_policy().await.ok(); + let mut cache = self.freeze_policy_cache.lock().await; + if let Some(policy) = loaded { + *cache = Some((policy.clone(), Instant::now())); + return policy; + } + cache + .as_ref() + .filter(|(_, loaded_at)| loaded_at.elapsed() < FREEZE_POLICY_CACHE_TTL) + .map_or_else(ProviderFreezePolicy::disabled, |(policy, _)| policy.clone()) + } + + /// 容量熔断入口：滑动窗口内累计容量类失败，达到阈值即写入带 + /// `CapacityFreeze` 类别的账号级冷却 + /// 调度侧立即屏蔽该账号，恢复由 + /// freeze-recovery worker 处理；本路径只依赖 Redis 可丢失事实 + pub async fn apply_capacity_failure(&self, account: &ProviderAccount, observed_at: SystemTime) { + let policy = self.freeze_policy().await; + if !policy.enabled() { + return; + } + let in_flight = self.current_in_flight(account.id()).await; + let Ok(count) = self + .cooldowns + .record_capacity_failure(account.id(), policy.window(), in_flight) + .await + else { + return; + }; + if count < policy.threshold() { + return; + } + let Some(until) = observed_at.checked_add(policy.freeze_duration()) else { + return; + }; + let cooldown = ProviderCooldown::new_with_kind( + account.id().clone(), + account.revision(), + until, + if policy.probe_enabled() { + ProviderCooldownKind::CapacityFreezeProbe + } else { + ProviderCooldownKind::CapacityFreeze + }, + ); + if self.cooldowns.put_if_later(cooldown).await.is_ok() { + tracing::warn!( + account_id = account.id().as_str(), + threshold = policy.threshold(), + window_seconds = policy.window().as_secs(), + freeze_seconds = policy.freeze_duration().as_secs(), + peak_in_flight = in_flight, + "账号容量熔断触发：冻结该账号一段时间", + ); + } + } + + async fn current_in_flight(&self, account_id: &ProviderAccountId) -> u32 { + self.leases + .account_in_flight(std::slice::from_ref(account_id)) + .await + .ok() + .and_then(|signals| signals.get(account_id).copied()) + .unwrap_or(0) + } + + /// 真实推理错误只更新额度访问事实，不伪造 Provider JSON 或展示百分比 + pub(crate) async fn record_confirmed_exhaustion( + &self, + account: &ProviderAccount, + evidence: QuotaEvidence, + reset_at: Option, + observed_at: SystemTime, + ) -> Result<(), CodexCredentialQuotaError> { + let outcome = self + .store + .apply_quota_access(QuotaAccessChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + state: QuotaState::exhausted(evidence, observed_at, reset_at), + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + Ok(()) + } + + /// 成功推理是额度可访问的权威证据，同时解除账号级 429 冷却 + pub async fn record_successful_inference( + &self, + account: &ProviderAccount, + observed_at: SystemTime, + ) -> Result<(), CodexCredentialQuotaError> { + // 已经发出的并发请求可能晚于耗尽事实成功返回；真实成功仍是独立的 + // Allowed 权威证据，但调度不会为了试探恢复而放行耗尽账号 + if account.quota().access() == QuotaAccessState::Exhausted { + let outcome = self + .store + .apply_quota_access(QuotaAccessChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + state: QuotaState::allowed(observed_at), + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + } + self.cooldowns + .clear_after_success(account.id(), account.revision()) + .await + .map_err(|error| CodexCredentialQuotaError::Store { + detail: error.to_string(), + })?; + Ok(()) + } + + pub async fn synchronize(&self) -> Result { + self.synchronize_at(SystemTime::now()).await + } + + /// 按本轮调度时刻选择到期账号；实际 HTTP 观察与落库仍使用发生时刻 + pub async fn synchronize_at( + &self, + now: SystemTime, + ) -> Result { + let mut accounts = self.repository.list_for_provider().await?; + accounts.retain(|account| { + account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH + }); + let mut summary = CodexQuotaSyncSummary::default(); + let account_ids = accounts + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let observed = self.store.get_quotas(&account_ids).await?; + let observed_snapshots = observed + .iter() + .filter_map(|obs| { + quota_snapshot_from_observation(obs) + .map(|snapshot| (obs.account_id.clone(), snapshot)) + }) + .collect::>(); + let observed_ids = observed + .into_iter() + .map(|observation| observation.account_id) + .collect::>(); + let initial = self.scheduling.reserve_initial_refreshes( + &accounts, + &observed_ids, + now, + self.initial_sync_delays.as_ref(), + ); + // 本轮首查后仍无快照的账号需要重排随机延迟。 + let mut pending_initial: BTreeSet = + initial.iter().map(|account| account.id().clone()).collect(); + let periodic = + self.scheduling + .reserve_periodic_refreshes(accounts, &observed_snapshots, now); + let accounts = initial + .into_iter() + .chain(periodic) + .fold(BTreeMap::new(), |mut unique, account| { + unique.insert(account.id().clone(), account); + unique + }) + .into_values() + .collect::>(); + if accounts.is_empty() { + return Ok(summary); + } + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + for account in accounts { + let observed_at = SystemTime::now(); + match self.fetch_usage(&client, &account).await { + Ok(FetchedCodexQuota { account, value }) => { + // 单账号解析或落库失败只影响该账号；其余账号继续同步 + match self + .apply_fetched_quota(&account, &value, observed_at, &mut summary) + .await + { + Ok(()) => { + pending_initial.remove(account.id()); + } + Err(error) => { + summary.transient += 1; + tracing::warn!( + account_id = %account.id(), + error = %error, + "OpenAI quota synchronization skipped one account" + ); + } + } + } + Err(CodexQuotaFetchError::InvalidCredential) => { + summary.stale += 1; + } + Err(CodexQuotaFetchError::Upstream { account, error }) => { + match classify_quota_endpoint_failure(&error) { + Some(QuotaEndpointFailure::Exhausted(evidence)) => { + summary.exhausted += 1; + if let Err(write_error) = self + .record_confirmed_exhaustion(&account, evidence, None, observed_at) + .await + { + summary.transient += 1; + tracing::warn!( + account_id = %account.id(), + error = %write_error, + "OpenAI quota exhaustion fact write failed" + ); + } + } + Some(QuotaEndpointFailure::Credential { state, reason }) => { + summary.banned += 1; + self.persist_credential_failure(&account, state, reason, observed_at) + .await; + } + None => { + summary.transient += 1; + tracing::warn!( + account_id = %account.id(), + error = %error, + upstream_status = upstream_error_status(&error), + upstream_code = ?upstream_error_code(&error), + "OpenAI quota upstream rejection; refresh cycle will retry later" + ); + } + } + } + } + } + // 首查失败或未能落库的账号重排随机延迟；已成功落库的账号由快照事实自然退出首查路径。 + self.scheduling + .defer_failed_initial_refreshes(&pending_initial, self.initial_sync_delays.as_ref()); + Ok(summary) + } + + #[must_use] + pub fn runtime_policy(&self) -> &Arc { + &self.runtime_policy + } + + /// 解析并 revision-fenced 落库单账号的 Provider quota JSON + async fn apply_fetched_quota( + &self, + account: &ProviderAccount, + value: &Value, + observed_at: SystemTime, + summary: &mut CodexQuotaSyncSummary, + ) -> Result<(), CodexCredentialQuotaError> { + let mut object = normalize_quota_window_placeholders( + value + .as_object() + .cloned() + .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, + ); + object.remove(RECOVERY_FIELD); + let mut snapshot = parse_account_quota_snapshot( + account.id().clone(), + account.revision(), + observed_at, + &Value::Object(object.clone()), + )?; + let previous = if account.quota().is_exhausted() { + self.read_snapshot_for(account).await? + } else { + None + }; + let state = reconcile_refresh( + account.quota(), + &mut snapshot, + previous.as_ref(), + &mut object, + )?; + let snapshot = snapshot.with_quota_state(state); + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(object), + observed_at, + state, + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + summary.stale += 1; + return Ok(()); + } + self.scheduling.observe(&snapshot); + if snapshot.quota().is_exhausted() { + summary.exhausted += 1; + } else { + summary.updated += 1; + } + Ok(()) + } + + /// 真实 429 的单一事实入口：写入 Redis 临时限流冷却（`until = now + retry_after`） + /// 凭据与额度主窗口（额度重置时间）都不改变——临时限流是独立维度， + /// 到期由 Redis key 过期自动解除，不污染配额耗尽状态 + /// 已有更晚的冷却不会被缩短（put_if_later） + pub async fn apply_rate_limit_429( + &self, + account: &ProviderAccount, + retry_after: Option, + observed_at: SystemTime, + ) -> Result<(), CodexCredentialQuotaError> { + let until = observed_at + .checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)) + .unwrap_or(observed_at); + self.cooldowns + .put_if_later(ProviderCooldown::new( + account.id().clone(), + account.revision(), + until, + )) + .await + .map_err(|error| CodexCredentialQuotaError::Store { + detail: error.to_string(), + })?; + Ok(()) + } + + /// 读取有效的账号冷却事实；等待恢复探测的冻结到期后仍有效 + pub async fn cooldown( + &self, + account_id: &ProviderAccountId, + ) -> Result, CodexCredentialQuotaError> { + let Some(cooldown) = self.cooldowns.read(account_id).await.map_err(|error| { + CodexCredentialQuotaError::Store { + detail: error.to_string(), + } + })? + else { + return Ok(None); + }; + let state = cooldown.scheduling_state(); + Ok(state.is_active(SystemTime::now()).then_some(state)) + } + + /// 读取单账号最后一次落库的 Provider quota，并由 Codex 域解析展示窗口 + pub async fn read_account( + &self, + account_id: &ProviderAccountId, + ) -> Result, CodexCredentialQuotaError> { + let account = self + .store + .get_account(account_id) + .await? + .filter(|account| account.provider().as_str() == "openai") + .ok_or(CodexCredentialQuotaError::NotFound)?; + self.read_snapshot_for(&account).await + } + + async fn read_snapshot_for( + &self, + account: &ProviderAccount, + ) -> Result, CodexCredentialQuotaError> { + if account.provider().as_str() != "openai" { + return Err(CodexCredentialQuotaError::NotFound); + } + let account_id = account.id(); + let Some(observation) = self + .store + .get_quotas(std::slice::from_ref(account_id)) + .await? + .into_iter() + .next() + else { + return Ok(None); + }; + if observation.account_id != *account_id + || observation.expected_revision != account.revision() + { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + if observation.quota.expose_to_provider().is_empty() { + return Ok(None); + } + let observed_at = observation.observed_at; + let snapshot = parse_account_quota_snapshot( + account_id.clone(), + account.revision(), + observed_at, + &Value::Object(observation.quota.expose_to_provider().clone()), + )? + .with_quota_state(observation.state); + self.scheduling.observe(&snapshot); + Ok(Some(snapshot)) + } + + /// 只刷新指定账号，revision-fenced 写入动态 Provider JSON 后返回解析快照 + pub async fn refresh_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::ObserveAccess) + .await + } + + /// 真实限额失败后的异步刷新只补齐展示快照，不允许 usage 快照撤销已确认的失败状态 + pub(crate) async fn refresh_account_after_failure( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::PreserveAccess) + .await + } + + async fn refresh_account_snapshot( + &self, + account_id: &ProviderAccountId, + authority: QuotaRefreshAuthority, + ) -> Result { + let account = self + .store + .get_account(account_id) + .await? + .filter(|account| account.provider().as_str() == "openai") + .ok_or(CodexCredentialQuotaError::NotFound)?; + if account.authentication_kind() != crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH { + return Err(CodexCredentialQuotaError::NotFound); + } + let observed_at = SystemTime::now(); + if !access_token_is_current(&account, observed_at) { + return Err(CodexCredentialQuotaError::CredentialRefreshRequired); + } + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + let FetchedCodexQuota { account, value } = match self.fetch_usage(&client, &account).await { + Ok(fetched) => fetched, + Err(CodexQuotaFetchError::InvalidCredential) => { + return Err(CodexCredentialQuotaError::InvalidCredentialData); + } + Err(CodexQuotaFetchError::Upstream { account, error }) => { + match classify_quota_endpoint_failure(&error) { + Some(QuotaEndpointFailure::Exhausted(evidence)) => { + self.record_confirmed_exhaustion(&account, evidence, None, observed_at) + .await?; + } + Some(QuotaEndpointFailure::Credential { state, reason }) => { + self.persist_credential_failure(&account, state, reason, observed_at) + .await; + } + None => {} + } + return Err(CodexCredentialQuotaError::Upstream { + detail: error.to_string(), + status: upstream_error_status(&error), + code: upstream_error_code(&error), + }); + } + }; + let mut object = normalize_quota_window_placeholders( + value + .as_object() + .cloned() + .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?, + ); + object.remove(RECOVERY_FIELD); + let mut snapshot = parse_account_quota_snapshot( + account.id().clone(), + account.revision(), + observed_at, + &Value::Object(object.clone()), + )?; + let previous = if account.quota().is_exhausted() { + self.read_snapshot_for(&account).await? + } else { + None + }; + if account.quota().is_exhausted() + && matches!(authority, QuotaRefreshAuthority::PreserveAccess) + { + if self + .store + .touch_quota_observation(QuotaObservationTouch { + account_id: account.id().clone(), + expected_revision: account.revision(), + observed_at, + }) + .await? + == QuotaWriteOutcome::Conflict + { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + return Ok(match previous { + Some(previous) => { + let previous = previous.with_observed_at(observed_at); + self.scheduling.observe(&previous); + previous + } + None => snapshot + .with_quota_state(account.quota()) + .with_observed_at(observed_at), + }); + } + let state = match authority { + QuotaRefreshAuthority::ObserveAccess => reconcile_refresh( + account.quota(), + &mut snapshot, + previous.as_ref(), + &mut object, + )?, + QuotaRefreshAuthority::PreserveAccess => account.quota(), + }; + let snapshot = snapshot.with_quota_state(state); + if self + .store + .compare_and_swap_quota(QuotaObservation { + plan_type: observed_account_plan(account.plan_type(), snapshot.plan_type()), + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new(object), + observed_at, + state, + }) + .await? + == QuotaWriteOutcome::Conflict + { + return Err(CodexCredentialQuotaError::RevisionConflict); + } + self.scheduling.observe(&snapshot); + Ok(snapshot) + } + + async fn persist_credential_failure( + &self, + account: &ProviderAccount, + credential_state: CredentialState, + reason: AccountErrorReason, + observed_at: SystemTime, + ) { + if let Err(error) = self + .repository + .apply_state_with_reason(account, credential_state, observed_at, Some(reason), None) + .await + { + tracing::warn!( + account_id = %account.id(), + ?credential_state, + reason = reason.as_str(), + error = %error, + "OpenAI quota credential fact write failed" + ); + } + } + + async fn fetch_usage( + &self, + client: &CodexBackendClient, + account: &ProviderAccount, + ) -> Result { + let credential = self + .repository + .load_runtime_credential(account) + .await + .map_err(|_| CodexQuotaFetchError::InvalidCredential)?; + let prepared = PreparedCodexRuntimeCredential { + account: account.clone(), + credential, + }; + let result = fetch_usage_with_5xx_retry(client, &prepared).await; + match result { + Ok(value) => Ok(FetchedCodexQuota { + account: prepared.account, + value, + }), + Err(CodexQuotaFetchAttemptError::InvalidCredential) => { + Err(CodexQuotaFetchError::InvalidCredential) + } + Err(CodexQuotaFetchAttemptError::Upstream(error)) => { + Err(CodexQuotaFetchError::Upstream { + account: Box::new(prepared.account), + error, + }) + } + } + } +} + +async fn fetch_usage_once( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, +) -> Result { + let authorization = prepared + .credential + .authentication + .authorization_header() + .map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?; + let request_id = format!("quota_{}", Uuid::now_v7().simple()); + client + .for_account(&prepared.account) + .map_err(CodexQuotaFetchAttemptError::Upstream)? + .fetch_usage(CodexRequestContext::auxiliary( + authorization.expose_secret(), + prepared.account.upstream_account_id(), + &request_id, + None, + )) + .await + .map_err(CodexQuotaFetchAttemptError::Upstream) +} + +/// 对 5xx 上游拒绝做有限次指数退避重试（1s/2s），吞掉瞬时抖动 +/// +/// 4xx（含 402/429）不重试：它们已经走额度状态转换，重试只会放大上游负载 +async fn fetch_usage_with_5xx_retry( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, +) -> Result { + let mut attempt = 0_u32; + loop { + let result = fetch_usage_once(client, prepared).await; + let retryable = match &result { + Ok(_) => false, + Err(CodexQuotaFetchAttemptError::Upstream(CodexClientError::Upstream { + status, + .. + })) => status.is_server_error(), + Err(_) => false, + }; + if !retryable || attempt >= QUOTA_FETCH_5XX_MAX_RETRIES { + return result; + } + attempt += 1; + let delay = QUOTA_FETCH_5XX_BASE_DELAY.saturating_mul(attempt); + tracing::warn!( + account_id = %prepared.account.id(), + retry_attempt = attempt, + retry_delay_ms = delay.as_millis(), + "OpenAI quota usage 5xx upstream rejection; retrying with backoff" + ); + tokio::time::sleep(delay).await; + } +} diff --git a/backend/crates/providers/openai/src/credential/quota/warmup.rs b/backend/crates/providers/openai/src/credential/quota/warmup.rs new file mode 100644 index 000000000..9bf45da8c --- /dev/null +++ b/backend/crates/providers/openai/src/credential/quota/warmup.rs @@ -0,0 +1,260 @@ +//! 账号额度窗口预热与完成事件校验 + +use super::*; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum WarmupTerminal { + Missing, + Completed, + Failed, +} + +fn observe_warmup_events(events: Vec, terminal: &mut WarmupTerminal) { + for event in events { + let parsed = serde_json::from_str::(&event.data).ok(); + let kind = parsed + .as_ref() + .and_then(|data| data.get("type")) + .and_then(Value::as_str) + .or(event.event.as_deref()); + match kind { + Some("response.completed") => { + // 只有明确完成的响应才算预热成功；失败事件即使随后出现完成帧也优先 + if *terminal != WarmupTerminal::Failed { + *terminal = if parsed.is_some() + && parsed + .as_ref() + .and_then(|data| data.pointer("/response/status")) + .and_then(Value::as_str) + .is_none_or(|status| status == "completed") + { + WarmupTerminal::Completed + } else { + WarmupTerminal::Failed + }; + } + } + Some("response.failed" | "response.incomplete" | "error") => { + *terminal = WarmupTerminal::Failed; + } + _ => {} + } + } +} + +async fn consume_warmup_sse( + response: &mut CodexBackendStreamingResponse, +) -> Result<(), &'static str> { + let mut decoder = SseEventDecoder::default(); + let mut terminal = WarmupTerminal::Missing; + let mut received_bytes = 0usize; + while let Some(chunk) = response.body.next().await { + let chunk = chunk.map_err(|_| "stream_read_failed")?; + received_bytes = received_bytes.saturating_add(chunk.len()); + if received_bytes > WARMUP_STREAM_MAX_BYTES { + return Err("stream_size_limit_exceeded"); + } + let events = decoder.push(&chunk).map_err(|_| "invalid_sse")?; + observe_warmup_events(events, &mut terminal); + if terminal != WarmupTerminal::Missing { + break; + } + } + if terminal == WarmupTerminal::Missing { + observe_warmup_events(decoder.finish().map_err(|_| "invalid_sse")?, &mut terminal); + } + match terminal { + WarmupTerminal::Completed => Ok(()), + WarmupTerminal::Missing => Err("completion_event_missing"), + WarmupTerminal::Failed => Err("response_failed"), + } +} + +impl CodexCredentialQuotaService { + /// 批量预激活 OAuth 账号的 5h 配额滑动窗口 + pub async fn execute_warmup( + &self, + model: &str, + ) -> Result { + let mut accounts = self.repository.list_for_provider().await?; + accounts.retain(|account| { + account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH + && account.enabled() + && matches!( + account.credential_state(), + CredentialState::Unknown | CredentialState::Ready + ) + }); + let mut summary = CodexWarmupSummary::default(); + if accounts.is_empty() { + return Ok(summary); + } + let account_ids = accounts + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let observed = self.store.get_quotas(&account_ids).await?; + let observed_snapshots = observed + .iter() + .filter_map(|obs| { + quota_snapshot_from_observation(obs) + .map(|snapshot| (obs.account_id.clone(), snapshot)) + }) + .collect::>(); + + let client = CodexBackendClient::new( + self.http.clone(), + self.base_url.clone(), + self.profile.clone(), + ); + + let now_utc = chrono::Utc::now(); + for account in accounts { + if let Some(snapshot) = observed_snapshots.get(account.id()) { + // 1. 周线已触顶或耗尽跳过 + let weekly_exhausted = snapshot + .windows() + .iter() + .any(|w| w.kind() == CodexQuotaWindowKind::Weekly && w.limit_reached()); + if weekly_exhausted { + summary.skipped_exhausted += 1; + continue; + } + // 2. 5h 窗口当前活跃且距重置时间 > 30 分钟跳过 + let has_active_5h = snapshot.windows().iter().any(|w| { + w.kind() == CodexQuotaWindowKind::ShortTerm + && w.reset_at().is_some_and(|reset_at| { + reset_at > now_utc + chrono::Duration::minutes(30) + }) + }); + if has_active_5h { + summary.skipped_active += 1; + continue; + } + } + + let credential = match self.repository.load_runtime_credential(&account).await { + Ok(cred) => cred, + Err(_) => { + summary.failed += 1; + continue; + } + }; + let authorization = match credential.authentication.authorization_header() { + Ok(auth) => auth, + Err(_) => { + summary.failed += 1; + continue; + } + }; + + let mut body = Map::new(); + body.insert("model".to_owned(), Value::String(model.to_owned())); + body.insert( + "input".to_owned(), + serde_json::json!([{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }]), + ); + body.insert("stream".to_owned(), Value::Bool(true)); + body.insert("store".to_owned(), Value::Bool(false)); + body.insert( + "service_tier".to_owned(), + Value::String("default".to_owned()), + ); + body.insert( + "reasoning".to_owned(), + serde_json::json!({"effort": "none"}), + ); + body.insert("text".to_owned(), serde_json::json!({"verbosity": "low"})); + + let upstream_request = CodexResponsesRequest::from_body(body); + let request_id = format!("warmup_{}", Uuid::now_v7().simple()); + let client_for_account = match client.for_account(&account) { + Ok(c) => c, + Err(error) => { + tracing::warn!(account_id = %account.id(), error = %error, "warmup client for account failed"); + summary.failed += 1; + continue; + } + }; + let context = crate::transport::CodexRequestContext::auxiliary( + authorization.expose_secret(), + account.upstream_account_id(), + &request_id, + None, + ); + + // HTTP 200 只确认响应头；必须消费 SSE 直到终态才能确认预热结果和额度事件 + let attempt = tokio::time::timeout(WARMUP_REQUEST_TIMEOUT, async { + let mut response = client_for_account + .create_response_stream_http_sse(&upstream_request, context) + .await + .map_err(|_| "request_rejected")?; + let terminal = consume_warmup_sse(&mut response).await; + let rate_limit_updates = match response.rate_limit_updates.as_ref() { + Some(updates) => std::mem::take(&mut *updates.lock().await), + None => Vec::new(), + }; + Ok::<_, &'static str>((terminal, response.rate_limit_headers, rate_limit_updates)) + }) + .await; + match attempt { + Ok(Ok((terminal, headers, rate_limit_updates))) => { + match terminal { + Ok(()) => { + // 被动额度同步会把成功请求的窗口事实标为可用，失败流不能复用该结论 + if !headers.is_empty() + && let Err(error) = + self.synchronize_passive_headers(&account, &headers).await + { + tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit header sync failed"); + } + if !rate_limit_updates.is_empty() + && let Err(error) = self + .synchronize_passive_rate_limits(&account, &rate_limit_updates) + .await + { + tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit event sync failed"); + } + summary.warmed_up += 1; + tracing::info!( + account_id = %account.id(), + model, + "OpenAI account warmed up successfully" + ); + } + Err(reason) => { + summary.failed += 1; + tracing::warn!( + account_id = %account.id(), + reason, + "OpenAI account warmup stream failed" + ); + } + } + } + Ok(Err(reason)) => { + summary.failed += 1; + tracing::warn!( + account_id = %account.id(), + reason, + "OpenAI account warmup request rejected by upstream" + ); + } + Err(_) => { + summary.failed += 1; + tracing::warn!( + account_id = %account.id(), + "OpenAI account warmup request timed out" + ); + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + + Ok(summary) + } +} diff --git a/backend/crates/providers/openai/src/credential/security.rs b/backend/crates/providers/openai/src/credential/security.rs index bd440045e..96a6c4789 100644 --- a/backend/crates/providers/openai/src/credential/security.rs +++ b/backend/crates/providers/openai/src/credential/security.rs @@ -46,6 +46,14 @@ pub enum CodexRuntimeAuthentication { } impl CodexRuntimeAuthentication { + /// 传输只消费公开连接设置，不读取凭据枚举或秘密 + pub(crate) fn responses_api_base_url(&self) -> Option<&str> { + match self { + Self::ApiKey(auth) => Some(&auth.configuration.base_url), + Self::OAuth(_) => None, + } + } + pub fn authorization_header(&self) -> Result { match self { Self::ApiKey(auth) => Ok(SecretString::from(format!( diff --git a/backend/crates/providers/openai/src/lib.rs b/backend/crates/providers/openai/src/lib.rs index d1506bbf0..17cda4d01 100644 --- a/backend/crates/providers/openai/src/lib.rs +++ b/backend/crates/providers/openai/src/lib.rs @@ -4,8 +4,12 @@ mod admin; pub mod config; mod jitter; mod provider; +mod request_identity; mod session_transport; +/// 持久账号使用的 API Key 认证种类，凭据与账号传输隔离共用 +pub(crate) const CODEX_AUTHENTICATION_KIND_API_KEY: &str = "api_key"; + use std::sync::Arc; use gateway_admin::ports::provider::ProviderAdmin; diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 145ec52eb..aa05fa48d 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -314,7 +314,7 @@ impl CodexProvider { .client_for_request(&context)? .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? - .with_authentication(lease.authentication()) + .with_responses_api_base_url(lease.authentication().responses_api_base_url()) .with_middleware_headers(middleware_headers), response_origin: request.response_origin, endpoint_path: request.endpoint_path, diff --git a/backend/crates/providers/openai/src/provider/live.rs b/backend/crates/providers/openai/src/provider/live.rs index e49c37e38..f8b2b57ea 100644 --- a/backend/crates/providers/openai/src/provider/live.rs +++ b/backend/crates/providers/openai/src/provider/live.rs @@ -367,7 +367,7 @@ impl CodexProvider { .client_for_request(&context)? .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? - .with_authentication(lease.authentication()) + .with_responses_api_base_url(lease.authentication().responses_api_base_url()) .with_middleware_headers(middleware_headers), registry: Arc::clone(&self.live_registry), response_origin: self.live_calls_url.clone(), diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 4066a523f..7121a80ad 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -59,12 +59,12 @@ use crate::credential::{ CodexCredentialRefreshService, CodexCredentialRepository, CodexCredentialSelector, CodexCyberPolicyScope, CodexQuotaRefreshPolicy, CodexSessionAffinity, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, SelectCodexProviderEndpointCredential, - account_session_with_headers, derive_codex_cyber_policy_session_key, - derive_codex_endpoint_session_affinity, derive_codex_session_affinity, - derive_codex_transport_key, derive_endpoint_affinity_with_headers, - derive_live_session_affinity, derive_previous_response_id_hash, derive_turn_alias, - turn_id_with_headers, + derive_codex_cyber_policy_session_key, derive_codex_endpoint_session_affinity, + derive_codex_session_affinity, derive_codex_transport_key, + derive_endpoint_affinity_with_headers, derive_live_session_affinity, + derive_previous_response_id_hash, derive_turn_alias, }; +use crate::request_identity::{account_session_with_headers, turn_id_with_headers}; use crate::session_transport::CodexSessionTransportRecovery; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, @@ -704,7 +704,7 @@ impl CodexProvider { validate_openai_reasoning(generate.protocol_payload().body())?; let mut upstream = encode_generate_request(&generate, upstream_model.as_str(), None) .map_err(map_request_error)?; - upstream.client_account_follow_only = crate::credential::follows_session_with_headers( + upstream.client_account_follow_only = crate::request_identity::follows_session_with_headers( generate.protocol_payload().body(), generate.protocol_payload().context(), &middleware_headers, @@ -1002,7 +1002,7 @@ impl CodexProvider { .client_for_request(&context)? .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? - .with_authentication(lease.authentication()) + .with_responses_api_base_url(lease.authentication().responses_api_base_url()) .with_connection_budget(context.connection_budget().clone()) .with_response_control(context.response_control().cloned()) .with_middleware_headers(middleware_headers), diff --git a/backend/crates/providers/openai/src/request_identity.rs b/backend/crates/providers/openai/src/request_identity.rs new file mode 100644 index 000000000..81e456d4e --- /dev/null +++ b/backend/crates/providers/openai/src/request_identity.rs @@ -0,0 +1,102 @@ +//! 请求正文与最终头部的会话身份解释，供账号亲和与传输编码共用 + +use serde_json::{Map, Value}; + +/// 只有能识别出的后代线程禁止迁移会话；没有身份的普通客户端沿用现有调度 +pub(crate) fn follows_session_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> bool { + let metadata = headers + .iter() + .find(|h| h.name().eq_ignore_ascii_case("x-codex-turn-metadata")) + .and_then(|h| std::str::from_utf8(h.value()).ok()); + let thread = account_thread_with_headers(body, context, headers); + if let Some((session, thread)) = + account_session_with_headers(body, context, headers).zip(thread) + { + return session != thread; + } + gateway_protocol::openai::codex_responses_request_semantics_with_turn_metadata( + body, + metadata.or_else(|| context.get("turn_metadata").and_then(Value::as_str)), + ) + .subagent_kind + .is_some() +} + +fn account_thread_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> Option { + let header = |name: &str| { + headers + .iter() + .find(|h| h.name().eq_ignore_ascii_case(name)) + .and_then(|h| std::str::from_utf8(h.value()).ok()) + }; + header("x-codex-turn-metadata") + .and_then(gateway_protocol::openai::turn_metadata_thread_id) + .or_else(|| { + header("thread-id") + .and_then(|value| non_empty(Some(value))) + .map(str::to_owned) + }) + .or_else(|| gateway_protocol::openai::codex_account_thread_id(body, context)) +} + +/// 中间件最终请求头按实际传输的覆盖语义参与身份复验 +pub(crate) fn account_session_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> Option { + let header = |name: &str| { + headers + .iter() + .find(|h| h.name().eq_ignore_ascii_case(name)) + .and_then(|h| std::str::from_utf8(h.value()).ok()) + }; + if let Some(session) = + header("x-codex-turn-metadata").and_then(gateway_protocol::openai::turn_metadata_session_id) + { + return Some(session); + } + let mut context = context.clone(); + for (name, field) in [ + ("session-id", "session_id"), + ("x-codex-turn-metadata", "turn_metadata"), + ] { + if let Some(value) = header(name) { + context.insert(field.to_owned(), Value::String(value.to_owned())); + } + } + gateway_protocol::openai::codex_account_session_id(body, &context) +} + +pub(crate) fn turn_id_with_headers( + body: &Map, + context: &Map, + headers: &[gateway_core::engine::middleware::MiddlewareHeader], +) -> Option { + headers + .iter() + .find(|header| header.name().eq_ignore_ascii_case("x-codex-turn-metadata")) + .and_then(|header| std::str::from_utf8(header.value()).ok()) + .and_then(gateway_protocol::openai::turn_metadata_turn_id) + .or_else(|| { + headers + .iter() + .find(|header| header.name().eq_ignore_ascii_case("x-client-turn-id")) + .and_then(|header| std::str::from_utf8(header.value()).ok()) + .and_then(|value| non_empty(Some(value))) + .map(str::to_owned) + }) + .or_else(|| gateway_protocol::openai::codex_turn_id(body, context)) +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index bcd130343..adc6c14ea 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -775,12 +775,9 @@ impl CodexBackendClient { fallback.send().await.map_err(CodexClientError::HttpJson) } - pub(crate) fn with_authentication( - mut self, - authentication: &crate::credential::CodexRuntimeAuthentication, - ) -> Self { - if let crate::credential::CodexRuntimeAuthentication::ApiKey(auth) = authentication { - self.base_url = auth.configuration.base_url.trim_end_matches('/').to_owned(); + pub(crate) fn with_responses_api_base_url(mut self, base_url: Option<&str>) -> Self { + if let Some(base_url) = base_url { + self.base_url = base_url.trim_end_matches('/').to_owned(); self.protocol = OpenAiUpstreamProtocol::ResponsesApi; self.websocket_origin_key = format!( "{}:{}", @@ -798,7 +795,7 @@ impl CodexBackendClient { let mut client = self.clone(); client.outbound_proxy = account.outbound_proxy().cloned(); client.egress_key = egress_key(account.id().as_str(), account.outbound_proxy()); - if account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_API_KEY { + if account.authentication_kind() == crate::CODEX_AUTHENTICATION_KIND_API_KEY { client .egress_key .push_str(&format!(":revision:{}", account.revision().get())); diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 6d0cad009..af338662a 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -730,7 +730,7 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map bool { true } @@ -424,7 +424,7 @@ impl gateway_core::runtime::extensions::ExtensionSetLease for SessionExtensionLe #[tokio::test] async fn child_binding_only_constrains_builtin_scheduling_and_preserves_plugin_choices() { use gateway_core::engine::policy::RequestPolicyContext; - use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetReference}; + use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetReference}; for explicit in [false, true] { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_subagent_a").await; diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 67c46dba2..8380a77c7 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -54,11 +54,19 @@ pub(crate) struct MemoryAccountStore { accounts: Mutex>, quota_reads: AtomicUsize, fail_provider_listing: AtomicBool, + pending_account_reads: bool, credential_load_hook: Mutex>>, credential_loads: AtomicUsize, } impl MemoryAccountStore { + pub(crate) fn with_pending_account_reads() -> Self { + Self { + pending_account_reads: true, + ..Self::default() + } + } + pub(crate) fn repository(self: &Arc) -> CodexCredentialRepository { CodexCredentialRepository::new(self.clone()) } @@ -236,6 +244,9 @@ impl ProviderAccountStore for MemoryAccountStore { &self, account: &ProviderAccountId, ) -> Result, StoreError> { + if self.pending_account_reads { + futures::future::pending::<()>().await; + } Ok(self .accounts .lock() diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs index 448d3bc66..fe6ee2cc8 100644 --- a/backend/crates/providers/xai/src/admin.rs +++ b/backend/crates/providers/xai/src/admin.rs @@ -1181,8 +1181,6 @@ fn project_quota( .map_or(Value::Null, |value| Value::String(value.to_owned())), ); for (key, value) in [ - ("monthlyLimitCents", billing.monthly_limit_cents()), - ("includedUsedCents", billing.included_used_cents()), ("onDemandCapCents", billing.on_demand_cap_cents()), ("onDemandUsedCents", billing.on_demand_used_cents()), ("prepaidBalanceCents", billing.prepaid_balance_cents()), diff --git a/backend/crates/providers/xai/src/credential/catalog.rs b/backend/crates/providers/xai/src/credential/catalog.rs index 35dfb578f..54969426d 100644 --- a/backend/crates/providers/xai/src/credential/catalog.rs +++ b/backend/crates/providers/xai/src/credential/catalog.rs @@ -3,536 +3,28 @@ use std::collections::BTreeMap; use std::fmt; use std::sync::{Arc, RwLock}; -use std::time::{Duration, Instant, SystemTime}; +use std::time::SystemTime; use async_trait::async_trait; use chrono::{DateTime, Utc}; use gateway_core::account::{ - AccountQuotaSignals, CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccount, - ProviderAccountId, QuotaObservation, QuotaState, + CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId, }; use gateway_core::provider_ports::{ ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope, }; use gateway_core::routing::{ProviderCatalogGeneration, ProviderKind}; -use tokio::sync::Mutex; use super::repository::{GrokCredentialRepository, LoadedGrokCredential}; use crate::XaiWireProfileState; use crate::transport::catalog::{MAX_CATALOG_MODELS, valid_model_slug, validate_etag}; use crate::{ - GrokBillingClient, GrokBillingTransport, GrokCatalogModel, GrokModelCatalogClient, - GrokModelCatalogSession, GrokModelCatalogSnapshot, GrokModelCatalogTransport, SecretValue, - parse_grok_billing, + GrokCatalogModel, GrokModelCatalogClient, GrokModelCatalogSession, GrokModelCatalogSnapshot, + GrokModelCatalogTransport, SecretValue, }; const CATALOG_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(5 * 60); const MAX_CATALOG_FETCH_ATTEMPTS: usize = 3; -const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); -const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); - -/// Grok Build Free 额度的滚动观察窗口 -pub const GROK_FREE_ROLLING_WINDOW_SECONDS: u64 = 86_400; - -/// xAI Provider 从动态 billing JSON 解析出的旧账号页安全投影 -#[derive(Debug, Clone, PartialEq)] -pub struct GrokBillingPresentation { - plan_type: Option, - used_percent: Option, - period_type: Option, - period_start: Option, - period_end: Option, - monthly_limit_cents: Option, - included_used_cents: Option, - on_demand_cap_cents: Option, - on_demand_used_cents: Option, - prepaid_balance_cents: Option, -} - -/// xAI credits 当前周期的官方语义 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum GrokQuotaPeriodKind { - Weekly, - Monthly, - Other, -} - -impl GrokBillingPresentation { - #[must_use] - pub fn plan_type(&self) -> Option<&str> { - self.plan_type.as_deref() - } - - #[must_use] - pub const fn used_percent(&self) -> Option { - self.used_percent - } - - /// 返回 billing 是否包含可直接用于账号额度展示的付费事实 - #[must_use] - pub fn has_authoritative_quota(&self) -> bool { - self.used_percent.is_some() - || [ - self.monthly_limit_cents, - self.included_used_cents, - self.on_demand_cap_cents, - self.on_demand_used_cents, - self.prepaid_balance_cents, - ] - .into_iter() - .flatten() - .any(|value| value > 0) - } - - #[must_use] - pub fn period_type(&self) -> Option<&str> { - self.period_type.as_deref() - } - - #[must_use] - pub fn period_kind(&self) -> GrokQuotaPeriodKind { - match self.period_type.as_deref().map(str::trim) { - Some(value) - if value.eq_ignore_ascii_case("USAGE_PERIOD_TYPE_WEEKLY") - || value.eq_ignore_ascii_case("WEEKLY") => - { - GrokQuotaPeriodKind::Weekly - } - Some(value) - if value.eq_ignore_ascii_case("USAGE_PERIOD_TYPE_MONTHLY") - || value.eq_ignore_ascii_case("MONTHLY") => - { - GrokQuotaPeriodKind::Monthly - } - None if self.monthly_limit_cents.is_some() || self.included_used_cents.is_some() => { - GrokQuotaPeriodKind::Monthly - } - _ => GrokQuotaPeriodKind::Other, - } - } - - #[must_use] - pub fn period_start(&self) -> Option<&str> { - self.period_start.as_deref() - } - - #[must_use] - pub fn period_end(&self) -> Option<&str> { - self.period_end.as_deref() - } - - #[must_use] - pub const fn monthly_limit_cents(&self) -> Option { - self.monthly_limit_cents - } - - #[must_use] - pub const fn included_used_cents(&self) -> Option { - self.included_used_cents - } - - #[must_use] - pub const fn on_demand_cap_cents(&self) -> Option { - self.on_demand_cap_cents - } - - #[must_use] - pub const fn on_demand_used_cents(&self) -> Option { - self.on_demand_used_cents - } - - #[must_use] - pub const fn prepaid_balance_cents(&self) -> Option { - self.prepaid_balance_cents - } -} - -/// 一个账号最近一次 xAI billing 观察结果 -#[derive(Debug, Clone, PartialEq)] -pub struct GrokQuotaSnapshot { - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - observed_at: DateTime, - billing: GrokBillingPresentation, -} - -impl GrokQuotaSnapshot { - #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision - } - - #[must_use] - pub const fn observed_at(&self) -> DateTime { - self.observed_at - } - - #[must_use] - pub const fn billing(&self) -> &GrokBillingPresentation { - &self.billing - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -pub enum GrokQuotaError { - #[error("xAI quota account is unavailable")] - AccountUnavailable, - #[error("xAI quota credential snapshot is stale")] - StaleCredentialSnapshot, - #[error("xAI quota credential or billing data is invalid")] - InvalidData, - #[error("xAI quota upstream request failed")] - Upstream, - #[error("xAI quota store is unavailable")] - Store, -} - -/// 官方 Grok Build billing 同步与 Provider-owned quota JSON 解析服务 -#[derive(Clone)] -pub struct GrokCredentialQuotaService { - repository: GrokCredentialRepository, - client: Arc, - scheduling: GrokQuotaSchedulingProjection, - wire_profile: XaiWireProfileState, -} - -#[derive(Clone, Default)] -struct GrokQuotaSchedulingProjection { - state: Arc>, - hydration: Arc>, -} - -#[derive(Default)] -struct GrokQuotaProjectionState { - next_version: u64, - entries: BTreeMap, -} - -#[derive(Debug, Clone, Copy)] -struct GrokQuotaSchedulingEntry { - version: u64, - revision: CredentialRevision, - expires_at: Instant, - signals: Option, -} - -#[derive(Clone)] -struct GrokQuotaHydrationTarget { - account: ProviderAccount, - expected_version: Option, -} - -impl GrokQuotaSchedulingProjection { - fn invalidate(&self, account_ids: &[ProviderAccountId]) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - for account_id in account_ids { - state.entries.remove(account_id); - } - } - - fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let now = Instant::now(); - accounts - .iter() - .filter_map(|account| { - let current = state.entries.get(account.id()); - current - .is_none_or(|entry| { - entry.revision != account.revision() || now >= entry.expires_at - }) - .then(|| GrokQuotaHydrationTarget { - account: account.clone(), - expected_version: current.map(|entry| entry.version), - }) - }) - .collect() - } - - fn observe(&self, snapshot: &GrokQuotaSnapshot) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { - return false; - }; - self.replace( - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - quota_scheduling_signals(snapshot.billing()), - ); - true - } - - fn observe_if_unchanged( - &self, - target: &GrokQuotaHydrationTarget, - snapshot: &GrokQuotaSnapshot, - ) -> bool { - let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { - return false; - }; - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return true; - } - insert_grok_projection_entry( - &mut state, - snapshot.account_id().clone(), - snapshot.credential_revision(), - remaining_ttl, - quota_scheduling_signals(snapshot.billing()), - ); - true - } - - fn mark_unknown_if_unchanged(&self, target: &GrokQuotaHydrationTarget, ttl: Duration) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state - .entries - .get(target.account.id()) - .map(|entry| entry.version) - != target.expected_version - { - return; - } - insert_grok_projection_entry( - &mut state, - target.account.id().clone(), - target.account.revision(), - ttl, - None, - ); - } - - fn replace( - &self, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, - ) { - let mut state = self - .state - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - insert_grok_projection_entry(&mut state, account_id, revision, ttl, signals); - } - - fn signals(&self, account: &ProviderAccount) -> Option { - let state = self - .state - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state - .entries - .get(account.id()) - .filter(|entry| { - entry.revision == account.revision() && Instant::now() < entry.expires_at - }) - .and_then(|entry| entry.signals) - } -} - -fn insert_grok_projection_entry( - state: &mut GrokQuotaProjectionState, - account_id: ProviderAccountId, - revision: CredentialRevision, - ttl: Duration, - signals: Option, -) { - state.next_version = state.next_version.saturating_add(1); - state.entries.insert( - account_id, - GrokQuotaSchedulingEntry { - version: state.next_version, - revision, - expires_at: Instant::now() + ttl, - signals, - }, - ); -} - -impl GrokCredentialQuotaService { - #[must_use] - pub fn new( - repository: GrokCredentialRepository, - transport: Arc, - wire_profile: XaiWireProfileState, - ) -> Self { - Self { - repository, - client: Arc::new(GrokBillingClient::new(transport)), - scheduling: GrokQuotaSchedulingProjection::default(), - wire_profile, - } - } - - /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 - pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { - if self.scheduling.hydration_targets(accounts).is_empty() { - return; - } - let _hydration = self.scheduling.hydration.lock().await; - let pending = self.scheduling.hydration_targets(accounts); - if pending.is_empty() { - return; - } - let pending_accounts = pending - .iter() - .map(|target| target.account.clone()) - .collect::>(); - let Ok(observations) = self.repository.quota_observations(&pending_accounts).await else { - for target in &pending { - self.scheduling - .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); - } - return; - }; - let observations = observations - .into_iter() - .map(|observation| (observation.account_id.clone(), observation)) - .collect::>(); - for target in pending { - let snapshot = observations - .get(target.account.id()) - .filter(|observation| observation.expected_revision == target.account.revision()) - .and_then(quota_snapshot_from_observation); - if !snapshot - .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) - { - self.scheduling - .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); - } - } - } - - #[must_use] - pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { - self.scheduling.signals(account) - } - - pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { - self.scheduling.invalidate(account_ids); - } - - /// 立即刷新一个账号的动态 billing document，并以 credential revision CAS 写回 - pub async fn refresh_account( - &self, - account_id: &ProviderAccountId, - ) -> Result { - let loaded = self - .repository - .load_current(account_id) - .await - .map_err(map_quota_repository_error)?; - if !loaded.account.enabled() - || loaded - .account - .access_token_expires_at() - .is_none_or(|expires_at| expires_at <= SystemTime::now()) - { - return Err(GrokQuotaError::AccountUnavailable); - } - let session = billing_session(&loaded, &self.wire_profile)?; - let billing = self - .client - .fetch(&session) - .await - .map_err(|_| GrokQuotaError::Upstream)?; - let mut document = billing.into_document(); - // 订阅查询失败不影响已取得的额度，也不能据此写入免费套餐 - match self.client.fetch_subscription(&session).await { - Ok(Some(plan_type)) => { - document.insert( - "subscriptionTier".to_owned(), - serde_json::Value::String(plan_type), - ); - } - Ok(None) => {} - Err(error) => tracing::warn!(%error, "xAI subscription query failed"), - } - let observed_at = Utc::now(); - let presentation = billing_presentation(&document)?; - let quota = loaded - .account - .quota() - .merge_observation(quota_access_observation(&presentation, observed_at.into())); - self.repository - .replace_quota( - loaded.account.id().clone(), - loaded.account.revision(), - document, - observed_at.into(), - quota, - ) - .await - .map_err(map_quota_repository_error)?; - let snapshot = GrokQuotaSnapshot { - account_id: loaded.account.id().clone(), - credential_revision: loaded.account.revision(), - observed_at, - billing: presentation, - }; - self.scheduling.observe(&snapshot); - Ok(snapshot) - } - - /// 读取并重新验证 Store 中的 Provider-owned quota JSON - pub async fn read_account( - &self, - account_id: &ProviderAccountId, - ) -> Result, GrokQuotaError> { - let Some(observation) = self - .repository - .quota(account_id) - .await - .map_err(map_quota_repository_error)? - else { - return Ok(None); - }; - let observed_at = observation.observed_at; - let document = observation.quota.into_inner(); - let snapshot = GrokQuotaSnapshot { - account_id: observation.account_id, - credential_revision: observation.expected_revision, - observed_at: observed_at.into(), - billing: billing_presentation(&document)?, - }; - self.scheduling.observe(&snapshot); - Ok(Some(snapshot)) - } -} - -impl fmt::Debug for GrokCredentialQuotaService { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("GrokCredentialQuotaService") - .field("repository", &self.repository) - .field("client", &self.client) - .finish() - } -} /// 官方 `/v1/models` 形成的一个 account 完整模型集合 #[derive(Clone, PartialEq, Eq)] @@ -1219,188 +711,3 @@ fn strict_model_union( } Ok(union.into_values().collect()) } - -fn billing_session( - loaded: &LoadedGrokCredential, - wire_profile: &XaiWireProfileState, -) -> Result { - let upstream_user_id = loaded - .account - .upstream_user_id() - .ok_or(GrokQuotaError::InvalidData)?; - GrokModelCatalogSession::new( - loaded.access_token.clone(), - SecretValue::new(upstream_user_id), - loaded - .account - .email() - .map(|value| SecretValue::new(value.to_owned())), - wire_profile.clone(), - ) - .map(|session| session.with_outbound_proxy(loaded.account.outbound_proxy().cloned())) - .map_err(|_| GrokQuotaError::InvalidData) -} - -fn billing_presentation( - document: &serde_json::Map, -) -> Result { - let body = serde_json::to_vec(document).map_err(|_| GrokQuotaError::InvalidData)?; - let snapshot = parse_grok_billing(&body).map_err(|_| GrokQuotaError::InvalidData)?; - let config = snapshot - .document() - .get("config") - .and_then(|value| value.as_object()); - let current_period = config - .and_then(|config| config.get("currentPeriod")) - .and_then(|value| value.as_object()); - let monthly_limit_cents = cent_value(config, "monthlyLimit"); - let included_used_cents = cent_value(config, "used"); - let used_percent = config - .and_then(|config| config.get("creditUsagePercent")) - .and_then(serde_json::Value::as_f64) - .or_else(|| match (included_used_cents, monthly_limit_cents) { - (Some(used), Some(limit)) if limit > 0 => { - Some((used.min(limit) as f64 / limit as f64) * 100.0) - } - _ => None, - }); - Ok(GrokBillingPresentation { - plan_type: dynamic_string(Some(snapshot.document()), "subscriptionTier"), - used_percent, - period_type: dynamic_string(current_period, "type"), - period_start: dynamic_string(current_period, "start") - .or_else(|| dynamic_string(config, "billingPeriodStart")), - period_end: dynamic_string(current_period, "end") - .or_else(|| dynamic_string(config, "billingPeriodEnd")), - monthly_limit_cents, - included_used_cents, - on_demand_cap_cents: cent_value(config, "onDemandCap"), - on_demand_used_cents: cent_value(config, "onDemandUsed"), - prepaid_balance_cents: cent_value(config, "prepaidBalance"), - }) -} - -fn quota_snapshot_from_observation(observation: &QuotaObservation) -> Option { - let observed_at = DateTime::::from(observation.observed_at); - let document = observation.quota.expose_to_provider(); - Some(GrokQuotaSnapshot { - account_id: observation.account_id.clone(), - credential_revision: observation.expected_revision, - observed_at, - billing: billing_presentation(document).ok()?, - }) -} - -fn quota_projection_ttl(observed_at: DateTime) -> Option { - let age = SystemTime::now() - .duration_since(SystemTime::from(observed_at)) - .unwrap_or(Duration::ZERO); - QUOTA_SCHEDULING_TTL - .checked_sub(age) - .filter(|remaining| !remaining.is_zero()) -} - -fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option { - let now = SystemTime::now(); - let period_end = billing - .period_end() - .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) - .map(|value| SystemTime::from(value.to_utc())); - let observation_expired = period_end.is_some_and(|period_end| period_end <= now); - let remaining_rank = (!observation_expired) - .then(|| { - billing - .used_percent() - .filter(|used| used.is_finite() && (0.0..=100.0).contains(used)) - .map(|used| ((100.0 - used) * 100.0).round() as u64) - }) - .flatten(); - let reset_at = period_end.filter(|reset_at| *reset_at > now); - (remaining_rank.is_some() || reset_at.is_some()) - .then(|| AccountQuotaSignals::new(reset_at, remaining_rank)) -} - -/// 把 xAI billing 文档归一化为访问结论 -/// 缺少正反证据时返回 Unknown， -/// 由 `QuotaState::merge_observation` 保留既有已确认结论 -fn quota_access_observation( - billing: &GrokBillingPresentation, - observed_at: SystemTime, -) -> QuotaState { - if !billing.has_authoritative_quota() { - return QuotaState::observed_unknown(observed_at); - } - if billing_proves_access_allowed(billing, observed_at) { - QuotaState::allowed(observed_at) - } else { - QuotaState::observed_unknown(observed_at) - } -} - -fn billing_proves_access_allowed( - billing: &GrokBillingPresentation, - observed_at: SystemTime, -) -> bool { - if billing - .period_end() - .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) - .map(|value| SystemTime::from(value.to_utc())) - .is_some_and(|period_end| period_end <= observed_at) - { - return false; - } - billing - .prepaid_balance_cents() - .is_some_and(|balance| balance > 0) - || billing - .used_percent() - .is_some_and(|used| used.is_finite() && used < 100.0) - || matches!( - (billing.on_demand_used_cents(), billing.on_demand_cap_cents()), - (Some(used), Some(cap)) if cap > 0 && used < cap - ) - || matches!( - (billing.included_used_cents(), billing.monthly_limit_cents()), - (Some(used), Some(limit)) if limit > 0 && used < limit - ) -} - -fn dynamic_string( - object: Option<&serde_json::Map>, - field: &str, -) -> Option { - object? - .get(field) - .and_then(serde_json::Value::as_str) - .map(str::to_owned) -} - -fn cent_value( - object: Option<&serde_json::Map>, - field: &str, -) -> Option { - let cent = object?.get(field)?.as_object()?; - match cent.get("val") { - Some(value) => value.as_i64(), - None => Some(0), - } -} - -fn map_quota_repository_error( - error: super::repository::GrokCredentialRepositoryError, -) -> GrokQuotaError { - use super::repository::GrokCredentialRepositoryError as RepositoryError; - match error { - RepositoryError::CredentialNotFound | RepositoryError::WrongProviderKind => { - GrokQuotaError::AccountUnavailable - } - RepositoryError::StaleCredentialRevision(_) | RepositoryError::Conflict(_) => { - GrokQuotaError::StaleCredentialSnapshot - } - RepositoryError::Store(_) => GrokQuotaError::Store, - RepositoryError::InvalidInput(_) - | RepositoryError::IdentityRebind - | RepositoryError::InvalidCredentialData(_) - | RepositoryError::RevisionOverflow => GrokQuotaError::InvalidData, - } -} diff --git a/backend/crates/providers/xai/src/credential/mod.rs b/backend/crates/providers/xai/src/credential/mod.rs index f30be95c2..7c1580334 100644 --- a/backend/crates/providers/xai/src/credential/mod.rs +++ b/backend/crates/providers/xai/src/credential/mod.rs @@ -10,6 +10,7 @@ mod http; mod import; mod oidc_verifier; mod pkce; +mod quota; mod refresh; mod repository; mod secret; @@ -26,11 +27,9 @@ pub use config::{ }; pub use catalog::{ - GROK_FREE_ROLLING_WINDOW_SECONDS, GrokBillingPresentation, GrokCatalogCache, - GrokCatalogCacheError, GrokCatalogScope, GrokCredentialCatalogCache, + GrokCatalogCache, GrokCatalogCacheError, GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialCatalogService, - GrokCredentialQuotaService, GrokPlanCatalog, GrokQuotaError, GrokQuotaPeriodKind, - GrokQuotaSnapshot, + GrokPlanCatalog, }; pub use discovery::DiscoveryDocument; pub use error::{ @@ -47,6 +46,10 @@ pub use import::{ }; pub use oidc_verifier::ReqwestOidcTokenVerifier; pub use pkce::Pkce; +pub use quota::{ + GROK_FREE_ROLLING_WINDOW_SECONDS, GrokBillingPresentation, GrokCredentialQuotaService, + GrokQuotaError, GrokQuotaPeriodKind, GrokQuotaSnapshot, +}; pub use refresh::{ DueGrokCredential, GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, diff --git a/backend/crates/providers/xai/src/credential/quota.rs b/backend/crates/providers/xai/src/credential/quota.rs new file mode 100644 index 000000000..7cdd18273 --- /dev/null +++ b/backend/crates/providers/xai/src/credential/quota.rs @@ -0,0 +1,669 @@ +//! Grok billing 额度事实、展示投影与请求调度缓存 + +use std::collections::BTreeMap; +use std::fmt; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant, SystemTime}; + +use chrono::{DateTime, Utc}; +use gateway_core::account::{ + AccountQuotaSignals, CredentialRevision, ProviderAccount, ProviderAccountId, QuotaObservation, + QuotaState, +}; +use tokio::sync::Mutex; + +use super::repository::{GrokCredentialRepository, LoadedGrokCredential}; +use crate::XaiWireProfileState; +use crate::{ + GrokBillingClient, GrokBillingTransport, GrokModelCatalogSession, SecretValue, + parse_grok_billing, +}; + +const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); +const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); + +/// Grok Build Free 额度的滚动观察窗口 +pub const GROK_FREE_ROLLING_WINDOW_SECONDS: u64 = 86_400; + +/// xAI Provider 从当前 credits 文档解析出的账号页安全投影 +#[derive(Debug, Clone, PartialEq)] +pub struct GrokBillingPresentation { + plan_type: Option, + used_percent: Option, + period_type: Option, + period_start: Option, + period_end: Option, + on_demand_cap_cents: Option, + on_demand_used_cents: Option, + prepaid_balance_cents: Option, +} + +/// xAI credits 当前周期的官方语义 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GrokQuotaPeriodKind { + Weekly, + Monthly, + Other, +} + +impl GrokBillingPresentation { + #[must_use] + pub fn plan_type(&self) -> Option<&str> { + self.plan_type.as_deref() + } + + #[must_use] + pub const fn used_percent(&self) -> Option { + self.used_percent + } + + /// 返回 billing 是否包含可直接用于账号额度展示的付费事实 + #[must_use] + pub fn has_authoritative_quota(&self) -> bool { + self.used_percent.is_some() + || self.has_prepaid_credits() + || [self.on_demand_cap_cents, self.on_demand_used_cents] + .into_iter() + .flatten() + .any(|value| value > 0) + } + + fn has_prepaid_credits(&self) -> bool { + // 官方账本可用负数表示充值余额，非零即有预付额度，保留原始符号 + self.prepaid_balance_cents + .is_some_and(|balance| balance != 0) + } + + #[must_use] + pub fn period_type(&self) -> Option<&str> { + self.period_type.as_deref() + } + + #[must_use] + pub fn period_kind(&self) -> GrokQuotaPeriodKind { + match self.period_type.as_deref() { + Some("USAGE_PERIOD_TYPE_WEEKLY") => GrokQuotaPeriodKind::Weekly, + Some("USAGE_PERIOD_TYPE_MONTHLY") => GrokQuotaPeriodKind::Monthly, + _ => GrokQuotaPeriodKind::Other, + } + } + + #[must_use] + pub fn period_start(&self) -> Option<&str> { + self.period_start.as_deref() + } + + #[must_use] + pub fn period_end(&self) -> Option<&str> { + self.period_end.as_deref() + } + + #[must_use] + pub const fn on_demand_cap_cents(&self) -> Option { + self.on_demand_cap_cents + } + + #[must_use] + pub const fn on_demand_used_cents(&self) -> Option { + self.on_demand_used_cents + } + + #[must_use] + pub const fn prepaid_balance_cents(&self) -> Option { + self.prepaid_balance_cents + } +} + +/// 一个账号最近一次 xAI billing 观察结果 +#[derive(Debug, Clone, PartialEq)] +pub struct GrokQuotaSnapshot { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + observed_at: DateTime, + billing: GrokBillingPresentation, +} + +impl GrokQuotaSnapshot { + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn observed_at(&self) -> DateTime { + self.observed_at + } + + #[must_use] + pub const fn billing(&self) -> &GrokBillingPresentation { + &self.billing + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum GrokQuotaError { + #[error("xAI quota account is unavailable")] + AccountUnavailable, + #[error("xAI quota credential snapshot is stale")] + StaleCredentialSnapshot, + #[error("xAI quota credential or billing data is invalid")] + InvalidData, + #[error("xAI quota upstream request failed")] + Upstream, + #[error("xAI quota store is unavailable")] + Store, +} + +/// 官方 Grok Build billing 同步与 Provider-owned quota JSON 解析服务 +#[derive(Clone)] +pub struct GrokCredentialQuotaService { + repository: GrokCredentialRepository, + client: Arc, + scheduling: GrokQuotaSchedulingProjection, + wire_profile: XaiWireProfileState, +} + +#[derive(Clone, Default)] +struct GrokQuotaSchedulingProjection { + state: Arc>, + hydration: Arc>, +} + +#[derive(Default)] +struct GrokQuotaProjectionState { + next_version: u64, + entries: BTreeMap, +} + +#[derive(Debug, Clone, Copy)] +struct GrokQuotaSchedulingEntry { + version: u64, + revision: CredentialRevision, + expires_at: Instant, + signals: Option, +} + +#[derive(Clone)] +struct GrokQuotaHydrationTarget { + account: ProviderAccount, + expected_version: Option, +} + +impl GrokQuotaSchedulingProjection { + fn invalidate(&self, account_ids: &[ProviderAccountId]) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for account_id in account_ids { + state.entries.remove(account_id); + } + } + + fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + accounts + .iter() + .filter_map(|account| { + let current = state.entries.get(account.id()); + current + .is_none_or(|entry| { + entry.revision != account.revision() || now >= entry.expires_at + }) + .then(|| GrokQuotaHydrationTarget { + account: account.clone(), + expected_version: current.map(|entry| entry.version), + }) + }) + .collect() + } + + fn observe(&self, snapshot: &GrokQuotaSnapshot) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + self.replace( + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.billing()), + ); + true + } + + fn observe_if_unchanged( + &self, + target: &GrokQuotaHydrationTarget, + snapshot: &GrokQuotaSnapshot, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return true; + } + insert_grok_projection_entry( + &mut state, + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.billing()), + ); + true + } + + fn mark_unknown_if_unchanged(&self, target: &GrokQuotaHydrationTarget, ttl: Duration) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return; + } + insert_grok_projection_entry( + &mut state, + target.account.id().clone(), + target.account.revision(), + ttl, + None, + ); + } + + fn replace( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, + ) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + insert_grok_projection_entry(&mut state, account_id, revision, ttl, signals); + } + + fn signals(&self, account: &ProviderAccount) -> Option { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .entries + .get(account.id()) + .filter(|entry| { + entry.revision == account.revision() && Instant::now() < entry.expires_at + }) + .and_then(|entry| entry.signals) + } +} + +fn insert_grok_projection_entry( + state: &mut GrokQuotaProjectionState, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, +) { + state.next_version = state.next_version.saturating_add(1); + state.entries.insert( + account_id, + GrokQuotaSchedulingEntry { + version: state.next_version, + revision, + expires_at: Instant::now() + ttl, + signals, + }, + ); +} + +impl GrokCredentialQuotaService { + #[must_use] + pub fn new( + repository: GrokCredentialRepository, + transport: Arc, + wire_profile: XaiWireProfileState, + ) -> Self { + Self { + repository, + client: Arc::new(GrokBillingClient::new(transport)), + scheduling: GrokQuotaSchedulingProjection::default(), + wire_profile, + } + } + + /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度 + pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { + if self.scheduling.hydration_targets(accounts).is_empty() { + return; + } + let _hydration = self.scheduling.hydration.lock().await; + let pending = self.scheduling.hydration_targets(accounts); + if pending.is_empty() { + return; + } + let pending_accounts = pending + .iter() + .map(|target| target.account.clone()) + .collect::>(); + let Ok(observations) = self.repository.quota_observations(&pending_accounts).await else { + for target in &pending { + self.scheduling + .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); + } + return; + }; + let observations = observations + .into_iter() + .map(|observation| (observation.account_id.clone(), observation)) + .collect::>(); + for target in pending { + let snapshot = observations + .get(target.account.id()) + .filter(|observation| observation.expected_revision == target.account.revision()) + .and_then(quota_snapshot_from_observation); + if !snapshot + .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) + { + self.scheduling + .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); + } + } + } + + #[must_use] + pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { + self.scheduling.signals(account) + } + + pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) { + self.scheduling.invalidate(account_ids); + } + + /// 立即刷新一个账号的动态 billing document，并以 credential revision CAS 写回 + pub async fn refresh_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + let loaded = self + .repository + .load_current(account_id) + .await + .map_err(map_quota_repository_error)?; + if !loaded.account.enabled() + || loaded + .account + .access_token_expires_at() + .is_none_or(|expires_at| expires_at <= SystemTime::now()) + { + return Err(GrokQuotaError::AccountUnavailable); + } + let session = billing_session(&loaded, &self.wire_profile)?; + let billing = self + .client + .fetch(&session) + .await + .map_err(|_| GrokQuotaError::Upstream)?; + let mut document = billing.into_document(); + // 订阅查询失败不影响已取得的额度，也不能据此写入免费套餐 + match self.client.fetch_subscription(&session).await { + Ok(Some(plan_type)) => { + document.insert( + "subscriptionTier".to_owned(), + serde_json::Value::String(plan_type), + ); + } + Ok(None) => {} + Err(error) => tracing::warn!(%error, "xAI subscription query failed"), + } + let observed_at = Utc::now(); + let presentation = billing_presentation(&document)?; + let quota = loaded + .account + .quota() + .merge_observation(quota_access_observation(&presentation, observed_at.into())); + self.repository + .replace_quota( + loaded.account.id().clone(), + loaded.account.revision(), + document, + observed_at.into(), + quota, + ) + .await + .map_err(map_quota_repository_error)?; + let snapshot = GrokQuotaSnapshot { + account_id: loaded.account.id().clone(), + credential_revision: loaded.account.revision(), + observed_at, + billing: presentation, + }; + self.scheduling.observe(&snapshot); + Ok(snapshot) + } + + /// 读取并重新验证 Store 中的 Provider-owned quota JSON + pub async fn read_account( + &self, + account_id: &ProviderAccountId, + ) -> Result, GrokQuotaError> { + let Some(observation) = self + .repository + .quota(account_id) + .await + .map_err(map_quota_repository_error)? + else { + return Ok(None); + }; + let observed_at = observation.observed_at; + let document = observation.quota.into_inner(); + let snapshot = GrokQuotaSnapshot { + account_id: observation.account_id, + credential_revision: observation.expected_revision, + observed_at: observed_at.into(), + billing: billing_presentation(&document)?, + }; + self.scheduling.observe(&snapshot); + Ok(Some(snapshot)) + } +} + +impl fmt::Debug for GrokCredentialQuotaService { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("GrokCredentialQuotaService") + .field("repository", &self.repository) + .field("client", &self.client) + .finish() + } +} + +fn billing_session( + loaded: &LoadedGrokCredential, + wire_profile: &XaiWireProfileState, +) -> Result { + let upstream_user_id = loaded + .account + .upstream_user_id() + .ok_or(GrokQuotaError::InvalidData)?; + GrokModelCatalogSession::new( + loaded.access_token.clone(), + SecretValue::new(upstream_user_id), + loaded + .account + .email() + .map(|value| SecretValue::new(value.to_owned())), + wire_profile.clone(), + ) + .map(|session| session.with_outbound_proxy(loaded.account.outbound_proxy().cloned())) + .map_err(|_| GrokQuotaError::InvalidData) +} + +fn billing_presentation( + document: &serde_json::Map, +) -> Result { + let body = serde_json::to_vec(document).map_err(|_| GrokQuotaError::InvalidData)?; + let snapshot = parse_grok_billing(&body).map_err(|_| GrokQuotaError::InvalidData)?; + let config = snapshot + .document() + .get("config") + .and_then(|value| value.as_object()); + let current_period = config + .and_then(|config| config.get("currentPeriod")) + .and_then(|value| value.as_object()); + let used_percent = config + .and_then(|config| config.get("creditUsagePercent")) + .and_then(serde_json::Value::as_f64); + Ok(GrokBillingPresentation { + plan_type: dynamic_string(Some(snapshot.document()), "subscriptionTier"), + used_percent, + period_type: dynamic_string(current_period, "type"), + period_start: dynamic_string(current_period, "start"), + period_end: dynamic_string(current_period, "end"), + on_demand_cap_cents: cent_value(config, "onDemandCap"), + on_demand_used_cents: cent_value(config, "onDemandUsed"), + prepaid_balance_cents: cent_value(config, "prepaidBalance"), + }) +} + +fn quota_snapshot_from_observation(observation: &QuotaObservation) -> Option { + let observed_at = DateTime::::from(observation.observed_at); + let document = observation.quota.expose_to_provider(); + Some(GrokQuotaSnapshot { + account_id: observation.account_id.clone(), + credential_revision: observation.expected_revision, + observed_at, + billing: billing_presentation(document).ok()?, + }) +} + +fn quota_projection_ttl(observed_at: DateTime) -> Option { + let age = SystemTime::now() + .duration_since(SystemTime::from(observed_at)) + .unwrap_or(Duration::ZERO); + QUOTA_SCHEDULING_TTL + .checked_sub(age) + .filter(|remaining| !remaining.is_zero()) +} + +fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option { + let now = SystemTime::now(); + let period_end = billing + .period_end() + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|value| SystemTime::from(value.to_utc())); + let observation_expired = period_end.is_some_and(|period_end| period_end <= now); + let remaining_rank = (!observation_expired) + .then(|| { + billing + .used_percent() + .filter(|used| used.is_finite() && (0.0..=100.0).contains(used)) + .map(|used| ((100.0 - used) * 100.0).round() as u64) + }) + .flatten(); + let reset_at = period_end.filter(|reset_at| *reset_at > now); + (remaining_rank.is_some() || reset_at.is_some()) + .then(|| AccountQuotaSignals::new(reset_at, remaining_rank)) +} + +/// 把 xAI billing 文档归一化为访问结论 +/// 缺少正反证据时返回 Unknown， +/// 由 `QuotaState::merge_observation` 保留既有已确认结论 +fn quota_access_observation( + billing: &GrokBillingPresentation, + observed_at: SystemTime, +) -> QuotaState { + if !billing.has_authoritative_quota() { + return QuotaState::observed_unknown(observed_at); + } + if billing_proves_access_allowed(billing, observed_at) { + QuotaState::allowed(observed_at) + } else { + QuotaState::observed_unknown(observed_at) + } +} + +fn billing_proves_access_allowed( + billing: &GrokBillingPresentation, + observed_at: SystemTime, +) -> bool { + if billing + .period_end() + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|value| SystemTime::from(value.to_utc())) + .is_some_and(|period_end| period_end <= observed_at) + { + return false; + } + billing.has_prepaid_credits() + || billing + .used_percent() + .is_some_and(|used| used.is_finite() && used < 100.0) + || matches!( + (billing.on_demand_used_cents(), billing.on_demand_cap_cents()), + (Some(used), Some(cap)) if cap > 0 && used < cap + ) +} + +fn dynamic_string( + object: Option<&serde_json::Map>, + field: &str, +) -> Option { + object? + .get(field) + .and_then(serde_json::Value::as_str) + .map(str::to_owned) +} + +fn cent_value( + object: Option<&serde_json::Map>, + field: &str, +) -> Option { + let cent = object?.get(field)?.as_object()?; + match cent.get("val") { + Some(value) => value.as_i64(), + None => Some(0), + } +} + +fn map_quota_repository_error( + error: super::repository::GrokCredentialRepositoryError, +) -> GrokQuotaError { + use super::repository::GrokCredentialRepositoryError as RepositoryError; + match error { + RepositoryError::CredentialNotFound | RepositoryError::WrongProviderKind => { + GrokQuotaError::AccountUnavailable + } + RepositoryError::StaleCredentialRevision(_) | RepositoryError::Conflict(_) => { + GrokQuotaError::StaleCredentialSnapshot + } + RepositoryError::Store(_) => GrokQuotaError::Store, + RepositoryError::InvalidInput(_) + | RepositoryError::IdentityRebind + | RepositoryError::InvalidCredentialData(_) + | RepositoryError::RevisionOverflow => GrokQuotaError::InvalidData, + } +} diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index f089575a0..e99c156e6 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -19,7 +19,8 @@ use gateway_core::provider_ports::{ }; use gateway_core::routing::{ProviderKind, UpstreamModelId}; -use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService}; +use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache}; +use super::quota::GrokCredentialQuotaService; use super::repository::{GrokCredentialRepository, GrokCredentialRepositoryError}; use super::types::UpdateGrokCredentialState; use crate::{ diff --git a/backend/crates/providers/xai/src/lib.rs b/backend/crates/providers/xai/src/lib.rs index 2ced85de2..96afa76fe 100644 --- a/backend/crates/providers/xai/src/lib.rs +++ b/backend/crates/providers/xai/src/lib.rs @@ -1,7 +1,7 @@ //! 官方 Grok Build OAuth Provider 边界 mod admin; -pub mod credential; +mod credential; mod provider; mod reasoning_replay; pub mod transport; diff --git a/backend/crates/providers/xai/src/provider/continuation.rs b/backend/crates/providers/xai/src/provider/continuation.rs index c9bb83b27..fe8d926fa 100644 --- a/backend/crates/providers/xai/src/provider/continuation.rs +++ b/backend/crates/providers/xai/src/provider/continuation.rs @@ -1,6 +1,7 @@ //! xAI continuation 与 reasoning replay 状态处理 use super::*; +use crate::transport::GrokReplayItem; #[derive(Clone, Serialize, Deserialize)] pub(super) struct XaiSessionState { @@ -147,7 +148,12 @@ fn apply_replay( ) -> Result<(), ProviderError> { let mut input = replay_input_for_account(previous, account.as_str(), true); input.reserve(current_input.len()); - input.extend(current_input.iter().cloned()); + input.extend( + current_input + .iter() + .cloned() + .map(GrokReplayItem::UpstreamInput), + ); request.set_replay_input(input).map_err(map_request_error)?; request.set_previous_response_id(None); request.inherit_session(None); @@ -158,21 +164,24 @@ pub(super) fn replay_input_for_account( state: &XaiSessionState, account_id: &str, force_portable: bool, -) -> Vec { +) -> Vec { state .transcript .iter() .filter_map(|item| match item { - XaiReplayItem::ClientInput(value) | XaiReplayItem::SanitizedOutput(value) => { - Some(value.clone()) + XaiReplayItem::ClientInput(value) => Some(GrokReplayItem::UpstreamInput(value.clone())), + XaiReplayItem::SanitizedOutput(value) => { + Some(GrokReplayItem::ClientOutput(value.clone())) } XaiReplayItem::AccountOutput { account_id: owner, item, } if owner == account_id && !force_portable => { - portable_output_item(item.clone(), false) + portable_output_item(item.clone(), false).map(GrokReplayItem::ClientOutput) + } + XaiReplayItem::AccountOutput { item, .. } => { + portable_output_item(item.clone(), true).map(GrokReplayItem::ClientOutput) } - XaiReplayItem::AccountOutput { item, .. } => portable_output_item(item.clone(), true), }) .collect() } diff --git a/backend/crates/providers/xai/src/provider/failure.rs b/backend/crates/providers/xai/src/provider/failure.rs index 8a038c0c0..3e6b493b1 100644 --- a/backend/crates/providers/xai/src/provider/failure.rs +++ b/backend/crates/providers/xai/src/provider/failure.rs @@ -383,15 +383,23 @@ pub(super) fn map_request_error(error: GrokRequestEncodeError) -> ProviderError GrokRequestEncodeError::InvalidProtocolPayload | GrokRequestEncodeError::InvalidRequestNormalization | GrokRequestEncodeError::InvalidRequestField { .. } => ProviderErrorKind::InvalidRequest, + GrokRequestEncodeError::UnsupportedPrewarm => ProviderErrorKind::Unsupported, GrokRequestEncodeError::Serialization => ProviderErrorKind::Protocol, }; let provider_error = provider_error(kind, UpstreamSendState::NotSent); - if kind != ProviderErrorKind::InvalidRequest { + if !matches!( + kind, + ProviderErrorKind::InvalidRequest | ProviderErrorKind::Unsupported + ) { return provider_error; } + let code = match error { + GrokRequestEncodeError::UnsupportedPrewarm => "unsupported_prewarm", + _ => "invalid_request_normalization", + }; provider_error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( error.to_string(), - Some("invalid_request_normalization".to_owned()), + Some(code.to_owned()), Some("invalid_request_error".to_owned()), )) } diff --git a/backend/crates/providers/xai/src/provider/mod.rs b/backend/crates/providers/xai/src/provider/mod.rs index 1c7053afe..f7b246741 100644 --- a/backend/crates/providers/xai/src/provider/mod.rs +++ b/backend/crates/providers/xai/src/provider/mod.rs @@ -81,7 +81,6 @@ use stream::*; pub(crate) use workers::worker_contributions; const HTTP_SSE_TRANSPORT: &str = "http_sse"; -const DEFAULT_GROK_MODEL: &str = "grok-4.5"; const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024; const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096; const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed"; @@ -205,12 +204,9 @@ impl Provider for GrokBuildProvider { async fn query_model_capabilities( &self, ) -> Result, ProviderError> { - let Ok(models) = self.catalog.query_models().await else { - return Ok(vec![default_grok_model_capabilities()?]); - }; - if models.is_empty() { - return Ok(vec![default_grok_model_capabilities()?]); - } + let models = self.catalog.query_models().await.map_err(|_| { + provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) + })?; Ok(models .into_iter() .map(compile_grok_model_capabilities) @@ -842,97 +838,41 @@ fn compile_grok_model_capabilities(model: GrokCatalogModel) -> ProviderModelCapa .with_presentation(grok_model_presentation(&model)) } -fn default_grok_model_capabilities() -> Result { - let model = UpstreamModelId::new(DEFAULT_GROK_MODEL.to_owned()) - .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; - let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None) - .with_upstream_feature_validation(); - Ok(ProviderModelCapabilities::new(model, capabilities) - .with_presentation(default_grok_model_presentation())) -} - -fn default_grok_model_presentation() -> ModelPresentation { - ModelPresentation::new( - Some("Grok 4.5".to_owned()), - Some("xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()), - ) - .with_context_window_tokens(Some(500_000)) - .with_max_context_window_tokens(Some(500_000)) - .with_image_input(true) - .with_agent_tools(true, true) -} - fn grok_model_presentation(model: &GrokCatalogModel) -> ModelPresentation { - let slug = model.request_model().as_str(); - let known_grok_4_5 = is_known_grok_4_5_model(slug); - let reasoning_evidence = model.capabilities().reasoning_effort(); - let catalog_reasoning_efforts = model + let reasoning_efforts = model .capabilities() .reasoning_efforts() .iter() .map(|effort| effort.as_str().to_owned()) .collect::>(); - let catalog_default_reasoning = model + let default_reasoning = model .capabilities() .default_reasoning_effort() .map(|effort| effort.as_str().to_owned()); - let (default_reasoning, reasoning_efforts) = match reasoning_evidence { - GrokCatalogCapabilityEvidence::DeclaredNative if !catalog_reasoning_efforts.is_empty() => { - let default_reasoning = catalog_default_reasoning - .filter(|default| { - catalog_reasoning_efforts - .iter() - .any(|effort| effort == default) - }) - .or_else(|| catalog_reasoning_efforts.first().cloned()); - (default_reasoning, catalog_reasoning_efforts) - } - GrokCatalogCapabilityEvidence::DeclaredNative => (None, Vec::new()), - GrokCatalogCapabilityEvidence::DeclaredUnsupported => { - (Some("none".to_owned()), vec!["none".to_owned()]) - } - GrokCatalogCapabilityEvidence::Unknown => (None, Vec::new()), - }; let context_window_tokens = model .limits() .context_window_tokens() - .map(std::num::NonZeroU64::get) - .or(known_grok_4_5.then_some(500_000)); + .map(std::num::NonZeroU64::get); let tool_evidence = model.capabilities().streaming_tool_calls(); ModelPresentation::new( - model - .display_name() - .map(str::to_owned) - .or_else(|| known_grok_4_5.then(|| "Grok 4.5".to_owned())), - model - .metadata() - .description() - .map(str::to_owned) - .or_else(|| { - known_grok_4_5 - .then(|| "xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()) - }), + model.display_name().map(str::to_owned), + model.metadata().description().map(str::to_owned), ) .with_reasoning(default_reasoning, reasoning_efforts) .with_context_window_tokens(context_window_tokens) - // Grok 目录只声明一个窗口，继续将它作为客户端可覆盖上限 - .with_max_context_window_tokens(context_window_tokens) - .with_image_input(known_grok_4_5) + .with_max_context_window_tokens( + model + .limits() + .max_context_window_tokens() + .map(std::num::NonZeroU64::get), + ) .with_agent_tools( tool_evidence != GrokCatalogCapabilityEvidence::DeclaredUnsupported, - tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative - || (known_grok_4_5 && tool_evidence == GrokCatalogCapabilityEvidence::Unknown), + tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative, ) .with_search_tool( model.capabilities().backend_search() == GrokCatalogCapabilityEvidence::DeclaredNative, ) .with_hidden(model.metadata().hidden().unwrap_or(false)) } - -fn is_known_grok_4_5_model(slug: &str) -> bool { - matches!( - slug, - DEFAULT_GROK_MODEL | "grok-4.5-latest" | "grok-4.5-build-free" | "grok-build-latest" - ) -} diff --git a/backend/crates/providers/xai/src/reasoning_replay.rs b/backend/crates/providers/xai/src/reasoning_replay.rs index 289bd4a7c..3a99b270f 100644 --- a/backend/crates/providers/xai/src/reasoning_replay.rs +++ b/backend/crates/providers/xai/src/reasoning_replay.rs @@ -8,6 +8,8 @@ use base64::Engine as _; use gateway_core::event::ProviderEvent; use serde_json::{Map, Value}; +use crate::transport::GrokReplayItem; + const REASONING_REPLAY_TTL: Duration = Duration::from_secs(60 * 60); const REASONING_REPLAY_MAX_ENTRIES: usize = 1_024; const REASONING_REPLAY_MAX_BYTES: usize = 8 * 1024 * 1024; @@ -82,7 +84,7 @@ impl GrokReasoningReplay { &self, key: &GrokReasoningReplayKey, input: &[Value], - ) -> Option> { + ) -> Option> { if input.is_empty() { return None; } @@ -560,17 +562,17 @@ fn replay_call_keys(item_type: &str, call_id: &str) -> Vec { .collect() } -fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec { +fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec { let insert_at = replay_insert_index(input, replay); let mut result = Vec::new(); for (index, item) in input.iter().enumerate() { if index == insert_at { - result.extend(replay.iter().cloned()); + result.extend(replay.iter().cloned().map(GrokReplayItem::ClientOutput)); } - result.push(item.clone()); + result.push(GrokReplayItem::UpstreamInput(item.clone())); } if insert_at == input.len() { - result.extend(replay.iter().cloned()); + result.extend(replay.iter().cloned().map(GrokReplayItem::ClientOutput)); } result } diff --git a/backend/crates/providers/xai/src/transport/canonical.rs b/backend/crates/providers/xai/src/transport/canonical.rs index 2fa90dc26..a4b7856a4 100644 --- a/backend/crates/providers/xai/src/transport/canonical.rs +++ b/backend/crates/providers/xai/src/transport/canonical.rs @@ -1041,7 +1041,10 @@ fn provider_reported_cost(response: &Value) -> Result FinishReason { .pointer("/incomplete_details/reason") .and_then(Value::as_str) { - Some("max_output_tokens" | "max_tokens") => FinishReason::Length, + Some("max_output_tokens" | "max_tokens" | "max_prompt_tokens" | "max_time_limit") => { + FinishReason::Length + } Some("content_filter") => FinishReason::ContentFilter, _ => FinishReason::Other, } diff --git a/backend/crates/providers/xai/src/transport/catalog.rs b/backend/crates/providers/xai/src/transport/catalog.rs index e457de139..8e270aa0c 100644 --- a/backend/crates/providers/xai/src/transport/catalog.rs +++ b/backend/crates/providers/xai/src/transport/catalog.rs @@ -570,17 +570,12 @@ fn validate_billing_config(config: &Map) -> Result<(), GrokBillin return Err(GrokBillingError::InvalidWire); } } - for field in [ - "monthlyLimit", - "used", - "onDemandCap", - "onDemandUsed", - "prepaidBalance", - ] { + for field in ["onDemandCap", "onDemandUsed", "prepaidBalance"] { if let Some(value) = config.get(field) && !value.is_null() { - validate_cent(value)?; + // 预付余额使用带符号的账本金额，限额和已使用量仍须非负 + validate_cent(value, field == "prepaidBalance")?; } } if let Some(period) = config.get("currentPeriod") @@ -610,12 +605,14 @@ fn validate_billing_config(config: &Map) -> Result<(), GrokBillin Ok(()) } -fn validate_cent(value: &Value) -> Result<(), GrokBillingError> { +fn validate_cent(value: &Value, allow_negative: bool) -> Result<(), GrokBillingError> { let Value::Object(cent) = value else { return Err(GrokBillingError::InvalidWire); }; if let Some(value) = cent.get("val") - && value.as_i64().is_none_or(|value| value < 0) + && value + .as_i64() + .is_none_or(|value| !allow_negative && value < 0) { return Err(GrokBillingError::InvalidWire); } @@ -761,6 +758,7 @@ impl GrokCatalogCapabilities { #[derive(Debug, Clone, PartialEq, Eq)] pub struct GrokCatalogLimits { context_window_tokens: Option, + max_context_window_tokens: Option, max_output_tokens: Option, } @@ -771,6 +769,12 @@ impl GrokCatalogLimits { self.context_window_tokens } + /// 返回目录默认窗口与可选窗口中的最大值；缺失表示未知 + #[must_use] + pub const fn max_context_window_tokens(&self) -> Option { + self.max_context_window_tokens + } + /// 返回明确最大输出 token；缺失表示未知 #[must_use] pub const fn max_output_tokens(&self) -> Option { @@ -958,69 +962,27 @@ enum GrokModelsObject { #[derive(Debug, Deserialize)] struct GrokModelWire { id: Option, - model: Option, - #[serde(rename = "modelId")] - model_id: Option, + model: String, name: Option, description: Option, - #[serde(rename = "contextWindow", alias = "context_window")] context_window: Option, - #[serde(rename = "maxCompletionTokens", alias = "max_completion_tokens")] + #[serde(default)] + context_windows: Vec, max_completion_tokens: Option, - #[serde(rename = "apiBackend", alias = "api_backend")] api_backend: Option, - #[serde(rename = "supportedInApi", alias = "supported_in_api")] supported_in_api: Option, - #[serde( - rename = "supportsReasoningEffort", - alias = "supports_reasoning_effort" - )] - supports_reasoning_effort: Option, - #[serde(rename = "reasoningEffort", alias = "reasoning_effort")] - reasoning_effort: Option, - #[serde(default, rename = "reasoningEfforts", alias = "reasoning_efforts")] + #[serde(default)] reasoning_efforts: Vec, - features: Option, - #[serde(rename = "supportsBackendSearch", alias = "supports_backend_search")] supports_backend_search: Option, - #[serde(rename = "streamToolCalls", alias = "stream_tool_calls")] stream_tool_calls: Option, hidden: Option, } #[derive(Debug, Deserialize)] -struct GrokModelFeaturesWire { - reasoning: Option, - #[serde(rename = "reasoningEffortOptions", alias = "reasoning_effort_options")] - reasoning_effort_options: Option, -} - -#[derive(Debug, Deserialize)] -struct GrokReasoningEffortOptionsWire { - #[serde(default, rename = "supportedEfforts", alias = "supported_efforts")] - supported_efforts: Vec, - #[serde(rename = "defaultEffort", alias = "default_effort")] - default_effort: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -enum GrokReasoningEffortOptionWire { - Bare(GrokCatalogReasoningEffort), - Detailed { - value: GrokCatalogReasoningEffort, - #[serde(default)] - default: bool, - }, -} - -impl GrokReasoningEffortOptionWire { - fn into_parts(self) -> (GrokCatalogReasoningEffort, bool) { - match self { - Self::Bare(value) => (value, false), - Self::Detailed { value, default } => (value, default), - } - } +struct GrokReasoningEffortOptionWire { + value: GrokCatalogReasoningEffort, + #[serde(default)] + default: bool, } /// 解析唯一受支持的 Grok 官方完整快照 @@ -1062,8 +1024,8 @@ pub fn parse_grok_model_catalog( } fn normalize_model(wire: GrokModelWire) -> Result { - let actual_model = wire.model.or(wire.model_id).or_else(|| wire.id.clone()); - let actual_model = actual_model.ok_or(GrokModelCatalogError::InvalidModelSlug)?; + let (context_window_tokens, max_context_window_tokens) = context_window_limits(&wire)?; + let actual_model = wire.model; if !valid_model_slug(&actual_model) { return Err(GrokModelCatalogError::InvalidModelSlug); } @@ -1082,55 +1044,28 @@ fn normalize_model(wire: GrokModelWire) -> Result>(), - options - .default_effort - .filter(|effort| *effort != GrokCatalogReasoningEffort::Unknown), - ) - }, - ); - let reasoning_efforts = if feature_reasoning_efforts.is_empty() { - top_level_reasoning_efforts - } else { - feature_reasoning_efforts - }; - let default_reasoning_effort = feature_default_reasoning_effort - .or_else(|| { - wire.reasoning_effort - .filter(|effort| *effort != GrokCatalogReasoningEffort::Unknown) + // 固定 CLI proxy 的对象菜单拥有档位与默认值，缺失声明时不替用户选择首项 + let mut default_reasoning_effort = None; + let reasoning_efforts = wire + .reasoning_efforts + .into_iter() + .filter_map(|option| { + if option.value == GrokCatalogReasoningEffort::Unknown { + return None; + } + if option.default && default_reasoning_effort.is_none() { + default_reasoning_effort = Some(option.value); + } + Some(option.value) }) - .or(top_level_menu_default) - .or_else(|| reasoning_efforts.first().copied()); - let default_reasoning_effort = default_reasoning_effort - .filter(|default| reasoning_efforts.iter().any(|effort| effort == default)) - .or_else(|| reasoning_efforts.first().copied()); + .collect::>(); let reasoning_effort = if reasoning_efforts.is_empty() { - GrokCatalogCapabilityEvidence::from_wire( - wire.supports_reasoning_effort.or(feature_reasoning), - ) + GrokCatalogCapabilityEvidence::Unknown } else { GrokCatalogCapabilityEvidence::DeclaredNative }; @@ -1149,6 +1084,7 @@ fn normalize_model(wire: GrokModelWire) -> Result Result, - menu_default: &mut Option, -) -> Vec { - values +fn context_window_limits( + wire: &GrokModelWire, +) -> Result<(Option, Option), GrokModelCatalogError> { + let default = + optional_positive(wire.context_window)?.or_else(|| wire.context_windows.first().copied()); + let maximum = default .into_iter() - .filter_map(|option| { - let (effort, is_default) = option.into_parts(); - if effort == GrokCatalogReasoningEffort::Unknown { - return None; - } - if is_default && menu_default.is_none() { - *menu_default = Some(effort); - } - Some(effort) - }) - .collect() + .chain(wire.context_windows.iter().copied()) + .max(); + Ok((default, maximum)) } fn responses_evidence( diff --git a/backend/crates/providers/xai/src/transport/mod.rs b/backend/crates/providers/xai/src/transport/mod.rs index d7fe899a2..6c7549b50 100644 --- a/backend/crates/providers/xai/src/transport/mod.rs +++ b/backend/crates/providers/xai/src/transport/mod.rs @@ -49,6 +49,7 @@ pub use network::{ ReqwestGrokModelCatalogTransport, ReqwestOAuthTransport, }; pub use profile::{XaiWireProfile, XaiWireProfileState}; +pub(crate) use request::GrokReplayItem; pub use request::{GrokRequestEncodeError, GrokResponsesRequest}; pub use session::{ GrokCredentialFailure, GrokCredentialFeedbackFuture, GrokSessionAffinityKey, diff --git a/backend/crates/providers/xai/src/transport/request/mod.rs b/backend/crates/providers/xai/src/transport/request/mod.rs index bcfa08aa2..c96434c99 100644 --- a/backend/crates/providers/xai/src/transport/request/mod.rs +++ b/backend/crates/providers/xai/src/transport/request/mod.rs @@ -12,7 +12,7 @@ use super::{GrokSessionAffinityKey, XAI_PROVIDER_NAME}; mod history; mod identity; -mod model; +mod parameters; mod response; mod schema; mod tools; @@ -22,7 +22,7 @@ pub(crate) use response::GrokResponseTransform; use history::*; use identity::*; -use model::*; +use parameters::*; use schema::*; use tools::*; @@ -40,6 +40,14 @@ pub struct GrokResponsesRequest { response_transform: GrokResponseTransform, } +/// 重放条目在进入 wire 编码前的来源，避免将已编码别名再次当作客户端名称解释 +pub(crate) enum GrokReplayItem { + /// 已由请求转换器编码的输入，保留上游别名 + UpstreamInput(Value), + /// 已向客户端回译的输出，需要按本轮工具声明重新编码 + ClientOutput(Value), +} + impl GrokResponsesRequest { /// 返回发送到 `/v1/responses` 的 JSON object #[must_use] @@ -99,18 +107,29 @@ impl GrokResponsesRequest { pub(crate) fn set_replay_input( &mut self, - input: Vec, + input: Vec, ) -> Result<(), GrokRequestEncodeError> { let mut normalizer = ToolNormalizer::for_replay(self.response_transform.clone()); let mut normalized = Vec::with_capacity(input.len()); for item in input { match item { - Value::Object(item) if string_field(&item, "type") == "custom_tool_call" => { + GrokReplayItem::ClientOutput(Value::Object(item)) + if string_field(&item, "type") == "custom_tool_call" => + { normalized.push(Value::Object( normalizer.normalize_custom_tool_call_input(&item)?, )); } - item => normalized.push(item), + GrokReplayItem::ClientOutput(Value::Object(item)) + if string_field(&item, "type") == "function_call" => + { + normalized.push(Value::Object( + normalizer.normalize_function_call_input(&item)?, + )); + } + GrokReplayItem::UpstreamInput(item) | GrokReplayItem::ClientOutput(item) => { + normalized.push(item); + } } } self.response_transform = normalizer.response; @@ -201,11 +220,16 @@ impl GrokResponsesRequest { if payload.protocol() != "openai" { return Err(GrokRequestEncodeError::InvalidProtocolPayload); } + // Grok Build 会把 Codex 的预热请求当作真实生成；必须在发送前拒绝， + // 不能删掉控制字段后继续调用，或在已经计费后隐藏响应与费用 + if payload.body().get("generate").and_then(Value::as_bool) == Some(false) { + return Err(GrokRequestEncodeError::UnsupportedPrewarm); + } let mut body = payload.body().clone(); if should_consume_terminal_compaction_trigger { consume_terminal_compaction_trigger(&mut body)?; } - let upstream_model = resolve_grok_text_responses_model_id(upstream_model); + let upstream_model = upstream_model.to_owned(); // 这些字段属于 Codex/OpenAI 侧请求控制，不是 xAI 上游协议字段 // OpenAI 透明路径会保留未知字段；这里只在 xAI adapter 内做最小剥离 body.remove("provider_options"); @@ -220,7 +244,7 @@ impl GrokResponsesRequest { ); sanitize_account_identity(&mut body); sanitize_client_metadata(&mut body); - normalize_build_request(&mut body, &upstream_model)?; + normalize_build_request(&mut body)?; let mut response_transform = normalize_responses_request(&mut body)?; response_transform.observe_client_cache_tools(); if enable_cache_route { @@ -258,33 +282,6 @@ impl GrokResponsesRequest { } } -fn resolve_grok_text_responses_model_id(model: &str) -> String { - let stripped = strip_grok_provider_prefix(model); - let normalized = stripped.to_ascii_lowercase(); - match normalized.as_str() { - "grok" | "grok-latest" | "grok-4.5" | "grok-4.5-latest" | "grok-build-latest" => { - "grok-4.5".to_owned() - } - "grok-4.6" | "grok-4.6-latest" => "grok-4.6".to_owned(), - "grok-4.3" | "grok-4.3-latest" => "grok-4.3".to_owned(), - "grok-3-mini" - | "grok-3-mini-fast" - | "grok-build-0.1" - | "grok-composer-2.5-fast" - | "grok-4.20-0309-reasoning" - | "grok-4.20-0309-non-reasoning" - | "grok-4.20-multi-agent-0309" => normalized, - "grok-build" => "grok-build-0.1".to_owned(), - "grok-composer" | "composer-2.5" => "grok-composer-2.5-fast".to_owned(), - "grok-4.20-reasoning" => "grok-4.20-0309-reasoning".to_owned(), - "grok-4.20-non-reasoning" => "grok-4.20-0309-non-reasoning".to_owned(), - "grok-4.20-multi-agent" | "grok-4.20-multi-agent-latest" => { - "grok-4.20-multi-agent-0309".to_owned() - } - _ => stripped.to_owned(), - } -} - fn consume_terminal_compaction_trigger( body: &mut Map, ) -> Result<(), GrokRequestEncodeError> { @@ -322,6 +319,9 @@ pub enum GrokRequestEncodeError { /// 数据面只接受 OpenAI adapter 保留的原始 Responses object #[error("Grok Build request is missing its OpenAI protocol payload")] InvalidProtocolPayload, + /// Grok Build 不支持 Codex 的非生成预热语义 + #[error("Grok Build does not support non-generating prewarm requests (generate=false)")] + UnsupportedPrewarm, /// JSON 序列化意外失败 #[error("Grok Build request serialization failed")] Serialization, diff --git a/backend/crates/providers/xai/src/transport/request/model.rs b/backend/crates/providers/xai/src/transport/request/model.rs deleted file mode 100644 index 28f805f02..000000000 --- a/backend/crates/providers/xai/src/transport/request/model.rs +++ /dev/null @@ -1,144 +0,0 @@ -//! 当前 Responses 参数到 Grok 模型能力的转换 - -use super::*; - -pub(super) fn normalize_build_request( - body: &mut Map, - upstream_model: &str, -) -> Result<(), GrokRequestEncodeError> { - apply_build_response_defaults(body)?; - normalize_build_reasoning_effort(body, upstream_model); - sanitize_build_model_fields(body, upstream_model); - Ok(()) -} - -fn apply_build_response_defaults( - body: &mut Map, -) -> Result<(), GrokRequestEncodeError> { - if body.get("store").is_none_or(Value::is_null) { - body.insert("store".to_owned(), Value::Bool(false)); - } - - let include = body - .entry("include".to_owned()) - .or_insert_with(|| Value::Array(Vec::new())); - if include.is_null() { - *include = Value::Array(Vec::new()); - } - let include = include - .as_array_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "include" })?; - if include.iter().any(|value| !value.is_string()) { - return Err(GrokRequestEncodeError::InvalidRequestField { field: "include" }); - } - if !include - .iter() - .any(|value| value.as_str() == Some("reasoning.encrypted_content")) - { - include.push(Value::String("reasoning.encrypted_content".to_owned())); - } - Ok(()) -} - -fn normalize_build_reasoning_effort(body: &mut Map, upstream_model: &str) { - let supports_effort = grok_model_supports_reasoning_effort(upstream_model); - if let Some(reasoning) = body.get_mut("reasoning").and_then(Value::as_object_mut) { - if let Some(effort) = reasoning.remove("effort") - && supports_effort - && let Some(effort) = normalize_grok_reasoning_effort_value(&effort, upstream_model) - { - reasoning.insert("effort".to_owned(), Value::String(effort.to_owned())); - } - if reasoning.is_empty() { - body.remove("reasoning"); - } - } - - if is_grok_composer_model(upstream_model) { - body.remove("reasoning"); - } -} - -fn normalize_grok_reasoning_effort_value(value: &Value, model: &str) -> Option<&'static str> { - let normalized = value.as_str()?.trim().to_ascii_lowercase(); - match normalized.as_str() { - "none" => Some("none"), - "minimal" | "low" => Some("low"), - "medium" => Some("medium"), - "xhigh" if matches!(model, "grok-4.6" | "grok-4.6-latest") => Some("xhigh"), - "high" | "xhigh" => Some("high"), - _ => None, - } -} - -pub(super) fn strip_grok_provider_prefix(model: &str) -> &str { - let model = model.trim(); - let lower = model.to_ascii_lowercase(); - ["xai/", "x-ai/", "grok/"] - .into_iter() - .find_map(|prefix| { - lower - .starts_with(prefix) - .then(|| model[prefix.len()..].trim()) - }) - .unwrap_or(model) -} - -pub(super) fn grok_model_last_slug(model: &str) -> String { - model - .trim() - .rsplit_once('/') - .map_or_else(|| model.trim(), |(_, slug)| slug.trim()) - .to_ascii_lowercase() -} - -pub(super) fn is_grok_composer_model(model: &str) -> bool { - matches!( - grok_model_last_slug(model).as_str(), - "grok-composer" | "grok-composer-2.5-fast" | "composer-2.5" - ) -} - -fn grok_model_supports_reasoning_effort(model: &str) -> bool { - matches!( - strip_grok_provider_prefix(model) - .to_ascii_lowercase() - .as_str(), - "grok-4.5" - | "grok-4.5-latest" - | "grok-4.6" - | "grok-4.6-latest" - | "grok-4.3" - | "grok-4.3-latest" - | "grok-3-mini" - | "grok-3-mini-fast" - | "grok-4.20-0309-reasoning" - | "grok-4.20-reasoning" - | "grok-4.20-multi-agent-0309" - ) -} - -fn sanitize_build_model_fields(body: &mut Map, upstream_model: &str) { - for field in ["prompt_cache_retention", "safety_identifier"] { - body.remove(field); - } - if upstream_model.trim().eq_ignore_ascii_case("grok-4.5") { - for field in [ - "presence_penalty", - "presencePenalty", - "frequency_penalty", - "frequencyPenalty", - "stop", - ] { - body.remove(field); - } - } - if ["grok-4.20", "grok-4.3", "grok-4.5", "grok-4.6"] - .iter() - .any(|prefix| grok_model_last_slug(upstream_model).starts_with(prefix)) - { - body.remove("logprobs"); - body.remove("top_logprobs"); - } - body.remove("external_web_access"); -} diff --git a/backend/crates/providers/xai/src/transport/request/parameters.rs b/backend/crates/providers/xai/src/transport/request/parameters.rs new file mode 100644 index 000000000..99c33ec88 --- /dev/null +++ b/backend/crates/providers/xai/src/transport/request/parameters.rs @@ -0,0 +1,76 @@ +//! 当前 Responses 参数默认值与 Grok wire 字段校验 + +use super::*; + +pub(super) fn normalize_build_request( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + apply_build_response_defaults(body)?; + normalize_build_reasoning_effort(body)?; + sanitize_build_fields(body); + Ok(()) +} + +fn apply_build_response_defaults( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + if body.get("store").is_none_or(Value::is_null) { + body.insert("store".to_owned(), Value::Bool(false)); + } + + let include = body + .entry("include".to_owned()) + .or_insert_with(|| Value::Array(Vec::new())); + if include.is_null() { + *include = Value::Array(Vec::new()); + } + let include = include + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "include" })?; + if include.iter().any(|value| !value.is_string()) { + return Err(GrokRequestEncodeError::InvalidRequestField { field: "include" }); + } + if !include + .iter() + .any(|value| value.as_str() == Some("reasoning.encrypted_content")) + { + include.push(Value::String("reasoning.encrypted_content".to_owned())); + } + Ok(()) +} + +fn normalize_build_reasoning_effort( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + let Some(reasoning) = body.get_mut("reasoning").filter(|value| !value.is_null()) else { + return Ok(()); + }; + let reasoning = reasoning + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "reasoning" })?; + let Some(effort) = reasoning.get_mut("effort").filter(|value| !value.is_null()) else { + return Ok(()); + }; + let normalized = effort + .as_str() + .map(|value| value.trim().to_ascii_lowercase()) + .filter(|value| { + matches!( + value.as_str(), + "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max" + ) + }) + .ok_or(GrokRequestEncodeError::InvalidRequestField { + field: "reasoning.effort", + })?; + // 菜单属于目录事实；请求编码只校验 wire，不根据 slug 改写用户选择的档位 + *effort = Value::String(normalized); + Ok(()) +} + +fn sanitize_build_fields(body: &mut Map) { + for field in ["prompt_cache_retention", "safety_identifier"] { + body.remove(field); + } + body.remove("external_web_access"); +} diff --git a/backend/crates/providers/xai/src/transport/request/tools.rs b/backend/crates/providers/xai/src/transport/request/tools.rs index e30e2e3ab..fdac10f73 100644 --- a/backend/crates/providers/xai/src/transport/request/tools.rs +++ b/backend/crates/providers/xai/src/transport/request/tools.rs @@ -121,13 +121,26 @@ impl ToolNormalizer { identity.kind as u8, identity.namespace, identity.name ); let mut alias = truncate_tool_alias(&base, &key); - if self + // Hosted tool 和 function 共用上游名称空间；缓存路由还会在归一化后注入搜索工具， + // 因此这两个官方 wire name 始终保留，避免同一函数随会话工具集变化而改名 + if is_hosted_search_wire_name(&alias) + || self + .response + .aliases + .get(&alias) + .is_some_and(|existing| existing != &identity) + { + alias = hashed_tool_alias(&base, &key); + } + let mut collision_index = 0; + while self .response .aliases .get(&alias) .is_some_and(|existing| existing != &identity) { - alias = hashed_tool_alias(&base, &key); + collision_index += 1; + alias = hashed_tool_alias(&base, &format!("{key}\0{collision_index}")); } self.response .aliases @@ -396,6 +409,10 @@ impl ToolNormalizer { } } +fn is_hosted_search_wire_name(name: &str) -> bool { + matches!(name, "web_search" | "x_search") +} + pub(super) fn optional_array( value: Option<&Value>, ) -> Result<(Vec, bool), GrokRequestEncodeError> { @@ -765,7 +782,7 @@ impl ToolNormalizer { normalized_tools: &[Value], ) -> Result<(), GrokRequestEncodeError> { if let Some(Value::Object(function)) = object.get_mut("function") { - rewrite_namespace_choice(function, &self.identity_aliases)?; + rewrite_function_choice(function, &self.identity_aliases)?; let name = string_field(function, "name").trim(); if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { payload.remove("tool_choice"); @@ -774,7 +791,7 @@ impl ToolNormalizer { payload.insert("tool_choice".to_owned(), Value::Object(object)); return Ok(()); } - rewrite_namespace_choice(&mut object, &self.identity_aliases)?; + rewrite_function_choice(&mut object, &self.identity_aliases)?; let name = string_field(&object, "name").trim(); if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { payload.remove("tool_choice"); @@ -788,27 +805,40 @@ impl ToolNormalizer { pub(super) fn normalize_web_search_filters( tool: &Map, ) -> Result>, GrokRequestEncodeError> { - let nested = match tool.get("filters") { - None | Some(Value::Null) => None, - Some(Value::Object(filters)) => filters - .get("allowed_domains") - .map(normalize_allowed_domains) - .transpose()?, - Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let top_level = tool - .get("allowed_domains") - .map(normalize_allowed_domains) - .transpose()?; - if nested.is_some() && top_level.is_some() && nested != top_level { + // 域名限制只接受当前 filters 合同；丢弃已知旧字段会把受限搜索变成无限制搜索 + if tool.contains_key("allowed_domains") { return Err(GrokRequestEncodeError::InvalidRequestNormalization); } - Ok(nested - .or(top_level) + let nested = match tool.get("filters") { + None | Some(Value::Null) => None, + Some(Value::Object(filters)) => Some(filters), + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let allowed = nested + .and_then(|filters| filters.get("allowed_domains")) + .map(normalize_web_search_domains) + .transpose()? + .filter(|domains| !domains.is_empty()); + let excluded = nested + .and_then(|filters| filters.get("excluded_domains")) + .map(normalize_web_search_domains) + .transpose()? + .filter(|domains| !domains.is_empty()); + if let Some(excluded) = excluded { + // Grok Build 的非空允许表和排除表互斥；空表不限制搜索，不能覆盖另一张表 + if allowed.is_some() { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + return Ok(Some(Map::from_iter([( + "excluded_domains".to_owned(), + Value::Array(excluded), + )]))); + } + Ok(allowed .map(|domains| Map::from_iter([("allowed_domains".to_owned(), Value::Array(domains))]))) } -pub(super) fn normalize_allowed_domains( +pub(super) fn normalize_web_search_domains( value: &Value, ) -> Result, GrokRequestEncodeError> { let domains = value @@ -862,21 +892,24 @@ pub(super) fn tools_of_type(tools: &[Value], kind: &str) -> Vec { .collect() } -pub(super) fn rewrite_namespace_choice( +pub(super) fn rewrite_function_choice( object: &mut Map, aliases: &BTreeMap, ) -> Result<(), GrokRequestEncodeError> { let name = string_field(object, "name").trim(); let namespace = string_field(object, "namespace").trim(); - if name.is_empty() || namespace.is_empty() { + if name.is_empty() { return Ok(()); } let identity = ToolIdentity::new(ToolKind::Function, namespace, name); - let alias = aliases - .get(&identity) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - object.insert("name".to_owned(), Value::String(alias)); + let Some(alias) = aliases.get(&identity) else { + return if namespace.is_empty() { + Ok(()) + } else { + Err(GrokRequestEncodeError::InvalidRequestNormalization) + }; + }; + object.insert("name".to_owned(), Value::String(alias.clone())); object.remove("namespace"); Ok(()) } @@ -1085,7 +1118,7 @@ impl ToolNormalizer { Ok(converted) } - fn normalize_function_call_input( + pub(super) fn normalize_function_call_input( &mut self, item: &Map, ) -> Result, GrokRequestEncodeError> { @@ -1093,8 +1126,11 @@ impl ToolNormalizer { let call_id = required_trimmed_string(item, "call_id")?; let arguments = encode_function_arguments(item.get("arguments"))?; let namespace = string_field(item, "namespace").trim(); - if !namespace.is_empty() { - name = self.alias(ToolIdentity::new(ToolKind::Function, namespace, &name)); + let identity = ToolIdentity::new(ToolKind::Function, namespace, &name); + if !namespace.is_empty() || is_hosted_search_wire_name(&name) { + name = self.alias(identity); + } else if let Some(alias) = self.identity_aliases.get(&identity) { + name.clone_from(alias); } let mut converted = item.clone(); strip_grok_internal_keys(&mut converted); diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index c0620025a..73031744b 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -485,7 +485,12 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() quota: OpaqueProviderData::new( serde_json::from_value(serde_json::json!({ "subscriptionTier": "Free", - "config": {} + "config": { + "creditUsagePercent": 25, + "onDemandCap": {"val": 500}, + "onDemandUsed": {"val": 100}, + "prepaidBalance": {"val": -500} + } })) .expect("quota document"), ), @@ -503,6 +508,19 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() .await .expect("cached subscription"); assert_eq!(quota.plan_type.as_deref(), Some("Free")); + assert_eq!(quota.windows.len(), 1); + assert_eq!(quota.windows[0].used_percent, Some(25.0)); + let window_data = quota.windows[0] + .provider_data + .as_ref() + .expect("quota details") + .expose_to_provider() + .expose_to_provider(); + assert_eq!(window_data["onDemandCapCents"], 500); + assert_eq!(window_data["onDemandUsedCents"], 100); + assert_eq!(window_data["prepaidBalanceCents"], -500); + assert!(!window_data.contains_key("monthlyLimitCents")); + assert!(!window_data.contains_key("includedUsedCents")); let models = admin .models(account.id(), false) .await diff --git a/backend/crates/providers/xai/tests/credential/catalog.rs b/backend/crates/providers/xai/tests/credential/catalog.rs index 4835c4ce1..ff8b9d4ff 100644 --- a/backend/crates/providers/xai/tests/credential/catalog.rs +++ b/backend/crates/providers/xai/tests/credential/catalog.rs @@ -5,31 +5,27 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; -use futures::future::{BoxFuture, join_all}; +use futures::future::BoxFuture; use gateway_core::account::{ - AccountStateChange, CredentialCasUpdate, CredentialRevision, CredentialState, - OpaqueProviderData, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, - QuotaAccessChange, QuotaEvidence, QuotaObservation, QuotaState, + CredentialRevision, OpaqueProviderData, ProviderAccountStore, ProviderAccountUpdate, + QuotaAccessChange, QuotaEvidence, QuotaState, }; use gateway_core::provider_ports::{ ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderStoreError, }; use provider_xai::{ - GROK_SUBSCRIPTION_URL, GrokBillingRequest, GrokBillingTransport, GrokBillingTransportError, - GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokBillingTransportResponse, GrokCatalogCache, GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialRepository, GrokModelCatalogRequest, GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind, GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokPlanCatalog, - GrokQuotaError, }; use crate::support::{ MemoryGrokCatalogCache, MemoryProviderAccountStore, account_id, create_input, seed_input, }; -const OFFICIAL_FIXTURE: &[u8] = - include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); +const CLI_PROXY_FIXTURE: &[u8] = + include_bytes!("../transport/catalog/fixtures/cli_proxy_models.json"); struct QueueCatalogTransport { calls: AtomicUsize, @@ -88,58 +84,6 @@ impl GrokModelCatalogTransport for QueueCatalogTransport { } } -struct QueueBillingTransport { - calls: AtomicUsize, - subscription: Option>, - responses: Mutex>>, -} - -impl QueueBillingTransport { - fn success(body: &[u8]) -> Arc { - Arc::new(Self { - calls: AtomicUsize::new(0), - subscription: None, - responses: Mutex::new(VecDeque::from([Ok(GrokBillingTransportResponse::new( - body, - ))])), - }) - } - - fn failure() -> Arc { - Arc::new(Self { - calls: AtomicUsize::new(0), - subscription: None, - responses: Mutex::new(VecDeque::from([Err(GrokBillingTransportError::new( - GrokBillingTransportErrorKind::Unavailable, - ))])), - }) - } -} - -impl GrokBillingTransport for QueueBillingTransport { - fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { - if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { - let response = self - .subscription - .clone() - .map(GrokBillingTransportResponse::new) - .ok_or_else(|| { - GrokBillingTransportError::new(GrokBillingTransportErrorKind::Unavailable) - }); - return Box::pin(async move { response }); - } - self.calls.fetch_add(1, Ordering::SeqCst); - let response = self - .responses - .lock() - .expect("billing response queue") - .pop_front() - .expect("one billing response"); - Box::pin(async move { response }) - } -} - -/// 每次读取都返回同一份无法解码的 catalog 文档的存储端口 struct CorruptCatalogCachePort; impl ProviderCatalogCachePort for CorruptCatalogCachePort { @@ -164,47 +108,6 @@ impl ProviderCatalogCachePort for CorruptCatalogCachePort { } } -enum BillingMutation { - State(AccountStateChange), - Credential(CredentialCasUpdate), -} - -struct MutatingBillingTransport { - store: Arc, - mutation: Mutex>, - body: Vec, -} - -impl GrokBillingTransport for MutatingBillingTransport { - fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { - if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { - return Box::pin(async { - Err(GrokBillingTransportError::new( - GrokBillingTransportErrorKind::Unavailable, - )) - }); - } - let store = Arc::clone(&self.store); - let mutation = self.mutation.lock().expect("mutation").take(); - let body = self.body.clone(); - Box::pin(async move { - match mutation.expect("one mutation per request") { - BillingMutation::State(change) => store - .apply_state_change(change) - .await - .expect("apply concurrent state"), - BillingMutation::Credential(update) => { - store - .compare_and_swap_credential(update) - .await - .expect("apply concurrent credential rotation"); - } - } - Ok(GrokBillingTransportResponse::new(body)) - }) - } -} - async fn repository_with_accounts( suffixes: &[(&str, &str)], ) -> (Arc, GrokCredentialRepository) { @@ -219,38 +122,6 @@ async fn repository_with_accounts( (store, repository) } -async fn set_account_state( - store: &MemoryProviderAccountStore, - id: &ProviderAccountId, - credential_state: CredentialState, -) { - store - .apply_state_change(AccountStateChange { - message: None, - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - credential_state, - error_reason: credential_state.error_reason(), - observed_at: SystemTime::now(), - }) - .await - .expect("set account state"); -} - -#[tokio::test] -async fn concurrent_cold_scheduling_hydration_reads_quota_once() { - let (store, repository) = - repository_with_accounts(&[("quota-hydration", "subject-hydration")]).await; - let account = store - .account(&account_id("quota-hydration")) - .expect("created account"); - let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); - - join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await; - - assert_eq!(store.quota_reads(), 1); -} - #[tokio::test] async fn catalog_query_caches_each_plan_and_returns_strict_union() { let (store, repository) = @@ -258,10 +129,10 @@ async fn catalog_query_caches_each_plan_and_returns_strict_union() { let cache = MemoryGrokCatalogCache::shared(); let cache_port: Arc = cache.clone(); let transport = QueueCatalogTransport::from_bodies([ - OFFICIAL_FIXTURE.to_vec(), - OFFICIAL_FIXTURE.to_vec(), - OFFICIAL_FIXTURE.to_vec(), - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), ]); let service = crate::support::grok_catalog_service(repository, transport, cache_port); assert_eq!(service.catalog_generation().get(), 0); @@ -293,7 +164,7 @@ async fn single_account_catalog_refresh_and_read_use_provider_cache_boundary() { let cache_port: Arc = cache; let service = crate::support::grok_catalog_service( repository, - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]), + QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]), cache_port, ); let refreshed = service @@ -324,7 +195,7 @@ async fn disabled_account_refresh_discovers_models_with_the_pinned_account() { .expect("disable account"); let service = crate::support::grok_catalog_service( repository, - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]), + QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]), MemoryGrokCatalogCache::shared(), ); @@ -374,7 +245,7 @@ async fn corrupt_catalog_cache_entry_reads_as_miss_and_refetches() { .is_none() ); - let transport = QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]); + let transport = QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]); let service = crate::support::grok_catalog_service(repository, transport.clone(), Arc::new(cache)); let catalog = service @@ -432,7 +303,7 @@ async fn manual_catalog_refresh_replaces_the_shared_plan_cache() { )) .await .expect("cache catalog"); - let transport = QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]); + let transport = QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]); let service = crate::support::grok_catalog_service(repository, transport.clone(), cache); let refreshed = service @@ -456,7 +327,7 @@ async fn catalog_refresh_falls_back_within_the_same_plan() { GrokModelCatalogTransportErrorKind::Unavailable, )), Ok(GrokModelCatalogTransportResponse::new( - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), None, )), ]); @@ -538,7 +409,7 @@ async fn quota_exhausted_account_remains_eligible_for_catalog_discovery() { .expect("mark quota exhausted"); let service = crate::support::grok_catalog_service( repository, - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]), + QueueCatalogTransport::from_bodies([CLI_PROXY_FIXTURE.to_vec()]), MemoryGrokCatalogCache::shared(), ); @@ -582,7 +453,7 @@ async fn failed_plan_scope_is_skipped_and_surviving_plans_still_cache() { GrokModelCatalogTransportErrorKind::Unavailable, )), Ok(GrokModelCatalogTransportResponse::new( - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), None, )), ]); @@ -638,12 +509,12 @@ async fn conflicting_facts_for_same_slug_fail_closed() { .await .expect("separate plan catalog"); let mut conflicting: serde_json::Value = - serde_json::from_slice(OFFICIAL_FIXTURE).expect("fixture JSON"); + serde_json::from_slice(CLI_PROXY_FIXTURE).expect("fixture JSON"); conflicting["data"][0]["name"] = serde_json::json!("Different name"); let service = crate::support::grok_catalog_service( repository, QueueCatalogTransport::from_bodies([ - OFFICIAL_FIXTURE.to_vec(), + CLI_PROXY_FIXTURE.to_vec(), serde_json::to_vec(&conflicting).expect("conflicting JSON"), ]), MemoryGrokCatalogCache::shared(), @@ -665,450 +536,3 @@ fn seed_rejects_duplicates_and_supports_exact_membership() { assert!(seed.permits("grok-4.5")); assert!(!seed.permits("grok-4")); } - -#[tokio::test] -async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fields() { - let (store, repository) = repository_with_accounts(&[("quota", "subject-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":37.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-13T00:00:00Z","end":"2026-07-20T00:00:00Z"},"prepaidBalance":{"val":2500},"futureWindow":{"kind":"rolling"}}}"#, - ); - let service = crate::support::grok_quota_service(repository, transport.clone()); - - let snapshot = service - .refresh_account(&account_id("quota")) - .await - .expect("refresh quota"); - let persisted = store - .get_quotas(&[account_id("quota")]) - .await - .expect("read persisted quota") - .pop() - .expect("quota exists"); - let document = persisted.quota; - - assert_eq!(transport.calls.load(Ordering::SeqCst), 1); - assert_eq!(snapshot.billing().used_percent(), Some(37.5)); - assert_eq!(snapshot.billing().plan_type(), None); - assert_eq!( - snapshot.billing().period_kind(), - provider_xai::GrokQuotaPeriodKind::Weekly - ); - assert_eq!( - snapshot.billing().period_type(), - Some("USAGE_PERIOD_TYPE_WEEKLY") - ); - assert_eq!(snapshot.billing().prepaid_balance_cents(), Some(2500)); - assert!( - document.expose_to_provider()["config"] - .get("futureWindow") - .is_some() - ); -} - -#[tokio::test] -async fn quota_refresh_persists_subscription_without_user_profile_and_reads_it_from_cache() { - let (store, repository) = - repository_with_accounts(&[("subscription", "subject-subscription")]).await; - let mut transport = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#); - Arc::get_mut(&mut transport).expect("unique transport").subscription = Some( - br#"{"userId":"verified-user","email":"private@example.com","subscriptionTier":"SuperGrokPro"}"#.to_vec(), - ); - let service = crate::support::grok_quota_service(repository, transport.clone()); - let snapshot = service - .refresh_account(&account_id("subscription")) - .await - .expect("refresh quota"); - assert_eq!(snapshot.billing().plan_type(), Some("SuperGrokPro")); - let cached = service - .read_account(&account_id("subscription")) - .await - .expect("read quota") - .expect("cached quota"); - assert_eq!(cached.billing().plan_type(), Some("SuperGrokPro")); - assert_eq!(transport.calls.load(Ordering::SeqCst), 1); - let stored = store - .get_quotas(&[account_id("subscription")]) - .await - .expect("persisted quota"); - let document = stored[0].quota.expose_to_provider(); - assert_eq!(document["subscriptionTier"], "SuperGrokPro"); - assert!(document.get("email").is_none()); - assert!(document.get("userId").is_none()); -} - -#[tokio::test] -async fn billing_percent_does_not_rewrite_quota_access_fact() { - let (store, repository) = - repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await; - let id = account_id("still-exhausted"); - store - .apply_quota_access(QuotaAccessChange { - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), - }) - .await - .expect("mark quota exhausted"); - - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":100}}"#), - ) - .refresh_account(&id) - .await - .expect("refresh exhausted quota"); - - assert_eq!( - store.account(&id).expect("account").quota().access(), - gateway_core::account::QuotaAccessState::Exhausted - ); -} - -#[tokio::test] -async fn authoritative_billing_refresh_clears_existing_quota_exhaustion() { - let (store, repository) = - repository_with_accounts(&[("recovered-quota", "recovered-quota")]).await; - let id = account_id("recovered-quota"); - store - .apply_quota_access(QuotaAccessChange { - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), - }) - .await - .expect("mark quota exhausted"); - - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, - ), - ) - .refresh_account(&id) - .await - .expect("refresh recovered quota"); - - assert_eq!( - store.account(&id).expect("account").quota().access(), - gateway_core::account::QuotaAccessState::Allowed - ); -} - -#[tokio::test] -async fn full_usage_display_does_not_invent_quota_exhaustion() { - let (store, repository) = - repository_with_accounts(&[("exhausted-quota", "exhausted-quota")]).await; - let id = account_id("exhausted-quota"); - - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":100,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, - ), - ) - .refresh_account(&id) - .await - .expect("refresh exhausted quota"); - - let quota = store.account(&id).expect("account").quota(); - assert_eq!( - quota.access(), - gateway_core::account::QuotaAccessState::Unknown - ); - assert!(quota.reset_at().is_none()); -} - -#[tokio::test] -async fn recovered_quota_does_not_clear_terminal_account_states() { - for (suffix, credential_state) in [ - ("keep-banned", CredentialState::Banned), - ("keep-expired", CredentialState::Expired), - ("keep-invalid", CredentialState::Invalid), - ] { - let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await; - let id = account_id(suffix); - set_account_state(&store, &id, credential_state).await; - crate::support::grok_quota_service( - repository, - QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#), - ) - .refresh_account(&id) - .await - .expect("refresh terminal account quota"); - - assert_eq!( - store.account(&id).expect("account").credential_state(), - credential_state - ); - } -} - -#[tokio::test] -async fn quota_refresh_preserves_ready_credential_state_across_concurrent_write() { - let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await; - let id = account_id("new-cooldown"); - let transport = Arc::new(MutatingBillingTransport { - store: Arc::clone(&store), - mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange { - message: None, - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - credential_state: CredentialState::Ready, - error_reason: None, - observed_at: SystemTime::now(), - }))), - body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), - }); - - crate::support::grok_quota_service(repository, transport) - .refresh_account(&id) - .await - .expect("refresh around concurrent state write"); - - let account = store.account(&id).expect("account"); - assert_eq!(account.credential_state(), CredentialState::Ready); -} - -#[tokio::test] -async fn quota_refresh_rejects_a_concurrent_credential_revision() { - let (store, repository) = repository_with_accounts(&[("new-revision", "new-revision")]).await; - let id = account_id("new-revision"); - let account = store.account(&id).expect("account"); - let update = CredentialCasUpdate::new( - id.clone(), - account.revision(), - ProviderAccountUpdate { - account_id: id.clone(), - name: account.name().to_owned(), - email: account.email().map(str::to_owned), - plan_type: account.plan_type().map(str::to_owned), - }, - store.credential(&id).expect("credential"), - account.has_refresh_token(), - account.access_token_expires_at(), - account.next_refresh_at(), - ) - .expect("credential update"); - let transport = Arc::new(MutatingBillingTransport { - store: Arc::clone(&store), - mutation: Mutex::new(Some(BillingMutation::Credential(update))), - body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), - }); - - assert!(matches!( - crate::support::grok_quota_service(repository, transport) - .refresh_account(&id) - .await, - Err(GrokQuotaError::StaleCredentialSnapshot) - )); - assert_eq!(store.account(&id).expect("account").revision().get(), 2); -} - -#[tokio::test] -async fn quota_projection_falls_back_to_legacy_monthly_usage() { - let (_, repository) = - repository_with_accounts(&[("monthly-quota", "subject-monthly-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"monthlyLimit":{"val":10000},"used":{"val":2500},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("monthly-quota")) - .await - .expect("refresh monthly quota"); - - assert_eq!(snapshot.billing().used_percent(), Some(25.0)); - assert_eq!( - snapshot.billing().period_kind(), - provider_xai::GrokQuotaPeriodKind::Monthly - ); - assert_eq!( - snapshot.billing().period_end(), - Some("2026-08-01T00:00:00Z") - ); -} - -#[tokio::test] -async fn quota_projection_preserves_unknown_period_for_dynamic_duration_fallback() { - let (_, repository) = - repository_with_accounts(&[("dynamic-quota", "subject-dynamic-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_FORTNIGHT","start":"2026-07-01T00:00:00Z","end":"2026-07-15T00:00:00Z"}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("dynamic-quota")) - .await - .expect("refresh dynamic quota"); - - assert_eq!(snapshot.billing().used_percent(), Some(12.5)); - assert_eq!( - snapshot.billing().period_kind(), - provider_xai::GrokQuotaPeriodKind::Other - ); - assert_eq!( - snapshot.billing().period_start(), - Some("2026-07-01T00:00:00Z") - ); - assert_eq!( - snapshot.billing().period_end(), - Some("2026-07-15T00:00:00Z") - ); -} - -#[tokio::test] -async fn expired_billing_window_does_not_participate_in_scheduling_rank() { - let (store, repository) = - repository_with_accounts(&[("expired-billing", "expired-billing")]).await; - let id = account_id("expired-billing"); - let account = store.account(&id).expect("account"); - let service = crate::support::grok_quota_service( - repository, - QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2000-01-01T00:00:00Z","end":"2000-01-08T00:00:00Z"}}}"#, - ), - ); - - service - .refresh_account(&id) - .await - .expect("refresh expired billing window"); - - assert_eq!(service.scheduling_signals(&account), None); - assert_eq!( - store - .account(&id) - .expect("refreshed account") - .quota() - .access(), - gateway_core::account::QuotaAccessState::Unknown - ); -} - -#[tokio::test] -async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustion() { - let (store, repository) = - repository_with_accounts(&[("free-quota", "subject-free-quota")]).await; - let id = account_id("free-quota"); - store - .apply_quota_access(QuotaAccessChange { - account_id: id.clone(), - expected_revision: CredentialRevision::new(1).expect("revision"), - state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), - }) - .await - .expect("mark quota exhausted"); - let transport = QueueBillingTransport::success( - br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&id) - .await - .expect("refresh Free quota"); - - assert!(!snapshot.billing().has_authoritative_quota()); - assert_eq!( - store.account(&id).expect("account").quota().access(), - gateway_core::account::QuotaAccessState::Exhausted - ); -} - -#[tokio::test] -async fn reported_zero_percent_is_authoritative_quota() { - let (_, repository) = - repository_with_accounts(&[("zero-percent", "subject-zero-percent")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"creditUsagePercent":0,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("zero-percent")) - .await - .expect("refresh reported quota"); - - assert!(snapshot.billing().has_authoritative_quota()); -} - -#[tokio::test] -async fn positive_prepaid_balance_is_authoritative_quota() { - let (_, repository) = - repository_with_accounts(&[("prepaid-quota", "subject-prepaid-quota")]).await; - let transport = QueueBillingTransport::success( - br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"prepaidBalance":{"val":500}}}"#, - ); - let snapshot = crate::support::grok_quota_service(repository, transport) - .refresh_account(&account_id("prepaid-quota")) - .await - .expect("refresh prepaid quota"); - - assert!(snapshot.billing().has_authoritative_quota()); -} - -#[tokio::test] -async fn quota_read_rejects_corrupt_provider_document() { - let (store, repository) = repository_with_accounts(&[("corrupt", "subject-corrupt")]).await; - let mut document = serde_json::Map::new(); - document.insert("config".to_owned(), serde_json::json!([])); - store - .compare_and_swap_quota(QuotaObservation { - plan_type: None, - account_id: account_id("corrupt"), - expected_revision: CredentialRevision::new(1).expect("revision"), - quota: OpaqueProviderData::new(document), - observed_at: SystemTime::now(), - state: QuotaState::unknown(), - }) - .await - .expect("seed corrupt quota"); - let service = crate::support::grok_quota_service( - repository, - QueueBillingTransport::success(br#"{"config":null}"#), - ); - - assert!(matches!( - service.read_account(&account_id("corrupt")).await, - Err(GrokQuotaError::InvalidData) - )); -} - -#[tokio::test] -async fn disabled_account_quota_refresh_never_calls_upstream() { - let (store, repository) = - repository_with_accounts(&[("disabled-quota", "subject-disabled")]).await; - store - .set_enabled(&account_id("disabled-quota"), false) - .await - .expect("disable account"); - let transport = QueueBillingTransport::success(br#"{"config":null}"#); - let service = crate::support::grok_quota_service(repository, transport.clone()); - - assert!(matches!( - service.refresh_account(&account_id("disabled-quota")).await, - Err(GrokQuotaError::AccountUnavailable) - )); - assert_eq!(transport.calls.load(Ordering::SeqCst), 0); -} - -#[tokio::test] -async fn failed_quota_refresh_does_not_replace_last_good_observation() { - let (store, repository) = repository_with_accounts(&[("stable", "subject-stable")]).await; - let good = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#); - crate::support::grok_quota_service(repository.clone(), good) - .refresh_account(&account_id("stable")) - .await - .expect("seed good observation"); - let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); - - assert!(matches!( - service.refresh_account(&account_id("stable")).await, - Err(GrokQuotaError::Upstream) - )); - let persisted = store - .get_quotas(&[account_id("stable")]) - .await - .expect("read quota") - .pop() - .expect("quota remains") - .quota; - assert_eq!( - persisted.expose_to_provider()["config"]["creditUsagePercent"].as_f64(), - Some(10.0), - ); -} diff --git a/backend/crates/providers/xai/tests/credential/mod.rs b/backend/crates/providers/xai/tests/credential/mod.rs index f48673bf5..df4f292fd 100644 --- a/backend/crates/providers/xai/tests/credential/mod.rs +++ b/backend/crates/providers/xai/tests/credential/mod.rs @@ -10,6 +10,7 @@ mod http; mod import; mod oidc_verifier; mod pkce; +mod quota; mod refresh; mod repository; mod secret; diff --git a/backend/crates/providers/xai/tests/credential/quota.rs b/backend/crates/providers/xai/tests/credential/quota.rs new file mode 100644 index 000000000..0efd180eb --- /dev/null +++ b/backend/crates/providers/xai/tests/credential/quota.rs @@ -0,0 +1,788 @@ +//! 验证 xAI 额度观察、revision隔离、权限状态与调度投影 + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::SystemTime; + +use futures::future::join_all; +use gateway_core::account::{ + AccountStateChange, CredentialCasUpdate, CredentialRevision, CredentialState, + OpaqueProviderData, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, + QuotaAccessChange, QuotaEvidence, QuotaObservation, QuotaState, +}; +use provider_xai::{ + GROK_SUBSCRIPTION_URL, GrokBillingRequest, GrokBillingTransport, GrokBillingTransportError, + GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokBillingTransportResponse, + GrokCredentialRepository, GrokQuotaError, +}; + +use crate::support::{MemoryProviderAccountStore, account_id, create_input, seed_input}; + +struct QueueBillingTransport { + calls: AtomicUsize, + subscription: Option>, + responses: Mutex>>, +} + +impl QueueBillingTransport { + fn success(body: &[u8]) -> Arc { + Arc::new(Self { + calls: AtomicUsize::new(0), + subscription: None, + responses: Mutex::new(VecDeque::from([Ok(GrokBillingTransportResponse::new( + body, + ))])), + }) + } + + fn failure() -> Arc { + Arc::new(Self { + calls: AtomicUsize::new(0), + subscription: None, + responses: Mutex::new(VecDeque::from([Err(GrokBillingTransportError::new( + GrokBillingTransportErrorKind::Unavailable, + ))])), + }) + } +} + +impl GrokBillingTransport for QueueBillingTransport { + fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { + if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { + let response = self + .subscription + .clone() + .map(GrokBillingTransportResponse::new) + .ok_or_else(|| { + GrokBillingTransportError::new(GrokBillingTransportErrorKind::Unavailable) + }); + return Box::pin(async move { response }); + } + self.calls.fetch_add(1, Ordering::SeqCst); + let response = self + .responses + .lock() + .expect("billing response queue") + .pop_front() + .expect("one billing response"); + Box::pin(async move { response }) + } +} + +/// 每次读取都返回同一份无法解码的 catalog 文档的存储端口 +enum BillingMutation { + State(AccountStateChange), + Credential(CredentialCasUpdate), +} + +struct MutatingBillingTransport { + store: Arc, + mutation: Mutex>, + body: Vec, +} + +impl GrokBillingTransport for MutatingBillingTransport { + fn execute(&self, request: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { + if request.endpoint().as_str() == GROK_SUBSCRIPTION_URL { + return Box::pin(async { + Err(GrokBillingTransportError::new( + GrokBillingTransportErrorKind::Unavailable, + )) + }); + } + let store = Arc::clone(&self.store); + let mutation = self.mutation.lock().expect("mutation").take(); + let body = self.body.clone(); + Box::pin(async move { + match mutation.expect("one mutation per request") { + BillingMutation::State(change) => store + .apply_state_change(change) + .await + .expect("apply concurrent state"), + BillingMutation::Credential(update) => { + store + .compare_and_swap_credential(update) + .await + .expect("apply concurrent credential rotation"); + } + } + Ok(GrokBillingTransportResponse::new(body)) + }) + } +} + +async fn repository_with_accounts( + suffixes: &[(&str, &str)], +) -> (Arc, GrokCredentialRepository) { + let store = MemoryProviderAccountStore::shared(); + let account_store: Arc = store.clone(); + let repository = GrokCredentialRepository::new(account_store); + for (suffix, subject) in suffixes { + seed_input(&store, &create_input(suffix, subject)) + .await + .expect("create account"); + } + (store, repository) +} + +async fn set_account_state( + store: &MemoryProviderAccountStore, + id: &ProviderAccountId, + credential_state: CredentialState, +) { + store + .apply_state_change(AccountStateChange { + message: None, + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + credential_state, + error_reason: credential_state.error_reason(), + observed_at: SystemTime::now(), + }) + .await + .expect("set account state"); +} + +#[tokio::test] +async fn concurrent_cold_scheduling_hydration_reads_quota_once() { + let (store, repository) = + repository_with_accounts(&[("quota-hydration", "subject-hydration")]).await; + let account = store + .account(&account_id("quota-hydration")) + .expect("created account"); + let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); + + join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await; + + assert_eq!(store.quota_reads(), 1); +} + +#[tokio::test] +async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fields() { + let (store, repository) = repository_with_accounts(&[("quota", "subject-quota")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":37.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-13T00:00:00Z","end":"2026-07-20T00:00:00Z"},"prepaidBalance":{"val":2500},"futureWindow":{"kind":"rolling"}}}"#, + ); + let service = crate::support::grok_quota_service(repository, transport.clone()); + + let snapshot = service + .refresh_account(&account_id("quota")) + .await + .expect("refresh quota"); + let persisted = store + .get_quotas(&[account_id("quota")]) + .await + .expect("read persisted quota") + .pop() + .expect("quota exists"); + let document = persisted.quota; + + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + assert_eq!(snapshot.billing().used_percent(), Some(37.5)); + assert_eq!(snapshot.billing().plan_type(), None); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Weekly + ); + assert_eq!( + snapshot.billing().period_type(), + Some("USAGE_PERIOD_TYPE_WEEKLY") + ); + assert_eq!(snapshot.billing().prepaid_balance_cents(), Some(2500)); + assert!( + document.expose_to_provider()["config"] + .get("futureWindow") + .is_some() + ); +} + +#[tokio::test] +async fn quota_refresh_persists_subscription_without_user_profile_and_reads_it_from_cache() { + let (store, repository) = + repository_with_accounts(&[("subscription", "subject-subscription")]).await; + let mut transport = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#); + Arc::get_mut(&mut transport).expect("unique transport").subscription = Some( + br#"{"userId":"verified-user","email":"private@example.com","subscriptionTier":"SuperGrokPro"}"#.to_vec(), + ); + let service = crate::support::grok_quota_service(repository, transport.clone()); + let snapshot = service + .refresh_account(&account_id("subscription")) + .await + .expect("refresh quota"); + assert_eq!(snapshot.billing().plan_type(), Some("SuperGrokPro")); + let cached = service + .read_account(&account_id("subscription")) + .await + .expect("read quota") + .expect("cached quota"); + assert_eq!(cached.billing().plan_type(), Some("SuperGrokPro")); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + let stored = store + .get_quotas(&[account_id("subscription")]) + .await + .expect("persisted quota"); + let document = stored[0].quota.expose_to_provider(); + assert_eq!(document["subscriptionTier"], "SuperGrokPro"); + assert!(document.get("email").is_none()); + assert!(document.get("userId").is_none()); +} + +#[tokio::test] +async fn billing_percent_does_not_rewrite_quota_access_fact() { + let (store, repository) = + repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await; + let id = account_id("still-exhausted"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), + }) + .await + .expect("mark quota exhausted"); + + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":100}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh exhausted quota"); + + assert_eq!( + store.account(&id).expect("account").quota().access(), + gateway_core::account::QuotaAccessState::Exhausted + ); +} + +#[tokio::test] +async fn billing_access_and_scheduling_use_current_quota_fields() { + use gateway_core::account::QuotaAccessState::{Allowed, Exhausted}; + + for (config, expected_percent, expected_access) in [ + ( + serde_json::json!({"creditUsagePercent": 100}), + Some(100.0), + Exhausted, + ), + ( + serde_json::json!({"creditUsagePercent": 25}), + Some(25.0), + Allowed, + ), + ( + serde_json::json!({"onDemandCap": {"val": 500}, "onDemandUsed": {"val": 100}}), + None, + Allowed, + ), + ( + serde_json::json!({ + "creditUsagePercent": 100, + "onDemandCap": {"val": 500}, "onDemandUsed": {"val": 100} + }), + Some(100.0), + Allowed, + ), + ] { + let (store, repository) = repository_with_accounts(&[("percent", "percent")]).await; + let id = account_id("percent"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + SystemTime::now(), + None, + ), + }) + .await + .expect("mark quota exhausted"); + let body = serde_json::to_vec(&serde_json::json!({"config": config})).expect("billing"); + let service = + crate::support::grok_quota_service(repository, QueueBillingTransport::success(&body)); + let snapshot = service.refresh_account(&id).await.expect("refresh quota"); + let account = store.account(&id).expect("refreshed account"); + + assert_eq!(snapshot.billing().used_percent(), expected_percent); + assert_eq!( + account.quota().access(), + expected_access, + "config: {config}" + ); + assert_eq!( + service + .scheduling_signals(&account) + .and_then(|signals| signals.remaining_rank()), + expected_percent.map(|percent| ((100.0 - percent) * 100.0).round() as u64) + ); + } +} + +#[tokio::test] +async fn signed_prepaid_balance_preserves_wire_and_recovers_access_when_current() { + use gateway_core::account::QuotaAccessState::{Allowed, Exhausted}; + + for (balance, period_end, expected_access) in [ + (-500, None, Allowed), + (i64::MIN, None, Allowed), + (0, None, Exhausted), + (500, None, Allowed), + (-500, Some("2000-01-01T00:00:00Z"), Exhausted), + (-500, Some("2099-01-01T00:00:00Z"), Allowed), + ] { + let (store, repository) = repository_with_accounts(&[("signed", "signed")]).await; + let id = account_id("signed"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + SystemTime::now(), + None, + ), + }) + .await + .expect("mark quota exhausted"); + let body = serde_json::to_vec(&serde_json::json!({ + "config": { + "prepaidBalance": {"val": balance}, "currentPeriod": {"end": period_end} + } + })) + .expect("billing"); + let transport = QueueBillingTransport::success(&body); + let service = crate::support::grok_quota_service(repository, transport.clone()); + + let snapshot = service + .refresh_account(&id) + .await + .expect("refresh signed balance"); + let cached = service + .read_account(&id) + .await + .expect("read quota") + .expect("cached quota"); + let persisted = store + .get_quotas(std::slice::from_ref(&id)) + .await + .expect("stored quota"); + + assert_eq!(snapshot.billing().prepaid_balance_cents(), Some(balance)); + assert_eq!(snapshot.billing().has_authoritative_quota(), balance != 0); + assert_eq!(cached.billing().prepaid_balance_cents(), Some(balance)); + assert_eq!( + persisted[0].quota.expose_to_provider()["config"]["prepaidBalance"]["val"], + balance + ); + assert_eq!( + store.account(&id).expect("account").quota().access(), + expected_access, + "balance: {balance}, period_end: {period_end:?}" + ); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn authoritative_billing_refresh_clears_existing_quota_exhaustion() { + let (store, repository) = + repository_with_accounts(&[("recovered-quota", "recovered-quota")]).await; + let id = account_id("recovered-quota"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), + }) + .await + .expect("mark quota exhausted"); + + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, + ), + ) + .refresh_account(&id) + .await + .expect("refresh recovered quota"); + + assert_eq!( + store.account(&id).expect("account").quota().access(), + gateway_core::account::QuotaAccessState::Allowed + ); +} + +#[tokio::test] +async fn full_usage_display_does_not_invent_quota_exhaustion() { + let (store, repository) = + repository_with_accounts(&[("exhausted-quota", "exhausted-quota")]).await; + let id = account_id("exhausted-quota"); + + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":100,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2099-07-22T00:00:00Z"}}}"#, + ), + ) + .refresh_account(&id) + .await + .expect("refresh exhausted quota"); + + let quota = store.account(&id).expect("account").quota(); + assert_eq!( + quota.access(), + gateway_core::account::QuotaAccessState::Unknown + ); + assert!(quota.reset_at().is_none()); +} + +#[tokio::test] +async fn recovered_quota_does_not_clear_terminal_account_states() { + for (suffix, credential_state) in [ + ("keep-banned", CredentialState::Banned), + ("keep-expired", CredentialState::Expired), + ("keep-invalid", CredentialState::Invalid), + ] { + let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await; + let id = account_id(suffix); + set_account_state(&store, &id, credential_state).await; + crate::support::grok_quota_service( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh terminal account quota"); + + assert_eq!( + store.account(&id).expect("account").credential_state(), + credential_state + ); + } +} + +#[tokio::test] +async fn quota_refresh_preserves_ready_credential_state_across_concurrent_write() { + let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await; + let id = account_id("new-cooldown"); + let transport = Arc::new(MutatingBillingTransport { + store: Arc::clone(&store), + mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange { + message: None, + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + credential_state: CredentialState::Ready, + error_reason: None, + observed_at: SystemTime::now(), + }))), + body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), + }); + + crate::support::grok_quota_service(repository, transport) + .refresh_account(&id) + .await + .expect("refresh around concurrent state write"); + + let account = store.account(&id).expect("account"); + assert_eq!(account.credential_state(), CredentialState::Ready); +} + +#[tokio::test] +async fn quota_refresh_rejects_a_concurrent_credential_revision() { + let (store, repository) = repository_with_accounts(&[("new-revision", "new-revision")]).await; + let id = account_id("new-revision"); + let account = store.account(&id).expect("account"); + let update = CredentialCasUpdate::new( + id.clone(), + account.revision(), + ProviderAccountUpdate { + account_id: id.clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + }, + store.credential(&id).expect("credential"), + account.has_refresh_token(), + account.access_token_expires_at(), + account.next_refresh_at(), + ) + .expect("credential update"); + let transport = Arc::new(MutatingBillingTransport { + store: Arc::clone(&store), + mutation: Mutex::new(Some(BillingMutation::Credential(update))), + body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), + }); + + assert!(matches!( + crate::support::grok_quota_service(repository, transport) + .refresh_account(&id) + .await, + Err(GrokQuotaError::StaleCredentialSnapshot) + )); + assert_eq!(store.account(&id).expect("account").revision().get(), 2); +} + +#[tokio::test] +async fn legacy_billing_fields_do_not_supply_quota_evidence_or_period() { + for (limit, used) in [ + ( + serde_json::json!({"val": 10000}), + serde_json::json!({"val": 2500}), + ), + ( + serde_json::json!("uninterpreted"), + serde_json::json!({"val": -1}), + ), + ] { + let (store, repository) = repository_with_accounts(&[("legacy", "legacy")]).await; + let id = account_id("legacy"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + SystemTime::now(), + None, + ), + }) + .await + .expect("mark quota exhausted"); + let body = serde_json::to_vec(&serde_json::json!({"config": { + "monthlyLimit": limit, "used": used, + "billingPeriodStart": "2026-07-01T00:00:00Z", + "billingPeriodEnd": "2099-08-01T00:00:00Z" + }})) + .expect("billing document"); + let service = + crate::support::grok_quota_service(repository, QueueBillingTransport::success(&body)); + let snapshot = service + .refresh_account(&id) + .await + .expect("refresh opaque billing"); + let account = store.account(&id).expect("account"); + + assert_eq!(snapshot.billing().used_percent(), None); + assert!(!snapshot.billing().has_authoritative_quota()); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Other + ); + assert_eq!(snapshot.billing().period_start(), None); + assert_eq!(snapshot.billing().period_end(), None); + assert_eq!(service.scheduling_signals(&account), None); + assert_eq!( + account.quota().access(), + gateway_core::account::QuotaAccessState::Exhausted + ); + } +} + +#[tokio::test] +async fn current_quota_period_does_not_fill_missing_dates_from_legacy_fields() { + let (_, repository) = repository_with_accounts(&[("current-period", "current-period")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY"},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2099-08-01T00:00:00Z"}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("current-period")) + .await + .expect("current quota"); + + assert_eq!(snapshot.billing().used_percent(), Some(25.0)); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Weekly + ); + assert_eq!(snapshot.billing().period_start(), None); + assert_eq!(snapshot.billing().period_end(), None); +} + +#[tokio::test] +async fn quota_projection_preserves_unknown_period_for_dynamic_duration_fallback() { + let (_, repository) = + repository_with_accounts(&[("dynamic-quota", "subject-dynamic-quota")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_FORTNIGHT","start":"2026-07-01T00:00:00Z","end":"2026-07-15T00:00:00Z"}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("dynamic-quota")) + .await + .expect("refresh dynamic quota"); + + assert_eq!(snapshot.billing().used_percent(), Some(12.5)); + assert_eq!( + snapshot.billing().period_kind(), + provider_xai::GrokQuotaPeriodKind::Other + ); + assert_eq!( + snapshot.billing().period_start(), + Some("2026-07-01T00:00:00Z") + ); + assert_eq!( + snapshot.billing().period_end(), + Some("2026-07-15T00:00:00Z") + ); +} + +#[tokio::test] +async fn expired_billing_window_does_not_participate_in_scheduling_rank() { + let (store, repository) = + repository_with_accounts(&[("expired-billing", "expired-billing")]).await; + let id = account_id("expired-billing"); + let account = store.account(&id).expect("account"); + let service = crate::support::grok_quota_service( + repository, + QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2000-01-01T00:00:00Z","end":"2000-01-08T00:00:00Z"}}}"#, + ), + ); + + service + .refresh_account(&id) + .await + .expect("refresh expired billing window"); + + assert_eq!(service.scheduling_signals(&account), None); + assert_eq!( + store + .account(&id) + .expect("refreshed account") + .quota() + .access(), + gateway_core::account::QuotaAccessState::Unknown + ); +} + +#[tokio::test] +async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustion() { + let (store, repository) = + repository_with_accounts(&[("free-quota", "subject-free-quota")]).await; + let id = account_id("free-quota"); + store + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None), + }) + .await + .expect("mark quota exhausted"); + let transport = QueueBillingTransport::success( + br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&id) + .await + .expect("refresh Free quota"); + + assert!(!snapshot.billing().has_authoritative_quota()); + assert_eq!( + store.account(&id).expect("account").quota().access(), + gateway_core::account::QuotaAccessState::Exhausted + ); +} + +#[tokio::test] +async fn reported_zero_percent_is_authoritative_quota() { + let (_, repository) = + repository_with_accounts(&[("zero-percent", "subject-zero-percent")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"creditUsagePercent":0,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("zero-percent")) + .await + .expect("refresh reported quota"); + + assert!(snapshot.billing().has_authoritative_quota()); +} + +#[tokio::test] +async fn positive_prepaid_balance_is_authoritative_quota() { + let (_, repository) = + repository_with_accounts(&[("prepaid-quota", "subject-prepaid-quota")]).await; + let transport = QueueBillingTransport::success( + br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"prepaidBalance":{"val":500}}}"#, + ); + let snapshot = crate::support::grok_quota_service(repository, transport) + .refresh_account(&account_id("prepaid-quota")) + .await + .expect("refresh prepaid quota"); + + assert!(snapshot.billing().has_authoritative_quota()); +} + +#[tokio::test] +async fn quota_read_rejects_corrupt_provider_document() { + let (store, repository) = repository_with_accounts(&[("corrupt", "subject-corrupt")]).await; + let mut document = serde_json::Map::new(); + document.insert("config".to_owned(), serde_json::json!([])); + store + .compare_and_swap_quota(QuotaObservation { + plan_type: None, + account_id: account_id("corrupt"), + expected_revision: CredentialRevision::new(1).expect("revision"), + quota: OpaqueProviderData::new(document), + observed_at: SystemTime::now(), + state: QuotaState::unknown(), + }) + .await + .expect("seed corrupt quota"); + let service = crate::support::grok_quota_service( + repository, + QueueBillingTransport::success(br#"{"config":null}"#), + ); + + assert!(matches!( + service.read_account(&account_id("corrupt")).await, + Err(GrokQuotaError::InvalidData) + )); +} + +#[tokio::test] +async fn disabled_account_quota_refresh_never_calls_upstream() { + let (store, repository) = + repository_with_accounts(&[("disabled-quota", "subject-disabled")]).await; + store + .set_enabled(&account_id("disabled-quota"), false) + .await + .expect("disable account"); + let transport = QueueBillingTransport::success(br#"{"config":null}"#); + let service = crate::support::grok_quota_service(repository, transport.clone()); + + assert!(matches!( + service.refresh_account(&account_id("disabled-quota")).await, + Err(GrokQuotaError::AccountUnavailable) + )); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn failed_quota_refresh_does_not_replace_last_good_observation() { + let (store, repository) = repository_with_accounts(&[("stable", "subject-stable")]).await; + let good = QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#); + crate::support::grok_quota_service(repository.clone(), good) + .refresh_account(&account_id("stable")) + .await + .expect("seed good observation"); + let service = crate::support::grok_quota_service(repository, QueueBillingTransport::failure()); + + assert!(matches!( + service.refresh_account(&account_id("stable")).await, + Err(GrokQuotaError::Upstream) + )); + let persisted = store + .get_quotas(&[account_id("stable")]) + .await + .expect("read quota") + .pop() + .expect("quota remains") + .quota; + assert_eq!( + persisted.expose_to_provider()["config"]["creditUsagePercent"].as_f64(), + Some(10.0), + ); +} diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index 54afb2912..53f9407db 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -31,8 +31,8 @@ use crate::support::{ rejecting_account_proxy, runtime_policy, seed_input, }; -const OFFICIAL_FIXTURE: &[u8] = - include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); +const CLI_PROXY_FIXTURE: &[u8] = + include_bytes!("../transport/catalog/fixtures/cli_proxy_models.json"); const OAUTH_BACKOFF_ATTEMPTS: u32 = 5; #[tokio::test] @@ -190,7 +190,7 @@ impl GrokModelCatalogTransport for StaticCatalogTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { Box::pin(async { Ok(GrokModelCatalogTransportResponse::new( - OFFICIAL_FIXTURE, + CLI_PROXY_FIXTURE, None, )) }) diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract/mod.rs similarity index 94% rename from backend/crates/providers/xai/tests/provider/contract.rs rename to backend/crates/providers/xai/tests/provider/contract/mod.rs index 90fe1b11e..742cbb997 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract/mod.rs @@ -1,5 +1,7 @@ //! 验证 xAI 原生执行的协议转换、账号选择与交付状态隔离 +mod replay_tool_names; + use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::num::NonZeroU32; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -38,12 +40,12 @@ use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingState, ProviderStoreError, }; +use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::routing::{ ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, SupportLevel, UpstreamModelId, }; -use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::upstream::UpstreamSendState; use provider_xai::{ GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport, @@ -68,7 +70,116 @@ use crate::support::{ const MODEL: &str = "grok-4.5"; const CATALOG_FIXTURE: &[u8] = - include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); + include_bytes!("../../transport/catalog/fixtures/cli_proxy_models.json"); + +#[tokio::test] +async fn native_xai_rejects_generate_false_before_upstream_io_or_usage() { + for store in [false, true] { + let transport = StubInferenceTransport::success(); + let selector = StubSelector::success(); + let provider = provider(Arc::clone(&selector), transport.clone()).await; + let request = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + object(json!({ + "model": MODEL, + "input": "hello", + "generate": false, + "store": store + })), + ) + .unwrap(), + )); + let error = match provider + .execute( + provider_request_with_operation("xai", request), + context(CancellationToken::new(), None), + ) + .await + { + Ok(_) => panic!("non-generating requests must fail before a usage-producing stream"), + Err(error) => error, + }; + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + let detail = error + .client_visible_upstream_error() + .expect("prewarm rejection must explain the unsupported request"); + assert_eq!(detail.code(), Some("unsupported_prewarm")); + assert_eq!(detail.error_type(), Some("invalid_request_error")); + assert!(detail.message().contains("generate=false")); + assert_eq!(selector.calls.load(Ordering::SeqCst), 0); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); + assert!(transport.requests.lock().unwrap().is_empty()); + } +} + +#[tokio::test] +async fn native_xai_rejects_translated_generate_false_before_upstream_io() { + let transport = StubInferenceTransport::success(); + let selector = StubSelector::success(); + let provider = provider(Arc::clone(&selector), transport.clone()).await; + let source = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("example-source", object(json!({"opaque": "source"}))) + .unwrap(), + )); + let translations = Arc::new(Mutex::new(0)); + let error = match provider + .execute( + provider_request_with_operation("xai", source), + context_with_middleware(Arc::new(TranslationMiddleware { + translations: Arc::clone(&translations), + target: object(json!({"model": MODEL, "input": "hello", "generate": false})), + })), + ) + .await + { + Ok(_) => panic!("translated prewarm must not start generation"), + Err(error) => error, + }; + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert_eq!(*translations.lock().unwrap(), 1); + assert_eq!(selector.calls.load(Ordering::SeqCst), 1); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); + assert!(transport.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn native_xai_preserves_generate_true_and_omitted_requests() { + for generate in [None, Some(true)] { + let transport = StubInferenceTransport::success(); + let provider = provider(StubSelector::success(), transport.clone()).await; + let mut body = object(json!({"model": MODEL, "input": "hello"})); + if let Some(generate) = generate { + body.insert("generate".to_owned(), json!(generate)); + } + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("openai", body).unwrap(), + )); + let mut stream = provider + .execute( + provider_request_with_operation("xai", operation), + context(CancellationToken::new(), None), + ) + .await + .expect("ordinary generation remains supported"); + let mut completed = false; + while let Some(event) = stream.next().await { + completed |= event + .expect("normal upstream response") + .into_parts() + .0 + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))); + } + assert!(completed); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + let requests = transport.requests.lock().unwrap(); + let sent: Value = serde_json::from_slice(requests[0].body()).unwrap(); + assert_eq!(sent.get("generate").and_then(Value::as_bool), generate); + } +} #[tokio::test] async fn native_xai_translates_a_non_native_source_before_encoding() { @@ -81,6 +192,7 @@ async fn native_xai_translates_a_non_native_source_before_encoding() { "example-source", Map::from_iter([ ("opaque".to_owned(), json!("source-only")), + ("generate".to_owned(), json!(false)), ("input".to_owned(), json!([{"type":"compaction_trigger"}])), ("prompt_cache_key".to_owned(), json!("must-not-be-read")), ]), @@ -539,7 +651,7 @@ async fn attempt_middleware_can_change_reasoning_before_native_validation() { let requests = transport.requests.lock().unwrap(); let body: Value = serde_json::from_slice(requests[0].body()).unwrap(); - assert_eq!(body["reasoning"]["effort"], "high"); + assert_eq!(body["reasoning"]["effort"], "xhigh"); assert_eq!(body["model"], MODEL); assert_eq!( observed.lock().unwrap()[0].1["reasoning"]["effort"], @@ -583,36 +695,45 @@ fn observed_transport_metrics() -> GrokInferenceTransportMetrics { const CATALOG_WITHOUT_FEATURE_METADATA: &[u8] = br#"{ "object": "list", - "data": [{ - "id": "grok-catalog-entry", - "model": "grok-catalog-entry", - "contextWindow": 1000000, - "maxCompletionTokens": 131072, - "apiBackend": "responses", - "supportedInApi": true, - "supportsReasoningEffort": true - }] + "data": [ + { + "id": "grok-catalog-entry", + "model": "grok-catalog-entry", + "context_window": 1000000, + "max_completion_tokens": 131072, + "api_backend": "responses", + "supported_in_api": true, + "supports_reasoning_effort": true + } + ] }"#; -const CATALOG_WITH_FEATURE_REASONING_OPTIONS: &[u8] = br#"{ +const CATALOG_WITH_REASONING_MENU: &[u8] = br#"{ "object": "list", - "data": [{ - "id": "grok-4.5", - "model": "grok-4.5", - "contextWindow": 500000, - "maxCompletionTokens": 131072, - "apiBackend": "responses", - "supportedInApi": true, - "reasoningEfforts": ["high", {"value": "medium", "default": true}, "low"], - "features": { - "reasoning": true, - "reasoningEffortOptions": { - "supportedEfforts": ["low", "medium", "high", "xhigh"], - "defaultEffort": "high" - } - }, - "supportsReasoningEffort": true, - "streamToolCalls": true - }] + "data": [ + { + "id": "grok-4.5", + "model": "grok-4.5", + "context_window": 500000, + "max_completion_tokens": 131072, + "api_backend": "responses", + "reasoning_efforts": [ + { + "value": "low" + }, + { + "value": "medium" + }, + { + "value": "high", + "default": true + }, + { + "value": "xhigh" + } + ], + "stream_tool_calls": true + } + ] }"#; const SUCCESS_SSE: &[u8] = concat!( "event: response.created\n", @@ -1076,13 +1197,13 @@ impl GrokModelCatalogTransport for CatalogWithoutFeatureMetadataTransport { } } -struct CatalogWithFeatureReasoningOptionsTransport; +struct CatalogWithReasoningMenuTransport; -impl GrokModelCatalogTransport for CatalogWithFeatureReasoningOptionsTransport { +impl GrokModelCatalogTransport for CatalogWithReasoningMenuTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { Box::pin(async { Ok(GrokModelCatalogTransportResponse::new( - CATALOG_WITH_FEATURE_REASONING_OPTIONS, + CATALOG_WITH_REASONING_MENU, None, )) }) @@ -1091,6 +1212,22 @@ impl GrokModelCatalogTransport for CatalogWithFeatureReasoningOptionsTransport { struct UnavailableCatalogTransport; +struct CatalogWithServerReasoningDefaultsTransport; + +impl GrokModelCatalogTransport for CatalogWithServerReasoningDefaultsTransport { + fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { + Box::pin(async { + Ok(GrokModelCatalogTransportResponse::new( + br#"{"object":"list","data":[ + {"model":"grok-4.7","api_backend":"responses","context_window":256000,"context_windows":[256000,500000],"reasoning_efforts":[{"value":"low"},{"value":"high"}]}, + {"model":"grok-future","api_backend":"responses","reasoning_efforts":[{"value":"low"},{"value":"high"},{"value":"unknown","default":true}]} + ]}"#, + None, + )) + }) + } +} + impl GrokModelCatalogTransport for UnavailableCatalogTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { Box::pin(async { @@ -1846,7 +1983,7 @@ fn diagnostic_context(request_id: &str, account_suffix: &str) -> AttemptContext } #[tokio::test] -async fn model_alias_should_be_canonical_before_account_selection() { +async fn account_selection_uses_the_exact_routed_model() { let selector = StubSelector::success(); let provider = provider(selector.clone(), StubInferenceTransport::success()).await; let mut stream = provider @@ -1866,7 +2003,7 @@ async fn model_alias_should_be_canonical_before_account_selection() { .lock() .expect("upstream models") .as_slice(), - ["grok-4.6"] + ["grok-4.6-latest"] ); } @@ -4064,7 +4201,7 @@ async fn provider_compiles_realtime_catalog_capabilities() { } #[tokio::test] -async fn provider_publishes_default_codex_profile_when_catalog_is_unavailable() { +async fn provider_propagates_catalog_failure_without_inventing_model_facts() { let provider = provider_with_catalog_transport( StubSelector::success(), StubInferenceTransport::success(), @@ -4072,32 +4209,21 @@ async fn provider_publishes_default_codex_profile_when_catalog_is_unavailable() Arc::new(UnavailableCatalogTransport), ) .await; - let capabilities = provider + let error = provider .query_model_capabilities() .await - .expect("fallback capabilities"); - - assert_eq!(capabilities.len(), 1); - assert_eq!(capabilities[0].upstream_model().as_str(), MODEL); - let presentation = capabilities[0] - .presentation() - .expect("fallback Grok model presentation"); - assert_eq!(presentation.display_name(), Some("Grok 4.5")); - assert_eq!(presentation.default_reasoning_effort(), None); - assert!(presentation.supported_reasoning_efforts().is_empty()); - assert_eq!(presentation.context_window_tokens(), Some(500_000)); - assert_eq!(presentation.max_context_window_tokens(), Some(500_000)); - assert!(presentation.agent_tools()); - assert!(presentation.parallel_tool_calls()); + .expect_err("catalog is unavailable"); + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); } #[tokio::test] -async fn provider_publishes_reasoning_efforts_from_feature_options() { +async fn provider_publishes_reasoning_efforts_from_current_menu() { let provider = provider_with_catalog_transport( StubSelector::success(), StubInferenceTransport::success(), StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable), - Arc::new(CatalogWithFeatureReasoningOptionsTransport), + Arc::new(CatalogWithReasoningMenuTransport), ) .await; let capabilities = provider @@ -4115,6 +4241,41 @@ async fn provider_publishes_reasoning_efforts_from_feature_options() { ); } +#[tokio::test] +async fn provider_keeps_server_reasoning_default_when_catalog_does_not_declare_one() { + let provider = provider_with_catalog_transport( + StubSelector::success(), + StubInferenceTransport::success(), + StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable), + Arc::new(CatalogWithServerReasoningDefaultsTransport), + ) + .await; + let capabilities = provider + .query_model_capabilities() + .await + .expect("capabilities"); + assert_eq!(capabilities.len(), 2); + for model in capabilities { + let presentation = model.presentation().expect("model presentation"); + assert_eq!( + presentation.supported_reasoning_efforts(), + ["low", "high"], + "{}", + model.upstream_model() + ); + assert_eq!( + presentation.default_reasoning_effort(), + None, + "{}", + model.upstream_model() + ); + if model.upstream_model().as_str() == "grok-4.7" { + assert_eq!(presentation.context_window_tokens(), Some(256_000)); + assert_eq!(presentation.max_context_window_tokens(), Some(500_000)); + } + } +} + #[tokio::test] async fn missing_catalog_tool_metadata_keeps_build_tools_routable() { let provider = provider_with_catalog_transport( @@ -4301,6 +4462,20 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe { #[tokio::test] async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_and_stays_cold() { + assert_upstream_adapter_stays_cold(operation()).await; + for protocol in ["openai", "example-source"] { + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + protocol, + object(json!({"model": MODEL, "input": "hello", "generate": false})), + ) + .unwrap(), + )); + assert_upstream_adapter_stays_cold(operation).await; + } +} + +async fn assert_upstream_adapter_stays_cold(operation: Operation) { let selector = StubSelector::success(); let transport = StubInferenceTransport::success(); let provider = provider(selector.clone(), transport.clone()).await; @@ -4344,7 +4519,7 @@ async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_an CancellationToken::new(), ); let mut stream = provider - .execute(provider_request("xai"), context) + .execute(provider_request_with_operation("xai", operation), context) .await .unwrap(); assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 0); diff --git a/backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs b/backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs new file mode 100644 index 000000000..574b9bfd1 --- /dev/null +++ b/backend/crates/providers/xai/tests/provider/contract/replay_tool_names.rs @@ -0,0 +1,185 @@ +//! 跨轮重放区分已编码输入与客户端可见输出，避免工具别名缺失或二次映射 + +use provider_xai::GrokResponsesRequest; + +use super::*; + +fn generate(body: Value) -> GenerateRequest { + GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("openai", body.as_object().expect("body").clone()) + .expect("payload"), + ) +} + +fn function(name: &str, label: &str) -> Value { + json!({"type": "function", "name": name, "description": label, + "parameters": {"type": "object"}}) +} + +fn history(names: [&str; 2], prefix: &str) -> Vec { + names.into_iter().enumerate().flat_map(|(index, name)| { + let call_id = format!("{prefix}_{index}"); + [ + json!({"type": "function_call", "name": name, "call_id": call_id, "arguments": "{}"}), + json!({"type": "function_call_output", "call_id": call_id, "output": "done"}), + ] + }).collect() +} + +fn declared_alias(body: &Value, label: &str) -> String { + body["tools"] + .as_array() + .expect("tools") + .iter() + .find(|tool| tool["description"] == label) + .expect("function")["name"] + .as_str() + .expect("alias") + .to_owned() +} + +async fn assert_replay_tool_names(attempt: Option) { + for hosted in ["web_search", "x_search"] { + for literal_first in [false, true] { + let key = ClientApiKeyId::new("key_xai_contract").expect("key"); + let isolated = GrokResponsesRequest::encode( + &generate(json!({"tools": [function(hosted, "reserved")]})), + MODEL, + &key, + ) + .expect("isolated function"); + let literal = isolated.body()["tools"][0]["name"] + .as_str() + .expect("generated alias"); + let mut tools = vec![function(hosted, "reserved"), function(literal, "literal")]; + if literal_first { + tools.reverse(); + } + if hosted == "web_search" { + tools.push(json!({"type": "web_search"})); + } + let first_body = json!({"model": "client-model", "store": false, + "prompt_cache_key": "replay_tool_names", "tools": tools, + "input": history([hosted, literal], "prior")}); + let encoded = GrokResponsesRequest::encode(&generate(first_body.clone()), MODEL, &key) + .expect("first request"); + let encoded_body = Value::Object(encoded.body().clone()); + let aliases = [ + declared_alias(&encoded_body, "reserved"), + declared_alias(&encoded_body, "literal"), + ]; + assert_ne!(aliases[0], aliases[1]); + let output: Vec<_> = aliases + .iter() + .enumerate() + .map(|(index, name)| { + json!({"type": "function_call", "id": format!("fc_returned_{index}"), + "name": name, "call_id": format!("returned_{index}"), "arguments": "{}"}) + }) + .collect(); + let response = [ + json!({"type": "response.created", "response": {"id": "resp_tool_names"}}), + json!({"type": "response.output_item.done", "output_index": 0, "item": output[0]}), + json!({"type": "response.output_item.done", "output_index": 1, "item": output[1]}), + json!({"type": "response.completed", "response": {"id": "resp_tool_names", + "status": "completed", "output": output}}), + ] + .into_iter() + .map(|event| format!("data: {event}\n\n")) + .collect::(); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(response.into_bytes()), + InferenceMode::Success, + ]); + let provider = provider(StubSelector::success(), transport.clone()).await; + let mut first = Arc::clone(&provider) + .execute( + provider_request_with_operation( + "xai", + Operation::Generate(generate(first_body)), + ), + context(CancellationToken::new(), None), + ) + .await + .expect("first stream"); + let state = collect_provider_state(&mut first) + .await + .expect("session state"); + + let mut second_input = history([hosted, literal], "current"); + for index in 0..2 { + second_input.push(json!({"type": "function_call_output", + "call_id": format!("returned_{index}"), "output": "returned result"})); + } + let mut second_body = json!({"model": "client-model", "store": false, + "prompt_cache_key": "replay_tool_names", "tools": tools, "input": second_input}); + let (operation, context) = if let Some(attempt) = attempt { + second_body["previous_response_id"] = json!("resp_gateway_tool_names"); + let pin = NativeContinuationPin::new( + PreviousResponseId::new("resp_gateway_tool_names"), + PreviousResponseId::new("resp_tool_names"), + key, + ProviderKind::new("xai").expect("provider"), + account_id("provider"), + ); + ( + operation_with_state(second_body, state), + context_with_continuation_attempt(ContinuationBinding::Pinned(pin), attempt), + ) + } else { + ( + Operation::Generate(generate(second_body)), + context(CancellationToken::new(), None), + ) + }; + let mut second = provider + .execute(provider_request_with_operation("xai", operation), context) + .await + .expect("second stream"); + while let Some(event) = second.next().await { + event.expect("second event"); + } + let requests = transport.requests.lock().expect("requests"); + let replay: Value = serde_json::from_slice(requests[1].body()).expect("replay body"); + assert!(replay.get("previous_response_id").is_none()); + for (index, expected) in aliases.iter().enumerate() { + assert_eq!( + declared_alias(&replay, if index == 0 { "reserved" } else { "literal" }), + *expected + ); + for prefix in ["prior", "returned", "current"] { + if attempt.is_none() && prefix == "prior" { + continue; + } + let call_id = format!("{prefix}_{index}"); + let call = replay["input"] + .as_array() + .expect("input") + .iter() + .find(|item| item["type"] == "function_call" && item["call_id"] == call_id) + .expect("replayed call"); + assert_eq!( + call["name"], *expected, + "{hosted}, literal_first={literal_first}, {call_id}" + ); + assert_eq!(call["arguments"], "{}"); + } + } + } + } +} + +#[tokio::test] +async fn replay_tool_names_reencode_outputs_on_store_false_native_continuation() { + assert_replay_tool_names(Some(ContinuationAttempt::Native)).await; +} + +#[tokio::test] +async fn replay_tool_names_reencode_outputs_on_replay_owner() { + assert_replay_tool_names(Some(ContinuationAttempt::ReplayOwner)).await; +} + +#[tokio::test] +async fn replay_tool_names_reencode_cached_reasoning_calls_without_reencoding_input() { + assert_replay_tool_names(None).await; +} diff --git a/backend/crates/providers/xai/tests/transport/canonical.rs b/backend/crates/providers/xai/tests/transport/canonical.rs index 5726beef1..51dac36ca 100644 --- a/backend/crates/providers/xai/tests/transport/canonical.rs +++ b/backend/crates/providers/xai/tests/transport/canonical.rs @@ -19,8 +19,9 @@ fn decoder_should_bill_the_sent_model_independently_of_the_response_model() { ("grok-4.5", Some("grok-4.3"), None, Some(2_175_000)), ("grok-4.5", None, None, Some(2_175_000)), ("grok-future", Some("grok-4.5"), None, None), + ("grok-latest", Some("grok-4.3"), Some(0), None), ("grok-4.5", Some("grok-future"), Some(123), None), - ("grok-4.5", Some("grok-future"), Some(0), None), + ("grok-4.5", Some("grok-future"), Some(0), Some(2_175_000)), ] { let created = serde_json::json!({"type":"response.created","response":{"id":"resp_model_cost","model":returned}}); let mut completed = serde_json::json!({ @@ -45,7 +46,10 @@ fn decoder_should_bill_the_sent_model_independently_of_the_response_model() { expected_cost, "{sent} -> {returned:?}" ); - assert_eq!(provider_cost_ticks(&facts), reported_cost); + assert_eq!( + provider_cost_ticks(&facts), + reported_cost.filter(|ticks| *ticks > 0) + ); assert_eq!(decoder.response_model(), returned); assert!(facts.iter().any(|event| matches!( event, @@ -72,6 +76,24 @@ fn terminal_cost_events( cached_tokens: u64, output_tokens: u64, provider_cost_ticks: Option, +) -> Vec { + terminal_cost_events_with_pricing( + model, + input_tokens, + cached_tokens, + output_tokens, + provider_cost_ticks, + None, + ) +} + +fn terminal_cost_events_with_pricing( + model: &str, + input_tokens: u64, + cached_tokens: u64, + output_tokens: u64, + provider_cost_ticks: Option, + pricing: Option, ) -> Vec { let provider_cost = provider_cost_ticks .map(|ticks| format!(",\"cost_in_usd_ticks\":{ticks}")) @@ -89,7 +111,13 @@ fn terminal_cost_events( cached_tokens = cached_tokens, provider_cost = provider_cost, ); - decode_canonical(body.as_bytes(), model).expect("canonical cost response") + GrokCanonicalDecoder::new(model) + .with_pricing(pricing) + .push(body.as_bytes()) + .expect("canonical cost response") + .into_iter() + .flat_map(|event| event.into_parts().0) + .collect() } fn decode_canonical(body: &[u8], upstream_model: &str) -> Result, ProviderError> { @@ -271,11 +299,16 @@ fn decoder_should_price_official_grok_45_build_free_variant() { } #[test] -fn zero_provider_cost_should_remain_an_authoritative_zero() { - let events = terminal_cost_events("grok-4.5", 1, 0, 1, Some(0)); +fn zero_provider_cost_should_use_available_token_pricing() { + for (model, calculated) in [("grok-4.5", Some(80_000)), ("grok-future", None)] { + let events = terminal_cost_events(model, 1, 0, 1, Some(0)); - assert_eq!(calculated_cost_ticks(&events), None); - assert_eq!(provider_cost_ticks(&events), Some(0)); + assert_eq!( + (calculated_cost_ticks(&events), provider_cost_ticks(&events)), + (calculated, None), + "{model}" + ); + } } #[test] @@ -401,7 +434,7 @@ fn billing_should_use_the_final_response_tier_and_preserve_provider_cost() { ("priority", None, Some(400_000_000), None), ("default", None, Some(200_000_000), None), ("future", None, None, None), - ("priority", Some(0), None, Some(0)), + ("priority", Some(0), Some(400_000_000), None), ("priority", Some(123), None, Some(123)), ] { let mut response = serde_json::json!({"service_tier": tier, "output": [], "usage": {"input_tokens": 10_000, "output_tokens": 0}}); @@ -427,12 +460,18 @@ fn billing_should_require_provider_totals_for_server_tools() { "image_generation_call", "future_tool_call", ] { - let response = serde_json::json!({"output": [{"type": kind, "id": "tool_1", "status": "completed"}], "usage": {"input_tokens": 100, "output_tokens": 1}}); - assert_eq!( - calculated_cost_ticks(&pricing_events(response, None)), - None, - "{kind}" - ); + for ticks in [None, Some(0)] { + let mut response = serde_json::json!({"output": [{"type": kind, "id": "tool_1", "status": "completed"}], "usage": {"input_tokens": 100, "output_tokens": 1}}); + if let Some(ticks) = ticks { + response["usage"]["cost_in_usd_ticks"] = serde_json::json!(ticks); + } + let events = pricing_events(response, None); + assert_eq!( + (calculated_cost_ticks(&events), provider_cost_ticks(&events)), + (None, None), + "{kind}: {ticks:?}" + ); + } } let request = tool_request(serde_json::json!({"input": "查资料", "tools": [{"type": "web_search"}]})); @@ -444,9 +483,10 @@ fn billing_should_require_provider_totals_for_server_tools() { None ); response["usage"]["cost_in_usd_ticks"] = serde_json::json!(0); + let events = pricing_events(response, Some(&request)); assert_eq!( - provider_cost_ticks(&pricing_events(response, Some(&request))), - Some(0) + (calculated_cost_ticks(&events), provider_cost_ticks(&events)), + (None, None) ); let response = serde_json::json!({"output": [], "usage": {"input_tokens": 100, "output_tokens": 1, "num_server_side_tool_calls": 1}}); assert_eq!(calculated_cost_ticks(&pricing_events(response, None)), None); @@ -495,6 +535,78 @@ fn decoder_should_leave_unpublished_model_pricing_unavailable() { assert_eq!(provider_cost_ticks(&events), None); } +#[test] +fn decoder_should_not_price_dynamic_or_compatibility_aliases_as_stable_models() { + for model in [ + "grok", + "grok-latest", + "grok-build-latest", + "grok-4.6-latest", + "grok-4.5-latest", + "grok-4.3-latest", + "grok-code-fast", + "grok-4.20-multi-agent-latest", + "grok-4.20-beta", + "grok-4.20-beta-reasoning", + "grok-4.20-beta-non-reasoning", + "grok-4.20-reasoning-gv2", + "grok-4.20-non-reasoning-gv2", + "xai/grok-4.5", + "grok-4.7", + ] { + for reported in [None, Some(0), Some(123)] { + let events = terminal_cost_events(model, 100, 25, 10, reported); + assert_eq!( + calculated_cost_ticks(&events), + None, + "{model}, reported: {reported:?}" + ); + assert_eq!( + provider_cost_ticks(&events), + reported.filter(|ticks| *ticks > 0).map(u128::from) + ); + assert!( + events + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))) + ); + } + } +} + +#[test] +fn explicit_alias_pricing_still_applies_when_provider_cost_is_unreported() { + let pricing: gateway_core::metering::ModelPriceOverride = + serde_json::from_value(serde_json::json!({ + "multiplierBps":15000,"bands":{"standard":{ + "input":"3", "output":"12", "cacheRead":"0", "cacheWrite":"0" + }} + })) + .expect("manual price"); + for reported in [None, Some(0), Some(123)] { + let events = terminal_cost_events_with_pricing( + "grok-latest", + 100, + 20, + 10, + reported, + Some(pricing.clone()), + ); + assert_eq!( + calculated_cost_ticks(&events), + if reported == Some(123) { + None + } else { + Some(5_400_000) + } + ); + assert_eq!( + provider_cost_ticks(&events), + reported.filter(|ticks| *ticks > 0).map(u128::from) + ); + } +} + #[test] fn decoder_should_leave_incomplete_usage_pricing_unavailable() { let body = concat!( @@ -1410,20 +1522,38 @@ fn decoder_should_require_terminal_response() { } #[test] -fn decoder_should_preserve_incomplete_length_reason() { - let body = concat!( - "event: response.created\n", - "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_short\"}}\n\n", - "event: response.incomplete\n", - "data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_short\",\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n", - ); - let events = decode_canonical(body.as_bytes(), "fallback").expect("incomplete response"); +fn decoder_should_classify_incomplete_reasons_and_preserve_native_wire() { + for (reason, expected) in [ + ("max_output_tokens", FinishReason::Length), + ("max_tokens", FinishReason::Length), + ("max_prompt_tokens", FinishReason::Length), + ("max_time_limit", FinishReason::Length), + ("content_filter", FinishReason::ContentFilter), + ("future_reason", FinishReason::Other), + ] { + let created = serde_json::json!({"type":"response.created","response":{"id":"resp_short"}}); + let incomplete = serde_json::json!({ + "type":"response.incomplete", + "response":{ + "id":"resp_short","status":"incomplete", + "incomplete_details":{"reason":reason} + } + }); + let body = format!("data: {created}\n\ndata: {incomplete}\n\n"); + let events = GrokCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect("incomplete response"); + let completed = events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .find_map(|fact| match fact { + GatewayEvent::Completed(meta) => meta.finish_reason(), + _ => None, + }); - assert!(matches!( - events.last(), - Some(GatewayEvent::Completed(meta)) - if meta.finish_reason() == Some(FinishReason::Length) - )); + assert_eq!(completed, Some(expected), "{reason}"); + assert_eq!(wire_events(&events).last().unwrap().data(), &incomplete); + } } #[test] diff --git a/backend/crates/providers/xai/tests/transport/catalog.rs b/backend/crates/providers/xai/tests/transport/catalog.rs index 4940717c8..bdebe5bb4 100644 --- a/backend/crates/providers/xai/tests/transport/catalog.rs +++ b/backend/crates/providers/xai/tests/transport/catalog.rs @@ -13,8 +13,7 @@ use provider_xai::{ MAX_GROK_MODEL_CATALOG_BYTES, SecretValue, parse_grok_billing, parse_grok_model_catalog, }; -const OFFICIAL_FIXTURE: &[u8] = - include_bytes!("catalog/fixtures/official_grok_models_snapshot.json"); +const CLI_PROXY_FIXTURE: &[u8] = include_bytes!("catalog/fixtures/cli_proxy_models.json"); struct CapturingTransport { calls: AtomicUsize, @@ -80,7 +79,7 @@ impl GrokBillingTransport for CapturingBillingTransport { #[tokio::test] async fn client_should_send_exact_oauth_headers_without_api_key() { let transport = Arc::new(CapturingTransport::success( - OFFICIAL_FIXTURE, + CLI_PROXY_FIXTURE, Some("\"grok-v1\""), )); let client = GrokModelCatalogClient::new(transport.clone()); @@ -128,7 +127,7 @@ async fn client_should_send_exact_oauth_headers_without_api_key() { #[tokio::test] async fn client_should_omit_email_when_verified_profile_has_none() { - let transport = Arc::new(CapturingTransport::success(OFFICIAL_FIXTURE, None)); + let transport = Arc::new(CapturingTransport::success(CLI_PROXY_FIXTURE, None)); let client = GrokModelCatalogClient::new(transport.clone()); client .fetch(&session(None)) @@ -262,10 +261,13 @@ async fn subscription_query_rejects_malformed_or_oversized_user_responses() { } #[test] -fn billing_parser_should_accept_credits_and_legacy_shapes() { +fn billing_parser_should_accept_current_credits_fields() { for body in [ br#"{"config":{"creditUsagePercent":31.25,"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-13T00:00:00Z","end":"2026-07-20T00:00:00Z"},"prepaidBalance":{"val":2500}}}"#.as_slice(), - br#"{"config":{"monthlyLimit":{"val":2000},"used":{"val":500},"onDemandCap":{"val":1000},"onDemandUsed":{"val":100}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-500}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-9223372036854775808}}}"#.as_slice(), + br#"{"config":{"onDemandCap":{},"onDemandUsed":{},"prepaidBalance":{}}}"#.as_slice(), + br#"{"config":{"onDemandCap":{"val":1000},"onDemandUsed":{"val":100}}}"#.as_slice(), br#"{"config":null}"#.as_slice(), ] { parse_grok_billing(body).expect("supported official billing shape"); @@ -278,7 +280,13 @@ fn billing_parser_should_reject_invalid_known_fields() { br#"[]"#.as_slice(), br#"{"config":[]}"#.as_slice(), br#"{"config":{"creditUsagePercent":101}}"#.as_slice(), - br#"{"config":{"used":{"val":-1}}}"#.as_slice(), + br#"{"config":{"onDemandCap":{"val":-1}}}"#.as_slice(), + br#"{"config":{"onDemandUsed":{"val":-1}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":null}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":"-500"}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-0.5}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":9223372036854775808}}}"#.as_slice(), + br#"{"config":{"prepaidBalance":{"val":-9223372036854775809}}}"#.as_slice(), br#"{"config":{"currentPeriod":"weekly"}}"#.as_slice(), br#"{"onDemandEnabled":"yes"}"#.as_slice(), ] { @@ -309,9 +317,9 @@ fn billing_snapshot_debug_should_not_print_values() { } #[test] -fn official_fixture_should_use_actual_model_and_whitelisted_metadata() { - let snapshot = parse_grok_model_catalog(OFFICIAL_FIXTURE, Some("W/\"grok-v1\"")) - .expect("official fixture should parse"); +fn cli_proxy_fixture_should_use_actual_model_and_whitelisted_metadata() { + let snapshot = parse_grok_model_catalog(CLI_PROXY_FIXTURE, Some("W/\"grok-v1\"")) + .expect("CLI proxy fixture should parse"); let model = &snapshot.models()[0]; assert_eq!( @@ -370,7 +378,7 @@ fn official_fixture_should_use_actual_model_and_whitelisted_metadata() { #[test] fn non_whitelisted_wire_fields_should_not_survive_normalization() { let snapshot = - parse_grok_model_catalog(OFFICIAL_FIXTURE, None).expect("official fixture should parse"); + parse_grok_model_catalog(CLI_PROXY_FIXTURE, None).expect("CLI proxy fixture should parse"); let debug = format!("{snapshot:?}"); assert!( @@ -382,16 +390,18 @@ fn non_whitelisted_wire_fields_should_not_survive_normalization() { #[test] fn invalid_etag_should_fail_the_entire_snapshot() { - let result = parse_grok_model_catalog(OFFICIAL_FIXTURE, Some("raw-unquoted-etag")); + let result = parse_grok_model_catalog(CLI_PROXY_FIXTURE, Some("raw-unquoted-etag")); assert!(matches!(result, Err(GrokModelCatalogError::InvalidEtag))); } #[test] fn missing_capability_fields_should_remain_unknown() { - let snapshot = - parse_grok_model_catalog(br#"{"object":"list","data":[{"id":"grok-unknown"}]}"#, None) - .expect("identity-only official entry should parse"); + let snapshot = parse_grok_model_catalog( + br#"{"object":"list","data":[{"id":"grok-unknown","model":"grok-unknown"}]}"#, + None, + ) + .expect("identity-only official entry should parse"); let model = &snapshot.models()[0]; assert_eq!( @@ -417,7 +427,7 @@ fn missing_capability_fields_should_remain_unknown() { #[test] fn reasoning_effort_menu_should_be_capability_evidence_without_legacy_flag() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-reasoning","reasoningEfforts":[{"value":"xhigh","default":true},"low"]}]}"#, + br#"{"object":"list","data":[{"id":"grok-reasoning","reasoning_efforts":[{"value":"xhigh","default":true},{"value":"low"}],"model":"grok-reasoning"}]}"#, None, ) .expect("reasoning effort menu"); @@ -444,38 +454,30 @@ fn reasoning_effort_menu_should_be_capability_evidence_without_legacy_flag() { } #[test] -fn reasoning_effort_options_should_read_official_features_shape() { +fn historical_and_other_endpoint_fields_do_not_supply_model_facts() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-4.5","reasoningEfforts":["high"],"features":{"reasoning":true,"reasoningEffortOptions":{"supportedEfforts":["low","medium","high","xhigh"],"defaultEffort":"high"}}}]}"#, + br#"{"object":"list","data":[{"model":"grok-future","contextWindow":500000,"contextWindows":[1000000],"apiBackend":"responses","reasoningEfforts":[{"value":"high","default":true}],"features":{"reasoning":true,"reasoningEffortOptions":{"supportedEfforts":["high"],"defaultEffort":"high"}},"_meta":{"contextWindow":500000,"reasoningEfforts":[{"value":"high","default":true}]},"capabilities":{"reasoning_effort":["high"],"default_reasoning_effort":"high"},"reasoning_effort":"high"}]}"#, None, - ) - .expect("official feature reasoning options should parse"); - let capabilities = snapshot.models()[0].capabilities(); - + ).expect("unknown fields do not define CLI proxy facts"); + let model = &snapshot.models()[0]; assert_eq!( - ( - capabilities.reasoning_effort(), - capabilities - .reasoning_efforts() - .iter() - .map(|effort| effort.as_str()) - .collect::>(), - capabilities - .default_reasoning_effort() - .map(|effort| effort.as_str()), - ), - ( - GrokCatalogCapabilityEvidence::DeclaredNative, - vec!["low", "medium", "high", "xhigh"], - Some("high"), - ) + model.capabilities().reasoning_effort(), + GrokCatalogCapabilityEvidence::Unknown + ); + assert!(model.capabilities().reasoning_efforts().is_empty()); + assert_eq!(model.capabilities().default_reasoning_effort(), None); + assert_eq!(model.limits().context_window_tokens(), None); + assert_eq!(model.limits().max_context_window_tokens(), None); + assert_eq!( + model.capabilities().responses_api(), + GrokCatalogCapabilityEvidence::Unknown ); } #[test] fn unknown_reasoning_effort_values_should_not_fail_the_snapshot() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-reasoning","reasoningEffort":"ultra","reasoningEfforts":[{"value":"ultra","default":true},"low","hyper"]}]}"#, + br#"{"object":"list","data":[{"id":"grok-reasoning","reasoning_effort":"ultra","reasoning_efforts":[{"value":"ultra","default":true},{"value":"low"},{"value":"hyper"}],"model":"grok-reasoning"}]}"#, None, ) .expect("unknown effort values must not break catalog parsing"); @@ -492,25 +494,143 @@ fn unknown_reasoning_effort_values_should_not_fail_the_snapshot() { .default_reasoning_effort() .map(|effort| effort.as_str()), ), - (vec!["low"], Some("low")) + (vec!["low"], None) ); } #[test] -fn model_id_should_take_priority_over_catalog_id() { +fn reasoning_menu_uses_only_known_explicit_default() { + for (menu, expected) in [ + ( + serde_json::json!([{"value":"low"},{"value":"high","default":true}]), + Some("high"), + ), + (serde_json::json!([{"value":"low"},{"value":"high"}]), None), + ( + serde_json::json!([{"value":"low"},{"value":"quantum","default":true}]), + None, + ), + ] { + let body = serde_json::to_vec(&serde_json::json!({ + "object":"list", "data":[{"model":"grok-4.7","reasoning_efforts":menu}] + })) + .expect("fixture"); + let snapshot = parse_grok_model_catalog(&body, None).expect("current menu"); + assert_eq!( + snapshot.models()[0] + .capabilities() + .default_reasoning_effort() + .map(|effort| effort.as_str()), + expected + ); + } +} + +#[test] +fn current_reasoning_menu_does_not_borrow_defaults_from_other_sources() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"catalog-entry","modelId":"grok-4-fast"}]}"#, + br#"{"object":"list","data":[{"model":"grok-future","reasoning_efforts":[{"value":"low"}],"reasoning_effort":"low","capabilities":{"reasoning_effort":["high"],"default_reasoning_effort":"high"}}]}"#, + None, + ).expect("current menu"); + let capabilities = snapshot.models()[0].capabilities(); + assert_eq!( + capabilities + .reasoning_efforts() + .iter() + .map(|effort| effort.as_str()) + .collect::>(), + ["low"] + ); + assert_eq!(capabilities.default_reasoning_effort(), None); +} + +#[test] +fn unknown_reasoning_menu_does_not_fall_back_to_other_sources() { + let snapshot = parse_grok_model_catalog( + br#"{"object":"list","data":[{"model":"grok-future","reasoning_efforts":[{"value":"quantum","default":true}],"capabilities":{"reasoning_effort":["high"],"default_reasoning_effort":"high"}}]}"#, + None, + ).expect("unknown effort"); + let capabilities = snapshot.models()[0].capabilities(); + assert!(capabilities.reasoning_efforts().is_empty()); + assert_eq!(capabilities.default_reasoning_effort(), None); +} + +#[test] +fn context_window_choices_supply_default_when_scalar_is_missing() { + let snapshot = parse_grok_model_catalog( + br#"{"object":"list","data":[{"model":"grok-4.7","context_windows":[500000,256000]}]}"#, None, ) - .expect("modelId is official fallback"); + .expect("window choices"); + assert_eq!( + snapshot.models()[0] + .limits() + .context_window_tokens() + .map(std::num::NonZeroU64::get), + Some(500000) + ); +} - assert_eq!(snapshot.models()[0].request_model().as_str(), "grok-4-fast"); +#[test] +fn context_window_choices_keep_scalar_default_and_publish_largest_supported_window() { + for (scalar, windows, maximum) in [ + (256000, vec![500000, 256000], 500000), + (1000000, vec![256000, 500000], 1000000), + ] { + let body = serde_json::to_vec(&serde_json::json!({"object":"list","data":[{"model":"grok-4.7","context_window":scalar,"context_windows":windows}]})).expect("fixture"); + let snapshot = parse_grok_model_catalog(&body, None).expect("window choices"); + let limits = snapshot.models()[0].limits(); + assert_eq!( + limits + .context_window_tokens() + .map(std::num::NonZeroU64::get), + Some(scalar) + ); + assert_eq!( + limits + .max_context_window_tokens() + .map(std::num::NonZeroU64::get), + Some(maximum) + ); + } +} + +#[test] +fn malformed_context_window_choices_are_rejected() { + for windows in [ + serde_json::json!([256000, "big"]), + serde_json::json!([0]), + serde_json::json!(500000), + serde_json::json!([-1, 500000]), + serde_json::Value::Null, + ] { + let body = serde_json::to_vec(&serde_json::json!({"object":"list","data":[{"model":"grok-future","context_window":300000,"context_windows":windows}]})).expect("fixture"); + assert!(matches!( + parse_grok_model_catalog(&body, None), + Err(GrokModelCatalogError::InvalidWire | GrokModelCatalogError::InvalidLimits) + )); + } +} + +#[test] +fn current_proxy_requires_the_explicit_request_model() { + for entry in [ + serde_json::json!({"id":"grok-4.7"}), + serde_json::json!({"modelId":"grok-4.7"}), + ] { + let body = serde_json::to_vec(&serde_json::json!({"object":"list","data":[entry]})) + .expect("fixture"); + assert!(matches!( + parse_grok_model_catalog(&body, None), + Err(GrokModelCatalogError::InvalidWire) + )); + } } #[test] fn official_responses_backend_should_be_native_without_redundant_supported_flag() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-responses","api_backend":"responses"}]}"#, + br#"{"object":"list","data":[{"id":"grok-responses","api_backend":"responses","model":"grok-responses"}]}"#, None, ) .expect("Responses backend is explicit capability evidence"); @@ -524,8 +644,8 @@ fn official_responses_backend_should_be_native_without_redundant_supported_flag( #[test] fn explicit_api_disable_and_non_responses_backend_should_remain_unsupported() { for body in [ - br#"{"object":"list","data":[{"id":"grok-disabled","api_backend":"responses","supported_in_api":false}]}"#.as_slice(), - br#"{"object":"list","data":[{"id":"grok-chat","api_backend":"chat_completions","supported_in_api":true}]}"#.as_slice(), + br#"{"object":"list","data":[{"id":"grok-disabled","api_backend":"responses","supported_in_api":false,"model":"grok-disabled"}]}"#.as_slice(), + br#"{"object":"list","data":[{"id":"grok-chat","api_backend":"chat_completions","supported_in_api":true,"model":"grok-chat"}]}"#.as_slice(), ] { let snapshot = parse_grok_model_catalog(body, None).expect("valid unsupported entry"); assert_eq!( @@ -538,8 +658,8 @@ fn explicit_api_disable_and_non_responses_backend_should_remain_unsupported() { #[test] fn list_discriminator_should_be_required_and_exact() { for body in [ - br#"{"data":[{"id":"grok-4"}]}"#.as_slice(), - br#"{"object":"collection","data":[{"id":"grok-4"}]}"#.as_slice(), + br#"{"data":[{"id":"grok-4","model":"grok-4"}]}"#.as_slice(), + br#"{"object":"collection","data":[{"id":"grok-4","model":"grok-4"}]}"#.as_slice(), ] { assert!(matches!( parse_grok_model_catalog(body, None), @@ -565,7 +685,7 @@ fn empty_data_should_fail_the_entire_snapshot() { #[test] fn duplicate_actual_models_should_fail_the_entire_snapshot() { let result = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"entry-a","model":"grok-4"},{"id":"entry-b","modelId":"grok-4"}]}"#, + br#"{"object":"list","data":[{"id":"entry-a","model":"grok-4"},{"id":"entry-b","model":"grok-4"}]}"#, None, ); @@ -578,7 +698,7 @@ fn duplicate_actual_models_should_fail_the_entire_snapshot() { #[test] fn unknown_top_level_pagination_fields_should_not_reject_the_snapshot() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","data":[{"id":"grok-4"}],"has_more":true,"cursor":"next"}"#, + br#"{"object":"list","data":[{"id":"grok-4","model":"grok-4"}],"has_more":true,"cursor":"next"}"#, None, ) .expect("unknown pagination fields are not part of the model contract"); @@ -589,7 +709,7 @@ fn unknown_top_level_pagination_fields_should_not_reject_the_snapshot() { #[test] fn unknown_api_backend_should_degrade_to_unknown_without_rejecting_the_snapshot() { let snapshot = parse_grok_model_catalog( - br#"{"object":"list","future_top_level":true,"data":[{"id":"grok-future","apiBackend":"future_backend","future_field":{"keep":true}}]}"#, + br#"{"object":"list","future_top_level":true,"data":[{"id":"grok-future","future_field":{"keep":true},"api_backend":"future_backend","model":"grok-future"}]}"#, None, ) .expect("future backend should not reject the snapshot"); @@ -605,7 +725,7 @@ fn unknown_api_backend_should_degrade_to_unknown_without_rejecting_the_snapshot( #[test] fn invalid_preferred_model_should_fail_without_falling_back() { let result = parse_grok_model_catalog( - br#"{"object":"list","data":[{"model":"https://evil.invalid/model","modelId":"grok-4","id":"entry"}]}"#, + br#"{"object":"list","data":[{"model":"https://evil.invalid/model","id":"entry"}]}"#, None, ); @@ -686,7 +806,7 @@ fn header_value<'a>(headers: &'a [(String, &str)], name: &str) -> Option<&'a str #[tokio::test] async fn captured_sensitive_headers_should_remain_typed_as_sensitive() { let session = session(Some("person@example.com")); - let transport = Arc::new(CapturingTransport::success(OFFICIAL_FIXTURE, None)); + let transport = Arc::new(CapturingTransport::success(CLI_PROXY_FIXTURE, None)); let client = GrokModelCatalogClient::new(transport.clone()); client.fetch(&session).await.expect("fetch fixture"); let request = transport.request.lock().expect("captured request"); diff --git a/backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json b/backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json new file mode 100644 index 000000000..92955a506 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/catalog/fixtures/cli_proxy_models.json @@ -0,0 +1,42 @@ +{ + "object": "list", + "data": [ + { + "id": "grok-4.5-catalog-entry", + "model": "grok-4.5", + "name": "Grok 4.5", + "description": "Official Grok Build coding model.", + "hidden": false, + "apiKey": "provider-only field must not enter catalog metadata", + "extraHeaders": { + "x-api-key": "provider-only field" + }, + "baseUrl": "https://provider-only.invalid/v1", + "context_window": 1000000, + "max_completion_tokens": 131072, + "api_backend": "responses", + "supported_in_api": true, + "supports_reasoning_effort": true, + "reasoning_efforts": [ + { + "value": "low" + }, + { + "value": "medium", + "label": "Balanced", + "default": true + }, + { + "value": "high" + }, + { + "id": "extended", + "value": "xhigh", + "label": "Extended" + } + ], + "supports_backend_search": true, + "stream_tool_calls": true + } + ] +} diff --git a/backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json b/backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json deleted file mode 100644 index fd8111699..000000000 --- a/backend/crates/providers/xai/tests/transport/catalog/fixtures/official_grok_models_snapshot.json +++ /dev/null @@ -1,32 +0,0 @@ -{ - "object": "list", - "data": [ - { - "id": "grok-4.5-catalog-entry", - "model": "grok-4.5", - "modelId": "grok-4.5-fallback", - "name": "Grok 4.5", - "description": "Official Grok Build coding model.", - "contextWindow": 1000000, - "maxCompletionTokens": 131072, - "apiBackend": "responses", - "supportedInApi": true, - "supportsReasoningEffort": true, - "reasoningEffort": "medium", - "reasoningEfforts": [ - "low", - { "value": "medium", "label": "Balanced", "default": true }, - "high", - { "id": "extended", "value": "xhigh", "label": "Extended" } - ], - "supportsBackendSearch": true, - "streamToolCalls": true, - "hidden": false, - "apiKey": "provider-only field must not enter catalog metadata", - "extraHeaders": { - "x-api-key": "provider-only field" - }, - "baseUrl": "https://provider-only.invalid/v1" - } - ] -} diff --git a/backend/crates/providers/xai/tests/transport/request/hosted_names.rs b/backend/crates/providers/xai/tests/transport/request/hosted_names.rs new file mode 100644 index 000000000..415bdda4a --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/hosted_names.rs @@ -0,0 +1,243 @@ +//! Hosted tool 与客户端函数同名时，验证声明、历史和响应使用同一别名合同 + +use gateway_core::event::ProviderEvent; +use provider_xai::GrokCanonicalDecoder; + +use super::*; + +fn function_tool(name: &str) -> Value { + json!({"type": "function", "name": name, "parameters": { + "type": "object", "properties": {"query": {"type": "string"}} + }}) +} + +fn function_history(name: &str) -> Value { + json!([ + {"type": "function_call", "id": "fc_previous", "call_id": "call_previous", + "name": name, "arguments": "{\"query\":\"previous\"}"}, + {"type": "function_call_output", "call_id": "call_previous", "output": "result"} + ]) +} + +fn assert_function_response_restored(request: &GrokResponsesRequest, alias: &str, name: &str) { + let item = json!({"type": "function_call", "id": "fc_next", "call_id": "call_next", + "name": alias, "arguments": "{\"query\":\"next\"}"}); + let frames = [ + json!({"type": "response.created", "response": {"id": "resp_names"}}), + json!({"type": "response.output_item.added", "output_index": 0, "item": item}), + json!({"type": "response.function_call_arguments.delta", "output_index": 0, + "item_id": "fc_next", "call_id": "call_next", "delta": "{\"query\":\"next\"}"}), + json!({"type": "response.output_item.done", "output_index": 0, "item": item}), + json!({"type": "response.completed", "response": {"id": "resp_names", "status": "completed", + "tools": request.body()["tools"], "output": [item]}}), + ] + .into_iter() + .map(|frame| format!("data: {frame}\n\n")) + .collect::(); + let events = GrokCanonicalDecoder::for_request("grok-4.7", request) + .push(frames.as_bytes()) + .expect("function response"); + let wire: Vec<_> = events + .iter() + .filter_map(ProviderEvent::wire_event) + .collect(); + let calls: Vec<_> = wire + .iter() + .filter_map(|event| event.data().get("item")) + .collect(); + assert_eq!( + calls.len(), + 2, + "client functions must not be filtered as cache calls" + ); + for call in calls { + assert_eq!(call["name"], name); + assert_eq!(call["type"], "function_call"); + assert_eq!(call["call_id"], "call_next"); + assert_eq!(call["arguments"], "{\"query\":\"next\"}"); + } + let completed = wire.last().expect("completed response").data(); + assert_eq!(completed["response"]["output"][0]["name"], name); + assert!( + completed["response"]["tools"] + .as_array() + .expect("visible tools") + .iter() + .any(|tool| tool["type"] == "function" && tool["name"] == name) + ); +} + +#[test] +fn hosted_names_preserve_explicit_function_declarations_choices_history_and_output() { + for hosted in ["web_search", "x_search"] { + for hosted_first in [false, true] { + for nested_choice in [false, true] { + let mut tools = vec![function_tool(hosted), json!({"type": hosted})]; + if hosted_first { + tools.reverse(); + } + let choice = if nested_choice { + json!({"type": "function", "function": {"name": hosted}}) + } else { + json!({"type": "function", "name": hosted}) + }; + let request = raw_request(json!({ + "input": function_history(hosted), "tools": tools, "tool_choice": choice + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("hosted and client function"); + let body = encoded.body(); + let tools = body["tools"].as_array().expect("tools"); + let function = tools + .iter() + .find(|tool| tool["type"] == "function") + .expect("function retained"); + let alias = function["name"].as_str().expect("function alias"); + assert_ne!(alias, hosted, "hosted wire name must be unique"); + assert_eq!(tools.len(), 2); + assert!(tools.iter().any(|tool| tool["type"] == hosted)); + let choice_name = if nested_choice { + &body["tool_choice"]["function"]["name"] + } else { + &body["tool_choice"]["name"] + }; + assert_eq!(choice_name, alias); + assert_eq!(body["input"][0]["name"], alias); + assert_eq!(body["input"][0]["call_id"], "call_previous"); + assert_eq!(body["input"][1]["output"], "result"); + assert_function_response_restored(&encoded, alias, hosted); + } + } + } +} + +#[test] +fn hosted_names_avoid_session_injected_x_search_and_keep_aliases_stable() { + for name in ["web_search", "x_search"] { + let base = json!({"input": function_history(name), "tools": [function_tool(name)], + "tool_choice": {"type": "function", "name": name}}); + let encoded = + GrokResponsesRequest::encode(&raw_request(base.clone()), "grok-4.7", &client_key()) + .expect("without session"); + let alias = encoded.body()["tools"][0]["name"].as_str().expect("alias"); + let mut with_session = base; + with_session["prompt_cache_key"] = json!("hosted_name_session"); + let session = + GrokResponsesRequest::encode(&raw_request(with_session), "grok-4.7", &client_key()) + .expect("with session"); + assert_eq!(session.body()["tools"][0]["name"], alias); + assert_eq!(session.body()["tools"][1], json!({"type": "x_search"})); + assert_ne!(alias, "x_search"); + assert_eq!(session.body()["tool_choice"]["name"], alias); + assert_eq!(session.body()["input"][0]["name"], alias); + assert_function_response_restored(&session, alias, name); + } +} + +#[test] +fn hosted_names_rewrite_history_when_empty_tools_trigger_both_cache_tools() { + for name in ["web_search", "x_search"] { + let request = raw_request(json!({"prompt_cache_key": "history_only_session", + "input": function_history(name)})); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("history only"); + assert_eq!( + encoded.body()["tools"], + json!([{"type": "web_search"}, {"type": "x_search"}]) + ); + assert_eq!(encoded.body()["tool_choice"], "none"); + assert_ne!(encoded.body()["input"][0]["name"], name); + assert_eq!(encoded.body()["input"][0]["call_id"], "call_previous"); + } +} + +#[test] +fn hosted_names_do_not_rename_similar_or_namespaced_client_functions() { + let request = raw_request(json!({"tools": [function_tool("web_search_local"), + {"type": "namespace", "name": "local", "tools": [function_tool("x_search")]}, + {"type": "web_search"}, {"type": "x_search"}], "input": "search"})); + let encoded = + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()).expect("ordinary names"); + assert_eq!(encoded.body()["tools"][0]["name"], "web_search_local"); + assert_eq!(encoded.body()["tools"][1]["name"], "local__x_search"); +} + +#[test] +fn hosted_names_preserve_custom_tools_lowered_to_functions() { + let request = raw_request(json!({ + "tools": [{"type": "custom", "name": "web_search"}, {"type": "web_search"}], + "tool_choice": {"type": "custom", "name": "web_search"}, + "input": [{"type": "custom_tool_call", "call_id": "call_previous", + "name": "web_search", "input": "previous"}] + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("custom and hosted search"); + let alias = encoded.body()["tools"][0]["name"] + .as_str() + .expect("custom alias"); + assert_ne!(alias, "web_search"); + assert_eq!(encoded.body()["tool_choice"]["name"], alias); + assert_eq!(encoded.body()["input"][0]["name"], alias); + let item = json!({"type": "function_call", "id": "fc_custom", "call_id": "call_custom", + "name": alias, "arguments": "{\"input\":\"next\"}"}); + let frames = [ + json!({"type": "response.created", "response": {"id": "resp_custom"}}), + json!({"type": "response.output_item.added", "output_index": 0, "item": item}), + json!({"type": "response.output_item.done", "output_index": 0, "item": item}), + json!({"type": "response.completed", "response": {"id": "resp_custom", + "status": "completed", "output": [item]}}), + ] + .into_iter() + .map(|frame| format!("data: {frame}\n\n")) + .collect::(); + let events = GrokCanonicalDecoder::for_request("grok-4.7", &encoded) + .push(frames.as_bytes()) + .expect("custom response"); + let completed = events + .iter() + .filter_map(ProviderEvent::wire_event) + .next_back() + .expect("completed response") + .data(); + let restored = &completed["response"]["output"][0]; + assert_eq!(restored["type"], "custom_tool_call"); + assert_eq!(restored["name"], "web_search"); + assert_eq!(restored["call_id"], "call_custom"); + assert_eq!(restored["input"], "next"); +} + +#[test] +fn hosted_names_keep_existing_function_names_distinct_from_generated_aliases() { + let first = raw_request(json!({"tools": [function_tool("x_search")]})); + let first = + GrokResponsesRequest::encode(&first, "grok-4.7", &client_key()).expect("initial alias"); + let generated = first.body()["tools"][0]["name"].as_str().expect("alias"); + assert_ne!(generated, "x_search"); + for generated_first in [false, true] { + let mut tools = vec![function_tool("x_search"), function_tool(generated)]; + if generated_first { + tools.reverse(); + } + let request = raw_request( + json!({"tools": tools, "prompt_cache_key": "collision_session", "input": "search"}), + ); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("alias collision"); + let first = encoded.body()["tools"][0]["name"] + .as_str() + .expect("first alias"); + let second = encoded.body()["tools"][1]["name"] + .as_str() + .expect("second alias"); + assert_ne!(first, second); + assert_ne!(first, "x_search"); + assert_ne!(second, "x_search"); + let original_names = if generated_first { + [generated, "x_search"] + } else { + ["x_search", generated] + }; + assert_function_response_restored(&encoded, first, original_names[0]); + assert_function_response_restored(&encoded, second, original_names[1]); + } +} diff --git a/backend/crates/providers/xai/tests/transport/request/mod.rs b/backend/crates/providers/xai/tests/transport/request/mod.rs index 223658c84..3f7813f27 100644 --- a/backend/crates/providers/xai/tests/transport/request/mod.rs +++ b/backend/crates/providers/xai/tests/transport/request/mod.rs @@ -6,6 +6,11 @@ use serde_json::{Map, Value, json}; use provider_xai::{GrokRequestEncodeError, GrokResponsesRequest}; +mod hosted_names; +mod model_identity; +mod parameters; +mod web_search; + fn raw_request(body: Value) -> GenerateRequest { let Value::Object(body) = body else { panic!("request fixture must be an object"); @@ -191,7 +196,7 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { "reasoning.encrypted_content" ])) ); - assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("high"))); + assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("xhigh"))); assert_eq!(body.pointer("/reasoning/summary"), Some(&json!("auto"))); let request = raw_request(json!({ @@ -200,9 +205,9 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { "reasoning": {"effort": "xhigh"} })); let encoded = GrokResponsesRequest::encode(&request, "xai/grok-4.6-latest", &client_key()) - .expect("aliased request"); + .expect("routed model request"); let body = Value::Object(encoded.body().clone()); - assert_eq!(body.pointer("/model"), Some(&json!("grok-4.6"))); + assert_eq!(body.pointer("/model"), Some(&json!("xai/grok-4.6-latest"))); assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("xhigh"))); assert_eq!(body.pointer("/reasoning_effort"), None); assert_eq!(body.pointer("/reasoningEffort"), None); @@ -215,7 +220,10 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { let encoded = GrokResponsesRequest::encode(&request, "grok-composer-2.5-fast", &client_key()) .expect("Composer request"); let body = Value::Object(encoded.body().clone()); - assert_eq!(body.pointer("/reasoning"), None); + assert_eq!( + body.pointer("/reasoning"), + Some(&json!({"effort":"high", "summary":"auto"})) + ); let request = raw_request(json!({ "model": "client-model", @@ -232,7 +240,7 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { body.pointer("/include"), Some(&json!(["reasoning.encrypted_content"])) ); - assert_eq!(body.pointer("/reasoning"), None); + assert_eq!(body.pointer("/reasoning"), Some(&json!({"effort":"max"}))); let request = raw_request(json!({ "model": "client-model", @@ -243,7 +251,7 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { .expect("Composer reasoning without effort"); assert_eq!( Value::Object(encoded.body().clone()).pointer("/reasoning"), - None + Some(&json!({})) ); let request = raw_request(json!({ @@ -255,12 +263,12 @@ fn encoder_should_apply_build_defaults_and_normalize_reasoning_effort() { .expect("unknown provider prefix"); assert_eq!( Value::Object(encoded.body().clone()).pointer("/reasoning/effort"), - None + Some(&json!("max")) ); } #[test] -fn encoder_should_strip_grok_unsupported_fields() { +fn encoder_should_strip_codex_only_fields_without_model_guesses() { let request = raw_request(json!({ "model": "client-model", "input": "hello", @@ -285,16 +293,11 @@ fn encoder_should_strip_grok_unsupported_fields() { })); let encoded = GrokResponsesRequest::encode(&request, "grok-latest", &client_key()) - .expect("sanitized 4.5 request"); + .expect("sanitized request"); let body = Value::Object(encoded.body().clone()); for pointer in [ "/prompt_cache_retention", "/safety_identifier", - "/presence_penalty", - "/presencePenalty", - "/frequency_penalty", - "/frequencyPenalty", - "/stop", "/external_web_access", "/metadata/external_web_access", "/tools/0/external_web_access", @@ -311,12 +314,9 @@ fn encoder_should_strip_grok_unsupported_fields() { let encoded = GrokResponsesRequest::encode(&request, "grok-4.20-reasoning", &client_key()) .expect("sanitized 4.20 request"); let body = Value::Object(encoded.body().clone()); - assert_eq!( - body.pointer("/model"), - Some(&json!("grok-4.20-0309-reasoning")) - ); - assert_eq!(body.pointer("/logprobs"), None); - assert_eq!(body.pointer("/top_logprobs"), None); + assert_eq!(body.pointer("/model"), Some(&json!("grok-4.20-reasoning"))); + assert_eq!(body.pointer("/logprobs"), Some(&json!(true))); + assert_eq!(body.pointer("/top_logprobs"), Some(&json!(5))); } #[test] @@ -1004,7 +1004,7 @@ fn hosted_tool_choice_should_preserve_normalized_web_search_choice() { "model": "client", "input": "search", "tools": [ - {"type": "web_search_preview", "allowed_domains": ["example.com"]}, + {"type": "web_search_preview", "filters": {"allowed_domains": ["example.com"]}}, {"type": "x_search"} ], "tool_choice": {"type": "web_search_preview"} diff --git a/backend/crates/providers/xai/tests/transport/request/model_identity.rs b/backend/crates/providers/xai/tests/transport/request/model_identity.rs new file mode 100644 index 000000000..d6ac9e0d2 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/model_identity.rs @@ -0,0 +1,54 @@ +//! 路由已选定真实模型，xAI 编码器不再解释旧别名或 Provider 前缀 + +use super::*; + +#[test] +fn routed_model_identity_is_preserved_without_provider_local_aliases() { + let request = raw_request(json!({"model": "client-model", "input": "hello"})); + for model in [ + "grok", + "grok-latest", + "grok-build-latest", + "grok-4.5-latest", + "grok-4.6-latest", + "grok-4.3-latest", + "grok-build", + "grok-composer", + "composer-2.5", + "grok-4.20-reasoning", + "grok-4.20-non-reasoning", + "grok-4.20-multi-agent", + "grok-4.20-multi-agent-latest", + "grok-4.7", + "Grok-4.6", + "xai/grok-4.7", + "x-ai/grok-4.7", + "grok/grok-4.7", + "tenant/custom-model", + ] { + let encoded = GrokResponsesRequest::encode(&request, model, &client_key()) + .expect("routed upstream model"); + assert_eq!(encoded.body()["model"], model); + } +} + +#[test] +fn routed_model_identity_does_not_infer_unsupported_fields_from_old_slugs() { + let fields = json!({"presence_penalty": 0.1, "frequency_penalty": 0.2, + "stop": ["done"], "logprobs": true, "top_logprobs": 5}); + let request = raw_request(fields.clone()); + for model in [ + "grok-4.5", + "grok-4.6", + "grok-4.3", + "grok-4.20-reasoning", + "tenant/grok-4.6", + "grok-4.7", + ] { + let encoded = + GrokResponsesRequest::encode(&request, model, &client_key()).expect("request"); + for (key, value) in fields.as_object().expect("fields") { + assert_eq!(encoded.body().get(key), Some(value), "{model} {key}"); + } + } +} diff --git a/backend/crates/providers/xai/tests/transport/request/parameters.rs b/backend/crates/providers/xai/tests/transport/request/parameters.rs new file mode 100644 index 000000000..4f969c9ea --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/parameters.rs @@ -0,0 +1,75 @@ +//! 验证推理档位遵守 Responses wire 而不由模型 slug 猜测 + +use super::*; + +#[test] +fn official_reasoning_efforts_are_preserved_for_current_and_future_models() { + for model in [ + "grok-4.7", + "grok-future", + "grok-4.5", + "grok-4.6", + "grok-3-mini", + "grok-composer-2.5-fast", + ] { + for effort in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] { + let request = raw_request(json!({ + "input": "hello", + "reasoning": {"effort": effort, "summary": "auto"} + })); + let encoded = GrokResponsesRequest::encode(&request, model, &client_key()) + .expect("official effort wire"); + assert_eq!( + encoded.body()["reasoning"], + json!({"effort": effort, "summary": "auto"}), + "model {model}" + ); + } + } +} + +#[test] +fn invalid_reasoning_effort_is_rejected_instead_of_silently_dropped() { + for effort in [json!("future-value"), json!(42), json!({}), json!([])] { + let request = raw_request(json!({"reasoning": {"effort": effort}})); + assert!( + matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { + field: "reasoning.effort" + }) + ), + "effort {effort}" + ); + } +} + +#[test] +fn reasoning_effort_keeps_existing_case_and_whitespace_normalization() { + let request = raw_request(json!({"reasoning": {"effort": " XHIGH "}})); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("normalized official effort"); + assert_eq!(encoded.body()["reasoning"]["effort"], "xhigh"); +} + +#[test] +fn reasoning_shape_is_validated_while_absent_effort_uses_server_default() { + for reasoning in [ + Value::Null, + json!({}), + json!({"effort": null}), + json!({"summary":"auto"}), + ] { + let request = raw_request(json!({"reasoning": reasoning})); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("no explicit effort"); + assert_eq!(encoded.body().get("reasoning"), Some(&reasoning)); + } + for reasoning in [json!("high"), json!(42), json!([])] { + let request = raw_request(json!({"reasoning": reasoning})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "reasoning" }) + )); + } +} diff --git a/backend/crates/providers/xai/tests/transport/request/web_search.rs b/backend/crates/providers/xai/tests/transport/request/web_search.rs new file mode 100644 index 000000000..a050a4117 --- /dev/null +++ b/backend/crates/providers/xai/tests/transport/request/web_search.rs @@ -0,0 +1,136 @@ +//! Web Search 域名过滤的协议转换与字段归属测试 + +use super::*; + +#[test] +fn web_search_filters_preserve_excluded_domains_for_supported_aliases() { + for kind in [ + "web_search", + "web_search_preview", + "web_search_preview_2025_03_11", + "web_search_2025_08_26", + ] { + let request = raw_request(json!({ + "input": "search", + "tools": [{"type": kind, "filters": {"excluded_domains": ["example.com"]}}] + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("excluded domains"); + + assert_eq!( + encoded.body()["tools"], + json!([{"type": "web_search", "filters": {"excluded_domains": ["example.com"]}}]), + "{kind}" + ); + } +} + +#[test] +fn web_search_filters_reject_combined_nonempty_allow_and_block_lists() { + for tool in [ + json!({"type": "web_search", "filters": { + "allowed_domains": ["allowed.example"], "excluded_domains": ["excluded.example"] + }}), + json!({"type": "web_search_preview", "allowed_domains": ["allowed.example"], + "filters": {"excluded_domains": ["excluded.example"]} + }), + ] { + let request = raw_request(json!({"input": "search", "tools": [tool]})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "tools" }) + )); + } +} + +#[test] +fn web_search_filters_keep_empty_lists_unbounded() { + for (filters, expected) in [ + ( + json!({"allowed_domains": [], "excluded_domains": ["example.com"]}), + json!({"excluded_domains": ["example.com"]}), + ), + ( + json!({"allowed_domains": ["example.com"], "excluded_domains": []}), + json!({"allowed_domains": ["example.com"]}), + ), + (json!({"excluded_domains": []}), Value::Null), + ( + json!({"allowed_domains": [], "excluded_domains": []}), + Value::Null, + ), + (json!({"allowed_domains": []}), Value::Null), + ] { + let request = raw_request(json!({ + "input": "search", "tools": [{"type": "web_search", "filters": filters}] + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()) + .expect("empty lists are unbounded"); + assert_eq!(encoded.body()["tools"][0]["filters"], expected); + } +} + +#[test] +fn web_search_filters_reject_retired_top_level_allowed_domains() { + for tool in [ + json!({"type": "web_search_preview", "allowed_domains": ["example.com"]}), + json!({"type": "web_search", "allowed_domains": ["example.com"], + "filters": {"allowed_domains": ["example.com"]}}), + json!({"type": "web_search", "allowed_domains": ["one.example"], + "filters": {"allowed_domains": ["other.example"]}}), + json!({"type": "web_search", "allowed_domains": []}), + ] { + let request = raw_request(json!({"input": "search", "tools": [tool]})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "tools" }) + )); + } +} + +#[test] +fn web_search_filters_reject_invalid_excluded_domains() { + for excluded in [ + Value::Null, + json!("example.com"), + json!([" "]), + json!([3]), + json!(["example.com", false]), + ] { + let request = raw_request(json!({"input": "search", "tools": [{ + "type": "web_search", "filters": {"excluded_domains": excluded} + }]})); + assert!(matches!( + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()), + Err(GrokRequestEncodeError::InvalidRequestField { field: "tools" }) + )); + } +} + +#[test] +fn web_search_filters_only_interpret_owned_tool_fields() { + let schema = json!({"type": "object", "properties": { + "filters": {"type": "string"}, "excluded_domains": {"type": "integer"} + }}); + let opaque = json!({"excluded_domains": "application data"}); + let request = raw_request(json!({ + "input": "search", + "excluded_domains": opaque, + "tools": [ + {"type": "function", "name": "lookup", "parameters": schema}, + {"type": "x_search", "excluded_domains": opaque}, + {"type": "web_search", "excluded_domains": opaque, + "filters": {"future_filter": opaque, "excluded_domains": ["example.com"]}} + ] + })); + let encoded = + GrokResponsesRequest::encode(&request, "grok-4.7", &client_key()).expect("field ownership"); + + assert_eq!(encoded.body()["excluded_domains"], opaque); + assert_eq!(encoded.body()["tools"][0]["parameters"], schema); + assert_eq!(encoded.body()["tools"][1]["excluded_domains"], opaque); + assert_eq!( + encoded.body()["tools"][2], + json!({"type": "web_search", "filters": {"excluded_domains": ["example.com"]}}) + ); +} diff --git a/docs/api.md b/docs/api.md index 49afe31d2..2b243ec22 100644 --- a/docs/api.md +++ b/docs/api.md @@ -322,6 +322,10 @@ OAuth 账号默认 `prefer_websocket`，客户端使用 HTTP/SSE 时仍可能选 客户端配置的 `supports_websockets` 只控制第一段连接，不是服务端传输策略开关。 上游在响应终态前发送 Close 1000 仍属于失败，不能按“正常关闭”计为成功 +内建 xAI 适配器不支持 `generate=false` 的非生成预热请求，在发送推理请求前返回 +`400 unsupported_prewarm`，不会将其当作普通生成调用。`generate=true` 或省略该字段正常执行； +客户端可跳过预热继续正式请求，外部插件 adapter 按各自的协议能力处理 + Codex OAuth backend 的候选上游为 WS 时，无 `previous_response_id` 的普通新链若规范化 `response.create` 达到 15 MiB，发送前选择 HTTP/SSE；HTTP 和 WS 入站均适用，下游交付协议不变。 该阈值为已观察到的上游消息大小边界预留余量，不是网关输入长度上限或 OpenAI 公布的统一限制； @@ -359,7 +363,9 @@ API Key 上游返回完整 Codex `models` 目录时沿用该合同；仅返回 未提供的推理能力保持未知，不补充推理档位。 模型别名仅替换 `slug`，不替换上游展示名、提示词、能力或 `priority`；保持原生模型顺序，新增别名附在后面。 目录按当前路由快照的模型存在性及账号模型政策过滤，避免公布已知无法路由的模型；新模型需待后台目录对账后进入列表。 -xAI 没有 Codex 原生目录，继续使用明确的通用画像适配 +xAI 没有 Codex 原生目录，使用官方 CLI proxy 当前目录生成通用画像。推理菜单读取 +`reasoning_efforts` 对象列表，默认值只取菜单中的显式标记；缺失时保持未知，不以首项代替， +也不回退旧字段或其他端点的目录形状 目录账号只能来自本次 Client Key 冻结的账号范围。OpenAI 的 OAuth 目录按账号 ID 排序，使用首个成功读取的 合格账号，最多尝试三个账号；API Key 目录按账号模型权限过滤后聚合。同名模型优先采用 OAuth 原生对象， @@ -380,7 +386,8 @@ xAI 没有 Codex 原生目录，继续使用明确的通用画像适配 Codex 专用目录中的 `context_window` 与 `max_context_window` 分别表示默认上下文窗口和客户端本地 覆盖的上限。OpenAI Provider 原样保留上游对应字段，缺失与 `null` 不互相转换；网关不通过部署配置 覆盖这些值。Codex 客户端配置 `model_context_window` 后，按该值与非空 `max_context_window` 的较小值 -使用窗口；上限为空时保留客户端本地值。xAI 目录只声明一个窗口，其 Provider 继续以该值作为客户端覆盖上限 +使用窗口；上限为空时保留客户端本地值。xAI 默认窗口取 `context_window`，缺失时取 `context_windows` +中的首个有效值；覆盖上限取默认窗口和可选窗口的最大值，只有标量时以该值作为上限 ### 透传与错误恢复 @@ -398,6 +405,16 @@ Responses wire 之间的协议转换层，转换只在 xAI Provider 内完成。 上游结构化错误的 message/code/type 按上述边界交付客户端，其中内嵌的账号指纹 UUID 已脱敏。模型映射是 全局精确映射，未命中时模型名原样交给候选 Provider；分组只限定账号集合，不参与模型改名 +xAI 编码器同样保留路由选定的完整模型名，不内置 `grok`、`grok-latest` 等旧型号别名，也不剥离模型前缀； +需要别名时使用显式模型映射。目录读取失败按错误处理，不用硬编码旧型号补全能力 + +xAI 的 `reasoning.effort` 接受 `none / minimal / low / medium / high / xhigh / max`，去除两端 +空白并转为小写后保留档位，不按模型名称降档或删除；模型是否支持该选择由上游判断。 +未知值或错误类型返回字段错误，其他 reasoning 字段保留。`web_search.filters` 支持 +`allowed_domains` 与 `excluded_domains`，两者不能同时为非空；客户端函数与托管搜索工具 +同名时使用内部别名，并在返回工具调用时还原客户端名称。搜索工具的旧顶层 `allowed_domains` +不再映射到 filters，直接返回请求错误 + OpenAI 明确返回 `server_is_overloaded`、`slow_down` 或模型容量不足错误时，代理在允许安全重放且 尚未交付输出的前提下，先做最多 3 次同账号重试，再通过现有调度换号。没有有效服务器建议时， 按 500ms 起步的指数退避等待，单次最多 8 秒；有效 `Retry-After` 优先，接受秒数、HTTP 日期和零延迟， @@ -982,6 +999,9 @@ OAuth start 使用： `quota_exhausted`。额度观测不会清除凭据过期、无效或封禁事实；这些事实统一投影为 `error`，并由 `errorReason` 区分。额度接口的 401/403 也不足以判定 refresh token 永久失效，credential 终态只由 OAuth refresh 的明确永久错误写入 +- xAI 额度百分比只读取 `creditUsagePercent`，缺失时保持未知，不由旧月额度与用量计算。 + 预付余额保留上游账本符号，非零余额及剩余按量额度可独立证明可用；窗口已过期或缺少可用证据时， + 不凭旧快照恢复账号 - 正常 Responses 请求会解析上游响应的 rate-limit headers，合并进同一 quota 快照并同步状态。Free、 K12 等套餐共用该状态机；套餐只参与账号展示和按套餐隔离的模型目录 cache，不存在 K12 专属额度路径 - 账号详情的 Token 统计、模型排行和列表 Token 汇总优先使用账号级周额度窗口，无可统计的周窗口时 @@ -1801,6 +1821,10 @@ OpenAI 与 xAI 的本地费用估算按实际发送给上游的请求模型（`u 用量计算；客户端请求 A、路由后发送 B 时按 B 计价，响应返回 C 不改变计价模型。实际发送模型缺少 定价时不估算，也不借用响应模型的价格。Provider 明确上报的已计费金额仍优先于本地估算 +xAI 的 `cost_in_usd_ticks=0` 表示上游没有报告费用，不证明请求免费；正数才作为上游账单。 +未报告时，只有已知价格且能覆盖本次费用的请求才使用本地估算，未知价格或无法覆盖的托管工具费用保持未知。 +动态 latest 与旧 beta/internal 别名不借用具体旧型号的内置价格；显式配置的模型价格仍可用于估算 + OpenAI Responses 用量记录的 `serviceTier` 与本地费用估算统一采用 Provider 最终发给上游的请求 `service_tier`，不使用响应档位覆盖或回退。例如发送 `priority`、响应回显 `default` 时，仍显示 `Fast` 并按 Priority 价格估算。未发送档位时，`serviceTier` 保持缺失，展示与估算按标准档处理。 diff --git a/docs/architecture.md b/docs/architecture.md index b88bce5df..f0f8d3fd1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -250,16 +250,22 @@ HTTP 字段见 [插件 API](api.md#12-插件管理)，更新与数据恢复见 [ | `account` | Provider 账号/credential/quota 值对象、持久化端口与请求级账号选择；`scope` 持有分组、账号目录和冻结账号范围 | | `policy` | Client API Key 准入、原始 Key 设置与客户端版本策略；只使用账号范围、Provider 身份和基础校验 | | `metering` | 标准化 Usage、金额、费用估算与费用明细；不表示账号或开票系统 | -| `upstream` | 跨 Engine、Event、Error 与 Provider 共用的 transport 名称、发送状态和不透明上游值 | +| `upstream` | 跨 Engine、Event、Error 与 Provider 共用的 transport 名称、尝试传输档位、发送状态和不透明上游值 | | `lifecycle` | 取消信号、可选执行截止与租约生命周期、连接注册与 drain 合同 | -| `engine` | attempt、发送/提交屏障、执行编排和持久化调用时序 | -| `routing` | 冻结路由事实、请求计划、Provider 只读目录合同以及运行时快照的表示与编译 | +| `middleware` | 与领域无关的类型化组合器及单次续体 | +| `engine` | attempt、发送/提交屏障、执行编排、具体中间件与上游适配计划、持久化调用时序 | +| `routing` | 冻结路由事实、请求计划、Provider 只读目录合同、扩展集合身份与保活合同、运行时快照及请求设置派生 | | `runtime` | 当前快照的发布、读取、revision 订阅与周期对账任务 | -| `settings` | 请求设置事实及其编译、宿主与 Key 默认值解析、显式覆盖来源；由 `routing` 快照组合并冻结 | +| `settings` | 设置值及纯编译规则；由 `routing` 快照组合并冻结 | `event` 通过 `validation` / `upstream` 使用基础值，不依赖承载原始事件的执行错误；账号值对象和 选择策略通过 `identity` / `account::scope` 使用身份与范围，不依赖路由计划。`routing` 和 `error` 的 -相关公开类型通过 re-export 引用上述定义，类型所有权和 Core 内部依赖归属定义模块。架构测试约束这些叶子依赖 +相关公开类型通过 re-export 引用上述定义，类型所有权和 Core 内部依赖归属定义模块。架构测试约束这些叶子依赖， +并检查一级 owner 的显式 `crate` / `self` / `super` 路径与重导出形成的依赖图；它不替代 Rust 名称解析和行为测试 + +`routing::request_settings` 统一解析宿主与 Key 默认值、显式覆盖来源及请求快照重算，不向全局发布请求改写。 +`routing::extensions` 只表达扩展集合的中立身份、静态目录和租约；`engine::extensions` 按冻结身份弱引用索引 +具体执行计划，插件 Runtime（`gateway-plugin/runtime`）的 `PreparedSet` 负责强持有计划。请求取得的冻结计划同时保活该集合，旧代次随最后一个使用者退出 Provider 模型能力、目录代次与 `ProviderCatalogPort` 由 `routing::catalog` 定义。快照编译和对账只消费 该只读合同，不反向依赖执行注册表；`ProviderRegistry` 实现目录端口，维护唯一的 Provider