feat(plugin): expose request lifecycles through onion middleware

Compose HTTP, WebSocket, request, attempt and public service calls through
shared onion middleware with scoped settings and invocation context.

Adopt trusted-process manifests, unified observation and forward-compatible
SDK contracts. Align management UI, documentation and the workbench reference.
Preserve native response bytes and WebSocket failure lifecycles.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-29 18:20:22 +08:00
co-authored by Codex
parent d83d3eb803
commit fe78c230e5
300 changed files with 18625 additions and 5633 deletions
+22
View File
@@ -1674,6 +1674,8 @@ dependencies = [
"gateway-core",
"heck",
"hex",
"prettyplease",
"quote",
"rand_core 0.6.4",
"secrecy",
"semver",
@@ -1681,6 +1683,7 @@ dependencies = [
"serde_json",
"sha2 0.11.0",
"subtle",
"syn 2.0.119",
"tempfile",
"thiserror",
"tokio",
@@ -1703,6 +1706,9 @@ dependencies = [
"gateway-core",
"gateway-protocol",
"headers",
"http",
"http-body",
"http-body-util",
"secrecy",
"serde",
"serde_json",
@@ -1728,6 +1734,9 @@ dependencies = [
"chrono-tz",
"futures",
"futures-timer",
"http",
"http-body",
"http-body-util",
"semver",
"serde",
"serde_json",
@@ -1801,6 +1810,9 @@ dependencies = [
"gateway-protocol",
"gateway-store",
"hex",
"http",
"http-body",
"http-body-util",
"image",
"jsonschema",
"roxmltree",
@@ -3147,6 +3159,16 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c"
[[package]]
name = "prettyplease"
version = "0.2.37"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
dependencies = [
"proc-macro2",
"syn 2.0.119",
]
[[package]]
name = "primefield"
version = "0.14.0"
+45 -1
View File
@@ -202,6 +202,23 @@ async fn launch(
plugin_runtime,
plugin_runtime.bind_model_ports(&nested_models, &affinity_lookup)
);
let settings = gateway_admin::initialize_settings(
store.admin_ports().settings(),
core.snapshot_control(),
admin_providers,
std::sync::Arc::new(gateway_host::pricing::ModelsDevPricing),
);
let mut services = gateway_admin::service::Registry::new({
let snapshots = core.snapshots();
let middleware = plugin_runtime.middleware_registry();
std::sync::Arc::new(move || {
let snapshot = snapshots.snapshot_for_diagnostics()?;
middleware.resolve(snapshot.extensions()?)
})
});
or_shutdown!(plugin_runtime, services.register_settings(&settings));
let services = std::sync::Arc::new(services);
or_shutdown!(plugin_runtime, plugin_runtime.bind_services(&services));
let commands = or_shutdown!(plugin_runtime, plugin_runtime.prepare_command_line().await);
or_shutdown!(plugin_runtime, store.start_command_line_writes());
let result = commands
@@ -277,6 +294,14 @@ async fn launch(
client,
store.admin_ports(),
gateway_admin::AdminRuntimePorts {
service_middleware: {
let snapshots = core.snapshots();
let middleware = plugin_runtime.middleware_registry();
std::sync::Arc::new(move || {
let snapshot = snapshots.snapshot_for_diagnostics()?;
middleware.resolve(snapshot.extensions()?)
})
},
providers: admin_providers,
plugin_preparation: plugin_runtime.clone(),
plugin_management: plugin_runtime.clone(),
@@ -330,9 +355,28 @@ async fn launch(
host.worker_health(),
host.connection_lifecycle(),
)
);
)
.with_middleware({
let snapshots = core.snapshots();
let middleware = plugin_runtime.middleware_registry();
move |snapshot| {
if let Some(snapshot) = snapshot {
return middleware.resolve(snapshot.extensions()?);
}
// 数据面未就绪时,管理与诊断路由仍可使用发布候选中的插件。
let diagnostic = snapshots.snapshot_for_diagnostics()?;
middleware.resolve(diagnostic.extensions()?)
}
});
let http_dispatcher = api.dispatcher();
or_shutdown!(plugin_runtime, plugin_runtime.bind_http(&http_dispatcher));
host.report_startup_ready("API");
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_services(&admin.services().public_services())
);
let mut plan = store.take_worker_contributions();
plan.extend(core.take_worker_contributions());
plan.extend(openai.take_worker_contributions());
+6 -1
View File
@@ -76,7 +76,12 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() {
"account/selection.rs" => Some(&["account", "concurrency", "identity", "validation"]),
"concurrency.rs" => Some(&["error"]),
"account/store.rs" => Some(&["account", "error", "identity", "validation"]),
path if path.starts_with("policy/") => Some(&["account", "policy", "validation"]),
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("account/") => Some(&["account", "identity", "validation"]),
_ => None,
};
+2 -1
View File
@@ -28,7 +28,8 @@ cpr-plugin package \
资源映射相对于作者清单所在目录解析。输出包括完整插件 ID、版本和目标 triple 命名的 `.tar.gz` 及 `.sha256`;归档根目录直接包含 `plugin.json`、可执行文件和已声明资源,不包含开发依赖
作者清单可省略普通贡献项的 `id`、`version` 和固定阶段;`middleware` 的 `request`/`attempt` 仍须显式选择。
作者清单可省略普通贡献项的 `id`、`version` 和固定阶段;`middleware` 必须显式声明版本 `3`,
并从 `http`、`websocket`、`service`、`request`、`attempt` 中选择挂载阶段。
CLI 与 SDK 的 `Manifest::from_author_slice`、`PluginBuilder::from_json` 共用规范化入口;归档中的
`plugin.json` 写入补全后的贡献项字段,不保留另一份作者声明。
完整规则见 [SDK 清单](../../crates/gateway-plugin/sdk/docs/manifest.md)
+3 -3
View File
@@ -237,7 +237,7 @@ fn assert_packaged_resources(icon_path: &str, icon_mime: &str, icon: &[u8]) {
fn source_manifest() -> &'static [u8] {
br#"{
"manifestVersion":1,
"manifestVersion":2,
"name":"request-workbench",
"displayName":"Request Workbench",
"publisher":"codex-proxy",
@@ -248,10 +248,10 @@ fn source_manifest() -> &'static [u8] {
"main":"bin/plugin",
"runtime":"trustedProcess",
"contributes":{
"middleware":{"stages":["request"],"inputFormats":["openai"],"outputFormats":["openai"]},
"middleware":{"version":3,"stages":["request"],"inputFormats":["openai"],"outputFormats":["openai"]},
"management":{}
},
"permissions":["requests","accounts"],
"icon":"assets/icon.png",
"resources":{"LICENSE":"text/plain","assets/icon.png":"image/png","web/index.html":"text/html","web/app.js":"text/javascript","web/app.css":"text/css"}
}"#
+3
View File
@@ -36,6 +36,9 @@ chrono-tz = "0.10"
tokio = { version = "1.52", features = ["rt", "time", "macros", "sync"] }
[dev-dependencies]
prettyplease = "0.2"
quote = "1"
syn = { version = "2", features = ["full", "visit-mut"] }
tempfile = "3"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "test-util"] }
+37 -6
View File
@@ -20,6 +20,7 @@ pub mod backup;
pub mod freeze_recovery;
pub mod model;
pub mod ports;
pub mod service;
mod use_case;
pub use use_case::plugins::{PluginDistributionPorts, PluginManagementService, PluginsService};
@@ -186,6 +187,7 @@ pub enum AdminConfigError {
/// 字段全部私有;调用方经 accessor 直接调用能力,不需要命名内部 `use_case` 模块。
#[derive(Clone)]
pub struct AdminServices {
public_services: Arc<service::Registry>,
plugins: Arc<PluginsService>,
plugin_management: Arc<PluginManagementService>,
proxies: Arc<dyn ProxiesService>,
@@ -205,6 +207,10 @@ pub struct AdminServices {
}
impl AdminServices {
pub fn public_services(&self) -> Arc<service::Registry> {
self.public_services.clone()
}
#[must_use]
pub fn plugin_management(&self) -> &PluginManagementService {
&self.plugin_management
@@ -267,6 +273,10 @@ impl AdminServices {
}
#[must_use]
pub fn settings_handle(&self) -> Arc<dyn SettingsService> {
self.settings.clone()
}
pub fn settings(&self) -> &dyn SettingsService {
self.settings.as_ref()
}
@@ -313,6 +323,7 @@ impl AdminBundle {
/// 组合根提供给控制面的运行能力;与配置和存储端口分别传入。
pub struct AdminRuntimePorts {
pub service_middleware: service::PlanSource,
pub plugin_preparation: Arc<dyn ports::plugins::PluginPreparation>,
pub plugin_management: Arc<dyn ports::plugin_management::PluginManagement>,
pub published_snapshot: gateway_core::runtime::RuntimeSnapshotHandle,
@@ -361,6 +372,7 @@ async fn initialize_inner(
plugin_accounts: Option<Arc<dyn PluginAccountAccess>>,
) -> Result<AdminBundle, AdminError> {
let AdminRuntimePorts {
service_middleware,
plugin_preparation,
plugin_management,
published_snapshot,
@@ -435,7 +447,16 @@ async fn initialize_inner(
});
let import_tasks = use_case::import_tasks::DefaultImportTasksService::new(credentials.clone());
let import_task = use_case::import_tasks::ImportTaskWorker(import_tasks.clone());
let settings = initialize_settings(
store.settings(),
snapshot.clone(),
registry.clone(),
pricing_source,
);
let mut public_services = service::Registry::new(service_middleware);
public_services.register_settings(&settings)?;
let services = AdminServices {
public_services: Arc::new(public_services),
plugin_management: Arc::new(PluginManagementService::new(
plugin_management,
store.plugins(),
@@ -476,12 +497,7 @@ async fn initialize_inner(
store.settings(),
registry.clone(),
)),
settings: Arc::new(DefaultSettingsService::new(
store.settings(),
snapshot.clone(),
registry,
pricing_source,
)),
settings,
system,
credentials,
plugin_accounts,
@@ -597,3 +613,18 @@ fn freeze_recovery_worker_contribution(
.map_err(|_| AdminError::internal("冻结恢复 Worker 注册信息不合法"))?;
Ok(vec![WorkerContribution::Registration(registration)])
}
/// 设置服务不依赖 Web 管理会话,CLI 与服务器通过同一用例执行事务和发布。
pub fn initialize_settings(
store: Arc<dyn ports::store::SettingsStore>,
snapshot: Arc<dyn SnapshotControl>,
providers: ProviderAdminRegistry,
pricing_source: Arc<dyn ports::pricing::PricingSource>,
) -> Arc<dyn SettingsService> {
Arc::new(DefaultSettingsService::new(
store,
snapshot,
providers,
pricing_source,
))
}
@@ -14,8 +14,14 @@ pub fn admin_session_actor_ref(admin_user_id: &str) -> String {
/// 已认证的管理主体。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AdminPrincipal {
Session { admin_user_id: String },
Session {
admin_user_id: String,
},
ApiKey,
/// 安装者完整信任的插件;身份仅由内部 HTTP 分派端口签发。
Plugin {
instance_id: String,
},
}
/// 传给管理用例的安全请求上下文。
@@ -33,6 +39,7 @@ impl AdminRequestContext {
admin_user_id: admin_user_id.clone(),
},
AdminPrincipal::ApiKey => MutationActor::AdminApiKey,
AdminPrincipal::Plugin { .. } => MutationActor::System,
};
MutationContext {
actor,
@@ -107,7 +107,10 @@ impl AdminError {
}
/// PostgreSQL 中所有正整数 revision 的管理层表示。
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
)]
#[serde(transparent)]
pub struct Revision(NonZeroU64);
impl Revision {
@@ -156,7 +159,8 @@ impl PageSize {
}
/// 可审计管理写操作的发起者。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum MutationActor {
AdminSession { admin_user_id: String },
AdminApiKey,
@@ -164,7 +168,8 @@ pub enum MutationActor {
}
/// 管理写操作必须携带的审计上下文。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub struct MutationContext {
pub actor: MutationActor,
pub request_id: String,
@@ -5,12 +5,6 @@ use serde::{Deserialize, Serialize};
use crate::model::Revision;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct PluginPermissionGrant {
pub permission: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PluginFailurePolicy {
@@ -26,6 +20,9 @@ pub struct PluginCapabilityBinding {
pub stage: String,
pub order: i32,
pub failure_policy: PluginFailurePolicy,
/// 观察绑定选择事件类型;其他能力不设置事件。
#[serde(default, skip_serializing_if = "Option::is_none")]
pub event: Option<String>,
#[serde(default)]
pub client_key_ids: Vec<String>,
#[serde(default)]
@@ -66,7 +63,6 @@ pub struct PluginInstance {
pub trusted_process: bool,
pub configuration: serde_json::Value,
pub secrets: BTreeMap<String, SecretString>,
pub grants: Vec<PluginPermissionGrant>,
pub bindings: Vec<PluginCapabilityBinding>,
pub revision: Revision,
}
@@ -78,6 +78,7 @@ pub struct PluginManagementRequest {
pub path: String,
pub query: String,
pub content_type: Option<String>,
pub headers: Vec<gateway_core::engine::middleware::MiddlewareHeader>,
pub body: Vec<u8>,
pub request_id: String,
}
@@ -85,5 +86,6 @@ pub struct PluginManagementRequest {
pub struct PluginManagementResponse {
pub status: u16,
pub content_type: String,
pub headers: Vec<gateway_core::engine::middleware::MiddlewareHeader>,
pub body: Arc<[u8]>,
}
@@ -29,19 +29,17 @@ pub struct PluginHostCompatibility {
pub manifest_schema_versions: Vec<u32>,
pub protocol_versions: Vec<u32>,
pub capabilities: Vec<PluginHostCapability>,
pub permissions: Vec<String>,
}
impl PluginHostCompatibility {
/// 拒绝空集合、重复项和不受控标识,避免发行清单把“未知”解释为通配。
#[must_use]
pub fn is_valid(&self) -> bool {
if self.schema_version != 1
if self.schema_version != 2
|| !valid_versions(&self.manifest_schema_versions)
|| !valid_versions(&self.protocol_versions)
|| self.capabilities.is_empty()
|| self.capabilities.len() > 64
|| self.permissions.len() > 64
{
return false;
}
@@ -53,10 +51,7 @@ impl PluginHostCompatibility {
}) {
return false;
}
let mut permissions = BTreeSet::new();
!self.permissions.iter().any(|permission| {
!valid_identifier(permission) || !permissions.insert(permission.as_str())
})
true
}
#[must_use]
@@ -65,11 +60,6 @@ impl PluginHostCompatibility {
.iter()
.any(|entry| entry.capability == capability && entry.versions.contains(&version))
}
#[must_use]
pub fn supports_permission(&self, permission: &str) -> bool {
self.permissions.iter().any(|entry| entry == permission)
}
}
fn valid_versions(versions: &[u32]) -> bool {
@@ -118,7 +108,6 @@ pub struct PluginCompatibilityRequirements {
pub manifest_schema_version: u32,
pub protocol_version: u32,
pub capabilities: Vec<(String, u32)>,
pub permissions: Vec<String>,
}
/// 一次下载使用的非敏感出站代理身份;地址与认证由受管代理记录持有。
@@ -177,12 +166,15 @@ impl PluginSource {
/// 管理端保留的单项贡献声明;map key 是能力标识,`id` 是实例绑定引用的扩展项身份。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[serde(rename_all = "camelCase")]
pub struct PluginContribution {
pub id: String,
pub version: u32,
#[serde(default)]
pub stages: Vec<String>,
#[serde(default)]
pub input_formats: Vec<String>,
#[serde(default)]
pub output_formats: Vec<String>,
}
@@ -195,7 +187,6 @@ pub enum PluginArtifactIcon {
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct PluginArtifactIconVariants {
pub light: String,
pub dark: String,
@@ -214,17 +205,10 @@ pub struct PluginArtifactIconResource {
pub body: Vec<u8>,
}
/// 管理端使用的静态描述,不承载 SDK 的运行时消息或宿主实现类型。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct PluginPermissionDescription {
pub permission: String,
pub label: String,
pub description: String,
}
/// 已校验制品的持久化投影;读取时忽略未知字段,可选字段缺省沿用清单语义。
/// 身份与执行所需事实仍为必填,包清单及协议版本由 Runtime 独立校验。
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[serde(rename_all = "camelCase")]
pub struct PluginArtifactMetadata {
pub plugin_id: String,
pub version: String,
@@ -239,13 +223,18 @@ pub struct PluginArtifactMetadata {
#[serde(default)]
pub icon: Option<PluginArtifactIcon>,
pub contributes: BTreeMap<String, PluginContribution>,
pub requested_permissions: Vec<String>,
#[serde(default = "empty_configuration_schema")]
pub configuration_schema: serde_json::Value,
#[serde(default)]
pub secret_fields: Vec<String>,
#[serde(default)]
pub state_namespaces: Vec<state::PluginStateSchema>,
}
fn empty_configuration_schema() -> serde_json::Value {
serde_json::json!({})
}
/// 由包检查端口产生的不可变制品;不派生 Debug,包内可能含敏感用户数据。
#[derive(Clone)]
pub struct InspectedPluginArtifact {
@@ -6,7 +6,7 @@ use crate::model::Revision;
/// 已由 Runtime 校验的命名空间 schema 与配额事实。
#[derive(Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[serde(rename_all = "camelCase")]
pub struct PluginStateSchema {
pub namespace: String,
pub schema_version: u32,
@@ -3,7 +3,7 @@
use gateway_core::metering::{ModelPriceOverride, PricingOverrides};
use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PricingCatalog {
pub defaults: PricingOverrides,
pub overrides: PricingOverrides,
@@ -27,6 +27,7 @@ pub struct PricingSyncPreview {
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[derive(serde::Serialize)]
pub struct SyncPricing {
pub preview: PricingSyncPreview,
pub models: BTreeMap<String, BTreeSet<String>>,
@@ -34,7 +35,7 @@ pub struct SyncPricing {
pub type PricingSyncChanges = BTreeMap<String, BTreeMap<String, Option<ModelPriceOverride>>>;
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum PricingChange {
Replace(ModelPriceOverride),
Multiplier(u32),
@@ -42,7 +43,7 @@ pub enum PricingChange {
Delete,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct UpdatePricing {
pub provider: String,
pub models: Vec<String>,
@@ -23,7 +23,8 @@ pub type ProviderRequestProfileUpdates =
pub use gateway_core::account::RotationStrategy;
/// 完整运行设置事实。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RuntimeSettings {
pub request_profiles: ProviderRequestProfiles,
pub config_revision: Revision,
@@ -59,8 +60,11 @@ pub struct RuntimeSettings {
}
/// 原子替换运行设置的命令。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ReplaceRuntimeSettings {
/// 读取设置时的版本;与写入在同一事务内比较,防止覆盖并发更新。
pub expected_revision: Revision,
/// 只覆盖提交的 Provider;未提交项保留当前持久值。
pub request_profile_updates: ProviderRequestProfileUpdates,
pub request_location_enabled: bool,
@@ -94,7 +98,8 @@ pub struct ReplaceRuntimeSettings {
}
/// 明文管理员 API Key;按产品约束明文落库,但禁止 Debug 泄漏。
#[derive(Clone, PartialEq, Eq)]
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(transparent)]
pub struct AdminApiKey(String);
impl AdminApiKey {
@@ -122,13 +127,16 @@ impl fmt::Debug for AdminApiKey {
}
/// 管理员 API Key 更新结果。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct AdminApiKeyMutation {
pub config_revision: Revision,
pub exists: bool,
}
/// 新管理员 API Key 的一次性返回结果。
#[derive(serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RegeneratedAdminApiKey {
pub mutation: AdminApiKeyMutation,
pub key: AdminApiKey,
@@ -154,3 +162,45 @@ impl RuntimeSettings {
.and_then(|provider| self.request_profiles.get(&provider))
}
}
impl From<RuntimeSettings> for ReplaceRuntimeSettings {
fn from(settings: RuntimeSettings) -> Self {
Self {
expected_revision: settings.config_revision,
request_profile_updates: settings
.request_profiles
.into_iter()
.map(|(provider, value)| (provider, Some(value)))
.collect(),
request_location_enabled: settings.request_location_enabled,
request_location: settings.request_location,
model_mappings: settings.model_mappings,
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: settings.refresh_concurrency,
max_concurrent_per_account: settings.max_concurrent_per_account,
request_interval_ms: settings.request_interval_ms,
max_waiting_per_key: settings.max_waiting_per_key,
max_waiting_per_account: settings.max_waiting_per_account,
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes,
smart_scheduling: settings.smart_scheduling,
rotation_strategy: settings.rotation_strategy,
min_codex_desktop_version: settings.min_codex_desktop_version,
min_codex_cli_version: settings.min_codex_cli_version,
usage_retention_days: settings.usage_retention_days,
ops_event_retention_days: settings.ops_event_retention_days,
audit_retention_days: settings.audit_retention_days,
account_auto_freeze_enabled: settings.account_auto_freeze_enabled,
account_auto_freeze_threshold: settings.account_auto_freeze_threshold,
account_auto_freeze_window_seconds: settings.account_auto_freeze_window_seconds,
account_auto_freeze_duration_seconds: settings.account_auto_freeze_duration_seconds,
account_auto_freeze_probe_enabled: settings.account_auto_freeze_probe_enabled,
account_auto_freeze_probe_model: settings.account_auto_freeze_probe_model,
account_auto_freeze_adaptive_concurrency: settings
.account_auto_freeze_adaptive_concurrency,
account_warmup_enabled: settings.account_warmup_enabled,
account_warmup_schedule_time: settings.account_warmup_schedule_time,
account_warmup_model: settings.account_warmup_model,
}
}
}
@@ -12,8 +12,8 @@ use crate::model::{
/// 仅按 Admin 固定的发布引用分派;Runtime 不自行选择当前代次,也不拥有管理员身份。
#[async_trait]
pub trait PluginManagement: Send + Sync {
/// 复核固定发布代次中的模型访问域;调用方已校验持久实例仍为当前目标。
async fn authorize_models(
/// 复核固定发布代次中的可用目标;调用方已校验持久实例仍为当前版本。
async fn validate_target(
&self,
published: &ExtensionSetReference,
target: &PluginManagementTarget,
@@ -176,23 +176,6 @@ pub trait PluginDistribution: Send + Sync {
/// Runtime 只解释插件包格式;安装事务与来源选择归 Admin。
#[async_trait]
pub trait PluginPackageInspector: Send + Sync {
/// 当前宿主对权限标识的说明,供展示投影使用,不属于不可变制品元数据。
fn permission_descriptions(
&self,
permissions: &[String],
) -> Vec<crate::model::plugins::PluginPermissionDescription> {
permissions
.iter()
.map(
|permission| crate::model::plugins::PluginPermissionDescription {
permission: permission.clone(),
label: permission.clone(),
description: String::new(),
},
)
.collect()
}
async fn inspect(
&self,
archive: Arc<[u8]>,
@@ -314,7 +297,7 @@ pub trait PluginStore: Send + Sync {
source: PluginSource,
context: &MutationContext,
) -> AdminStoreResult<PluginArtifactMutation>;
/// 接受不可变制品声明的全部能力域;接受事实不随实例设置变化。
/// 确认信任并安装精确摘要的制品;安装事实不随实例设置变化。
async fn accept_artifact(
&self,
digest: &str,
@@ -0,0 +1,129 @@
//! 主动服务调用的统一组合入口;普通业务方法由所属 HTTP 或执行边界包裹。
mod registry;
mod settings;
pub use registry::Registry;
use std::{future::Future, sync::Arc};
use futures::future::BoxFuture;
use gateway_core::{
engine::middleware::FrozenMiddlewarePlan,
middleware::{
compose,
service::{Context, Error, Value},
},
};
use crate::model::{AdminError, AdminErrorKind};
pub type PlanSource = Arc<dyn Fn() -> Option<FrozenMiddlewarePlan> + Send + Sync>;
/// 空计划也是快照,只有未解析的外部入口才读取当前发布集合。
#[derive(Clone, Default)]
pub enum Plan {
#[default]
Current,
Frozen(Option<FrozenMiddlewarePlan>),
}
/// 非服务入口也有父调用身份和取消信号;发布计划在首次进入公开服务时解析。
#[derive(Clone)]
pub struct Origin {
pub request_id: String,
pub call_id: String,
pub cancellation: gateway_core::lifecycle::CancellationToken,
pub extensions: gateway_core::engine::extensions::ExtensionCallScope,
pub plan: Plan,
}
impl From<Context> for Origin {
fn from(context: Context) -> Self {
Self {
request_id: context.request_id,
call_id: context.call_id,
cancellation: context.cancellation,
extensions: context.extensions,
plan: Plan::Frozen(Some(context.plan)),
}
}
}
tokio::task_local! { static CURRENT: Origin; }
/// task local 随 future poll 进入/退出,不泄漏给并行任务。
pub async fn scope<T>(context: impl Into<Origin>, future: impl Future<Output = T>) -> T {
CURRENT.scope(context.into(), future).await
}
pub(crate) async fn invoke(
source: &PlanSource,
operation: &'static str,
input: Value,
terminal: impl FnOnce(Value) -> BoxFuture<'static, Result<Value, Error>> + Send + 'static,
) -> Result<Value, Error> {
let parent = CURRENT.try_with(Clone::clone).ok();
let plan = match parent.as_ref().map(|parent| &parent.plan) {
Some(Plan::Frozen(plan)) => plan.clone(),
_ => source(),
};
let call_id = uuid::Uuid::now_v7().to_string();
let parent_call_id = parent.as_ref().map(|parent| parent.call_id.clone());
let origin = Origin {
request_id: parent
.as_ref()
.map_or_else(|| call_id.clone(), |parent| parent.request_id.clone()),
cancellation: parent
.as_ref()
.map(|parent| parent.cancellation.clone())
.unwrap_or_default(),
extensions: parent.map(|parent| parent.extensions).unwrap_or_default(),
call_id,
plan: Plan::Frozen(plan.clone()),
};
let Some(plan) = plan.filter(FrozenMiddlewarePlan::has_service) else {
return tokio::select! {
biased;
() = origin.cancellation.cancelled() => Err(Error::unavailable("服务调用已取消")),
result = scope(origin.clone(), terminal(input)) => result,
};
};
let context = Context {
operation,
request_id: origin.request_id,
call_id: origin.call_id,
parent_call_id,
cancellation: origin.cancellation,
extensions: origin.extensions,
plan: plan.clone(),
};
let inner = context.clone();
let next = compose(Vec::new(), move |input| {
Box::pin(scope(inner, terminal(input)))
});
tokio::select! {
biased;
() = context.cancellation.cancelled() => Err(Error::unavailable("服务调用已取消")),
result = plan.handle_service(context.clone(), input, next) => result,
}
}
pub(crate) fn encode_error(error: AdminError) -> Error {
Error {
kind: 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",
}
.into(),
message: error.message().to_owned(),
details: None,
}
}
@@ -0,0 +1,76 @@
//! 公开操作登记在所属服务;注册表只校验类型和分派,不实现业务。
use crate::model::AdminError;
use futures::future::BoxFuture;
use gateway_core::middleware::service::{Error, Value};
use serde::{Serialize, de::DeserializeOwned};
use std::{collections::BTreeMap, sync::Arc};
type Handler = Arc<dyn Fn(Value) -> BoxFuture<'static, Result<Value, Error>> + Send + Sync>;
pub struct Registry {
middleware: super::PlanSource,
operations: BTreeMap<&'static str, Handler>,
}
impl Registry {
pub fn new(middleware: super::PlanSource) -> Self {
Self {
middleware,
operations: BTreeMap::new(),
}
}
/// 网关与命令行复用同一组设置服务;注册不会执行读取或写入。
pub fn register_settings(
&mut self,
service: &Arc<dyn crate::SettingsService>,
) -> Result<(), AdminError> {
super::settings::register(self, service)
}
pub(crate) fn register<I, O>(
&mut self,
name: &'static str,
handler: impl Fn(I) -> BoxFuture<'static, Result<O, AdminError>> + Send + Sync + 'static,
) -> Result<(), AdminError>
where
I: DeserializeOwned + Send + 'static,
O: Serialize + Send + 'static,
{
let std::collections::btree_map::Entry::Vacant(entry) = self.operations.entry(name) else {
return Err(AdminError::internal(format!("公开服务重复登记:{name}")));
};
let handler = Arc::new(handler);
entry.insert(Arc::new(move |input| {
let handler = handler.clone();
Box::pin(async move {
let input = serde_json::from_value(input)
.map_err(|_| Error::invalid("服务输入类型无效"))?;
let output = handler(input).await.map_err(super::encode_error)?;
serde_json::to_value(output).map_err(|_| Error::unavailable("服务结果编码失败"))
})
}));
Ok(())
}
pub async fn call(&self, operation: &str, input: Value) -> Result<Value, Error> {
let (operation, handler) =
self.operations
.get_key_value(operation)
.ok_or_else(|| Error {
details: None,
kind: "not_found".into(),
message: format!("公开服务不存在:{operation}"),
})?;
let handler = handler.clone();
super::invoke(&self.middleware, operation, input, move |input| {
handler(input)
})
.await
}
pub fn operations(&self) -> impl Iterator<Item = &'static str> + '_ {
self.operations.keys().copied()
}
}
@@ -0,0 +1,85 @@
//! 设置服务的唯一公开操作声明;SDK 操作和数据合同由对应合同测试生成。
use crate::{
SettingsService,
model::{
AdminError, MutationContext,
pricing::{SyncPricing, UpdatePricing},
settings::ReplaceRuntimeSettings,
},
};
use std::sync::Arc;
pub(super) fn register(
registry: &mut super::Registry,
service: &Arc<dyn SettingsService>,
) -> Result<(), AdminError> {
// 声明保留 SDK 类型名和原生参数转换,分派名称由业务方法名直接派生。
macro_rules! register {
($sdk:ident, $method:ident, $input:ty, $pattern:pat $(, $argument:expr)* $(,)?) => {{
let service = service.clone();
registry.register(concat!("settings.", stringify!($method)), move |$pattern: $input| {
let service = service.clone();
Box::pin(async move { service.$method($($argument),*).await })
})?;
}};
}
register!(Load, load, (), ());
register!(
Replace,
replace,
(MutationContext, ReplaceRuntimeSettings),
(context, command),
&context,
command
);
register!(ApiKeyExists, admin_api_key_exists, (), ());
register!(
RegenerateApiKey,
regenerate_admin_api_key,
MutationContext,
context,
&context
);
register!(
DeleteApiKey,
delete_admin_api_key,
MutationContext,
context,
&context
);
register!(Pricing, pricing, (), ());
register!(PreviewPricingSync, preview_pricing_sync, (), ());
register!(
Sync,
sync_pricing,
(MutationContext, SyncPricing),
(context, command),
&context,
command
);
register!(
Update,
update_pricing,
(MutationContext, UpdatePricing),
(context, command),
&context,
command
);
register!(
ClientProfileOptions,
client_profile_options,
String,
provider,
&provider
);
register!(
PreviewClientProfile,
preview_client_profile,
(String, Option<gateway_core::account::OpaqueProviderData>),
(provider, configuration),
&provider,
configuration.as_ref()
);
Ok(())
}
@@ -85,10 +85,6 @@ impl SystemUpdatePreflight for PluginSystemUpdatePreflight {
.any(|(capability, version)| {
!compatibility.supports_capability(capability, *version)
})
|| requirements
.permissions
.iter()
.any(|permission| !compatibility.supports_permission(permission))
{
return Err(incompatible(&instance.id));
}
@@ -14,7 +14,7 @@ use crate::{
distribution::VerifiedPluginArtifact,
instances::{
ConfigurePluginInstance, PluginCapabilityBinding, PluginFailurePolicy,
PluginInstance, PluginPermissionGrant,
PluginInstance,
},
},
},
@@ -52,13 +52,6 @@ pub struct PluginsService {
}
impl PluginsService {
pub fn permission_descriptions(
&self,
permissions: &[String],
) -> Vec<crate::model::plugins::PluginPermissionDescription> {
self.inspector.permission_descriptions(permissions)
}
#[must_use]
pub fn new(
store: Arc<dyn PluginStore>,
@@ -384,11 +377,14 @@ pub(super) fn default_bindings(metadata: &PluginArtifactMetadata) -> Vec<PluginC
| "maintenance"
)
})
.flat_map(|(_, contribution)| {
contribution
.stages
.iter()
.map(|stage| PluginCapabilityBinding {
.flat_map(|(capability, contribution)| {
contribution.stages.iter().flat_map(move |stage| {
let events: &[Option<&str>] = if capability == "observer" {
&[Some("request_completed"), Some("websocket_response")]
} else {
&[None]
};
events.iter().map(move |event| PluginCapabilityBinding {
contribution: contribution.id.clone(),
stage: stage.clone(),
order: 0,
@@ -403,8 +399,10 @@ pub(super) fn default_bindings(metadata: &PluginArtifactMetadata) -> Vec<PluginC
account_group_ids: Vec::new(),
provider_ids: Vec::new(),
models: Vec::new(),
event: event.map(str::to_owned),
identity_bindings: Vec::new(),
})
})
})
.collect()
}
@@ -425,12 +423,6 @@ fn default_instance(
trusted_process: true,
configuration,
secrets: Default::default(),
grants: metadata
.requested_permissions
.iter()
.cloned()
.map(|permission| PluginPermissionGrant { permission })
.collect(),
bindings,
revision,
}
@@ -10,7 +10,7 @@ use crate::{
plugins::instances::{
ConfigurePluginInstance, PluginInstance, PluginInstanceMutation,
PluginInstanceReplacement, PluginInstanceRuntime, PluginInstanceRuntimeStatus,
PluginInstanceSnapshot, PluginInstanceView, PluginPermissionGrant,
PluginInstanceSnapshot, PluginInstanceView,
},
plugins::state::{PluginStateCommit, PluginStateConfiguration},
},
@@ -225,13 +225,6 @@ impl PluginsService {
.as_ref()
.map_or_else(Default::default, |instance| instance.secrets.clone())
}),
grants: artifact
.metadata
.requested_permissions
.iter()
.cloned()
.map(|permission| PluginPermissionGrant { permission })
.collect(),
bindings: input.bindings,
revision: candidate_revision,
};
@@ -690,7 +683,6 @@ fn validate_input(instance: &PluginInstance) -> Result<(), AdminError> {
|| !instance.configuration.is_object()
|| serde_json::to_vec(&instance.configuration).map_or(true, |bytes| bytes.len() > 48 * 1024)
|| instance.secrets.len() > 64
|| instance.grants.len() > 32
|| instance.bindings.len() > 64
{
return Err(AdminError::invalid("插件实例配置不合法或超过大小限制"));
@@ -16,7 +16,7 @@ use crate::{
ports::{plugin_management::PluginManagement, plugins::PluginStore},
};
/// 菜单是发布视图的投影;每次业务动作仍复核持久实例,旧页面不能保留已撤销权限。
/// 菜单来自发布视图;每次业务动作复核持久实例,旧页面不能继续使用已停用或换版的目标。
pub struct PluginManagementService {
runtime: Arc<dyn PluginManagement>,
store: Arc<dyn PluginStore>,
@@ -25,12 +25,9 @@ pub struct PluginManagementService {
impl PluginManagementService {
/// 页面发起模型请求前复核持久目标与当前发布代次;长响应可重复调用以收敛撤销。
pub async fn authorize_models(
&self,
target: &PluginManagementTarget,
) -> Result<(), AdminError> {
pub async fn validate_target(&self, target: &PluginManagementTarget) -> Result<(), AdminError> {
let reference = self.authorize(target).await?;
self.runtime.authorize_models(&reference, target).await
self.runtime.validate_target(&reference, target).await
}
pub async fn start_callback(
@@ -252,10 +252,11 @@ fn version_bindings(
let previous = capability.and_then(|capability| current.contributes.get(capability));
previous.is_none_or(|contribution| {
!contribution.stages.contains(&binding.stage)
|| instance
.bindings
.iter()
.any(|old| old.contribution == contribution.id && old.stage == binding.stage)
|| instance.bindings.iter().any(|old| {
old.contribution == contribution.id
&& old.stage == binding.stage
&& old.event == binding.event
})
})
});
for previous in &instance.bindings {
@@ -276,7 +277,9 @@ fn version_bindings(
let mut binding = previous.clone();
binding.contribution.clone_from(&contribution.id);
if let Some(default) = bindings.iter_mut().find(|default| {
default.contribution == binding.contribution && default.stage == binding.stage
default.contribution == binding.contribution
&& default.stage == binding.stage
&& default.event == binding.event
}) {
*default = binding;
} else {
@@ -102,6 +102,50 @@ impl DefaultSettingsService {
}
}
fn validate_settings(command: &ReplaceRuntimeSettings) -> Result<(), AdminError> {
let valid = command.request_location.validate().is_ok()
&& command.responses_max_decompressed_body_bytes > 0
&& isize::try_from(command.responses_max_decompressed_body_bytes).is_ok()
&& command.refresh_margin_seconds > 0
&& command.refresh_concurrency > 0
&& command.max_waiting_per_key <= 1_000
&& command.max_waiting_per_account <= 1_000
&& (1..=120).contains(&command.concurrency_wait_timeout_seconds)
&& command.usage_retention_days >= 31
&& command.ops_event_retention_days > 0
&& command.audit_retention_days > 0
&& valid_client_version(command.min_codex_desktop_version.as_deref())
&& valid_client_version(command.min_codex_cli_version.as_deref())
&& valid_probe_model(command.account_auto_freeze_probe_model.as_deref())
&& gateway_core::provider_ports::valid_warmup_schedule_time(
&command.account_warmup_schedule_time,
)
&& valid_probe_model(command.account_warmup_model.as_deref())
&& (!command.account_warmup_enabled || command.account_warmup_model.is_some())
&& i64::try_from(command.request_interval_ms).is_ok()
&& (2..=1_000).contains(&command.account_auto_freeze_threshold)
&& (60..=3_600).contains(&command.account_auto_freeze_window_seconds)
&& (300..=604_800).contains(&command.account_auto_freeze_duration_seconds);
if valid {
Ok(())
} else {
Err(AdminError::invalid("运行时设置不满足约束"))
}
}
fn valid_client_version(value: Option<&str>) -> bool {
value.is_none_or(|value| CodexClientVersion::parse(value).is_ok())
}
fn valid_probe_model(value: Option<&str>) -> bool {
value.is_none_or(|value| {
!value.is_empty()
&& value.len() <= 128
&& value == value.trim()
&& !value.bytes().any(|byte| byte.is_ascii_control())
})
}
#[async_trait]
impl SettingsService for DefaultSettingsService {
async fn preview_pricing_sync(
@@ -365,7 +409,17 @@ impl SettingsService for DefaultSettingsService {
.store
.replace_runtime_settings(command, context)
.await
.map_err(|error| map_store_error(error, "runtime settings"))?;
.map_err(|error| {
if matches!(
error.kind(),
crate::ports::store::AdminStoreErrorKind::Conflict
| crate::ports::store::AdminStoreErrorKind::StaleRevision
) {
AdminError::conflict("运行设置已被其他调用修改,请重新读取并确认后提交")
} else {
map_store_error(error, "runtime settings")
}
})?;
publish_committed(self.snapshot.as_ref(), settings.config_revision).await?;
Ok(settings)
}
@@ -406,47 +460,3 @@ impl SettingsService for DefaultSettingsService {
Ok(mutation)
}
}
fn validate_settings(command: &ReplaceRuntimeSettings) -> Result<(), AdminError> {
let valid = command.request_location.validate().is_ok()
&& command.responses_max_decompressed_body_bytes > 0
&& isize::try_from(command.responses_max_decompressed_body_bytes).is_ok()
&& command.refresh_margin_seconds > 0
&& command.refresh_concurrency > 0
&& command.max_waiting_per_key <= 1_000
&& command.max_waiting_per_account <= 1_000
&& (1..=120).contains(&command.concurrency_wait_timeout_seconds)
&& command.usage_retention_days >= 31
&& command.ops_event_retention_days > 0
&& command.audit_retention_days > 0
&& valid_client_version(command.min_codex_desktop_version.as_deref())
&& valid_client_version(command.min_codex_cli_version.as_deref())
&& valid_probe_model(command.account_auto_freeze_probe_model.as_deref())
&& gateway_core::provider_ports::valid_warmup_schedule_time(
&command.account_warmup_schedule_time,
)
&& valid_probe_model(command.account_warmup_model.as_deref())
&& (!command.account_warmup_enabled || command.account_warmup_model.is_some())
&& i64::try_from(command.request_interval_ms).is_ok()
&& (2..=1_000).contains(&command.account_auto_freeze_threshold)
&& (60..=3_600).contains(&command.account_auto_freeze_window_seconds)
&& (300..=604_800).contains(&command.account_auto_freeze_duration_seconds);
if valid {
Ok(())
} else {
Err(AdminError::invalid("运行时设置不满足约束"))
}
}
fn valid_client_version(value: Option<&str>) -> bool {
value.is_none_or(|value| CodexClientVersion::parse(value).is_ok())
}
fn valid_probe_model(value: Option<&str>) -> bool {
value.is_none_or(|value| {
!value.is_empty()
&& value.len() <= 128
&& value == value.trim()
&& !value.bytes().any(|byte| byte.is_ascii_control())
})
}
@@ -1,6 +1,7 @@
//! 模型模块测试:镜像 src/model/ 中带逻辑的模块。
mod backup;
mod plugins;
mod provider_credentials;
mod quota_forecast;
mod quota_forecast_sampling;
@@ -0,0 +1,105 @@
use gateway_admin::model::plugins::PluginArtifactMetadata;
use serde_json::{Value, json};
fn metadata() -> Value {
json!({
"pluginId": "test.example", "version": "1.0.0", "name": "example",
"displayName": "Example", "publisher": "test", "description": "Example plugin",
"license": "MIT", "sha256": "a".repeat(64), "platforms": ["linux-x86_64"],
"contributes": {"middleware": {
"id": "test.example.middleware", "version": 3,
"stages": [], "inputFormats": [], "outputFormats": []
}},
"configurationSchema": {}, "secretFields": [], "stateNamespaces": []
})
}
#[test]
fn metadata_ignores_unknown_fields_at_each_object_boundary() {
let mut original = metadata();
original["icon"] = json!({"light": "light.svg", "dark": "dark.svg"});
original["stateNamespaces"] = json!([{
"namespace": "cache", "schemaVersion": 1, "schemaSha256": "b".repeat(64),
"schema": {}, "maximumRecords": 10, "maximumBytes": 1024,
"maximumValueBytes": 128, "migratesFrom": []
}]);
let expected: PluginArtifactMetadata = serde_json::from_value(original.clone()).unwrap();
for path in ["", "/contributes/middleware", "/icon", "/stateNamespaces/0"] {
let mut extended = original.clone();
extended.pointer_mut(path).unwrap()["futureField"] = json!({"enabled": true});
assert_eq!(
serde_json::from_value::<PluginArtifactMetadata>(extended).unwrap(),
expected,
"unknown field at {path}"
);
}
}
#[test]
fn metadata_defaults_only_optional_fields() {
let original = metadata();
let expected: PluginArtifactMetadata = serde_json::from_value(original.clone()).unwrap();
let mut older = original;
for field in ["configurationSchema", "secretFields", "stateNamespaces"] {
older.as_object_mut().unwrap().remove(field);
}
for field in ["stages", "inputFormats", "outputFormats"] {
older["contributes"]["middleware"]
.as_object_mut()
.unwrap()
.remove(field);
}
assert_eq!(
serde_json::from_value::<PluginArtifactMetadata>(older).unwrap(),
expected
);
}
#[test]
fn metadata_still_requires_identity_and_execution_facts() {
for field in [
"pluginId",
"version",
"name",
"publisher",
"sha256",
"platforms",
"contributes",
] {
let mut incomplete = metadata();
incomplete.as_object_mut().unwrap().remove(field);
assert!(
serde_json::from_value::<PluginArtifactMetadata>(incomplete).is_err(),
"missing {field}"
);
}
for field in ["id", "version"] {
let mut incomplete = metadata();
incomplete["contributes"]["middleware"]
.as_object_mut()
.unwrap()
.remove(field);
assert!(
serde_json::from_value::<PluginArtifactMetadata>(incomplete).is_err(),
"missing contribution {field}"
);
}
}
#[test]
fn metadata_still_rejects_wrong_types_in_known_fields() {
for (path, value) in [
("/version", json!(1)),
("/secretFields", json!(false)),
("/stateNamespaces", json!({})),
("/contributes/middleware/version", json!("3")),
("/contributes/middleware/stages", json!("http")),
] {
let mut invalid = metadata();
*invalid.pointer_mut(path).unwrap() = value;
assert!(
serde_json::from_value::<PluginArtifactMetadata>(invalid).is_err(),
"wrong type at {path}"
);
}
}
@@ -105,6 +105,7 @@ pub(super) struct AdminHarness {
plugin_store: Arc<dyn gateway_admin::ports::plugins::PluginStore>,
plugin_inspector: Arc<dyn gateway_admin::ports::plugins::PluginPackageInspector>,
client_key_verifier: Arc<dyn ClientKeyVerifier>,
service_middleware: gateway_admin::service::PlanSource,
}
impl AdminHarness {
@@ -132,6 +133,7 @@ impl AdminHarness {
plugin_store: Arc::new(plugins::TestPluginPorts),
plugin_inspector: Arc::new(plugins::TestPluginPorts),
client_key_verifier: Arc::new(UnavailableClientKeyVerifier),
service_middleware: Arc::new(|| None),
}
}
@@ -190,6 +192,11 @@ impl AdminHarness {
self
}
pub(super) fn service_middleware(mut self, source: gateway_admin::service::PlanSource) -> Self {
self.service_middleware = source;
self
}
pub(super) fn backup(mut self, backup: BackupStorePorts) -> Self {
self.backup = backup;
self
@@ -261,6 +268,7 @@ impl AdminHarness {
Arc::new(plugins::TestPluginPorts),
),
gateway_admin::AdminRuntimePorts {
service_middleware: self.service_middleware,
plugin_preparation: Arc::new(plugins::TestPluginPorts),
plugin_management: Arc::new(plugins::TestPluginPorts),
published_snapshot: gateway_core::runtime::RuntimeSnapshotHandle::default(),
@@ -1101,3 +1109,4 @@ fn unavailable_system() -> SystemOperationError {
"unavailable in this test",
)
}
mod service_contract;
@@ -103,7 +103,7 @@ impl PluginStore for Fixture {
trusted_process: true,
configuration: serde_json::json!({}),
secrets: BTreeMap::new(),
grants: Vec::new(),
bindings: Vec::new(),
revision: Revision::new(1).expect("revision"),
}],
@@ -126,7 +126,7 @@ impl PluginStore for Fixture {
platforms: Vec::new(),
icon: None,
contributes: BTreeMap::new(),
requested_permissions: Vec::new(),
configuration_schema: serde_json::json!({}),
secret_fields: Vec::new(),
state_namespaces: Vec::new(),
@@ -285,11 +285,11 @@ impl SystemOperations for PreflightingSystem {
#[tokio::test]
async fn system_update_accepts_compatible_enabled_plugins_at_the_same_revision() {
let fixture = Fixture::new(requirements("^1.0", "executor", "log"), &[7]);
let fixture = Fixture::new(requirements("^1.0", "executor"), &[7]);
let services = super::AdminHarness::new()
.plugins(fixture.clone(), fixture)
.system(Arc::new(PreflightingSystem {
candidate: candidate("executor", "log"),
candidate: candidate("executor"),
}))
.build()
.await;
@@ -303,11 +303,11 @@ async fn system_update_accepts_compatible_enabled_plugins_at_the_same_revision()
#[tokio::test]
async fn system_update_rejects_missing_target_capability_and_revision_changes() {
let incompatible = Fixture::new(requirements("^1.0", "executor", "log"), &[7]);
let incompatible = Fixture::new(requirements("^1.0", "executor"), &[7]);
let services = super::AdminHarness::new()
.plugins(incompatible.clone(), incompatible)
.system(Arc::new(PreflightingSystem {
candidate: candidate("models", "log"),
candidate: candidate("models"),
}))
.build()
.await;
@@ -319,11 +319,11 @@ async fn system_update_rejects_missing_target_capability_and_revision_changes()
.is_err()
);
let stale = Fixture::new(requirements("^1.0", "executor", "log"), &[7, 8]);
let stale = Fixture::new(requirements("^1.0", "executor"), &[7, 8]);
let services = super::AdminHarness::new()
.plugins(stale.clone(), stale)
.system(Arc::new(PreflightingSystem {
candidate: candidate("executor", "log"),
candidate: candidate("executor"),
}))
.build()
.await;
@@ -336,32 +336,26 @@ async fn system_update_rejects_missing_target_capability_and_revision_changes()
);
}
fn requirements(
host_version: &str,
capability: &str,
permission: &str,
) -> PluginCompatibilityRequirements {
fn requirements(host_version: &str, capability: &str) -> PluginCompatibilityRequirements {
PluginCompatibilityRequirements {
host_version: host_version.into(),
manifest_schema_version: 2,
protocol_version: 3,
capabilities: vec![(capability.into(), 1)],
permissions: vec![permission.into()],
}
}
fn candidate(capability: &str, permission: &str) -> SystemUpdateCandidate {
fn candidate(capability: &str) -> SystemUpdateCandidate {
let manifest = serde_json::to_vec(&serde_json::json!({
"schema_version": 1,
"sealed": true,
"gateway_version": "1.2.0",
"gateway_git_sha": "a".repeat(40),
"plugin_host": {
"schema_version": 1,
"schema_version": 2,
"manifest_schema_versions": [2],
"protocol_versions": [3],
"capabilities": [{ "capability": capability, "versions": [1] }],
"permissions": [permission],
},
"plugins": [],
}))
@@ -337,7 +337,7 @@ async fn existing_plugin_instance_prevents_an_extra_default_instance() {
trusted_process: false,
configuration: json!({}),
secrets: BTreeMap::new(),
grants: Vec::new(),
bindings: Vec::new(),
revision: revision(1),
};
@@ -389,7 +389,7 @@ fn metadata(with_required_default: bool) -> PluginArtifactMetadata {
contribution("test.example.middleware", &["observation", "request"]),
),
]),
requested_permissions: vec!["network".into()],
configuration_schema: json!({
"type":"object",
"properties": {
@@ -519,7 +519,7 @@ fn preview_fixture() -> Fixture {
platforms: vec!["linux-x86_64".into()],
icon: None,
contributes: Default::default(),
requested_permissions: vec!["log".into()],
configuration_schema: serde_json::json!({"type":"object"}),
secret_fields: vec![],
state_namespaces: vec![],
@@ -635,7 +635,6 @@ async fn verifying_remote_package_uses_inspector_and_returns_fixed_facts_without
.await
.unwrap();
assert_eq!(verified.metadata.sha256, "a".repeat(64));
assert_eq!(verified.metadata.requested_permissions, vec!["log"]);
assert!(
matches!(verified.source, PluginSource::Github { ref repository, ref tag, ref asset, .. }
if repository == "owner/plugin" && tag == "v2.0.0" && asset == "plugin.tar.gz")
@@ -20,8 +20,7 @@ use gateway_admin::{
ConfigurePluginInstance, PluginCapabilityBinding, PluginFailurePolicy,
PluginInstance, PluginInstanceMutation, PluginInstanceReplacement,
PluginInstanceRuntime, PluginInstanceRuntimeFailure, PluginInstanceRuntimeStatus,
PluginInstanceSnapshot, PluginPermissionGrant, PluginVersionConfiguration,
RollbackPluginInstance,
PluginInstanceSnapshot, PluginVersionConfiguration, RollbackPluginInstance,
},
state::{
ApplyPluginStateMigration, DeletePluginState, PluginStateCommit,
@@ -103,7 +102,7 @@ impl LifecycleFixture {
trusted_process: true,
configuration: json!({}),
secrets: Default::default(),
grants: vec![],
bindings: vec![],
revision: revision(10),
}],
@@ -599,7 +598,7 @@ fn artifact(digest: &str, plugin_id: &str, version: &str) -> InstalledPluginArti
output_formats: Vec::new(),
},
)]),
requested_permissions: vec![],
configuration_schema: json!({"type":"object"}),
secret_fields: vec![],
state_namespaces: vec![],
@@ -998,13 +997,12 @@ async fn failed_state_migration_never_overwrites_a_concurrent_configuration_chan
}
#[tokio::test]
async fn rollback_restores_target_configuration_secrets_and_bindings_and_derives_permissions() {
async fn rollback_restores_target_configuration_secrets_and_bindings() {
use secrecy::ExposeSecret as _;
let fixture = LifecycleFixture::new(MigrationBehavior::Succeed);
{
let mut data = fixture.data.lock().unwrap();
data.artifacts[1].metadata.requested_permissions = vec!["network".into()];
let instance = &mut data.snapshot.instances[0];
instance.configuration = json!({"label":"retained"});
instance
@@ -1019,6 +1017,7 @@ async fn rollback_restores_target_configuration_secrets_and_bindings_and_derives
account_group_ids: Vec::new(),
provider_ids: Vec::new(),
models: Vec::new(),
event: None,
identity_bindings: Vec::new(),
});
}
@@ -1031,12 +1030,6 @@ async fn rollback_restores_target_configuration_secrets_and_bindings_and_derives
assert_eq!(result.instance.name, before.name);
assert_eq!(result.instance.enabled, before.enabled);
assert_eq!(result.instance.configuration, json!({"oldVersion": true}));
assert_eq!(
result.instance.grants,
[PluginPermissionGrant {
permission: "network".into()
}]
);
assert!(result.instance.bindings.is_empty());
assert_eq!(
result.instance.secrets["token"].expose_secret(),
@@ -1275,6 +1268,7 @@ async fn version_plan_keeps_explicit_values_adds_defaults_and_remaps_bindings_wi
account_group_ids: vec![],
provider_ids: vec![],
models: vec!["test-model".into()],
event: None,
identity_bindings: vec![],
}];
}
@@ -1399,3 +1393,56 @@ async fn switching_versions_uses_saved_settings_and_rejects_stale_retry() {
assert_eq!(error.kind(), AdminErrorKind::Conflict);
assert_eq!(fixture.snapshot().config_revision, revision);
}
#[tokio::test]
async fn version_plan_preserves_observer_event_scopes_and_disabled_subscriptions() {
for include_websocket in [false, true] {
let fixture = LifecycleFixture::new(MigrationBehavior::Succeed);
{
let mut data = fixture.data.lock().unwrap();
data.history.clear();
for (index, artifact) in data.artifacts.iter_mut().enumerate() {
artifact.metadata.contributes.clear();
artifact.metadata.contributes.insert(
"observer".into(),
PluginContribution {
id: format!("observer-{index}"),
version: 1,
stages: vec!["observation".into()],
input_formats: vec![],
output_formats: vec![],
},
);
}
data.snapshot.instances[0].bindings = ["request_completed", "websocket_response"]
.into_iter()
.take(if include_websocket { 2 } else { 1 })
.map(|event| PluginCapabilityBinding {
contribution: "observer-0".into(),
stage: "observation".into(),
event: Some(event.into()),
order: 7,
failure_policy: PluginFailurePolicy::Observe,
client_key_ids: vec![],
account_group_ids: vec![],
provider_ids: vec![],
models: vec![format!("model-{event}")],
identity_bindings: vec![],
})
.collect();
}
let before = fixture.snapshot();
let plan = service(fixture.clone(), Arc::new(Published::default()))
.version_plan(&before.instances[0].id, NEW_ARTIFACT)
.await
.unwrap();
assert_eq!(plan.bindings.len(), before.instances[0].bindings.len());
for (binding, previous) in plan.bindings.iter().zip(&before.instances[0].bindings) {
assert_eq!(binding.contribution, "observer-1");
assert_eq!(binding.event, previous.event);
assert_eq!(binding.models, previous.models);
assert_eq!(binding.order, previous.order);
}
assert!(fixture.data.lock().unwrap().saves.is_empty());
}
}
@@ -107,7 +107,7 @@ impl gateway_admin::ports::proxy::ProxyStore for TestPluginPorts {
#[async_trait]
impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts {
async fn authorize_models(
async fn validate_target(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
@@ -142,7 +142,7 @@ impl PluginPackageInspector for Fixture {
)],
icon: None,
contributes: Default::default(),
requested_permissions: Vec::new(),
configuration_schema: serde_json::json!({}),
secret_fields: Vec::new(),
state_namespaces: Vec::new(),
@@ -372,11 +372,10 @@ fn identity() -> OfficialPluginReleaseIdentity {
fn plugin_host() -> serde_json::Value {
serde_json::json!({
"schema_version": 1,
"schema_version": 2,
"manifest_schema_versions": [2],
"protocol_versions": [3],
"capabilities": [{ "capability": "executor", "versions": [1] }],
"permissions": ["log"],
})
}
@@ -0,0 +1,290 @@
//! 从宿主类型和操作声明生成独立 SDK 合同;普通测试阻止两侧声明漂移。
use std::{collections::BTreeMap, path::Path};
use quote::quote;
use syn::{Item, Type, parse::Parse, parse_quote, visit_mut::VisitMut};
struct Operation {
name: syn::Ident,
method: syn::Ident,
input: Type,
}
impl Parse for Operation {
fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
let name = input.parse()?;
input.parse::<syn::Token![,]>()?;
let method = input.parse()?;
input.parse::<syn::Token![,]>()?;
let arguments = input.parse()?;
// 后续模式和参数表达式只属于宿主调用,不进入线协议。
input.parse::<syn::Token![,]>()?;
let _ = syn::Pat::parse_single(input)?;
while !input.is_empty() {
input.parse::<syn::Token![,]>()?;
if !input.is_empty() {
let _ = input.parse::<syn::Expr>()?;
}
}
Ok(Self {
name,
method,
input: arguments,
})
}
}
fn source(root: &Path, path: &str) -> syn::File {
syn::parse_file(&std::fs::read_to_string(root.join(path)).unwrap()).unwrap()
}
// SDK 保留 wire 值,不依赖宿主的领域构造器、时区库或价格校验实现。
struct WireTypes;
impl VisitMut for WireTypes {
fn visit_type_mut(&mut self, ty: &mut Type) {
if let Type::Path(path) = ty {
assert!(path.qself.is_none(), "服务合同不支持关联类型");
if path
.path
.segments
.first()
.is_some_and(|segment| segment.ident == "serde_json")
{
syn::visit_mut::visit_type_mut(self, ty);
return;
}
let last = path.path.segments.last().unwrap().clone();
*ty = match last.ident.to_string().as_str() {
"DateTime" | "Tz" | "TokenPrice" | "AdminApiKey" => parse_quote!(String),
"OpaqueProviderData" => parse_quote!(serde_json::Map<String, serde_json::Value>),
_ => Type::Path(syn::TypePath {
qself: None,
path: last.into(),
}),
};
}
syn::visit_mut::visit_type_mut(self, ty);
}
}
fn contract_type(mut item: Item, name: &str) -> Item {
let (attrs, ident, fields, comparable) = match &mut item {
Item::Struct(item) => {
item.vis = parse_quote!(pub);
let comparable = !matches!(
name,
"RuntimeSettings"
| "ReplaceRuntimeSettings"
| "AdminApiKeyMutation"
| "RegeneratedAdminApiKey"
);
(
&mut item.attrs,
&mut item.ident,
Some(&mut item.fields),
comparable,
)
}
Item::Enum(item) => (&mut item.attrs, &mut item.ident, None, true),
_ => panic!("公开合同必须是 struct 或 enum"),
};
*ident = syn::parse_str(name).unwrap();
attrs.retain(|attr| attr.path().is_ident("serde"));
attrs.insert(
0,
if name == "SmartSchedulingConfig" {
parse_quote!(#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)])
} else if comparable {
parse_quote!(#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)])
} else {
parse_quote!(#[derive(Clone, Serialize, Deserialize)])
},
);
if let Some(fields) = fields {
for field in fields {
field.vis = parse_quote!(pub);
}
}
WireTypes.visit_item_mut(&mut item);
item
}
#[test]
fn sdk_settings_contract_matches_host_declarations() {
let root = Path::new(env!("CARGO_MANIFEST_DIR"));
let mut types = Vec::new();
let mut replacement_fields = Vec::new();
for (path, names) in [
(
"src/model/settings.rs",
&[
"RuntimeSettings",
"ReplaceRuntimeSettings",
"AdminApiKeyMutation",
"RegeneratedAdminApiKey",
][..],
),
(
"src/model/mod.rs",
&["MutationActor", "MutationContext"][..],
),
(
"../gateway-core/src/account/location.rs",
&["RequestLocation"][..],
),
(
"../gateway-core/src/account/smart_scheduling.rs",
&["SmartSchedulingValues"][..],
),
(
"../gateway-core/src/metering/pricing.rs",
&["TokenPriceOverride", "ModelPriceOverride"][..],
),
(
"src/model/pricing.rs",
&[
"PricingCatalog",
"PricingSyncPreview",
"SyncPricing",
"PricingChange",
"UpdatePricing",
][..],
),
] {
let items = source(root, path).items;
for name in names {
let item = items
.iter()
.find(|item| match item {
Item::Struct(item) => item.ident == name,
Item::Enum(item) => item.ident == name,
_ => false,
})
.unwrap_or_else(|| panic!("缺少合同类型 {name}"));
if let Item::Struct(item) = item
&& item.ident == "ReplaceRuntimeSettings"
{
replacement_fields = item
.fields
.iter()
.map(|field| field.ident.clone().unwrap())
.collect();
}
types.push(contract_type(
item.clone(),
if *name == "SmartSchedulingValues" {
"SmartSchedulingConfig"
} else {
name
},
));
}
}
let service = source(root, "src/use_case/settings.rs")
.items
.into_iter()
.find_map(|item| match item {
Item::Trait(item) if item.ident == "SettingsService" => Some(item),
_ => None,
})
.unwrap();
let outputs: BTreeMap<_, _> = service
.items
.into_iter()
.filter_map(|item| match item {
syn::TraitItem::Fn(method) => {
let syn::ReturnType::Type(_, output) = method.sig.output else {
panic!("服务必须返回 Result")
};
let Type::Path(result) = *output else {
panic!("服务必须返回 Result")
};
let segment = result.path.segments.last().unwrap();
assert_eq!(segment.ident, "Result");
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
panic!("Result 缺少类型参数")
};
let syn::GenericArgument::Type(mut output) = args.args.first().unwrap().clone()
else {
panic!("Result 缺少输出类型")
};
WireTypes.visit_type_mut(&mut output);
Some((method.sig.ident.to_string(), output))
}
_ => None,
})
.collect();
let register = source(root, "src/service/settings.rs")
.items
.into_iter()
.find_map(|item| match item {
Item::Fn(item) if item.sig.ident == "register" => Some(item),
_ => None,
})
.unwrap();
let operations: Vec<_> = register
.block
.stmts
.into_iter()
.filter_map(|stmt| {
let syn::Stmt::Macro(stmt) = stmt else {
return None;
};
if !stmt.mac.path.is_ident("register") {
return None;
}
let Operation {
name,
method,
mut input,
} = syn::parse2(stmt.mac.tokens).unwrap();
WireTypes.visit_type_mut(&mut input);
let output = outputs.get(&method.to_string()).unwrap();
let operation = format!("settings.{method}");
Some(quote! {
pub struct #name;
impl Operation for #name {
const NAME: &'static str = #operation;
type Input = #input;
type Output = #output;
}
})
})
.collect();
assert_eq!(operations.len(), outputs.len(), "公开服务方法必须完整声明");
let assignments = replacement_fields.iter().map(|field| match field.to_string().as_str() {
"expected_revision" => quote!(expected_revision: settings.config_revision),
"request_profile_updates" => quote!(request_profile_updates: settings.request_profiles.into_iter().map(|(provider, value)| (provider, Some(value))).collect()),
_ => quote!(#field: settings.#field),
});
let generated = quote! {
use super::Operation;
use serde::{Deserialize, Serialize};
use std::{collections::{BTreeMap, BTreeSet}, num::NonZeroU64};
pub type ProviderRequestProfiles = BTreeMap<String, serde_json::Map<String, serde_json::Value>>;
pub type ProviderRequestProfileUpdates = BTreeMap<String, Option<serde_json::Map<String, serde_json::Value>>>;
pub type ModelMappings = BTreeMap<String, String>;
pub type Revision = NonZeroU64;
pub type RotationStrategy = String;
pub type PricingOverrides = BTreeMap<String, BTreeMap<String, ModelPriceOverride>>;
#(#types)*
#(#operations)*
impl From<RuntimeSettings> for ReplaceRuntimeSettings {
fn from(settings: RuntimeSettings) -> Self {
Self { #(#assignments),* }
}
}
};
let output = root.join("../gateway-plugin/sdk/src/call/services/settings.rs");
let expected = prettyplease::unparse(&syn::parse2::<syn::File>(generated).unwrap());
if std::env::var_os("CPR_UPDATE_SERVICE_CONTRACT").is_some() {
std::fs::write(&output, format!("// 从宿主设置类型与 service/settings.rs 生成;更新命令见 SDK 维护说明。\n{expected}")).unwrap();
}
let actual = syn::parse_file(&std::fs::read_to_string(output).unwrap()).unwrap();
assert!(
prettyplease::unparse(&actual) == expected,
"SDK 服务合同已过期:CPR_UPDATE_SERVICE_CONTRACT=1 cargo test -p gateway-admin --test main sdk_settings_contract_matches_host_declarations"
);
}
@@ -24,6 +24,66 @@ struct PricingSettingsStore {
updated: Mutex<Option<UpdatePricing>>,
}
#[tokio::test]
async fn explicit_service_calls_freeze_empty_plan_and_cancel_before_entering_terminal() {
use gateway_admin::service::{Origin, Plan, scope};
use std::sync::atomic::{AtomicUsize, Ordering};
let resolutions = Arc::new(AtomicUsize::new(0));
let source_calls = resolutions.clone();
let store = Arc::new(PricingSettingsStore::default());
let services = super::AdminHarness::new()
.settings(store.clone())
.service_middleware(Arc::new(move || {
source_calls.fetch_add(1, Ordering::SeqCst);
None
}))
.build()
.await;
let context = MutationContext {
actor: gateway_admin::model::MutationActor::System,
request_id: "empty-plan".into(),
};
let command = UpdatePricing {
provider: "xai".into(),
models: vec!["example-model".into()],
change: PricingChange::Reset,
};
for expected in 1..=2 {
services
.public_services()
.call(
"settings.update_pricing",
serde_json::to_value((&context, command.clone())).unwrap(),
)
.await
.unwrap();
// 内部 pricing 不再单独进入插件链;显式调用只解析一次发布集合。
assert_eq!(resolutions.load(Ordering::SeqCst), expected);
}
store.updated.lock().unwrap().take();
let cancellation = gateway_core::lifecycle::CancellationToken::new();
cancellation.cancel();
let error = scope(
Origin {
request_id: "cancelled-native".into(),
call_id: "entry".into(),
cancellation,
extensions: Default::default(),
plan: Plan::Frozen(None),
},
services.public_services().call(
"settings.update_pricing",
serde_json::to_value((context, command)).unwrap(),
),
)
.await
.unwrap_err();
assert_eq!(error.kind, "unavailable");
assert!(store.updated.lock().unwrap().is_none());
assert_eq!(resolutions.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn client_profile_preview_rejects_unknown_provider_before_loading_settings() {
let services = super::AdminHarness::new()
@@ -213,6 +273,7 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
request_id: "request-settings".to_owned(),
},
ReplaceRuntimeSettings {
expected_revision: gateway_admin::model::Revision::new(1).unwrap(),
request_profile_updates: Default::default(),
request_location_enabled: false,
request_location: Default::default(),
+4 -1
View File
@@ -18,6 +18,9 @@ bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["serde"] }
flate2 = "1.1.9"
futures = "0.3.32"
http = "1.5.0"
http-body = "1.0.1"
http-body-util = "0.1.5"
gateway-admin = { path = "../gateway-admin" }
gateway-core = { path = "../gateway-core" }
gateway-protocol = { path = "../gateway-protocol" }
@@ -28,6 +31,7 @@ serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve
sha2 = "0.11.0"
thiserror = "2.0.18"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "sync", "time"] }
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.7.0", features = ["cors", "fs", "request-id", "timeout", "trace"] }
tracing = "0.1.44"
url = "2.5.8"
@@ -39,7 +43,6 @@ async-trait = "0.1.89"
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread", "test-util"] }
tokio-tungstenite = "0.29.0"
tower = { version = "0.5.3", features = ["util"] }
[lints]
workspace = true
@@ -32,6 +32,12 @@ where
type Rejection = AdminError;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
// 仅内部 Dispatcher 写入此上下文;普通客户端 headers 无法声明插件身份。
if let Some(context) = parts.extensions.get::<AdminRequestContext>() {
return Ok(Self {
context: context.clone(),
});
}
let principal = require_admin_auth(state, &parts.headers).await?;
let request_id = admin_request_id(parts).ok_or_else(AdminError::internal)?;
Ok(Self {
@@ -8,7 +8,7 @@ use axum::{
};
use gateway_admin::model::plugins::{
InstalledPluginArtifact, PluginArtifactMetadata, PluginIconTheme, PluginInstallResult,
PluginPermissionDescription, PluginSource, distribution::VerifiedPluginArtifact,
PluginSource, distribution::VerifiedPluginArtifact,
};
use serde::{Deserialize, Serialize};
@@ -23,16 +23,16 @@ use crate::{
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct ArtifactView {
metadata: MetadataView,
metadata: PluginArtifactMetadata,
source: PluginSource,
installed_at: String,
accepted_at: Option<String>,
}
impl ArtifactView {
fn new(value: InstalledPluginArtifact, service: &gateway_admin::PluginsService) -> Self {
fn new(value: InstalledPluginArtifact) -> Self {
Self {
metadata: MetadataView::new(value.metadata, service),
metadata: value.metadata,
source: value.source,
installed_at: value.installed_at.to_rfc3339(),
accepted_at: value.accepted_at.map(|value| value.to_rfc3339()),
@@ -40,37 +40,16 @@ impl ArtifactView {
}
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct MetadataView {
#[serde(flatten)]
metadata: PluginArtifactMetadata,
permission_descriptions: Vec<PluginPermissionDescription>,
}
impl MetadataView {
fn new(metadata: PluginArtifactMetadata, service: &gateway_admin::PluginsService) -> Self {
Self {
permission_descriptions: service
.permission_descriptions(&metadata.requested_permissions),
metadata,
}
}
}
#[derive(Serialize)]
pub(super) struct VerifiedArtifactView {
metadata: MetadataView,
metadata: PluginArtifactMetadata,
source: PluginSource,
}
impl VerifiedArtifactView {
pub(super) fn new(
value: VerifiedPluginArtifact,
service: &gateway_admin::PluginsService,
) -> Self {
pub(super) fn new(value: VerifiedPluginArtifact) -> Self {
Self {
metadata: MetadataView::new(value.metadata, service),
metadata: value.metadata,
source: value.source,
}
}
@@ -86,9 +65,9 @@ pub(super) struct InstallResultView {
}
impl InstallResultView {
pub(super) fn new(value: PluginInstallResult, service: &gateway_admin::PluginsService) -> Self {
pub(super) fn new(value: PluginInstallResult) -> Self {
Self {
artifact: ArtifactView::new(value.mutation.artifact, service),
artifact: ArtifactView::new(value.mutation.artifact),
config_revision: value.mutation.config_revision.get(),
default_instance_id: value.default_instance_id,
configuration_required: value.configuration_required,
@@ -173,7 +152,7 @@ async fn list<S: SessionState + Send + Sync>(
.await
.map_err(map_admin_service_error)?
.into_iter()
.map(|artifact| ArtifactView::new(artifact, state.admin_services().plugins()))
.map(ArtifactView::new)
.collect();
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(items)))
}
@@ -200,10 +179,7 @@ async fn verify_upload<S: SessionState + Send + Sync>(
.map_err(map_admin_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(VerifiedArtifactView::new(
result,
state.admin_services().plugins(),
)),
AdminEnvelope::ok(VerifiedArtifactView::new(result)),
))
}
@@ -228,10 +204,7 @@ async fn upload<S: SessionState + Send + Sync>(
.map_err(map_admin_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(InstallResultView::new(
result,
state.admin_services().plugins(),
)),
AdminEnvelope::ok(InstallResultView::new(result)),
))
}
@@ -254,10 +227,7 @@ async fn accept<S: SessionState + Send + Sync>(
.map_err(map_admin_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(InstallResultView::new(
result,
state.admin_services().plugins(),
)),
AdminEnvelope::ok(InstallResultView::new(result)),
))
}
@@ -149,10 +149,7 @@ async fn verify<S: SessionState + Send + Sync>(
.map_err(map_admin_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(VerifiedArtifactView::new(
result,
state.admin_services().plugins(),
)),
AdminEnvelope::ok(VerifiedArtifactView::new(result)),
))
}
@@ -169,10 +166,7 @@ async fn install<S: SessionState + Send + Sync>(
.map_err(map_admin_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(InstallResultView::new(
result,
state.admin_services().plugins(),
)),
AdminEnvelope::ok(InstallResultView::new(result)),
))
}
@@ -137,7 +137,6 @@ impl From<PluginInstanceView> for InstanceView {
trusted_process: _,
configuration,
secrets,
grants: _,
bindings,
revision,
} = value.instance;
@@ -6,7 +6,7 @@ use axum::{
extract::{
ConnectInfo, DefaultBodyLimit, Extension, Path, Request, State, rejection::BytesRejection,
},
http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header},
http::{HeaderMap, HeaderValue, StatusCode, header},
response::{IntoResponse, Response},
routing::{any, get, post},
};
@@ -24,6 +24,7 @@ use headers::{ETag, HeaderMapExt as _, IfNoneMatch};
use serde::Deserialize;
use sha2::{Digest as _, Sha256};
use crate::middleware::headers::{decode_headers, encode_headers};
use crate::{
admin::{
AdminAuth, AdminEnvelope, AdminError, AdminJson, AdminQuery, AdminResponse,
@@ -99,27 +100,28 @@ async fn model_responses(
state
.admin_services()
.plugin_management()
.authorize_models(&target)
.validate_target(&target)
.await
.map_err(map_admin_service_error)?;
let client_key_id = ClientApiKeyId::new(query.client_key_id)
.map_err(|_| AdminError::bad_request("Client Key 标识不合法"))?;
let headers = model_request_headers(&parts.headers);
// 推理身份由显式 Key 决定;原始 HTTP 头仍完整交给插件,原生上游自行构造认证。
let headers = parts.headers;
let (decoded, middleware_body) =
match decode_request_with_body(&body, &headers, MAXIMUM_MODEL_BODY_BYTES) {
Ok(decoded) => decoded,
Err(error) => {
return Ok(harden_model_response(protocol_error_response(
return Ok(protocol_error_response(
StatusCode::BAD_REQUEST,
error.protocol_body(),
)));
));
}
};
let service = state.openai().clone();
let execution = service.execution();
let prepared = match execution.prepare_plugin_execution(&client_key_id).await {
Ok(prepared) => prepared,
Err(error) => return Ok(harden_model_response(gateway_error_response(&error))),
Err(error) => return Ok(gateway_error_response(&error)),
};
let authorization: Arc<dyn ResponseAuthorization> =
Arc::new(PluginModelsAuthorization { state, target });
@@ -139,12 +141,11 @@ async fn model_responses(
headers,
decoded,
middleware_body,
maximum_body_bytes: MAXIMUM_MODEL_BODY_BYTES,
},
Some(authorization),
)
.await;
Ok(harden_model_response(response))
Ok(response)
}
struct PluginModelsAuthorization {
@@ -158,7 +159,7 @@ impl ResponseAuthorization for PluginModelsAuthorization {
self.state
.admin_services()
.plugin_management()
.authorize_models(&self.target)
.validate_target(&self.target)
.await
.map_err(plugin_model_authorization_error)
})
@@ -180,51 +181,6 @@ fn plugin_model_authorization_error(error: gateway_admin::model::AdminError) ->
GatewayError::new(kind, message)
}
fn model_request_headers(headers: &HeaderMap) -> HeaderMap {
let mut sanitized = HeaderMap::new();
for name in [
header::ACCEPT,
header::CONTENT_ENCODING,
header::CONTENT_TYPE,
header::USER_AGENT,
HeaderName::from_static("cf-connecting-ip"),
HeaderName::from_static("x-forwarded-for"),
HeaderName::from_static("x-real-ip"),
] {
for value in headers.get_all(&name) {
sanitized.append(name.clone(), value.clone());
}
}
sanitized
}
fn harden_model_response(mut response: Response) -> Response {
for name in [
header::SET_COOKIE,
HeaderName::from_static("set-cookie2"),
header::WWW_AUTHENTICATE,
header::PROXY_AUTHENTICATE,
header::AUTHORIZATION,
header::PROXY_AUTHORIZATION,
] {
response.headers_mut().remove(name);
}
let headers = response.headers_mut();
headers.insert(
header::CONTENT_SECURITY_POLICY,
HeaderValue::from_static("default-src 'none'; frame-ancestors 'none'; sandbox"),
);
headers.insert(
header::X_CONTENT_TYPE_OPTIONS,
HeaderValue::from_static("nosniff"),
);
headers.insert(
header::REFERRER_POLICY,
HeaderValue::from_static("no-referrer"),
);
response
}
async fn start_callback<S: SessionState + Send + Sync>(
auth: AdminAuth,
State(state): State<S>,
@@ -271,6 +227,7 @@ async fn callback<S: SessionState + Send + Sync>(
.map(|(_, value)| value.into_owned())
.filter(|value| !value.is_empty())
.ok_or_else(|| AdminError::bad_request("插件登录回调缺少 state"))?;
let headers = encode_headers(request.headers());
if values.next().is_some() || to_bytes(request.into_body(), 0).await.is_err() {
return Err(AdminError::bad_request("插件登录回调 state 重复或包含正文"));
}
@@ -285,6 +242,7 @@ async fn callback<S: SessionState + Send + Sync>(
path: path.path,
query,
content_type: None,
headers,
body: Vec::new(),
request_id: uuid::Uuid::now_v7().to_string(),
},
@@ -340,6 +298,7 @@ async fn handle<S: SessionState + Send + Sync>(
let target = path.target();
let method = request.method().as_str().to_owned();
let query = request.uri().query().unwrap_or_default().to_owned();
let headers = encode_headers(request.headers());
let body = to_bytes(request.into_body(), MAXIMUM_BODY_BYTES)
.await
.map_err(|_| {
@@ -355,6 +314,7 @@ async fn handle<S: SessionState + Send + Sync>(
path: path.path,
query,
content_type,
headers,
body: body.to_vec(),
request_id: auth.context().request_id.clone(),
},
@@ -452,5 +412,7 @@ fn raw_response(result: PluginManagementResponse) -> Result<Response, AdminError
"permissions-policy",
HeaderValue::from_static("camera=(), microphone=(), geolocation=(), payment=()"),
);
// 宿主提供默认值,插件显式字段随后覆盖;保留同名多值,只验证 HTTP 语法。
headers.extend(decode_headers(result.headers).map_err(|_| AdminError::bad_gateway())?);
Ok(response)
}
@@ -37,6 +37,7 @@ pub type ProviderRequestProfileUpdates =
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeSettingsView {
pub config_revision: u64,
pub smart_scheduling_defaults: gateway_core::account::SmartSchedulingConfig,
pub provider_request_profiles: ProviderRequestProfiles,
/// 固定兼容字段;值始终从 provider_request_profiles 派生。
@@ -78,6 +79,7 @@ pub struct RuntimeSettingsView {
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateRuntimeSettingsRequest {
pub config_revision: u64,
#[serde(default)]
pub provider_request_profiles: ProviderRequestProfileUpdates,
/// 兼容既有 wire;与泛化字段冲突时拒绝整个请求。
@@ -213,6 +215,8 @@ impl UpdateRuntimeSettingsRequest {
self.xai_client_profile,
)?;
Ok(ReplaceRuntimeSettings {
expected_revision: gateway_admin::model::Revision::new(self.config_revision)
.map_err(|_| WireValidationError::new("configRevision"))?,
request_profile_updates,
request_location_enabled: self.request_location_enabled,
request_location: self
@@ -264,6 +268,7 @@ impl From<RuntimeSettings> for RuntimeSettingsView {
.map(|(provider, profile)| (provider.as_str().to_owned(), profile.into_inner()))
.collect::<ProviderRequestProfiles>();
Self {
config_revision: settings.config_revision.get(),
openai_client_profile: provider_request_profiles.get("openai").cloned(),
xai_client_profile: provider_request_profiles.get("xai").cloned(),
provider_request_profiles,
+64 -11
View File
@@ -9,7 +9,7 @@ use std::sync::Arc;
use std::time::Duration;
use axum::Router;
use axum::http::{HeaderName, HeaderValue, Method, Request, StatusCode};
use axum::http::{HeaderName, HeaderValue, Method, Request};
use axum::routing::get;
use gateway_admin::AdminServices;
use gateway_core::engine::execution::ExecutionService;
@@ -19,7 +19,6 @@ use serde::Deserialize;
use tower_http::cors::CorsLayer;
use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer};
use tower_http::services::{ServeDir, ServeFile};
use tower_http::timeout::TimeoutLayer;
use tower_http::trace::TraceLayer;
use url::Url;
@@ -32,6 +31,7 @@ pub mod admin;
pub mod auth;
mod health;
mod key_usage;
mod middleware;
pub mod openai;
mod provider;
mod session_cookie;
@@ -112,11 +112,50 @@ pub enum ApiConfigError {
/// 完成组装的唯一 API router。
pub struct ApiBundle {
router: Router,
settings: middleware::SettingsSource,
middleware: Option<middleware::PlanSource>,
timeout: Option<Duration>,
request_id_header: HeaderName,
}
impl ApiBundle {
/// 组合根从请求已冻结的快照解析计划;响应流持有同一代次。
#[must_use]
pub fn with_middleware(
mut self,
source: impl Fn(
Option<&gateway_core::routing::RuntimeSnapshot>,
) -> Option<gateway_core::engine::middleware::FrozenMiddlewarePlan>
+ Send
+ Sync
+ 'static,
) -> Self {
self.middleware = Some(Arc::new(source));
self
}
/// 内部调用使用同一总路由;组合根持有强引用,Runtime 只保存 Weak。
pub fn dispatcher(&self) -> Arc<dyn gateway_core::middleware::http::Dispatcher> {
Arc::new(middleware::RouterDispatcher {
router: middleware::wrap(
self.router.clone(),
self.middleware.clone(),
self.timeout,
self.request_id_header.clone(),
self.settings.clone(),
),
settings: self.settings.clone(),
})
}
pub fn router(self) -> Router {
self.router
middleware::wrap(
self.router,
self.middleware,
self.timeout,
self.request_id_header,
self.settings,
)
}
}
@@ -133,6 +172,17 @@ pub fn initialize(
config.validate().map_err(ApiError::Config)?;
let request_id_header = HeaderName::from_str(&config.request_id_header)
.map_err(|_| ApiError::Config(ApiConfigError::InvalidRequestIdHeader))?;
let settings: middleware::SettingsSource = {
let execution = execution.clone();
let timeout_ms = config
.request_timeout_seconds
.map(|seconds| seconds.saturating_mul(1000));
Arc::new(move || {
execution
.request_settings()
.map(|settings| settings.with_http_timeout(timeout_ms))
})
};
let state = ApiState {
admin,
openai: OpenAiService::new(execution, lifecycle),
@@ -177,12 +227,6 @@ pub fn initialize(
.allow_credentials(true),
);
}
if let Some(seconds) = config.request_timeout_seconds {
router = router.layer(TimeoutLayer::with_status_code(
StatusCode::REQUEST_TIMEOUT,
Duration::from_secs(seconds),
));
}
// Trace 必须在 SetRequestId 内侧,span 才能捕获本服务生成的 request_id;
// 默认 DEBUG span 会被 info 日志过滤器丢弃,这里显式用 info_span。
let trace_layer = TraceLayer::new_for_http().make_span_with({
@@ -211,9 +255,18 @@ pub fn initialize(
let router = router
.layer(PropagateRequestIdLayer::new(request_id_header.clone()))
.layer(trace_layer)
.layer(SetRequestIdLayer::new(request_id_header, MakeRequestUuid))
.layer(SetRequestIdLayer::new(
request_id_header.clone(),
MakeRequestUuid,
))
.with_state(state);
Ok(ApiBundle { router })
Ok(ApiBundle {
settings,
router,
middleware: None,
timeout: config.request_timeout_seconds.map(Duration::from_secs),
request_id_header,
})
}
async fn static_cache_control(mut response: axum::response::Response) -> axum::response::Response {
@@ -0,0 +1,27 @@
//! HTTP HeaderMap 与插件边界值的无损转换。
use axum::http::{HeaderMap, HeaderName, HeaderValue};
use bytes::Bytes;
use gateway_core::engine::middleware::{MiddlewareError, MiddlewareHeader};
/// 完整保留多值与非 UTF-8 header,Runtime 只转换 wire 编码。
pub(crate) fn encode_headers(headers: &HeaderMap) -> Vec<MiddlewareHeader> {
headers
.iter()
.map(|(name, value)| {
MiddlewareHeader::new(name.as_str(), Bytes::copy_from_slice(value.as_bytes()))
})
.collect()
}
pub(crate) fn decode_headers(headers: Vec<MiddlewareHeader>) -> Result<HeaderMap, MiddlewareError> {
let mut result = HeaderMap::new();
for header in headers {
let name = HeaderName::from_bytes(header.name().as_bytes())
.map_err(|_| MiddlewareError::InvalidState)?;
let value =
HeaderValue::from_bytes(header.value()).map_err(|_| MiddlewareError::InvalidState)?;
result.append(name, value);
}
Ok(result)
}
@@ -0,0 +1,359 @@
//! 路由外层只适配 HTTP 值和正文,不解释插件协议或业务策略。
pub(crate) mod headers;
pub(crate) mod websocket;
use std::{
convert::Infallible,
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Duration,
};
use axum::{
Router,
body::Body,
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 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<gateway_core::settings::RequestSettings>; }
pub(crate) fn current_settings() -> Option<gateway_core::settings::RequestSettings> {
SETTINGS.try_with(Clone::clone).ok().flatten()
}
pub(crate) fn current() -> Option<contract::Context> {
CURRENT.try_with(Clone::clone).ok()
}
pub(crate) async fn scope<T>(
context: Option<contract::Context>,
future: impl std::future::Future<Output = T>,
) -> T {
match context {
Some(context) => CURRENT.scope(context, future).await,
None => future.await,
}
}
pub(crate) type SettingsSource =
Arc<dyn Fn() -> Option<gateway_core::settings::RequestSettings> + Send + Sync>;
pub(crate) struct RouterDispatcher {
pub(crate) router: Router,
pub(crate) settings: SettingsSource,
}
impl contract::Dispatcher for RouterDispatcher {
fn request_settings(&self) -> Option<gateway_core::settings::RequestSettings> {
(self.settings)()
}
fn dispatch(
&self,
context: contract::Context,
mut request: contract::Request,
) -> futures::future::BoxFuture<
'static,
Result<contract::Response, gateway_core::engine::middleware::MiddlewareError>,
> {
let router = self.router.clone();
let cancellation = context.cancellation.clone();
if let Some(instance_id) = &context.plugin_instance_id {
request
.extensions_mut()
.insert(gateway_admin::model::auth::AdminRequestContext {
principal: gateway_admin::model::auth::AdminPrincipal::Plugin {
instance_id: instance_id.clone(),
},
request_id: format!(
"plugin:{instance_id}:call:{}:request:{}",
context.call_id, context.request_id
),
});
}
request.extensions_mut().insert(context);
Box::pin(async move {
let response = tokio::select! {
biased;
() = cancellation.cancelled() => return Err(gateway_core::engine::middleware::MiddlewareError::Gateway(gateway_core::error::GatewayError::new(gateway_core::error::GatewayErrorKind::Cancelled, "HTTP child cancelled"))),
result = router.oneshot(request.map(Body::new)) => match result { Ok(response) => response, Err(never) => match never {} },
};
Ok(response.map(|body| body.map_err(|error| Box::new(error) as _).boxed_unsync()))
})
}
}
pub(crate) type PlanSource = Arc<
dyn Fn(Option<&gateway_core::routing::RuntimeSnapshot>) -> Option<FrozenMiddlewarePlan>
+ Send
+ Sync,
>;
pub(crate) fn wrap(
router: Router,
source: Option<PlanSource>,
timeout: Option<Duration>,
request_id_header: HeaderName,
settings: SettingsSource,
) -> Router {
// Axum 的 route layer 在匹配后运行;外层 fallback service 才能让改写后的 URI 重新匹配内层路由。
Router::new().fallback_service(tower::service_fn(move |mut request: Request<Body>| {
let router = router.clone();
let inherited = request.extensions().get::<contract::Context>().is_some();
let configuration = request.extensions_mut().get_or_insert(Settings {
timeout,
runtime: None,
});
if configuration.runtime.is_none() && (source.is_some() || inherited) {
configuration.runtime = settings();
}
let snapshot = configuration
.runtime
.as_ref()
.map(|settings| settings.snapshot());
let plan = request
.extensions()
.get::<contract::Context>()
.and_then(|context| context.plan.clone())
.or_else(|| {
source
.as_ref()
.and_then(|source| source(snapshot.as_deref()))
});
let request_id_header = request_id_header.clone();
async move { Ok::<_, Infallible>(handle(router, request, plan, request_id_header).await) }
}))
}
async fn handle(
router: Router,
mut request: Request<Body>,
plan: Option<FrozenMiddlewarePlan>,
request_id_header: HeaderName,
) -> Response {
let started = Instant::now();
let inherited = request.extensions_mut().remove::<contract::Context>();
if let Some(plan) = plan.as_ref().filter(|plan| plan.has_websocket()) {
request.extensions_mut().insert(plan.clone());
}
if inherited.is_none()
&& !plan
.as_ref()
.is_some_and(|plan| plan.has_http() || plan.has_service())
{
let settings = request
.extensions()
.get::<Settings>()
.and_then(|settings| settings.runtime.clone());
return SETTINGS
.scope(settings, dispatch(router, request, started))
.await;
}
let context = inherited.unwrap_or_else(|| {
let request_id = request
.headers()
.get(&request_id_header)
.and_then(|value| value.to_str().ok())
.map(str::to_owned)
.unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
contract::Context {
plugin_instance_id: None,
call_id: request_id.clone(),
request_id,
parent_call_id: None,
extensions: Default::default(),
cancellation: CancellationToken::new(),
plan: plan.clone(),
}
});
let context = contract::Context {
plan: plan.clone(),
..context
};
if let Ok(value) = HeaderValue::from_str(&context.request_id) {
request.headers_mut().insert(request_id_header, value);
}
let guard = Arc::new(CallLifetime {
cancellation: context.cancellation.clone(),
_plan: plan.clone(),
});
request.extensions_mut().insert(guard.clone());
request.extensions_mut().insert(context.clone());
let service_context = gateway_admin::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()),
};
request
.extensions_mut()
.insert(Arc::new(websocket::Upgrade {
context: context.clone(),
}) as Arc<dyn contract::upgrade::WebSocketUpgrade>);
let request = request.map(|body| body.map_err(|error| Box::new(error) as _).boxed_unsync());
let terminal_context = context.clone();
let next = compose(Vec::new(), move |request: contract::Request| {
Box::pin(async move {
let settings = request
.extensions()
.get::<Settings>()
.and_then(|settings| settings.runtime.clone());
let response = CURRENT
.scope(
terminal_context,
SETTINGS.scope(
settings,
gateway_admin::service::scope(
service_context,
dispatch(router, request.map(Body::new), started),
),
),
)
.await;
Ok(response.map(|body| body.map_err(|error| Box::new(error) as _).boxed_unsync()))
})
});
let result = match plan {
Some(plan) => plan.handle_http(context, request, next).await,
None => next.run(request).await,
};
match result {
Ok(mut response) => {
if let Some(upgraded) = response
.extensions_mut()
.remove::<contract::upgrade::Upgraded>()
{
if response.status() != upgraded.status {
return StatusCode::BAD_GATEWAY.into_response();
}
let Some(task) = upgraded.take_task() else {
return StatusCode::BAD_GATEWAY.into_response();
};
let lifetime = guard.clone();
tokio::spawn(async move {
let cancellation = lifetime.cancellation.clone();
let _lifetime = lifetime;
tokio::select! {
biased;
() = cancellation.cancelled() => {},
_ = task => {},
}
});
}
response.map(|body| Body::new(ResponseBody::new(body, guard)))
}
Err(error) => {
let status = if error.is_rejected() {
StatusCode::FORBIDDEN
} else {
StatusCode::BAD_GATEWAY
};
(status, axum::Json(serde_json::json!({"error": {"message": "HTTP middleware failed", "type": "middleware_error"}}))).into_response()
}
}
}
async fn dispatch(router: Router, request: Request<Body>, started: Instant) -> Response {
let timeout = request
.extensions()
.get::<Settings>()
.and_then(|settings| settings.timeout);
let dispatch = router.oneshot(request);
let result = match timeout {
Some(timeout) => {
let Some(remaining) = timeout.checked_sub(started.elapsed()) else {
return StatusCode::REQUEST_TIMEOUT.into_response();
};
match tokio::time::timeout(remaining, dispatch).await {
Ok(result) => result,
Err(_) => return StatusCode::REQUEST_TIMEOUT.into_response(),
}
}
None => dispatch.await,
};
match result {
Ok(response) => response,
Err(never) => match never {},
}
}
pub(crate) struct CallLifetime {
cancellation: CancellationToken,
_plan: Option<FrozenMiddlewarePlan>,
}
impl Drop for CallLifetime {
fn drop(&mut self) {
self.cancellation.cancel();
}
}
struct ResponseBody {
body: contract::Body,
_lifetime: Arc<CallLifetime>,
cancellation: futures::future::BoxFuture<'static, ()>,
finished: bool,
}
impl ResponseBody {
fn new(body: contract::Body, lifetime: Arc<CallLifetime>) -> Self {
let cancellation = lifetime.cancellation.clone();
Self {
body,
_lifetime: lifetime,
cancellation: Box::pin(async move { cancellation.cancelled().await }),
finished: false,
}
}
}
impl HttpBody for ResponseBody {
type Data = bytes::Bytes;
type Error = Box<dyn std::error::Error + Send + Sync>;
fn poll_frame(
mut self: Pin<&mut Self>,
context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
if self.finished {
return Poll::Ready(None);
}
if self.cancellation.as_mut().poll(context).is_ready() {
self.finished = true;
self.body = contract::empty_body();
return Poll::Ready(Some(Err(Box::new(gateway_core::error::GatewayError::new(
gateway_core::error::GatewayErrorKind::Cancelled,
"HTTP body cancelled",
)))));
}
let frame = Pin::new(&mut self.body).poll_frame(context);
if matches!(frame, Poll::Ready(None)) {
self.finished = true;
}
frame
}
fn is_end_stream(&self) -> bool {
self.finished || self.body.is_end_stream()
}
fn size_hint(&self) -> SizeHint {
self.body.size_hint()
}
}
@@ -0,0 +1,273 @@
//! 服务器握手与双工传输;自定义协议及多轮执行由接管会话的插件编排。
use std::sync::Arc;
use axum::{
body::Body,
extract::{
FromRequestParts as _,
ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade},
},
};
use futures::{
FutureExt as _, SinkExt as _, StreamExt as _,
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 http_body_util::BodyExt as _;
pub(super) struct Upgrade {
pub(super) context: http::Context,
}
impl upgrade::WebSocketUpgrade for Upgrade {
fn accept(
&self,
mut parts: ::http::request::Parts,
protocols: Vec<String>,
) -> BoxFuture<'static, Result<(http::Response, upgrade::PendingSession), MiddlewareError>>
{
let context = self.context.clone();
Box::pin(async move {
let websocket = WebSocketUpgrade::from_request_parts(&mut parts, &())
.await
.map_err(|_| MiddlewareError::InvalidState)?;
let headers = parts
.headers
.iter()
.map(|(name, value)| {
MiddlewareHeader::new(
name.as_str(),
bytes::Bytes::copy_from_slice(value.as_bytes()),
)
})
.collect::<Vec<_>>()
.into();
let (sender, receiver) = tokio::sync::oneshot::channel();
let response = websocket
.protocols(protocols)
.max_message_size(usize::MAX)
.max_frame_size(usize::MAX)
.on_upgrade(move |socket| async move {
let _ = sender.send(socket);
});
let pending = async move {
let socket = tokio::select! {
biased;
() = context.cancellation.cancelled() => return Err(Arc::new(MiddlewareError::Fault)),
socket = receiver => socket.map_err(|_| Arc::new(MiddlewareError::Fault))?,
};
let (writer, reader) = socket.split();
let writer = Arc::new(Writer { socket: tokio::sync::Mutex::new(writer), cancellation: context.cancellation.clone() });
let (incoming, received) = tokio::sync::mpsc::channel(1);
let cancellation = context.cancellation.clone();
let reader_task = tokio::spawn(read_messages(reader, incoming, cancellation));
Ok(Arc::new(Session {
reader: tokio::sync::Mutex::new(received), reader_task: reader_task.abort_handle(), writer,
plan: context.plan, connection_id: context.request_id, headers,
cancellation: context.cancellation,
}) as Arc<dyn core::Session>)
}.boxed().shared();
Ok((
response
.map(|body: Body| body.map_err(|error| Box::new(error) as _).boxed_unsync()),
pending,
))
})
}
}
/// 原生协议与插件接管会话共用消息链;调用取消不等待插件释放 transport 锁。
pub(crate) async fn transform(
context: core::Context,
message: core::Message,
) -> Result<Option<core::Message>, MiddlewareError> {
let Some(plan) = context
.plan
.as_ref()
.filter(|plan| plan.has_websocket())
.cloned()
else {
return Ok(Some(message));
};
let cancellation = context.cancellation.clone();
let next = compose(Vec::new(), |message| Box::pin(async { Ok(Some(message)) }));
tokio::select! {
biased;
() = cancellation.cancelled() => Err(MiddlewareError::Fault),
result = plan.handle_websocket(context, message, next) => result,
}
}
struct Session {
reader: tokio::sync::Mutex<tokio::sync::mpsc::Receiver<Message>>,
reader_task: tokio::task::AbortHandle,
writer: Arc<Writer>,
plan: Option<FrozenMiddlewarePlan>,
connection_id: String,
headers: Arc<[MiddlewareHeader]>,
cancellation: CancellationToken,
}
impl Session {
async fn transform(
&self,
message: core::Message,
direction: core::Direction,
) -> Result<Option<core::Message>, MiddlewareError> {
transform(
core::Context {
plan: self.plan.clone(),
connection_id: self.connection_id.clone(),
direction,
headers: self.headers.clone(),
cancellation: self.cancellation.clone(),
sender: self.writer.clone(),
},
message,
)
.await
}
}
impl core::Session for Session {
fn receive(&self) -> BoxFuture<'_, Result<Option<core::Message>, MiddlewareError>> {
Box::pin(async move {
loop {
let message = tokio::select! {
biased;
() = self.cancellation.cancelled() => return Err(MiddlewareError::Fault),
message = async { self.reader.lock().await.recv().await } => message,
};
let Some(message) = message else {
return Ok(None);
};
let message = into_core(message);
// transport 锁只保护一次读取;插件可以在收包处理中主动发送。
if let Some(message) = self.transform(message, core::Direction::Incoming).await? {
return Ok(Some(message));
}
}
})
}
}
impl core::Sender for Session {
fn send(&self, message: core::Message) -> BoxFuture<'_, Result<(), MiddlewareError>> {
Box::pin(async move {
if let Some(message) = self.transform(message, core::Direction::Outgoing).await? {
self.writer.send(message).await?;
}
Ok(())
})
}
}
struct Writer {
socket: tokio::sync::Mutex<SplitSink<WebSocket, Message>>,
cancellation: CancellationToken,
}
impl core::Sender for Writer {
fn send(&self, message: core::Message) -> BoxFuture<'_, Result<(), MiddlewareError>> {
Box::pin(async move {
let message = from_core(message)?;
tokio::select! {
biased;
() = self.cancellation.cancelled() => Err(MiddlewareError::Fault),
result = async {
let mut socket = self.socket.lock().await;
let mut guard = CancelOnDrop(Some(self.cancellation.clone()));
socket.send(message).await.map_err(|_| MiddlewareError::Fault)?;
guard.0.take();
Ok(())
} => result,
}
})
}
}
pub(crate) struct CancelOnDrop(pub(crate) Option<CancellationToken>);
impl Drop for CancelOnDrop {
fn drop(&mut self) {
if let Some(cancellation) = &self.0 {
cancellation.cancel();
}
}
}
pub(crate) fn into_core(message: Message) -> core::Message {
let (kind, payload) = match message {
Message::Text(data) => (core::Kind::Text, data.into()),
Message::Binary(data) => (core::Kind::Binary, data),
Message::Ping(data) => (core::Kind::Ping, data),
Message::Pong(data) => (core::Kind::Pong, data),
Message::Close(Some(frame)) => (
core::Kind::Close {
code: Some(frame.code),
},
frame.reason.into(),
),
Message::Close(None) => (core::Kind::Close { code: None }, bytes::Bytes::new()),
};
core::Message { kind, payload }
}
pub(crate) fn from_core(message: core::Message) -> Result<Message, MiddlewareError> {
Ok(match message.kind {
core::Kind::Text => Message::Text(
message
.payload
.try_into()
.map_err(|_| MiddlewareError::InvalidState)?,
),
core::Kind::Binary => Message::Binary(message.payload),
core::Kind::Ping if message.payload.len() <= 125 => Message::Ping(message.payload),
core::Kind::Pong if message.payload.len() <= 125 => Message::Pong(message.payload),
core::Kind::Close { code: Some(code) } if message.payload.len() <= 123 => {
Message::Close(Some(CloseFrame {
code,
reason: message
.payload
.try_into()
.map_err(|_| MiddlewareError::InvalidState)?,
}))
}
core::Kind::Close { code: None } if message.payload.is_empty() => Message::Close(None),
_ => return Err(MiddlewareError::InvalidState),
})
}
impl Drop for Session {
fn drop(&mut self) {
self.reader_task.abort();
}
}
async fn read_messages(
mut socket: SplitStream<WebSocket>,
incoming: tokio::sync::mpsc::Sender<Message>,
cancellation: CancellationToken,
) {
let _guard = CancelOnDrop(Some(cancellation.clone()));
loop {
let message = tokio::select! {
biased;
() = cancellation.cancelled() => return,
message = socket.next() => message,
};
let Some(Ok(message)) = message else {
return;
};
tokio::select! {
biased;
() = cancellation.cancelled() => return,
result = incoming.send(message) => if result.is_err() { return; },
}
}
}
@@ -140,8 +140,12 @@ fn client_authentication_request(
.ok_or(ClientApiKeyAuthError::MissingAuthorization)?
.to_str()
.map_err(|_| ClientApiKeyAuthError::MalformedAuthorization)?;
ClientAuthenticationRequest::new(raw.to_owned())
.map_err(|_| ClientApiKeyAuthError::MalformedAuthorization)
let request = ClientAuthenticationRequest::new(raw.to_owned())
.map_err(|_| ClientApiKeyAuthError::MalformedAuthorization)?;
Ok(match crate::middleware::current_settings() {
Some(settings) => request.with_settings(settings),
None => request,
})
}
pub(crate) fn client_access_error_response(error: ClientAccessError) -> Response {
@@ -16,10 +16,10 @@ use gateway_core::{
operation::Operation,
};
use crate::middleware::headers::encode_headers;
use crate::openai::middleware::{
ExpectedBody, HttpMiddlewareInput, PendingExecution, buffered_response, error_response,
into_http_response, invoke_http_middleware, pending_execution_response, request_headers,
request_parts,
ExpectedBody, PendingExecution, RequestInput, buffered_response, error_response,
into_http_response, invoke_request, pending_execution_response, request_parts,
};
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
@@ -33,7 +33,7 @@ use super::{
pub(super) async fn provider_endpoint_response<F>(
service: OpenAiService,
client: AuthenticatedClient,
input: HttpMiddlewareInput,
input: RequestInput,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
decode: F,
@@ -48,7 +48,7 @@ where
};
let request_id = prepared.request_id().clone();
let endpoint = input.endpoint.clone();
let result = invoke_http_middleware(
let result = invoke_request(
execution,
prepared,
input,
@@ -125,7 +125,7 @@ async fn collect_raw_json_response(
Ok(pending_execution_response(
"openai".to_owned(),
parts.status.as_u16(),
request_headers(&parts.headers),
encode_headers(&parts.headers),
MiddlewareFrame::new(bytes, framing, true).with_transformed(transformed),
execution,
))
@@ -16,7 +16,8 @@ use gateway_core::operation::{
use serde_json::Value;
use crate::ApiState;
use crate::openai::middleware::{HttpMiddlewareInput, request_headers};
use crate::middleware::headers::encode_headers;
use crate::openai::middleware::RequestInput;
use crate::openai::{
auth::{authenticate_client, client_access_error_response},
endpoint::provider_endpoint_response,
@@ -83,13 +84,13 @@ async fn handle_image_request(
provider_endpoint_response(
service.clone(),
client,
HttpMiddlewareInput {
RequestInput {
endpoint: endpoint.to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: Some(OperationKind::GenerateImage),
transport: ClientTransport::HttpJson,
model_hint: None,
headers: request_headers(&headers),
headers: encode_headers(&headers),
body,
},
client_ip,
@@ -1,6 +1,6 @@
//! 未提交响应的 HTTP 交付;只验证协议边界,不解释或重算 Core 的业务事实。
use super::request::{decode_headers, request_headers};
use crate::middleware::headers::{decode_headers, encode_headers};
use crate::openai::error::{
gateway_error_from_engine, gateway_error_response, openai_error_response,
};
@@ -18,7 +18,6 @@ use gateway_core::{
},
error::GatewayError,
};
use gateway_protocol::openai::response_header_is_forwardable;
/// 只读目录/用量与协议错误共用的单帧响应;没有 Provider attempt 或计量副作用。
pub(crate) async fn buffered_response(
@@ -37,7 +36,7 @@ pub(crate) async fn buffered_response(
Ok(MiddlewareResponse::new(
protocol.to_owned(),
parts.status.as_u16(),
request_headers(&parts.headers),
encode_headers(&parts.headers),
Box::new(BufferedBody(Some(MiddlewareFrame::new(
bytes, framing, true,
)))),
@@ -221,29 +220,8 @@ fn response_head(
if status.is_informational() {
return Err(MiddlewareError::InvalidState);
}
let mut result = decode_headers(headers)?;
let connection_options = result
.get_all(header::CONNECTION)
.iter()
.filter_map(|value| value.to_str().ok())
.flat_map(|value| value.split(','))
.map(str::trim)
.filter(|name| !name.is_empty())
.map(str::to_ascii_lowercase)
.collect::<Vec<_>>();
// Content-Type 由协议 adapter 生成;其余字段在最终交付边界统一剥离认证与传输信息。
let blocked = result
.keys()
.filter(|name| {
*name != header::CONTENT_TYPE
&& !response_header_is_forwardable(name.as_str(), &connection_options)
})
.cloned()
.collect::<Vec<_>>();
for name in blocked {
result.remove(name);
}
Ok((status, result))
// 原生上游头已在协议终端处理;插件返回后的字段只做 HTTP 类型校验。
Ok((status, decode_headers(headers)?))
}
async fn fail_before_commit(mut body: Box<dyn MiddlewareBody>, error: MiddlewareError) -> Response {
@@ -262,19 +240,23 @@ pub(crate) fn error_response(error: MiddlewareError) -> Response {
MiddlewareError::Provider(error) => {
gateway_error_response(&GatewayError::from_provider(&error))
}
MiddlewareError::Rejected => openai_error_response(
StatusCode::FORBIDDEN,
"Request rejected by middleware",
"invalid_request_error",
"middleware_rejected",
)
.into_response(),
MiddlewareError::Fault | MiddlewareError::InvalidState => openai_error_response(
StatusCode::BAD_GATEWAY,
"Request middleware failed",
"server_error",
"middleware_failed",
)
.into_response(),
MiddlewareError::Rejected | MiddlewareError::Remote { rejected: true, .. } => {
openai_error_response(
StatusCode::FORBIDDEN,
"Request rejected by middleware",
"invalid_request_error",
"middleware_rejected",
)
.into_response()
}
MiddlewareError::Fault | MiddlewareError::Remote { .. } | MiddlewareError::InvalidState => {
openai_error_response(
StatusCode::BAD_GATEWAY,
"Request middleware failed",
"server_error",
"middleware_failed",
)
.into_response()
}
}
}
@@ -5,7 +5,5 @@ mod request;
mod session;
pub(crate) use http::{ExpectedBody, buffered_response, error_response, into_http_response};
pub(crate) use request::{
HttpMiddlewareInput, invoke_http_middleware, query_response, request_headers, request_parts,
};
pub(crate) use request::{RequestInput, invoke_request, query_response, request_parts};
pub(crate) use session::{PendingExecution, finalize_session, pending_execution_response};
@@ -1,10 +1,9 @@
//! 请求身份冻结、单次 next 与调用代次保活;HTTP 编码在 http 子模块中完成。
use crate::middleware::headers::{decode_headers, encode_headers};
use crate::openai::error::gateway_error_response;
use axum::{
http::{HeaderMap, HeaderName, HeaderValue},
response::Response,
};
use axum::{http::HeaderMap, response::Response};
use bytes::Bytes;
use futures::future::BoxFuture;
use gateway_core::{
@@ -12,11 +11,10 @@ use gateway_core::{
execution::{
AuthenticatedClient, ClientTransport, ExecutionService, PreparedRootExecution,
},
extensions::ExtensionCallScope,
middleware::{
FrozenMiddlewarePlan, MiddlewareBody, MiddlewareError, MiddlewareFrame,
MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewareRequest,
MiddlewareResponse, MiddlewareTarget,
MiddlewareHeader, MiddlewareMount, MiddlewareRequest, MiddlewareResponse,
MiddlewareTarget,
},
},
lifecycle::CancellationToken,
@@ -26,7 +24,7 @@ use std::{future::Future, sync::Arc};
use super::http::{ExpectedBody, buffered_response, error_response, into_http_response};
pub(crate) struct HttpMiddlewareInput {
pub(crate) struct RequestInput {
pub endpoint: String,
pub protocol: String,
pub operation: Option<OperationKind>,
@@ -36,7 +34,7 @@ pub(crate) struct HttpMiddlewareInput {
pub body: Bytes,
}
impl HttpMiddlewareInput {
impl RequestInput {
pub(crate) fn query(endpoint: String, model_hint: Option<String>, headers: &HeaderMap) -> Self {
Self {
endpoint,
@@ -44,7 +42,7 @@ impl HttpMiddlewareInput {
operation: None,
transport: ClientTransport::HttpJson,
model_hint,
headers: request_headers(headers),
headers: encode_headers(headers),
body: Bytes::new(),
}
}
@@ -54,7 +52,7 @@ impl HttpMiddlewareInput {
pub(crate) async fn query_response<F, Fut>(
execution: Arc<dyn ExecutionService>,
client: AuthenticatedClient,
input: HttpMiddlewareInput,
input: RequestInput,
terminal: F,
) -> Response
where
@@ -66,7 +64,7 @@ where
Err(error) => return gateway_error_response(&error),
};
let protocol = input.protocol.clone();
let result = invoke_http_middleware(
let result = invoke_request(
execution,
prepared,
input,
@@ -86,7 +84,7 @@ where
}
}
pub(crate) type HttpMiddlewareTerminal = Box<
pub(crate) type RequestTerminal = Box<
dyn FnOnce(
PreparedRootExecution,
MiddlewareRequest,
@@ -94,13 +92,19 @@ pub(crate) type HttpMiddlewareTerminal = Box<
+ Send,
>;
pub(crate) async fn invoke_http_middleware(
pub(crate) async fn invoke_request(
execution: Arc<dyn ExecutionService>,
prepared: PreparedRootExecution,
input: HttpMiddlewareInput,
terminal: HttpMiddlewareTerminal,
mut prepared: PreparedRootExecution,
input: RequestInput,
terminal: RequestTerminal,
) -> Result<MiddlewareResponse, MiddlewareError> {
let plan = execution.middleware_plan(&prepared);
let origin = crate::middleware::current();
if let Some(origin) = &origin {
prepared = prepared.with_extension_scope(origin.extensions.clone());
}
let plan = origin
.and_then(|origin| origin.plan)
.or_else(|| execution.middleware_plan(&prepared));
let lifetime = RequestLifetime {
cancellation: prepared.cancellation(),
_plan: plan.clone(),
@@ -117,10 +121,22 @@ pub(crate) async fn invoke_http_middleware(
model: input.model_hint,
account_id: None,
},
ExtensionCallScope::default(),
prepared.extension_scope(),
);
let request = MiddlewareRequest::new(input.protocol, input.headers, input.body);
let next = Box::new(HttpNext { prepared, terminal });
let mut request = MiddlewareRequest::new(input.protocol, input.headers, input.body);
if plan.is_some() {
request = request.with_settings(prepared.request_settings());
}
let next = gateway_core::middleware::compose(Vec::new(), move |request: MiddlewareRequest| {
Box::pin(async move {
if let Some(settings) = request.settings() {
prepared
.apply_settings(settings)
.map_err(MiddlewareError::Gateway)?;
}
terminal(prepared, request).await
})
});
let response = match plan {
Some(plan) => plan.handle(context, request, next).await?,
None => next.run(request).await?,
@@ -138,20 +154,6 @@ pub(crate) async fn invoke_http_middleware(
Ok(response)
}
struct HttpNext {
prepared: PreparedRootExecution,
terminal: HttpMiddlewareTerminal,
}
impl MiddlewareNext for HttpNext {
fn run(
self: Box<Self>,
request: MiddlewareRequest,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
(self.terminal)(self.prepared, request)
}
}
struct RequestLifetime {
cancellation: CancellationToken,
_plan: Option<FrozenMiddlewarePlan>,
@@ -198,16 +200,6 @@ impl MiddlewareBody for RequestBody {
}
}
/// 保留多值与非 UTF-8 header;是否向插件公开由 Runtime 按授权投影。
pub(crate) fn request_headers(headers: &HeaderMap) -> Vec<MiddlewareHeader> {
headers
.iter()
.map(|(name, value)| {
MiddlewareHeader::new(name.as_str(), Bytes::copy_from_slice(value.as_bytes()))
})
.collect()
}
/// 终点重新解码插件改写的请求;冻结认证身份不从这些 header 再次取得。
pub(crate) fn request_parts(
request: MiddlewareRequest,
@@ -215,15 +207,3 @@ pub(crate) fn request_parts(
let (protocol, headers, body) = request.into_parts();
Ok((protocol, decode_headers(headers)?, body))
}
pub(super) fn decode_headers(headers: Vec<MiddlewareHeader>) -> Result<HeaderMap, MiddlewareError> {
let mut result = HeaderMap::new();
for header in headers {
let name = HeaderName::from_bytes(header.name().as_bytes())
.map_err(|_| MiddlewareError::InvalidState)?;
let value =
HeaderValue::from_bytes(header.value()).map_err(|_| MiddlewareError::InvalidState)?;
result.append(name, value);
}
Ok(result)
}
@@ -15,7 +15,7 @@ use serde_json::{Value, json};
use crate::{
ApiState,
openai::middleware::{self, HttpMiddlewareInput},
openai::middleware::{self, RequestInput},
};
use super::{
@@ -46,7 +46,7 @@ pub(crate) async fn models(
middleware::query_response(
service.execution(),
client,
HttpMiddlewareInput::query(super::router::MODELS_PATH.to_owned(), None, &headers),
RequestInput::query(super::router::MODELS_PATH.to_owned(), None, &headers),
move |client| models_response(state, client, query),
)
.await
@@ -152,7 +152,7 @@ pub(crate) async fn model_detail(
middleware::query_response(
service.execution(),
client,
HttpMiddlewareInput::query(
RequestInput::query(
format!("{}/{model_id}", super::router::MODELS_PATH),
Some(model_id.clone()),
&headers,
@@ -32,10 +32,10 @@ use gateway_protocol::openai::sse::{DONE_SSE_FRAME, response_failed_sse_event_wi
use tokio::time::Instant;
use crate::ApiState;
use crate::middleware::headers::encode_headers;
use crate::openai::middleware::{
ExpectedBody, HttpMiddlewareInput, PendingExecution, error_response, finalize_session,
into_http_response, invoke_http_middleware, pending_execution_response, request_headers,
request_parts,
ExpectedBody, PendingExecution, RequestInput, error_response, finalize_session,
into_http_response, invoke_request, pending_execution_response, request_parts,
};
use crate::openai::service::OpenAiService;
use crate::openai::{
@@ -98,7 +98,6 @@ pub(crate) async fn responses(
headers,
decoded,
middleware_body,
maximum_body_bytes,
},
None,
)
@@ -112,7 +111,6 @@ pub(crate) struct ResponsesHttpRequest {
pub headers: HeaderMap,
pub decoded: DecodedResponsesRequest,
pub middleware_body: Bytes,
pub maximum_body_bytes: usize,
}
pub(crate) async fn execute_prepared_responses(
@@ -127,7 +125,6 @@ pub(crate) async fn execute_prepared_responses(
headers,
decoded,
middleware_body,
maximum_body_bytes,
} = request;
let (client_ip, user_agent) = request_client_context(&headers, peer_address);
let stream = decoded.metadata().stream();
@@ -146,19 +143,19 @@ pub(crate) async fn execute_prepared_responses(
// 中间件正文已经由入口按上限解压;传输编码和旧长度不能与改写后的正文混用。
middleware_headers.remove(CONTENT_ENCODING);
middleware_headers.remove(CONTENT_LENGTH);
let input = HttpMiddlewareInput {
let input = RequestInput {
endpoint: crate::openai::router::RESPONSES_PATH.to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: Some(OperationKind::Generate),
transport,
model_hint,
headers: request_headers(&middleware_headers),
headers: encode_headers(&middleware_headers),
body: middleware_body,
};
let terminal_service = service.clone();
let validation = ResponseValidationFacts::default();
let terminal_validation = validation.clone();
let middleware = invoke_http_middleware(
let middleware = invoke_request(
execution,
prepared,
input,
@@ -171,7 +168,10 @@ pub(crate) async fn execute_prepared_responses(
let decoded = decode_request_with_headers(
&request_body,
&request_headers,
maximum_body_bytes,
prepared
.client()
.snapshot()
.responses_max_decompressed_body_bytes(),
)
.map_err(|_| MiddlewareError::Rejected)?
.with_client_context(client_ip, user_agent)
@@ -704,7 +704,7 @@ async fn guarded_http_response(
Ok(pending_execution_response(
OPENAI_PROTOCOL.to_owned(),
parts.status.as_u16(),
request_headers(&parts.headers),
encode_headers(&parts.headers),
MiddlewareFrame::new(bytes, MiddlewareFraming::RawBytes, true),
execution,
))
@@ -0,0 +1,83 @@
//! 插件消息组合在 transport 锁之外执行,主动发送复用唯一 writer。
use super::{ConnectionWriteError, PumpContext, write_message};
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;
pub(super) struct Writer<S> {
socket: tokio::sync::Mutex<S>,
context: Arc<PumpContext>,
}
impl<S, E> Writer<S>
where
S: Sink<Message, Error = E> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
{
pub(super) fn new(socket: S, context: Arc<PumpContext>) -> Self {
Self {
socket: tokio::sync::Mutex::new(socket),
context,
}
}
pub(super) async fn write(&self, message: Message) -> Result<(), ConnectionWriteError> {
// 预算包含排队与真实写入,取消不等待其他调用释放锁。
tokio::select! {
biased;
() = self.context.cancellation.cancelled() => Err(ConnectionWriteError::Closed),
result = tokio::time::timeout(self.context.config.write_timeout, async {
let mut socket = self.socket.lock().await;
// 发送中断后不能继续复用可能含半写入帧的连接。
let mut in_flight = CancelOnDrop(Some(self.context.cancellation.clone()));
let result = write_message(&mut *socket, message, &self.context).await;
if result.is_ok() { in_flight.0.take(); }
result
}) => result.unwrap_or(Err(ConnectionWriteError::Timeout { timeout: self.context.config.write_timeout })),
}
}
}
impl<S, E> core::Sender for Writer<S>
where
S: Sink<Message, Error = E> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
{
fn send(&self, message: core::Message) -> BoxFuture<'_, Result<(), MiddlewareError>> {
Box::pin(async move {
let closing = matches!(message.kind, core::Kind::Close { .. });
self.write(from_core(message)?)
.await
.map_err(|_| MiddlewareError::Fault)?;
if closing {
self.context.cancellation.cancel();
}
Ok(())
})
}
}
pub(super) async fn transform(
message: Message,
direction: core::Direction,
sender: Arc<dyn core::Sender>,
context: &PumpContext,
) -> Result<Option<Message>, MiddlewareError> {
let Some(plan) = &context.middleware else {
return Ok(Some(message));
};
crate::middleware::websocket::transform(
core::Context {
plan: Some(plan.clone()),
connection_id: context.connection_id.to_string(),
direction,
headers: context.headers.clone(),
cancellation: context.cancellation.clone(),
sender,
},
into_core(message),
)
.await?
.map(from_core)
.transpose()
}
@@ -1,5 +1,7 @@
//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界。
mod middleware;
use std::{
collections::VecDeque,
fmt,
@@ -12,7 +14,11 @@ use std::{
use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
use futures::{Sink, SinkExt, Stream, StreamExt};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader},
lifecycle::CancellationToken,
middleware::websocket as contract,
};
use thiserror::Error;
use tokio::{
sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot},
@@ -145,6 +151,7 @@ pub enum PumpExitReason {
InboundOverload,
ServerClose,
PumpStopped,
MiddlewareFailed,
}
impl PumpExitReason {
@@ -161,6 +168,7 @@ impl PumpExitReason {
Self::InboundOverload => "inbound_overload",
Self::ServerClose => "server_close",
Self::PumpStopped => "pump_stopped",
Self::MiddlewareFailed => "middleware_failed",
}
}
}
@@ -176,10 +184,17 @@ pub enum ConnectionWriteError {
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<Result<(), ConnectionWriteError>>,
acknowledged: oneshot::Sender<Result<WriteOutcome, ConnectionWriteError>>,
}
#[derive(Default)]
@@ -187,7 +202,6 @@ struct ConnectionStats {
command_queue_high_water: AtomicUsize,
command_backpressure_count: AtomicU64,
ping_received_count: AtomicU64,
pong_written_count: AtomicU64,
ping_written_count: AtomicU64,
pong_received_count: AtomicU64,
last_read_ms: AtomicU64,
@@ -243,12 +257,16 @@ impl ResponsesWebSocketConnection {
socket: WebSocket,
connection_id: String,
cancellation: CancellationToken,
middleware: Option<FrozenMiddlewarePlan>,
headers: Arc<[MiddlewareHeader]>,
) -> Self {
spawn_connection(
spawn_with_middleware(
socket,
Arc::<str>::from(connection_id),
cancellation,
ConnectionConfig::PRODUCTION,
middleware,
headers,
)
}
@@ -337,7 +355,7 @@ impl ResponsesWebSocketConnection {
&mut self,
payload: String,
context: WriteContext,
) -> Result<(), ConnectionWriteError> {
) -> Result<WriteOutcome, ConnectionWriteError> {
self.send(Message::Text(payload.into()), context).await
}
@@ -392,7 +410,7 @@ impl ResponsesWebSocketConnection {
&mut self,
message: Message,
context: WriteContext,
) -> Result<(), ConnectionWriteError> {
) -> Result<WriteOutcome, ConnectionWriteError> {
let Some(commands) = self.commands.as_ref().cloned() else {
return Err(ConnectionWriteError::Closed);
};
@@ -423,19 +441,20 @@ impl ResponsesWebSocketConnection {
};
let duration = started_at.elapsed();
match &result {
Ok(()) if phase.is_milestone() => tracing::info!(
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(()) => tracing::debug!(
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(),
"Responses WebSocket frame write succeeded"
?outcome,
"Responses WebSocket frame processed"
),
Err(error) => tracing::info!(
websocket_connection_id = %self.connection_id,
@@ -487,7 +506,6 @@ impl ResponsesWebSocketConnection {
.command_backpressure_count
.load(Ordering::Relaxed),
ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed),
pong_written_count = self.stats.pong_written_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,
@@ -506,13 +524,35 @@ impl Drop for ResponsesWebSocketConnection {
}
}
/// 为一个 socket 启动单 owner pump,并返回协调层句柄。
/// 为自动处理 Ping/Pong 的 WebSocket transport 启动单 owner pump。
pub fn spawn_connection<S, E>(
socket: S,
connection_id: Arc<str>,
cancellation: CancellationToken,
config: ConnectionConfig,
) -> ResponsesWebSocketConnection
where
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin + Send + 'static,
E: fmt::Display + Send + 'static,
{
spawn_with_middleware(
socket,
connection_id,
cancellation,
config,
None,
Arc::from([]),
)
}
fn spawn_with_middleware<S, E>(
socket: S,
connection_id: Arc<str>,
cancellation: CancellationToken,
config: ConnectionConfig,
middleware: Option<FrozenMiddlewarePlan>,
headers: Arc<[MiddlewareHeader]>,
) -> ResponsesWebSocketConnection
where
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin + Send + 'static,
E: fmt::Display + Send + 'static,
@@ -534,6 +574,8 @@ where
Arc::clone(&expired),
Arc::clone(&stats),
config,
middleware,
headers,
));
ResponsesWebSocketConnection {
connection_id,
@@ -552,8 +594,8 @@ where
#[expect(clippy::too_many_arguments)]
async fn run_pump<S, E>(
mut socket: S,
mut commands: mpsc::Receiver<ConnectionCommand>,
socket: S,
commands: mpsc::Receiver<ConnectionCommand>,
incoming: mpsc::Sender<PendingConnectionEvent>,
exited: oneshot::Sender<PumpExitReason>,
connection_id: Arc<str>,
@@ -562,150 +604,206 @@ async fn run_pump<S, E>(
expired: Arc<AtomicBool>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
middleware: Option<FrozenMiddlewarePlan>,
headers: Arc<[MiddlewareHeader]>,
) where
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin,
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + 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<dyn contract::Sender> = 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<FrozenMiddlewarePlan>,
headers: Arc<[MiddlewareHeader]>,
connection_id: Arc<str>,
cancellation: CancellationToken,
opened_at: Instant,
expired: Arc<AtomicBool>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
}
async fn read_connection<S, E>(
socket: &mut S,
incoming: &mpsc::Sender<PendingConnectionEvent>,
sender: Arc<dyn contract::Sender>,
context: &PumpContext,
) -> PumpExitReason
where
S: Stream<Item = Result<Message, E>> + Unpin,
E: fmt::Display,
{
let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER));
let deadline = tokio::time::sleep_until(opened_at + config.max_age);
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<S, E>(
writer: Arc<middleware::Writer<S>>,
mut commands: mpsc::Receiver<ConnectionCommand>,
context: &PumpContext,
) -> PumpExitReason
where
S: Sink<Message, Error = E> + Unpin + Send + 'static,
E: fmt::Display + Send + 'static,
{
let mut heartbeat = tokio::time::interval_at(
opened_at + DOWNSTREAM_PING_INTERVAL,
context.opened_at + DOWNSTREAM_PING_INTERVAL,
DOWNSTREAM_PING_INTERVAL,
);
heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut ping_sequence = 0_u64;
let reason = loop {
tokio::select! {
loop {
let message = tokio::select! {
biased;
() = cancellation.cancelled() => break PumpExitReason::LifecycleShutdown,
() = &mut deadline, if !deadline_elapsed => {
deadline_elapsed = true;
expired.store(true, Ordering::Release);
// 满队列本身已能唤醒空闲协调层;到期标志仍阻止启动下一轮。
// 不能因无法追加 Expired 通知而中断正在收尾的响应。
match emit_incoming(&incoming, &event_slots, ConnectionEvent::Expired) {
Ok(()) | Err(PumpExitReason::InboundOverload) => {}
Err(reason) => break reason,
}
}
() = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown,
command = commands.recv() => {
let Some(command) = command else {
break PumpExitReason::CoordinatorDropped;
};
let ConnectionCommand {
message,
context,
acknowledged,
} = command;
let Some(command) = command else { return PumpExitReason::CoordinatorDropped; };
let ConnectionCommand { message, context: write_context, acknowledged } = command;
let closing = matches!(message, Message::Close(_));
let result = socket.send(message).await.map_err(|error| {
let message = error.to_string();
tracing::debug!(
websocket_connection_id = %connection_id,
request_id = context.request_id.as_deref().unwrap_or(""),
frame_phase = context.phase.as_str(),
error = %message,
"Responses WebSocket pump transport write failed"
);
ConnectionWriteError::Transport { message }
});
let failed = result.is_err();
if !failed {
stats.record_write(opened_at);
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 failed {
break PumpExitReason::WriteError;
}
if closing {
break PumpExitReason::ServerClose;
}
if let Some(reason) = reason { return reason; }
if closing { return PumpExitReason::ServerClose; }
continue;
}
// 官方 Codex WsStream 在独立 pump 中回复 Pong,且不将控制帧交给业务流。
// 这里只为下游链路保活,不设置 Pong deadline,也不延长业务请求预算。
_ = heartbeat.tick() => {
ping_sequence = ping_sequence.wrapping_add(1);
let ping = Message::Ping(ping_sequence.to_be_bytes().to_vec().into());
let result = tokio::select! {
biased;
() = cancellation.cancelled() => break PumpExitReason::LifecycleShutdown,
result = timeout(config.write_timeout, socket.send(ping)) => result,
};
match result {
Ok(Ok(())) => {
stats.ping_written_count.fetch_add(1, Ordering::Relaxed);
stats.record_write(opened_at);
}
Ok(Err(error)) => {
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"Responses WebSocket Ping write failed"
);
break PumpExitReason::WriteError;
}
Err(_) => break PumpExitReason::WriteTimeout,
}
Message::Ping(ping_sequence.to_be_bytes().to_vec().into())
}
message = socket.next() => {
if matches!(&message, Some(Ok(_))) {
stats.record_read(opened_at);
}
match message {
Some(Ok(Message::Text(payload))) => {
if let Err(reason) = emit_incoming(&incoming, &event_slots, ConnectionEvent::Text(payload.to_string())) {
break reason;
}
}
Some(Ok(Message::Binary(_))) => {
if let Err(reason) = emit_incoming(&incoming, &event_slots, ConnectionEvent::Binary) {
break reason;
}
}
Some(Ok(Message::Ping(payload))) => {
stats.ping_received_count.fetch_add(1, Ordering::Relaxed);
match timeout(config.write_timeout, socket.send(Message::Pong(payload))).await {
Ok(Ok(())) => {
stats.pong_written_count.fetch_add(1, Ordering::Relaxed);
stats.record_write(opened_at);
}
Ok(Err(error)) => {
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"Responses WebSocket Pong write failed"
);
break PumpExitReason::WriteError;
}
Err(_) => break PumpExitReason::WriteTimeout,
}
}
Some(Ok(Message::Pong(_))) => {
stats.pong_received_count.fetch_add(1, Ordering::Relaxed);
}
Some(Ok(Message::Close(_))) => break PumpExitReason::ClientClose,
Some(Err(error)) => {
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"Responses WebSocket receive failed"
);
break PumpExitReason::ReadError;
}
None => break PumpExitReason::PeerEof,
}
},
};
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);
}
};
// 退出不能排在待执行请求之后,也不能因业务队列已满而丢失。
let _ = exited.send(reason);
tracing::debug!(
websocket_connection_id = %connection_id,
pump_exit_reason = reason.as_str(),
"Responses WebSocket pump exited"
);
context
.stats
.ping_written_count
.fetch_add(1, Ordering::Relaxed);
}
}
async fn write_message<S, E>(
socket: &mut S,
message: Message,
context: &PumpContext,
) -> Result<(), ConnectionWriteError>
where
S: Sink<Message, Error = E> + 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(
@@ -18,7 +18,7 @@ use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::ProviderSessionState;
use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
use crate::openai::middleware::PendingExecution;
use crate::openai::middleware::{PendingExecution, pending_execution_response};
use crate::openai::responses::validation::{ResponseValidationFacts, ResponsesDeliveryValidator};
use super::{
@@ -92,17 +92,17 @@ pub(super) async fn execution_response(
let headers = body.response_headers().to_vec();
let status = middleware_error_status(&error);
let _ = body.record_client_status(status).await;
Box::new(body).close().await;
if let MiddlewareError::Engine(error) = error {
let response_headers = headers
.iter()
.map(|header| MiddlewareHeader::new(header.name(), header.value().clone()))
.collect();
return Ok(MiddlewareResponse::new(
// 失败也要保留会话到洋葱链返回,否则释放会话会取消仍在等待 next 的插件。
return Ok(pending_execution_response(
"openai".to_owned(),
status,
response_headers,
Box::new(SingleFrameBody(Some(MiddlewareFrame::new(
MiddlewareFrame::new(
Bytes::from(initial_engine_error_event(
&error,
request_id.as_ref(),
@@ -110,9 +110,11 @@ pub(super) async fn execution_response(
)),
MiddlewareFraming::JsonDocument,
true,
)))),
),
body.execution,
));
}
Box::new(body).close().await;
return Err(error);
}
let headers = body
@@ -128,18 +130,6 @@ pub(super) async fn execution_response(
))
}
struct SingleFrameBody(Option<MiddlewareFrame>);
impl MiddlewareBody for SingleFrameBody {
fn next_frame(&mut self) -> BoxFuture<'_, Result<Option<MiddlewareFrame>, MiddlewareError>> {
Box::pin(async { Ok(self.0.take()) })
}
fn close(self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move { drop(self) })
}
}
pub(super) async fn forward_response(
connection: &mut ResponsesWebSocketConnection,
response: MiddlewareResponse,
@@ -272,15 +262,17 @@ pub(super) async fn forward_response(
(true, true) => FramePhase::Terminal,
(true, false) => FramePhase::Data,
};
if connection
match connection
.send_text(message, WriteContext::request(&request_id, phase))
.await
.is_err()
{
detach_body(body);
return ForwardOutcome::Disconnect;
Ok(super::connection::WriteOutcome::Written) => first_frame_written = true,
Ok(super::connection::WriteOutcome::Suppressed) => {}
Err(_) => {
detach_body(body);
return ForwardOutcome::Disconnect;
}
}
first_frame_written = true;
terminal_seen = terminal;
}
}
@@ -369,14 +361,18 @@ fn middleware_gateway_error(error: MiddlewareError) -> GatewayError {
MiddlewareError::Gateway(error) => error,
MiddlewareError::Engine(error) => gateway_error_from_engine(&error),
MiddlewareError::Provider(error) => GatewayError::from_provider(&error),
MiddlewareError::Rejected => GatewayError::new(
GatewayErrorKind::PolicyDenied,
"request middleware rejected the request",
),
MiddlewareError::Fault | MiddlewareError::InvalidState => GatewayError::new(
GatewayErrorKind::Internal,
"request middleware returned an invalid response",
),
MiddlewareError::Rejected | MiddlewareError::Remote { rejected: true, .. } => {
GatewayError::new(
GatewayErrorKind::PolicyDenied,
"request middleware rejected the request",
)
}
MiddlewareError::Fault | MiddlewareError::Remote { .. } | MiddlewareError::InvalidState => {
GatewayError::new(
GatewayErrorKind::Internal,
"request middleware returned an invalid response",
)
}
}
}
@@ -393,8 +389,10 @@ fn middleware_error_status(error: &MiddlewareError) -> u16 {
.0
.as_u16()
}
MiddlewareError::Rejected => StatusCode::FORBIDDEN.as_u16(),
MiddlewareError::Fault | MiddlewareError::InvalidState => {
MiddlewareError::Rejected | MiddlewareError::Remote { rejected: true, .. } => {
StatusCode::FORBIDDEN.as_u16()
}
MiddlewareError::Fault | MiddlewareError::Remote { .. } | MiddlewareError::InvalidState => {
StatusCode::INTERNAL_SERVER_ERROR.as_u16()
}
}
@@ -23,18 +23,16 @@ use gateway_core::{
diagnostics::{TraceContext, diagnostic_json},
engine::{
execution::{AuthenticatedClient, ClientTransport},
middleware::MiddlewareError,
middleware::{FrozenMiddlewarePlan, MiddlewareError},
},
lifecycle::{ConnectionGuard, ConnectionLifecycle},
operation::OperationKind,
};
use crate::middleware::headers::encode_headers;
use crate::{
ApiState,
openai::middleware::{
HttpMiddlewareInput, invoke_http_middleware, request_headers as middleware_request_headers,
request_parts,
},
openai::middleware::{RequestInput, invoke_request, request_parts},
openai::{
auth::{authenticate_client, client_access_error_response},
error::runtime_unavailable_response,
@@ -61,6 +59,8 @@ const CONNECTION_LIMIT_CLOSE_REASON: &str = "Responses websocket connection limi
pub(crate) async fn responses_websocket(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
middleware: Option<Extension<FrozenMiddlewarePlan>>,
lifetime: Option<Extension<Arc<crate::middleware::CallLifetime>>>,
headers: HeaderMap,
websocket: WebSocketUpgrade,
) -> Response {
@@ -81,6 +81,8 @@ pub(crate) async fn responses_websocket(
user_agent,
request_headers,
headers,
middleware.map(|Extension(plan)| plan),
lifetime.map(|Extension(lifetime)| lifetime),
)
}
@@ -97,6 +99,7 @@ impl ResponsesWebSocketAdapter {
Self { service }
}
#[expect(clippy::too_many_arguments)]
fn upgrade_with_client_context(
&self,
websocket: WebSocketUpgrade,
@@ -105,6 +108,8 @@ impl ResponsesWebSocketAdapter {
user_agent: Option<String>,
request_headers: OpenAiRequestHeaders,
raw_headers: HeaderMap,
middleware: Option<FrozenMiddlewarePlan>,
lifetime: Option<Arc<crate::middleware::CallLifetime>>,
) -> Response {
let connection_guard = match self.service.try_register_connection() {
Ok(guard) => guard,
@@ -128,14 +133,17 @@ impl ResponsesWebSocketAdapter {
raw_headers,
lifecycle: self.service.lifecycle(),
connection_guard,
middleware,
};
let origin = crate::middleware::current();
websocket
// 覆盖 axum/tungstenite 的私有 64 MiB message 与 16 MiB frame 默认值。
// Responses JSON 的协议可接受性由上游决定,代理不另设 wire 长度上限。
.max_message_size(usize::MAX)
.max_frame_size(usize::MAX)
.on_upgrade(move |socket| async move {
serve_responses_websocket(socket, session).await;
let _lifetime = lifetime;
crate::middleware::scope(origin, serve_responses_websocket(socket, session)).await;
})
}
}
@@ -150,6 +158,7 @@ struct ResponsesWebSocketSession {
raw_headers: HeaderMap,
lifecycle: Arc<dyn ConnectionLifecycle>,
connection_guard: Box<dyn ConnectionGuard>,
middleware: Option<FrozenMiddlewarePlan>,
}
async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocketSession) {
@@ -163,6 +172,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
raw_headers,
lifecycle,
connection_guard,
middleware,
} = session;
tracing::info!(
websocket_connection_id = %connection_id,
@@ -172,7 +182,13 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
let cancellation = lifecycle.cancellation();
let request_headers =
request_headers.with_downstream_websocket_connection_id(connection_id.clone());
let mut connection = ResponsesWebSocketConnection::new(socket, connection_id, cancellation);
let mut connection = ResponsesWebSocketConnection::new(
socket,
connection_id,
cancellation,
middleware,
encode_headers(&raw_headers).into(),
);
let mut request_count = 0_u64;
let mut replay = ConnectionReplaySnapshot::default();
@@ -235,7 +251,27 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
break;
}
let execution = service.execution();
let prepared = match execution.prepare_execution(client.clone()).await {
let preparation = async {
// 握手的显式改写可继承,Key 策略与宿主设置仍在每轮执行前刷新。
let mut client = client.clone();
if let Some(settings) = client.request_settings() {
let current = execution.request_settings().ok_or_else(|| {
gateway_core::error::GatewayError::new(
gateway_core::error::GatewayErrorKind::Internal,
"runtime snapshot is unavailable",
)
})?;
let rebased = settings.rebase(current.snapshot()).map_err(|_| {
gateway_core::error::GatewayError::new(
gateway_core::error::GatewayErrorKind::InvalidRequest,
"request settings are invalid",
)
})?;
client = client.with_request_settings(rebased);
}
execution.prepare_execution(client).await
};
let prepared = match preparation.await {
Ok(prepared) => prepared,
Err(error) => {
trace_rejected_request(
@@ -257,13 +293,13 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
let response_control = prepared.response_control();
let capture = new_replay_capture();
let validation = ResponseValidationFacts::default();
let input = HttpMiddlewareInput {
let input = RequestInput {
endpoint: crate::openai::router::RESPONSES_PATH.to_owned(),
protocol: "openai".to_owned(),
operation: Some(OperationKind::Generate),
transport: ClientTransport::WebSocket,
model_hint: Some(decoded.metadata().requested_model().to_owned()),
headers: middleware_request_headers(&raw_headers),
headers: encode_headers(&raw_headers),
body: Bytes::from(payload.clone()),
};
let service_for_terminal = service.clone();
@@ -273,7 +309,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
let request_id_for_terminal = Arc::clone(&request_id);
let capture_for_terminal = Arc::clone(&capture);
let validation_for_terminal = validation.clone();
let invoke = invoke_http_middleware(
let invoke = invoke_request(
execution,
prepared,
input,
@@ -13,7 +13,8 @@ use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{Operation, OperationKind, RawJsonPayload, StandaloneSearchRequest};
use crate::ApiState;
use crate::openai::middleware::{HttpMiddlewareInput, request_headers};
use crate::middleware::headers::encode_headers;
use crate::openai::middleware::RequestInput;
use crate::openai::{
auth::{authenticate_client, client_access_error_response},
endpoint::provider_endpoint_response,
@@ -42,13 +43,13 @@ pub(crate) async fn standalone_search(
provider_endpoint_response(
service.clone(),
client,
HttpMiddlewareInput {
RequestInput {
endpoint: SEARCH_PATH.to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: Some(OperationKind::Search),
transport: ClientTransport::HttpJson,
model_hint: None,
headers: request_headers(&headers),
headers: encode_headers(&headers),
body,
},
client_ip,
@@ -18,7 +18,7 @@ use serde_json::{Value, json};
use crate::{
ApiState,
auth::SessionState as _,
openai::middleware::{self as plugin_middleware, HttpMiddlewareInput},
openai::middleware::{self as plugin_middleware, RequestInput},
};
use super::{
@@ -48,7 +48,7 @@ async fn usage(
plugin_middleware::query_response(
state.openai().execution(),
client,
HttpMiddlewareInput::query(super::router::USAGE_PATH.to_owned(), None, &headers),
RequestInput::query(super::router::USAGE_PATH.to_owned(), None, &headers),
move |client| usage_response(state, client, uri),
)
.await
@@ -184,6 +184,7 @@ impl AdminTestFixture {
ClientConfig::default(),
stores,
gateway_admin::AdminRuntimePorts {
service_middleware: std::sync::Arc::new(|| None),
plugin_preparation: plugin_ports.clone(),
plugin_management: plugin_ports.clone(),
published_snapshot: published_snapshot.clone(),
@@ -535,6 +536,13 @@ impl SettingsStore for MemorySettingsStore {
_: &MutationContext,
) -> AdminStoreResult<RuntimeSettings> {
let mut settings = self.settings.lock().expect("settings");
if command.expected_revision != settings.config_revision {
return Err(AdminStoreError::new(
AdminStoreErrorKind::Conflict,
"runtime settings",
"settings revision changed",
));
}
let mut request_profiles = settings.request_profiles.clone();
for (provider, profile) in command.request_profile_updates {
if let Some(profile) = profile {
@@ -120,10 +120,19 @@ impl MiddlewarePlan for ModelMiddleware {
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
_: Box<dyn MiddlewareNext>,
_: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
for header in request.headers() {
assert!(!["cookie", "authorization", "x-api-key"].contains(&header.name()));
for (name, value) in [
("cookie", "cpr_session=management-session".to_owned()),
("authorization", "Bearer admin-fixture".to_owned()),
("x-api-key", format!("admin-{}", "a".repeat(64))),
] {
assert!(
request
.headers()
.iter()
.any(|header| { header.name() == name && header.value() == value.as_bytes() })
);
}
assert!(
request
@@ -158,6 +167,7 @@ impl MiddlewarePlan for ModelMiddleware {
vec![
MiddlewareHeader::new("set-cookie", "fixture=value".into()),
MiddlewareHeader::new("www-authenticate", "Bearer".into()),
MiddlewareHeader::new("x-plugin-response", "complete".into()),
],
Box::new(ModelBody {
frame: Some(MiddlewareFrame::new(bytes.into(), framing, terminal)),
@@ -204,7 +214,61 @@ fn model_request(stream: bool) -> Request<Body> {
}
#[tokio::test]
async fn model_bridge_executes_and_protects_request_and_response_headers() {
async fn management_calls_and_callbacks_preserve_headers_and_override_defaults() {
let fixture = published_fixture().await;
let target = format!("plugin-instance-diagnostic/{}/7", "a".repeat(64));
for uri in [
format!("/api/admin/plugins/extensions/{target}/api/echo"),
format!("/plugins/callbacks/{target}/oauth?state=nonce"),
] {
let response = crate::openai::api_router_with_admin_and_execution(
fixture.services.clone(),
ModelExecution::new(ModelReply::Complete),
)
.oneshot(
Request::get(uri)
.header("cookie", "cpr_session=management-session")
.header("authorization", "Bearer management-fixture")
.header(
"x-plugin-binary",
axum::http::HeaderValue::from_bytes(&[0xff]).unwrap(),
)
.header("set-cookie", "first=1")
.header("set-cookie", "second=2")
.header("content-security-policy", "default-src 'self'")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers()["cookie"],
"cpr_session=management-session"
);
assert_eq!(
response.headers()["authorization"],
"Bearer management-fixture"
);
assert_eq!(response.headers()["x-plugin-binary"].as_bytes(), &[0xff]);
assert_eq!(
response
.headers()
.get_all("set-cookie")
.iter()
.map(|value| value.to_str().unwrap())
.collect::<Vec<_>>(),
["first=1", "second=2"]
);
assert_eq!(
response.headers()["content-security-policy"],
"default-src 'self'"
);
}
}
#[tokio::test]
async fn model_bridge_preserves_complete_request_and_response_headers() {
let fixture = published_fixture().await;
let execution = ModelExecution::new(ModelReply::Complete);
let response = crate::openai::api_router_with_admin_and_execution(fixture.services, execution)
@@ -212,10 +276,9 @@ async fn model_bridge_executes_and_protects_request_and_response_headers() {
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert!(!response.headers().contains_key("set-cookie"));
assert!(!response.headers().contains_key("www-authenticate"));
assert_eq!(response.headers()["x-content-type-options"], "nosniff");
assert_eq!(response.headers()["referrer-policy"], "no-referrer");
assert_eq!(response.headers()["set-cookie"], "fixture=value");
assert_eq!(response.headers()["www-authenticate"], "Bearer");
assert_eq!(response.headers()["x-plugin-response"], "complete");
assert_eq!(response.headers()["cache-control"], "no-store");
let body = axum::body::to_bytes(response.into_body(), 4096)
.await
@@ -290,6 +353,7 @@ pub(super) fn resource_fixture(
path: &str,
) -> gateway_admin::model::plugins::management::PluginManagementResponse {
gateway_admin::model::plugins::management::PluginManagementResponse {
headers: Vec::new(),
status: if path == "missing" { 404 } else { 200 },
content_type: if path.ends_with(".css") {
"text/css"
@@ -303,11 +367,10 @@ pub(super) fn resource_fixture(
async fn published_fixture() -> AdminTestFixture {
use gateway_core::{
account::{AccountSelectionPolicy, RotationStrategy},
routing::{ConfigRevision, RuntimeSnapshot},
runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
};
use std::{num::NonZeroU32, sync::Arc, time::Duration};
use std::sync::Arc;
struct Lease;
impl ExtensionSetLease for Lease {
fn is_ready(&self) -> bool {
@@ -322,10 +385,13 @@ async fn published_fixture() -> AdminTestFixture {
fixture.published_snapshot.publish(
RuntimeSnapshot::new(
ConfigRevision::new(7).unwrap(),
AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(1).unwrap(),
Duration::ZERO,
gateway_core::settings::SettingsValues::new(
1,
0,
"smart",
Default::default(),
None,
None,
),
vec![],
vec![],
@@ -510,6 +576,8 @@ async fn model_responses_requires_the_exact_authorized_target_before_execution()
"/api/admin/plugins/extensions/{target}/models/responses?{query}"
))
.header("cookie", "cpr_session=management-session")
.header("authorization", "Bearer admin-fixture")
.header("x-api-key", format!("admin-{}", "a".repeat(64)))
.header("content-type", "application/json")
.header("x-request-id", "management-model-request")
.body(Body::from(
@@ -61,7 +61,7 @@ fn artifact_metadata_wire_exposes_v2_identity_and_contributions() {
output_formats: vec!["openai".into()],
},
)]),
requested_permissions: vec![],
configuration_schema: serde_json::json!({}),
secret_fields: vec![],
state_namespaces: vec![],
@@ -414,13 +414,15 @@ async fn artifact_acceptance_returns_the_flat_install_contract() {
assert_eq!(result["defaultInstanceId"], serde_json::Value::Null);
assert_eq!(result["configurationRequired"], false);
assert_eq!(result["artifact"]["metadata"]["sha256"], "a".repeat(64));
assert_eq!(
result["artifact"]["metadata"]["permissionDescriptions"][0]["permission"],
"models"
assert!(
result["artifact"]["metadata"]
.get("permissionDescriptions")
.is_none()
);
assert_eq!(
result["artifact"]["metadata"]["permissionDescriptions"][0]["label"],
"当前宿主标签"
assert!(
result["artifact"]["metadata"]
.get("requestedPermissions")
.is_none()
);
assert!(result["artifact"]["acceptedAt"].is_string());
assert!(result.get("mutation").is_none());
@@ -540,7 +542,7 @@ pub(super) struct TestPluginPorts {
#[async_trait]
impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts {
async fn authorize_models(
async fn validate_target(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
@@ -563,10 +565,12 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: &str,
_: gateway_admin::model::plugins::management::PluginManagementRequest,
request: gateway_admin::model::plugins::management::PluginManagementRequest,
) -> Result<gateway_admin::model::plugins::management::PluginManagementResponse, AdminError>
{
Ok(management::resource_fixture("ui/index.html"))
let mut response = management::resource_fixture("ui/index.html");
response.headers = request.headers;
Ok(response)
}
async fn views(
&self,
@@ -592,10 +596,12 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: gateway_admin::model::plugins::management::PluginManagementRequest,
request: gateway_admin::model::plugins::management::PluginManagementRequest,
) -> Result<gateway_admin::model::plugins::management::PluginManagementResponse, AdminError>
{
Ok(management::resource_fixture("api/echo"))
let mut response = management::resource_fixture("api/echo");
response.headers = request.headers;
Ok(response)
}
}
@@ -677,7 +683,7 @@ fn icon_artifact(digest: &str) -> InspectedPluginArtifact {
icon: (digest.as_bytes()[0] == b'f')
.then(|| PluginArtifactIcon::Path("assets/icon.png".into())),
contributes: BTreeMap::new(),
requested_permissions: Vec::new(),
configuration_schema: serde_json::json!({}),
secret_fields: Vec::new(),
state_namespaces: Vec::new(),
@@ -687,8 +693,7 @@ fn icon_artifact(digest: &str) -> InspectedPluginArtifact {
}
fn accepted_artifact() -> InstalledPluginArtifact {
let mut metadata = icon_artifact(&"a".repeat(64)).metadata;
metadata.requested_permissions = vec!["models".into()];
let metadata = icon_artifact(&"a".repeat(64)).metadata;
InstalledPluginArtifact {
metadata,
source: PluginSource::Upload,
@@ -802,7 +807,7 @@ impl PluginStore for TestPluginPorts {
trusted_process: true,
configuration: serde_json::json!({}),
secrets: BTreeMap::new(),
grants: Vec::new(),
bindings: Vec::new(),
revision: Revision::new(7).unwrap(),
}]
@@ -895,22 +900,6 @@ impl PluginStore for TestPluginPorts {
#[async_trait]
impl PluginPackageInspector for TestPluginPorts {
fn permission_descriptions(
&self,
permissions: &[String],
) -> Vec<gateway_admin::model::plugins::PluginPermissionDescription> {
permissions
.iter()
.map(
|permission| gateway_admin::model::plugins::PluginPermissionDescription {
permission: permission.clone(),
label: "当前宿主标签".into(),
description: "当前宿主说明".into(),
},
)
.collect()
}
async fn inspect(
&self,
_: Arc<[u8]>,
@@ -37,6 +37,7 @@ async fn response_json(response: axum::response::Response) -> Value {
fn update_body() -> Value {
json!({
"configRevision": 7,
"requestLocationEnabled": false,
"requestLocation": {"country":"US", "region":"Ohio", "city":"Piketon", "timezone":"America/New_York"},
"modelMappings": {
@@ -247,6 +248,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
assert_eq!(
value,
json!({
"configRevision": 7,
"providerRequestProfiles": {},
"openaiClientProfile": null,
"xaiClientProfile": null,
@@ -420,7 +422,7 @@ async fn settings_post_should_replace_global_model_mappings() {
.expect("settings update response");
let data = response_json(response).await["data"].clone();
assert!(data.get("configRevision").is_none());
assert_eq!(data["configRevision"], 8);
assert_eq!(data["modelMappings"]["gpt-5.4"], "gpt-5.5");
assert_eq!(data["modelMappings"]["grok-latest"], "grok-4.5");
}
@@ -608,12 +610,12 @@ async fn request_location_should_normalize_toggle_and_round_trip() {
.oneshot(request(Method::GET, "/api/admin/settings", None))
.await
.unwrap();
assert_eq!(
response_json(response).await["data"]["requestLocation"],
expected
);
let data = response_json(response).await["data"].clone();
assert_eq!(data["requestLocation"], expected);
let mut revision = data["configRevision"].clone();
for enabled in [false, true] {
let mut body = update_body();
body["configRevision"] = revision.clone();
body["requestLocationEnabled"] = json!(enabled);
body["requestLocation"] = expected.clone();
let response = app(fixture.state())
@@ -630,6 +632,7 @@ async fn request_location_should_normalize_toggle_and_round_trip() {
.await
.unwrap();
let data = response_json(response).await["data"].clone();
revision = data["configRevision"].clone();
assert_eq!(data["requestLocationEnabled"], json!(enabled));
assert_eq!(data["requestLocation"], expected);
}
@@ -1215,3 +1218,53 @@ async fn pricing_endpoints_require_administrator_authentication() {
);
}
}
#[tokio::test]
async fn settings_update_rejects_a_stale_version_without_replacing_the_saved_value() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let router = app(fixture.state());
let first = router
.clone()
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(update_body()),
))
.await
.unwrap();
assert_eq!(first.status(), StatusCode::OK);
let first = response_json(first).await["data"].clone();
let mut stale = update_body();
stale["refreshMarginSeconds"] = json!(9999);
let conflict = router
.clone()
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(stale.clone()),
))
.await
.unwrap();
assert_eq!(conflict.status(), StatusCode::CONFLICT);
let current = router
.clone()
.oneshot(request(Method::GET, "/api/admin/settings", None))
.await
.unwrap();
assert_eq!(response_json(current).await["data"], first);
stale["configRevision"] = first["configRevision"].clone();
let retry = router
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(stale),
))
.await
.unwrap();
assert_eq!(retry.status(), StatusCode::OK);
assert_eq!(
response_json(retry).await["data"]["refreshMarginSeconds"],
9999
);
}
@@ -66,6 +66,9 @@ fn source_tree_should_match_frozen_machine_manifest() {
"src/key_usage/query.rs",
"src/key_usage/presenter.rs",
"src/lib.rs",
"src/middleware/headers.rs",
"src/middleware/mod.rs",
"src/middleware/websocket/mod.rs",
"src/openai/middleware/http.rs",
"src/openai/middleware/mod.rs",
"src/openai/middleware/request.rs",
@@ -83,7 +86,8 @@ fn source_tree_should_match_frozen_machine_manifest() {
"src/openai/responses/request.rs",
"src/openai/responses/response.rs",
"src/openai/responses/validation.rs",
"src/openai/responses/websocket/connection.rs",
"src/openai/responses/websocket/connection/mod.rs",
"src/openai/responses/websocket/connection/middleware/mod.rs",
"src/openai/responses/websocket/forward.rs",
"src/openai/responses/websocket/mod.rs",
"src/openai/responses/websocket/protocol.rs",
@@ -131,6 +135,8 @@ fn test_tree_should_match_frozen_rust_mirror() {
"tests/key_usage/mod.rs",
"tests/key_usage/fixtures.rs",
"tests/main.rs",
"tests/middleware/mod.rs",
"tests/middleware/websocket/mod.rs",
"tests/openai/middleware/mod.rs",
"tests/openai/auth.rs",
"tests/openai/error.rs",
@@ -141,7 +147,8 @@ fn test_tree_should_match_frozen_rust_mirror() {
"tests/openai/responses/http.rs",
"tests/openai/responses/mod.rs",
"tests/openai/responses/request.rs",
"tests/openai/responses/websocket/connection.rs",
"tests/openai/responses/websocket/connection/mod.rs",
"tests/openai/responses/websocket/connection/middleware/mod.rs",
"tests/openai/responses/websocket/forward.rs",
"tests/openai/responses/websocket/interrupt.rs",
"tests/openai/responses/websocket/mod.rs",
+1
View File
@@ -3,6 +3,7 @@ mod architecture;
mod auth;
mod health;
mod key_usage;
mod middleware;
mod openai;
mod provider;
mod support;
@@ -0,0 +1,620 @@
mod websocket;
use std::{
sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use axum::{
body::Body,
http::{Request, StatusCode},
};
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 http_body::Frame;
use http_body_util::{BodyExt as _, StreamBody};
use tower::ServiceExt as _;
#[derive(Clone, Copy, Debug)]
enum Mode {
Rewrite,
Upgrade {
reject: bool,
idle: bool,
block_messages: bool,
},
Service,
Echo,
Slow {
clear_timeout: bool,
},
}
#[derive(Debug)]
struct Plan {
mode: Mode,
calls: Arc<AtomicUsize>,
cancellation: Arc<Mutex<Option<CancellationToken>>>,
message_lifetime: Arc<(tokio::sync::Notify, tokio::sync::Notify)>,
}
struct Lease;
impl ExtensionSetLease for Lease {
fn is_ready(&self) -> bool {
true
}
}
impl MiddlewarePlan for Plan {
fn has_http(&self) -> bool {
!matches!(self.mode, Mode::Service)
}
fn has_service(&self) -> bool {
matches!(self.mode, Mode::Service)
}
fn has_websocket(&self) -> bool {
matches!(self.mode, Mode::Upgrade { .. })
}
fn handle_websocket(
&self,
context: gateway_core::middleware::websocket::Context,
message: gateway_core::middleware::websocket::Message,
next: gateway_core::middleware::websocket::Next,
) -> BoxFuture<
'static,
Result<Option<gateway_core::middleware::websocket::Message>, MiddlewareError>,
> {
assert!(context.plan.is_some());
self.calls.fetch_add(1, Ordering::SeqCst);
if !matches!(
self.mode,
Mode::Upgrade {
block_messages: true,
..
}
) {
return next.run(message);
}
let lifetime = self.message_lifetime.clone();
Box::pin(async move {
struct NotifyOnDrop(Arc<(tokio::sync::Notify, tokio::sync::Notify)>);
impl Drop for NotifyOnDrop {
fn drop(&mut self) {
self.0.1.notify_one();
}
}
let _guard = NotifyOnDrop(lifetime.clone());
lifetime.0.notify_one();
std::future::pending().await
})
}
fn handle_service(
&self,
context: gateway_core::middleware::service::Context,
input: serde_json::Value,
next: gateway_core::middleware::service::Next,
) -> BoxFuture<'static, Result<serde_json::Value, gateway_core::middleware::service::Error>>
{
self.calls.fetch_add(1, Ordering::SeqCst);
assert_eq!(context.operation, "settings.load");
assert_eq!(context.request_id, "service-request");
assert_eq!(context.parent_call_id.as_deref(), Some("service-request"));
assert!(input.is_null());
Box::pin(async move {
let mut result = next.run(input).await?;
result["request_interval_ms"] = serde_json::json!(4321);
Ok(result)
})
}
fn handle_http(
&self,
context: contract::Context,
mut request: contract::Request,
next: contract::Next,
) -> BoxFuture<'static, Result<contract::Response, MiddlewareError>> {
if matches!(self.mode, Mode::Service) {
return next.run(request);
}
self.calls.fetch_add(1, Ordering::SeqCst);
*self.cancellation.lock().unwrap() = Some(context.cancellation);
let mode = self.mode;
Box::pin(async move {
match mode {
Mode::Service => unreachable!("service-only plan bypasses HTTP layers"),
Mode::Upgrade { reject, idle, .. } => {
let upgrade = request
.extensions_mut()
.remove::<Arc<dyn contract::upgrade::WebSocketUpgrade>>()
.unwrap();
let (parts, _) = request.into_parts();
let (mut response, session) =
upgrade.accept(parts, vec!["test-session".into()]).await?;
let status = response.status();
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};
let receive = session.receive();
tokio::pin!(receive);
tokio::select! {
message = &mut receive => panic!("must await the client: {}", message.is_ok()),
() = tokio::time::sleep(Duration::from_millis(30)) => {},
}
session.send(Message { kind: Kind::Text, payload: Bytes::from_static(b"ready") }).await?;
let message = receive.await?.unwrap();
session.send(message).await?;
session.send(Message { kind: Kind::Close { code: Some(1000) }, payload: Bytes::from_static(b"done") }).await
})));
if reject {
*response.status_mut() = StatusCode::FORBIDDEN;
}
Ok(response)
}
Mode::Rewrite => {
assert_eq!(request.uri(), "/unrecognized/client/path?raw=%2f%FF");
assert_eq!(request.headers()["authorization"], "Bearer fixture-secret");
assert_eq!(request.headers().get_all("x-repeat").iter().count(), 2);
*request.uri_mut() = "/healthz".parse().unwrap();
*request.method_mut() = http::Method::GET;
next.run(request).await
}
Mode::Echo => {
let (parts, body) = request.into_parts();
let mut response = contract::Response::new(body);
*response.status_mut() = StatusCode::CREATED;
*response.headers_mut() = parts.headers;
response.extensions_mut().insert(123_u32);
Ok(response)
}
Mode::Slow { clear_timeout } => {
assert_eq!(
request
.extensions()
.get::<contract::Settings>()
.unwrap()
.timeout,
Some(Duration::from_secs(1))
);
if clear_timeout {
request
.extensions_mut()
.get_mut::<contract::Settings>()
.unwrap()
.timeout = None;
}
tokio::time::sleep(Duration::from_secs(2)).await;
next.run(request).await
}
}
})
}
fn handle(
&self,
_: MiddlewareContext,
_: MiddlewareRequest,
_: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
Box::pin(async { panic!("HTTP 入口不能伪装为已认证模型请求") })
}
}
async fn router(mode: Mode, timeout: Option<u64>) -> (axum::Router, Arc<Plan>) {
let admin = crate::admin::AdminTestFixture::new().await;
admin.auth.insert_session("service-fixture");
let plan = Arc::new(Plan {
mode,
calls: Arc::default(),
cancellation: Arc::default(),
message_lifetime: Arc::default(),
});
let frozen = FrozenMiddlewarePlan::new(
plan.clone(),
ExtensionSetReference::new(
ExtensionSetId::new("http-test".into()).unwrap(),
Arc::new(Lease),
),
);
let router = crate::openai::api_bundle(
admin.services,
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: Vec::new(),
request_timeout_seconds: timeout,
request_id_header: "x-request-id".into(),
},
)
.with_middleware(move |_| Some(frozen.clone()))
.router();
(router, plan)
}
#[tokio::test]
async fn rewrites_unknown_path_before_routing_without_reading_request_body() {
let (router, plan) = router(Mode::Rewrite, None).await;
let reads = Arc::new(AtomicUsize::new(0));
let body_reads = reads.clone();
let body = Body::from_stream(stream::poll_fn(move |_| {
body_reads.fetch_add(1, Ordering::SeqCst);
std::task::Poll::Ready(Some(Ok::<_, std::io::Error>(Bytes::from_static(
b"must remain lazy",
))))
}));
let response = router
.oneshot(
Request::builder()
.method("POST")
.uri("/unrecognized/client/path?raw=%2f%FF")
.header("authorization", "Bearer fixture-secret")
.header("x-repeat", "one")
.header("x-repeat", "two")
.body(body)
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::NO_CONTENT);
assert_eq!(plan.calls.load(Ordering::SeqCst), 1);
assert_eq!(reads.load(Ordering::SeqCst), 0);
assert!(
!plan
.cancellation
.lock()
.unwrap()
.as_ref()
.unwrap()
.is_cancelled()
);
drop(response);
assert!(
plan.cancellation
.lock()
.unwrap()
.as_ref()
.unwrap()
.is_cancelled()
);
}
#[tokio::test]
async fn short_circuit_preserves_binary_headers_trailers_extensions_and_backpressure() {
let (router, _) = router(Mode::Echo, None).await;
let mut trailers = http::HeaderMap::new();
trailers.append("x-end", "one".parse().unwrap());
trailers.append("x-end", "two".parse().unwrap());
let reads = Arc::new(AtomicUsize::new(0));
let observed = reads.clone();
let mut frames = vec![
Frame::data(Bytes::from_static(b"first")),
Frame::data(Bytes::from_static(b"second")),
Frame::trailers(trailers),
]
.into_iter();
let body = Body::new(StreamBody::new(stream::poll_fn(move |_| {
observed.fetch_add(1, Ordering::SeqCst);
std::task::Poll::Ready(frames.next().map(Ok::<_, std::io::Error>))
})));
let response = router
.oneshot(
Request::builder()
.uri("/arbitrary/plugin/route")
.header("x-binary", http::HeaderValue::from_bytes(&[0xFF]).unwrap())
.body(body)
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(response.headers()["x-binary"].as_bytes(), &[0xFF]);
assert_eq!(response.extensions().get::<u32>(), Some(&123));
assert_eq!(reads.load(Ordering::SeqCst), 0);
let mut body = response.into_body();
assert_eq!(
body.frame().await.unwrap().unwrap().into_data().unwrap(),
"first"
);
assert_eq!(reads.load(Ordering::SeqCst), 1);
assert_eq!(
body.frame().await.unwrap().unwrap().into_data().unwrap(),
"second"
);
let end = body
.frame()
.await
.unwrap()
.unwrap()
.into_trailers()
.unwrap();
assert_eq!(end.get_all("x-end").iter().count(), 2);
assert!(body.frame().await.is_none());
}
#[tokio::test(start_paused = true)]
async fn host_timeout_is_a_baseline_and_plugin_override_is_consumed_by_terminal() {
for clear_timeout in [false, true] {
let (router, _) = router(Mode::Slow { clear_timeout }, Some(1)).await;
let response = router
.oneshot(
Request::builder()
.uri("/healthz")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(
response.status(),
if clear_timeout {
StatusCode::NO_CONTENT
} else {
StatusCode::REQUEST_TIMEOUT
}
);
}
}
#[tokio::test]
async fn admin_http_does_not_invoke_service_middleware_for_native_business_methods() {
let (router, plan) = router(Mode::Service, None).await;
let response = router
.oneshot(
Request::builder()
.uri("/api/admin/settings")
.header("cookie", "cpr_session=service-fixture")
.header("x-request-id", "service-request")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let bytes = axum::body::to_bytes(response.into_body(), 1024 * 1024)
.await
.unwrap();
let result: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_ne!(result["data"]["requestIntervalMs"], 4321);
assert_eq!(plan.calls.load(Ordering::SeqCst), 0);
}
#[derive(Debug)]
struct Inspect(Arc<Mutex<Vec<contract::Context>>>);
impl gateway_core::engine::middleware::MiddlewarePlan for Inspect {
fn has_http(&self) -> bool {
true
}
fn handle_http(
&self,
context: contract::Context,
mut request: contract::Request,
next: contract::Next,
) -> BoxFuture<'static, Result<contract::Response, MiddlewareError>> {
self.0.lock().unwrap().push(context.clone());
assert_eq!(
request
.extensions()
.get::<contract::Context>()
.unwrap()
.call_id,
context.call_id
);
assert!(context.extensions.contains("origin-plugin"));
assert_eq!(request.headers()["x-request-id"], "parent-request");
assert_eq!(
request
.extensions()
.get::<contract::Settings>()
.unwrap()
.timeout,
None
);
*request.uri_mut() = "/healthz".parse().unwrap();
next.run(request)
}
fn handle(
&self,
_: MiddlewareContext,
_: MiddlewareRequest,
_: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
Box::pin(async { panic!("health request has no model boundary") })
}
}
#[tokio::test]
async fn internal_dispatch_uses_the_same_router_and_the_frozen_parent_context() {
let admin = crate::admin::AdminTestFixture::new().await;
let calls = Arc::new(Mutex::new(Vec::new()));
let frozen = FrozenMiddlewarePlan::new(
Arc::new(Inspect(calls.clone())),
ExtensionSetReference::new(
ExtensionSetId::new("frozen-child".into()).unwrap(),
Arc::new(Lease),
),
);
let bundle = crate::openai::api_bundle(
admin.services,
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: Vec::new(),
request_timeout_seconds: Some(60),
request_id_header: "x-request-id".into(),
},
);
let dispatcher = bundle.dispatcher();
let parent = CancellationToken::new();
let child = parent.child_token();
let context = contract::Context {
plugin_instance_id: None,
request_id: "parent-request".into(),
call_id: "child-call".into(),
parent_call_id: Some("parent-call".into()),
extensions: Default::default(),
plan: Some(frozen),
cancellation: child.clone(),
};
let context = contract::Context {
extensions: context
.extensions
.extending("origin-plugin".into())
.unwrap(),
..context
};
let mut request = http::Request::builder()
.uri("/plugin/unknown/route")
.body(contract::empty_body())
.unwrap();
request.extensions_mut().insert(contract::Settings {
timeout: None,
runtime: None,
});
let response = dispatcher.dispatch(context, request).await.unwrap();
assert_eq!(response.status(), StatusCode::NO_CONTENT);
assert!(!child.is_cancelled());
let body = response.into_body().collect().await.unwrap().to_bytes();
assert!(body.is_empty());
assert!(child.is_cancelled());
assert!(!parent.is_cancelled());
let calls = calls.lock().unwrap();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].call_id, "child-call");
assert_eq!(calls[0].parent_call_id.as_deref(), Some("parent-call"));
}
#[tokio::test]
async fn parent_cancellation_wakes_a_pending_child_response_body() {
let admin = crate::admin::AdminTestFixture::new().await;
let plan = Arc::new(Plan {
mode: Mode::Echo,
calls: Arc::default(),
cancellation: Arc::default(),
message_lifetime: Arc::default(),
});
let frozen = FrozenMiddlewarePlan::new(
plan,
ExtensionSetReference::new(
ExtensionSetId::new("cancel-child".into()).unwrap(),
Arc::new(Lease),
),
);
let bundle = crate::openai::api_bundle(
admin.services,
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: Vec::new(),
request_timeout_seconds: None,
request_id_header: "x-request-id".into(),
},
);
let parent = CancellationToken::new();
let context = contract::Context {
plugin_instance_id: None,
request_id: "parent".into(),
call_id: "child".into(),
parent_call_id: Some("parent".into()),
extensions: Default::default(),
plan: Some(frozen),
cancellation: parent.child_token(),
};
let entered = Arc::new(tokio::sync::Notify::new());
let observed = entered.clone();
let body = StreamBody::new(stream::poll_fn(move |_| {
observed.notify_one();
std::task::Poll::Pending::<Option<Result<Frame<Bytes>, std::io::Error>>>
}))
.map_err(|error| Box::new(error) as _)
.boxed_unsync();
let response = bundle
.dispatcher()
.dispatch(context, contract::Request::new(body))
.await
.unwrap();
let task = tokio::spawn(async move { response.into_body().collect().await });
entered.notified().await;
parent.cancel();
assert!(
tokio::time::timeout(Duration::from_secs(1), task)
.await
.unwrap()
.unwrap()
.is_err()
);
}
#[tokio::test]
async fn plugin_http_dispatch_uses_host_identity_without_client_credentials() {
let admin = crate::admin::AdminTestFixture::new().await;
let bundle = crate::openai::api_bundle(
admin.services,
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: Vec::new(),
request_timeout_seconds: None,
request_id_header: "x-request-id".into(),
},
);
let dispatcher = bundle.dispatcher();
let context = contract::Context {
plugin_instance_id: Some("installed-plugin".into()),
request_id: "model-parent".into(),
call_id: "management-child".into(),
parent_call_id: Some("parent-call".into()),
extensions: Default::default(),
plan: None,
cancellation: CancellationToken::new(),
};
let response = dispatcher
.dispatch(
context.clone(),
http::Request::builder()
.uri("/api/admin/settings")
.body(contract::empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let data: serde_json::Value =
serde_json::from_slice(&response.into_body().collect().await.unwrap().to_bytes()).unwrap();
assert!(data["data"].is_object());
let mutation = contract::Context {
cancellation: CancellationToken::new(),
call_id: "mutation-child".into(),
..context
};
let response = dispatcher
.dispatch(
mutation,
http::Request::builder()
.method("POST")
.uri("/api/admin/settings/admin-api-key/delete")
.body(contract::empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
response.into_body().collect().await.unwrap();
let response = bundle
.router()
.oneshot(
Request::builder()
.uri("/api/admin/settings")
.header("x-plugin-instance-id", "installed-plugin")
.header("x-parent-call-id", "parent-call")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
@@ -0,0 +1,160 @@
use super::*;
use futures::{SinkExt as _, StreamExt as _};
use tokio_tungstenite::{
connect_async,
tungstenite::{Message, client::IntoClientRequest as _},
};
#[tokio::test]
async fn custom_upgrade_keeps_the_call_alive_and_allows_send_during_receive() {
let (router, plan) = router(
Mode::Upgrade {
reject: false,
idle: false,
block_messages: false,
},
None,
)
.await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
tokio::time::timeout(Duration::from_secs(5), async {
let mut request = format!("ws://{address}/plugin-owned/session")
.into_client_request()
.unwrap();
request.headers_mut().insert(
"sec-websocket-protocol",
"other, test-session".parse().unwrap(),
);
let (mut socket, response) = connect_async(request).await.unwrap();
assert_eq!(response.headers()["sec-websocket-protocol"], "test-session");
assert_eq!(
socket.next().await.unwrap().unwrap(),
Message::Text("ready".into())
);
let cancellation = plan.cancellation.lock().unwrap().clone().unwrap();
assert!(!cancellation.is_cancelled(), "101 正文结束不能取消会话");
socket
.send(Message::Binary(vec![0, 255, 128].into()))
.await
.unwrap();
assert_eq!(
socket.next().await.unwrap().unwrap(),
Message::Binary(vec![0, 255, 128].into())
);
let Message::Close(Some(close)) = socket.next().await.unwrap().unwrap() else {
panic!("close frame expected");
};
assert_eq!(close.reason, "done");
cancellation.cancelled().await;
assert_eq!(
plan.calls.load(Ordering::SeqCst),
5,
"HTTP 握手及四次消息收发各组合一次"
);
})
.await
.unwrap();
server.abort();
}
#[tokio::test]
async fn replacing_upgrade_status_drops_the_session_before_handshake() {
let (router, plan) = router(
Mode::Upgrade {
reject: true,
idle: false,
block_messages: false,
},
None,
)
.await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
let result = tokio::time::timeout(
Duration::from_secs(5),
connect_async(format!("ws://{address}/plugin-owned/session")),
)
.await
.unwrap();
assert!(
matches!(result, Err(tokio_tungstenite::tungstenite::Error::Http(response)) if response.status() == StatusCode::BAD_GATEWAY)
);
assert!(
plan.cancellation
.lock()
.unwrap()
.as_ref()
.unwrap()
.is_cancelled()
);
server.abort();
}
#[tokio::test]
async fn disconnect_cancels_an_idle_session_without_a_plugin_receive_call() {
let (router, plan) = router(
Mode::Upgrade {
reject: false,
idle: true,
block_messages: false,
},
None,
)
.await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
let (socket, _) = connect_async(format!("ws://{address}/idle")).await.unwrap();
let cancellation = plan.cancellation.lock().unwrap().clone().unwrap();
assert!(!cancellation.is_cancelled());
drop(socket);
tokio::time::timeout(Duration::from_secs(5), cancellation.cancelled())
.await
.unwrap();
server.abort();
}
#[tokio::test]
async fn cancelling_custom_session_drops_a_blocked_message_middleware() {
let (router, plan) = router(
Mode::Upgrade {
reject: false,
idle: false,
block_messages: true,
},
None,
)
.await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
let (socket, _) = connect_async(format!("ws://{address}/blocked"))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(5), plan.message_lifetime.0.notified())
.await
.unwrap();
let cancellation = plan.cancellation.lock().unwrap().clone().unwrap();
cancellation.cancel();
tokio::time::timeout(Duration::from_secs(5), plan.message_lifetime.1.notified())
.await
.unwrap();
assert_eq!(
plan.calls.load(Ordering::SeqCst),
2,
"取消后不会重新进入消息链"
);
drop(socket);
server.abort();
}
@@ -409,3 +409,165 @@ fn chatgpt_remote_suffix_created_by_header_truncation_should_not_skip_the_gate()
assert_eq!(client.kind(), CodexClientKind::Desktop);
assert!(client.version().is_none());
}
#[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::*,
};
struct Lease;
impl ExtensionSetLease for Lease {
fn is_ready(&self) -> bool {
true
}
}
#[derive(Debug)]
struct Entry;
impl MiddlewarePlan for Entry {
fn has_http(&self) -> bool {
true
}
fn handle_http(
&self,
_: contract::Context,
mut request: contract::Request,
next: contract::Next,
) -> BoxFuture<'static, Result<contract::Response, MiddlewareError>> {
let settings = request
.extensions_mut()
.get_mut::<contract::Settings>()
.unwrap();
let context = settings.runtime.as_ref().unwrap();
let mut values = serde_json::to_value(context.values()).unwrap();
assert_eq!(values["responses_max_decompressed_body_bytes"], 32);
assert_eq!(values["min_codex_cli_version"], "0.40.0");
values["responses_max_decompressed_body_bytes"] = serde_json::json!(1024);
values["min_codex_cli_version"] = serde_json::Value::Null;
settings.runtime = Some(
context
.replace(serde_json::from_value(values).unwrap(), "entry")
.unwrap(),
);
next.run(request)
}
fn handle(
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
next.run(request)
}
}
let snapshot = super::snapshot("sk_entry", "openai");
let settings = snapshot
.settings()
.clone()
.with_min_codex_client_versions(super::CodexClientMinVersions::new(
None,
Some(super::CodexClientVersion::parse("0.40.0").unwrap()),
))
.with_responses_max_decompressed_body_bytes(32);
let snapshot = snapshot.with_settings(&settings).unwrap();
let snapshots = super::RuntimeSnapshotHandle::new(snapshot);
let execution = Arc::new(super::DefaultExecutionService::new(
snapshots.clone(),
Arc::new(super::UnusedExecutionStore),
super::ProviderRegistry::default(),
Arc::new(super::UnusedAdmissions),
Arc::new(super::UnusedContinuation),
Arc::new(super::IgnoredClientApiKeyUsage),
));
let admin = crate::admin::AdminTestFixture::new().await;
let bundle = gateway_api::initialize(
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: vec![],
request_timeout_seconds: None,
request_id_header: "x-request-id".into(),
},
execution,
admin.services,
vec![],
Arc::new(super::EmptyWorkerHealth),
Arc::new(super::TestLifecycle::default()),
)
.unwrap();
let baseline = bundle.dispatcher();
let plan = FrozenMiddlewarePlan::new(
Arc::new(Entry),
ExtensionSetReference::new(
ExtensionSetId::new("entry-settings".into()).unwrap(),
Arc::new(Lease),
),
);
let publisher = snapshots.clone();
let router = bundle
.with_middleware(move |snapshot| {
let snapshot = snapshot.unwrap();
assert_eq!(snapshot.responses_max_decompressed_body_bytes(), 32);
publisher.publish(
snapshot
.with_settings(
&snapshot
.settings()
.clone()
.with_responses_max_decompressed_body_bytes(64),
)
.unwrap(),
);
Some(plan.clone())
})
.router();
let invalid = format!("{}!", " ".repeat(100));
let compressed = zstd::stream::encode_all(invalid.as_bytes(), 0).unwrap();
let request = || {
axum::http::Request::post("/v1/responses")
.header(AUTHORIZATION, "Bearer sk_entry")
.header("user-agent", "codex_cli_rs/0.1.0")
.header("content-encoding", "zstd")
.body(axum::body::Body::from(compressed.clone()))
.unwrap()
};
let response = router.oneshot(request()).await.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body: serde_json::Value = serde_json::from_slice(
&axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap(),
)
.unwrap();
assert_eq!(
body["error"]["code"], "invalid_json",
"the expanded body must reach the decoder"
);
let context = contract::Context {
request_id: "baseline".into(),
call_id: "baseline".into(),
parent_call_id: None,
plugin_instance_id: None,
plan: None,
extensions: Default::default(),
cancellation: Default::default(),
};
use http_body_util::BodyExt as _;
let response = baseline
.dispatch(
context,
request().map(|body| body.map_err(|error| Box::new(error) as _).boxed_unsync()),
)
.await
.unwrap();
assert_eq!(
response.status(),
StatusCode::UPGRADE_REQUIRED,
"request changes must not alter the published baseline"
);
assert_eq!(
snapshots
.acquire()
.unwrap()
.responses_max_decompressed_body_bytes(),
64
);
}
@@ -703,8 +703,12 @@ mod model_routing {
),
] {
for stream in [false, true] {
let snapshot = crate::openai::snapshot("sk_model_routing", "openai")
let snapshot = crate::openai::snapshot("sk_model_routing", "openai");
let settings = snapshot
.settings()
.clone()
.with_model_mappings(mappings.clone());
let snapshot = snapshot.with_settings(&settings).unwrap();
let response = request_model(snapshot, model, stream).await;
assert_eq!(
@@ -28,6 +28,7 @@ pub(super) struct RequestMiddleware {
pub(super) response_actions: Mutex<VecDeque<ResponseFrameAction>>,
pub(super) response_headers: Vec<MiddlewareHeader>,
pub(super) short_circuit_json: Option<Bytes>,
pub(super) settings: Option<fn(&mut gateway_core::settings::ExecutionSettings)>,
}
#[derive(Debug, Clone)]
@@ -92,8 +93,8 @@ impl MiddlewarePlan for RequestMiddleware {
fn handle(
&self,
context: MiddlewareContext,
request: MiddlewareRequest,
next: Box<dyn MiddlewareNext>,
mut request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
assert_eq!(context.mount(), MiddlewareMount::Request);
assert_eq!(context.operation(), self.expected_operation);
@@ -107,6 +108,18 @@ impl MiddlewarePlan for RequestMiddleware {
let actions = std::mem::take(&mut *self.response_actions.lock().unwrap());
let response_headers = self.response_headers.clone();
let short_circuit_json = self.short_circuit_json.clone();
if let Some(update) = self.settings {
let settings = request.settings().expect("effective request settings");
let mut values = settings.execution_values().expect("execution scope");
update(&mut values);
let updated = settings
.replace_execution(&values, "test-middleware")
.map_err(|_| MiddlewareError::InvalidState);
request = match updated {
Ok(settings) => request.with_settings(settings),
Err(error) => return Box::pin(async move { Err(error) }),
};
}
Box::pin(async move {
if let Some(body) = short_circuit_json {
return Ok(MiddlewareResponse::new(
+37 -16
View File
@@ -10,13 +10,12 @@ mod search;
mod usage;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::time::SystemTime;
use async_trait::async_trait;
use futures::future::BoxFuture;
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use gateway_core::account::ProviderAccountId;
use gateway_core::engine::admission::{
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
ClientAdmissionRequest, ClientAdmissionRestoreResult,
@@ -98,6 +97,13 @@ pub(super) fn api_router_with_config(
admin: gateway_admin::AdminServices,
config: gateway_api::ApiConfig,
) -> axum::Router {
api_bundle(admin, config).router()
}
pub(super) fn api_bundle(
admin: gateway_admin::AdminServices,
config: gateway_api::ApiConfig,
) -> gateway_api::ApiBundle {
let execution = Arc::new(DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot("unused-client-route-key", "openai")),
Arc::new(UnusedExecutionStore),
@@ -106,7 +112,15 @@ pub(super) fn api_router_with_config(
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
));
api_router_with_config_and_execution(admin, config, execution)
gateway_api::initialize(
config,
execution,
admin,
Vec::new(),
Arc::new(EmptyWorkerHealth),
Arc::new(TestLifecycle::default()),
)
.unwrap()
}
fn api_router_with_config_and_execution(
@@ -184,8 +198,12 @@ pub(super) fn authenticated_client_for_provider_with_limit(
provider_name: &str,
bytes: usize,
) -> AuthenticatedClient {
let snapshot = snapshot(plaintext, provider_name)
.with_responses_max_decompressed_body_bytes(std::num::NonZeroUsize::new(bytes).unwrap());
let snapshot = snapshot(plaintext, provider_name);
let settings = snapshot
.settings()
.clone()
.with_responses_max_decompressed_body_bytes(bytes as u64);
let snapshot = snapshot.with_settings(&settings).unwrap();
let source = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
Arc::new(UnusedExecutionStore),
@@ -204,11 +222,18 @@ pub(super) fn authenticated_client_with_min_versions(
desktop: Option<&str>,
cli: Option<&str>,
) -> AuthenticatedClient {
let snapshot =
snapshot(plaintext, "openai").with_min_codex_client_versions(CodexClientMinVersions::new(
desktop.map(|version| CodexClientVersion::parse(version).expect("Desktop min version")),
cli.map(|version| CodexClientVersion::parse(version).expect("CLI min version")),
));
let snapshot = snapshot(plaintext, "openai");
let settings =
snapshot
.settings()
.clone()
.with_min_codex_client_versions(CodexClientMinVersions::new(
desktop.map(|version| {
CodexClientVersion::parse(version).expect("Desktop min version")
}),
cli.map(|version| CodexClientVersion::parse(version).expect("CLI min version")),
));
let snapshot = snapshot.with_settings(&settings).unwrap();
let source = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
Arc::new(UnusedExecutionStore),
@@ -248,11 +273,7 @@ fn snapshot_with_client_key(
);
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).expect("concurrency"),
Duration::from_millis(1),
),
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
vec![provider.clone()],
["model-a", "model-b"]
.into_iter()
@@ -156,10 +156,10 @@ impl ExecutionService for ModelsExecution {
}
}
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId> {
["model-a", "model-b"]
.into_iter()
.map(|model| PublicModelId::new(model).expect("model"))
.map(|model| PublicModelId::new(client.snapshot().mapped_model(model)).expect("model"))
.collect()
}
@@ -196,6 +196,40 @@ fn authorized_request(path: &str) -> Request<Body> {
.expect("build models request")
}
#[tokio::test]
async fn request_settings_reach_query_terminal_without_mutating_authenticated_baseline() {
let mut execution = ModelsExecution::new();
Arc::get_mut(&mut execution).unwrap().middleware =
Some(Arc::new(crate::openai::middleware::RequestMiddleware {
settings: Some(|settings| {
let mut runtime = serde_json::to_value(&settings.runtime).unwrap();
runtime["model_mappings"] = serde_json::json!({"model-a":"plugin-model"});
runtime["concurrency_wait_timeout_seconds"] = serde_json::json!(30);
settings.runtime = serde_json::from_value(runtime).unwrap();
settings.disable_fast = true;
}),
..Default::default()
}));
let app = api_router(execution.clone()).await;
for _ in 0..2 {
let response = app
.clone()
.oneshot(authorized_request("/v1/models"))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(body["data"][0]["id"], "plugin-model");
assert_eq!(body["data"][1]["id"], "model-b");
assert_eq!(
execution.client.snapshot().mapped_model("model-a"),
"model-a"
);
assert!(!execution.client.policy().account_scope().disable_fast());
}
}
#[tokio::test]
async fn model_queries_enter_the_same_chain_after_auth_and_hold_generation_until_body_end() {
use std::sync::atomic::Ordering;
@@ -2382,7 +2382,7 @@ async fn buffered_response_never_merges_items_from_a_different_wire_response() {
}
#[tokio::test]
async fn buffered_response_middleware_replaces_body_and_filters_unsafe_headers() {
async fn buffered_response_middleware_replaces_body_and_preserves_explicit_headers() {
let admin = crate::admin::AdminTestFixture::new().await.services;
let trace = Arc::new(Trace::default());
let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()]);
@@ -2395,7 +2395,7 @@ async fn buffered_response_middleware_replaces_body_and_filters_unsafe_headers()
)])
.with_response_headers(vec![
MiddlewareHeader::new("x-policy-result", Bytes::from_static(b"replaced")),
MiddlewareHeader::new("authorization", Bytes::from_static(b"hidden")),
MiddlewareHeader::new("authorization", Bytes::from_static(b"plugin-value")),
]),
);
@@ -2403,7 +2403,7 @@ async fn buffered_response_middleware_replaces_body_and_filters_unsafe_headers()
response_with_middleware(&admin, "openai", session, false, Some(middleware)).await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.headers()["x-policy-result"], "replaced");
assert!(response.headers().get(AUTHORIZATION).is_none());
assert_eq!(response.headers()[AUTHORIZATION], "plugin-value");
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("read policy JSON body");
@@ -0,0 +1,378 @@
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},
};
#[derive(Debug, Default)]
struct Plan {
block_terminal: bool,
control_received: Arc<tokio::sync::Notify>,
incoming: AtomicUsize,
outgoing: AtomicUsize,
http_cancellation: Mutex<Option<CancellationToken>>,
session_cancellation: Mutex<Option<CancellationToken>>,
}
struct Lease;
impl ExtensionSetLease for Lease {
fn is_ready(&self) -> bool {
true
}
}
impl MiddlewarePlan for Plan {
fn has_http(&self) -> bool {
true
}
fn has_websocket(&self) -> bool {
true
}
fn handle_http(
&self,
context: http_contract::Context,
mut request: http_contract::Request,
next: http_contract::Next,
) -> BoxFuture<'static, Result<http_contract::Response, MiddlewareError>> {
*self.http_cancellation.lock().unwrap() = Some(context.cancellation);
assert_eq!(request.uri(), "/custom/websocket");
*request.uri_mut() = "/v1/responses".parse().unwrap();
next.run(request)
}
fn handle_websocket(
&self,
context: core::Context,
message: core::Message,
next: core::Next,
) -> BoxFuture<'static, Result<Option<core::Message>, MiddlewareError>> {
assert!(
context.plan.is_some(),
"消息回调必须继承连接冻结的服务组合计划"
);
assert!(
context
.headers
.iter()
.any(|header| header.name() == "authorization"
&& header.value().as_ref() == b"Bearer sk_ws_crossing_limit")
);
*self.session_cancellation.lock().unwrap() = Some(context.cancellation.clone());
match context.direction {
core::Direction::Incoming => self.incoming.fetch_add(1, Ordering::SeqCst),
core::Direction::Outgoing => self.outgoing.fetch_add(1, Ordering::SeqCst),
};
let control_received = self.control_received.clone();
let block_terminal = self.block_terminal;
Box::pin(async move {
if message.payload == "signal" {
control_received.notify_one();
return Ok(None);
}
if context.direction == core::Direction::Incoming && message.payload == "custom.control"
{
context
.sender
.send(core::Message {
kind: core::Kind::Binary,
payload: Bytes::from_static(&[0, 255, 1]),
})
.await?;
return Ok(None);
}
let mut output = next.run(message).await?;
if context.direction == core::Direction::Outgoing
&& let Some(message) = output.as_mut()
&& message.kind == core::Kind::Text
{
let mut value: Value = serde_json::from_slice(&message.payload).unwrap();
if block_terminal && value["type"] == "response.completed" {
message.kind = core::Kind::Binary;
message.payload = Bytes::from(vec![1; 64 * 1024 * 1024]);
return Ok(output);
}
value["plugin"] = Value::Bool(true);
message.payload = Bytes::from(serde_json::to_vec(&value).unwrap());
}
Ok(output)
})
}
fn handle(
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
// 统一计划延续到连接内的模型调用,没有模型绑定时直接委托。
next.run(request)
}
}
#[tokio::test]
async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake_overrides() {
use gateway_core::{
engine::execution::DefaultExecutionService, engine::provider::ProviderRegistry,
runtime::RuntimeSnapshotHandle,
};
#[derive(Debug, Default)]
struct SettingsPlan(Mutex<Vec<Value>>);
impl MiddlewarePlan for SettingsPlan {
fn has_http(&self) -> bool {
true
}
fn handle_http(
&self,
_: http_contract::Context,
mut request: http_contract::Request,
next: http_contract::Next,
) -> BoxFuture<'static, Result<http_contract::Response, MiddlewareError>> {
let settings = request
.extensions_mut()
.get_mut::<http_contract::Settings>()
.unwrap();
let baseline = settings.runtime.as_ref().unwrap();
let mut values = serde_json::to_value(baseline.values()).unwrap();
values["request_interval_ms"] = json!(0);
settings.runtime = Some(
baseline
.replace(serde_json::from_value(values).unwrap(), "handshake")
.unwrap(),
);
next.run(request)
}
fn handle(
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
_: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
self.0.lock().unwrap().push(
serde_json::to_value(request.settings().unwrap().execution_values().unwrap())
.unwrap(),
);
Box::pin(async { Err(MiddlewareError::Rejected) })
}
}
let initial = crate::openai::snapshot("sk_ws_settings", "openai");
let snapshots = RuntimeSnapshotHandle::new(initial.clone());
let execution = Arc::new(DefaultExecutionService::new(
snapshots.clone(),
Arc::new(crate::openai::UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(crate::openai::UnusedAdmissions),
Arc::new(crate::openai::UnusedContinuation),
Arc::new(crate::openai::IgnoredClientApiKeyUsage),
));
let admin = crate::admin::AdminTestFixture::new().await;
let plan = Arc::new(SettingsPlan::default());
let frozen = FrozenMiddlewarePlan::new(
plan.clone(),
ExtensionSetReference::new(
ExtensionSetId::new("ws-settings".into()).unwrap(),
Arc::new(Lease),
),
);
let router = gateway_api::initialize(
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: vec![],
request_timeout_seconds: None,
request_id_header: "x-request-id".into(),
},
execution,
admin.services,
vec![],
Arc::new(crate::openai::EmptyWorkerHealth),
Arc::new(crate::openai::TestLifecycle::default()),
)
.unwrap()
.with_middleware(move |_| Some(frozen.clone()))
.router();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.unwrap();
request
.headers_mut()
.insert(AUTHORIZATION, "Bearer sk_ws_settings".parse().unwrap());
let (mut socket, _) = connect_async(request).await.unwrap();
for round in 0..3 {
if round == 1 {
let mut settings = serde_json::to_value(initial.settings()).unwrap();
settings["request_interval_ms"] = json!(50);
settings["responses_max_decompressed_body_bytes"] = json!(2048);
snapshots.publish(
initial
.clone()
.with_settings(&serde_json::from_value(settings).unwrap())
.unwrap(),
);
} else if round == 2 {
snapshots.publish(crate::openai::snapshot("sk_replaced", "openai"));
}
socket
.send(ClientMessage::Text(
json!({"type":"response.create", "model":"model-a", "input":"hello"})
.to_string()
.into(),
))
.await
.unwrap();
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.unwrap()
.unwrap()
.unwrap();
let ClientMessage::Text(text) = message else {
panic!("expected error event")
};
let event: Value = serde_json::from_str(&text).unwrap();
assert_eq!(event["type"], "error");
assert_eq!(
plan.0.lock().unwrap().len(),
(round + 1).min(2),
"revoked Key must not reach model middleware"
);
}
{
let observed = plan.0.lock().unwrap();
assert_eq!(
observed[0]["runtime"]["responses_max_decompressed_body_bytes"],
64 * 1024 * 1024
);
assert_eq!(
observed[1]["runtime"]["responses_max_decompressed_body_bytes"],
2048
);
assert_eq!(observed[0]["runtime"]["request_interval_ms"], 0);
assert_eq!(observed[1]["runtime"]["request_interval_ms"], 0);
}
socket.close(None).await.unwrap();
server.abort();
}
#[tokio::test]
async fn rewrite_upgrade_and_control_active_response_keep_http_context_until_connection_ends() {
let plan = Arc::new(Plan::default());
let frozen = FrozenMiddlewarePlan::new(
plan.clone(),
ExtensionSetReference::new(
ExtensionSetId::new("ws-middleware-test".into()).unwrap(),
Arc::new(Lease),
),
);
let (trace, mut socket, server) =
start_active_response_with_middleware(Some(frozen), "/custom/websocket").await;
// 升级后的默认会话仍持有入口上下文;101 正文结束不再提前取消它。
assert!(
!plan
.http_cancellation
.lock()
.unwrap()
.as_ref()
.unwrap()
.is_cancelled()
);
assert!(
!plan
.session_cancellation
.lock()
.unwrap()
.as_ref()
.unwrap()
.is_cancelled()
);
socket
.send(ClientMessage::Text("custom.control".into()))
.await
.unwrap();
let reply = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.unwrap()
.unwrap()
.unwrap();
assert!(matches!(reply, ClientMessage::Binary(bytes) if bytes.as_ref() == [0,255,1]));
assert_eq!(trace.starts.load(Ordering::SeqCst), 1);
assert!(!trace.cancelled.load(Ordering::SeqCst));
trace.release_terminal.notify_one();
let reply = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.unwrap()
.unwrap()
.unwrap();
let ClientMessage::Text(text) = reply else {
panic!("terminal text")
};
let event: Value = serde_json::from_str(&text).unwrap();
assert_eq!(event["type"], "response.completed");
assert_eq!(event["plugin"], true);
assert!(plan.incoming.load(Ordering::SeqCst) >= 2);
assert!(plan.outgoing.load(Ordering::SeqCst) >= 3);
socket.close(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(2), async {
loop {
if plan
.session_cancellation
.lock()
.unwrap()
.as_ref()
.unwrap()
.is_cancelled()
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let http_cancellation = plan.http_cancellation.lock().unwrap().clone().unwrap();
tokio::time::timeout(Duration::from_secs(2), http_cancellation.cancelled())
.await
.unwrap();
server.abort();
}
#[tokio::test]
async fn input_plugin_runs_while_a_real_socket_write_is_blocked() {
let plan = Arc::new(Plan {
block_terminal: true,
..Plan::default()
});
let frozen = FrozenMiddlewarePlan::new(
plan.clone(),
ExtensionSetReference::new(
ExtensionSetId::new("ws-backpressure-test".into()).unwrap(),
Arc::new(Lease),
),
);
let (trace, mut socket, server) =
start_active_response_with_middleware(Some(frozen), "/custom/websocket").await;
trace.release_terminal.notify_one();
// 客户端不读取 64 MiB 输出,真实 TCP 写入将受接收窗口背压。
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(!trace.finalized.load(Ordering::SeqCst));
socket
.send(ClientMessage::Text("signal".into()))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(2), plan.control_received.notified())
.await
.expect("control middleware must run without waiting for the blocked output");
assert!(!trace.finalized.load(Ordering::SeqCst));
socket.close(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(2), async {
loop {
if trace.finalized.load(Ordering::SeqCst) {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("blocked response must finalize on disconnect");
server.abort();
}
@@ -1,3 +1,5 @@
mod middleware;
use std::{
fmt,
pin::Pin,
@@ -237,7 +239,14 @@ impl Stream for TestSocket {
type Item = Result<Message, TestSocketError>;
fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.incoming.poll_recv(context)
let event = self.incoming.poll_recv(context);
// 模拟 Axum transport 的自动 Pong;pump 不再重复实现协议层应答。
if let Poll::Ready(Some(Ok(Message::Ping(payload)))) = &event {
self.written
.send(Message::Pong(payload.clone()))
.expect("automatic Pong");
}
event
}
}
@@ -719,13 +728,44 @@ async fn start_active_response() -> (
Arc<CrossingLimitTrace>,
TestClientSocket,
tokio::task::JoinHandle<()>,
) {
start_active_response_with_middleware(None, "/v1/responses").await
}
async fn start_active_response_with_middleware(
plan: Option<gateway_core::engine::middleware::FrozenMiddlewarePlan>,
path: &str,
) -> (
Arc<CrossingLimitTrace>,
TestClientSocket,
tokio::task::JoinHandle<()>,
) {
let trace = Arc::new(CrossingLimitTrace::default());
let execution = Arc::new(CrossingLimitExecution {
client: authenticated_client("sk_ws_crossing_limit"),
trace: Arc::clone(&trace),
});
let app = api_router(execution).await;
let app = if let Some(plan) = plan {
let admin = crate::admin::AdminTestFixture::new().await;
gateway_api::initialize(
gateway_api::ApiConfig {
asset_directory: std::env::temp_dir(),
cors_allowed_origins: Vec::new(),
request_timeout_seconds: None,
request_id_header: "x-request-id".into(),
},
execution,
admin.services,
Vec::new(),
Arc::new(crate::openai::EmptyWorkerHealth),
Arc::new(crate::openai::TestLifecycle::default()),
)
.unwrap()
.with_middleware(move |_| Some(plan.clone()))
.router()
} else {
api_router(execution).await
};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
@@ -735,7 +775,7 @@ async fn start_active_response() -> (
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
let mut request = format!("ws://{address}{path}")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
@@ -1097,3 +1137,33 @@ async fn active_response_receives_heartbeats_during_two_hundred_seconds_of_silen
socket.close(None).await.expect("close WebSocket");
server.abort();
}
#[tokio::test(start_paused = true)]
async fn client_close_interrupts_a_blocked_write() {
let PumpHarness {
mut connection,
incoming,
dropped,
..
} = test_connection(true);
let write = connection.send_text(
"blocked".to_owned(),
WriteContext::connection(FramePhase::Data),
);
tokio::pin!(write);
assert!(
tokio::time::timeout(Duration::from_millis(1), &mut write)
.await
.is_err()
);
incoming
.send(Ok(Message::Close(None)))
.expect("client close");
assert!(matches!(
tokio::time::timeout(Duration::from_secs(1), write)
.await
.expect("read side observes close while write is blocked"),
Err(ConnectionWriteError::Closed)
));
assert!(dropped.load(Ordering::Acquire));
}
@@ -1,15 +1,23 @@
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, atomic::Ordering};
use std::sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
};
use std::time::Duration;
use axum::http::{StatusCode, header::AUTHORIZATION};
use bytes::Bytes;
use futures::{SinkExt, StreamExt};
use futures::{SinkExt, StreamExt, future::BoxFuture};
use gateway_core::engine::middleware::{
FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan,
MiddlewareRequest, MiddlewareResponse,
};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
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::upstream::{OpaqueUpstreamValue, UpstreamSendState};
use serde_json::{Value, json};
use tokio::net::TcpStream;
@@ -142,6 +150,63 @@ fn upstream_failure(
)
}
#[derive(Debug)]
struct InspectFailureMiddleware(Arc<AtomicUsize>);
impl ExtensionSetLease for InspectFailureMiddleware {
fn is_ready(&self) -> bool {
true
}
}
impl MiddlewarePlan for InspectFailureMiddleware {
fn handle(
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
let drops = Arc::clone(&self.0);
Box::pin(async move {
let response = next.run(request).await?;
assert_eq!(
drops.load(Ordering::Acquire),
0,
"next must retain the failed execution until the onion chain returns"
);
assert_eq!(response.status_code(), 503);
Ok(response)
})
}
}
#[tokio::test]
async fn initial_failure_retains_execution_until_request_middleware_returns() {
let drops = Arc::new(AtomicUsize::new(0));
let middleware = Arc::new(InspectFailureMiddleware(Arc::clone(&drops)));
let trace = Arc::new(AtomicFailureTrace {
initial_errors: Mutex::new(VecDeque::from([EngineError::EmptyRoutingPlan])),
session_drops: Arc::clone(&drops),
middleware: Some(FrozenMiddlewarePlan::new(
middleware.clone(),
ExtensionSetReference::new(
ExtensionSetId::new("failure-inspector".to_owned()).unwrap(),
middleware,
),
)),
..AtomicFailureTrace::default()
});
let (mut socket, _server) = connect(Arc::clone(&trace), vec![]).await;
send_request(&mut socket).await;
let error = next_error(&mut socket).await;
assert_eq!(error["status"], 503);
assert_eq!(error["error"]["code"], "no_available_provider");
assert_no_duplicate_failure(&mut socket).await;
assert_eq!(drops.load(Ordering::Acquire), 1);
assert!(!trace.committed.load(Ordering::Acquire));
socket.close(None).await.unwrap();
}
#[tokio::test]
async fn initial_http_or_opening_failure_preserves_status_and_structured_error() {
// 分类故意都用 Unavailable:不能把收到的 401/429/500 等按网关分类改成 502。
@@ -308,6 +308,8 @@ struct AtomicFailureTrace {
downstream_websocket_connection_ids: Mutex<Vec<String>>,
initial_errors: Mutex<std::collections::VecDeque<EngineError>>,
first_batch: Mutex<Option<CoordinatedEvent>>,
session_drops: Arc<AtomicUsize>,
middleware: Option<gateway_core::engine::middleware::FrozenMiddlewarePlan>,
}
struct AtomicFailureSession {
@@ -316,6 +318,12 @@ struct AtomicFailureSession {
fail_before_first_event: bool,
}
impl Drop for AtomicFailureSession {
fn drop(&mut self) {
self.trace.session_drops.fetch_add(1, Ordering::AcqRel);
}
}
impl ExecutionSession for AtomicFailureSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
@@ -391,6 +399,13 @@ struct AtomicFailureExecution {
}
impl ExecutionService for AtomicFailureExecution {
fn middleware_plan(
&self,
_: &gateway_core::engine::execution::PreparedRootExecution,
) -> Option<gateway_core::engine::middleware::FrozenMiddlewarePlan> {
self.trace.middleware.clone()
}
fn authenticate(
&self,
plaintext: &str,
+4 -1
View File
@@ -17,8 +17,11 @@ base64 = "0.23.0"
bytes = "1.11.1"
chrono-tz = { version = "0.10", features = ["serde"] }
futures = "0.3.32"
http = "1.5.0"
http-body = "1.0.1"
http-body-util = "0.1.5"
futures-timer = "3.0.3"
serde = { version = "1.0.228", features = ["derive"] }
serde = { version = "1.0.228", features = ["derive", "rc"] }
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
sha2 = "0.11.0"
semver = "1.0.27"
@@ -245,7 +245,8 @@ impl fmt::Debug for PlaintextCredential {
}
/// Provider-owned 的任意 JSON object;公共层只搬运、不读取内部 key。
#[derive(Clone, PartialEq, Eq)]
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(transparent)]
pub struct OpaqueProviderData(Map<String, Value>);
impl OpaqueProviderData {
@@ -272,6 +272,11 @@ impl FrozenAccountScope {
self.request_profiles.get(provider)
}
#[must_use]
pub fn request_profiles(&self) -> &BTreeMap<ProviderKind, super::OpaqueProviderData> {
&self.request_profiles
}
/// Key 绑定分组的冻结 Fast 限制,与账号成员资格无关。
#[must_use]
pub const fn with_disable_fast(mut self, disable_fast: bool) -> Self {
@@ -14,7 +14,8 @@ use super::{
};
/// `runtime_settings.rotation_strategy` 的稳定值。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RotationStrategy {
Smart,
QuotaResetPriority,
@@ -18,11 +18,27 @@ pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024;
/// API 从数据面 Authorization 头构造的有界认证信封。
#[derive(Clone, PartialEq, Eq)]
pub struct ClientAuthenticationRequest {
settings: Option<crate::settings::RequestSettings>,
authorization: Arc<str>,
native_bearer: Option<PlaintextClientApiKey>,
}
impl ClientAuthenticationRequest {
#[must_use]
pub fn with_settings(mut self, settings: crate::settings::RequestSettings) -> Self {
self.settings = Some(settings);
self
}
#[must_use]
pub fn settings(&self) -> Option<&crate::settings::RequestSettings> {
self.settings.as_ref()
}
pub(crate) fn take_settings(&mut self) -> Option<crate::settings::RequestSettings> {
self.settings.take()
}
/// 创建入口认证信封;HTTP 语法与缺失头仍由 API owner 负责。
pub fn new(authorization: impl Into<String>) -> Result<Self, ClientAuthenticationRequestError> {
let authorization = authorization.into();
@@ -39,6 +55,7 @@ impl ClientAuthenticationRequest {
.filter(|value| !value.is_empty())
.and_then(|value| PlaintextClientApiKey::new(value.to_owned()).ok());
Ok(Self {
settings: None,
authorization: authorization.into(),
native_bearer,
})
@@ -65,10 +65,10 @@ use crate::routing::{
RoutingContext, RuntimeSnapshot, UpstreamModelId,
};
use crate::runtime::{RuntimeSnapshotHandle, RuntimeSnapshotPublisher};
use crate::settings::RequestSettings;
const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60);
const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100);
const MAX_NESTED_DEPTH: usize = 4;
const MAX_NESTED_EXECUTIONS: usize = 16;
const MAX_CONCURRENT_NESTED_EXECUTIONS: usize = 4;
@@ -108,12 +108,32 @@ pub struct ExecutionRequestMetadata {
#[derive(Clone)]
pub struct AuthenticatedClient {
settings: Option<RequestSettings>,
snapshot: Arc<RuntimeSnapshot>,
policy: ClientPolicy,
authentication: Option<ClientAuthenticationRequest>,
}
impl AuthenticatedClient {
fn execution_timeout(&self) -> Duration {
self.settings
.as_ref()
.and_then(|settings| settings.execution_timeout(self.policy.key_id()))
.unwrap_or(MODEL_REQUEST_DEADLINE)
}
#[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<RuntimeSnapshot> {
&self.snapshot
@@ -157,6 +177,7 @@ pub struct StartExecution {
/// 中间件短路时本值直接释放;调用 `next` 后必须原样交回
/// [`ExecutionService::start_prepared`],不能重新认证为另一个 Key。
pub struct PreparedRootExecution {
extension_scope: ExtensionCallScope,
response_control: super::response_control::ResponseControl,
client: AuthenticatedClient,
request_id: ModelRequestId,
@@ -168,6 +189,51 @@ pub struct PreparedRootExecution {
}
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,
duration_ms(
self.deadline_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()
@@ -176,13 +242,14 @@ impl PreparedRootExecution {
fn new(client: AuthenticatedClient) -> Result<Self, GatewayError> {
let started_at = SystemTime::now();
let deadline_at = started_at
.checked_add(MODEL_REQUEST_DEADLINE)
.checked_add(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()?,
@@ -433,13 +500,18 @@ pub struct BoundModelExecutionContext {
authority: Arc<BoundModelExecutionAuthority>,
}
impl BoundModelExecutionContext {
#[must_use]
pub fn request_settings(&self) -> Option<&RequestSettings> {
self.authority.client.request_settings()
}
}
struct BoundModelExecutionAuthority {
client: AuthenticatedClient,
account_scope: Arc<FrozenAccountScope>,
deadline_at: SystemTime,
cancellation: CancellationToken,
extension_scope: ExtensionCallScope,
graph: Arc<NestedExecutionGraph>,
initiating_plugin_instance_id: String,
}
@@ -452,7 +524,6 @@ impl fmt::Debug for BoundModelExecutionContext {
"initiating_plugin_instance_id",
&self.authority.initiating_plugin_instance_id,
)
.field("deadline_at", &self.authority.deadline_at)
.finish_non_exhaustive()
}
}
@@ -523,6 +594,11 @@ pub trait ExecutionSession: Send {
}
pub trait ExecutionService: Send + Sync {
/// 未就绪时管理与健康接口仍可运行,模型认证按既有路径报告不可用。
fn request_settings(&self) -> Option<RequestSettings> {
None
}
fn authenticate(
&self,
plaintext: &str,
@@ -771,7 +847,9 @@ impl DefaultExecutionService {
.find(|policy| policy.authorize().is_ok())
.cloned()
.ok_or(ClientAuthenticationError::InvalidKey)?;
let policy = RequestSettings::new(snapshot.clone()).apply_policy(policy);
Ok(AuthenticatedClient {
settings: None,
snapshot,
policy,
authentication: Some(authentication),
@@ -780,12 +858,15 @@ impl DefaultExecutionService {
async fn authenticate_request_without_usage(
&self,
authentication: ClientAuthenticationRequest,
mut authentication: ClientAuthenticationRequest,
) -> Result<AuthenticatedClient, ClientAuthenticationError> {
let snapshot = self
.snapshots
.acquire()
.map_err(|_| ClientAuthenticationError::SnapshotUnavailable)?;
let snapshot = match authentication.settings() {
Some(settings) => settings.snapshot(),
None => self
.snapshots
.acquire()
.map_err(|_| ClientAuthenticationError::SnapshotUnavailable)?,
};
let frontend = self
.frontend_authentication
.as_ref()
@@ -819,7 +900,12 @@ impl DefaultExecutionService {
} else {
native_policy(&snapshot, &authentication)?
};
let policy = match authentication.settings() {
Some(settings) => settings.apply_policy(policy),
None => RequestSettings::new(snapshot.clone()).apply_policy(policy),
};
Ok(AuthenticatedClient {
settings: authentication.take_settings(),
snapshot,
policy,
authentication: Some(authentication),
@@ -879,6 +965,10 @@ impl DefaultExecutionService {
"bound execution identity cannot enter the public start path",
)
})?;
let authentication = match client.settings.clone() {
Some(settings) => authentication.with_settings(settings),
None => authentication,
};
client = self
.authenticate_request_without_usage(authentication)
.await
@@ -901,6 +991,7 @@ impl DefaultExecutionService {
metadata: ExecutionRequestMetadata,
) -> Result<StartedExecution, GatewayError> {
let PreparedRootExecution {
extension_scope,
client,
response_control,
request_id,
@@ -915,7 +1006,7 @@ impl DefaultExecutionService {
account_scope: Arc::clone(client.policy.account_scope()),
deadline_at,
cancellation,
extension_scope: ExtensionCallScope::default(),
extension_scope,
required_provider: None,
required_account: None,
nested: None,
@@ -1496,7 +1587,7 @@ impl DefaultExecutionService {
parent_account,
} = request;
let parent = self.active_request(&parent_request_id)?;
if parent.extension_scope.len() >= MAX_NESTED_DEPTH {
if parent.extension_scope.len() >= ExtensionCallScope::MAXIMUM_DEPTH {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"nested model execution depth is exhausted",
@@ -1560,8 +1651,11 @@ impl DefaultExecutionService {
async fn client_for_key_id(
&self,
client_key_id: &ClientApiKeyId,
settings: Option<RequestSettings>,
) -> Result<AuthenticatedClient, GatewayError> {
if let Some(refresh) = &self.snapshot_refresh {
if settings.is_none()
&& let Some(refresh) = &self.snapshot_refresh
{
refresh.refresh().await.map_err(|_| {
GatewayError::new(
GatewayErrorKind::Internal,
@@ -1569,12 +1663,16 @@ impl DefaultExecutionService {
)
})?;
}
let snapshot = self.snapshots.acquire().map_err(|_| {
GatewayError::new(
GatewayErrorKind::Internal,
"current runtime snapshot is unavailable",
)
})?;
let snapshot = settings
.as_ref()
.map(|settings| settings.snapshot())
.map_or_else(|| self.snapshots.acquire(), Ok)
.map_err(|_| {
GatewayError::new(
GatewayErrorKind::Internal,
"current runtime snapshot is unavailable",
)
})?;
let policy = snapshot
.client_policy(client_key_id)
.cloned()
@@ -1588,8 +1686,12 @@ impl DefaultExecutionService {
GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled")
})?;
Ok(AuthenticatedClient {
policy: match &settings {
Some(settings) => settings.apply_policy(policy),
None => RequestSettings::new(snapshot.clone()).apply_policy(policy),
},
settings,
snapshot,
policy,
authentication: None,
})
}
@@ -1598,19 +1700,10 @@ impl DefaultExecutionService {
&self,
binding: BoundModelExecutionBinding,
) -> Result<BoundModelExecutionContext, GatewayError> {
if binding.timeout.is_zero() {
return Err(GatewayError::new(
GatewayErrorKind::InvalidRequest,
"bound model execution timeout is invalid",
));
}
let client = self.client_for_key_id(&binding.client_key_id).await?;
let deadline_at = SystemTime::now()
.checked_add(binding.timeout.min(MODEL_REQUEST_DEADLINE))
.ok_or_else(|| {
GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid")
})?;
if binding.extension_scope.len() >= MAX_NESTED_DEPTH {
let client = self
.client_for_key_id(&binding.client_key_id, binding.settings)
.await?;
if binding.extension_scope.len() >= ExtensionCallScope::MAXIMUM_DEPTH {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"plugin model execution depth exceeded",
@@ -1630,12 +1723,8 @@ impl DefaultExecutionService {
authority: Arc::new(BoundModelExecutionAuthority {
client,
account_scope,
deadline_at,
cancellation: binding.cancellation,
extension_scope,
graph: Arc::new(NestedExecutionGraph::new(Arc::new(
ExecutionEffects::default(),
))),
initiating_plugin_instance_id: binding.initiating_plugin_instance_id,
}),
})
@@ -1660,25 +1749,25 @@ impl DefaultExecutionService {
"bound plugin invocation was cancelled",
));
}
if authority
.deadline_at
.duration_since(SystemTime::now())
.unwrap_or_default()
.is_zero()
{
return Err(GatewayError::new(
GatewayErrorKind::Timeout,
"bound plugin invocation deadline elapsed",
));
}
let now = SystemTime::now();
let deadline_at = now
.checked_add(authority.client.execution_timeout())
.ok_or_else(|| {
GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid")
})?;
let (account_scope, required_provider) = restrict_model_execution_scope(
authority.account_scope.as_ref(),
provider,
account.as_ref(),
None,
)?;
let permit = authority.graph.acquire()?;
let execution_effects_baseline = authority.graph.effects.epoch();
// 连接承载多个独立执行;递归预算和副作用只在本次执行的子调用图内共享。
// 父连接的取消仍通过 authority 传播,不能因新建执行丢失扩展调用链。
let graph = Arc::new(NestedExecutionGraph::new(Arc::new(
ExecutionEffects::default(),
)));
let permit = graph.acquire()?;
let execution_effects_baseline = graph.effects.epoch();
metadata.endpoint = "host.model".to_owned();
metadata.transport = ClientTransport::InternalPlugin;
metadata.client_ip = None;
@@ -1695,7 +1784,7 @@ impl DefaultExecutionService {
AuthorizedExecution {
account_scope,
response_control: None,
deadline_at: authority.deadline_at,
deadline_at,
cancellation: authority.cancellation.child_token(),
extension_scope: authority.extension_scope.clone(),
required_provider,
@@ -1704,11 +1793,11 @@ impl DefaultExecutionService {
bound: Some(BoundExecutionFacts {
initiating_plugin_instance_id: authority.initiating_plugin_instance_id.clone(),
}),
graph: Some(Arc::clone(&authority.graph)),
graph: Some(graph),
execution_effects_baseline,
nested_permit: Some(permit),
},
SystemTime::now(),
now,
None,
)
.await
@@ -2084,6 +2173,10 @@ impl ExecutionStore for TransientExecutionStore {
}
impl ExecutionService for DefaultExecutionService {
fn request_settings(&self) -> Option<RequestSettings> {
self.snapshots.acquire().ok().map(RequestSettings::new)
}
fn authenticate(
&self,
plaintext: &str,
@@ -2224,7 +2317,7 @@ impl ExecutionService for DefaultExecutionService {
) -> BoxFuture<'_, Result<PreparedRootExecution, GatewayError>> {
let client_key_id = client_key_id.clone();
Box::pin(async move {
let client = self.client_for_key_id(&client_key_id).await?;
let client = self.client_for_key_id(&client_key_id, None).await?;
self.client_api_key_usage
.record_used(client.policy.key_id());
PreparedRootExecution::new(client)
@@ -2293,12 +2386,8 @@ impl NestedModelExecutionPort for DefaultExecutionService {
) -> BoxFuture<'_, Result<Vec<PublicModelId>, GatewayError>> {
Box::pin(async move {
let authority = &context.authority;
let remaining = authority
.deadline_at
.duration_since(SystemTime::now())
.unwrap_or_default();
let cancellation = authority.cancellation.cancelled().fuse();
let timeout = Delay::new(remaining).fuse();
let timeout = Delay::new(MODEL_REQUEST_DEADLINE).fuse();
let catalog = self
.client_model_catalog(&authority.client, &protocol, &client_version)
.fuse();
@@ -12,6 +12,8 @@ pub struct ExtensionCallScope {
}
impl ExtensionCallScope {
pub const MAXIMUM_DEPTH: usize = 4;
#[must_use]
pub fn contains(&self, instance_id: &str) -> bool {
self.instance_ids.contains(instance_id)
@@ -28,10 +28,17 @@ use crate::operation::{
};
use crate::policy::ClientApiKeyId;
use crate::runtime::extensions::{ExtensionSetId, ExtensionSetReference};
use crate::settings::RequestSettings;
/// 中间件在一次逻辑请求中的挂载位置。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum MiddlewareMount {
/// 所有 HTTP 接口的认证、路由和正文解析之前。
Http,
/// 下游 WebSocket 消息在默认解析和写入之前组合。
WebSocket,
/// 类型化公开服务在执行默认业务逻辑前组合。
Service,
/// 客户端请求外层;一次逻辑请求只执行一次。
Request,
/// 受管 Provider attempt 内层;每次真实 retry 都重新建立。
@@ -94,6 +101,7 @@ impl fmt::Debug for MiddlewareHeader {
/// 一次中间件调用可修改的请求事实。
#[derive(Clone, PartialEq, Eq)]
pub struct MiddlewareRequest {
settings: Option<RequestSettings>,
protocol: String,
headers: Vec<MiddlewareHeader>,
body: Bytes,
@@ -114,6 +122,7 @@ impl MiddlewareRequest {
pub fn new(protocol: impl Into<String>, headers: Vec<MiddlewareHeader>, body: Bytes) -> Self {
let protocol = protocol.into();
Self {
settings: None,
original_protocol: protocol.clone(),
original_body: body.clone(),
capabilities: None,
@@ -123,6 +132,17 @@ impl MiddlewareRequest {
}
}
#[must_use]
pub fn with_settings(mut self, settings: RequestSettings) -> Self {
self.settings = Some(settings);
self
}
#[must_use]
pub fn settings(&self) -> Option<&RequestSettings> {
self.settings.as_ref()
}
#[must_use]
pub fn protocol(&self) -> &str {
&self.protocol
@@ -265,12 +285,18 @@ impl MiddlewareFrame {
/// 返回正文是否经过至少一个中间件的替换、展开或丢弃映射。
///
/// 该事实只在宿主进程内传播,不进入插件 RPC wire,插件不能自行伪造。
/// 插件可读取该事实;对快照的修改不改变宿主确认的来源。
#[must_use]
pub const fn transformed(&self) -> bool {
self.transformed
}
/// 返回来源事件的只读视图;读取不转移原结算和会话状态的所有权。
#[must_use]
pub fn event(&self) -> Option<&ProviderEvent> {
self.envelope.as_ref().map(MiddlewareFrameEnvelope::event)
}
/// 继承宿主确认的改写事实;`false` 不能清除内层已经记录的改写。
#[must_use]
pub fn with_transformed(mut self, transformed: bool) -> Self {
@@ -286,7 +312,7 @@ impl MiddlewareFrame {
/// 拆出插件可见正文与只能由宿主原样搬运的事实信封。
///
/// Runtime 在一对多转换时必须把信封附到最后一个输出;这样 `Completed` 不会在
/// 前置输出交付前提前终结。信封不能被读取、序列化、复制或伪造。
/// 前置输出交付前提前终结。读取快照不会转移信封的计量所有权。
#[must_use]
pub fn into_parts(
self,
@@ -347,13 +373,20 @@ impl fmt::Debug for MiddlewareFrame {
}
}
/// Runtime 可移动但不能检查或构造的 Core 事实信封。
/// 可读取完整事件、只能随正文移动一次的 Core 事实信封。
///
/// 该类型故意不实现 `Clone`,防止同一 usage/cost 在一对多响应中被复制。
pub struct MiddlewareFrameEnvelope {
event: ProviderEvent,
}
impl MiddlewareFrameEnvelope {
#[must_use]
pub const fn event(&self) -> &ProviderEvent {
&self.event
}
}
impl fmt::Debug for MiddlewareFrameEnvelope {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
@@ -453,7 +486,7 @@ impl MiddlewareResponse {
)
}
/// 搬运宿主 terminal 的不透明来源事实;Runtime 不得序列化或自行构造。
/// 搬运宿主 terminal 的来源事实;读取快照不转移其所有权。
#[must_use]
pub fn with_envelope(mut self, envelope: MiddlewareResponseEnvelope) -> Self {
self.envelope = Some(envelope);
@@ -485,6 +518,11 @@ pub struct MiddlewareResponseEnvelope {
}
impl MiddlewareResponseEnvelope {
#[must_use]
pub const fn metadata(&self) -> &ProviderCallMetadata {
&self.metadata
}
pub(crate) fn into_provider_metadata(self) -> ProviderCallMetadata {
self.metadata
}
@@ -498,9 +536,10 @@ impl fmt::Debug for MiddlewareResponseEnvelope {
}
}
/// Core 已冻结、但不会序列化敏感身份的调用事实。
/// Core 冻结的调用事实;完整业务身份可由插件读取。
#[derive(Clone)]
pub struct MiddlewareContext {
plan: Option<FrozenMiddlewarePlan>,
request_id: ModelRequestId,
mount: MiddlewareMount,
attempt_index: Option<NonZeroU32>,
@@ -519,6 +558,11 @@ pub struct MiddlewareContext {
}
impl MiddlewareContext {
#[must_use]
pub fn plan(&self) -> Option<FrozenMiddlewarePlan> {
self.plan.clone()
}
#[must_use]
pub fn new(target: MiddlewareTarget, authority: MiddlewareAuthority) -> Self {
let MiddlewareTarget {
@@ -541,6 +585,7 @@ impl MiddlewareContext {
execution_effects,
} = authority;
Self {
plan: None,
request_id,
mount,
attempt_index,
@@ -604,13 +649,13 @@ impl MiddlewareContext {
self.account_id.as_ref()
}
/// 只供宿主 Runtime 匹配冻结 binding;不得序列化给插件。
/// 返回冻结的业务归属,也用于匹配 binding。
#[must_use]
pub const fn client_key_id(&self) -> &ClientApiKeyId {
&self.client_key_id
}
/// 只供宿主 Runtime 匹配冻结 binding;不得序列化给插件。
/// 返回冻结的业务归属,也用于匹配 binding。
#[must_use]
pub fn account_group_ids(&self) -> &[AccountGroupId] {
&self.account_group_ids
@@ -651,7 +696,7 @@ pub struct MiddlewareTarget {
pub account_id: Option<ProviderAccountId>,
}
/// Runtime 只用于授权求交的冻结主体;字段不得进入插件 wire。
/// 冻结的业务主体与生命周期;不用于裁剪插件权限。
pub struct MiddlewareAuthority {
pub client_key_id: ClientApiKeyId,
pub account_group_ids: Arc<[AccountGroupId]>,
@@ -684,6 +729,11 @@ impl fmt::Debug for MiddlewareContext {
pub enum MiddlewareError {
Rejected,
Fault,
/// 保留跨进程失败,供外层插件读取;客户端响应由协议 owner 决定。
Remote {
source: Box<dyn std::error::Error + Send + Sync>,
rejected: bool,
},
InvalidState,
Gateway(GatewayError),
Engine(EngineError),
@@ -695,6 +745,7 @@ impl fmt::Display for MiddlewareError {
formatter.write_str(match self {
Self::Rejected => "middleware rejected the request",
Self::Fault => "middleware call failed",
Self::Remote { source, .. } => return fmt::Display::fmt(source, formatter),
Self::InvalidState => "middleware state is invalid",
Self::Gateway(_) => "gateway middleware terminal failed",
Self::Engine(_) => "execution middleware terminal failed",
@@ -703,7 +754,24 @@ impl fmt::Display for MiddlewareError {
}
}
impl std::error::Error for MiddlewareError {}
impl std::error::Error for MiddlewareError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Remote { source, .. } => Some(source.as_ref()),
Self::Gateway(error) => Some(error),
Self::Engine(error) => Some(error),
Self::Provider(error) => Some(error),
_ => None,
}
}
}
impl MiddlewareError {
#[must_use]
pub const fn is_rejected(&self) -> bool {
matches!(self, Self::Rejected | Self::Remote { rejected: true, .. })
}
}
impl From<GatewayError> for MiddlewareError {
fn from(error: GatewayError) -> Self {
@@ -724,20 +792,58 @@ impl From<ProviderError> for MiddlewareError {
}
/// 绑定当前调用且只能消费一次的宿主续体。
pub trait MiddlewareNext: Send {
fn run(
self: Box<Self>,
request: MiddlewareRequest,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>>;
}
pub type MiddlewareNext =
crate::middleware::Next<MiddlewareRequest, MiddlewareResponse, MiddlewareError>;
/// 一个发布代次冻结的统一数据面中间件计划。
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<crate::middleware::service::Value, crate::middleware::service::Error>,
> {
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<Option<crate::middleware::websocket::Message>, 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<crate::middleware::http::Response, MiddlewareError>> {
next.run(request)
}
fn handle(
&self,
context: MiddlewareContext,
request: MiddlewareRequest,
next: Box<dyn MiddlewareNext>,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>>;
}
@@ -749,6 +855,51 @@ pub struct FrozenMiddlewarePlan {
}
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<crate::middleware::service::Value, crate::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::middleware::websocket::Context,
message: crate::middleware::websocket::Message,
next: crate::middleware::websocket::Next,
) -> BoxFuture<'static, Result<Option<crate::middleware::websocket::Message>, 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<crate::middleware::http::Response, MiddlewareError>> {
context.plan = Some(self.clone());
self.plan.handle_http(context, request, next)
}
#[must_use]
pub fn new(plan: Arc<dyn MiddlewarePlan>, generation: ExtensionSetReference) -> Self {
Self {
@@ -759,10 +910,11 @@ impl FrozenMiddlewarePlan {
pub fn handle(
&self,
context: MiddlewareContext,
mut context: MiddlewareContext,
request: MiddlewareRequest,
next: Box<dyn MiddlewareNext>,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
context.plan = Some(self.clone());
self.plan.handle(context, request, next)
}
}
@@ -7,7 +7,6 @@
use std::{
fmt,
sync::atomic::{AtomicUsize, Ordering},
time::Duration,
};
use futures::future::BoxFuture;
@@ -58,14 +57,15 @@ pub struct NestedModelExecutionRequest {
pub parent_account: Option<ProviderAccountId>,
}
/// 管理页或 CLI 为一次插件调用显式绑定的模型执行身份。
/// 插件为当前调用或连接显式绑定的模型执行身份。
///
/// Key ID 由已获模型域访问的调用选择,发起实例由宿主确定。Core 在绑定时冻结当前
/// Key ID 由插件选择,发起实例由宿主确定。Core 在绑定时冻结当前
/// Key 策略;每个实际模型请求仍独立经过准入、预算、账本与计费。
pub struct BoundModelExecutionBinding {
pub settings: Option<crate::settings::RequestSettings>,
pub client_key_id: ClientApiKeyId,
pub initiating_plugin_instance_id: String,
pub timeout: Duration,
/// 身份随父调用或连接回收;每个模型请求使用独立的执行期限。
pub cancellation: CancellationToken,
/// 宿主保留的扩展调用链,异步观察不能通过新建模型请求清空防递归事实。
pub extension_scope: super::extensions::ExtensionCallScope,
@@ -80,7 +80,6 @@ impl fmt::Debug for BoundModelExecutionBinding {
"initiating_plugin_instance_id",
&self.initiating_plugin_instance_id,
)
.field("timeout", &self.timeout)
.finish_non_exhaustive()
}
}
@@ -135,7 +134,7 @@ pub trait NestedModelExecutionPort: Send + Sync {
client_version: String,
) -> BoxFuture<'_, Result<Vec<PublicModelId>, GatewayError>>;
/// 为单次受保护管理调用或 CLI 命令冻结显式执行身份与共享调用图。
/// 为插件调用冻结显式执行身份;取消父调用会取消其所有模型执行。
fn bind(
&self,
binding: BoundModelExecutionBinding,
@@ -52,6 +52,11 @@ impl ModelRouteInput {
&self.account_group_ids
}
#[must_use]
pub fn extension_scope(&self) -> &ExtensionCallScope {
&self.extension_scope
}
#[must_use]
pub fn suppresses_plugin(&self, instance_id: &str) -> bool {
self.extension_scope.contains(instance_id)
@@ -200,6 +205,11 @@ impl AccountScheduleInput {
&self.account_group_ids
}
#[must_use]
pub fn extension_scope(&self) -> &ExtensionCallScope {
&self.extension_scope
}
#[must_use]
pub fn suppresses_plugin(&self, instance_id: &str) -> bool {
self.extension_scope.contains(instance_id)
@@ -468,29 +468,22 @@ const MAX_ATTEMPT_HEADER_NAME_BYTES: usize = 128;
const MAX_ATTEMPT_HEADER_VALUE_BYTES: usize = 16 * 1024;
const MAX_ATTEMPT_HEADER_TOTAL_BYTES: usize = 64 * 1024;
struct ProviderMiddlewareNext {
fn provider_middleware_next(
operation: Operation,
transport: ClientTransport,
terminal: Option<ProviderMiddlewareTerminal>,
}
impl MiddlewareNext for ProviderMiddlewareNext {
fn run(
mut self: Box<Self>,
request: MiddlewareRequest,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
terminal: ProviderMiddlewareTerminal,
) -> MiddlewareNext {
crate::middleware::compose(Vec::new(), move |request| {
Box::pin(async move {
let fallback_protocol = self.operation.protocol().to_owned();
let fallback_protocol = operation.protocol().to_owned();
if request.has_capability_declaration() {
return Err(MiddlewareError::InvalidState);
}
let (protocol, headers, body) = request.into_parts();
validate_attempt_headers(&headers)?;
let operation = self
.operation
let operation = operation
.replace_middleware_wire(protocol, body)
.map_err(|_| MiddlewareError::InvalidState)?;
let terminal = self.terminal.take().ok_or(MiddlewareError::InvalidState)?;
let stream = terminal(operation, headers)
.await
.map_err(MiddlewareError::Provider)?;
@@ -499,14 +492,14 @@ impl MiddlewareNext for ProviderMiddlewareNext {
let body = ProviderMiddlewareBody {
stream,
pending: VecDeque::new(),
transport: self.transport,
transport,
};
Ok(
MiddlewareResponse::new(protocol, 200, Vec::new(), Box::new(body))
.with_provider_metadata(metadata),
)
})
}
})
}
fn validate_attempt_headers(headers: &[MiddlewareHeader]) -> Result<(), MiddlewareError> {
@@ -516,7 +509,6 @@ fn validate_attempt_headers(headers: &[MiddlewareHeader]) -> Result<(), Middlewa
let mut total = 0_usize;
for header in headers {
let name = header.name();
let normalized = name.to_ascii_lowercase();
if name.is_empty()
|| name.len() > MAX_ATTEMPT_HEADER_NAME_BYTES
|| !name
@@ -527,7 +519,6 @@ fn validate_attempt_headers(headers: &[MiddlewareHeader]) -> Result<(), Middlewa
.value()
.iter()
.any(|byte| *byte != b'\t' && (*byte < b' ' || *byte == 0x7f))
|| attempt_header_is_protected(&normalized)
{
return Err(MiddlewareError::InvalidState);
}
@@ -542,48 +533,6 @@ fn validate_attempt_headers(headers: &[MiddlewareHeader]) -> Result<(), Middlewa
Ok(())
}
fn attempt_header_is_protected(name: &str) -> bool {
name.contains("auth")
|| name.contains("credential")
|| name.contains("secret")
|| name.contains("token")
|| name.contains("cookie")
|| name.contains("session")
|| name.contains("conversation")
|| name.contains("thread")
|| name.contains("account")
|| name.contains("organization")
|| name.contains("project")
|| name.contains("tenant")
|| name.contains("principal")
|| name.contains("identity")
|| name.contains("user-id")
|| name.ends_with("-key")
|| name.ends_with("_key")
|| name.starts_with("sec-websocket-")
|| matches!(
name,
"connection"
| "keep-alive"
| "proxy-connection"
| "proxy-authenticate"
| "proxy-authorization"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
| "host"
| "content-length"
| "content-type"
| "content-encoding"
| "accept"
| "accept-encoding"
| "user-agent"
| "x-request-id"
| "x-gateway-request-id"
)
}
struct ProviderMiddlewareBody {
stream: ProviderStream,
pending: VecDeque<(ProviderEvent, bool)>,
@@ -713,11 +662,7 @@ pub async fn execute_attempt_middleware(
.handle(
context,
request,
Box::new(ProviderMiddlewareNext {
operation,
transport,
terminal: Some(terminal),
}),
provider_middleware_next(operation, transport, terminal),
)
.await
.map_err(middleware_prepare_error)?;
@@ -806,11 +751,14 @@ fn middleware_frame_to_provider_event(
fn middleware_prepare_error(error: MiddlewareError) -> ProviderError {
match error {
MiddlewareError::Provider(error) => error,
MiddlewareError::Rejected => ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::NotSent,
),
MiddlewareError::Rejected | MiddlewareError::Remote { rejected: true, .. } => {
ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::NotSent,
)
}
MiddlewareError::Fault
| MiddlewareError::Remote { .. }
| MiddlewareError::InvalidState
| MiddlewareError::Gateway(_)
| MiddlewareError::Engine(_) => middleware_protocol_error(UpstreamSendState::NotSent),
@@ -820,11 +768,14 @@ fn middleware_prepare_error(error: MiddlewareError) -> ProviderError {
fn middleware_body_error(error: MiddlewareError) -> ProviderError {
match error {
MiddlewareError::Provider(error) => error,
MiddlewareError::Rejected => ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::Ambiguous,
),
MiddlewareError::Rejected | MiddlewareError::Remote { rejected: true, .. } => {
ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::Ambiguous,
)
}
MiddlewareError::Fault
| MiddlewareError::Remote { .. }
| MiddlewareError::InvalidState
| MiddlewareError::Gateway(_)
| MiddlewareError::Engine(_) => middleware_protocol_error(UpstreamSendState::Ambiguous),
+2 -2
View File
@@ -842,13 +842,13 @@ impl ProviderEvent {
/// 返回协议表达是否由中间件替换、展开或丢弃映射产生。
///
/// 该事实只供宿主最终协议边界决定是否启用严格复核,不进入插件 wire。
/// 宿主最终协议边界据此启用严格复核,插件可读取其快照。
#[must_use]
pub const fn middleware_transformed(&self) -> bool {
self.middleware_transformed
}
/// 返回中间件第一次改写前的协议表达;仅宿主协议复核可见。
/// 返回中间件第一次改写前的协议表达,供协议复核和插件读取。
#[must_use]
pub fn middleware_origin_wire(&self) -> Option<&ProtocolWireEvent> {
self.middleware_origin_wire.as_deref()
+9 -1
View File
@@ -4,9 +4,17 @@ use crate::validation::{IdentifierError, validate_text};
use std::fmt;
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize)]
#[serde(transparent)]
pub struct ProviderKind(String);
impl<'de> serde::Deserialize<'de> for ProviderKind {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
Self::new(<String as serde::Deserialize>::deserialize(deserializer)?)
.map_err(serde::de::Error::custom)
}
}
impl ProviderKind {
/// 校验 Provider slug。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
+2
View File
@@ -13,11 +13,13 @@ pub mod health;
pub mod identity;
pub mod lifecycle;
pub mod metering;
pub mod middleware;
pub mod operation;
pub mod policy;
pub mod provider_ports;
pub mod routing;
pub mod runtime;
pub mod settings;
pub mod task;
pub mod upstream;
pub mod validation;
@@ -0,0 +1,55 @@
//! HTTP 中间件的值与流合同;不包含服务器、客户端或业务路由实现。
pub mod upgrade;
use std::time::Duration;
use bytes::Bytes;
use http_body_util::{BodyExt as _, Empty, combinators::UnsyncBoxBody};
use crate::{engine::middleware::MiddlewareError, lifecycle::CancellationToken};
/// 正文保留数据帧、trailers 和背压;不要求完整读取或重新编码。
pub type Body = UnsyncBoxBody<Bytes, Box<dyn std::error::Error + Send + Sync>>;
pub type Request = http::Request<Body>;
pub type Response = http::Response<Body>;
pub type Next = super::Next<Request, Response, MiddlewareError>;
/// 路由前调用没有 Client Key;身份在默认认证流程中产生。
#[derive(Clone, Debug)]
pub struct Context {
pub request_id: String,
pub plugin_instance_id: Option<String>,
pub call_id: String,
pub parent_call_id: Option<String>,
pub extensions: crate::engine::extensions::ExtensionCallScope,
pub plan: Option<crate::engine::middleware::FrozenMiddlewarePlan>,
pub cancellation: CancellationToken,
}
/// 内部子请求复用 API 总路由;端口不解释业务路径,也不经网络回环。
pub trait Dispatcher: Send + Sync {
fn request_settings(&self) -> Option<crate::settings::RequestSettings> {
None
}
fn dispatch(
&self,
context: Context,
request: Request,
) -> futures::future::BoxFuture<'static, Result<Response, MiddlewareError>>;
}
/// 宿主解析后放入请求 extensions,插件显式覆盖后由终端消费。
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Settings {
pub timeout: Option<Duration>,
pub runtime: Option<crate::settings::RequestSettings>,
}
#[must_use]
pub fn empty_body() -> Body {
Empty::<Bytes>::new()
.map_err(|never| match never {})
.boxed_unsync()
}
@@ -0,0 +1,44 @@
//! 升级保留 HTTP 返回路径;会话任务由最外层传输在提交握手时接管。
use std::sync::{Arc, Mutex};
use futures::future::{BoxFuture, Shared};
use crate::{engine::middleware::MiddlewareError, middleware::websocket::Session};
pub type PendingSession =
Shared<BoxFuture<'static, Result<Arc<dyn Session>, Arc<MiddlewareError>>>>;
pub type SessionTask = BoxFuture<'static, Result<(), MiddlewareError>>;
/// 具体服务器负责握手与 socket;Core 不依赖 HTTP 服务器实现。
pub trait WebSocketUpgrade: Send + Sync {
fn accept(
&self,
parts: http::request::Parts,
protocols: Vec<String>,
) -> BoxFuture<'static, Result<(super::Response, PendingSession), MiddlewareError>>;
}
/// 响应被替换或丢弃时任务随之释放,不启动脱离请求的插件调用。
#[derive(Clone)]
pub struct Upgraded {
pub status: http::StatusCode,
task: Arc<Mutex<Option<SessionTask>>>,
}
impl Upgraded {
#[must_use]
pub fn new(status: http::StatusCode, task: SessionTask) -> Self {
Self {
status,
task: Arc::new(Mutex::new(Some(task))),
}
}
pub fn take_task(&self) -> Option<SessionTask> {
self.task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take()
}
}
@@ -0,0 +1,47 @@
//! 与协议无关的洋葱组合器;输入、结果与错误保留调用方的领域类型。
pub mod http;
pub mod service;
pub mod websocket;
use futures::future::BoxFuture;
/// 单次调用中的一层;上下文放在类型化输入中或由闭包捕获。
pub type Middleware<Input, Output, Error> = Box<
dyn FnOnce(Input, Next<Input, Output, Error>) -> BoxFuture<'static, Result<Output, Error>>
+ Send,
>;
type Terminal<Input, Output, Error> =
Box<dyn FnOnce(Input) -> BoxFuture<'static, Result<Output, Error>> + Send>;
/// 消费自身后进入下一层,所有权保证同一个续体不能运行两次。
///
/// 返回流或会话时,资源寿命由结果对象持有,组合器不后台驱动或缓冲结果。
pub struct Next<Input, Output, Error> {
layers: std::vec::IntoIter<Middleware<Input, Output, Error>>,
terminal: Terminal<Input, Output, Error>,
}
impl<Input, Output, Error> Next<Input, Output, Error> {
/// 直接返回下一层 future,不启动与父调用脱离的任务。
pub fn run(mut self, input: Input) -> BoxFuture<'static, Result<Output, Error>> {
match self.layers.next() {
Some(layer) => layer(input, self),
None => (self.terminal)(input),
}
}
}
/// 按给定顺序进入中间件,响应与错误沿等待栈逆序返回。
///
/// 空链直接调用终端;输入不需要 Clone,不经过 JSON 或字符串操作分派。
pub fn compose<Input, Output, Error>(
layers: Vec<Middleware<Input, Output, Error>>,
terminal: impl FnOnce(Input) -> BoxFuture<'static, Result<Output, Error>> + Send + 'static,
) -> Next<Input, Output, Error> {
Next {
layers: layers.into_iter(),
terminal: Box::new(terminal),
}
}
@@ -0,0 +1,51 @@
//! 公开服务的跨进程组合边界;原生调用仍使用所属领域的输入输出类型。
use crate::{
engine::{extensions::ExtensionCallScope, middleware::FrozenMiddlewarePlan},
lifecycle::CancellationToken,
};
/// 值只在进入插件链时编码;操作的领域校验与副作用由终端所属服务执行。
pub type Value = serde_json::Value;
pub type Next = super::Next<Value, Value, Error>;
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, thiserror::Error)]
#[error("{message}")]
pub struct Error {
pub kind: String,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub details: Option<Value>,
}
impl Error {
#[must_use]
pub fn invalid(message: impl Into<String>) -> Self {
Self {
kind: "invalid".into(),
message: message.into(),
details: None,
}
}
#[must_use]
pub fn unavailable(message: impl Into<String>) -> Self {
Self {
kind: "unavailable".into(),
message: message.into(),
details: None,
}
}
}
/// 子服务继承同一发布计划及取消信号;防递归集合不承担字段或权限过滤。
#[derive(Clone, Debug)]
pub struct Context {
pub operation: &'static str,
pub request_id: String,
pub call_id: String,
pub parent_call_id: Option<String>,
pub cancellation: CancellationToken,
pub extensions: ExtensionCallScope,
pub plan: FrozenMiddlewarePlan,
}
@@ -0,0 +1,47 @@
//! 双工消息边界,不解析消息正文或第三方控制协议。
use crate::{
engine::middleware::{MiddlewareError, MiddlewareHeader},
lifecycle::CancellationToken,
};
use bytes::Bytes;
use futures::future::BoxFuture;
use std::sync::Arc;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
Incoming,
Outgoing,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
Text,
Binary,
Ping,
Pong,
Close { code: Option<u16> },
}
#[derive(Clone, Debug)]
pub struct Message {
pub kind: Kind,
pub payload: Bytes,
}
/// 等待实际 transport 写入;与接收端分开持有,不在插件调用期间锁住 writer。
pub trait Sender: Send + Sync {
fn send(&self, message: Message) -> BoxFuture<'_, Result<(), MiddlewareError>>;
}
/// 一个接收端和独立发送端;接收等待不能阻塞发送。
pub trait Session: Sender {
fn receive(&self) -> BoxFuture<'_, Result<Option<Message>, MiddlewareError>>;
}
#[derive(Clone)]
pub struct Context {
pub plan: Option<crate::engine::middleware::FrozenMiddlewarePlan>,
pub connection_id: String,
pub direction: Direction,
pub headers: Arc<[MiddlewareHeader]>,
pub cancellation: CancellationToken,
pub sender: Arc<dyn Sender>,
}
pub type Next = super::Next<Message, Option<Message>, MiddlewareError>;
+46 -3
View File
@@ -8,10 +8,12 @@ pub use client_version::{
ClientVersionRejection, CodexClientKind, CodexClientMinVersions, CodexClientVersion,
};
use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use crate::account::scope::FrozenAccountScope;
use crate::account::{OpaqueProviderData, scope::FrozenAccountScope};
use crate::identity::ProviderKind;
use crate::validation::{IdentifierError, PolicyError, validate_text};
/// `client_api_keys.id` 的核心值对象。
@@ -89,7 +91,8 @@ impl fmt::Debug for PlaintextClientApiKey {
}
/// 零表示对应维度不额外限制。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RateLimits {
pub max_concurrency: u64,
pub requests_per_minute: u64,
@@ -105,9 +108,18 @@ impl RateLimits {
}
}
/// Key 自身的设置默认值;请求派生策略保留本值,不能把宿主默认或插件覆盖写回。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientSettings {
pub request_profiles: BTreeMap<ProviderKind, OpaqueProviderData>,
pub disable_fast: bool,
pub limits: RateLimits,
}
/// 从 `client_api_keys` 冻结的公开准入事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientPolicy {
defaults: Arc<ClientSettings>,
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
account_scope: Arc<FrozenAccountScope>,
@@ -117,7 +129,33 @@ pub struct ClientPolicy {
impl ClientPolicy {
#[must_use]
pub const fn new(
pub fn defaults(&self) -> &Arc<ClientSettings> {
&self.defaults
}
pub(crate) fn with_settings(
mut self,
profiles: &BTreeMap<ProviderKind, OpaqueProviderData>,
disable_fast: bool,
limits: RateLimits,
) -> Self {
if self.account_scope.request_profiles() != profiles
|| self.account_scope.disable_fast() != disable_fast
{
self.account_scope = Arc::new(
self.account_scope
.as_ref()
.clone()
.with_request_profiles(profiles.clone())
.with_disable_fast(disable_fast),
);
}
self.limits = limits;
self
}
#[must_use]
pub fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
account_scope: Arc<FrozenAccountScope>,
@@ -125,6 +163,11 @@ impl ClientPolicy {
limits: RateLimits,
) -> Self {
Self {
defaults: Arc::new(ClientSettings {
request_profiles: account_scope.request_profiles().clone(),
disable_fast: account_scope.disable_fast(),
limits,
}),
key_id,
plaintext_key,
account_scope,

Some files were not shown because too many files have changed in this diff Show More