Compare commits

...
30 Commits
Author SHA1 Message Date
ZyyandCodex 705f0fad06 chore(release): v3.22.2
Co-authored-by: Codex <noreply@openai.com>
2026-10-09 18:11:24 +08:00
ZyyandCodex f0c60f501b fix(client): enable remote compaction in Codex config
Declare the V2 remote compaction capability in generated Codex configuration and deployment examples. Document that CC Switch imports do not preserve the full capability configuration.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 18:05:07 +08:00
ZyyandCodex d1febdd249 fix(openai): remove precommit byte limit
Keep large structural events buffered until semantic output, terminal completion, EOF, or the replay grace deadline. Retain byte counts for diagnostics and preserve historical byte_limit exports.

Refs #369

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 18:05:01 +08:00
ZyyandCodex 53d4086cc0 fix(openai): preserve opaque request fields and websocket messages
Forward numeric and extended reasoning effort values and preserve standalone endpoint business headers.

Retain business metadata and original WebSocket text through provider, middleware, and client delivery. Validate rewritten plugin output without rejecting unchanged heartbeats or opaque messages.

Add HTTP, WebSocket, and plugin process regressions for the passthrough boundaries.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 17:23:40 +08:00
ZyyandCodex ce03e8119d fix(diagnostics): retain health and xAI worker failure causes
Record health component failures through operational diagnostics with transition deduplication. Preserve quota and catalog failure context without interrupting the account batch.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 17:10:29 +08:00
ZyyandCodex 3f05b6450c fix(usage): preserve precommit facts in diagnostics export
Export allowlisted release reasons, prefetch bytes and wait times, and document diagnostics schema v4.

Refs #369

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 16:55:12 +08:00
ZyyandCodex a4b73099b5 fix(openai): pass through websocket control frames
Forward non-create text frames unchanged over the original upstream WebSocket while preserving response.create admission and serialization. Keep the control transport available after completion and relay idle upstream events without creating another inference attempt.

Cover unknown frames, late interrupts, terminal settlement, binding revocation, and connection reuse.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 16:26:31 +08:00
ZyyandCodex d875a62652 chore(release): v3.22.1
Co-authored-by: Codex <noreply@openai.com>
2026-10-09 14:10:51 +08:00
ZyyandCodex 98704dfde8 fix(usage): parse response critical path timings
Read response-scoped engine and API overhead metrics while excluding logical-turn totals and incomplete observations. Validate nested response IDs and cover the observed payload over SSE and WebSocket.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 13:51:40 +08:00
ZyyandCodex b07bca312a fix(upgrade): handle long migrations and unify forced shutdown
Extend the startup healthcheck grace period to five minutes and report elapsed migration time every fifteen seconds without changing published migrations.

Use one shutdown path for OS signals, internal cancellation, and managed restarts. Stop listening without waiting for active sessions, preserve bounded worker cleanup, and remove the obsolete connection drain timeout.

Cover shutdown entry points and worker cleanup with real HTTP lifecycle tests and document the deployment behavior.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 13:51:14 +08:00
ZyyandCodex 60966d935e chore(release): v3.22.0
Co-authored-by: Codex <noreply@openai.com>
2026-10-09 13:03:32 +08:00
ZyyandCodex 52981f40f9 fix(providers): align Codex 0.162 protocol handling
Apply latest incremental tool declarations during Grok normalization and align nested Retry-After header parsing with official Codex semantics.

Cover developer instructions in input, Responses Lite tool declarations, partial_answer phases, and end_turn preservation.

Pin the upstream SQLx fix for cancelled BEGIN rollback and add a deterministic pooled-connection regression test.

Report safe Redis error categories and timeout status when the cooldown persistence regression fails.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 12:10:02 +08:00
ZyyandCodex a51c4cd640 refactor(plugins): simplify lifecycle and reuse unchanged instances
Share unchanged prepared instances across extension sets and close each
process when its last owner releases it. Keep inflight request plans intact.

Unify migration recovery and publication, centralize authentication binding
validation, avoid duplicate package inspection, and separate middleware plans.

Validated with 292 admin, 182 runtime, 37 plugin storage and 15 architecture
tests, rustfmt, and strict workspace Clippy.

Co-authored-by: Codex <noreply@openai.com>
2026-10-09 00:12:29 +08:00
ZyyandCodex bda6d973a8 docs: streamline documentation around current behavior
Remove version migration narratives, consolidate repeated explanations, and organize settings and architecture references.

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 22:26:12 +08:00
ZyyandCodex 66a7e4a695 fix(client): align Codex config export with CC Switch 4.0.4
Export config.toml with provider bearer authentication and send supported connection and usage fields to CC Switch. Remove the separate auth.json output and size the configuration preview to its content.

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 22:09:02 +08:00
think-backandCodex 9483f2c91b perf(groups): avoid aggregate queries when loading group selectors
Load account and API key group selectors through a lightweight admin options
endpoint while preserving pagination, filters, disabled groups and revision
metadata. Keep aggregate statistics on the group management endpoint.

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 18:20:24 +08:00
ZyyandCodex faf50c28ad feat(usage): unify request observations and upstream performance
Consolidate request diagnostics into typed request_observation_json with
migration 0025, preserving historical facts and query constraints.

Use upstream response timestamps for output throughput, align TTFT event
boundaries, and surface official timing metrics in the latency popover.

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 18:17:09 +08:00
ZyyandCodex cb397ae53e chore(release): v3.21.3
Co-authored-by: Codex <noreply@openai.com>
2026-10-08 12:12:34 +08:00
ZyyandCodex 63802f9b1a fix(usage): exclude cached tokens from displayed input
Reuse the fresh input calculation across usage records, token details and key usage metrics.

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 11:43:26 +08:00
ZyyandCodex de599ed768 fix(openai): apply timezone overrides without content kinds
Recognize complete Codex environment contexts when custom providers omit
content classifications, while keeping explicit classifications authoritative.

Cover HTTP and WebSocket overrides and proxy location precedence, and
document timezone and client privacy boundaries.

Refs #366

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 11:16:31 +08:00
ZyyandCodex 3315ddbff6 fix(openai): isolate guardian concurrency from ordinary requests
Give Guardian requests a separate per-account lease pool and wait queue while preserving ordinary concurrency and shared request intervals. Use the ordinary pool when the Guardian limit is zero, and align scheduling, diagnostics, settings copy, and documentation.

Co-authored-by: Codex <noreply@openai.com>
2026-10-08 10:38:52 +08:00
ZyyandCodex 6c583e341a fix(openai): reject invalid continuation replay before account selection
Check the recorded native continuation scope before selecting an account
or committing a session binding. Reuse the same scope validation after
middleware while preserving native connection checks and migration policy.

Cover ReplayOwner and ReplayAny with persisted provider metadata or a
connection-local pin, with and without an existing session binding.
Rejected replays must not acquire leases, write bindings or aliases, send
upstream requests, or redirect subsequent child requests.

Refs #364

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 22:45:42 +08:00
ZyyandCodex fb15f3e36c chore(release): v3.21.2
Co-authored-by: Codex <noreply@openai.com>
2026-10-07 20:21:38 +08:00
ZyyandCodex a940eb599f fix(openai): allow atomic first binding from any session request
Allow requests with a reliable session identity to claim a missing account
binding, including child threads and independent background tasks. Remove
the missing-binding wait queue while preserving existing migration rights
and bound-account waiting behavior.

Cover concurrent root/child claims, memory prewarm and turn reuse,
middleware identity changes, cancellation, and bound-account timeouts.
Update the documented strict-affinity contract.

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 20:18:03 +08:00
ZyyandCodex 22a848b360 refactor(backend): clarify ownership and align xai contracts
Separate execution, provider ports, administration, and plugin lifecycle responsibilities.

Fix shutdown ordering, account lock reclamation, multipart cancellation, budget windows, and WebSocket close handling.

Align xAI catalog, reasoning, tools, replay, quota, and cost handling with current Grok contracts and remove legacy fallbacks.

Update API and architecture documentation with regression coverage. Validated backend suites, strict workspace Clippy, isolated PostgreSQL/Redis integration, and real Codex/Grok client flows.

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 19:34:10 +08:00
ZyyandCodex 754da83b02 refactor(backend): simplify query ownership and resource handling
Move observability query policies into Admin and reuse account state reads.
Bound Redis reads to measured batches and share protocol and transport helpers.
Close backup process and staging cleanup gaps, and remove unused entry points,
redundant models, dependencies, and obsolete tests.

Validated 3,526 workspace tests, then all 352 Store tests after batch tuning,
plus strict workspace Clippy and formatting checks. Local release benchmarks
cover concurrent account reads and controlled transport delays.

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 15:43:29 +08:00
ZyyandCodex f174320e2a chore(release): v3.21.1
Co-authored-by: Codex <noreply@openai.com>
2026-10-07 03:57:15 +08:00
ZyyandCodex 9db8528758 fix(metrics): align TTFT and throughput with sub2api
Count non-preamble Responses events for first-token timing while preserving
semantic output boundaries and request-level timing origins. Calculate
output throughput over the complete request duration in usage rows and
aggregate queries.

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 03:52:59 +08:00
ZyyandCodex 119032c067 feat(settings): add preferred account affinity mode
Route relaxed requests through the configured scheduler and keep the session primary during preferred-mode spillover. Default new settings to strict while preserving saved modes.

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 03:50:59 +08:00
ZyyandCodex 3d1dc79f5d fix(auth): refresh sessions only after authentication failures
Use read-only status checks for initial authentication and remove activity-driven renewal. Keep failed and newly queued requests waiting for the current refresh before replaying with the updated cookie.

Co-authored-by: Codex <noreply@openai.com>
2026-10-07 03:10:50 +08:00
470 changed files with 33384 additions and 23246 deletions
+2 -3
View File
@@ -21,7 +21,7 @@
## 快速预览
无需部署,打开 [快速预览服务](https://codex-proxy-rs.ainz.cc) 即可体验管理端的系统概览、账号分组、代理管理与用量统计
无需部署即可体验管理端的系统概览、账号分组、代理管理与用量统计:
| 登录信息 | 默认值 |
| --- | --- |
@@ -53,8 +53,7 @@ curl -fsSL https://raw.githubusercontent.com/zyycn/codex-proxy-rs/main/deploy/in
[安装脚本](deploy/install.sh) 默认安装到当前目录下的 `codex-proxy-rs/`,下载同一正式 Release 的部署文件,
自动生成密码、设置目录权限并启动服务
完成后会显示访问地址和管理员密码,请保存密码,再按下方步骤
[添加账号与客户端密钥](#添加账号与客户端密钥)
完成后保存显示的访问地址和管理员密码,再[登录管理端](#登录管理端)并[添加账号与客户端密钥](#添加账号与客户端密钥)
可在执行 `bash install.sh` 时传入环境变量:
+59 -44
View File
@@ -303,7 +303,7 @@ dependencies = [
"lru",
"percent-encoding",
"regex-lite",
"sha2 0.11.0",
"sha2",
"tracing",
"url",
]
@@ -325,7 +325,7 @@ dependencies = [
"hmac",
"http",
"percent-encoding",
"sha2 0.11.0",
"sha2",
"time",
"tracing",
]
@@ -358,7 +358,7 @@ dependencies = [
"md-5",
"pin-project-lite",
"sha1 0.11.0",
"sha2 0.11.0",
"sha2",
"tracing",
]
@@ -877,7 +877,7 @@ dependencies = [
"gateway-plugin-sdk",
"hex",
"serde_json",
"sha2 0.11.0",
"sha2",
"tar",
"tempfile",
"thiserror",
@@ -1688,7 +1688,7 @@ dependencies = [
"semver",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"subtle",
"syn 3.0.6",
"tempfile",
@@ -1719,7 +1719,7 @@ dependencies = [
"secrecy",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"thiserror",
"tokio",
"tokio-tungstenite 0.29.0",
@@ -1745,12 +1745,11 @@ dependencies = [
"futures",
"futures-timer",
"http",
"http-body",
"http-body-util",
"semver",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"thiserror",
"tracing",
"url",
@@ -1790,7 +1789,7 @@ dependencies = [
"serde_ignored",
"serde_json",
"serde_path_to_error",
"sha2 0.11.0",
"sha2",
"tar",
"tempfile",
"thiserror",
@@ -1831,7 +1830,7 @@ dependencies = [
"semver",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"sqlx",
"tar",
"tempfile",
@@ -1859,6 +1858,7 @@ dependencies = [
name = "gateway-protocol"
version = "0.0.0"
dependencies = [
"http",
"httpdate",
"memchr",
"serde",
@@ -1887,7 +1887,7 @@ dependencies = [
"secrecy",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"sqlx",
"tempfile",
"thiserror",
@@ -2525,6 +2525,15 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libredox"
version = "0.1.25"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61ff90caf6077a803a240f62fdbe88645a890bbca49ef8174c3cb0404362171d"
dependencies = [
"libc",
]
[[package]]
name = "libsqlite3-sys"
version = "0.37.0"
@@ -2842,6 +2851,24 @@ dependencies = [
"libc",
]
[[package]]
name = "objc2-core-foundation"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536"
dependencies = [
"bitflags",
]
[[package]]
name = "objc2-system-configuration"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396"
dependencies = [
"objc2-core-foundation",
]
[[package]]
name = "once_cell"
version = "1.21.4"
@@ -2930,7 +2957,7 @@ dependencies = [
"elliptic-curve",
"primefield",
"primeorder",
"sha2 0.11.0",
"sha2",
]
[[package]]
@@ -3271,7 +3298,7 @@ dependencies = [
"semver",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"subtle",
"tar",
"tempfile",
@@ -3313,7 +3340,7 @@ dependencies = [
"semver",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"thiserror",
"tokio",
"tracing",
@@ -4024,17 +4051,6 @@ version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
dependencies = [
"cfg-if",
"cpufeatures 0.2.17",
"digest 0.10.7",
]
[[package]]
name = "sha2"
version = "0.11.0"
@@ -4155,8 +4171,7 @@ dependencies = [
[[package]]
name = "sqlx"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "378620ccc25c62c89d8be1c819e76a88d59bdcc3304733330788948e619bfd71"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"sqlx-core",
"sqlx-macros",
@@ -4168,8 +4183,7 @@ dependencies = [
[[package]]
name = "sqlx-core"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"base64 0.22.1",
"bytes",
@@ -4192,7 +4206,7 @@ dependencies = [
"rustls",
"serde",
"serde_json",
"sha2 0.10.9",
"sha2",
"smallvec",
"thiserror",
"tokio",
@@ -4206,8 +4220,7 @@ dependencies = [
[[package]]
name = "sqlx-macros"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd2b84f2bc39a5705ef27ec785a11c934a41bbd4a24941e257927cddc26b60bf"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"proc-macro2",
"quote",
@@ -4219,8 +4232,7 @@ dependencies = [
[[package]]
name = "sqlx-macros-core"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fb8d96de5fdc85a5c4ec813432b523ec637e80ba98f046555f75f7908ddac7c3"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"cfg-if",
"dotenvy",
@@ -4231,7 +4243,7 @@ dependencies = [
"quote",
"serde",
"serde_json",
"sha2 0.10.9",
"sha2",
"sqlx-core",
"sqlx-mysql",
"sqlx-postgres",
@@ -4245,8 +4257,7 @@ dependencies = [
[[package]]
name = "sqlx-mysql"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "90b8020fe17c5f2c245bfa2505d7ef59c5604839527c740266ad2214acebea27"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"bitflags",
"byteorder",
@@ -4263,7 +4274,7 @@ dependencies = [
"percent-encoding",
"serde",
"sha1 0.11.0",
"sha2 0.11.0",
"sha2",
"sqlx-core",
"thiserror",
"tracing",
@@ -4273,8 +4284,7 @@ dependencies = [
[[package]]
name = "sqlx-postgres"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"atoi",
"base64 0.22.1",
@@ -4297,7 +4307,7 @@ dependencies = [
"rand 0.10.2",
"serde",
"serde_json",
"sha2 0.11.0",
"sha2",
"smallvec",
"sqlx-core",
"stringprep",
@@ -4310,8 +4320,7 @@ dependencies = [
[[package]]
name = "sqlx-sqlite"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "488e99c397a62007e4229aec669a179816339afc6d2620ca6fa420dbee2e982c"
source = "git+https://github.com/transact-rs/sqlx?rev=f1e94ec34dd5f11ad6488a3a8ac7988439a3f757#f1e94ec34dd5f11ad6488a3a8ac7988439a3f757"
dependencies = [
"atoi",
"chrono",
@@ -4494,7 +4503,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
@@ -5183,6 +5192,12 @@ name = "whoami"
version = "2.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d"
dependencies = [
"libc",
"libredox",
"objc2-system-configuration",
"web-sys",
]
[[package]]
name = "winapi"
+2
View File
@@ -59,3 +59,5 @@ panic = "abort"
[patch.crates-io]
# Upstream SOCKS IPv6 fix; remove once a release includes hyperium/hyper-util#302.
hyper-util = { git = "https://github.com/hyperium/hyper-util", rev = "191b1acf0a6b18e807e4abf5a6d4a7c47979ecd6" }
# 固定上游 BEGIN 取消回滚修复,避免未结束事务复用;见 transact-rs/sqlx#4394
sqlx = { git = "https://github.com/transact-rs/sqlx", rev = "f1e94ec34dd5f11ad6488a3a8ac7988439a3f757" }
@@ -93,7 +93,7 @@ async fn execute(
Some(plugin_runtime.clone()),
Some(plugin_runtime.observer_registry()),
Some(plugin_runtime.policy_registry()),
Some(plugin_runtime.middleware_registry()),
Some(plugin_runtime.execution_registry()),
);
let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control());
or_shutdown!(plugin_runtime, management.bind(&plugin_runtime));
@@ -110,7 +110,7 @@ async fn execute(
providers.admin,
std::sync::Arc::new(gateway_host::pricing::ModelsDevPricing),
);
let mut services = gateway_admin::service::Registry::new(service_middleware(
let mut services = gateway_admin::public_service::Registry::new(service_middleware(
core.snapshots(),
&plugin_runtime,
));
@@ -72,11 +72,11 @@ impl ManagementPorts {
pub(super) fn service_middleware(
snapshots: RuntimeSnapshotHandle,
runtime: &PluginRuntime,
) -> gateway_admin::service::PlanSource {
let middleware = runtime.middleware_registry();
) -> gateway_admin::public_service::PlanSource {
let middleware = runtime.execution_registry();
Arc::new(move || {
let snapshot = snapshots.snapshot_for_diagnostics()?;
middleware.resolve(snapshot.extensions()?)
middleware.middleware(snapshot.extensions()?)
})
}
+5 -5
View File
@@ -1,4 +1,4 @@
//! 服务模式完成全部组装后交给 Host 监听,并在请求和 Worker 排空后关闭插件
//! 服务模式完成全部组装后交给 Host 监听,并在停止监听与 Worker 收尾后关闭插件
use gateway_host::ConfigError;
@@ -28,7 +28,7 @@ pub async fn run() -> Result<(), BootstrapError> {
Some(plugin_runtime.clone()),
Some(plugin_runtime.observer_registry()),
Some(plugin_runtime.policy_registry()),
Some(plugin_runtime.middleware_registry()),
Some(plugin_runtime.execution_registry()),
Some(plugin_runtime.frontend_authentication_registry()),
);
let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control());
@@ -122,14 +122,14 @@ pub async fn run() -> Result<(), BootstrapError> {
)
.with_middleware({
let snapshots = core.snapshots();
let middleware = plugin_runtime.middleware_registry();
let middleware = plugin_runtime.execution_registry();
move |snapshot| {
if let Some(snapshot) = snapshot {
return middleware.resolve(snapshot.extensions()?);
return middleware.middleware(snapshot.extensions()?);
}
// 数据面未就绪时,管理与诊断路由仍可使用发布候选中的插件
let diagnostic = snapshots.snapshot_for_diagnostics()?;
middleware.resolve(diagnostic.extensions()?)
middleware.middleware(diagnostic.extensions()?)
}
});
let http_dispatcher = api.dispatcher();
+142 -29
View File
@@ -79,18 +79,25 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() {
path if path.starts_with("policy/") => {
Some(&["account", "identity", "policy", "validation"])
}
"settings/values.rs" | "settings/compiled.rs" => {
Some(&["account", "concurrency", "identity", "metering", "policy"])
}
path if path.starts_with("settings/") => Some(&[
"account",
"concurrency",
"identity",
"metering",
"policy",
"settings",
]),
path if path.starts_with("account/") => Some(&["account", "identity", "validation"]),
_ => None,
};
let Some(allowed) = allowed else { continue };
let source = fs::read_to_string(root.join(&relative)).expect("read core source");
let syntax = syn::parse_file(&source).expect("parse core source");
let mut references = CrateReferences::default();
let mut references = CrateReferences::in_module(&relative);
references.visit_file(&syntax);
for dependency in references.0 {
// 相对路径中的 self 指向当前 owner,不是跨 owner 依赖
references.owners.remove(&references.module_path[0]);
for dependency in references.owners {
assert!(
allowed.contains(&dependency.as_str()),
"{} must not depend on crate::{dependency}",
@@ -100,49 +107,155 @@ fn core_value_owners_do_not_depend_on_execution_or_routing() {
}
}
#[test]
fn core_owner_paths_form_a_dag() {
let root = backend_root().join("crates/gateway-core/src");
let mut graph = BTreeMap::<String, BTreeSet<String>>::new();
for relative in super::rust_files(&root) {
if relative == Path::new("lib.rs") {
continue;
}
let mut references = CrateReferences::in_module(&relative);
let owner = references.module_path[0].clone();
let source = fs::read_to_string(root.join(&relative)).expect("read core source");
references.visit_file(&syn::parse_file(&source).expect("parse core source"));
references.owners.remove(&owner);
graph.entry(owner).or_default().extend(references.owners);
}
for owner in graph.keys() {
let mut pending = graph[owner].iter().cloned().collect::<Vec<_>>();
let mut visited = BTreeSet::new();
while let Some(dependency) = pending.pop() {
assert_ne!(
&dependency, owner,
"Core owner dependency cycle reaches {owner}"
);
if visited.insert(dependency.clone())
&& let Some(next) = graph.get(&dependency)
{
pending.extend(next.iter().cloned());
}
}
}
}
#[test]
fn core_reference_scan_resolves_groups_reexports_and_relative_paths() {
let syntax = syn::parse_file(
r#"
pub use crate::{engine::AttemptContext, settings::SettingsValues as Values};
use super::super::runtime as publication;
use self::compiled::Value;
// crate::ignored::Comment
const EXAMPLE: &str = "crate::ignored::Literal";
"#,
)
.expect("parse reference fixture");
let mut references = CrateReferences::in_module(Path::new("routing/snapshot.rs"));
references.visit_file(&syntax);
assert_eq!(
references.owners,
BTreeSet::from([
"engine".to_owned(),
"settings".to_owned(),
"runtime".to_owned(),
"routing".to_owned(),
])
);
}
#[derive(Default)]
struct CrateReferences(BTreeSet<String>);
struct CrateReferences {
owners: BTreeSet<String>,
module_path: Vec<String>,
}
impl<'ast> Visit<'ast> for CrateReferences {
fn visit_path(&mut self, path: &'ast syn::Path) {
let mut segments = path.segments.iter();
if segments
.next()
.is_some_and(|segment| segment.ident == "crate")
&& let Some(owner) = segments.next()
{
self.0.insert(owner.ident.to_string());
}
self.record(
&path
.segments
.iter()
.map(|part| part.ident.to_string())
.collect::<Vec<_>>(),
);
syn::visit::visit_path(self, path);
}
fn visit_item_use(&mut self, item: &'ast syn::ItemUse) {
self.use_tree(&item.tree, false);
self.use_tree(&item.tree, &mut Vec::new());
}
}
impl CrateReferences {
fn use_tree(&mut self, tree: &syn::UseTree, crate_root: bool) {
fn in_module(relative: &Path) -> Self {
let mut module = relative
.with_extension("")
.components()
.map(|part| {
part.as_os_str()
.to_str()
.expect("UTF-8 module path")
.to_owned()
})
.collect::<Vec<_>>();
if module.last().is_some_and(|part| part == "mod") {
module.pop();
}
Self {
owners: BTreeSet::new(),
module_path: module,
}
}
fn record(&mut self, path: &[String]) {
let Some(first) = path.first() else { return };
let owner = if first == "crate" {
path.get(1).cloned()
} else if first == "self" || first == "super" {
let mut module = self.module_path.clone();
let mut rest = path;
while rest
.first()
.is_some_and(|part| part == "super" || part == "self")
{
if rest[0] == "super" {
module.pop();
}
rest = &rest[1..];
}
module.first().or_else(|| rest.first()).cloned()
} else {
None
};
if let Some(owner) = owner {
self.owners.insert(owner);
}
}
fn use_tree(&mut self, tree: &syn::UseTree, prefix: &mut Vec<String>) {
match tree {
syn::UseTree::Path(path) if crate_root => {
self.0.insert(path.ident.to_string());
syn::UseTree::Path(path) => {
prefix.push(path.ident.to_string());
self.use_tree(&path.tree, prefix);
prefix.pop();
}
syn::UseTree::Name(name) if crate_root => {
self.0.insert(name.ident.to_string());
syn::UseTree::Name(name) => {
prefix.push(name.ident.to_string());
self.record(prefix);
prefix.pop();
}
syn::UseTree::Rename(rename) if crate_root => {
self.0.insert(rename.ident.to_string());
syn::UseTree::Rename(rename) => {
prefix.push(rename.ident.to_string());
self.record(prefix);
prefix.pop();
}
syn::UseTree::Path(path) => self.use_tree(&path.tree, path.ident == "crate"),
syn::UseTree::Group(group) => {
for item in &group.items {
self.use_tree(item, crate_root);
self.use_tree(item, prefix);
}
}
syn::UseTree::Glob(_) if crate_root => {
self.0.insert("*".to_owned());
}
_ => {}
syn::UseTree::Glob(_) => self.record(prefix),
}
}
}
@@ -221,7 +334,7 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[
"crates/providers/openai",
&["config", "credential", "transport"],
),
("crates/providers/xai", &["credential", "transport"]),
("crates/providers/xai", &["transport"]),
];
/// 不对应单一生产模块、而是校验 crate/workspace 整体契约的根级测试场景
@@ -76,11 +76,6 @@ fn config_loader_should_reject_missing_runtime_data_dir() {
assert!(parse_config(&config).is_err());
}
#[test]
fn config_loader_should_inject_connection_passwords_into_urls() {
parse_config(&valid_config()).expect("Store validates password injection into both URLs");
}
#[test]
fn config_loader_should_apply_only_explicit_topology_overrides() {
let invalid = valid_config()
+2 -5
View File
@@ -28,11 +28,8 @@ cpr-plugin package \
资源映射相对于作者清单所在目录解析。输出包括完整插件 ID、版本和目标 triple 命名的 `.tar.gz` 及 `.sha256`;归档根目录直接包含 `plugin.json`、可执行文件和已声明资源,不包含开发依赖
作者清单可省略普通贡献项的 `id`、`version` 和固定阶段;`middleware` 必须显式声明版本 `4`,
并从 `http`、`websocket`、`service`、`request`、`attempt` 中选择挂载阶段;`upstream_adapter` 必须显式声明版本 `2`。
CLI 与 SDK 的 `Manifest::from_author_slice`、`PluginBuilder::from_json` 共用规范化入口;归档中的
`plugin.json` 写入补全后的贡献项字段,不保留另一份作者声明。
完整规则见 [SDK 清单](../../crates/gateway-plugin/sdk/docs/manifest.md)
作者清单的省略规则、能力版本和挂载阶段见 [SDK 扩展项声明](../../crates/gateway-plugin/sdk/docs/manifest.md#扩展项简写)。
CLI 与 SDK 的 `Manifest::from_author_slice`、`PluginBuilder::from_json` 共用规范化入口,归档只保留补全字段后的 `plugin.json`
支持 `x86_64-unknown-linux-gnu`、`aarch64-unknown-linux-gnu` 和 `aarch64-apple-darwin`。
完整开发示例在独立仓库 `codex-proxy-plugins` 的 `examples/workbench`,可通过[子模块](../../../docs/development.md#源码联调)检出;本工具不负责源码构建、上传或启用插件
@@ -281,7 +281,6 @@ impl BackupTask {
.await?;
}
None => {
let _ = self.dump.cleanup_staging(&record.id).await;
self.fail_task(record, code::PG_DUMP_FAILED, "导出进程中断,暂存归档不完整")
.await?;
}
@@ -307,7 +306,8 @@ impl BackupTask {
&& record.sha256.as_deref() == Some(remote.sha256.as_str())
{
let now = Utc::now();
self.repository
if self
.repository
.transition_status(
&record.id,
status_transition(BackupStatus::Uploading, BackupStatus::Completed)?,
@@ -315,7 +315,11 @@ impl BackupTask {
now,
)
.await
.map_err(repo_error)?;
.map_err(repo_error)?
.is_some()
{
self.cleanup_staging(&record.id).await;
}
return Ok(());
}
match self.dump.inspect_staging(&record.id).await? {
@@ -455,9 +459,7 @@ impl BackupTask {
now,
cancellation,
)
.await?;
let _ = self.dump.cleanup_staging(&record.id).await;
Ok(())
.await
}
/// 上传 + 远端校验,内部完成终态迁移
@@ -518,6 +520,7 @@ impl BackupTask {
size_bytes,
"备份完成"
);
self.cleanup_staging(&record.id).await;
}
}
Ok(Some(remote)) => {
@@ -596,10 +599,16 @@ impl BackupTask {
)
.await
.map_err(repo_error)?;
let _ = self.dump.cleanup_staging(&record.id).await;
self.cleanup_staging(&record.id).await;
Ok(())
}
async fn cleanup_staging(&self, backup_id: &str) {
if let Err(error) = self.dump.cleanup_staging(backup_id).await {
warn!(backup_id, error = %error, "清理备份暂存失败");
}
}
/// 执行一小批到期保留清理:先处理 `expires_at` 已到期的记录,再按天数/份数清理计划备份
async fn run_retention_batch(&self, now: DateTime<Utc>) -> Result<(), WorkerTaskError> {
let settings = self.repository.load_settings().await.map_err(repo_error)?;
+8 -19
View File
@@ -20,7 +20,7 @@ pub mod backup;
pub mod freeze_recovery;
pub mod model;
pub mod ports;
pub mod service;
pub mod public_service;
mod use_case;
pub use use_case::plugins::{PluginDistributionPorts, PluginManagementService, PluginsService};
@@ -205,7 +205,7 @@ pub enum AdminConfigError {
#[derive(Clone)]
pub struct AdminServices {
timezone: gateway_core::time::DeploymentTimeZone,
public_services: Arc<service::Registry>,
public_services: Arc<public_service::Registry>,
plugins: Arc<PluginsService>,
plugin_management: Arc<PluginManagementService>,
proxies: Arc<dyn ProxiesService>,
@@ -219,7 +219,8 @@ pub struct AdminServices {
settings: Arc<dyn SettingsService>,
system: Arc<dyn SystemService>,
credentials: Arc<CredentialsService>,
plugin_accounts: Arc<dyn PluginAccountAccess>,
// Runtime 只持有 Weak,这个强引用保障账号端口与 Admin 服务同寿命
_plugin_accounts: Arc<dyn PluginAccountAccess>,
backups: Arc<dyn BackupService>,
import_tasks: Arc<dyn ImportTasksService>,
}
@@ -230,7 +231,7 @@ impl AdminServices {
self.timezone
}
pub fn public_services(&self) -> Arc<service::Registry> {
pub fn public_services(&self) -> Arc<public_service::Registry> {
self.public_services.clone()
}
@@ -254,12 +255,6 @@ impl AdminServices {
self.key_usage.as_ref()
}
/// 取得账号服务的共享句柄;后台编排(冻结恢复 worker)需要持有 Arc
#[must_use]
pub fn accounts_handle(&self) -> Arc<dyn AccountsService> {
Arc::clone(&self.accounts)
}
#[must_use]
pub fn proxies(&self) -> &dyn ProxiesService {
self.proxies.as_ref()
@@ -309,12 +304,6 @@ impl AdminServices {
self.credentials.as_ref()
}
/// Runtime 只持有该窄端口的 Weak;AdminBundle 保持实际生命周期
#[must_use]
pub fn plugin_accounts_handle(&self) -> Arc<dyn PluginAccountAccess> {
Arc::clone(&self.plugin_accounts)
}
#[must_use]
pub fn backups(&self) -> &dyn BackupService {
self.backups.as_ref()
@@ -342,7 +331,7 @@ impl AdminBundle {
/// 组合根提供给控制面的运行能力;与配置和存储端口分别传入
pub struct AdminRuntimePorts {
pub timezone: gateway_core::time::DeploymentTimeZone,
pub service_middleware: service::PlanSource,
pub service_middleware: public_service::PlanSource,
pub plugin_preparation: Arc<dyn ports::plugins::PluginPreparation>,
pub plugin_management: Arc<dyn ports::plugin_management::PluginManagement>,
pub published_snapshot: gateway_core::runtime::RuntimeSnapshotHandle,
@@ -477,7 +466,7 @@ async fn initialize_inner(
registry.clone(),
pricing_source,
);
let mut public_services = service::Registry::new(service_middleware);
let mut public_services = public_service::Registry::new(service_middleware);
public_services.register_settings(&settings)?;
let services = AdminServices {
timezone,
@@ -526,7 +515,7 @@ async fn initialize_inner(
settings,
system,
credentials,
plugin_accounts,
_plugin_accounts: plugin_accounts,
import_tasks,
backups,
};
@@ -126,6 +126,16 @@ pub struct AccountGroupPage {
pub page_size: u16,
}
/// 分组选择器使用的轻量分页结果
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupOptionsPage {
pub config_revision: Revision,
pub items: Vec<AccountGroupRef>,
pub total: u64,
pub page: u32,
pub page_size: u16,
}
/// 创建账号分组
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateAccountGroup {
@@ -177,6 +177,32 @@ pub struct UsageQuery {
pub page_size: ObservabilityPageSize,
}
/// Admin 确定概览的数量、趋势粒度及最近请求筛选,Store 只执行该选择
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DashboardQuery {
pub range: TimeRange,
pub granularity: Granularity,
pub account_limit: u16,
pub recent_request_limit: u16,
pub recent_request_filter: UsageFilter,
}
impl DashboardQuery {
#[must_use]
pub fn new(range: TimeRange) -> Self {
Self {
range,
granularity: Granularity::for_range(range),
account_limit: 4,
recent_request_limit: 10,
recent_request_filter: UsageFilter {
outcome: Some(RequestOutcome::Succeeded),
..UsageFilter::default()
},
}
}
}
/// 运维错误过滤条件
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OpsErrorFilter {
@@ -477,6 +503,30 @@ pub enum Granularity {
Day,
}
impl Granularity {
/// 管理趋势按范围选择粒度,费用校验与请求指标必须使用同一策略
#[must_use]
pub fn for_range(range: TimeRange) -> Self {
let seconds = range.end.signed_duration_since(range.start).num_seconds();
if seconds <= 2 * 24 * 60 * 60 {
Self::FifteenMinutes
} else if seconds <= 31 * 24 * 60 * 60 {
Self::Hour
} else {
Self::Day
}
}
#[must_use]
pub const fn seconds(self) -> i64 {
match self {
Self::FifteenMinutes => 15 * 60,
Self::Hour => 60 * 60,
Self::Day => 24 * 60 * 60,
}
}
}
/// 一段时间桶内的请求指标
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RequestMetricPoint {
@@ -589,13 +639,14 @@ pub struct DashboardObservation {
pub range: TimeRange,
pub totals: DashboardTotals,
pub provider_accounts: AccountPoolMetrics,
pub runtime_slots: Option<DashboardRuntimeSlots>,
pub trend: Vec<RequestMetricPoint>,
pub account_usage: Vec<DashboardAccountUsage>,
pub recent_requests: Vec<UsageListRecord>,
}
/// 使用记录表格的窄读模型
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct UsageListRecord {
pub client_api_key_name: Option<String>,
pub id: String,
@@ -636,6 +687,13 @@ pub struct UsageListRecord {
pub first_text_ms: Option<u64>,
pub first_token_ms: Option<u64>,
pub provider_processing_ms: Option<u64>,
pub upstream_response_ms: Option<u64>,
pub upstream_api_overhead_ms: Option<f64>,
pub upstream_engine_ms: Option<f64>,
pub upstream_engine_iapi_ttft_ms: Option<f64>,
pub upstream_engine_service_ttft_ms: Option<f64>,
pub upstream_engine_iapi_tbt_ms: Option<f64>,
pub upstream_engine_service_tbt_ms: Option<f64>,
pub latency_ms: Option<u64>,
pub admission_decision_ms: Option<u64>,
pub account_selection_wait_ms: Option<u64>,
@@ -651,7 +709,7 @@ pub struct UsageListRecord {
}
/// 一次完整模型请求的公共观测记录
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct UsageRecord {
pub id: String,
pub client_api_key_ref: String,
@@ -710,6 +768,13 @@ pub struct UsageRecord {
pub first_text_ms: Option<u64>,
pub first_token_ms: Option<u64>,
pub provider_processing_ms: Option<u64>,
pub upstream_response_ms: Option<u64>,
pub upstream_api_overhead_ms: Option<f64>,
pub upstream_engine_ms: Option<f64>,
pub upstream_engine_iapi_ttft_ms: Option<f64>,
pub upstream_engine_service_ttft_ms: Option<f64>,
pub upstream_engine_iapi_tbt_ms: Option<f64>,
pub upstream_engine_service_tbt_ms: Option<f64>,
pub latency_ms: Option<u64>,
pub admission_decision_ms: Option<u64>,
pub account_selection_wait_ms: Option<u64>,
@@ -730,7 +795,7 @@ pub struct UsageRecord {
}
/// 用量分页结果
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct UsagePage {
pub items: Vec<UsageListRecord>,
pub current_page: u32,
@@ -776,7 +841,7 @@ pub struct UsageAttempt {
}
/// 一条请求及其全部尝试
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct UsageDetail {
pub trace: Option<serde_json::Value>,
pub related_requests: Vec<serde_json::Value>,
@@ -1,7 +1,7 @@
//! 插件管理页面、资源访问与授权回调的运行时端口
use async_trait::async_trait;
use gateway_core::runtime::extensions::ExtensionSetReference;
use gateway_core::routing::extensions::ExtensionSetReference;
use crate::model::{
AdminError,
@@ -26,7 +26,7 @@ use crate::model::{
PluginCompatibilityRequirements, PluginSource,
},
};
use gateway_core::runtime::extensions::ExtensionSetReference;
use gateway_core::routing::extensions::ExtensionSetReference;
/// Admin 准备完整候选并保活到提交后的发布结束;进程状态不能代替持久启用状态
#[async_trait]
@@ -175,12 +175,13 @@ pub trait PluginDistribution: Send + Sync {
/// Runtime 只解释插件包格式;安装事务与来源选择归 Admin
#[async_trait]
pub trait PluginPackageInspector: Send + Sync {
/// 可解析但版本范围未经宿主承诺的诊断,不决定启动资格
/// 完整校验包并返回版本范围诊断;校验失败与可加载但未经宿主承诺的提醒分开
async fn compatibility_warning(
&self,
_archive: Arc<[u8]>,
_expected_sha256: String,
archive: Arc<[u8]>,
expected_sha256: String,
) -> Result<Option<String>, AdminError> {
self.inspect(archive, Some(expected_sha256)).await?;
Ok(None)
}
+19 -17
View File
@@ -12,8 +12,9 @@ use super::backup::BackupStorePorts;
use crate::model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation,
AccountGroupOptionsPage, AccountGroupPage, DeleteAccountGroup, NewAccountGroup,
SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountPageItem, AccountRuntimeSnapshot,
@@ -27,9 +28,9 @@ use crate::model::{
UpdateClientKey, UpdateClientKeyBudgetLimits,
},
observability::{
DashboardObservation, DashboardRuntimeSlots, DiagnosticDimension, DiagnosticsObservation,
OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageCalculatedBillingFact,
UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
DashboardObservation, DashboardQuery, DiagnosticDimension, DiagnosticsObservation,
Granularity, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange,
UsageCalculatedBillingFact, UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
},
provider_credentials::{
AuthorizationCommit, CredentialDetails, CredentialImportCommit, CredentialImportResult,
@@ -385,6 +386,11 @@ pub trait AccountGroupStore: Send + Sync {
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage>;
async fn list_account_group_options(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupOptionsPage>;
async fn load_account_group_members(
&self,
group_ids: &[gateway_core::routing::AccountGroupId],
@@ -425,27 +431,21 @@ pub trait ObservabilityStore: Send + Sync {
/// 返回历史统计区间和指定观测时刻下的实时账号状态
async fn dashboard_summary(
&self,
range: TimeRange,
query: DashboardQuery,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation>;
/// 返回 Dashboard 可选的实时槽位事实
///
/// 该状态来自可丢失的运行时存储;无实现或运行时存储不可用时返回 `None`,不影响
/// 持久观测数据的读取
async fn dashboard_runtime_slots(
async fn dashboard_trend(
&self,
_observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<DashboardRuntimeSlots>> {
Ok(None)
}
async fn dashboard_trend(&self, range: TimeRange) -> AdminStoreResult<Vec<RequestMetricPoint>>;
range: TimeRange,
granularity: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>>;
async fn usage_trend(
&self,
range: TimeRange,
filter: UsageFilter,
granularity: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>>;
/// 流式返回可由 Provider 重新校验的已计算费用事实,不保证顺序
@@ -454,6 +454,7 @@ pub trait ObservabilityStore: Send + Sync {
&self,
range: TimeRange,
filter: UsageFilter,
granularity: Granularity,
) -> UsageCalculatedBillingStream<'_>;
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage>;
@@ -471,6 +472,7 @@ pub trait ObservabilityStore: Send + Sync {
range: TimeRange,
filter: UsageFilter,
dimension: DiagnosticDimension,
limit: u16,
) -> AdminStoreResult<DiagnosticsObservation>;
async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage>;
@@ -8,11 +8,11 @@ use std::{future::Future, sync::Arc};
use futures::future::BoxFuture;
use gateway_core::{
engine::middleware::FrozenMiddlewarePlan,
middleware::{
compose,
engine::middleware::{
FrozenMiddlewarePlan,
service::{Context, Error, Value},
},
middleware::compose,
};
use crate::model::{AdminError, AdminErrorKind};
@@ -2,7 +2,7 @@
use crate::model::AdminError;
use futures::future::BoxFuture;
use gateway_core::middleware::service::{Error, Value};
use gateway_core::engine::middleware::service::{Error, Value};
use serde::{Serialize, de::DeserializeOwned};
use std::{collections::BTreeMap, sync::Arc};
@@ -19,9 +19,9 @@ use crate::{
AdminError, MutationContext,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupListQuery,
AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
CreateAccountGroup, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
AccountGroupMemberFact, AccountGroupMutation, AccountGroupOptionsPage,
AccountGroupPage, AccountGroupRecord, CreateAccountGroup, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
},
ports::store::{AccountGroupStore, AccountRuntimeStore},
@@ -33,6 +33,10 @@ use super::{map_store_error, publish_committed};
#[async_trait]
pub trait AccountGroupService: Send + Sync {
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError>;
async fn list_options(
&self,
query: AccountGroupListQuery,
) -> Result<AccountGroupOptionsPage, AdminError>;
async fn create(
&self,
context: &MutationContext,
@@ -130,6 +134,16 @@ impl AccountGroupService for DefaultAccountGroupService {
Ok(page)
}
async fn list_options(
&self,
query: AccountGroupListQuery,
) -> Result<AccountGroupOptionsPage, AdminError> {
self.store
.list_account_group_options(query)
.await
.map_err(|error| map_store_error(error, "account group options"))
}
async fn create(
&self,
context: &MutationContext,
@@ -0,0 +1,371 @@
//! 凭据变更共享的绑定校验、事务结算、失效与提交后观察
use crate::{
model::{
AdminError, MutationContext,
accounts::DeleteAccounts,
provider_credentials::{
AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult,
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
PendingAuthorizationMutation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
ProviderQuotaRequest, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
};
use gateway_core::{account::ProviderAccountId, routing::ProviderKind, runtime::SnapshotControl};
use super::{map_store_error, publish_committed};
pub(super) async fn publish_credentials_and_observe_quota(
provider: &std::sync::Arc<dyn crate::ports::provider::ProviderAdmin>,
snapshot: &dyn SnapshotControl,
revision: crate::model::Revision,
account_ids: &[ProviderAccountId],
request_id: &str,
) -> Result<(), AdminError> {
provider.account_facts_changed(account_ids).await;
publish_committed(snapshot, revision).await?;
// 凭据已经提交;额度只是可重建的观察,不能拖住管理请求或回滚提交结果
let provider = provider.clone();
let account_ids = account_ids.to_vec();
let request_id = request_id.to_owned();
tokio::spawn(async move {
for account_id in account_ids {
if let Err(error) = provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await
{
tracing::warn!(
request_id,
provider = %provider.provider_kind().as_str(),
account_id = %account_id.as_str(),
quota_error = ?error.kind(),
"initial quota observation failed"
);
}
}
});
Ok(())
}
pub(super) async fn required_credential(
accounts: &dyn AccountStore,
provider_kind: &ProviderKind,
account_id: &ProviderAccountId,
resource: &'static str,
) -> Result<CredentialDetails, AdminError> {
accounts
.credential_details(provider_kind, account_id)
.await
.map_err(|error| map_store_error(error, resource))?
.ok_or_else(|| AdminError::not_found("Provider 凭据不存在"))
}
pub(super) async fn required_plugin_credential(
accounts: &dyn AccountStore,
account_id: &ProviderAccountId,
resource: &'static str,
) -> Result<CredentialDetails, AdminError> {
accounts
.credential_details_by_id(account_id)
.await
.map_err(|error| map_store_error(error, resource))?
.ok_or_else(|| AdminError::not_found("Provider 凭据不存在"))
}
pub(super) async fn import_proxy_binding(
proxies: &dyn crate::ports::proxy::ProxyStore,
id: Option<&str>,
) -> Result<Option<crate::ports::proxy::ProxyImportReservation>, AdminError> {
let Some(id) = id else { return Ok(None) };
let reservation = proxies
.reserve_import(id)
.await
.map_err(|error| map_store_error(error, "proxy"))?;
Ok(Some(reservation))
}
pub(super) async fn pending_authorization(
accounts: &dyn AccountStore,
proxies: &dyn crate::ports::proxy::ProxyStore,
provider_kind: &ProviderKind,
command: &StartAuthorization,
resource: &'static str,
) -> Result<PendingAuthorizationMutation, AdminError> {
let proxy;
let mut proxy_id = None;
let target = match &command.reauthorization {
Some(account_id) => {
proxy = required_credential(accounts, provider_kind, account_id, resource)
.await?
.credential
.outbound_proxy;
AuthorizationMutationTarget::Reauthorize {
account_id: account_id.clone(),
}
}
None => {
use crate::model::proxies::AccountProxySelection;
proxy = match &command.outbound_proxy {
Some(AccountProxySelection::Url(value)) => Some(value.clone()),
Some(AccountProxySelection::Saved(id)) => {
let record = proxies
.get(id)
.await
.map_err(|error| map_store_error(error, "proxy"))?;
proxy_id = Some(record.id);
Some(record.proxy)
}
Some(AccountProxySelection::Direct) | None => None,
};
AuthorizationMutationTarget::Create {
name: command.name.clone(),
}
}
};
Ok(PendingAuthorizationMutation::new(
provider_kind.clone(),
target,
crate::model::provider_credentials::AuthorizationOwnerBinding::from_context(
&command.context,
),
)
.with_outbound_proxy(proxy)
.with_outbound_proxy_id(proxy_id))
}
pub(super) fn validate_prepared_import(
provider_kind: &ProviderKind,
prepared: &PreparedCredentialImport,
_resource: &'static str,
) -> Result<(), AdminError> {
if prepared.provider_kind != *provider_kind
|| prepared
.credentials
.iter()
.any(|credential| credential.provider_kind != *provider_kind)
{
return Err(AdminError::conflict("Provider 准备结果与请求范围不一致"));
}
Ok(())
}
pub(super) fn validate_prepared_rotation(
account: &crate::model::accounts::AccountRecord,
prepared: &PreparedCredentialRotation,
_resource: &'static str,
) -> Result<(), AdminError> {
validate_prepared_rotation_facts(account, prepared.facts())
}
pub(super) fn validate_prepared_rotation_facts(
account: &crate::model::accounts::AccountRecord,
facts: &crate::model::provider_credentials::PreparedCredentialRotationFacts,
) -> Result<(), AdminError> {
if facts.account_id.as_str() != account.id.as_str()
|| facts.provider_kind != account.provider_kind
{
return Err(AdminError::conflict("Provider 准备结果与当前凭据不一致"));
}
Ok(())
}
pub(super) async fn validate_authorization_commit(
provider_kind: &ProviderKind,
context: &MutationContext,
prepared: PreparedAuthorizationCommit,
resource: &'static str,
) -> Result<PreparedAuthorizationCommit, AdminError> {
if prepared.pending.provider_kind() != provider_kind
|| !prepared.pending.owner_binding().matches_context(context)
{
let validation_error = AdminError::conflict("OAuth 授权绑定已失效,请重新发起授权");
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
let matches_target = match (prepared.pending.target(), &prepared.credential) {
(
AuthorizationMutationTarget::Create { .. },
PreparedAuthorizationCredential::Create(credential),
) => credential.provider_kind == *provider_kind,
(
AuthorizationMutationTarget::Reauthorize { account_id },
PreparedAuthorizationCredential::Reauthorize(credential),
) => {
let facts = credential.facts();
facts.provider_kind == *provider_kind && facts.account_id == *account_id
}
_ => false,
};
if !matches_target {
let validation_error = AdminError::conflict("OAuth 凭据与待提交目标不一致");
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
Ok(prepared)
}
pub(super) async fn commit_authorization(
accounts: &dyn AccountStore,
prepared: PreparedAuthorizationCommit,
key: crate::model::provider_credentials::AuthorizationReceiptKey,
settings: Option<crate::model::accounts::AccountImportSettings>,
context: &MutationContext,
resource: &'static str,
) -> Result<crate::model::provider_credentials::AuthorizationCommitResult, AdminError> {
if settings.is_some()
&& matches!(
&prepared.credential,
PreparedAuthorizationCredential::Reauthorize(_)
)
{
prepared.abort().await?;
return Err(AdminError::invalid("重新授权不能修改账号设置"));
}
let crate::model::provider_credentials::AuthorizationCommitSettlement {
command,
credential_guard,
authorization_guard,
} = prepared.into_commit(settings, key);
match accounts.commit_authorization(command, context).await {
Ok(result) => {
if let Some(guard) = credential_guard
&& result.newly_committed
{
guard.finish();
}
if let Some(guard) = authorization_guard
&& let Err(error) = guard.commit().await
{
// 账号与回执已经原子提交,Redis 清理失败不能把确定的成功改写为失败
tracing::warn!(resource, settlement_error = %error, "authorization committed but pending cleanup failed");
}
Ok(result)
}
Err(error) => {
drop(credential_guard);
if let Some(guard) = authorization_guard
&& let Err(settlement_error) = guard.abort().await
{
tracing::warn!(
resource,
store_error = %error,
settlement_error = %settlement_error,
"OAuth authorization claim release failed after Store commit failure"
);
return Err(settlement_error);
}
Err(map_store_error(error, resource))
}
}
}
pub(super) async fn commit_credential_rotation(
accounts: &dyn AccountStore,
prepared: PreparedCredentialRotation,
settings: Option<crate::model::accounts::UpdateAccount>,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (facts, guard) = prepared.into_parts();
match accounts
.commit_credential_rotation(
CredentialRotationCommit {
prepared: facts,
settings,
},
context,
)
.await
{
Ok(result) => {
guard.finish();
Ok(result)
}
Err(error) => {
drop(guard);
Err(map_store_error(error, resource))
}
}
}
pub(super) async fn commit_credential_refresh(
accounts: &dyn AccountStore,
prepared: PreparedCredentialRotation,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (facts, guard) = prepared.into_parts();
match accounts
.commit_credential_refresh(
CredentialRotationCommit {
prepared: facts,
settings: None,
},
context,
)
.await
{
Ok(result) => {
guard.finish();
Ok(result)
}
Err(error) => {
drop(guard);
Err(map_store_error(error, resource))
}
}
}
pub(super) async fn delete_credentials(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
command: CredentialDeletion,
resource: &'static str,
) -> Result<CredentialDeletionResult, AdminError> {
for account_id in &command.account_ids {
required_credential(accounts, provider.provider_kind(), account_id, resource).await?;
}
let account_ids = command.account_ids;
let revision = accounts
.delete_accounts(
DeleteAccounts {
account_ids: account_ids
.iter()
.map(|account_id| account_id.as_str().to_owned())
.collect(),
},
&command.context,
)
.await
.map_err(|error| map_store_error(error, resource))?;
for account_id in &account_ids {
provider.account_unavailable(account_id).await;
}
provider.account_facts_changed(&account_ids).await;
Ok(CredentialDeletionResult {
config_revision: revision,
account_ids,
})
}
@@ -0,0 +1,67 @@
//! 管理端口错误到用例错误的统一分类与来源保留
use crate::model::{AdminError, AdminErrorKind};
use crate::ports::store::{AdminStoreError, AdminStoreErrorKind};
pub(super) fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError {
let kind = match error.kind() {
AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid,
AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound,
AdminStoreErrorKind::StaleRevision
| AdminStoreErrorKind::DuplicateName
| AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict,
AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable,
};
let message = match kind {
AdminErrorKind::Invalid => "请求参数不合法",
AdminErrorKind::NotFound => "请求的资源不存在",
AdminErrorKind::Conflict => "当前资源状态冲突,请刷新后重试",
AdminErrorKind::Unavailable => "依赖服务暂不可用",
_ => "服务内部错误",
};
AdminError::new(kind, message).with_source(OperationFailure {
operation: resource,
source: error.into(),
})
}
pub(super) fn map_provider_error(
error: crate::ports::provider::ProviderAdminError,
resource: &'static str,
) -> AdminError {
use crate::ports::provider::ProviderAdminErrorKind;
let kind = match error.kind() {
ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid,
ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid,
ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound,
ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict,
ProviderAdminErrorKind::Ambiguous => AdminErrorKind::UpstreamResultUnknown,
ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway,
ProviderAdminErrorKind::Internal => AdminErrorKind::Internal,
};
let message = match kind {
AdminErrorKind::Invalid => "Provider 请求不合法",
AdminErrorKind::NotFound => "Provider 资源不存在",
AdminErrorKind::Conflict => "Provider 资源状态冲突,请刷新后重试",
AdminErrorKind::BadGateway => "上游服务请求失败",
AdminErrorKind::UpstreamResultUnknown => "上游执行结果未知,请刷新状态后再决定是否重试",
AdminErrorKind::Unavailable => "Provider 服务暂不可用",
AdminErrorKind::Internal => "服务内部错误",
_ => "Provider 操作失败",
};
AdminError::new(kind, error.public_message().unwrap_or(message)).with_source(OperationFailure {
operation: resource,
source: error.into(),
})
}
// 用例的操作名补充端口资源上下文,底层分类与原始原因仍由来源持有
#[derive(Debug, thiserror::Error)]
#[error("admin operation failed: {operation}")]
struct OperationFailure {
operation: &'static str,
source: gateway_core::error::ErrorSource,
}
@@ -14,7 +14,9 @@ use crate::{
KeyUsageOverview, KeyUsageQuery, KeyUsageRecordKind, KeyUsageRecords,
KeyUsageRecordsQuery,
},
observability::{OpsErrorFilter, OpsErrorQuery, TimeRange, UsageFilter, UsageQuery},
observability::{
Granularity, OpsErrorFilter, OpsErrorQuery, TimeRange, UsageFilter, UsageQuery,
},
system::SystemVersion,
},
ports::store::{ClientKeyStore, ObservabilityStore},
@@ -192,9 +194,13 @@ impl KeyUsageService for DefaultKeyUsageService {
};
let (overview, trend, health_points) = futures::try_join!(
self.observations.usage_summary(query.range, filter.clone()),
self.observations.usage_trend(query.range, filter),
self.observations
.usage_trend(today, usage_filter(&id, None)),
.usage_trend(query.range, filter, Granularity::for_range(query.range)),
self.observations.usage_trend(
today,
usage_filter(&id, None),
Granularity::for_range(today)
),
)
.map_err(|error| map_store_error(error, "key usage overview"))?;
Ok(Some(KeyUsageOverview {
+12 -446
View File
@@ -19,450 +19,16 @@ pub mod proxies;
pub mod settings;
pub mod system;
use crate::{
model::{
AdminError, AdminErrorKind, MutationContext,
accounts::DeleteAccounts,
provider_credentials::{
AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult,
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
PendingAuthorizationMutation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
ProviderQuotaRequest, StartAuthorization,
},
},
ports::{
provider::ProviderAdmin,
store::{AccountStore, AdminStoreError, AdminStoreErrorKind},
},
mod credential_mutation;
mod error;
mod publication;
use credential_mutation::{
commit_authorization, commit_credential_refresh, commit_credential_rotation,
delete_credentials, import_proxy_binding, pending_authorization,
publish_credentials_and_observe_quota, required_credential, required_plugin_credential,
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
validate_prepared_rotation_facts,
};
use gateway_core::{
account::ProviderAccountId,
routing::{ConfigRevision, ProviderKind},
runtime::SnapshotControl,
};
fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError {
let kind = match error.kind() {
AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid,
AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound,
AdminStoreErrorKind::StaleRevision
| AdminStoreErrorKind::DuplicateName
| AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict,
AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable,
};
let message = match kind {
AdminErrorKind::Invalid => "请求参数不合法",
AdminErrorKind::NotFound => "请求的资源不存在",
AdminErrorKind::Conflict => "当前资源状态冲突,请刷新后重试",
AdminErrorKind::Unavailable => "依赖服务暂不可用",
_ => "服务内部错误",
};
AdminError::new(kind, message).with_source(OperationFailure {
operation: resource,
source: error.into(),
})
}
fn map_provider_error(
error: crate::ports::provider::ProviderAdminError,
resource: &'static str,
) -> AdminError {
use crate::ports::provider::ProviderAdminErrorKind;
let kind = match error.kind() {
ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid,
ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid,
ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound,
ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict,
ProviderAdminErrorKind::Ambiguous => AdminErrorKind::UpstreamResultUnknown,
ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway,
ProviderAdminErrorKind::Internal => AdminErrorKind::Internal,
};
let message = match kind {
AdminErrorKind::Invalid => "Provider 请求不合法",
AdminErrorKind::NotFound => "Provider 资源不存在",
AdminErrorKind::Conflict => "Provider 资源状态冲突,请刷新后重试",
AdminErrorKind::BadGateway => "上游服务请求失败",
AdminErrorKind::UpstreamResultUnknown => "上游执行结果未知,请刷新状态后再决定是否重试",
AdminErrorKind::Unavailable => "Provider 服务暂不可用",
AdminErrorKind::Internal => "服务内部错误",
_ => "Provider 操作失败",
};
AdminError::new(kind, error.public_message().unwrap_or(message)).with_source(OperationFailure {
operation: resource,
source: error.into(),
})
}
async fn publish_committed(
snapshot: &dyn SnapshotControl,
revision: crate::model::Revision,
) -> Result<(), AdminError> {
let revision = ConfigRevision::new(revision.get())
.map_err(|_| AdminError::internal("已提交的配置版本不合法"))?;
snapshot.publish_committed(revision).await;
Ok(())
}
async fn publish_credentials_and_observe_quota(
provider: &std::sync::Arc<dyn crate::ports::provider::ProviderAdmin>,
snapshot: &dyn SnapshotControl,
revision: crate::model::Revision,
account_ids: &[ProviderAccountId],
request_id: &str,
) -> Result<(), AdminError> {
provider.account_facts_changed(account_ids).await;
publish_committed(snapshot, revision).await?;
// 凭据已经提交;额度只是可重建的观察,不能拖住管理请求或回滚提交结果
let provider = provider.clone();
let account_ids = account_ids.to_vec();
let request_id = request_id.to_owned();
tokio::spawn(async move {
for account_id in account_ids {
if let Err(error) = provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: true,
rolling_usage: None,
})
.await
{
tracing::warn!(
request_id,
provider = %provider.provider_kind().as_str(),
account_id = %account_id.as_str(),
quota_error = ?error.kind(),
"initial quota observation failed"
);
}
}
});
Ok(())
}
async fn required_credential(
accounts: &dyn AccountStore,
provider_kind: &ProviderKind,
account_id: &ProviderAccountId,
resource: &'static str,
) -> Result<CredentialDetails, AdminError> {
accounts
.credential_details(provider_kind, account_id)
.await
.map_err(|error| map_store_error(error, resource))?
.ok_or_else(|| AdminError::not_found("Provider 凭据不存在"))
}
async fn required_plugin_credential(
accounts: &dyn AccountStore,
account_id: &ProviderAccountId,
resource: &'static str,
) -> Result<CredentialDetails, AdminError> {
accounts
.credential_details_by_id(account_id)
.await
.map_err(|error| map_store_error(error, resource))?
.ok_or_else(|| AdminError::not_found("Provider 凭据不存在"))
}
async fn import_proxy_binding(
proxies: &dyn crate::ports::proxy::ProxyStore,
id: Option<&str>,
) -> Result<Option<crate::ports::proxy::ProxyImportReservation>, AdminError> {
let Some(id) = id else { return Ok(None) };
let reservation = proxies
.reserve_import(id)
.await
.map_err(|error| map_store_error(error, "proxy"))?;
Ok(Some(reservation))
}
async fn pending_authorization(
accounts: &dyn AccountStore,
proxies: &dyn crate::ports::proxy::ProxyStore,
provider_kind: &ProviderKind,
command: &StartAuthorization,
resource: &'static str,
) -> Result<PendingAuthorizationMutation, AdminError> {
let proxy;
let mut proxy_id = None;
let target = match &command.reauthorization {
Some(account_id) => {
proxy = required_credential(accounts, provider_kind, account_id, resource)
.await?
.credential
.outbound_proxy;
AuthorizationMutationTarget::Reauthorize {
account_id: account_id.clone(),
}
}
None => {
use crate::model::proxies::AccountProxySelection;
proxy = match &command.outbound_proxy {
Some(AccountProxySelection::Url(value)) => Some(value.clone()),
Some(AccountProxySelection::Saved(id)) => {
let record = proxies
.get(id)
.await
.map_err(|error| map_store_error(error, "proxy"))?;
proxy_id = Some(record.id);
Some(record.proxy)
}
Some(AccountProxySelection::Direct) | None => None,
};
AuthorizationMutationTarget::Create {
name: command.name.clone(),
}
}
};
Ok(PendingAuthorizationMutation::new(
provider_kind.clone(),
target,
crate::model::provider_credentials::AuthorizationOwnerBinding::from_context(
&command.context,
),
)
.with_outbound_proxy(proxy)
.with_outbound_proxy_id(proxy_id))
}
fn validate_prepared_import(
provider_kind: &ProviderKind,
prepared: &PreparedCredentialImport,
_resource: &'static str,
) -> Result<(), AdminError> {
if prepared.provider_kind != *provider_kind
|| prepared
.credentials
.iter()
.any(|credential| credential.provider_kind != *provider_kind)
{
return Err(AdminError::conflict("Provider 准备结果与请求范围不一致"));
}
Ok(())
}
fn validate_prepared_rotation(
account: &crate::model::accounts::AccountRecord,
prepared: &PreparedCredentialRotation,
_resource: &'static str,
) -> Result<(), AdminError> {
validate_prepared_rotation_facts(account, prepared.facts())
}
fn validate_prepared_rotation_facts(
account: &crate::model::accounts::AccountRecord,
facts: &crate::model::provider_credentials::PreparedCredentialRotationFacts,
) -> Result<(), AdminError> {
if facts.account_id.as_str() != account.id.as_str()
|| facts.provider_kind != account.provider_kind
{
return Err(AdminError::conflict("Provider 准备结果与当前凭据不一致"));
}
Ok(())
}
async fn validate_authorization_commit(
provider_kind: &ProviderKind,
context: &MutationContext,
prepared: PreparedAuthorizationCommit,
resource: &'static str,
) -> Result<PreparedAuthorizationCommit, AdminError> {
if prepared.pending.provider_kind() != provider_kind
|| !prepared.pending.owner_binding().matches_context(context)
{
let validation_error = AdminError::conflict("OAuth 授权绑定已失效,请重新发起授权");
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
let matches_target = match (prepared.pending.target(), &prepared.credential) {
(
AuthorizationMutationTarget::Create { .. },
PreparedAuthorizationCredential::Create(credential),
) => credential.provider_kind == *provider_kind,
(
AuthorizationMutationTarget::Reauthorize { account_id },
PreparedAuthorizationCredential::Reauthorize(credential),
) => {
let facts = credential.facts();
facts.provider_kind == *provider_kind && facts.account_id == *account_id
}
_ => false,
};
if !matches_target {
let validation_error = AdminError::conflict("OAuth 凭据与待提交目标不一致");
if let Err(error) = prepared.abort().await {
tracing::warn!(
resource,
settlement_error = %error,
"OAuth authorization claim release failed after preparation validation"
);
return Err(error);
}
return Err(validation_error);
}
Ok(prepared)
}
async fn commit_authorization(
accounts: &dyn AccountStore,
prepared: PreparedAuthorizationCommit,
key: crate::model::provider_credentials::AuthorizationReceiptKey,
settings: Option<crate::model::accounts::AccountImportSettings>,
context: &MutationContext,
resource: &'static str,
) -> Result<crate::model::provider_credentials::AuthorizationCommitResult, AdminError> {
if settings.is_some()
&& matches!(
&prepared.credential,
PreparedAuthorizationCredential::Reauthorize(_)
)
{
prepared.abort().await?;
return Err(AdminError::invalid("重新授权不能修改账号设置"));
}
let crate::model::provider_credentials::AuthorizationCommitSettlement {
command,
credential_guard,
authorization_guard,
} = prepared.into_commit(settings, key);
match accounts.commit_authorization(command, context).await {
Ok(result) => {
if let Some(guard) = credential_guard
&& result.newly_committed
{
guard.finish();
}
if let Some(guard) = authorization_guard
&& let Err(error) = guard.commit().await
{
// 账号与回执已经原子提交,Redis 清理失败不能把确定的成功改写为失败
tracing::warn!(resource, settlement_error = %error, "authorization committed but pending cleanup failed");
}
Ok(result)
}
Err(error) => {
drop(credential_guard);
if let Some(guard) = authorization_guard
&& let Err(settlement_error) = guard.abort().await
{
tracing::warn!(
resource,
store_error = %error,
settlement_error = %settlement_error,
"OAuth authorization claim release failed after Store commit failure"
);
return Err(settlement_error);
}
Err(map_store_error(error, resource))
}
}
}
async fn commit_credential_rotation(
accounts: &dyn AccountStore,
prepared: PreparedCredentialRotation,
settings: Option<crate::model::accounts::UpdateAccount>,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (facts, guard) = prepared.into_parts();
match accounts
.commit_credential_rotation(
CredentialRotationCommit {
prepared: facts,
settings,
},
context,
)
.await
{
Ok(result) => {
guard.finish();
Ok(result)
}
Err(error) => {
drop(guard);
Err(map_store_error(error, resource))
}
}
}
async fn commit_credential_refresh(
accounts: &dyn AccountStore,
prepared: PreparedCredentialRotation,
context: &MutationContext,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
let (facts, guard) = prepared.into_parts();
match accounts
.commit_credential_refresh(
CredentialRotationCommit {
prepared: facts,
settings: None,
},
context,
)
.await
{
Ok(result) => {
guard.finish();
Ok(result)
}
Err(error) => {
drop(guard);
Err(map_store_error(error, resource))
}
}
}
async fn delete_credentials(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
command: CredentialDeletion,
resource: &'static str,
) -> Result<CredentialDeletionResult, AdminError> {
for account_id in &command.account_ids {
required_credential(accounts, provider.provider_kind(), account_id, resource).await?;
}
let account_ids = command.account_ids;
let revision = accounts
.delete_accounts(
DeleteAccounts {
account_ids: account_ids
.iter()
.map(|account_id| account_id.as_str().to_owned())
.collect(),
},
&command.context,
)
.await
.map_err(|error| map_store_error(error, resource))?;
for account_id in &account_ids {
provider.account_unavailable(account_id).await;
}
provider.account_facts_changed(&account_ids).await;
Ok(CredentialDeletionResult {
config_revision: revision,
account_ids,
})
}
// 用例的操作名补充端口资源上下文,底层分类与原始原因仍由来源持有
#[derive(Debug, thiserror::Error)]
#[error("admin operation failed: {operation}")]
struct OperationFailure {
operation: &'static str,
source: gateway_core::error::ErrorSource,
}
use error::{map_provider_error, map_store_error};
use publication::publish_committed;
@@ -17,14 +17,14 @@ use crate::{
AdminError,
observability::{
CostCoverage, CurrencyCost, DashboardAccountUsage, DashboardCapacity,
DashboardPeriodMetrics, DashboardResult, DecimalAmount, DiagnosticDimension,
DiagnosticsItem, DiagnosticsResult, HealthStatus, HealthTimeline, HealthTimelinePoint,
OpsErrorPage, OpsErrorQuery, ProviderBillingInput, RequestMetricPoint, RequestMetrics,
TimeRange, Trend, TrendKind, TrendPoint, TrendSummary, UsageBilling,
UsageCalculatedBillingFact, UsageDetail, UsageFilter, UsageInsights, UsageInsightsCost,
UsageInsightsCostPoint, UsageInsightsHealth, UsageInsightsHealthPoint,
UsageInsightsPerformance, UsageInsightsPerformancePoint, UsageOverview, UsagePage,
UsageQuery, UsageSummary,
DashboardPeriodMetrics, DashboardQuery, DashboardResult, DecimalAmount,
DiagnosticDimension, DiagnosticsItem, DiagnosticsResult, Granularity, HealthStatus,
HealthTimeline, HealthTimelinePoint, OpsErrorPage, OpsErrorQuery, ProviderBillingInput,
RequestMetricPoint, RequestMetrics, TimeRange, Trend, TrendKind, TrendPoint,
TrendSummary, UsageBilling, UsageCalculatedBillingFact, UsageDetail, UsageFilter,
UsageInsights, UsageInsightsCost, UsageInsightsCostPoint, UsageInsightsHealth,
UsageInsightsHealthPoint, UsageInsightsPerformance, UsageInsightsPerformancePoint,
UsageOverview, UsagePage, UsageQuery, UsageSummary,
},
provider_credentials::ProviderQuotaRequest,
},
@@ -157,12 +157,13 @@ impl DefaultObservabilityService {
range: TimeRange,
) -> Result<DashboardResult, AdminError> {
let observed_at = Utc::now();
let (mut observation, settings, runtime_slots) = futures::try_join!(
self.store.dashboard_summary(range, observed_at),
let (mut observation, settings) = futures::try_join!(
self.store
.dashboard_summary(DashboardQuery::new(range), observed_at),
self.settings.load_runtime_settings(),
self.store.dashboard_runtime_slots(observed_at),
)
.map_err(|error| map_store_error(error, "dashboard"))?;
let runtime_slots = observation.runtime_slots;
self.enrich_list_records(&mut observation.recent_requests);
self.enrich_dashboard_quotas(&mut observation.account_usage)
.await;
@@ -265,7 +266,7 @@ impl ObservabilityService for DefaultObservabilityService {
) -> Result<Trend, AdminError> {
let points = self
.store
.dashboard_trend(range)
.dashboard_trend(range, Granularity::for_range(range))
.await
.map_err(|error| map_store_error(error, "dashboard trend"))?;
trend(kind, points)
@@ -319,10 +320,11 @@ impl ObservabilityService for DefaultObservabilityService {
range: TimeRange,
filter: UsageFilter,
) -> Result<UsageInsights, AdminError> {
let granularity = Granularity::for_range(range);
let aggregates = async {
futures::try_join!(
self.store.usage_summary(range, filter.clone()),
self.store.usage_trend(range, filter.clone()),
self.store.usage_trend(range, filter.clone(), granularity),
)
.map_err(|error| map_store_error(error, "usage insights"))
};
@@ -331,7 +333,7 @@ impl ObservabilityService for DefaultObservabilityService {
recover_standard_costs(
&self.providers,
self.store
.usage_calculated_billing_facts(range, filter.clone()),
.usage_calculated_billing_facts(range, filter.clone(), granularity),
),
)?;
build_usage_insights(overview, trend, standard_costs)
@@ -345,7 +347,7 @@ impl ObservabilityService for DefaultObservabilityService {
) -> Result<DiagnosticsResult, AdminError> {
let observation = self
.store
.usage_diagnostics(range, filter, dimension)
.usage_diagnostics(range, filter, dimension, 100)
.await
.map_err(|error| map_store_error(error, "usage diagnostics"))?;
let items = observation
@@ -2,7 +2,10 @@
use std::collections::{BTreeMap, BTreeSet};
use gateway_core::{policy::ClientApiKeyId, routing::AccountGroupId};
use gateway_core::{
policy::ClientApiKeyId,
routing::{AccountGroupId, extensions::ExtensionSetReference},
};
use secrecy::ExposeSecret;
use super::PluginsService;
@@ -41,14 +44,11 @@ impl PluginsService {
.map_err(|error| map_store_error(error, "plugin"))?;
let warning = match self
.inspector
.inspect(artifact.archive.clone(), Some(digest.to_owned()))
.compatibility_warning(artifact.archive, digest.to_owned())
.await
{
Ok(_) => PackageStatus {
warning: self
.inspector
.compatibility_warning(artifact.archive, digest.to_owned())
.await?,
Ok(warning) => PackageStatus {
warning,
load_error: None,
},
Err(error) if error.kind() == AdminErrorKind::Invalid => PackageStatus {
@@ -368,25 +368,8 @@ impl PluginsService {
)
.await
.map_err(|e| map_store_error(e, "plugin"))?;
let prepared = if result.config_revision == candidate_revision
&& result.instance.revision == candidate_revision
{
prepared
} else {
snapshot.config_revision = result.config_revision;
snapshot
.instances
.retain(|item| item.id != result.instance.id);
snapshot.instances.push(result.instance.clone());
apply_replacements(&mut snapshot, &replacements);
self.preparation.prepare(snapshot).await?
};
self.preparation
.activate_state(&prepared, &result.instance)
self.publish_instance(snapshot, prepared, &result, &replacements)
.await?;
// prepared 的强引用跨过提交与唯一发布入口,防止候选在被读取之前回收
publish_committed(self.snapshots.as_ref(), result.config_revision).await?;
drop(prepared);
Ok(result)
}
@@ -463,48 +446,45 @@ impl PluginsService {
.await);
}
};
if let Err(error) = self
.preparation
.migrate_state(&migration_prepared, transition.clone())
.await
{
return Err(self
.recover_failed_state_transition(
Some(&transition.id),
&previous,
&previous_state,
disabled_revision,
// 迁移开始后的所有提交前失败都走同一恢复出口;提交成功后不再回滚旧配置
let committed = async {
self.preparation
.migrate_state(&migration_prepared, transition.clone())
.await?;
let mut current = self
.store
.load_instances()
.await
.map_err(|error| map_store_error(error, "plugin"))?;
if current.config_revision != expected_revision {
return Err(AdminError::conflict("迁移期间插件配置已发生变化"));
}
let candidate_revision = next_revision(expected_revision)?;
target.revision = candidate_revision;
current.instances.retain(|item| item.id != target.id);
current.instances.push(target.clone());
current.config_revision = candidate_revision;
apply_replacements(&mut current, replacements);
let prepared = self.preparation.prepare(current.clone()).await?;
let result = self
.store
.save_instance_replacing(
target,
expected_revision,
PluginStateCommit {
configuration: target_state,
transition_id: Some(transition.id.clone()),
},
replacements,
context,
error,
)
.await);
.await
.map_err(|error| map_store_error(error, "plugin"))?;
Ok((current, prepared, result))
}
let mut current = self
.store
.load_instances()
.await
.map_err(|error| map_store_error(error, "plugin"))?;
if current.config_revision != expected_revision {
return Err(self
.recover_failed_state_transition(
Some(&transition.id),
&previous,
&previous_state,
disabled_revision,
context,
AdminError::conflict("迁移期间插件配置已发生变化"),
)
.await);
}
let candidate_revision = next_revision(expected_revision)?;
target.revision = candidate_revision;
current.instances.retain(|item| item.id != target.id);
current.instances.push(target.clone());
current.config_revision = candidate_revision;
apply_replacements(&mut current, replacements);
let prepared = match self.preparation.prepare(current.clone()).await {
Ok(prepared) => prepared,
.await;
let (current, prepared, result) = match committed {
Ok(committed) => committed,
Err(error) => {
return Err(self
.recover_failed_state_transition(
@@ -518,53 +498,9 @@ impl PluginsService {
.await);
}
};
let result = match self
.store
.save_instance_replacing(
target,
expected_revision,
PluginStateCommit {
configuration: target_state,
transition_id: Some(transition.id.clone()),
},
replacements,
context,
)
.await
{
Ok(result) => result,
Err(error) => {
return Err(self
.recover_failed_state_transition(
Some(&transition.id),
&previous,
&previous_state,
disabled_revision,
context,
map_store_error(error, "plugin"),
)
.await);
}
};
let prepared = if result.config_revision == candidate_revision
&& result.instance.revision == candidate_revision
{
prepared
} else {
current.config_revision = result.config_revision;
current
.instances
.retain(|item| item.id != result.instance.id);
current.instances.push(result.instance.clone());
apply_replacements(&mut current, replacements);
self.preparation.prepare(current).await?
};
self.preparation
.activate_state(&prepared, &result.instance)
self.publish_instance(current, prepared, &result, replacements)
.await?;
publish_committed(self.snapshots.as_ref(), result.config_revision).await?;
drop(migration_prepared);
drop(prepared);
Ok(result)
}
@@ -615,22 +551,8 @@ impl PluginsService {
)
.await
.map_err(|error| map_store_error(error, "plugin"))?;
let prepared = if mutation.config_revision == candidate_revision
&& mutation.instance.revision == candidate_revision
{
prepared
} else {
snapshot.config_revision = mutation.config_revision;
snapshot
.instances
.retain(|item| item.id != mutation.instance.id);
snapshot.instances.push(mutation.instance.clone());
self.preparation.prepare(snapshot).await?
};
self.preparation
.activate_state(&prepared, &mutation.instance)
self.publish_instance(snapshot, prepared, &mutation, &[])
.await?;
publish_committed(self.snapshots.as_ref(), mutation.config_revision).await?;
Ok::<(), AdminError>(())
}
.await;
@@ -646,6 +568,35 @@ impl PluginsService {
}
}
/// 以事务返回的版本激活并发布,候选强引用必须覆盖整个提交后阶段
async fn publish_instance(
&self,
mut snapshot: PluginInstanceSnapshot,
prepared: ExtensionSetReference,
mutation: &PluginInstanceMutation,
replacements: &[PluginInstanceReplacement],
) -> Result<(), AdminError> {
let prepared = if mutation.config_revision == snapshot.config_revision
&& mutation.instance.revision == snapshot.config_revision
{
prepared
} else {
snapshot.config_revision = mutation.config_revision;
snapshot
.instances
.retain(|item| item.id != mutation.instance.id);
snapshot.instances.push(mutation.instance.clone());
apply_replacements(&mut snapshot, replacements);
self.preparation.prepare(snapshot).await?
};
self.preparation
.activate_state(&prepared, &mutation.instance)
.await?;
publish_committed(self.snapshots.as_ref(), mutation.config_revision).await?;
drop(prepared);
Ok(())
}
pub async fn delete_instance(
&self,
id: &str,
@@ -2,7 +2,8 @@
use std::sync::Arc;
use gateway_core::runtime::{RuntimeSnapshotHandle, extensions::ExtensionSetReference};
use gateway_core::routing::extensions::ExtensionSetReference;
use gateway_core::runtime::RuntimeSnapshotHandle;
use crate::{
model::{
@@ -0,0 +1,14 @@
//! 管理事务提交后的运行快照发布入口
use crate::model::AdminError;
use gateway_core::{routing::ConfigRevision, runtime::SnapshotControl};
pub(super) async fn publish_committed(
snapshot: &dyn SnapshotControl,
revision: crate::model::Revision,
) -> Result<(), AdminError> {
let revision = ConfigRevision::new(revision.get())
.map_err(|_| AdminError::internal("已提交的配置版本不合法"))?;
snapshot.publish_committed(revision).await;
Ok(())
}
@@ -519,9 +519,16 @@ impl DatabaseDumpPort for FakeDumpPort {
async fn cleanup_staging(
&self,
_backup_id: &str,
backup_id: &str,
) -> Result<(), gateway_admin::model::backup::BackupError> {
Ok(())
match std::fs::remove_file(self.dir.path().join(format!("{backup_id}.dump"))) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(_) => Err(gateway_admin::model::backup::BackupError::new(
gateway_admin::model::backup::code::PG_DUMP_FAILED,
"cleanup failed".to_owned(),
)),
}
}
}
@@ -55,6 +55,7 @@ async fn backup_task_runs_full_dump_upload_verify_pipeline() {
.object(&seed.object_key)
.expect("uploaded object exists");
assert_eq!(stored.as_slice(), DUMP_CONTENT);
assert!(dump.inspect_staging(&seed.id).await.unwrap().is_none());
}
#[tokio::test]
@@ -259,3 +260,67 @@ async fn timezone_change_rebases_future_cursor_without_running_old_due_backup()
task.run_cycle(&CancellationToken::new()).await.unwrap();
assert!(repository.all_records().is_empty());
}
#[tokio::test]
async fn completed_recovery_removes_staging_for_dumping_and_uploading_tasks() {
for (status, remote_exists) in [
(BackupStatus::Dumping, false),
(BackupStatus::Uploading, false),
(BackupStatus::Uploading, true),
] {
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
let dump = Arc::new(FakeDumpPort::new());
let object_store = Arc::new(FakeObjectStore::new());
let seed = BackupRecordSeed {
id: backup_id("recovery-cleanup"),
trigger_kind: BackupTriggerKind::Manual,
scheduled_at: None,
object_key: "recovery.dump".to_owned(),
expires_at: None,
};
repository.insert_backup_record(seed.clone()).await.unwrap();
let artifact = dump
.dump(DumpRequest {
backup_id: seed.id.clone(),
cancellation: CancellationToken::new(),
})
.await
.unwrap();
repository.force_status(&seed.id, BackupStatus::Dumping);
if status == BackupStatus::Uploading {
repository
.transition_status(
&seed.id,
gateway_admin::model::backup::BackupStatusTransition::try_new(
BackupStatus::Dumping,
status,
)
.unwrap(),
gateway_admin::ports::backup::StatusTransitionUpdate {
size_bytes: Some(artifact.size_bytes),
sha256: Some(artifact.sha256),
..Default::default()
},
Utc::now(),
)
.await
.unwrap();
}
if remote_exists {
object_store
.objects
.lock()
.unwrap()
.insert(seed.object_key.clone(), DUMP_CONTENT.to_vec());
}
BackupTask::new(repository.clone(), dump.clone(), object_store)
.run_cycle(&CancellationToken::new())
.await
.unwrap();
assert_eq!(repository.all_records()[0].status, BackupStatus::Completed);
assert!(
dump.inspect_staging(&seed.id).await.unwrap().is_none(),
"completed recovery must remove staging: {status:?}"
);
}
}
@@ -70,12 +70,6 @@ fn request_outcome_filter_should_preserve_known_and_bounded_other_values() {
assert!(RequestOutcome::new("future\nstate").is_err());
}
#[test]
fn time_range_should_accept_exactly_366_days() {
let end = Utc::now();
assert!(TimeRange::new(end - TimeDelta::days(366), end).is_ok());
}
#[test]
fn time_range_should_reject_more_than_366_days() {
let end = Utc::now();
@@ -10,12 +10,13 @@ use async_trait::async_trait;
use chrono::{DateTime, Duration, Utc};
use gateway_admin::{
model::{
MutationContext, PageSize, Revision,
AdminErrorKind, MutationContext, PageSize, Revision,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup, NewAccountGroup,
SetAccountGroupEnabled, UpdateAccountGroup,
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation,
AccountGroupOptionsPage, AccountGroupPage, AccountGroupRecord, AccountGroupRef,
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
accounts::AccountRuntimeSnapshot,
observability::DecimalAmount,
@@ -90,6 +91,62 @@ async fn group_query_service_enriches_only_current_page_members_with_runtime_fac
);
}
#[tokio::test]
async fn group_options_query_should_skip_member_and_runtime_enrichment() {
let groups = Arc::new(FakeGroupStore::default());
let runtime = Arc::new(FakeRuntimeStore::default());
let service = AdminHarness::new()
.account_groups(groups.clone())
.account_runtime(runtime.clone())
.build()
.await;
let page = service
.account_groups()
.list_options(AccountGroupListQuery {
page: 1,
page_size: PageSize::new(20).expect("page size"),
search: None,
enabled: None,
})
.await
.expect("list account group options");
assert_eq!(page.items[0].id.as_str(), GROUP_ID);
assert!(
groups
.requested_groups
.lock()
.expect("requested groups")
.is_empty()
);
assert!(
runtime
.requested_accounts
.lock()
.expect("requested accounts")
.is_empty()
);
}
#[tokio::test]
async fn group_options_query_should_preserve_store_failures() {
let service = AdminHarness::new().build().await;
let error = service
.account_groups()
.list_options(AccountGroupListQuery {
page: 1,
page_size: PageSize::new(20).expect("page size"),
search: None,
enabled: None,
})
.await
.expect_err("unavailable options store must fail");
assert_eq!(error.kind(), AdminErrorKind::Unavailable);
}
#[derive(Default)]
struct FakeGroupStore {
requested_groups: Mutex<Vec<String>>,
@@ -144,6 +201,24 @@ impl AccountGroupStore for FakeGroupStore {
})
}
async fn list_account_group_options(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupOptionsPage> {
Ok(AccountGroupOptionsPage {
config_revision: Revision::new(1).expect("revision"),
items: vec![AccountGroupRef {
id: group_id(),
name: "Primary".to_owned(),
color: AccountGroupColor::parse("#2563EBFF").expect("color"),
enabled: true,
}],
total: 1,
page: query.page,
page_size: query.page_size.get(),
})
}
async fn load_account_group_members(
&self,
group_ids: &[AccountGroupId],
@@ -300,7 +375,7 @@ fn unused() -> AdminStoreError {
#[tokio::test]
async fn invalid_group_fields_are_rejected_before_store_mutations() {
use gateway_admin::model::{AdminErrorKind, MutationActor, account_groups::CreateAccountGroup};
use gateway_admin::model::{MutationActor, account_groups::CreateAccountGroup};
let service = AdminHarness::new()
.account_groups(Arc::new(FakeGroupStore::default()))
.build()
@@ -1538,10 +1538,13 @@ async fn plugin_account_save_rejects_stale_authentication_kind_before_commit() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = accounts_service(provider, store.clone()).await;
let access = gateway_admin::initialize_plugin_accounts(
ProviderAdminRegistry::new([provider as Arc<dyn ProviderAdmin>]).unwrap(),
store.clone(),
Arc::new(RecordingPluginAccountPublication(events.clone())),
);
let error = services
.plugin_accounts_handle()
let error = access
.save(
PreparedPluginAccountSave::Replace {
facts: plugin_rotation_facts(&account_record("openai")),
@@ -5,7 +5,6 @@ use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use chrono::{DateTime, Duration as TimeDelta, Utc};
use gateway_admin::freeze_recovery::{FreezeRecoveryDeps, FreezeRecoveryTask};
use gateway_admin::model::MutationContext;
use gateway_admin::model::accounts::{AccountFreeze, AccountRuntimeSnapshot};
use gateway_admin::model::settings::{
@@ -18,7 +17,9 @@ use gateway_core::engine::probe::{
AccountProbe, AccountProbeError, AccountProbeRequest, AccountProbeResult,
};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::task::{ScheduledTask as _, WorkerCycleContext, WorkerId, WorkerKind};
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerId, WorkerKind, WorkerRunnable,
};
use super::AdminHarness;
use super::accounts::{FakeAccountStore, FakeProviderAdmin, account_record, events, revision};
@@ -246,7 +247,7 @@ async fn recovery_task(
settings: RuntimeSettings,
runtime: Arc<FreezeRuntimeStore>,
probe: Arc<dyn AccountProbe>,
) -> (FreezeRecoveryTask, Arc<FakeAccountStore>) {
) -> (Box<dyn ScheduledTask>, Arc<FakeAccountStore>) {
recovery_task_with_store(settings, runtime, probe).await
}
@@ -254,9 +255,9 @@ async fn recovery_task_with_store(
settings: RuntimeSettings,
runtime: Arc<FreezeRuntimeStore>,
probe: Arc<dyn AccountProbe>,
) -> (FreezeRecoveryTask, Arc<FakeAccountStore>) {
) -> (Box<dyn ScheduledTask>, Arc<FakeAccountStore>) {
let store = FakeAccountStore::new("openai", events());
let services = AdminHarness::new()
let bundle = AdminHarness::new()
.accounts(Arc::clone(&store) as Arc<dyn AccountStore>)
.account_runtime(Arc::clone(&runtime) as Arc<dyn AccountRuntimeStore>)
.settings(Arc::new(FreezeSettingsStore {
@@ -264,20 +265,32 @@ async fn recovery_task_with_store(
}))
.provider(FakeProviderAdmin::new("openai", events()))
.probe(probe)
.build()
.build_bundle()
.await;
let task = FreezeRecoveryTask::new(FreezeRecoveryDeps {
accounts: services.accounts_handle(),
store: store.clone(),
runtime,
settings: Arc::new(FreezeSettingsStore {
settings: settings.clone(),
}),
});
let task = freeze_worker(bundle);
(task, store)
}
async fn run_cycle(task: &FreezeRecoveryTask) {
fn freeze_worker(mut bundle: gateway_admin::AdminBundle) -> Box<dyn ScheduledTask> {
bundle
.take_worker_contributions()
.into_iter()
.find_map(|contribution| {
let WorkerContribution::Registration(registration) = contribution else {
return None;
};
if registration.id.kind() != WorkerKind::AccountFreezeRecovery {
return None;
}
match registration.runnable {
WorkerRunnable::Scheduled { task, .. } => Some(task),
WorkerRunnable::Daemon { .. } => None,
}
})
.expect("registered freeze recovery worker")
}
async fn run_cycle(task: &dyn ScheduledTask) {
let worker = WorkerId::try_new(WorkerKind::AccountFreezeRecovery, "test").expect("worker id");
let context = WorkerCycleContext::new(worker, None, CancellationToken::new());
task.run_cycle(context)
@@ -311,7 +324,7 @@ async fn probe_success_clears_freeze() {
)
.await;
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
assert_eq!(runtime.cleared(), vec!["acct_test".to_owned()]);
assert!(runtime.extended().is_empty());
@@ -327,7 +340,7 @@ async fn probe_failure_postpones_freeze() {
)
.await;
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
assert!(runtime.cleared().is_empty());
let extended = runtime.extended();
@@ -351,7 +364,7 @@ async fn disabled_policy_releases_due_freeze_without_probing() {
)
.await;
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
assert_eq!(runtime.cleared(), vec!["acct_test".to_owned()]);
assert!(runtime.extended().is_empty());
@@ -370,7 +383,7 @@ async fn adaptive_concurrency_lowers_limit_to_observed_peak() {
)
.await;
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
// 峰值 4 × 0.8 = 3(不低于下限 2),低于全局默认 5,应下调到 3
assert_eq!(
@@ -395,7 +408,7 @@ async fn adaptive_concurrency_never_raises_limit() {
)
.await;
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
// 峰值 40 的目标 32 高于全局默认 5;只降不升,不下发任何更新
assert!(store.update_commands().is_empty());
@@ -415,7 +428,7 @@ async fn probe_skips_freezes_far_from_expiry() {
)
.await;
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
assert!(runtime.cleared().is_empty());
assert!(runtime.extended().is_empty());
@@ -431,7 +444,7 @@ async fn disabled_accounts_are_skipped() {
let mut record = account_record("openai");
record.enabled = false;
let store = FakeAccountStore::with_account(record, events());
let services = AdminHarness::new()
let bundle = AdminHarness::new()
.accounts(Arc::clone(&store) as Arc<dyn AccountStore>)
.account_runtime(Arc::clone(&runtime) as Arc<dyn AccountRuntimeStore>)
.settings(Arc::new(FreezeSettingsStore {
@@ -439,18 +452,11 @@ async fn disabled_accounts_are_skipped() {
}))
.provider(FakeProviderAdmin::new("openai", events()))
.probe(Arc::new(SuccessfulProbe))
.build()
.build_bundle()
.await;
let task = FreezeRecoveryTask::new(FreezeRecoveryDeps {
accounts: services.accounts_handle(),
store: store.clone(),
runtime: Arc::clone(&runtime) as Arc<dyn AccountRuntimeStore>,
settings: Arc::new(FreezeSettingsStore {
settings: runtime_settings(true, true, true),
}),
});
let task = freeze_worker(bundle);
run_cycle(&task).await;
run_cycle(task.as_ref()).await;
assert!(runtime.cleared().is_empty());
assert!(runtime.extended().is_empty());
@@ -29,8 +29,9 @@ use gateway_admin::{
model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation,
AccountGroupOptionsPage, AccountGroupPage, DeleteAccountGroup, NewAccountGroup,
SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRuntimeSnapshot, AccountUpdateResult,
@@ -44,10 +45,10 @@ use gateway_admin::{
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
},
observability::{
DashboardDesktopRelease, DashboardObservation, DashboardWireAttribute,
DashboardDesktopRelease, DashboardObservation, DashboardQuery, DashboardWireAttribute,
DashboardWireProfile, DashboardWireTarget, DesktopReleaseStatus, DiagnosticDimension,
DiagnosticsObservation, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange,
UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
DiagnosticsObservation, Granularity, OpsErrorPage, OpsErrorQuery, RequestMetricPoint,
TimeRange, UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
},
provider_credentials::{
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
@@ -107,7 +108,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,
service_middleware: gateway_admin::public_service::PlanSource,
}
impl AdminHarness {
@@ -194,7 +195,10 @@ impl AdminHarness {
self
}
pub(super) fn service_middleware(mut self, source: gateway_admin::service::PlanSource) -> Self {
pub(super) fn service_middleware(
mut self,
source: gateway_admin::public_service::PlanSource,
) -> Self {
self.service_middleware = source;
self
}
@@ -421,6 +425,13 @@ impl AccountGroupStore for UnavailableAccountGroupStore {
Err(unavailable("account groups"))
}
async fn list_account_group_options(
&self,
_: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupOptionsPage> {
Err(unavailable("account group options"))
}
async fn load_account_group_members(
&self,
_: &[gateway_core::routing::AccountGroupId],
@@ -738,13 +749,17 @@ impl ClientKeyStore for UnavailableStore {
impl ObservabilityStore for UnavailableStore {
async fn dashboard_summary(
&self,
_: TimeRange,
_: DashboardQuery,
_: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
Err(unavailable("dashboard"))
}
async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult<Vec<RequestMetricPoint>> {
async fn dashboard_trend(
&self,
_: TimeRange,
_: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
Err(unavailable("dashboard trend"))
}
@@ -752,6 +767,7 @@ impl ObservabilityStore for UnavailableStore {
&self,
_: TimeRange,
_: UsageFilter,
_: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
Err(unavailable("usage trend"))
}
@@ -760,6 +776,7 @@ impl ObservabilityStore for UnavailableStore {
&self,
_: TimeRange,
_: UsageFilter,
_: Granularity,
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
Box::pin(futures::stream::once(async {
Err(unavailable("usage billing facts"))
@@ -783,6 +800,7 @@ impl ObservabilityStore for UnavailableStore {
_: TimeRange,
_: UsageFilter,
_: DiagnosticDimension,
_: u16,
) -> AdminStoreResult<DiagnosticsObservation> {
Err(unavailable("usage diagnostics"))
}
@@ -18,7 +18,7 @@ use gateway_admin::{
MutationContext, Revision,
observability::{
AccountPoolMetrics, AttemptMetrics, CostCoverage, CurrencyCost, DashboardAccountUsage,
DashboardObservation, DashboardRuntimeSlots, DiagnosticDimension,
DashboardObservation, DashboardQuery, DashboardRuntimeSlots, DiagnosticDimension,
DiagnosticObservation, DiagnosticsObservation, Granularity, HealthStatus,
LatencyPercentiles, ObservabilityPageSize, OpsErrorPage, OpsErrorQuery,
PercentileMilliseconds, RequestMetricPoint, RequestMetrics, TimeRange, TrendKind,
@@ -282,7 +282,7 @@ async fn dashboard_capacity_distinguishes_unlimited_inheritance_finite_overrides
}
#[tokio::test]
async fn dashboard_summary_should_share_one_current_observation_time_across_runtime_facts() {
async fn dashboard_summary_should_observe_current_accounts_for_a_historical_range() {
let historical_end = Utc::now() - Duration::hours(1);
let store = Arc::new(FixtureObservabilityStore::new(observation_range(
historical_end,
@@ -295,8 +295,7 @@ async fn dashboard_summary_should_share_one_current_observation_time_across_runt
.await
.expect("dashboard summary");
let (summary_observed_at, slots_observed_at) = store.observed_times();
assert_eq!(summary_observed_at, slots_observed_at);
let summary_observed_at = *store.summary_observed_at.lock().unwrap();
assert!(summary_observed_at.is_some_and(|value| value > historical_end));
}
@@ -789,7 +788,8 @@ struct FixtureObservabilityStore {
diagnostics: Mutex<DiagnosticsObservation>,
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
summary_observed_at: Mutex<Option<DateTime<Utc>>>,
slots_observed_at: Mutex<Option<DateTime<Utc>>>,
summary_query: Mutex<Option<DashboardQuery>>,
usage_granularities: Mutex<Vec<Granularity>>,
usage_records: Mutex<Vec<UsageListRecord>>,
account_usage: Mutex<Vec<DashboardAccountUsage>>,
dashboard_delay: Mutex<StdDuration>,
@@ -811,7 +811,8 @@ impl FixtureObservabilityStore {
diagnostics: Mutex::new(DiagnosticsObservation::default()),
runtime_slots: Mutex::new(None),
summary_observed_at: Mutex::new(None),
slots_observed_at: Mutex::new(None),
summary_query: Mutex::new(None),
usage_granularities: Mutex::new(Vec::new()),
usage_records: Mutex::new(Vec::new()),
account_usage: Mutex::new(Vec::new()),
dashboard_delay: Mutex::new(StdDuration::ZERO),
@@ -850,16 +851,6 @@ impl FixtureObservabilityStore {
*self.runtime_slots.lock().expect("runtime slots") = runtime_slots;
}
fn observed_times(&self) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
(
*self
.summary_observed_at
.lock()
.expect("summary observed at"),
*self.slots_observed_at.lock().expect("slots observed at"),
)
}
fn replace_usage_records(&self, records: Vec<UsageListRecord>) {
*self.usage_records.lock().expect("usage records") = records;
}
@@ -869,10 +860,11 @@ impl FixtureObservabilityStore {
impl ObservabilityStore for FixtureObservabilityStore {
async fn dashboard_summary(
&self,
range: TimeRange,
query: DashboardQuery,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
self.dashboard_summary_calls.fetch_add(1, Ordering::Relaxed);
*self.summary_query.lock().unwrap() = Some(query.clone());
*self
.summary_observed_at
.lock()
@@ -882,24 +874,21 @@ impl ObservabilityStore for FixtureObservabilityStore {
tokio::time::sleep(dashboard_delay).await;
}
Ok(DashboardObservation {
range,
range: query.range,
totals: Default::default(),
provider_accounts: AccountPoolMetrics::default(),
runtime_slots: *self.runtime_slots.lock().expect("runtime slots"),
trend: self.trend.lock().expect("trend").clone(),
account_usage: self.account_usage.lock().expect("account usage").clone(),
recent_requests: Vec::new(),
})
}
async fn dashboard_runtime_slots(
async fn dashboard_trend(
&self,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<DashboardRuntimeSlots>> {
*self.slots_observed_at.lock().expect("slots observed at") = Some(observed_at);
Ok(*self.runtime_slots.lock().expect("runtime slots"))
}
async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult<Vec<RequestMetricPoint>> {
_: TimeRange,
_: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
Ok(self.trend.lock().expect("trend").clone())
}
@@ -907,7 +896,9 @@ impl ObservabilityStore for FixtureObservabilityStore {
&self,
_: TimeRange,
_: UsageFilter,
granularity: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
self.usage_granularities.lock().unwrap().push(granularity);
Ok(self.trend.lock().expect("trend").clone())
}
@@ -915,7 +906,9 @@ impl ObservabilityStore for FixtureObservabilityStore {
&self,
_: TimeRange,
_: UsageFilter,
granularity: Granularity,
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
self.usage_granularities.lock().unwrap().push(granularity);
let facts = self
.calculated_billing_facts
.lock()
@@ -954,6 +947,7 @@ impl ObservabilityStore for FixtureObservabilityStore {
_: TimeRange,
_: UsageFilter,
_: DiagnosticDimension,
_: u16,
) -> AdminStoreResult<DiagnosticsObservation> {
Ok(self.diagnostics.lock().expect("diagnostics").clone())
}
@@ -1144,6 +1138,13 @@ fn total_record(
first_text_ms: None,
first_token_ms: None,
provider_processing_ms: None,
upstream_response_ms: None,
upstream_api_overhead_ms: None,
upstream_engine_ms: None,
upstream_engine_iapi_ttft_ms: None,
upstream_engine_service_ttft_ms: None,
upstream_engine_iapi_tbt_ms: None,
upstream_engine_service_tbt_ms: None,
latency_ms: Some(100),
admission_decision_ms: None,
account_selection_wait_ms: None,
@@ -1195,3 +1196,41 @@ fn quarter_hour_start(value: DateTime<Utc>) -> DateTime<Utc> {
let elapsed = value.timestamp().rem_euclid(15 * 60);
value - Duration::seconds(elapsed) - Duration::nanoseconds(i64::from(value.nanosecond()))
}
#[tokio::test]
async fn admin_selects_dashboard_defaults_and_matching_usage_buckets_at_range_boundaries() {
let now = Utc::now();
for (duration, expected) in [
(Duration::days(2), Granularity::FifteenMinutes),
(Duration::days(2) + Duration::seconds(1), Granularity::Hour),
(Duration::days(31), Granularity::Hour),
(Duration::days(31) + Duration::seconds(1), Granularity::Day),
] {
let range = TimeRange::new(now - duration, now).unwrap();
let store = Arc::new(FixtureObservabilityStore::new(range));
let services = observability_services(store.clone()).await;
services
.observability()
.dashboard_summary(range, TrendKind::Usage)
.await
.unwrap();
let query = store.summary_query.lock().unwrap().clone().unwrap();
assert_eq!(query.range, range);
assert_eq!(query.granularity, expected);
assert_eq!(query.account_limit, 4);
assert_eq!(query.recent_request_limit, 10);
assert_eq!(
query.recent_request_filter.outcome,
Some(gateway_admin::model::observability::RequestOutcome::Succeeded)
);
services
.observability()
.usage_insights(range, UsageFilter::default())
.await
.unwrap();
assert_eq!(
*store.usage_granularities.lock().unwrap(),
vec![expected, expected]
);
}
}
@@ -24,12 +24,10 @@ use gateway_admin::{
store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
},
};
use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference};
use gateway_core::{
routing::ConfigRevision,
runtime::{
RuntimeSnapshotHandle, SnapshotControl,
extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
},
runtime::{RuntimeSnapshotHandle, SnapshotControl},
};
use serde_json::json;
@@ -117,7 +115,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for Fixture {
&self,
_: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot,
_: Option<u64>,
_: Option<&gateway_core::runtime::extensions::ExtensionSetReference>,
_: Option<&gateway_core::routing::extensions::ExtensionSetReference>,
) -> Option<
std::collections::BTreeMap<
String,
@@ -132,7 +130,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for Fixture {
impl gateway_admin::ports::plugins::PluginStateLifecycle for Fixture {
async fn activate_state(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::instances::PluginInstance,
) -> Result<(), AdminError> {
Ok(())
@@ -142,7 +140,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for Fixture {
}
async fn migrate_state(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: gateway_admin::model::plugins::state::PluginStateTransition,
) -> Result<(), AdminError> {
panic!("unexpected state migration")
@@ -40,12 +40,10 @@ use gateway_admin::{
store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
},
};
use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference};
use gateway_core::{
routing::ConfigRevision,
runtime::{
RuntimeSnapshotHandle, SnapshotControl,
extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
},
runtime::{RuntimeSnapshotHandle, SnapshotControl},
};
use serde_json::json;
@@ -87,9 +85,10 @@ impl gateway_admin::ports::plugins::PluginPackageInspector for LifecycleFixture
}
async fn compatibility_warning(
&self,
_: Arc<[u8]>,
_: String,
archive: Arc<[u8]>,
digest: String,
) -> Result<Option<String>, AdminError> {
self.inspect(archive, Some(digest)).await?;
Ok(self.data.lock().unwrap().compatibility_warning.clone())
}
async fn inspect(
@@ -132,6 +131,9 @@ struct FixtureData {
saves: Vec<SavedInstance>,
history: BTreeMap<String, PluginVersionConfiguration>,
aborted: usize,
transition_active: bool,
fail_reload_after_migration: bool,
fail_next_load: bool,
activated: usize,
quiesced: usize,
diagnostics: Option<BTreeMap<String, PluginInstanceRuntime>>,
@@ -185,6 +187,9 @@ impl LifecycleFixture {
},
)]),
aborted: 0,
transition_active: false,
fail_reload_after_migration: false,
fail_next_load: false,
activated: 0,
quiesced: 0,
diagnostics: None,
@@ -237,6 +242,9 @@ impl LifecycleFixture {
.instances
.retain(|current| current.id != instance.id);
data.snapshot.instances.push(instance.clone());
if transition_id.is_some() {
data.transition_active = false;
}
data.saves.push(SavedInstance {
artifact_sha256: instance.artifact_sha256.clone(),
enabled: instance.enabled,
@@ -336,7 +344,11 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for LifecycleFixture {
_: PluginStateTransition,
) -> Result<(), AdminError> {
match self.behavior {
MigrationBehavior::Succeed => Ok(()),
MigrationBehavior::Succeed => {
let mut data = self.data.lock().unwrap();
data.fail_next_load = std::mem::take(&mut data.fail_reload_after_migration);
Ok(())
}
MigrationBehavior::Fail => Err(AdminError::invalid("migration fixture failed")),
MigrationBehavior::ConcurrentChangeThenFail => {
let mut data = self.data.lock().unwrap();
@@ -400,9 +412,10 @@ impl PluginStateStore for LifecycleFixture {
artifact_sha256: &str,
_: PluginStateConfiguration,
) -> PluginStateStoreResult<PluginStateTransition> {
let data = self.data.lock().unwrap();
let mut data = self.data.lock().unwrap();
let current = &data.snapshot.instances[0];
if current.enabled
if data.transition_active
|| current.enabled
|| current.id != instance_id
|| current.revision != expected_instance_revision
|| artifact_sha256 != NEW_ARTIFACT
@@ -411,6 +424,7 @@ impl PluginStateStore for LifecycleFixture {
PluginStateStoreErrorKind::Conflict,
));
}
data.transition_active = true;
Ok(PluginStateTransition {
id: "00000000-0000-7000-8000-000000000002".into(),
instance_id: instance_id.into(),
@@ -440,7 +454,9 @@ impl PluginStateStore for LifecycleFixture {
}
async fn abort_transition(&self, _: &str) -> PluginStateStoreResult<()> {
self.data.lock().unwrap().aborted += 1;
let mut data = self.data.lock().unwrap();
data.aborted += 1;
data.transition_active = false;
Ok(())
}
}
@@ -465,7 +481,11 @@ impl PluginStore for LifecycleFixture {
}))
}
async fn load_instances(&self) -> AdminStoreResult<PluginInstanceSnapshot> {
Ok(self.snapshot())
let mut data = self.data.lock().unwrap();
if std::mem::take(&mut data.fail_next_load) {
return Err(admin_error(AdminStoreErrorKind::Unavailable));
}
Ok(data.snapshot.clone())
}
async fn load_version_configuration(
@@ -1048,6 +1068,36 @@ async fn failed_state_migration_restores_the_previous_enabled_version_with_cas()
assert_eq!(published.revisions.lock().unwrap().as_slice(), [11, 12]);
}
#[tokio::test]
async fn reload_failure_after_migration_aborts_restores_and_allows_retry() {
let fixture = LifecycleFixture::new(MigrationBehavior::Succeed);
fixture.data.lock().unwrap().fail_reload_after_migration = true;
let published = Arc::new(Published::default());
let service = service(fixture.clone(), published.clone());
let id = "00000000-0000-7000-8000-000000000001";
let error = service
.configure_instance(Some(id), input(), &context())
.await
.err()
.unwrap();
assert_eq!(error.kind(), AdminErrorKind::Unavailable);
{
let data = fixture.data.lock().unwrap();
assert_eq!(data.aborted, 1);
assert!(!data.transition_active);
assert!(data.snapshot.instances[0].enabled);
assert_eq!(data.snapshot.instances[0].artifact_sha256, OLD_ARTIFACT);
assert_eq!(*published.revisions.lock().unwrap(), [11, 12]);
}
let result = service
.configure_instance(Some(id), input(), &context())
.await
.unwrap();
assert_eq!(result.instance.artifact_sha256, NEW_ARTIFACT);
assert!(result.instance.enabled);
assert!(!fixture.data.lock().unwrap().transition_active);
}
#[tokio::test]
async fn failed_state_migration_never_overwrites_a_concurrent_configuration_change() {
let fixture = LifecycleFixture::new(MigrationBehavior::ConcurrentChangeThenFail);
@@ -111,7 +111,7 @@ impl gateway_admin::ports::proxy::ProxyStore for TestPluginPorts {
impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts {
async fn validate_target(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
) -> Result<(), AdminError> {
Err(AdminError::not_found("unused plugin fixture"))
@@ -119,7 +119,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
async fn start_callback(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: gateway_admin::model::plugins::management::StartPluginManagementCallback,
_: &gateway_admin::model::auth::AdminRequestContext,
@@ -129,7 +129,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
}
async fn callback(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: &str,
_: gateway_admin::model::plugins::management::PluginManagementRequest,
@@ -139,14 +139,14 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
}
async fn views(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
) -> Result<Vec<gateway_admin::model::plugins::management::PluginManagementView>, AdminError>
{
Ok(Vec::new())
}
async fn resource(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: &str,
_: bool,
@@ -156,7 +156,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
}
async fn handle(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: gateway_admin::model::plugins::management::PluginManagementRequest,
) -> Result<gateway_admin::model::plugins::management::PluginManagementResponse, AdminError>
@@ -171,7 +171,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts
&self,
_: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot,
_: Option<u64>,
_: Option<&gateway_core::runtime::extensions::ExtensionSetReference>,
_: Option<&gateway_core::routing::extensions::ExtensionSetReference>,
) -> Option<
std::collections::BTreeMap<
String,
@@ -186,7 +186,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts
impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts {
async fn activate_state(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::instances::PluginInstance,
) -> Result<(), AdminError> {
panic!("unexpected state activation")
@@ -196,7 +196,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts {
}
async fn migrate_state(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: gateway_admin::model::plugins::state::PluginStateTransition,
) -> Result<(), AdminError> {
panic!("unexpected state migration")
@@ -222,7 +222,7 @@ impl gateway_admin::ports::plugins::PluginPreparation for TestPluginPorts {
async fn prepare(
&self,
_: gateway_admin::model::plugins::instances::PluginInstanceSnapshot,
) -> Result<gateway_core::runtime::extensions::ExtensionSetReference, AdminError> {
) -> Result<gateway_core::routing::extensions::ExtensionSetReference, AdminError> {
Err(AdminError::invalid("unused plugin fixture"))
}
}
@@ -247,7 +247,7 @@ fn sdk_settings_contract_matches_host_declarations() {
_ => None,
})
.collect();
let register = source(root, "src/service/settings.rs")
let register = source(root, "src/public_service/settings.rs")
.items
.into_iter()
.find_map(|item| match item {
@@ -293,7 +293,7 @@ fn sdk_settings_contract_matches_host_declarations() {
let generated = quote! {
//! 宿主设置服务的操作标识与请求、响应数据合同
//!
//! 从宿主设置类型与 service/settings.rs 生成;更新命令见 SDK 维护说明
//! 从宿主设置类型与 public_service/settings.rs 生成;更新命令见 SDK 维护说明
use super::Operation;
use serde::{Deserialize, Serialize};
@@ -28,7 +28,7 @@ struct PricingSettingsStore {
#[tokio::test]
async fn explicit_service_calls_freeze_empty_plan_and_cancel_before_entering_terminal() {
use gateway_admin::service::{Origin, Plan, scope};
use gateway_admin::public_service::{Origin, Plan, scope};
use std::sync::atomic::{AtomicUsize, Ordering};
let resolutions = Arc::new(AtomicUsize::new(0));
@@ -16,8 +16,9 @@ use gateway_admin::model::{
PageSize,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor, AccountGroupListQuery,
AccountGroupMutation, AccountGroupPage, AccountGroupRecord, AccountGroupUsage,
CreateAccountGroup, DeleteAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
AccountGroupMutation, AccountGroupOptionsPage, AccountGroupPage, AccountGroupRecord,
AccountGroupUsage, CreateAccountGroup, DeleteAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
};
use gateway_core::{account::FastMode, routing::AccountGroupId};
@@ -212,6 +213,38 @@ struct AccountGroupPageData {
config_revision: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupOptionsPageData {
items: Vec<AccountGroupRefView>,
page: PageMeta,
config_revision: u64,
}
impl From<AccountGroupOptionsPage> for AccountGroupOptionsPageData {
fn from(page: AccountGroupOptionsPage) -> Self {
let total_pages = if page.total == 0 {
0
} else {
page.total.div_ceil(u64::from(page.page_size))
};
Self {
items: page
.items
.into_iter()
.map(AccountGroupRefView::from)
.collect(),
page: PageMeta::new(
page.page,
u32::from(page.page_size),
page.total,
u32::try_from(total_pages).unwrap_or(u32::MAX),
),
config_revision: page.config_revision.get(),
}
}
}
impl From<(AccountGroupPage, crate::time::TimePresenter)> for AccountGroupPageData {
fn from((page, time): (AccountGroupPage, crate::time::TimePresenter)) -> Self {
let total_pages = if page.total == 0 {
@@ -263,6 +296,7 @@ where
{
Router::new()
.route("/api/admin/account-groups", get(list::<S>))
.route("/api/admin/account-groups/options", get(list_options::<S>))
.route("/api/admin/account-groups/create", post(create::<S>))
.route("/api/admin/account-groups/update", post(update::<S>))
.route("/api/admin/account-groups/enable", post(enable::<S>))
@@ -270,6 +304,26 @@ where
.route("/api/admin/account-groups/delete", post(delete::<S>))
}
async fn list_options<S>(
_auth: AdminAuth,
State(state): State<S>,
AdminQuery(query): AdminQuery<ListAccountGroupsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: SessionState + Send + Sync,
{
let result = state
.admin_services()
.account_groups()
.list_options(query.into_command().map_err(map_wire_error)?)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountGroupOptionsPageData::from(result)),
))
}
async fn list<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -307,8 +307,14 @@ pub(crate) fn usage_list_record_view(
first_event_ms: record.first_event_ms,
first_reasoning_ms: record.first_reasoning_ms,
first_text_ms: record.first_text_ms,
first_token_ms: record.first_token_ms,
openai_processing_ms: record.provider_processing_ms,
upstream_response_ms: record.upstream_response_ms,
upstream_api_overhead_ms: record.upstream_api_overhead_ms,
upstream_engine_ms: record.upstream_engine_ms,
upstream_engine_iapi_ttft_ms: record.upstream_engine_iapi_ttft_ms,
upstream_engine_service_ttft_ms: record.upstream_engine_service_ttft_ms,
upstream_engine_iapi_tbt_ms: record.upstream_engine_iapi_tbt_ms,
upstream_engine_service_tbt_ms: record.upstream_engine_service_tbt_ms,
},
first_token_latency_ms: record.first_token_ms,
latency_ms: record.latency_ms,
@@ -345,7 +351,6 @@ pub(crate) fn usage_record_view(
.error_message
.clone()
.unwrap_or_else(|| outcome.clone());
let first_token_display = display_duration(record.first_token_ms);
let latency_display = display_duration(record.latency_ms);
let cost_coverage = match record.cost_source.as_str() {
"provider_reported" | "calculated" => CostCoverageView {
@@ -409,15 +414,20 @@ pub(crate) fn usage_record_view(
first_event_ms: record.first_event_ms,
first_reasoning_ms: record.first_reasoning_ms,
first_text_ms: record.first_text_ms,
first_token_ms: record.first_token_ms,
openai_processing_ms: record.provider_processing_ms,
upstream_response_ms: record.upstream_response_ms,
upstream_api_overhead_ms: record.upstream_api_overhead_ms,
upstream_engine_ms: record.upstream_engine_ms,
upstream_engine_iapi_ttft_ms: record.upstream_engine_iapi_ttft_ms,
upstream_engine_service_ttft_ms: record.upstream_engine_service_ttft_ms,
upstream_engine_iapi_tbt_ms: record.upstream_engine_iapi_tbt_ms,
upstream_engine_service_tbt_ms: record.upstream_engine_service_tbt_ms,
},
attempt_index: None,
attempt_count: u64::from(record.attempt_count),
response_id: record.client_response_id,
upstream_request_id: record.upstream_request_id,
latency_ms: record.latency_ms,
first_token_ms: record.first_token_ms,
input_tokens: record.input_tokens,
output_tokens: record.output_tokens,
cached_tokens: record.cached_tokens,
@@ -441,7 +451,6 @@ pub(crate) fn usage_record_view(
costs,
cost_coverage,
first_token_latency_ms: record.first_token_ms,
first_token_latency_ms_display: first_token_display,
latency_ms_display: latency_display,
logical_outcome: outcome,
}
@@ -756,7 +765,7 @@ pub(crate) fn trend_summary_view(
ratio: None,
},
TrendSummaryView {
label: "吞吐 P50".to_owned(),
label: "速率 P50".to_owned(),
value: summary
.minimum_output_throughput_p50
.map_or_else(|| "—".to_owned(), |value| format!("{value} tok/s")),
@@ -904,6 +913,7 @@ pub(crate) fn dashboard_view(
rotation_strategy,
} = result;
let domain::DashboardObservation {
runtime_slots: _,
range,
totals,
provider_accounts,
@@ -54,15 +54,7 @@ pub struct UsageQuery {
impl UsageQuery {
/// 校验 Element Plus 风格的页码分页字段
pub fn validate_pagination(&self) -> Result<(u32, u16), WireValidationError> {
let current_page = self.current_page.unwrap_or(1);
if current_page == 0 {
return Err(WireValidationError::new("currentPage"));
}
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
if DomainPageSize::new(page_size).is_err() {
return Err(WireValidationError::new("pageSize"));
}
Ok((current_page, page_size))
validate_pagination(self.current_page, self.page_size)
}
}
@@ -131,15 +123,7 @@ pub struct OpsQuery {
impl OpsQuery {
/// 校验 Element Plus 风格的页码分页字段
pub fn validate_pagination(&self) -> Result<(u32, u16), WireValidationError> {
let current_page = self.current_page.unwrap_or(1);
if current_page == 0 {
return Err(WireValidationError::new("currentPage"));
}
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
if DomainPageSize::new(page_size).is_err() {
return Err(WireValidationError::new("pageSize"));
}
Ok((current_page, page_size))
validate_pagination(self.current_page, self.page_size)
}
}
@@ -396,3 +380,18 @@ pub(crate) fn non_empty(value: Option<String>) -> Option<String> {
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}
fn validate_pagination(
current_page: Option<u32>,
page_size: Option<u16>,
) -> Result<(u32, u16), WireValidationError> {
let current_page = current_page.unwrap_or(1);
if current_page == 0 {
return Err(WireValidationError::new("currentPage"));
}
let page_size = page_size.unwrap_or(DEFAULT_PAGE_SIZE);
if DomainPageSize::new(page_size).is_err() {
return Err(WireValidationError::new("pageSize"));
}
Ok((current_page, page_size))
}
@@ -160,7 +160,6 @@ pub struct UsageRecordView {
pub response_id: Option<String>,
pub upstream_request_id: Option<String>,
pub latency_ms: Option<u64>,
pub first_token_ms: Option<u64>,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
@@ -185,7 +184,6 @@ pub struct UsageRecordView {
pub costs: Vec<CostView>,
pub cost_coverage: CostCoverageView,
pub first_token_latency_ms: Option<u64>,
pub first_token_latency_ms_display: String,
pub latency_ms_display: String,
pub logical_outcome: String,
}
@@ -222,9 +220,21 @@ pub struct UsageLatencyDetailsView {
#[serde(skip_serializing_if = "Option::is_none")]
pub first_text_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub first_token_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub openai_processing_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_response_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_api_overhead_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_engine_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_engine_iapi_ttft_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_engine_service_ttft_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_engine_iapi_tbt_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_engine_service_tbt_ms: Option<f64>,
}
/// 单次上游尝试展示
+89 -15
View File
@@ -1,10 +1,12 @@
//! `/healthz` 对 Core、Store 与 Host 关键 worker 健康事实的聚合
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use axum::{extract::State, http::StatusCode};
use futures::future::join_all;
use gateway_core::diagnostics::{OperationalDiagnostics, OperationalFailure};
use gateway_core::error::{ErrorDetails, ErrorSource};
use gateway_core::health::{
HealthProbe, HealthState, WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource,
WorkerRuntimeState,
@@ -19,6 +21,16 @@ const HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(2);
pub(crate) struct HealthStatus {
probes: Arc<[Arc<dyn HealthProbe>]>,
workers: Arc<dyn WorkerHealthSource>,
diagnostics: Arc<dyn OperationalDiagnostics>,
failures: Arc<Mutex<Vec<HealthFailure>>>,
}
#[derive(Clone, PartialEq, Eq, Debug, thiserror::Error)]
#[error("{component} is {state}: {reason}")]
struct HealthFailure {
component: String,
state: &'static str,
reason: String,
}
impl HealthStatus {
@@ -26,31 +38,93 @@ impl HealthStatus {
pub(crate) fn new(
probes: Vec<Arc<dyn HealthProbe>>,
workers: Arc<dyn WorkerHealthSource>,
diagnostics: Arc<dyn OperationalDiagnostics>,
) -> Self {
Self {
probes: probes.into(),
workers,
diagnostics,
failures: Arc::default(),
}
}
pub(crate) async fn healthy(&self) -> bool {
if !self
let mut failures: Vec<_> = self
.workers
.snapshot()
.iter()
.filter(|worker| worker_affects_healthz(worker))
.all(worker_is_healthy)
{
return false;
}
let checks = self.probes.iter().map(|probe| probe.check());
tokio::time::timeout(HEALTH_CHECK_TIMEOUT, join_all(checks))
.await
.is_ok_and(|states| {
states
.into_iter()
.all(|state| matches!(state, HealthState::Healthy))
.into_iter()
.filter(|worker| worker_affects_healthz(worker) && !worker_is_healthy(worker))
.map(|worker| HealthFailure {
component: match worker.key {
WorkerHealthKey::Task(id) => format!("worker:{id}"),
WorkerHealthKey::Disabled(kind) => format!("worker:{kind:?}"),
},
state: "unhealthy",
reason: format!(
"state={:?}; last_error={}",
worker.state,
worker.last_error.as_deref().unwrap_or("none")
),
})
.collect();
// 逐个探针共享截止时间,保留已完成结果并指出具体超时组件
let deadline = tokio::time::Instant::now() + HEALTH_CHECK_TIMEOUT;
let checks = self.probes.iter().map(|probe| async move {
let (state, reason) = match tokio::time::timeout_at(deadline, probe.check()).await {
Ok(HealthState::Healthy) => return None,
Ok(HealthState::Degraded(reason)) => ("degraded", reason),
Ok(HealthState::Unhealthy(reason)) => ("unhealthy", reason),
Err(_) => ("timed_out", "health probe exceeded 2s deadline".to_owned()),
};
Some(HealthFailure {
component: format!("probe:{}", probe.name()),
state,
reason,
})
});
failures.extend(join_all(checks).await.into_iter().flatten());
let healthy = failures.is_empty();
let changed = {
let mut previous = self
.failures
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let changed = failures
.iter()
.filter(|failure| !previous.contains(failure))
.cloned()
.collect::<Vec<_>>();
for recovered in previous
.iter()
.filter(|old| !failures.iter().any(|new| new.component == old.component))
{
tracing::info!(component = %recovered.component, "health component recovered");
}
*previous = failures;
changed
};
// 只记录故障变化;使用既有有界观测队列,文件日志失效时仍有独立的运维证据
for failure in changed {
tracing::warn!(component = %failure.component, state = failure.state, "health check failed");
let mut event = OperationalFailure::new(
"health",
"check",
"health_check_failed",
"Health check failed",
);
event.correlation_id = Some(failure.component.clone());
event.details =
ErrorDetails::capture(Some(&ErrorSource::new(failure.clone())), None, false);
if self.diagnostics.record_failure(event).await.is_err() {
tracing::warn!("health diagnostic could not be recorded");
// 入队失败不能抑制下次探测的记录机会,也不改变本次健康判定
self.failures
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.retain(|previous| previous != &failure);
}
}
healthy
}
}
@@ -201,6 +201,20 @@ pub(super) struct RecordView {
#[derive(Default, Serialize)]
#[serde(rename_all = "camelCase")]
struct OutputTimingView {
#[serde(skip_serializing_if = "Option::is_none")]
upstream_response_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_api_overhead_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_engine_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_engine_iapi_ttft_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_engine_service_ttft_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_engine_iapi_tbt_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_engine_service_tbt_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
first_event_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
@@ -226,6 +240,13 @@ fn success_record(value: UsageListRecord, time: crate::time::TimePresenter) -> R
latency_ms: value.latency_ms,
first_token_latency_ms: value.first_token_ms,
latency_details: OutputTimingView {
upstream_response_ms: value.upstream_response_ms,
upstream_api_overhead_ms: value.upstream_api_overhead_ms,
upstream_engine_ms: value.upstream_engine_ms,
upstream_engine_iapi_ttft_ms: value.upstream_engine_iapi_ttft_ms,
upstream_engine_service_ttft_ms: value.upstream_engine_service_ttft_ms,
upstream_engine_iapi_tbt_ms: value.upstream_engine_iapi_tbt_ms,
upstream_engine_service_tbt_ms: value.upstream_engine_service_tbt_ms,
first_event_ms: value.first_event_ms,
first_reasoning_ms: value.first_reasoning_ms,
first_text_ms: value.first_text_ms,
+2 -2
View File
@@ -138,7 +138,7 @@ impl ApiBundle {
}
/// 内部调用使用同一总路由;组合根持有强引用,Runtime 只保存 Weak
pub fn dispatcher(&self) -> Arc<dyn gateway_core::middleware::http::Dispatcher> {
pub fn dispatcher(&self) -> Arc<dyn gateway_core::engine::middleware::http::Dispatcher> {
Arc::new(middleware::RouterDispatcher {
router: middleware::wrap(
self.router.clone(),
@@ -190,7 +190,7 @@ pub fn initialize(
let state = ApiState {
admin,
openai: OpenAiService::new(execution, lifecycle),
health: HealthStatus::new(probes, worker_health),
health: HealthStatus::new(probes, worker_health, diagnostics.clone()),
};
let index = config.asset_directory.join("index.html");
let mut router = Router::new()
@@ -17,23 +17,21 @@ use axum::{
http::{HeaderName, HeaderValue, Request, StatusCode},
response::{IntoResponse as _, Response},
};
use gateway_core::{
engine::middleware::FrozenMiddlewarePlan,
lifecycle::CancellationToken,
middleware::{
compose,
http::{self as contract, Settings},
},
};
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
use gateway_core::engine::middleware::http as contract;
use gateway_core::engine::middleware::http::Settings;
use gateway_core::lifecycle::CancellationToken;
use gateway_core::middleware::compose;
use http_body::{Body as HttpBody, Frame, SizeHint};
use http_body_util::BodyExt as _;
use tokio::time::Instant;
use tower::ServiceExt as _;
tokio::task_local! { static CURRENT: contract::Context; }
tokio::task_local! { static SETTINGS: Option<gateway_core::settings::RequestSettings>; }
tokio::task_local! { static SETTINGS: Option<gateway_core::routing::request_settings::RequestSettings>; }
pub(crate) fn current_settings() -> Option<gateway_core::settings::RequestSettings> {
pub(crate) fn current_settings() -> Option<gateway_core::routing::request_settings::RequestSettings>
{
SETTINGS.try_with(Clone::clone).ok().flatten()
}
@@ -52,7 +50,7 @@ pub(crate) async fn scope<T>(
}
pub(crate) type SettingsSource =
Arc<dyn Fn() -> Option<gateway_core::settings::RequestSettings> + Send + Sync>;
Arc<dyn Fn() -> Option<gateway_core::routing::request_settings::RequestSettings> + Send + Sync>;
pub(crate) struct RouterDispatcher {
pub(crate) router: Router,
@@ -60,7 +58,7 @@ pub(crate) struct RouterDispatcher {
}
impl contract::Dispatcher for RouterDispatcher {
fn request_settings(&self) -> Option<gateway_core::settings::RequestSettings> {
fn request_settings(&self) -> Option<gateway_core::routing::request_settings::RequestSettings> {
(self.settings)()
}
@@ -195,12 +193,12 @@ async fn handle(
});
request.extensions_mut().insert(guard.clone());
request.extensions_mut().insert(context.clone());
let service_context = gateway_admin::service::Origin {
let service_context = gateway_admin::public_service::Origin {
request_id: context.request_id.clone(),
call_id: context.call_id.clone(),
cancellation: context.cancellation.clone(),
extensions: context.extensions.clone(),
plan: gateway_admin::service::Plan::Frozen(plan.clone()),
plan: gateway_admin::public_service::Plan::Frozen(plan.clone()),
};
request
.extensions_mut()
@@ -220,7 +218,7 @@ async fn handle(
terminal_context,
SETTINGS.scope(
settings,
gateway_admin::service::scope(
gateway_admin::public_service::scope(
service_context,
dispatch(router, request.map(Body::new), started),
),
@@ -14,15 +14,14 @@ use futures::{
future::BoxFuture,
stream::{SplitSink, SplitStream},
};
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareError, MiddlewareHeader},
lifecycle::CancellationToken,
middleware::{
compose,
http::{self as http, upgrade},
websocket as core,
},
};
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareError;
use gateway_core::engine::middleware::MiddlewareHeader;
use gateway_core::engine::middleware::http;
use gateway_core::engine::middleware::http::upgrade;
use gateway_core::engine::middleware::websocket as core;
use gateway_core::lifecycle::CancellationToken;
use gateway_core::middleware::compose;
use http_body_util::BodyExt as _;
pub(super) struct Upgrade {
@@ -84,30 +84,6 @@ impl ClientApiKeyAuthError {
}
}
/// 从请求头提取 Bearer Client API key
///
/// # Errors
///
/// Header 缺失、Bearer 语法错误或 Key 不能作为 HTTP Bearer 值时返回稳定错误
pub fn bearer_client_api_key(headers: &HeaderMap) -> Result<&str, ClientApiKeyAuthError> {
let raw = headers
.get(AUTHORIZATION)
.ok_or(ClientApiKeyAuthError::MissingAuthorization)?
.to_str()
.map_err(|_| ClientApiKeyAuthError::MalformedAuthorization)?;
let token = raw
.strip_prefix("Bearer ")
.ok_or(ClientApiKeyAuthError::MalformedAuthorization)?
.trim();
if token.is_empty() {
return Err(ClientApiKeyAuthError::MalformedAuthorization);
}
if gateway_core::policy::PlaintextClientApiKey::validate(token).is_err() {
return Err(ClientApiKeyAuthError::InvalidKeyFormat);
}
Ok(token)
}
pub(crate) async fn authenticate_client(
service: &OpenAiService,
headers: &HeaderMap,
@@ -95,15 +95,25 @@ impl OpenAiRequestHeaders {
}
}
/// 独立端点只消费会话身份与 turn metadata,不携带 Responses 的连接状态
/// 独立端点保留业务扩展头与会话身份,不携带 Responses 的连接状态
pub(crate) fn session_context(&self) -> Map<String, Value> {
let mut context = Map::new();
insert_protocol_context(&mut context, "session_id", self.session_id.as_ref());
insert_protocol_context(&mut context, "thread_id", self.thread_id.as_ref());
insert_protocol_context(&mut context, "turn_metadata", self.turn_metadata.as_ref());
self.insert_passthrough_headers(&mut context);
context
}
fn insert_passthrough_headers(&self, context: &mut Map<String, Value>) {
if !self.passthrough_headers.is_empty() {
context.insert(
PASSTHROUGH_HEADERS_CONTEXT_KEY.to_owned(),
Value::Array(self.passthrough_headers.clone()),
);
}
}
fn protocol_context(
&self,
use_websocket: Option<bool>,
@@ -153,12 +163,7 @@ impl OpenAiRequestHeaders {
insert_protocol_context(&mut context, "turn_id", self.turn_id.as_ref());
insert_protocol_context(&mut context, "responses_lite", self.responses_lite.as_ref());
insert_protocol_context(&mut context, "memgen_request", self.memgen_request.as_ref());
if !self.passthrough_headers.is_empty() {
context.insert(
PASSTHROUGH_HEADERS_CONTEXT_KEY.to_owned(),
Value::Array(self.passthrough_headers.clone()),
);
}
self.insert_passthrough_headers(&mut context);
if let Some(use_websocket) = use_websocket {
context.insert("use_websocket".to_owned(), Value::Bool(use_websocket));
}
@@ -47,7 +47,7 @@ impl OpenAiResponsesEncoder {
return Vec::new();
};
self.observe_wire(wire);
let projected = client_failure_payload(wire);
let projected = client_payload(wire);
let data = projected.as_ref().unwrap_or_else(|| wire.data());
// 当前 Codex 不消费 Responses `error` event,会在 EOF 时丢失失败原因
// 在客户端 SSE 边界统一投影成它能识别的 `response.failed`
@@ -70,22 +70,29 @@ impl OpenAiResponsesEncoder {
{
return vec![raw_sse_frame.clone()];
}
let data = if projected.is_none() {
wire.raw_websocket_message()
.map(str::to_owned)
.unwrap_or_else(|| data.to_string())
} else {
data.to_string()
};
vec![Bytes::from(encode_sse_event_with_metadata(
effective_event_type(wire).unwrap_or_default(),
&data.to_string(),
&data,
wire.sse_id(),
wire.sse_retry(),
))]
}
/// 消费一个 Provider event,并返回 WebSocket JSON messages
/// 消费一个 Provider event,并返回完整 WebSocket 文本消息
pub fn push_websocket(&mut self, event: &ProviderEvent) -> Vec<String> {
self.observe_canonical_identity(event);
let Some(wire) = openai_wire(event).filter(|wire| wire.has_json_data()) else {
let Some(wire) = openai_wire(event) else {
return Vec::new();
};
self.observe_wire(wire);
let projected = client_failure_payload(wire);
let projected = client_payload(wire);
let data = projected.as_ref().unwrap_or_else(|| wire.data());
// WS 客户端只对它无法消费的裸 `error` 帧做投影:codex 的 WS 端点
// 会静默忽略缺少 status 且不含内置可重试错误码的 `error` 帧,客户
@@ -104,6 +111,14 @@ impl OpenAiResponsesEncoder {
{
return vec![data.to_string()];
}
if projected.is_none()
&& let Some(raw) = wire.raw_websocket_message()
{
return vec![raw.to_owned()];
}
if !wire.has_json_data() {
return Vec::new();
}
vec![data.to_string()]
}
@@ -248,15 +263,43 @@ fn effective_event_type(wire: &ProtocolWireEvent) -> Option<&str> {
.or_else(|| wire.data().get("type").and_then(Value::as_str))
}
fn client_failure_payload(wire: &ProtocolWireEvent) -> Option<Value> {
fn client_payload(wire: &ProtocolWireEvent) -> Option<Value> {
if matches!(
effective_event_type(wire),
Some("error" | "response.failed")
) {
capacity_error_for_client(wire.data())
} else {
None
return capacity_error_for_client(wire.data());
}
if !matches!(
effective_event_type(wire),
Some("response.metadata" | "codex.response.metadata")
) {
return None;
}
let headers = wire.data().get("headers")?.as_object()?;
let connection_options: Vec<String> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("connection"))
.flat_map(|(_, value)| match value {
Value::Array(values) => values.iter().filter_map(Value::as_str).collect::<Vec<_>>(),
_ => value.as_str().into_iter().collect(),
})
.flat_map(|value| value.split(','))
.map(|name| name.trim().to_ascii_lowercase())
.collect();
if headers
.keys()
.all(|name| super::response_header_is_forwardable(name, &connection_options))
{
return None;
}
// 只剥离既有响应头隔离策略禁止的字段,metadata 业务扩展保持不透明
let mut projected = wire.data().clone();
projected
.get_mut("headers")?
.as_object_mut()?
.retain(|name, _| super::response_header_is_forwardable(name, &connection_options));
Some(projected)
}
/// 判断 WS 上游的 `error` 帧是否已经是客户端可直接消费的形状
@@ -209,6 +209,11 @@ impl ResponsesDeliveryValidator {
transformed: bool,
facts: &ResponseValidationFacts,
) -> Result<(), ResponseValidationError> {
// 未改写的未知载荷不参与本地关联判断,不能因旁路解析失败截断原流
// 插件实际改写的载荷仍须满足交付合同
if parsed.is_err() && !transformed {
return Ok(());
}
if !self.active && !transformed {
if parsed.and_then(|frame| self.state.apply(frame)).is_err() {
self.invalid_prefix = true;
@@ -253,7 +258,7 @@ impl ResponsesDeliveryValidator {
return Err(ResponseValidationError);
}
}
ParsedFrame::Event { .. } | ParsedFrame::Done => {}
ParsedFrame::Event { .. } | ParsedFrame::Done | ParsedFrame::NonEvent => {}
}
self.state.apply(parsed)
}
@@ -292,6 +297,7 @@ impl ResponsesDeliveryValidator {
enum ParsedFrame {
Event { event_type: String, value: Value },
Done,
NonEvent,
}
impl ParsedFrame {
@@ -301,6 +307,10 @@ impl ParsedFrame {
return Ok(Self::Done);
}
let mut events = parse_sse_events(frame).map_err(|_| ResponseValidationError)?;
// 注释和仅含 SSE 元数据的帧不承载响应业务事实
if events.is_empty() {
return Ok(Self::NonEvent);
}
if events.len() != 1 {
return Err(ResponseValidationError);
}
@@ -363,6 +373,9 @@ struct DeliveryState {
impl DeliveryState {
fn apply(&mut self, frame: ParsedFrame) -> Result<(), ResponseValidationError> {
if matches!(frame, ParsedFrame::NonEvent) {
return Ok(());
}
if self.done {
return Err(ResponseValidationError);
}
@@ -0,0 +1,171 @@
//! 下游连接的帧、写入结果、退出原因与队列预算
use std::{sync::Arc, time::Duration};
use thiserror::Error;
use tokio::sync::OwnedSemaphorePermit;
pub(super) const OUTBOUND_COMMAND_BUFFER: usize = 32;
pub(super) const INBOUND_EVENT_BUFFER: usize = 32;
pub(super) const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
pub(super) const DOWNSTREAM_CLOSE_TIMEOUT: Duration = Duration::from_secs(1);
pub(super) const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60);
pub(super) const DOWNSTREAM_PING_INTERVAL: Duration = Duration::from_secs(25);
/// WebSocket pump 的写入与生命周期预算
#[derive(Clone, Copy)]
pub struct ConnectionConfig {
pub(super) write_timeout: Duration,
pub(super) max_age: Duration,
}
impl ConnectionConfig {
/// 生产环境固定预算
pub const PRODUCTION: Self = Self {
write_timeout: DOWNSTREAM_WRITE_TIMEOUT,
max_age: CONNECTION_MAX_AGE,
};
}
/// 业务层可观察的客户端输入;Ping/Pong 始终由 pump 消费
pub enum ConnectionEvent {
Text(String),
Binary,
Expired,
Exited(PumpExitReason),
}
/// 接收队列与活动期间暂存的请求共用容量;移动帧不释放它占用的名额
pub(in super::super) struct PendingConnectionEvent {
pub(in super::super) event: ConnectionEvent,
pub(super) _permit: Option<OwnedSemaphorePermit>,
}
impl PendingConnectionEvent {
pub(super) fn exited(reason: PumpExitReason) -> Self {
Self {
event: ConnectionEvent::Exited(reason),
_permit: None,
}
}
}
/// 下游写入阶段;名称刻意使用 write,而不是暗示客户端已消费的 delivery
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FramePhase {
Metadata,
First,
Data,
Terminal,
FirstAndTerminal,
Error,
ConnectionLimit,
Close,
}
impl FramePhase {
pub(super) const fn as_str(self) -> &'static str {
match self {
Self::Metadata => "metadata",
Self::First => "first",
Self::Data => "data",
Self::Terminal => "terminal",
Self::FirstAndTerminal => "first_and_terminal",
Self::Error => "error",
Self::ConnectionLimit => "connection_limit",
Self::Close => "close",
}
}
pub(super) const fn is_milestone(self) -> bool {
matches!(
self,
Self::First
| Self::Terminal
| Self::FirstAndTerminal
| Self::Error
| Self::ConnectionLimit
| Self::Close
)
}
}
/// 一次下游写入的请求归属与协议阶段
#[derive(Clone)]
pub struct WriteContext {
pub(super) request_id: Option<Arc<str>>,
pub(super) phase: FramePhase,
}
impl WriteContext {
/// 创建归属于某个请求的写入上下文
#[must_use]
pub fn request(request_id: &Arc<str>, phase: FramePhase) -> Self {
Self {
request_id: Some(Arc::clone(request_id)),
phase,
}
}
/// 创建连接级写入上下文
#[must_use]
pub const fn connection(phase: FramePhase) -> Self {
Self {
request_id: None,
phase,
}
}
}
/// WebSocket pump 停止的稳定原因
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PumpExitReason {
ClientClose,
PeerEof,
ReadError,
WriteError,
WriteTimeout,
LifecycleShutdown,
ConnectionMaxAge,
CoordinatorDropped,
InboundOverload,
ServerClose,
PumpStopped,
MiddlewareFailed,
}
impl PumpExitReason {
pub(in super::super) const fn as_str(self) -> &'static str {
match self {
Self::ClientClose => "client_close",
Self::PeerEof => "peer_eof",
Self::ReadError => "read_error",
Self::WriteError => "write_error",
Self::WriteTimeout => "write_timeout",
Self::LifecycleShutdown => "lifecycle_shutdown",
Self::ConnectionMaxAge => "connection_max_age",
Self::CoordinatorDropped => "coordinator_dropped",
Self::InboundOverload => "inbound_overload",
Self::ServerClose => "server_close",
Self::PumpStopped => "pump_stopped",
Self::MiddlewareFailed => "middleware_failed",
}
}
}
/// 下游 WebSocket 写入失败
#[derive(Debug, Error)]
pub enum ConnectionWriteError {
#[error("downstream WebSocket pump is closed")]
Closed,
#[error("downstream WebSocket write timed out after {timeout:?}")]
Timeout { timeout: Duration },
#[error("downstream WebSocket transport write failed: {message}")]
Transport { message: String },
}
/// 只有实际 transport 写入才记为 Written;插件丢弃消息不伪造写入成功
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WriteOutcome {
Written,
Suppressed,
}
@@ -0,0 +1,362 @@
//! 协调层持有的连接句柄与唯一 pump 启动入口
use super::{
frame::*,
pump::run_pump,
state::{ConnectionCommand, ConnectionStats},
};
use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
use futures::{Sink, Stream};
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader},
lifecycle::CancellationToken,
};
use std::{
collections::VecDeque,
fmt,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use tokio::{
sync::{mpsc, oneshot},
task::JoinHandle,
time::{Instant, timeout},
};
/// 协调层持有的单 owner WebSocket 连接句柄
pub struct ResponsesWebSocketConnection {
connection_id: Arc<str>,
opened_at: Instant,
expired: Arc<AtomicBool>,
commands: Option<mpsc::Sender<ConnectionCommand>>,
incoming: mpsc::Receiver<PendingConnectionEvent>,
deferred: VecDeque<PendingConnectionEvent>,
exited: oneshot::Receiver<PumpExitReason>,
pump_task: Option<JoinHandle<()>>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
exit_reason: Option<PumpExitReason>,
}
impl ResponsesWebSocketConnection {
pub(in super::super) fn new(
socket: WebSocket,
connection_id: String,
cancellation: CancellationToken,
middleware: Option<FrozenMiddlewarePlan>,
headers: Arc<[MiddlewareHeader]>,
) -> Self {
spawn_connection(
socket,
Arc::<str>::from(connection_id),
cancellation,
ConnectionConfig::PRODUCTION,
middleware,
headers,
)
}
pub(in super::super) fn id(&self) -> &str {
&self.connection_id
}
/// 返回连接是否已经达到生命周期上限
#[must_use]
pub fn is_expired(&self) -> bool {
self.expired.load(Ordering::Acquire)
}
pub(in super::super) fn age(&self) -> Duration {
self.opened_at.elapsed()
}
/// 等待下一个需要业务层处理的客户端事件
pub async fn next_event(&mut self) -> Option<ConnectionEvent> {
if self.exit_reason.is_none() {
match self.exited.try_recv() {
Ok(reason) => {
self.exit_reason = Some(reason);
}
Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
self.exit_reason = Some(PumpExitReason::PumpStopped);
}
Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {}
}
}
if let Some(reason) = self.exit_reason {
return Some(ConnectionEvent::Exited(reason));
}
if let Some(event) = self.deferred.pop_front() {
return Some(event.event);
}
self.next_active_event().await.map(|event| event.event)
}
/// 活动响应期间只消费新到的帧,已排队的下一轮请求继续保持串行
pub(in super::super) async fn next_active_event(&mut self) -> Option<PendingConnectionEvent> {
if let Some(reason) = self.exit_reason {
return Some(PendingConnectionEvent::exited(reason));
}
let event = tokio::select! {
biased;
reason = &mut self.exited => {
Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
}
event = self.incoming.recv() => event,
};
if let Some(PendingConnectionEvent {
event: ConnectionEvent::Exited(reason),
..
}) = event.as_ref()
{
self.exit_reason.get_or_insert(*reason);
} else if event.is_none() {
self.exit_reason.get_or_insert(PumpExitReason::PumpStopped);
}
event
}
pub(in super::super) fn defer(&mut self, event: PendingConnectionEvent) {
self.deferred.push_back(event);
}
/// 等待连接退出,不消费留给后续串行请求的业务帧
pub async fn wait_for_exit(&mut self) -> PumpExitReason {
if let Some(reason) = self.exit_reason {
return reason;
}
let reason = (&mut self.exited)
.await
.unwrap_or(PumpExitReason::PumpStopped);
self.exit_reason = Some(reason);
reason
}
/// 串行写入文本帧,并等待 pump 确认 transport 写入结果
///
/// # Errors
///
/// pump 已关闭、写入超时或 transport 失败时返回稳定错误
pub async fn send_text(
&mut self,
payload: String,
context: WriteContext,
) -> Result<WriteOutcome, ConnectionWriteError> {
self.send(Message::Text(payload.into()), context).await
}
pub(in super::super) async fn close_policy(
&mut self,
reason: &'static str,
request_id: Option<&Arc<str>>,
) {
let context = request_id.map_or_else(
|| WriteContext::connection(FramePhase::Close),
|request_id| WriteContext::request(request_id, FramePhase::Close),
);
self.close(
CloseFrame {
code: close_code::POLICY,
reason: reason.into(),
},
context,
PumpExitReason::ServerClose,
)
.await;
}
pub(in super::super) async fn close_for_connection_limit(&mut self, reason: &'static str) {
self.close(
CloseFrame {
code: close_code::NORMAL,
reason: reason.into(),
},
WriteContext::connection(FramePhase::Close),
PumpExitReason::ConnectionMaxAge,
)
.await;
}
async fn close(
&mut self,
frame: CloseFrame,
context: WriteContext,
exit_reason: PumpExitReason,
) {
if self
.send(Message::Close(Some(frame)), context)
.await
.is_ok()
{
self.exit_reason.get_or_insert(exit_reason);
}
}
async fn send(
&mut self,
message: Message,
context: WriteContext,
) -> Result<WriteOutcome, ConnectionWriteError> {
let Some(commands) = self.commands.as_ref().cloned() else {
return Err(ConnectionWriteError::Closed);
};
self.stats.observe_command_queue(&commands);
let (acknowledged, acknowledgement) = oneshot::channel();
let started_at = Instant::now();
let request_id = context.request_id.clone();
let phase = context.phase;
let command = ConnectionCommand {
message,
context,
acknowledged,
};
let write = async move {
commands
.send(command)
.await
.map_err(|_| ConnectionWriteError::Closed)?;
acknowledgement
.await
.map_err(|_| ConnectionWriteError::Closed)?
};
let result = match timeout(self.config.write_timeout, write).await {
Ok(result) => result,
Err(_) => Err(ConnectionWriteError::Timeout {
timeout: self.config.write_timeout,
}),
};
let duration = started_at.elapsed();
match &result {
Ok(WriteOutcome::Written) if phase.is_milestone() => tracing::info!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
"Responses WebSocket frame write succeeded"
),
Ok(outcome) => tracing::debug!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
?outcome,
"Responses WebSocket frame processed"
),
Err(error) => tracing::info!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
error = %error,
"Responses WebSocket frame write failed"
),
}
if let Err(error) = &result {
let reason = match error {
ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout,
ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError,
ConnectionWriteError::Closed => PumpExitReason::PumpStopped,
};
self.terminate(reason);
}
result
}
fn terminate(&mut self, reason: PumpExitReason) {
self.exit_reason.get_or_insert(reason);
self.commands.take();
if let Some(task) = self.pump_task.take() {
task.abort();
}
}
pub(in super::super) fn log_summary(&self, request_count: u64) {
let connection_age_ms = self.age().as_millis();
let read_idle_ms = connection_age_ms
.saturating_sub(u128::from(self.stats.last_read_ms.load(Ordering::Relaxed)));
let write_idle_ms = connection_age_ms
.saturating_sub(u128::from(self.stats.last_write_ms.load(Ordering::Relaxed)));
tracing::info!(
websocket_connection_id = %self.connection_id,
request_count,
connection_age_ms,
pump_exit_reason = self
.exit_reason
.map_or("unknown", PumpExitReason::as_str),
command_queue_high_water = self
.stats
.command_queue_high_water
.load(Ordering::Relaxed),
command_backpressure_count = self
.stats
.command_backpressure_count
.load(Ordering::Relaxed),
ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed),
ping_written_count = self.stats.ping_written_count.load(Ordering::Relaxed),
pong_received_count = self.stats.pong_received_count.load(Ordering::Relaxed),
read_idle_ms,
write_idle_ms,
"Responses WebSocket disconnected"
);
}
}
impl Drop for ResponsesWebSocketConnection {
fn drop(&mut self) {
self.commands.take();
if let Some(task) = self.pump_task.take() {
task.abort();
}
}
}
/// 为 WebSocket transport 启动包含冻结中间件计划的单 owner pump
pub fn spawn_connection<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,
{
let opened_at = Instant::now();
let expired = Arc::new(AtomicBool::new(false));
let stats = Arc::new(ConnectionStats::default());
let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER);
let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER);
let (exit_tx, exit_rx) = oneshot::channel();
let pump_task = tokio::spawn(run_pump(
socket,
command_rx,
incoming_tx,
exit_tx,
Arc::clone(&connection_id),
cancellation,
opened_at,
Arc::clone(&expired),
Arc::clone(&stats),
config,
middleware,
headers,
));
ResponsesWebSocketConnection {
connection_id,
opened_at,
expired,
commands: Some(command_tx),
incoming: incoming_rx,
deferred: VecDeque::new(),
exited: exit_rx,
pump_task: Some(pump_task),
stats,
config,
exit_reason: None,
}
}
@@ -1,11 +1,13 @@
//! 插件消息组合在 transport 锁之外执行,主动发送复用唯一 writer
use super::{ConnectionWriteError, PumpContext, write_message};
use super::{ConnectionWriteError, frame::DOWNSTREAM_CLOSE_TIMEOUT, state::PumpContext};
pub(super) use crate::middleware::websocket::CancelOnDrop;
use crate::middleware::websocket::{from_core, into_core};
use axum::extract::ws::Message;
use futures::{Sink, future::BoxFuture};
use gateway_core::{engine::middleware::MiddlewareError, middleware::websocket as core};
use std::sync::Arc;
use futures::{Sink, SinkExt, future::BoxFuture};
use gateway_core::engine::middleware::{MiddlewareError, websocket as core};
use gateway_core::lifecycle::CancellationToken;
use std::{fmt, sync::Arc};
use tokio::time::timeout;
pub(super) struct Writer<S> {
socket: tokio::sync::Mutex<S>,
@@ -37,6 +39,36 @@ where
}) => result.unwrap_or(Err(ConnectionWriteError::Timeout { timeout: self.context.config.write_timeout })),
}
}
pub(super) async fn finish_peer_close(
&self,
cancellation: &CancellationToken,
) -> Result<(), ConnectionWriteError> {
// 主动 Sender 也可能在 pump 之外持锁;先独占 transport 再取消业务,
// 有在途写或已取消的半帧时立即跳过,不能等待它释放锁后继续刷新
let socket = self
.socket
.try_lock()
.ok()
.filter(|_| !self.context.cancellation.is_cancelled());
self.context.cancellation.cancel();
let Some(mut socket) = socket else {
return Err(ConnectionWriteError::Closed);
};
// 收到 Close 后 transport 已排队应答;只刷新它,不能再发送一份 Close
// 业务调用已取消,清理只受宿主取消与独立短预算约束,不沿用业务写入的长超时
tokio::select! {
biased;
() = cancellation.cancelled() => Err(ConnectionWriteError::Closed),
result = timeout(DOWNSTREAM_CLOSE_TIMEOUT, async {
socket.flush().await.map_err(|error| {
ConnectionWriteError::Transport { message: error.to_string() }
})
}) => result.unwrap_or(Err(ConnectionWriteError::Timeout {
timeout: DOWNSTREAM_CLOSE_TIMEOUT,
})),
}
}
}
impl<S, E> core::Sender for Writer<S>
where
@@ -81,3 +113,25 @@ pub(super) async fn transform(
.map(from_core)
.transpose()
}
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 }),
}
}
}
}
@@ -1,826 +1,13 @@
//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界
//! 下游客户端 WebSocket 的连接合同、句柄与单 owner pump
mod frame;
mod handle;
mod middleware;
mod pump;
mod state;
use std::{
collections::VecDeque,
fmt,
sync::{
Arc,
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
},
time::Duration,
pub use frame::{
ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase, PumpExitReason,
WriteContext, WriteOutcome,
};
use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
use futures::{Sink, SinkExt, Stream, StreamExt};
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader},
lifecycle::CancellationToken,
middleware::websocket as contract,
};
use thiserror::Error;
use tokio::{
sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot},
task::JoinHandle,
time::{Instant, timeout},
};
const OUTBOUND_COMMAND_BUFFER: usize = 32;
const INBOUND_EVENT_BUFFER: usize = 32;
const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60);
const DOWNSTREAM_PING_INTERVAL: Duration = Duration::from_secs(25);
/// WebSocket pump 的写入与生命周期预算
#[derive(Clone, Copy)]
pub struct ConnectionConfig {
write_timeout: Duration,
max_age: Duration,
}
impl ConnectionConfig {
/// 生产环境固定预算
pub const PRODUCTION: Self = Self {
write_timeout: DOWNSTREAM_WRITE_TIMEOUT,
max_age: CONNECTION_MAX_AGE,
};
}
/// 业务层可观察的客户端输入;Ping/Pong 始终由 pump 消费
pub enum ConnectionEvent {
Text(String),
Binary,
Expired,
Exited(PumpExitReason),
}
/// 接收队列与活动期间暂存的请求共用容量;移动帧不释放它占用的名额
pub(super) struct PendingConnectionEvent {
pub(super) event: ConnectionEvent,
_permit: Option<OwnedSemaphorePermit>,
}
impl PendingConnectionEvent {
fn exited(reason: PumpExitReason) -> Self {
Self {
event: ConnectionEvent::Exited(reason),
_permit: None,
}
}
}
/// 下游写入阶段;名称刻意使用 write,而不是暗示客户端已消费的 delivery
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FramePhase {
Metadata,
First,
Data,
Terminal,
FirstAndTerminal,
Error,
ConnectionLimit,
Close,
}
impl FramePhase {
const fn as_str(self) -> &'static str {
match self {
Self::Metadata => "metadata",
Self::First => "first",
Self::Data => "data",
Self::Terminal => "terminal",
Self::FirstAndTerminal => "first_and_terminal",
Self::Error => "error",
Self::ConnectionLimit => "connection_limit",
Self::Close => "close",
}
}
const fn is_milestone(self) -> bool {
matches!(
self,
Self::First
| Self::Terminal
| Self::FirstAndTerminal
| Self::Error
| Self::ConnectionLimit
| Self::Close
)
}
}
/// 一次下游写入的请求归属与协议阶段
#[derive(Clone)]
pub struct WriteContext {
request_id: Option<Arc<str>>,
phase: FramePhase,
}
impl WriteContext {
/// 创建归属于某个请求的写入上下文
#[must_use]
pub fn request(request_id: &Arc<str>, phase: FramePhase) -> Self {
Self {
request_id: Some(Arc::clone(request_id)),
phase,
}
}
/// 创建连接级写入上下文
#[must_use]
pub const fn connection(phase: FramePhase) -> Self {
Self {
request_id: None,
phase,
}
}
}
/// WebSocket pump 停止的稳定原因
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PumpExitReason {
ClientClose,
PeerEof,
ReadError,
WriteError,
WriteTimeout,
LifecycleShutdown,
ConnectionMaxAge,
CoordinatorDropped,
InboundOverload,
ServerClose,
PumpStopped,
MiddlewareFailed,
}
impl PumpExitReason {
pub(super) const fn as_str(self) -> &'static str {
match self {
Self::ClientClose => "client_close",
Self::PeerEof => "peer_eof",
Self::ReadError => "read_error",
Self::WriteError => "write_error",
Self::WriteTimeout => "write_timeout",
Self::LifecycleShutdown => "lifecycle_shutdown",
Self::ConnectionMaxAge => "connection_max_age",
Self::CoordinatorDropped => "coordinator_dropped",
Self::InboundOverload => "inbound_overload",
Self::ServerClose => "server_close",
Self::PumpStopped => "pump_stopped",
Self::MiddlewareFailed => "middleware_failed",
}
}
}
/// 下游 WebSocket 写入失败
#[derive(Debug, Error)]
pub enum ConnectionWriteError {
#[error("downstream WebSocket pump is closed")]
Closed,
#[error("downstream WebSocket write timed out after {timeout:?}")]
Timeout { timeout: Duration },
#[error("downstream WebSocket transport write failed: {message}")]
Transport { message: String },
}
/// 只有实际 transport 写入才记为 Written;插件丢弃消息不伪造写入成功
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WriteOutcome {
Written,
Suppressed,
}
struct ConnectionCommand {
message: Message,
context: WriteContext,
acknowledged: oneshot::Sender<Result<WriteOutcome, ConnectionWriteError>>,
}
#[derive(Default)]
struct ConnectionStats {
command_queue_high_water: AtomicUsize,
command_backpressure_count: AtomicU64,
ping_received_count: AtomicU64,
ping_written_count: AtomicU64,
pong_received_count: AtomicU64,
last_read_ms: AtomicU64,
last_write_ms: AtomicU64,
}
impl ConnectionStats {
fn record_read(&self, opened_at: Instant) {
self.last_read_ms.store(
u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
}
fn record_write(&self, opened_at: Instant) {
self.last_write_ms.store(
u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
}
fn observe_command_queue(&self, sender: &mpsc::Sender<ConnectionCommand>) {
let queued = OUTBOUND_COMMAND_BUFFER
.saturating_sub(sender.capacity())
.saturating_add(1)
.min(OUTBOUND_COMMAND_BUFFER);
self.command_queue_high_water
.fetch_max(queued, Ordering::Relaxed);
if sender.capacity() == 0 {
self.command_backpressure_count
.fetch_add(1, Ordering::Relaxed);
}
}
}
/// 协调层持有的单 owner WebSocket 连接句柄
pub struct ResponsesWebSocketConnection {
connection_id: Arc<str>,
opened_at: Instant,
expired: Arc<AtomicBool>,
commands: Option<mpsc::Sender<ConnectionCommand>>,
incoming: mpsc::Receiver<PendingConnectionEvent>,
deferred: VecDeque<PendingConnectionEvent>,
exited: oneshot::Receiver<PumpExitReason>,
pump_task: Option<JoinHandle<()>>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
exit_reason: Option<PumpExitReason>,
}
impl ResponsesWebSocketConnection {
pub(super) fn new(
socket: WebSocket,
connection_id: String,
cancellation: CancellationToken,
middleware: Option<FrozenMiddlewarePlan>,
headers: Arc<[MiddlewareHeader]>,
) -> Self {
spawn_with_middleware(
socket,
Arc::<str>::from(connection_id),
cancellation,
ConnectionConfig::PRODUCTION,
middleware,
headers,
)
}
pub(super) fn id(&self) -> &str {
&self.connection_id
}
/// 返回连接是否已经达到生命周期上限
#[must_use]
pub fn is_expired(&self) -> bool {
self.expired.load(Ordering::Acquire)
}
pub(super) fn age(&self) -> Duration {
self.opened_at.elapsed()
}
/// 等待下一个需要业务层处理的客户端事件
pub async fn next_event(&mut self) -> Option<ConnectionEvent> {
if self.exit_reason.is_none() {
match self.exited.try_recv() {
Ok(reason) => {
self.exit_reason = Some(reason);
}
Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
self.exit_reason = Some(PumpExitReason::PumpStopped);
}
Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {}
}
}
if let Some(reason) = self.exit_reason {
return Some(ConnectionEvent::Exited(reason));
}
if let Some(event) = self.deferred.pop_front() {
return Some(event.event);
}
self.next_active_event().await.map(|event| event.event)
}
/// 活动响应期间只消费新到的帧,已排队的下一轮请求继续保持串行
pub(super) async fn next_active_event(&mut self) -> Option<PendingConnectionEvent> {
if let Some(reason) = self.exit_reason {
return Some(PendingConnectionEvent::exited(reason));
}
let event = tokio::select! {
biased;
reason = &mut self.exited => {
Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
}
event = self.incoming.recv() => event,
};
if let Some(PendingConnectionEvent {
event: ConnectionEvent::Exited(reason),
..
}) = event.as_ref()
{
self.exit_reason.get_or_insert(*reason);
} else if event.is_none() {
self.exit_reason.get_or_insert(PumpExitReason::PumpStopped);
}
event
}
pub(super) fn defer(&mut self, event: PendingConnectionEvent) {
self.deferred.push_back(event);
}
/// 等待连接退出,不消费留给后续串行请求的业务帧
pub async fn wait_for_exit(&mut self) -> PumpExitReason {
if let Some(reason) = self.exit_reason {
return reason;
}
let reason = (&mut self.exited)
.await
.unwrap_or(PumpExitReason::PumpStopped);
self.exit_reason = Some(reason);
reason
}
/// 串行写入文本帧,并等待 pump 确认 transport 写入结果
///
/// # Errors
///
/// pump 已关闭、写入超时或 transport 失败时返回稳定错误
pub async fn send_text(
&mut self,
payload: String,
context: WriteContext,
) -> Result<WriteOutcome, ConnectionWriteError> {
self.send(Message::Text(payload.into()), context).await
}
pub(super) async fn close_policy(
&mut self,
reason: &'static str,
request_id: Option<&Arc<str>>,
) {
let context = request_id.map_or_else(
|| WriteContext::connection(FramePhase::Close),
|request_id| WriteContext::request(request_id, FramePhase::Close),
);
self.close(
CloseFrame {
code: close_code::POLICY,
reason: reason.into(),
},
context,
PumpExitReason::ServerClose,
)
.await;
}
pub(super) async fn close_for_connection_limit(&mut self, reason: &'static str) {
self.close(
CloseFrame {
code: close_code::NORMAL,
reason: reason.into(),
},
WriteContext::connection(FramePhase::Close),
PumpExitReason::ConnectionMaxAge,
)
.await;
}
async fn close(
&mut self,
frame: CloseFrame,
context: WriteContext,
exit_reason: PumpExitReason,
) {
if self
.send(Message::Close(Some(frame)), context)
.await
.is_ok()
{
self.exit_reason.get_or_insert(exit_reason);
}
}
async fn send(
&mut self,
message: Message,
context: WriteContext,
) -> Result<WriteOutcome, ConnectionWriteError> {
let Some(commands) = self.commands.as_ref().cloned() else {
return Err(ConnectionWriteError::Closed);
};
self.stats.observe_command_queue(&commands);
let (acknowledged, acknowledgement) = oneshot::channel();
let started_at = Instant::now();
let request_id = context.request_id.clone();
let phase = context.phase;
let command = ConnectionCommand {
message,
context,
acknowledged,
};
let write = async move {
commands
.send(command)
.await
.map_err(|_| ConnectionWriteError::Closed)?;
acknowledgement
.await
.map_err(|_| ConnectionWriteError::Closed)?
};
let result = match timeout(self.config.write_timeout, write).await {
Ok(result) => result,
Err(_) => Err(ConnectionWriteError::Timeout {
timeout: self.config.write_timeout,
}),
};
let duration = started_at.elapsed();
match &result {
Ok(WriteOutcome::Written) if phase.is_milestone() => tracing::info!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
"Responses WebSocket frame write succeeded"
),
Ok(outcome) => tracing::debug!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
?outcome,
"Responses WebSocket frame processed"
),
Err(error) => tracing::info!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
error = %error,
"Responses WebSocket frame write failed"
),
}
if let Err(error) = &result {
let reason = match error {
ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout,
ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError,
ConnectionWriteError::Closed => PumpExitReason::PumpStopped,
};
self.terminate(reason);
}
result
}
fn terminate(&mut self, reason: PumpExitReason) {
self.exit_reason.get_or_insert(reason);
self.commands.take();
if let Some(task) = self.pump_task.take() {
task.abort();
}
}
pub(super) fn log_summary(&self, request_count: u64) {
let connection_age_ms = self.age().as_millis();
let read_idle_ms = connection_age_ms
.saturating_sub(u128::from(self.stats.last_read_ms.load(Ordering::Relaxed)));
let write_idle_ms = connection_age_ms
.saturating_sub(u128::from(self.stats.last_write_ms.load(Ordering::Relaxed)));
tracing::info!(
websocket_connection_id = %self.connection_id,
request_count,
connection_age_ms,
pump_exit_reason = self
.exit_reason
.map_or("unknown", PumpExitReason::as_str),
command_queue_high_water = self
.stats
.command_queue_high_water
.load(Ordering::Relaxed),
command_backpressure_count = self
.stats
.command_backpressure_count
.load(Ordering::Relaxed),
ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed),
ping_written_count = self.stats.ping_written_count.load(Ordering::Relaxed),
pong_received_count = self.stats.pong_received_count.load(Ordering::Relaxed),
read_idle_ms,
write_idle_ms,
"Responses WebSocket disconnected"
);
}
}
impl Drop for ResponsesWebSocketConnection {
fn drop(&mut self) {
self.commands.take();
if let Some(task) = self.pump_task.take() {
task.abort();
}
}
}
/// 为自动处理 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,
{
let opened_at = Instant::now();
let expired = Arc::new(AtomicBool::new(false));
let stats = Arc::new(ConnectionStats::default());
let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER);
let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER);
let (exit_tx, exit_rx) = oneshot::channel();
let pump_task = tokio::spawn(run_pump(
socket,
command_rx,
incoming_tx,
exit_tx,
Arc::clone(&connection_id),
cancellation,
opened_at,
Arc::clone(&expired),
Arc::clone(&stats),
config,
middleware,
headers,
));
ResponsesWebSocketConnection {
connection_id,
opened_at,
expired,
commands: Some(command_tx),
incoming: incoming_rx,
deferred: VecDeque::new(),
exited: exit_rx,
pump_task: Some(pump_task),
stats,
config,
exit_reason: None,
}
}
#[expect(clippy::too_many_arguments)]
async fn run_pump<S, E>(
socket: S,
commands: mpsc::Receiver<ConnectionCommand>,
incoming: mpsc::Sender<PendingConnectionEvent>,
exited: oneshot::Sender<PumpExitReason>,
connection_id: Arc<str>,
cancellation: CancellationToken,
opened_at: Instant,
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 + 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(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(
context.opened_at + DOWNSTREAM_PING_INTERVAL,
DOWNSTREAM_PING_INTERVAL,
);
heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut ping_sequence = 0_u64;
loop {
let message = tokio::select! {
biased;
() = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown,
command = commands.recv() => {
let Some(command) = command else { return PumpExitReason::CoordinatorDropped; };
let ConnectionCommand { message, context: write_context, acknowledged } = command;
let closing = matches!(message, Message::Close(_));
let result = match middleware::transform(message, contract::Direction::Outgoing, writer.clone(), context).await {
Ok(Some(message)) => writer.write(message).await.map(|()| WriteOutcome::Written),
Ok(None) => Ok(WriteOutcome::Suppressed),
Err(error) => Err(ConnectionWriteError::Transport { message: error.to_string() }),
};
let reason = result.as_ref().err().map(|error| write_exit_reason(error, context));
if let Err(error) = &result {
tracing::debug!(websocket_connection_id = %context.connection_id, request_id = write_context.request_id.as_deref().unwrap_or(""), frame_phase = write_context.phase.as_str(), error = %error, "Responses WebSocket pump transport write failed");
}
let _ = acknowledged.send(result);
if let Some(reason) = reason { return reason; }
if closing { return PumpExitReason::ServerClose; }
continue;
}
_ = heartbeat.tick() => {
ping_sequence = ping_sequence.wrapping_add(1);
Message::Ping(ping_sequence.to_be_bytes().to_vec().into())
}
};
let message = match middleware::transform(
message,
contract::Direction::Outgoing,
writer.clone(),
context,
)
.await
{
Ok(Some(message)) => message,
Ok(None) => continue,
Err(_) => return PumpExitReason::MiddlewareFailed,
};
if let Err(error) = writer.write(message).await {
tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket control write failed");
return write_exit_reason(&error, context);
}
context
.stats
.ping_written_count
.fetch_add(1, Ordering::Relaxed);
}
}
async fn write_message<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(
incoming: &mpsc::Sender<PendingConnectionEvent>,
slots: &Arc<Semaphore>,
event: ConnectionEvent,
) -> Result<(), PumpExitReason> {
let permit = Arc::clone(slots)
.try_acquire_owned()
.map_err(|_| PumpExitReason::InboundOverload)?;
incoming
.try_send(PendingConnectionEvent {
event,
_permit: Some(permit),
})
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
})
}
pub use handle::{ResponsesWebSocketConnection, spawn_connection};
@@ -0,0 +1,230 @@
//! 单 owner 驱动连接读写、心跳与有界入站队列
use super::{
frame::*,
middleware,
state::{ConnectionCommand, ConnectionStats, PumpContext},
};
use axum::extract::ws::Message;
use futures::{Sink, Stream, StreamExt};
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader, websocket as contract},
lifecycle::CancellationToken,
};
use std::{
fmt,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
};
use tokio::{
sync::{Semaphore, mpsc, oneshot},
time::Instant,
};
#[expect(clippy::too_many_arguments)]
pub(super) async fn run_pump<S, E>(
socket: S,
commands: mpsc::Receiver<ConnectionCommand>,
incoming: mpsc::Sender<PendingConnectionEvent>,
exited: oneshot::Sender<PumpExitReason>,
connection_id: Arc<str>,
cancellation: CancellationToken,
opened_at: Instant,
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 + 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,
};
if matches!(reason, PumpExitReason::ClientClose) {
if let Err(error) = writer.finish_peer_close(&cancellation).await {
tracing::debug!(websocket_connection_id = %context.connection_id, %error, "Responses WebSocket close acknowledgement could not be flushed");
}
} else {
context.cancellation.cancel();
}
let _ = exited.send(reason);
tracing::debug!(websocket_connection_id = %context.connection_id, pump_exit_reason = reason.as_str(), "Responses WebSocket pump exited");
}
async fn read_connection<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(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(
context.opened_at + DOWNSTREAM_PING_INTERVAL,
DOWNSTREAM_PING_INTERVAL,
);
heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut ping_sequence = 0_u64;
loop {
let message = tokio::select! {
biased;
() = context.cancellation.cancelled() => return PumpExitReason::LifecycleShutdown,
command = commands.recv() => {
let Some(command) = command else { return PumpExitReason::CoordinatorDropped; };
let ConnectionCommand { message, context: write_context, acknowledged } = command;
let closing = matches!(message, Message::Close(_));
let result = match middleware::transform(message, contract::Direction::Outgoing, writer.clone(), context).await {
Ok(Some(message)) => writer.write(message).await.map(|()| WriteOutcome::Written),
Ok(None) => Ok(WriteOutcome::Suppressed),
Err(error) => Err(ConnectionWriteError::Transport { message: error.to_string() }),
};
let reason = result.as_ref().err().map(|error| write_exit_reason(error, context));
if let Err(error) = &result {
tracing::debug!(websocket_connection_id = %context.connection_id, request_id = write_context.request_id.as_deref().unwrap_or(""), frame_phase = write_context.phase.as_str(), error = %error, "Responses WebSocket pump transport write failed");
}
let _ = acknowledged.send(result);
if let Some(reason) = reason { return reason; }
if closing { return PumpExitReason::ServerClose; }
continue;
}
_ = heartbeat.tick() => {
ping_sequence = ping_sequence.wrapping_add(1);
Message::Ping(ping_sequence.to_be_bytes().to_vec().into())
}
};
let message = match middleware::transform(
message,
contract::Direction::Outgoing,
writer.clone(),
context,
)
.await
{
Ok(Some(message)) => message,
Ok(None) => continue,
Err(_) => return PumpExitReason::MiddlewareFailed,
};
if let Err(error) = writer.write(message).await {
tracing::info!(websocket_connection_id = %context.connection_id, error = %error, "Responses WebSocket control write failed");
return write_exit_reason(&error, context);
}
context
.stats
.ping_written_count
.fetch_add(1, Ordering::Relaxed);
}
}
fn write_exit_reason(error: &ConnectionWriteError, context: &PumpContext) -> PumpExitReason {
match error {
ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout,
ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError,
ConnectionWriteError::Closed if context.cancellation.is_cancelled() => {
PumpExitReason::LifecycleShutdown
}
ConnectionWriteError::Closed => PumpExitReason::PumpStopped,
}
}
fn emit_incoming(
incoming: &mpsc::Sender<PendingConnectionEvent>,
slots: &Arc<Semaphore>,
event: ConnectionEvent,
) -> Result<(), PumpExitReason> {
let permit = Arc::clone(slots)
.try_acquire_owned()
.map_err(|_| PumpExitReason::InboundOverload)?;
incoming
.try_send(PendingConnectionEvent {
event,
_permit: Some(permit),
})
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
})
}
@@ -0,0 +1,75 @@
//! 单连接共享的命令、运行预算与观测状态
use super::frame::{
ConnectionConfig, ConnectionWriteError, OUTBOUND_COMMAND_BUFFER, WriteContext, WriteOutcome,
};
use axum::extract::ws::Message;
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader},
lifecycle::CancellationToken,
};
use std::sync::{
Arc,
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
};
use tokio::{
sync::{mpsc, oneshot},
time::Instant,
};
pub(super) struct ConnectionCommand {
pub(super) message: Message,
pub(super) context: WriteContext,
pub(super) acknowledged: oneshot::Sender<Result<WriteOutcome, ConnectionWriteError>>,
}
#[derive(Default)]
pub(super) struct ConnectionStats {
pub(super) command_queue_high_water: AtomicUsize,
pub(super) command_backpressure_count: AtomicU64,
pub(super) ping_received_count: AtomicU64,
pub(super) ping_written_count: AtomicU64,
pub(super) pong_received_count: AtomicU64,
pub(super) last_read_ms: AtomicU64,
pub(super) last_write_ms: AtomicU64,
}
impl ConnectionStats {
pub(super) fn record_read(&self, opened_at: Instant) {
self.last_read_ms.store(
u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
}
pub(super) fn record_write(&self, opened_at: Instant) {
self.last_write_ms.store(
u64::try_from(opened_at.elapsed().as_millis()).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
}
pub(super) fn observe_command_queue(&self, sender: &mpsc::Sender<ConnectionCommand>) {
let queued = OUTBOUND_COMMAND_BUFFER
.saturating_sub(sender.capacity())
.saturating_add(1)
.min(OUTBOUND_COMMAND_BUFFER);
self.command_queue_high_water
.fetch_max(queued, Ordering::Relaxed);
if sender.capacity() == 0 {
self.command_backpressure_count
.fetch_add(1, Ordering::Relaxed);
}
}
}
pub(super) struct PumpContext {
pub(super) middleware: Option<FrozenMiddlewarePlan>,
pub(super) headers: Arc<[MiddlewareHeader]>,
pub(super) connection_id: Arc<str>,
pub(super) cancellation: CancellationToken,
pub(super) opened_at: Instant,
pub(super) expired: Arc<AtomicBool>,
pub(super) stats: Arc<ConnectionStats>,
pub(super) config: ConnectionConfig,
}
@@ -12,7 +12,7 @@ use gateway_core::engine::middleware::{
MiddlewareBody, MiddlewareError, MiddlewareFrame, MiddlewareFraming, MiddlewareHeader,
MiddlewareResponse,
};
use gateway_core::engine::response_control::ResponseControl;
use gateway_core::engine::response_control::{ResponseControl, ResponseControlUnavailable};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::ProviderSessionState;
@@ -25,7 +25,7 @@ use super::{
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody},
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
protocol::{
decode_response_interrupt, error_event, initial_engine_error_event, response_metadata_event,
error_event, initial_engine_error_event, is_response_create, response_metadata_event,
},
};
@@ -137,7 +137,7 @@ pub(super) async fn forward_response(
replay: &mut ConnectionReplaySnapshot,
capture: ReplayCaptureHandle,
validation: ResponseValidationFacts,
response_control: ResponseControl,
response_control: &ResponseControl,
) -> ForwardOutcome {
let (protocol, status, headers, mut body, _) = response.into_parts();
if protocol != "openai" || StatusCode::from_u16(status).is_err() {
@@ -150,8 +150,14 @@ pub(super) async fn forward_response(
ProviderResponseHeader::new(name, value)
})
.collect::<Vec<_>>();
let first = match next_body_input(connection, body.as_mut(), &response_control, &request_id)
.await
let first = match next_body_input(
connection,
body.as_mut(),
response_control,
&request_id,
false,
)
.await
{
BodyInput::Frame(Ok(Some(frame))) => frame,
BodyInput::Frame(Ok(None)) => {
@@ -197,8 +203,14 @@ pub(super) async fn forward_response(
let frame = match current.take() {
Some(frame) => frame,
None => {
match next_body_input(connection, body.as_mut(), &response_control, &request_id)
.await
match next_body_input(
connection,
body.as_mut(),
response_control,
&request_id,
terminal_seen,
)
.await
{
BodyInput::Frame(Ok(Some(frame))) => frame,
BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => {
@@ -284,8 +296,10 @@ fn detach_body(body: Box<dyn MiddlewareBody>) {
}
fn validate_frame(frame: &MiddlewareFrame) -> Result<(), MiddlewareError> {
if frame.framing() != MiddlewareFraming::JsonDocument
|| serde_json::from_slice::<serde::de::IgnoredAny>(frame.bytes()).is_err()
if !matches!(
frame.framing(),
MiddlewareFraming::JsonDocument | MiddlewareFraming::RawBytes
) || std::str::from_utf8(frame.bytes()).is_err()
{
return Err(MiddlewareError::InvalidState);
}
@@ -302,6 +316,7 @@ async fn next_body_input(
body: &mut dyn MiddlewareBody,
response_control: &ResponseControl,
request_id: &Arc<str>,
terminal_seen: bool,
) -> BodyInput {
// 处理控制帧时继续持有同一个读取 future,不能取消正在加工正文的中间件
let frame = body.next_frame();
@@ -312,19 +327,11 @@ async fn next_body_input(
let Some(event) = event else { return BodyInput::Disconnect; };
match &event.event {
ConnectionEvent::Text(payload) => {
let error = match decode_response_interrupt(payload) {
Ok(Some(response_id)) => response_control.interrupt(&response_id).err().map(|_| {
super::super::RequestDecodeError::InvalidValue { field: "response_id".to_owned() }.protocol_body()
}),
Ok(None) => {
connection.defer(event);
continue;
}
Err(error) => Some(error.protocol_body()),
};
if let Some(error) = error
&& send_protocol_error(connection, StatusCode::BAD_REQUEST, error, request_id).await == ForwardOutcome::Disconnect
{
if is_response_create(payload) {
connection.defer(event);
continue;
}
if send_control(connection, response_control, payload, request_id).await == ForwardOutcome::Disconnect {
return BodyInput::Disconnect;
}
}
@@ -337,6 +344,11 @@ async fn next_body_input(
ConnectionEvent::Exited(_) => return BodyInput::Disconnect,
}
}
event = response_control.receive(), if terminal_seen => {
if forward_control_event(connection, response_control, event).await == ForwardOutcome::Disconnect {
return BodyInput::Disconnect;
}
}
frame = &mut frame => return BodyInput::Frame(frame),
}
}
@@ -537,19 +549,26 @@ impl WebSocketExecutionBody {
self.transformed_pending = transformed;
continue;
}
messages.extend(encoded.into_iter().map(|message| (message, transformed)));
let framing = if event.wire_event().is_some_and(|wire| {
wire.raw_websocket_message().is_some() && !wire.has_json_data()
}) {
MiddlewareFraming::RawBytes
} else {
MiddlewareFraming::JsonDocument
};
messages.extend(
encoded
.into_iter()
.map(|message| (message, transformed, framing)),
);
self.transformed_pending = false;
}
let terminal = self.encoder.is_completed() || self.encoder.has_wire_failure();
let last = messages.len().saturating_sub(1);
self.pending.extend(messages.into_iter().enumerate().map(
|(index, (message, transformed))| {
MiddlewareFrame::new(
Bytes::from(message),
MiddlewareFraming::JsonDocument,
terminal && index == last,
)
.with_transformed(transformed)
|(index, (message, transformed, framing))| {
MiddlewareFrame::new(Bytes::from(message), framing, terminal && index == last)
.with_transformed(transformed)
},
));
self.awaiting_terminal_eof = terminal;
@@ -720,3 +739,46 @@ async fn send_metadata(
.await
.is_ok()
}
pub(super) async fn send_control(
connection: &mut ResponsesWebSocketConnection,
control: &ResponseControl,
payload: &str,
request_id: &Arc<str>,
) -> ForwardOutcome {
let result = tokio::select! {
_ = connection.wait_for_exit() => return ForwardOutcome::Disconnect,
result = control.send(payload) => result,
};
if result.is_err() {
return send_gateway_error(
connection,
&GatewayError::new(
GatewayErrorKind::InvalidRequest,
"No upstream WebSocket connection is available for control messages",
),
request_id,
)
.await;
}
ForwardOutcome::Continue
}
pub(super) async fn forward_control_event(
connection: &mut ResponsesWebSocketConnection,
control: &ResponseControl,
event: Result<String, ResponseControlUnavailable>,
) -> ForwardOutcome {
let Ok(payload) = event else {
control.clear();
return ForwardOutcome::Continue;
};
if connection
.send_text(payload, WriteContext::connection(FramePhase::Data))
.await
.is_err()
{
return ForwardOutcome::Disconnect;
}
ForwardOutcome::Continue
}
@@ -24,6 +24,7 @@ use gateway_core::{
engine::{
execution::{AuthenticatedClient, ClientTransport},
middleware::{FrozenMiddlewarePlan, MiddlewareError},
response_control::ResponseControl,
},
lifecycle::{ConnectionGuard, ConnectionLifecycle},
operation::OperationKind,
@@ -46,11 +47,12 @@ use super::{
};
use connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext};
use forward::{
ConnectionReplaySnapshot, ForwardOutcome, execution_response, forward_response,
new_replay_capture, send_gateway_error, send_middleware_error, send_protocol_error,
ConnectionReplaySnapshot, ForwardOutcome, execution_response, forward_control_event,
forward_response, new_replay_capture, send_control, send_gateway_error, send_middleware_error,
send_protocol_error,
};
use protocol::connection_limit_event;
pub use protocol::{ResponseCreateFrameError, decode_response_create_with_context};
use protocol::{connection_limit_event, is_response_create};
const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only";
const CONNECTION_LIMIT_CLOSE_REASON: &str = "Responses websocket connection limit reached";
@@ -191,10 +193,20 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
);
let mut request_count = 0_u64;
let mut replay = ConnectionReplaySnapshot::default();
let mut response_control = ResponseControl::default();
loop {
let Some(event) = connection.next_event().await else {
break;
let event = tokio::select! {
event = connection.next_event() => {
let Some(event) = event else { break };
event
}
event = response_control.receive() => {
if forward_control_event(&mut connection, &response_control, event).await == ForwardOutcome::Disconnect {
break;
}
continue;
}
};
let payload = match event {
ConnectionEvent::Text(payload) => payload,
@@ -212,8 +224,22 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
expire_connection(&mut connection).await;
break;
}
request_count = request_count.saturating_add(1);
let correlation_id = Arc::<str>::from(service.next_request_id());
if !is_response_create(&payload) {
if send_control(
&mut connection,
&response_control,
&payload,
&correlation_id,
)
.await
== ForwardOutcome::Disconnect
{
break;
}
continue;
}
request_count = request_count.saturating_add(1);
// 初步解码只保留路由事实,避免整份正文跨越准入等待和响应交付
let model_hint = match decode_response_create_with_context(&payload, &request_headers) {
Ok(decoded) => Some(decoded.metadata().requested_model().to_owned()),
@@ -251,6 +277,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
expire_connection(&mut connection).await;
break;
}
response_control.clear();
let execution = service.execution();
let preparation = async {
// 握手的显式改写可继承,Key 策略与宿主设置仍在每轮执行前刷新
@@ -291,7 +318,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
}
};
let request_id = Arc::<str>::from(prepared.request_id().to_string());
let response_control = prepared.response_control();
response_control = prepared.response_control();
let capture = new_replay_capture();
let validation = ResponseValidationFacts::default();
let input = RequestInput {
@@ -381,7 +408,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
&mut replay,
capture,
validation,
response_control,
&response_control,
)
.await
== ForwardOutcome::Disconnect
@@ -31,42 +31,22 @@ pub fn decode_response_create_with_context(
decode_response_create_inner(payload, request_headers)
}
/// 中断只引用活动响应;普通创建帧仍交给原有解码与串行准入路径
pub(super) fn decode_response_interrupt(
payload: &str,
) -> Result<Option<String>, ResponseCreateFrameError> {
/// 只识别需要本地准入的创建帧,其余文本不在网关解释
pub(super) fn is_response_create(payload: &str) -> bool {
#[derive(serde::Deserialize)]
struct FrameType {
#[serde(rename = "type")]
message_type: Option<String>,
message_type: String,
}
// 非控制帧仍在原来的串行请求边界校验,不能抢先拒绝排队的普通请求
// 只投影类型,避免为排队的大型 response.create 再构造完整 JSON 树
let Ok(frame) = serde_json::from_str::<FrameType>(payload) else {
return Ok(None);
};
if frame.message_type.as_deref() != Some("response.interrupt") {
return Ok(None);
match serde_json::from_str::<FrameType>(payload) {
Ok(frame) => frame.message_type == "response.create",
// struct 投影拒绝重复 type;沿用创建解码器的 JSON object 语义,
// 避免这类创建请求绕过准入进入透传路径
Err(_) => serde_json::from_str::<Value>(payload).is_ok_and(|frame| {
frame.get("type").and_then(Value::as_str) == Some("response.create")
}),
}
let value: Value =
serde_json::from_str(payload).map_err(|_| ResponseCreateFrameError::InvalidJson)?;
if value.get("mode").and_then(Value::as_str) != Some("discard_partial_items") {
return Err(ResponseCreateFrameError::Request(
RequestDecodeError::InvalidValue {
field: "mode".to_owned(),
},
));
}
let id = value
.get("response_id")
.and_then(Value::as_str)
.filter(|id| !id.is_empty())
.ok_or_else(|| {
ResponseCreateFrameError::Request(RequestDecodeError::InvalidValue {
field: "response_id".to_owned(),
})
})?;
Ok(Some(id.to_owned()))
}
fn decode_response_create_inner(
@@ -49,6 +49,36 @@ async fn list_route_should_keep_camel_case_group_and_page_wire() {
assert!(value["data"]["items"][0].get("member_count").is_none());
}
#[tokio::test]
async fn options_route_should_return_only_lightweight_group_fields() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::GET,
"/api/admin/account-groups/options?page=1&pageSize=1&search=alpha&enabled=true",
None,
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(
value["data"]["items"][0],
json!({
"id": PRIMARY_GROUP_ID,
"name": "Alpha routing",
"color": "#2563EBFF",
"enabled": true
})
);
assert_eq!(value["data"]["page"]["page"], 1);
assert_eq!(value["data"]["page"]["pageSize"], 1);
assert_eq!(value["data"]["page"]["total"], 1);
assert_eq!(value["data"]["page"]["totalPages"], 1);
assert_eq!(value["data"]["configRevision"], 7);
}
#[tokio::test]
async fn create_route_should_create_an_empty_group() {
let fixture = authenticated_fixture().await;
@@ -163,6 +193,11 @@ async fn list_route_should_reject_unknown_and_invalid_pagination_fields() {
"/api/admin/account-groups?page=0",
"/api/admin/account-groups?pageSize=0",
"/api/admin/account-groups?pageSize=201",
"/api/admin/account-groups/options?other=true",
"/api/admin/account-groups/options?page=0",
"/api/admin/account-groups/options?pageSize=0",
"/api/admin/account-groups/options?pageSize=201",
"/api/admin/account-groups/options?search=%00",
] {
let response = request(router(&fixture), Method::GET, uri, None, true).await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{uri}");
@@ -219,6 +254,14 @@ async fn account_group_routes_should_require_admin_authentication() {
false,
)
.await;
let options_response = request(
router(&fixture),
Method::GET,
"/api/admin/account-groups/options",
None,
false,
)
.await;
let post_response = request(
router(&fixture),
Method::POST,
@@ -229,9 +272,29 @@ async fn account_group_routes_should_require_admin_authentication() {
.await;
assert_eq!(get_response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(options_response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(post_response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn options_route_should_reject_client_key_identity() {
let fixture = crate::support::key_fixture().await;
let response = router(&fixture)
.oneshot(
Request::builder()
.method(Method::GET)
.uri("/api/admin/account-groups/options")
.header("x-api-key", crate::support::RAW_KEY)
.header("x-request-id", "req_account_group_options_client_key")
.body(Body::empty())
.expect("client key options request"),
)
.await
.expect("client key options response");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
async fn authenticated_fixture() -> AdminTestFixture {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
+60 -9
View File
@@ -18,9 +18,10 @@ use gateway_admin::{
MutationContext, Revision,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup, NewAccountGroup,
SetAccountGroupEnabled, UpdateAccountGroup,
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation,
AccountGroupOptionsPage, AccountGroupPage, AccountGroupRecord, AccountGroupRef,
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountPageItem, AccountRuntimeSnapshot,
@@ -37,9 +38,10 @@ use gateway_admin::{
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
},
observability::{
DashboardObservation, DecimalAmount, DiagnosticDimension, DiagnosticsObservation,
OpsError, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail,
UsageFilter, UsageListRecord, UsageOverview, UsagePage, UsageQuery,
DashboardObservation, DashboardQuery, DecimalAmount, DiagnosticDimension,
DiagnosticsObservation, Granularity, OpsError, OpsErrorPage, OpsErrorQuery,
RequestMetricPoint, TimeRange, UsageDetail, UsageFilter, UsageListRecord,
UsageOverview, UsagePage, UsageQuery,
},
provider_credentials::{
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
@@ -788,6 +790,48 @@ impl AccountGroupStore for MemoryAccountGroupStore {
})
}
async fn list_account_group_options(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupOptionsPage> {
let state = self.state.lock().expect("account groups");
let search = query.search.as_ref().map(|value| value.to_lowercase());
let mut matching = state
.groups
.values()
.filter(|record| {
query
.enabled
.is_none_or(|enabled| record.enabled == enabled)
})
.filter(|record| {
search
.as_ref()
.is_none_or(|search| record.name.to_lowercase().contains(search))
})
.map(|record| AccountGroupRef {
id: record.id.clone(),
name: record.name.clone(),
color: record.color.clone(),
enabled: record.enabled,
})
.collect::<Vec<_>>();
matching.sort_by(|left, right| left.name.cmp(&right.name));
let total = matching.len() as u64;
let page_size = usize::from(query.page_size.get());
let offset = usize::try_from(query.page.saturating_sub(1))
.unwrap_or(usize::MAX)
.saturating_mul(page_size);
let items = matching.into_iter().skip(offset).take(page_size).collect();
Ok(AccountGroupOptionsPage {
config_revision: state.revision,
items,
total,
page: query.page,
page_size: query.page_size.get(),
})
}
async fn load_account_group_members(
&self,
group_ids: &[gateway_core::routing::AccountGroupId],
@@ -1303,13 +1347,13 @@ impl AccountRuntimeStore for UnusedStore {
impl ObservabilityStore for UnusedStore {
async fn dashboard_summary(
&self,
range: TimeRange,
query: DashboardQuery,
_: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
*self
.dashboard_summary_range
.lock()
.expect("dashboard summary range") = Some(range);
.expect("dashboard summary range") = Some(query.range);
self.dashboard_observation
.lock()
.expect("dashboard observation")
@@ -1317,7 +1361,11 @@ impl ObservabilityStore for UnusedStore {
.ok_or_else(|| unavailable("dashboard"))
}
async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult<Vec<RequestMetricPoint>> {
async fn dashboard_trend(
&self,
_: TimeRange,
_: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
Err(unavailable("dashboard trend"))
}
@@ -1325,6 +1373,7 @@ impl ObservabilityStore for UnusedStore {
&self,
range: TimeRange,
filter: UsageFilter,
_: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
let mut data = self.observations.lock().expect("observations");
data.trends.push((range, filter));
@@ -1335,6 +1384,7 @@ impl ObservabilityStore for UnusedStore {
&self,
_: TimeRange,
_: UsageFilter,
_: Granularity,
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
Box::pin(futures::stream::once(async {
Err(unavailable("usage billing facts"))
@@ -1382,6 +1432,7 @@ impl ObservabilityStore for UnusedStore {
_: TimeRange,
_: UsageFilter,
_: DiagnosticDimension,
_: u16,
) -> AdminStoreResult<DiagnosticsObservation> {
Ok(self.diagnostics.lock().expect("diagnostics").clone())
}
@@ -25,6 +25,7 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
runtime_slots: None,
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
totals: Default::default(),
provider_accounts: AccountPoolMetrics {
@@ -107,6 +108,7 @@ async fn dashboard_summary_should_use_lifetime_totals_for_card_footers() {
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
runtime_slots: None,
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
totals: DashboardTotals {
request_count: 42,
@@ -178,6 +180,7 @@ async fn dashboard_summary_should_default_to_current_deployment_day() {
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
runtime_slots: None,
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
totals: Default::default(),
provider_accounts: AccountPoolMetrics::default(),
@@ -230,6 +233,22 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let now = Utc::now();
let mut record = usage_record_with_account(
"req_detail",
"acct_snap_a",
"Snapshot Alpha",
"alpha@example.invalid",
"oauth",
now,
);
record.first_token_ms = Some(20);
record.upstream_response_ms = Some(1_000);
record.upstream_api_overhead_ms = Some(120.25);
record.upstream_engine_ms = Some(6400.0);
record.upstream_engine_iapi_ttft_ms = Some(650.5);
record.upstream_engine_service_ttft_ms = Some(720.25);
record.upstream_engine_iapi_tbt_ms = Some(18.45);
record.upstream_engine_service_tbt_ms = Some(20.12);
fixture
.usage_detail
.lock()
@@ -237,14 +256,7 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
.replace(UsageDetail {
trace: None,
related_requests: Vec::new(),
request: usage_record_with_account(
"req_detail",
"acct_snap_a",
"Snapshot Alpha",
"alpha@example.invalid",
"oauth",
now,
),
request: record,
attempts: vec![UsageAttempt {
source: "ops_event".to_owned(),
id: "ops_detail".to_owned(),
@@ -297,6 +309,39 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
.await
.expect("usage detail body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("usage detail JSON");
assert_eq!(value["data"]["firstTokenLatencyMs"], 20);
assert_eq!(value["data"]["latencyDetails"]["upstreamResponseMs"], 1_000);
assert_eq!(
value["data"]["latencyDetails"]["upstreamApiOverheadMs"],
json!(120.25)
);
assert_eq!(
value["data"]["latencyDetails"]["upstreamEngineMs"],
json!(6400.0)
);
assert_eq!(
value["data"]["latencyDetails"]["upstreamEngineIapiTtftMs"],
json!(650.5)
);
assert_eq!(
value["data"]["latencyDetails"]["upstreamEngineServiceTtftMs"],
json!(720.25)
);
assert_eq!(
value["data"]["latencyDetails"]["upstreamEngineIapiTbtMs"],
json!(18.45)
);
assert_eq!(
value["data"]["latencyDetails"]["upstreamEngineServiceTbtMs"],
json!(20.12)
);
assert!(value["data"].get("firstTokenMs").is_none());
assert!(value["data"].get("firstTokenLatencyMsDisplay").is_none());
assert!(
value["data"]["latencyDetails"]
.get("firstTokenMs")
.is_none()
);
assert_eq!(
serde_json::json!({
"accountId": value["data"]["accountId"],
@@ -699,6 +744,13 @@ fn usage_record_with_account(
first_text_ms: None,
first_token_ms: None,
provider_processing_ms: None,
upstream_response_ms: None,
upstream_api_overhead_ms: None,
upstream_engine_ms: None,
upstream_engine_iapi_ttft_ms: None,
upstream_engine_service_ttft_ms: None,
upstream_engine_iapi_tbt_ms: None,
upstream_engine_service_tbt_ms: None,
latency_ms: None,
admission_decision_ms: None,
account_selection_wait_ms: None,
@@ -831,6 +883,13 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
first_text_ms: Some(19),
first_token_ms: Some(20),
provider_processing_ms: Some(21),
upstream_response_ms: Some(2_000),
upstream_api_overhead_ms: None,
upstream_engine_ms: None,
upstream_engine_iapi_ttft_ms: None,
upstream_engine_service_ttft_ms: None,
upstream_engine_iapi_tbt_ms: None,
upstream_engine_service_tbt_ms: None,
latency_ms: Some(31),
admission_decision_ms: Some(1),
account_selection_wait_ms: Some(2),
@@ -967,8 +1026,8 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
"firstEventMs": 17,
"firstReasoningMs": 18,
"firstTextMs": 19,
"firstTokenMs": 20,
"openaiProcessingMs": 21,
"upstreamResponseMs": 2000,
},
"firstTokenLatencyMs": 20,
"latencyMs": 31,
@@ -1042,6 +1101,7 @@ async fn calendar_queries_and_display_follow_deployment_timezone_with_one_anchor
let fixture = AdminTestFixture::with_timezone(timezone).await;
fixture.auth.insert_session("valid-session");
*fixture.dashboard_observation.lock().unwrap() = Some(DashboardObservation {
runtime_slots: None,
range,
totals: Default::default(),
provider_accounts: Default::default(),
@@ -5,24 +5,31 @@ use axum::{
http::{Request, StatusCode},
};
use futures::future::BoxFuture;
use gateway_core::{
engine::{
execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService,
PreparedRootExecution, StartExecution, StartProviderExecution, StartedExecution,
},
middleware::{
FrozenMiddlewarePlan, MiddlewareBody, MiddlewareContext, MiddlewareError,
MiddlewareFrame, MiddlewareFraming, MiddlewareHeader, MiddlewareNext, MiddlewarePlan,
MiddlewareRequest, MiddlewareResponse,
},
},
error::GatewayError,
lifecycle::CancellationToken,
policy::ClientApiKeyId,
routing::PublicModelId,
runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
};
use gateway_core::engine::execution::AuthenticatedClient;
use gateway_core::engine::execution::ClientAuthenticationError;
use gateway_core::engine::execution::ExecutionService;
use gateway_core::engine::execution::PreparedRootExecution;
use gateway_core::engine::execution::StartExecution;
use gateway_core::engine::execution::StartProviderExecution;
use gateway_core::engine::execution::StartedExecution;
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareBody;
use gateway_core::engine::middleware::MiddlewareContext;
use gateway_core::engine::middleware::MiddlewareError;
use gateway_core::engine::middleware::MiddlewareFrame;
use gateway_core::engine::middleware::MiddlewareFraming;
use gateway_core::engine::middleware::MiddlewareHeader;
use gateway_core::engine::middleware::MiddlewareNext;
use gateway_core::engine::middleware::MiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareRequest;
use gateway_core::engine::middleware::MiddlewareResponse;
use gateway_core::error::GatewayError;
use gateway_core::lifecycle::CancellationToken;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::PublicModelId;
use gateway_core::routing::extensions::ExtensionSetId;
use gateway_core::routing::extensions::ExtensionSetLease;
use gateway_core::routing::extensions::ExtensionSetReference;
use std::sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
@@ -368,10 +375,11 @@ pub(super) fn resource_fixture(
}
async fn published_fixture() -> AdminTestFixture {
use gateway_core::{
routing::{ConfigRevision, RuntimeSnapshot},
runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
};
use gateway_core::routing::ConfigRevision;
use gateway_core::routing::RuntimeSnapshot;
use gateway_core::routing::extensions::ExtensionSetId;
use gateway_core::routing::extensions::ExtensionSetLease;
use gateway_core::routing::extensions::ExtensionSetReference;
use std::sync::Arc;
struct Lease;
impl ExtensionSetLease for Lease {
@@ -546,7 +546,7 @@ pub(super) struct TestPluginPorts {
impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPorts {
async fn validate_target(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
) -> Result<(), AdminError> {
Ok(())
@@ -554,7 +554,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
async fn start_callback(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: gateway_admin::model::plugins::management::StartPluginManagementCallback,
_: &gateway_admin::model::auth::AdminRequestContext,
@@ -564,7 +564,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
}
async fn callback(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
_: &str,
request: gateway_admin::model::plugins::management::PluginManagementRequest,
@@ -576,14 +576,14 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
}
async fn views(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
) -> Result<Vec<gateway_admin::model::plugins::management::PluginManagementView>, AdminError>
{
Ok(Vec::new())
}
async fn resource(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
path: &str,
public: bool,
@@ -596,7 +596,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
}
async fn handle(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::management::PluginManagementTarget,
request: gateway_admin::model::plugins::management::PluginManagementRequest,
) -> Result<gateway_admin::model::plugins::management::PluginManagementResponse, AdminError>
@@ -611,7 +611,7 @@ impl gateway_admin::ports::plugin_management::PluginManagement for TestPluginPor
impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts {
async fn activate_state(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: &gateway_admin::model::plugins::instances::PluginInstance,
) -> Result<(), AdminError> {
Ok(())
@@ -621,7 +621,7 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for TestPluginPorts {
}
async fn migrate_state(
&self,
_: &gateway_core::runtime::extensions::ExtensionSetReference,
_: &gateway_core::routing::extensions::ExtensionSetReference,
_: gateway_admin::model::plugins::state::PluginStateTransition,
) -> Result<(), AdminError> {
panic!("unexpected state migration")
@@ -647,7 +647,7 @@ impl gateway_admin::ports::plugins::PluginPreparation for TestPluginPorts {
async fn prepare(
&self,
_: gateway_admin::model::plugins::instances::PluginInstanceSnapshot,
) -> Result<gateway_core::runtime::extensions::ExtensionSetReference, AdminError> {
) -> Result<gateway_core::routing::extensions::ExtensionSetReference, AdminError> {
Err(AdminError::invalid("unused plugin fixture"))
}
}
@@ -658,7 +658,7 @@ impl gateway_admin::ports::plugins::PluginRuntimeDiagnostics for TestPluginPorts
&self,
snapshot: &gateway_admin::model::plugins::instances::PluginInstanceSnapshot,
_: Option<u64>,
_: Option<&gateway_core::runtime::extensions::ExtensionSetReference>,
_: Option<&gateway_core::routing::extensions::ExtensionSetReference>,
) -> Option<BTreeMap<String, gateway_admin::model::plugins::instances::PluginInstanceRuntime>>
{
Some(
@@ -1392,7 +1392,7 @@ async fn guardian_reservation_round_trips_and_rejects_invalid_values() {
#[tokio::test]
async fn account_affinity_and_rotation_budget_round_trip_and_reject_invalid_values() {
for (mode, rotations, ttl) in [("relaxed", 0, 1), ("strict", 31, 720)] {
for (mode, rotations, ttl) in [("relaxed", 0, 1), ("preferred", 3, 24), ("strict", 31, 720)] {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut body = update_body();
@@ -1,4 +1,4 @@
//! 检查 API 层依赖、源码与测试模块边界,以及管理路由约定
//! 检查 API 层依赖与管理路由约定,模块布局和测试镜像由 workspace 架构检查覆盖
use std::{fs, path::Path};
@@ -29,150 +29,6 @@ fn workspace_should_include_gateway_api() {
assert!(workspace.contains("\"crates/gateway-api\""));
}
#[test]
fn source_tree_should_match_frozen_machine_manifest() {
let root = Path::new(env!("CARGO_MANIFEST_DIR"));
let mut actual = rust_files(&root.join("src"));
actual.sort();
let mut expected = vec![
"src/admin/account_groups.rs",
"src/admin/accounts/credentials.rs",
"src/admin/accounts/handlers.rs",
"src/admin/accounts/import_tasks.rs",
"src/admin/accounts/mod.rs",
"src/admin/accounts/presenter.rs",
"src/admin/accounts/wire.rs",
"src/admin/auth.rs",
"src/admin/backups.rs",
"src/admin/client_keys.rs",
"src/admin/diagnostics.rs",
"src/admin/extract.rs",
"src/admin/mod.rs",
"src/admin/observability/mod.rs",
"src/admin/observability/presenter.rs",
"src/admin/observability/query.rs",
"src/admin/observability/routes.rs",
"src/admin/observability/wire.rs",
"src/admin/plugins/artifacts.rs",
"src/admin/plugins/distribution.rs",
"src/admin/plugins/instances.rs",
"src/admin/plugins/management.rs",
"src/admin/plugins/mod.rs",
"src/admin/presenter.rs",
"src/admin/proxies.rs",
"src/admin/settings.rs",
"src/admin/system.rs",
"src/admin/wire.rs",
"src/auth.rs",
"src/health.rs",
"src/key_usage/mod.rs",
"src/key_usage/query.rs",
"src/key_usage/presenter.rs",
"src/lib.rs",
"src/middleware/headers.rs",
"src/middleware/mod.rs",
"src/middleware/websocket.rs",
"src/openai/middleware/http.rs",
"src/openai/middleware/mod.rs",
"src/openai/middleware/request.rs",
"src/openai/middleware/session.rs",
"src/openai/auth.rs",
"src/openai/error.rs",
"src/openai/images/http.rs",
"src/openai/images/mod.rs",
"src/openai/live/http.rs",
"src/openai/live/mod.rs",
"src/openai/live/websocket.rs",
"src/openai/mod.rs",
"src/openai/models.rs",
"src/openai/endpoint.rs",
"src/openai/responses/error.rs",
"src/openai/responses/http.rs",
"src/openai/responses/mod.rs",
"src/openai/responses/request.rs",
"src/openai/responses/response.rs",
"src/openai/responses/validation.rs",
"src/openai/responses/websocket/connection/mod.rs",
"src/openai/responses/websocket/connection/middleware.rs",
"src/openai/responses/websocket/forward.rs",
"src/openai/responses/websocket/mod.rs",
"src/openai/responses/websocket/protocol.rs",
"src/openai/router.rs",
"src/openai/search/http.rs",
"src/openai/search/mod.rs",
"src/openai/service.rs",
"src/openai/usage.rs",
"src/provider.rs",
"src/session_cookie.rs",
"src/time.rs",
];
expected.sort_unstable();
assert_eq!(actual, expected);
}
#[test]
fn test_tree_should_match_frozen_rust_mirror() {
let root = Path::new(env!("CARGO_MANIFEST_DIR"));
let mut actual = rust_files(&root.join("tests"));
actual.sort();
let mut expected = vec![
"tests/admin/account_groups.rs",
"tests/admin/accounts/handlers.rs",
"tests/admin/accounts/import_tasks.rs",
"tests/admin/accounts/mod.rs",
"tests/admin/accounts/presenter.rs",
"tests/admin/auth.rs",
"tests/admin/client_keys.rs",
"tests/admin/errors.rs",
"tests/admin/error_diagnostics.rs",
"tests/admin/mod.rs",
"tests/admin/observability/mod.rs",
"tests/admin/observability/query.rs",
"tests/admin/observability/response.rs",
"tests/admin/plugins/management.rs",
"tests/admin/plugins/mod.rs",
"tests/admin/proxies.rs",
"tests/admin/settings.rs",
"tests/admin/system.rs",
"tests/admin/wire.rs",
"tests/architecture.rs",
"tests/auth.rs",
"tests/support.rs",
"tests/health.rs",
"tests/key_usage/mod.rs",
"tests/key_usage/fixtures.rs",
"tests/main.rs",
"tests/middleware/mod.rs",
"tests/middleware/websocket.rs",
"tests/openai/middleware/mod.rs",
"tests/openai/auth.rs",
"tests/openai/error.rs",
"tests/openai/images/mod.rs",
"tests/openai/mod.rs",
"tests/openai/models.rs",
"tests/openai/endpoint.rs",
"tests/openai/responses/http.rs",
"tests/openai/responses/mod.rs",
"tests/openai/responses/request.rs",
"tests/openai/responses/websocket/connection/mod.rs",
"tests/openai/responses/websocket/connection/middleware.rs",
"tests/openai/responses/websocket/forward.rs",
"tests/openai/responses/websocket/interrupt.rs",
"tests/openai/responses/websocket/mod.rs",
"tests/openai/responses/websocket/protocol.rs",
"tests/openai/live/mod.rs",
"tests/openai/live/websocket.rs",
"tests/openai/router.rs",
"tests/openai/search/mod.rs",
"tests/openai/usage.rs",
"tests/provider.rs",
];
expected.sort_unstable();
assert_eq!(actual, expected);
}
#[test]
fn admin_routes_should_use_only_get_post_and_static_paths() {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/admin");
@@ -192,20 +48,6 @@ fn admin_routes_should_use_only_get_post_and_static_paths() {
);
}
fn rust_files(root: &Path) -> Vec<String> {
let manifest = Path::new(env!("CARGO_MANIFEST_DIR"));
all_files(root)
.into_iter()
.filter(|path| path.extension().and_then(|value| value.to_str()) == Some("rs"))
.map(|path| {
path.strip_prefix(manifest)
.expect("path below manifest")
.to_string_lossy()
.replace('\\', "/")
})
.collect()
}
fn all_files(root: &Path) -> Vec<std::path::PathBuf> {
let mut files = Vec::new();
let mut pending = vec![root.to_path_buf()];
+215 -2
View File
@@ -1,6 +1,6 @@
//! 验证健康接口按基础设施与关键 Worker 状态判定可用性
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use axum::{
body::Body,
@@ -13,12 +13,17 @@ use gateway_core::{
StartProviderExecution, StartedExecution,
},
error::GatewayError,
health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState},
health::{
HealthProbe, HealthState, WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource,
WorkerRuntimeState,
},
routing::PublicModelId,
task::{WorkerId, WorkerKind},
};
use tower::ServiceExt;
use crate::support::RecordingDiagnostics;
#[tokio::test]
async fn healthz_should_return_no_content_when_all_inputs_are_healthy() {
let response = crate::openai::api_router(Arc::new(UnusedExecution))
@@ -81,6 +86,214 @@ async fn healthz_should_reject_critical_worker_failure() {
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}
#[tokio::test]
async fn healthz_records_probe_changes_once_and_records_recurrence_after_recovery() {
let probe = Arc::new(MutableProbe(Mutex::new(HealthState::Unhealthy(
"file logging rejected records".to_owned(),
))));
let diagnostics = Arc::new(RecordingDiagnostics::default());
let router = health_router(vec![probe.clone()], Vec::new(), diagnostics.clone()).await;
let responses =
futures::future::join_all((0..8).map(|_| health_response(router.clone()))).await;
for response in responses {
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert!(
axum::body::to_bytes(response.into_body(), 1024)
.await
.unwrap()
.is_empty()
);
}
{
let events = diagnostics.0.lock().unwrap();
assert_eq!(events.len(), 1);
assert_eq!(events[0].component, "health");
assert_eq!(
events[0].correlation_id.as_deref(),
Some("probe:file_logging")
);
assert_eq!(events[0].message, "Health check failed");
assert!(
events[0]
.details
.as_ref()
.unwrap()
.as_str()
.contains("file logging rejected records")
);
}
*probe.0.lock().unwrap() = HealthState::Degraded("file logging maintenance failed".to_owned());
assert_eq!(
health_response(router.clone()).await.status(),
StatusCode::SERVICE_UNAVAILABLE
);
assert_eq!(diagnostics.0.lock().unwrap().len(), 2);
*probe.0.lock().unwrap() = HealthState::Healthy;
assert_eq!(
health_response(router.clone()).await.status(),
StatusCode::NO_CONTENT
);
*probe.0.lock().unwrap() = HealthState::Degraded("file logging maintenance failed".to_owned());
assert_eq!(
health_response(router).await.status(),
StatusCode::SERVICE_UNAVAILABLE
);
assert_eq!(diagnostics.0.lock().unwrap().len(), 3);
}
#[tokio::test(start_paused = true)]
async fn healthz_retains_worker_and_completed_probe_failures_when_another_probe_times_out() {
let diagnostics = Arc::new(RecordingDiagnostics::default());
let router = health_router(
vec![
Arc::new(MutableProbe(Mutex::new(HealthState::Unhealthy(
"file output failed".to_owned(),
)))),
Arc::new(PendingProbe),
],
vec![
worker_snapshot(
WorkerKind::RuntimeSnapshotReconciliation,
WorkerRuntimeState::BackingOff,
),
worker_snapshot(
WorkerKind::QuotaCatalogHealth,
WorkerRuntimeState::BackingOff,
),
],
diagnostics.clone(),
)
.await;
assert_eq!(
health_response(router).await.status(),
StatusCode::SERVICE_UNAVAILABLE
);
let events = diagnostics.0.lock().unwrap();
assert_eq!(events.len(), 3);
for (component, reason) in [
(
"worker:runtime_snapshot_reconciliation:health-test",
"test worker failure",
),
("probe:file_logging", "file output failed"),
("probe:pending", "timed_out"),
] {
let event = events
.iter()
.find(|event| event.correlation_id.as_deref() == Some(component))
.expect(component);
assert!(event.details.as_ref().unwrap().as_str().contains(reason));
}
}
#[tokio::test]
async fn healthz_retries_failed_diagnostic_admission_without_changing_health_result() {
let diagnostics = Arc::new(FailingOnceDiagnostics::default());
let router = health_router(
vec![Arc::new(MutableProbe(Mutex::new(HealthState::Unhealthy(
"file output failed".to_owned(),
))))],
Vec::new(),
diagnostics.clone(),
)
.await;
assert_eq!(
health_response(router.clone()).await.status(),
StatusCode::SERVICE_UNAVAILABLE
);
assert!(diagnostics.recorded.0.lock().unwrap().is_empty());
assert_eq!(
health_response(router).await.status(),
StatusCode::SERVICE_UNAVAILABLE
);
assert_eq!(diagnostics.recorded.0.lock().unwrap().len(), 1);
}
#[derive(Default)]
struct FailingOnceDiagnostics {
failed: std::sync::atomic::AtomicBool,
recorded: RecordingDiagnostics,
}
#[async_trait::async_trait]
impl gateway_core::diagnostics::OperationalDiagnostics for FailingOnceDiagnostics {
async fn record_failure(
&self,
failure: gateway_core::diagnostics::OperationalFailure,
) -> Result<(), gateway_core::error::StoreError> {
if !self.failed.swap(true, std::sync::atomic::Ordering::SeqCst) {
return Err(gateway_core::error::StoreError::new(
gateway_core::error::StoreErrorKind::Unavailable,
));
}
self.recorded.record_failure(failure).await
}
}
async fn health_router(
probes: Vec<Arc<dyn HealthProbe>>,
workers: Vec<WorkerHealthSnapshot>,
diagnostics: Arc<dyn gateway_core::diagnostics::OperationalDiagnostics>,
) -> axum::Router {
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".to_owned(),
},
Arc::new(UnusedExecution),
admin.services,
probes,
Arc::new(StaticWorkerHealth(workers)),
Arc::new(crate::openai::TestLifecycle::default()),
diagnostics,
)
.unwrap()
.router()
}
async fn health_response(router: axum::Router) -> axum::response::Response {
router
.oneshot(
Request::builder()
.uri("/healthz")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap()
}
struct MutableProbe(Mutex<HealthState>);
impl HealthProbe for MutableProbe {
fn name(&self) -> &'static str {
"file_logging"
}
fn check(&self) -> BoxFuture<'_, HealthState> {
Box::pin(std::future::ready(self.0.lock().unwrap().clone()))
}
}
struct PendingProbe;
impl HealthProbe for PendingProbe {
fn name(&self) -> &'static str {
"pending"
}
fn check(&self) -> BoxFuture<'_, HealthState> {
Box::pin(std::future::pending())
}
}
#[derive(Clone)]
struct StaticWorkerHealth(Vec<WorkerHealthSnapshot>);
@@ -135,6 +135,13 @@ fn usage_record() -> UsageListRecord {
first_text_ms: Some(600),
first_token_ms: Some(210),
provider_processing_ms: None,
upstream_response_ms: Some(1_000),
upstream_api_overhead_ms: None,
upstream_engine_ms: None,
upstream_engine_iapi_ttft_ms: None,
upstream_engine_service_ttft_ms: None,
upstream_engine_iapi_tbt_ms: Some(2.450638),
upstream_engine_service_tbt_ms: None,
latency_ms: Some(1250),
admission_decision_ms: None,
account_selection_wait_ms: None,
@@ -305,9 +305,20 @@ async fn records_keep_pagination_and_hide_admin_and_upstream_data() {
assert_eq!(record["firstTokenLatencyMs"], 210);
assert_fields(
&record["latencyDetails"],
&["firstEventMs", "firstReasoningMs", "firstTextMs"],
&[
"upstreamResponseMs",
"upstreamEngineIapiTbtMs",
"firstEventMs",
"firstReasoningMs",
"firstTextMs",
],
);
assert_eq!(record["latencyDetails"]["firstEventMs"], 100);
assert_eq!(record["latencyDetails"]["upstreamResponseMs"], 1000);
assert_eq!(
record["latencyDetails"]["upstreamEngineIapiTbtMs"],
2.450638
);
assert!(record["statusCode"].is_null());
} else {
assert_eq!(record["statusCode"], 502);
@@ -16,15 +16,18 @@ use axum::{
};
use bytes::Bytes;
use futures::{future::BoxFuture, stream};
use gateway_core::{
engine::middleware::{
FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan,
MiddlewareRequest, MiddlewareResponse,
},
lifecycle::CancellationToken,
middleware::http as contract,
runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
};
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareContext;
use gateway_core::engine::middleware::MiddlewareError;
use gateway_core::engine::middleware::MiddlewareNext;
use gateway_core::engine::middleware::MiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareRequest;
use gateway_core::engine::middleware::MiddlewareResponse;
use gateway_core::engine::middleware::http as contract;
use gateway_core::lifecycle::CancellationToken;
use gateway_core::routing::extensions::ExtensionSetId;
use gateway_core::routing::extensions::ExtensionSetLease;
use gateway_core::routing::extensions::ExtensionSetReference;
use http_body::Frame;
use http_body_util::{BodyExt as _, StreamBody};
use tower::ServiceExt as _;
@@ -70,12 +73,12 @@ impl MiddlewarePlan for Plan {
}
fn handle_websocket(
&self,
context: gateway_core::middleware::websocket::Context,
message: gateway_core::middleware::websocket::Message,
next: gateway_core::middleware::websocket::Next,
context: gateway_core::engine::middleware::websocket::Context,
message: gateway_core::engine::middleware::websocket::Message,
next: gateway_core::engine::middleware::websocket::Next,
) -> BoxFuture<
'static,
Result<Option<gateway_core::middleware::websocket::Message>, MiddlewareError>,
Result<Option<gateway_core::engine::middleware::websocket::Message>, MiddlewareError>,
> {
assert!(context.plan.is_some());
self.calls.fetch_add(1, Ordering::SeqCst);
@@ -103,11 +106,13 @@ impl MiddlewarePlan for Plan {
}
fn handle_service(
&self,
context: gateway_core::middleware::service::Context,
context: gateway_core::engine::middleware::service::Context,
input: serde_json::Value,
next: gateway_core::middleware::service::Next,
) -> BoxFuture<'static, Result<serde_json::Value, gateway_core::middleware::service::Error>>
{
next: gateway_core::engine::middleware::service::Next,
) -> BoxFuture<
'static,
Result<serde_json::Value, gateway_core::engine::middleware::service::Error>,
> {
self.calls.fetch_add(1, Ordering::SeqCst);
assert_eq!(context.operation, "settings.load");
assert_eq!(context.request_id, "service-request");
@@ -146,7 +151,7 @@ impl MiddlewarePlan for Plan {
response.extensions_mut().insert(contract::upgrade::Upgraded::new(status, Box::pin(async move {
let session = session.await.map_err(|_| MiddlewareError::Fault)?;
if idle { return std::future::pending().await; }
use gateway_core::middleware::websocket::{Kind, Message};
use gateway_core::engine::middleware::websocket::{Kind, Message};
let receive = session.receive();
tokio::pin!(receive);
tokio::select! {
+71 -111
View File
@@ -11,9 +11,7 @@ use axum::http::{
};
use futures::future::BoxFuture;
use gateway_api::openai::auth::{
ClientApiKeyAuthError, bearer_client_api_key, identify_codex_client,
};
use gateway_api::openai::auth::{ClientApiKeyAuthError, identify_codex_client};
use gateway_core::{
engine::{
authentication::ClientAuthenticationRequest,
@@ -118,114 +116,76 @@ async fn plugin_authentication_envelope_contains_only_authorization() {
assert_eq!(execution.calls.load(Ordering::SeqCst), 1);
}
#[test]
fn bearer_client_api_key_should_reject_missing_authorization() {
assert_eq!(
bearer_client_api_key(&HeaderMap::new()),
Err(ClientApiKeyAuthError::MissingAuthorization)
#[tokio::test]
async fn native_authentication_checks_the_current_envelope_at_the_http_entry() {
let fixture = crate::admin::AdminTestFixture::new().await;
let maximum_key = "x".repeat(
gateway_core::engine::authentication::MAXIMUM_AUTHORIZATION_BYTES - "Bearer ".len(),
);
}
#[test]
fn actor_authorization_marker_should_not_replace_client_api_key() {
let mut headers = HeaderMap::new();
headers.insert(
"x-openai-actor-authorization",
HeaderValue::from_static("proxy-managed"),
);
assert_eq!(
bearer_client_api_key(&headers),
Err(ClientApiKeyAuthError::MissingAuthorization)
);
}
#[test]
fn actor_authorization_marker_should_preserve_existing_bearer_authentication() {
let mut headers = HeaderMap::new();
headers.insert(
"x-openai-actor-authorization",
HeaderValue::from_static("proxy-managed"),
);
headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer sk_client"));
assert_eq!(bearer_client_api_key(&headers), Ok("sk_client"));
}
#[test]
fn bearer_client_api_key_should_reject_non_utf8_authorization() {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_bytes(&[0xff]).expect("opaque header value"),
);
assert_eq!(
bearer_client_api_key(&headers),
Err(ClientApiKeyAuthError::MalformedAuthorization)
);
}
#[test]
fn bearer_client_api_key_should_require_exact_bearer_scheme() {
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, HeaderValue::from_static("bearer sk_client"));
assert_eq!(
bearer_client_api_key(&headers),
Err(ClientApiKeyAuthError::MalformedAuthorization)
);
}
#[test]
fn bearer_client_api_key_should_reject_empty_bearer_token() {
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer "));
assert_eq!(
bearer_client_api_key(&headers),
Err(ClientApiKeyAuthError::MalformedAuthorization)
);
}
#[test]
fn bearer_client_api_key_should_accept_migrated_formats_without_rewriting() {
for key in [
"q".to_owned(),
"sk-old/key+value=:!@".to_owned(),
"x".repeat(8192),
] {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {key}")).unwrap(),
);
assert_eq!(bearer_client_api_key(&headers), Ok(key.as_str()));
let oversized_key = format!("{maximum_key}x");
let mut cases = vec![
("missing", "sk_client", None, StatusCode::UNAUTHORIZED),
(
"opaque",
"sk_client",
Some(HeaderValue::from_bytes(&[0xff]).unwrap()),
StatusCode::UNAUTHORIZED,
),
(
"lowercase scheme",
"sk_client",
Some(HeaderValue::from_static("bearer sk_client")),
StatusCode::UNAUTHORIZED,
),
(
"empty token",
"sk_client",
Some(HeaderValue::from_static("Bearer ")),
StatusCode::UNAUTHORIZED,
),
(
"whitespace inside token",
"sk_client",
Some(HeaderValue::from_static("Bearer key with spaces")),
StatusCode::UNAUTHORIZED,
),
(
"trimmed token",
"sk_client",
Some(HeaderValue::from_static("Bearer sk_client ")),
StatusCode::OK,
),
];
for key in ["q", "sk-old/key+value=:!@", maximum_key.as_str()] {
cases.push((
"configured native key",
key,
Some(HeaderValue::from_str(&format!("Bearer {key}")).unwrap()),
StatusCode::OK,
));
}
}
cases.push((
"oversized envelope",
oversized_key.as_str(),
Some(HeaderValue::from_str(&format!("Bearer {oversized_key}")).unwrap()),
StatusCode::UNAUTHORIZED,
));
#[test]
fn bearer_client_api_key_should_reject_whitespace_inside_tokens() {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer key with spaces"),
);
assert_eq!(
bearer_client_api_key(&headers),
Err(ClientApiKeyAuthError::InvalidKeyFormat)
);
}
#[test]
fn bearer_client_api_key_should_return_trimmed_gateway_key() {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer sk_client_secret "),
);
assert_eq!(bearer_client_api_key(&headers), Ok("sk_client_secret"));
for (case, key, authorization, expected) in cases {
let app =
super::api_router_with_admin_and_client(fixture.services.clone(), key, "key-auth");
let mut request = crate::support::empty_request(Method::GET, "/v1/models");
// actor 标记既不能代替认证,也不能干扰已有的原生 Key
request.headers_mut().insert(
"x-openai-actor-authorization",
HeaderValue::from_static("proxy-managed"),
);
if let Some(authorization) = authorization {
request.headers_mut().insert(AUTHORIZATION, authorization);
}
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), expected, "{case}");
}
}
#[test]
@@ -414,9 +374,9 @@ fn chatgpt_remote_suffix_created_by_header_truncation_should_not_skip_the_gate()
#[tokio::test]
async fn http_settings_freeze_before_plan_resolution_and_apply_before_admission() {
use gateway_core::{
engine::middleware::*, middleware::http as contract, runtime::extensions::*,
};
use gateway_core::engine::middleware::http as contract;
use gateway_core::engine::middleware::*;
use gateway_core::routing::extensions::*;
struct Lease;
impl ExtensionSetLease for Lease {
fn is_ready(&self) -> bool {
@@ -10,6 +10,7 @@ use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header::AUTHORIZATION},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use bytes::Bytes;
use futures::future::BoxFuture;
use gateway_core::engine::execution::{
@@ -300,6 +301,11 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio
.header("x-codex-image-turn-id", "turn_image_route")
.header("session-id", "root-image-session")
.header("thread-id", "child-image-thread")
.header("x-future-business", "first")
.header("x-future-business", "second")
.header("cookie", "client-secret")
.header("connection", "x-private-hop")
.header("x-private-hop", "client-hop")
.body(Body::from(body.to_vec()))
.expect("image request"),
)
@@ -338,8 +344,37 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio
assert_eq!(captured.transport, ClientTransport::HttpJson);
assert_eq!(captured.kind, *kind);
assert_eq!(captured.body.as_ref(), *body);
let mut context = captured.context.clone();
let headers = context
.as_object_mut()
.unwrap()
.remove("opaque_request_headers")
.unwrap();
let headers: Vec<_> = headers
.as_array()
.unwrap()
.iter()
.map(|entry| {
(
entry[0].as_str().unwrap(),
STANDARD.decode(entry[1].as_str().unwrap()).unwrap(),
)
})
.collect();
assert_eq!(
captured.context,
headers
.iter()
.filter(|(name, _)| *name == "x-future-business")
.map(|(_, value)| value.as_slice())
.collect::<Vec<_>>(),
[b"first".as_slice(), b"second".as_slice()]
);
assert!(!headers.iter().any(|(name, _)| matches!(
*name,
"authorization" | "cookie" | "connection" | "x-private-hop"
)));
assert_eq!(
context,
json!({"image_turn_id": "turn_image_route", "session_id": "root-image-session", "thread_id": "child-image-thread"})
);
}
@@ -10,15 +10,22 @@ use std::{
use bytes::Bytes;
use futures::future::BoxFuture;
use gateway_core::{
engine::middleware::{
FrozenMiddlewarePlan, MiddlewareBody, MiddlewareContext, MiddlewareError, MiddlewareFrame,
MiddlewareFraming, MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewarePlan,
MiddlewareRequest, MiddlewareResponse,
},
operation::OperationKind,
runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
};
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareBody;
use gateway_core::engine::middleware::MiddlewareContext;
use gateway_core::engine::middleware::MiddlewareError;
use gateway_core::engine::middleware::MiddlewareFrame;
use gateway_core::engine::middleware::MiddlewareFraming;
use gateway_core::engine::middleware::MiddlewareHeader;
use gateway_core::engine::middleware::MiddlewareMount;
use gateway_core::engine::middleware::MiddlewareNext;
use gateway_core::engine::middleware::MiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareRequest;
use gateway_core::engine::middleware::MiddlewareResponse;
use gateway_core::operation::OperationKind;
use gateway_core::routing::extensions::ExtensionSetId;
use gateway_core::routing::extensions::ExtensionSetLease;
use gateway_core::routing::extensions::ExtensionSetReference;
/// OpenAI API 入口共用的中间件夹具;授权和 RPC 由 Runtime 行为测试覆盖
#[derive(Debug, Default)]
@@ -30,7 +37,8 @@ pub(super) struct RequestMiddleware {
pub(super) response_actions: Mutex<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)>,
pub(super) settings:
Option<fn(&mut gateway_core::routing::request_settings::ExecutionSettings)>,
}
#[derive(Debug, Clone)]
@@ -323,7 +323,7 @@ fn snapshot_with_client_key(
}
#[derive(Default)]
struct TestLifecycle {
pub(crate) struct TestLifecycle {
cancellation: gateway_core::lifecycle::CancellationToken,
}
@@ -2279,6 +2279,7 @@ async fn buffered_response_commits_only_after_complete_json_is_encoded() {
async fn buffered_response_collects_completed_output_items_without_rewriting_streams() {
let message = json!({
"id": "msg_test", "type": "message", "role": "assistant", "status": "completed",
"phase": "partial_answer",
"content": [{"type": "output_text", "text": "完整文本", "annotations": []}],
"future_item_field": {"keep": true}
});
@@ -2288,6 +2289,7 @@ async fn buffered_response_collects_completed_output_items_without_rewriting_str
});
let terminal = json!({
"id": "resp_test", "status": "completed", "output": [],
"end_turn": false,
"usage": {"input_tokens": 20, "output_tokens": 8, "total_tokens": 28},
"future_terminal_field": {"keep": true}
});
@@ -3020,3 +3022,67 @@ async fn compressed_http_request_above_default_limit_should_reach_execution_afte
}
}
}
#[tokio::test]
async fn streaming_middleware_preserves_heartbeats_before_and_after_rewriting() {
for opaque in [
b": ping\n\n".as_slice(),
b"id: next\nretry: 1000\n\n".as_slice(),
b"event: vendor.future\ndata: future non-JSON payload\n\n".as_slice(),
] {
let admin = crate::admin::AdminTestFixture::new().await.services;
let trace = Arc::new(Trace::default());
let heartbeat = Bytes::copy_from_slice(opaque);
let future = Bytes::from_static(
b"event: response.future\ndata: {\"type\":\"response.future\",\"extension\":true}\n\n",
);
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::Event(delivery(started(), CommitRequirement::CommitBeforeDelivery)),
NextStep::Event(delivery_provider(
ProviderEvent::wire(
ProtocolWireEvent::raw_sse("openai", heartbeat.clone()).unwrap(),
),
CommitRequirement::AlreadyCommitted,
)),
NextStep::Event(delivery_provider(
ProviderEvent::wire(
ProtocolWireEvent::raw_sse("openai", future.clone()).unwrap(),
),
CommitRequirement::AlreadyCommitted,
)),
NextStep::Event(delivery_provider(
ProviderEvent::wire(
ProtocolWireEvent::raw_sse("openai", heartbeat.clone()).unwrap(),
),
CommitRequirement::AlreadyCommitted,
)),
NextStep::Event(delivery(completed(), CommitRequirement::AlreadyCommitted)),
NextStep::FinalizeSuccess,
],
);
let replacement = Bytes::from_static(b"event: response.future\ndata: {\"type\":\"response.future\",\"extension\":\"changed\"}\n\n");
let middleware = Arc::new(responses_middleware().with_response_actions(vec![
ResponseFrameAction::Continue,
ResponseFrameAction::Continue,
ResponseFrameAction::Replace(replacement.clone()),
ResponseFrameAction::Continue,
ResponseFrameAction::Continue,
ResponseFrameAction::Continue,
]));
let response =
response_with_middleware(&admin, "openai", session, true, Some(middleware)).await;
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let text = std::str::from_utf8(&body).unwrap();
assert_eq!(
text.matches(std::str::from_utf8(opaque).unwrap()).count(),
2
);
assert!(text.contains(std::str::from_utf8(&replacement).unwrap()));
assert!(text.contains("response.completed"));
assert!(!text.contains("middleware_failed"));
assert!(text.ends_with("data: [DONE]\n\n"));
assert!(!trace.is_cancelled());
}
}
@@ -1071,3 +1071,35 @@ fn transparent_encoder_should_only_require_wire_terminal_for_buffered_conversion
ResponseEncodeError::MissingWireTerminal
);
}
#[test]
fn metadata_keeps_business_extensions_and_filters_only_protected_headers() {
let data = json!({
"type":"response.metadata",
"headers":{"x-future-header":["first","second"],"authorization":"private","set-cookie":"private","chatgpt-account-id":"private","connection":"x-private-hop","x-private-hop":"private"},
"metadata":{"type":"safety_buffering","use_cases":["cyber"],"reasons":["user_risk"],"openai_verification_recommendation":{"future":true},"openai_chatgpt_moderation_metadata":{"future":true}},
"future_extension":{"keep":true}
});
let raw = data.to_string();
let event = ProviderEvent::wire(
ProtocolWireEvent::json("openai", Some("response.metadata".to_owned()), data.clone())
.unwrap()
.with_raw_websocket_message(raw.as_str()),
);
let mut encoder = OpenAiResponsesEncoder::new();
let ws = encoder.push_websocket(&event);
let delivered: Value = serde_json::from_str(&ws[0]).unwrap();
assert_eq!(delivered["metadata"], data["metadata"]);
assert_eq!(delivered["future_extension"], data["future_extension"]);
assert_eq!(
delivered["headers"],
json!({"x-future-header":["first","second"]})
);
let sse = encoder.push_sse(&event);
assert!(!std::str::from_utf8(&sse[0]).unwrap().contains("private"));
assert!(
std::str::from_utf8(&sse[0])
.unwrap()
.contains("safety_buffering")
);
}
@@ -2,14 +2,100 @@
use super::*;
use bytes::Bytes;
use gateway_core::{
engine::middleware::{
FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareNext, MiddlewarePlan,
MiddlewareRequest, MiddlewareResponse,
},
middleware::{http as http_contract, websocket as core},
runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
};
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareContext;
use gateway_core::engine::middleware::MiddlewareError;
use gateway_core::engine::middleware::MiddlewareNext;
use gateway_core::engine::middleware::MiddlewarePlan;
use gateway_core::engine::middleware::MiddlewareRequest;
use gateway_core::engine::middleware::MiddlewareResponse;
use gateway_core::engine::middleware::http as http_contract;
use gateway_core::engine::middleware::websocket as core;
use gateway_core::routing::extensions::ExtensionSetId;
use gateway_core::routing::extensions::ExtensionSetLease;
use gateway_core::routing::extensions::ExtensionSetReference;
#[tokio::test(start_paused = true)]
async fn client_close_does_not_flush_a_retained_sender_with_an_inflight_write() {
#[derive(Default)]
struct RetainedSenderPlan(Mutex<Option<Arc<dyn core::Sender>>>);
impl fmt::Debug for RetainedSenderPlan {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("RetainedSenderPlan")
}
}
impl MiddlewarePlan for RetainedSenderPlan {
fn has_websocket(&self) -> bool {
true
}
fn handle_websocket(
&self,
context: core::Context,
message: core::Message,
next: core::Next,
) -> BoxFuture<'static, Result<Option<core::Message>, MiddlewareError>> {
*self.0.lock().unwrap() = Some(context.sender);
next.run(message)
}
fn handle(
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
next.run(request)
}
}
let plan = Arc::new(RetainedSenderPlan::default());
let frozen = FrozenMiddlewarePlan::new(
plan.clone(),
ExtensionSetReference::new(
ExtensionSetId::new("retained-writer".into()).unwrap(),
Arc::new(Lease),
),
);
let (incoming, received) = unbounded_channel();
let (written, mut output) = unbounded_channel();
let flush_started = Arc::new(AtomicBool::new(false));
let socket = TestSocket {
incoming: received,
written,
stall_writes: true,
flush_started: flush_started.clone(),
dropped: Arc::new(AtomicBool::new(false)),
};
let mut connection = spawn_connection(
socket,
Arc::from("retained-writer"),
CancellationToken::new(),
ConnectionConfig::PRODUCTION,
Some(frozen),
Arc::from([]),
);
incoming.send(Ok(Message::Text("capture".into()))).unwrap();
assert!(matches!(
connection.next_event().await,
Some(ConnectionEvent::Text(_))
));
let sender = plan.0.lock().unwrap().clone().unwrap();
let mut write = sender.send(core::Message {
kind: core::Kind::Text,
payload: Bytes::from_static(b"pending"),
});
assert!(futures::poll!(write.as_mut()).is_pending());
// 写 future 在 pump 的 select 之外持锁,而且暂不继续 poll 以响应取消
incoming.send(Ok(Message::Close(None))).unwrap();
let reason = tokio::time::timeout(Duration::from_millis(50), connection.wait_for_exit())
.await
.expect("peer close must not wait for a retained sender");
assert_eq!(reason, PumpExitReason::ClientClose);
assert!(!flush_started.load(Ordering::Acquire));
assert!(write.await.is_err());
assert!(matches!(output.try_recv(), Err(TryRecvError::Empty)));
plan.0.lock().unwrap().take();
}
#[derive(Debug, Default)]
struct Plan {
block_terminal: bool,
@@ -234,6 +234,7 @@ struct TestSocket {
incoming: UnboundedReceiver<Result<Message, TestSocketError>>,
written: UnboundedSender<Message>,
stall_writes: bool,
flush_started: Arc<AtomicBool>,
dropped: Arc<AtomicBool>,
}
@@ -268,7 +269,12 @@ impl Sink<Message> for TestSocket {
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
self.flush_started.store(true, Ordering::Release);
if self.stall_writes {
Poll::Pending
} else {
Poll::Ready(Ok(()))
}
}
fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
@@ -286,6 +292,7 @@ struct PumpHarness {
connection: ResponsesWebSocketConnection,
incoming: UnboundedSender<Result<Message, TestSocketError>>,
written: UnboundedReceiver<Message>,
flush_started: Arc<AtomicBool>,
dropped: Arc<AtomicBool>,
cancellation: CancellationToken,
}
@@ -294,11 +301,13 @@ fn test_connection(stall_writes: bool) -> PumpHarness {
let (incoming_tx, incoming_rx) = unbounded_channel();
let (written_tx, written_rx) = unbounded_channel();
let dropped = Arc::new(AtomicBool::new(false));
let flush_started = Arc::new(AtomicBool::new(false));
let cancellation = CancellationToken::new();
let socket = TestSocket {
incoming: incoming_rx,
written: written_tx,
stall_writes,
flush_started: Arc::clone(&flush_started),
dropped: Arc::clone(&dropped),
};
let connection = spawn_connection(
@@ -306,11 +315,14 @@ fn test_connection(stall_writes: bool) -> PumpHarness {
Arc::from("ws_test"),
cancellation.clone(),
ConnectionConfig::PRODUCTION,
None,
Arc::from([]),
);
PumpHarness {
connection,
incoming: incoming_tx,
written: written_rx,
flush_started,
dropped,
cancellation,
}
@@ -402,6 +414,62 @@ async fn client_close_should_take_priority_over_queued_business_frames() {
));
}
#[tokio::test(start_paused = true)]
async fn client_close_flush_has_a_short_deadline_without_sending_another_frame() {
let PumpHarness {
mut connection,
incoming,
mut written,
flush_started,
dropped,
..
} = test_connection(true);
incoming.send(Ok(Message::Close(None))).expect("close peer");
let started = tokio::time::Instant::now();
let reason = tokio::time::timeout(Duration::from_secs(2), connection.wait_for_exit())
.await
.expect("peer close cannot consume the ordinary write timeout");
assert_eq!(reason, PumpExitReason::ClientClose);
assert_eq!(started.elapsed(), Duration::from_secs(1));
assert!(flush_started.load(Ordering::Acquire));
assert!(dropped.load(Ordering::Acquire));
assert!(
written.recv().await.is_none(),
"do not send a duplicate Close"
);
}
#[tokio::test(start_paused = true)]
async fn lifecycle_shutdown_interrupts_a_pending_client_close_flush() {
let PumpHarness {
mut connection,
incoming,
flush_started,
cancellation,
dropped,
..
} = test_connection(true);
incoming.send(Ok(Message::Close(None))).expect("close peer");
for _ in 0..100 {
if flush_started.load(Ordering::Acquire) {
break;
}
tokio::task::yield_now().await;
}
assert!(
flush_started.load(Ordering::Acquire),
"close flush must start"
);
cancellation.cancel();
let reason = tokio::time::timeout(Duration::from_millis(50), connection.wait_for_exit())
.await
.expect("host cancellation must not wait for the close timeout");
assert_eq!(reason, PumpExitReason::ClientClose);
assert!(dropped.load(Ordering::Acquire));
}
#[tokio::test]
async fn full_business_queue_should_not_block_lifecycle_shutdown() {
let PumpHarness {
@@ -986,6 +1054,40 @@ async fn requests_deferred_during_active_response_share_the_inbound_queue_limit(
server.abort();
}
#[tokio::test]
async fn client_normal_close_after_completed_response_receives_normal_close_reply() {
let (trace, mut socket, server) = start_active_response().await;
trace.release_terminal.notify_one();
loop {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("terminal response timeout")
.expect("response remains open")
.expect("valid response frame");
if let ClientMessage::Text(text) = message {
let event: Value = serde_json::from_str(&text).expect("response JSON");
if event["type"] == "response.completed" {
break;
}
}
}
socket
.close(Some(tokio_tungstenite::tungstenite::protocol::CloseFrame {
code: CloseCode::Normal,
reason: "client finished".into(),
}))
.await
.expect("send normal close");
// close() 只发送本端关闭帧;必须读取对端应答,才能证明握手正常完成
let reply = tokio::time::timeout(Duration::from_secs(2), socket.next()).await;
server.abort();
assert!(
matches!(&reply, Ok(Some(Ok(ClientMessage::Close(Some(frame)))))
if frame.code == CloseCode::Normal && frame.reason == "client finished"),
"client must receive the close acknowledgement, got {reply:?}"
);
}
#[tokio::test]
async fn client_close_should_cancel_active_execution_without_starting_queued_requests() {
let (trace, mut socket, server) = start_active_response().await;
@@ -1007,6 +1109,12 @@ async fn client_close_should_cancel_active_execution_without_starting_queued_req
.expect("cancel active execution promptly");
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
assert!(trace.cancelled.load(Ordering::Acquire));
let reply = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("close reply timeout")
.expect("close reply")
.expect("clean close handshake");
assert!(matches!(reply, ClientMessage::Close(None)));
server.abort();
}
@@ -1071,6 +1179,7 @@ async fn stalled_heartbeat_can_be_cancelled_without_waiting_for_write_timeout()
written: _written,
dropped,
cancellation,
..
} = test_connection(true);
tokio::time::advance(Duration::from_secs(25)).await;
tokio::task::yield_now().await;
@@ -19,7 +19,7 @@ use gateway_core::error::{
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ProviderError, ProviderErrorKind,
};
use gateway_core::event::{ProtocolWireEvent, ProviderEvent, ProviderResponseHeader};
use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference};
use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference};
use gateway_core::upstream::{OpaqueUpstreamValue, UpstreamSendState};
use serde_json::{Value, json};
use tokio::net::TcpStream;
@@ -106,10 +106,18 @@ async fn next_error(socket: &mut TestSocket) -> Value {
}
async fn assert_no_duplicate_failure(socket: &mut TestSocket) {
// 后续入站帧只会在本轮收敛后处理,比固定等待窗口更可靠地检查没有重复失败终态
socket.send(Message::Text("{".into())).await.unwrap();
// 无效创建帧仍走串行队列,以本地业务校验作为本轮收敛后的屏障
socket
.send(Message::Text(
json!({"type":"response.create","model":"model-a","stream":false})
.to_string()
.into(),
))
.await
.unwrap();
let error = next_event(socket).await;
assert_eq!(error["error"]["code"], "invalid_json");
assert_eq!(error["error"]["code"], "invalid_value");
assert_eq!(error["error"]["param"], "stream");
assert_eq!(error["status"], 400);
assert_eq!(error["headers"]["x-request-id"], error["request_id"]);
assert_eq!(
@@ -1,12 +1,70 @@
//! 验证 Responses WebSocket 中断向当前执行传递并释放相关资源
//! 验证未知控制消息、迟到中断、结算期间收发和创建请求的串行准入
use super::*;
use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError};
use gateway_core::engine::response_control::{
ResponseControlTransport, ResponseControlUnavailable,
};
struct ControlTransport {
id: String,
active: std::sync::atomic::AtomicBool,
requests: Mutex<Vec<String>>,
sender: tokio::sync::mpsc::Sender<Value>,
receiver: tokio::sync::Mutex<tokio::sync::mpsc::Receiver<Value>>,
}
impl ControlTransport {
fn new(id: String) -> Self {
let (sender, receiver) = tokio::sync::mpsc::channel(32);
Self {
id,
active: std::sync::atomic::AtomicBool::new(true),
requests: Mutex::default(),
sender,
receiver: tokio::sync::Mutex::new(receiver),
}
}
}
#[async_trait]
impl ResponseControlTransport for ControlTransport {
async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable> {
self.requests.lock().unwrap().push(payload.to_owned());
let original: Value =
serde_json::from_str(payload).unwrap_or_else(|_| Value::String(payload.to_owned()));
let event = if original["type"] == "future.reject" {
upstream_control_error()
} else if original["type"] == "response.interrupt"
&& original["response_id"] == self.id
&& original["mode"] == "discard_partial_items"
&& self.active.load(Ordering::SeqCst)
{
json!({"type":"response.incomplete", "response":{"id":self.id,"model":"model-a","status":"incomplete","output":[],"incomplete_details":{"reason":"interrupted"}}})
} else {
json!({"type":"future.control.ack", "original":original})
};
self.sender
.send(event)
.await
.map_err(|_| ResponseControlUnavailable)
}
async fn receive(&self) -> Result<String, ResponseControlUnavailable> {
self.receiver
.lock()
.await
.recv()
.await
.map(|event| event.to_string())
.ok_or(ResponseControlUnavailable)
}
}
#[derive(Default)]
struct InterruptProvider {
calls: AtomicUsize,
controls: Mutex<Vec<ResponseControl>>,
transport: Mutex<Option<Arc<ControlTransport>>>,
complete: tokio::sync::Notify,
initial_messages: Mutex<std::collections::VecDeque<ProtocolWireEvent>>,
}
#[async_trait]
@@ -30,12 +88,10 @@ impl Provider for InterruptProvider {
) -> Result<ProviderStream, ProviderError> {
let number = self.calls.fetch_add(1, Ordering::SeqCst);
let id = format!("resp_interrupt_{number}");
let control = context
.response_control()
.expect("root control propagated through Core")
.clone();
let active = control.activate(id.clone()).unwrap();
self.controls.lock().unwrap().push(control);
let transport = Arc::new(ControlTransport::new(id));
let owner: Arc<dyn ResponseControlTransport> = transport.clone();
context.response_control().unwrap().bind(&owner);
*self.transport.lock().unwrap() = Some(transport.clone());
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
@@ -44,41 +100,55 @@ impl Provider for InterruptProvider {
UpstreamTransport::new("websocket").unwrap(),
);
let body = futures::stream::unfold(
(0, Some(active), id),
|(phase, mut active, id)| async move {
let (event_type, status, event) = match phase {
(0, transport, self),
|(phase, transport, provider)| async move {
let pending = if phase == 1 {
provider.initial_messages.lock().unwrap().pop_front()
} else {
None
};
if let Some(wire) = pending {
return Some((Ok(ProviderEvent::wire(wire)), (phase, transport, provider)));
}
let (wire, events, next_phase) = match phase {
0 => (
"response.created",
"in_progress",
GatewayEvent::Started(ResponseMeta::new(&id, "model-a")),
json!({"type":"response.created","response":{"id":transport.id,"model":"model-a","status":"in_progress","output":[]}}),
vec![GatewayEvent::Started(ResponseMeta::new(
&transport.id,
"model-a",
))],
1,
),
1 => {
active.as_ref().unwrap().requested().await;
drop(active.take());
(
"response.incomplete",
"incomplete",
GatewayEvent::Completed(ResponseMeta::new(&id, "model-a")),
)
let wire = tokio::select! {
reply = transport.receive() => serde_json::from_str::<Value>(&reply.unwrap()).unwrap(),
() = provider.complete.notified() => json!({"type":"response.completed","response":{"id":transport.id,"model":"model-a","status":"completed","output":[]}}),
};
if matches!(
wire["type"].as_str(),
Some("response.completed" | "response.incomplete")
) {
transport.active.store(false, Ordering::SeqCst);
(
wire,
vec![GatewayEvent::Completed(ResponseMeta::new(
&transport.id,
"model-a",
))],
2,
)
} else {
(wire, Vec::new(), 1)
}
}
_ => return None,
};
let mut response = json!({"id":id,"model":"model-a","status":status,"output":[]});
if phase == 1 {
response["incomplete_details"] = json!({"reason":"interrupted"});
}
let event_type = wire["type"].as_str().unwrap().to_owned();
let event = ProviderEvent::canonical_with_wire(
vec![event],
ProtocolWireEvent::json(
"openai",
Some(event_type.to_owned()),
json!({
"type":event_type, "response":response
}),
)
.unwrap(),
events,
ProtocolWireEvent::json("openai", Some(event_type), wire).unwrap(),
);
Some((Ok(event), (phase + 1, active, id)))
Some((Ok(event), (next_phase, transport, provider)))
},
);
Ok(ProviderStream::new(metadata, body, ()))
@@ -88,6 +158,8 @@ impl Provider for InterruptProvider {
#[derive(Default)]
struct InterruptAdmissions {
active: AtomicUsize,
release_started: tokio::sync::Notify,
release_gate: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
}
impl ClientAdmissionPort for InterruptAdmissions {
@@ -109,6 +181,11 @@ impl ClientAdmissionPort for InterruptAdmissions {
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async {
let gate = self.release_gate.lock().unwrap().take();
self.release_started.notify_one();
if let Some(gate) = gate {
gate.await.unwrap();
}
self.active.fetch_sub(1, Ordering::SeqCst);
Ok(true)
})
@@ -151,15 +228,15 @@ async fn next(socket: &mut TestSocket) -> Value {
}
}
#[tokio::test]
async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
let provider = Arc::new(InterruptProvider::default());
let admissions = Arc::new(InterruptAdmissions::default());
async fn connect(
provider: Arc<InterruptProvider>,
admissions: Arc<InterruptAdmissions>,
) -> (TestSocket, Server) {
let execution = Arc::new(DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot("sk_ws_interrupt", "openai")),
Arc::new(SettlementPorts::default()),
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
admissions.clone(),
ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap(),
admissions,
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
Arc::new(crate::support::RecordingDiagnostics::default()),
@@ -167,7 +244,7 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let _server = Server(tokio::spawn(async move {
let server = Server(tokio::spawn(async move {
axum::serve(listener, app).await.unwrap()
}));
let mut request = format!("ws://{address}/v1/responses")
@@ -177,29 +254,136 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
AUTHORIZATION,
HeaderValue::from_static("Bearer sk_ws_interrupt"),
);
let (mut socket, _) = connect_async(request).await.unwrap();
let (socket, _) = connect_async(request).await.unwrap();
(socket, server)
}
fn upstream_control_error() -> Value {
json!({"type":"error", "status":400, "error":{"type":"invalid_request_error","code":"future_rejected","message":"Synthetic upstream rejection","param":"type"},"vendor_extension":true})
}
fn interrupt(id: &str, mode: &str) -> Value {
json!({"type":"response.interrupt","response_id":id,"mode":mode,"future_extension":true})
}
#[tokio::test]
async fn late_interrupt_after_terminal_is_sent_upstream_and_keeps_the_connection_reusable() {
let provider = Arc::new(InterruptProvider::default());
let admissions = Arc::new(InterruptAdmissions::default());
let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await;
let create = json!({"type":"response.create","model":"model-a","input":"hello"});
for (number, terminal) in [(0, "response.completed"), (1, "response.incomplete")] {
let id = format!("resp_interrupt_{number}");
send(&mut socket, create.clone()).await;
assert_eq!(next(&mut socket).await["type"], "response.created");
if number == 0 {
provider.complete.notify_one();
} else {
send(&mut socket, interrupt(&id, "discard_partial_items")).await;
}
assert_eq!(next(&mut socket).await["type"], terminal);
tokio::time::timeout(Duration::from_secs(5), async {
while admissions.active.load(Ordering::SeqCst) != 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
for control in [
interrupt(&id, "discard_partial_items"),
interrupt(&id, "future_mode"),
json!({"type":"future.control","extra":[1,2]}),
] {
send(&mut socket, control.clone()).await;
assert_eq!(
next(&mut socket).await,
json!({"type":"future.control.ack","original":control})
);
}
send(&mut socket, json!({"type":"future.reject"})).await;
assert_eq!(next(&mut socket).await, upstream_control_error());
assert_eq!(provider.calls.load(Ordering::SeqCst), number + 1);
assert_eq!(admissions.active.load(Ordering::SeqCst), 0);
}
socket.close(None).await.unwrap();
}
#[tokio::test]
async fn late_interrupt_during_terminal_settlement_does_not_abort_or_reorder_requests() {
let provider = Arc::new(InterruptProvider::default());
let (release, wait_for_release) = tokio::sync::oneshot::channel();
let admissions = Arc::new(InterruptAdmissions {
release_gate: Mutex::new(Some(wait_for_release)),
..Default::default()
});
let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await;
let create = json!({"type":"response.create","model":"model-a","input":"hello"});
let interrupt =
|id: &str, mode: &str| json!({"type":"response.interrupt","response_id":id,"mode":mode});
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
assert_eq!(next(&mut socket).await["type"], "error");
assert_eq!(provider.calls.load(Ordering::SeqCst), 0);
send(&mut socket, create.clone()).await;
assert_eq!(next(&mut socket).await["type"], "response.created");
provider.complete.notify_one();
assert_eq!(next(&mut socket).await["type"], "response.completed");
tokio::time::timeout(
Duration::from_secs(5),
admissions.release_started.notified(),
)
.await
.unwrap();
send(&mut socket, create).await;
for invalid in [
interrupt("resp_other", "discard_partial_items"),
interrupt("resp_interrupt_0", "unknown"),
interrupt("", "discard_partial_items"),
let late = interrupt("resp_interrupt_0", "discard_partial_items");
send(&mut socket, late.clone()).await;
assert_eq!(
next(&mut socket).await,
json!({"type":"future.control.ack", "original":late})
);
assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
assert_eq!(admissions.active.load(Ordering::SeqCst), 1);
release.send(()).unwrap();
assert_eq!(
next(&mut socket).await["response"]["id"],
"resp_interrupt_1"
);
socket.close(None).await.unwrap();
}
#[tokio::test]
async fn unknown_controls_preserve_bytes_and_do_not_admit_queued_creates() {
let provider = Arc::new(InterruptProvider::default());
let admissions = Arc::new(InterruptAdmissions::default());
let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await;
let create = json!({"type":"response.create","model":"model-a","input":"hello"});
send(&mut socket, json!({"type":"future.control"})).await;
let error = next(&mut socket).await;
assert_eq!(error["type"], "error");
assert_ne!(error["error"]["param"], "type");
assert_eq!(provider.calls.load(Ordering::SeqCst), 0);
send(&mut socket, create).await;
assert_eq!(next(&mut socket).await["type"], "response.created");
// 重复 type 仍是创建请求,不能借投影解码失败绕过串行准入
socket.send(ClientMessage::Text(r#"{"type":"response.create","type":"response.create","model":"model-a","input":"next"}"#.into())).await.unwrap();
for raw in [
r#"{ "type": "future.control", "mode": "unknown", "extension": [1, 2] }"#,
r#"{"extension":"no local type contract"}"#,
r#"{"type":42,"extension":true}"#,
"not-json",
] {
send(&mut socket, invalid).await;
let error = next(&mut socket).await;
assert_eq!(error["type"], "error");
assert_eq!(error["status"], 400);
socket.send(ClientMessage::Text(raw.into())).await.unwrap();
let expected =
serde_json::from_str::<Value>(raw).unwrap_or_else(|_| Value::String(raw.to_owned()));
assert_eq!(next(&mut socket).await["original"], expected);
assert_eq!(
provider
.transport
.lock()
.unwrap()
.as_ref()
.unwrap()
.requests
.lock()
.unwrap()
.last()
.unwrap(),
raw
);
assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
}
send(
@@ -207,27 +391,12 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
let terminal = next(&mut socket).await;
assert_eq!(terminal["type"], "response.incomplete");
assert_eq!(
terminal["response"]["incomplete_details"]["reason"],
"interrupted"
);
assert_eq!(next(&mut socket).await["type"], "response.incomplete");
assert_eq!(
next(&mut socket).await["response"]["id"],
"resp_interrupt_1"
);
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
assert_eq!(next(&mut socket).await["type"], "error");
assert_eq!(
provider.controls.lock().unwrap()[0].interrupt("resp_interrupt_0"),
Err(ResponseInterruptError::Unavailable)
);
socket.close(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(5), async {
while admissions.active.load(Ordering::SeqCst) != 0 {
@@ -236,8 +405,46 @@ async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
})
.await
.expect("disconnect releases active execution");
assert_eq!(
provider.controls.lock().unwrap()[1].interrupt("resp_interrupt_1"),
Err(ResponseInterruptError::Unavailable)
);
}
#[tokio::test]
async fn active_websocket_delivery_preserves_unparsed_messages_and_business_metadata() {
let messages = vec![
"{ \"type\": \"future.event\", \"number\": 1e3, \"escaped\": \"\\u0061\" }\r\n".to_owned(),
r#"{"type":"response.metadata","metadata":{"type":"safety_buffering","use_cases":["cyber"],"reasons":["user_risk"],"future":{"keep":true}}}"#.to_owned(),
format!("{{\"type\":\"future.deep\",\"extension\":{}0{}}}", "[".repeat(140), "]".repeat(140)),
"future non-JSON text".to_owned(), "[DONE]".to_owned(),
];
let provider = Arc::new(InterruptProvider::default());
*provider.initial_messages.lock().unwrap() = messages
.iter()
.map(|raw| match serde_json::from_str::<Value>(raw) {
Ok(value) => ProtocolWireEvent::json(
"openai",
value.get("type").and_then(Value::as_str).map(str::to_owned),
value,
)
.unwrap()
.with_raw_websocket_message(raw.as_str()),
Err(_) => ProtocolWireEvent::raw_websocket("openai", raw.as_str()).unwrap(),
})
.collect();
let (mut socket, _server) =
connect(provider.clone(), Arc::new(InterruptAdmissions::default())).await;
send(
&mut socket,
json!({"type":"response.create","model":"model-a","input":"synthetic"}),
)
.await;
assert_eq!(next(&mut socket).await["type"], "response.created");
for expected in &messages {
let actual = tokio::time::timeout(Duration::from_secs(5), socket.next())
.await
.unwrap()
.unwrap()
.unwrap();
assert_eq!(actual.to_text().unwrap(), expected);
}
provider.complete.notify_one();
assert_eq!(next(&mut socket).await["type"], "response.completed");
}
@@ -7,6 +7,7 @@ use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header::AUTHORIZATION},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use bytes::Bytes;
use futures::future::BoxFuture;
use gateway_core::engine::ModelRequestId;
@@ -181,6 +182,11 @@ async fn search_route_should_preserve_request_and_success_response_bytes() {
.header("x-codex-turn-metadata", turn_metadata)
.header("session-id", "root-session")
.header("thread-id", "child-thread")
.header("x-future-business", "first")
.header("x-future-business", "second")
.header("cookie", "client-secret")
.header("connection", "x-private-hop")
.header("x-private-hop", "client-hop")
.body(Body::from(request_body.to_vec()))
.expect("search request"),
)
@@ -213,8 +219,37 @@ async fn search_route_should_preserve_request_and_success_response_bytes() {
assert_eq!(captured[0].endpoint, "/v1/alpha/search");
assert_eq!(captured[0].transport, ClientTransport::HttpJson);
assert_eq!(captured[0].body.as_ref(), request_body);
let mut context = captured[0].context.clone();
let headers = context
.as_object_mut()
.unwrap()
.remove("opaque_request_headers")
.unwrap();
let headers: Vec<_> = headers
.as_array()
.unwrap()
.iter()
.map(|entry| {
(
entry[0].as_str().unwrap(),
STANDARD.decode(entry[1].as_str().unwrap()).unwrap(),
)
})
.collect();
assert_eq!(
captured[0].context,
headers
.iter()
.filter(|(name, _)| *name == "x-future-business")
.map(|(_, value)| value.as_slice())
.collect::<Vec<_>>(),
[b"first".as_slice(), b"second".as_slice()]
);
assert!(!headers.iter().any(|(name, _)| matches!(
*name,
"authorization" | "cookie" | "connection" | "x-private-hop"
)));
assert_eq!(
context,
json!({"turn_metadata": turn_metadata, "session_id": "root-session", "thread_id": "child-thread"})
);
assert_eq!(
-1
View File
@@ -20,7 +20,6 @@ chrono-tz = { version = "0.10", features = ["serde"] }
event-listener = "5.4.2"
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", "rc"] }
@@ -90,18 +90,6 @@ impl AccountConcurrency {
Self::Limited(limit) => Some(limit),
}
}
/// 扣除当前请求不可占用的预留名额;不限并发不受影响,有限上限至少保留 1 个名额
#[must_use]
pub const fn excluding_reserved(self, reserved: u32) -> Self {
match self {
Self::Unlimited => Self::Unlimited,
Self::Limited(limit) => match NonZeroU32::new(limit.get().saturating_sub(reserved)) {
Some(remaining) => Self::Limited(remaining),
None => Self::Limited(NonZeroU32::MIN),
},
}
}
}
impl From<NonZeroU32> for AccountConcurrency {
@@ -588,18 +576,6 @@ impl AccountStatus {
_ => None,
}
}
/// 管理页稳定业务排序:正常、限流、耗尽、错误、停用
#[must_use]
pub const fn sort_rank(self) -> u8 {
match self {
Self::Normal => 0,
Self::RateLimited => 1,
Self::QuotaExhausted => 2,
Self::Error => 3,
Self::Disabled => 4,
}
}
}
/// `error` 状态下的稳定原因码
@@ -52,12 +52,13 @@ pub const MAX_ACCOUNT_ROTATIONS: u32 = 31;
/// 会话绑定与轮次关联共用的保留时长上限,单位小时
pub const MAX_SESSION_AFFINITY_TTL_HOURS: u32 = 720;
/// OpenAI 会话账号亲和;具体线程身份由 Provider 解释
/// OpenAI 会话账号亲和;请求身份和模式行为由 Provider 解释
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AccountAffinity {
#[default]
Relaxed,
Preferred,
#[default]
Strict,
}
@@ -66,6 +67,7 @@ impl AccountAffinity {
pub const fn as_str(self) -> &'static str {
match self {
Self::Relaxed => "relaxed",
Self::Preferred => "preferred",
Self::Strict => "strict",
}
}
@@ -74,6 +76,7 @@ impl AccountAffinity {
pub fn parse(value: &str) -> Option<Self> {
match value {
"relaxed" => Some(Self::Relaxed),
"preferred" => Some(Self::Preferred),
"strict" => Some(Self::Strict),
_ => None,
}
@@ -107,7 +110,7 @@ impl AccountSelectionPolicy {
max_concurrent_per_account: max_concurrent_per_account.into(),
request_interval,
openai_guardian_reserved_concurrency: 0,
openai_account_affinity: AccountAffinity::Relaxed,
openai_account_affinity: AccountAffinity::default(),
max_account_rotations: 3,
openai_session_affinity_ttl: Duration::from_secs(24 * 3600),
queue_policy: ConcurrencyQueuePolicy {
@@ -203,6 +206,7 @@ impl AccountSelectionPolicy {
/// Store 提供并发事实,Provider 叠加自己解释的额度事实;全部信号均可重建
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRuntimeSignals {
/// 本次调度容量池的在途数,不合并其他独立池
pub in_flight: u32,
pub last_started_at: Option<SystemTime>,
pub quota_reset_at: Option<SystemTime>,
@@ -520,7 +524,7 @@ pub struct AccountSelectionContext {
pub round_robin_cursor: u64,
pub eligibility: AccountEligibilityPolicy,
pub account_scope: Option<std::sync::Arc<crate::account::scope::FrozenAccountScope>>,
/// 本请求不可占用的每账号预留并发名额,由 Provider 按请求类别决定;0 表示不预留
/// 本请求使用的独立预留池上限,由 Provider 按请求类别决定;0 表示使用普通账号池
pub reserved_concurrency: u32,
}
@@ -528,9 +532,11 @@ impl AccountSelectionContext {
/// 本请求在该账号上可使用的并发上限;资格判断、租约与策略投影必须共用这一口径
#[must_use]
pub fn concurrency_limit(&self, account: &ProviderAccount) -> AccountConcurrency {
account
.effective_concurrency(self.policy.max_concurrent_per_account())
.excluding_reserved(self.reserved_concurrency)
if self.reserved_concurrency > 0 {
AccountConcurrency::new(self.reserved_concurrency)
} else {
account.effective_concurrency(self.policy.max_concurrent_per_account())
}
}
}
@@ -720,7 +726,6 @@ impl AccountSelector {
let candidate = match context.policy.strategy() {
RotationStrategy::QuotaResetPriority => {
let default_concurrency = context.policy.max_concurrent_per_account();
eligible.sort_by(|left, right| {
(
left.signals.quota_reset_at.is_none(),
@@ -731,8 +736,8 @@ impl AccountSelector {
right.signals.quota_reset_at,
))
.then_with(|| {
capacity_utilization(left, default_concurrency)
.total_cmp(&capacity_utilization(right, default_concurrency))
capacity_utilization(left, context)
.total_cmp(&capacity_utilization(right, context))
})
.then_with(|| {
left.signals
@@ -748,13 +753,7 @@ impl AccountSelector {
let index = context.round_robin_cursor as usize % eligible.len();
eligible.get(index).copied()?
}
RotationStrategy::Smart => select_smart_candidate(
&eligible,
context.policy.max_concurrent_per_account(),
context.round_robin_cursor,
context.policy.smart_scheduling(),
context.now,
)?,
RotationStrategy::Smart => select_smart_candidate(&eligible, context)?,
RotationStrategy::Sticky => {
eligible.sort_by_key(|candidate| {
(
@@ -840,17 +839,7 @@ impl AccountSelector {
// 账号信号在队列锁外评分,锁内只读取实时队长并查表,避免扫描候选阻塞其他等待者
let scores = candidates
.iter()
.map(|candidate| {
(
candidate.account.id(),
smart_score(
candidate,
context.policy.max_concurrent_per_account(),
config,
context.now,
),
)
})
.map(|candidate| (candidate.account.id(), smart_score(candidate, context)))
.collect::<HashMap<_, _>>();
waiting
.wait_with_priority(keys, |key, count| {
@@ -911,13 +900,9 @@ const SMART_LATENCY_HALF_SCORE_MS: f64 = 10_000.0;
// 距重置一小时时得分减半;未知和已过期时间不提供重置奖励
const SMART_RESET_HALF_SCORE_SECONDS: f64 = 3_600.0;
fn capacity_utilization(
candidate: &AccountCandidate,
default_concurrency: AccountConcurrency,
) -> f64 {
candidate
.account
.effective_concurrency(default_concurrency)
fn capacity_utilization(candidate: &AccountCandidate, context: &AccountSelectionContext) -> f64 {
context
.concurrency_limit(&candidate.account)
.limit()
.map_or(0.0, |limit| {
f64::from(candidate.signals.in_flight) / f64::from(limit.get())
@@ -926,38 +911,29 @@ fn capacity_utilization(
fn select_smart_candidate<'a>(
candidates: &[&'a AccountCandidate],
default_concurrency: AccountConcurrency,
cursor: u64,
config: SmartSchedulingConfig,
now: SystemTime,
context: &AccountSelectionContext,
) -> Option<&'a AccountCandidate> {
let mut ranked = candidates
.iter()
.map(|candidate| {
(
*candidate,
smart_score(candidate, default_concurrency, config, now),
)
})
.map(|candidate| (*candidate, smart_score(candidate, context)))
.collect::<Vec<_>>();
let best_score = ranked
.iter()
.map(|(_, score)| *score)
.max_by(f64::total_cmp)?;
ranked.retain(|(_, score)| best_score - score <= config.score_tolerance());
ranked.retain(|(_, score)| {
best_score - score <= context.policy.smart_scheduling().score_tolerance()
});
// 轮换顺序保持稳定,避免分数轻微交错与 cursor 同步后仍反复命中同一账号
ranked.sort_unstable_by(|(left, _), (right, _)| left.account.id().cmp(right.account.id()));
let index = (cursor % ranked.len() as u64) as usize;
let index = (context.round_robin_cursor % ranked.len() as u64) as usize;
Some(ranked[index].0)
}
pub(crate) fn smart_score(
candidate: &AccountCandidate,
default_concurrency: AccountConcurrency,
config: SmartSchedulingConfig,
now: SystemTime,
) -> f64 {
let load = 1.0 - capacity_utilization(candidate, default_concurrency).clamp(0.0, 1.0);
pub(crate) fn smart_score(candidate: &AccountCandidate, context: &AccountSelectionContext) -> f64 {
let config = context.policy.smart_scheduling();
let now = context.now;
let load = 1.0 - capacity_utilization(candidate, context).clamp(0.0, 1.0);
let quota = candidate
.signals
.quota_remaining_rank
@@ -41,17 +41,14 @@ impl TraceContext {
"blocker": AccountSelector.scheduling_blocker(candidate, context)
.map(|blocker| format!("{blocker:?}")),
"inFlight": signals.in_flight,
"concurrencyLimit": candidate.account
.effective_concurrency(context.policy.max_concurrent_per_account()).get(),
"concurrencyLimit": context.concurrency_limit(&candidate.account).get(),
"quotaRemainingBasisPoints": signals.quota_remaining_rank,
"quotaResetAtUnixMs": signals.quota_reset_at.and_then(|reset|
reset.duration_since(std::time::UNIX_EPOCH).ok()
).map(|duration| duration.as_millis()),
"failureRateBasisPoints": signals.failure_rate_basis_points,
"firstOutputLatencyMs": signals.first_output_latency_ms,
"smartScore": smart.then(|| smart_score(
candidate, context.policy.max_concurrent_per_account(), context.policy.smart_scheduling(), context.now
)),
"smartScore": smart.then(|| smart_score(candidate, context)),
})
})
.collect::<Vec<_>>();
@@ -10,7 +10,7 @@ use futures::future::BoxFuture;
use crate::{
policy::{ClientApiKeyId, PlaintextClientApiKey},
runtime::extensions::{ExtensionSetId, ExtensionSetReference},
routing::extensions::{ExtensionSetId, ExtensionSetReference},
};
pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024;
@@ -18,24 +18,29 @@ pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024;
/// API 从数据面 Authorization 头构造的有界认证信封
#[derive(Clone, PartialEq, Eq)]
pub struct ClientAuthenticationRequest {
settings: Option<crate::settings::RequestSettings>,
settings: Option<crate::routing::request_settings::RequestSettings>,
authorization: Arc<str>,
native_bearer: Option<PlaintextClientApiKey>,
}
impl ClientAuthenticationRequest {
#[must_use]
pub fn with_settings(mut self, settings: crate::settings::RequestSettings) -> Self {
pub fn with_settings(
mut self,
settings: crate::routing::request_settings::RequestSettings,
) -> Self {
self.settings = Some(settings);
self
}
#[must_use]
pub fn settings(&self) -> Option<&crate::settings::RequestSettings> {
pub fn settings(&self) -> Option<&crate::routing::request_settings::RequestSettings> {
self.settings.as_ref()
}
pub(crate) fn take_settings(&mut self) -> Option<crate::settings::RequestSettings> {
pub(crate) fn take_settings(
&mut self,
) -> Option<crate::routing::request_settings::RequestSettings> {
self.settings.take()
}
@@ -0,0 +1,71 @@
//! 请求准入租约及费用结算的释放边界
use super::failure::record_resource_failure;
use crate::{
engine::{
ModelRequestId,
admission::ClientAdmissionPort,
budget::{ClientBudgetCharge, ClientBudgetPort},
},
lifecycle::LeaseGuard,
policy::ClientApiKeyId,
};
use std::sync::Arc;
pub(super) struct AdmissionLease {
pub(super) renewal: Option<Box<dyn LeaseGuard>>,
pub(super) armed: bool,
pub(super) port: Arc<dyn ClientAdmissionPort>,
pub(super) diagnostics: Arc<dyn crate::diagnostics::OperationalDiagnostics>,
pub(super) client_api_key_id: ClientApiKeyId,
pub(super) model_request_id: ModelRequestId,
}
pub(super) async fn settle_budget(
port: &dyn ClientBudgetPort,
charge: ClientBudgetCharge,
diagnostics: &dyn crate::diagnostics::OperationalDiagnostics,
) {
let request_id = charge.request_id.clone();
if let Err(error) = port.settle(charge).await {
record_resource_failure(
diagnostics,
&request_id,
"settle_client_budget",
"Client budget settlement failed; storage will retry on the next request",
error,
)
.await;
}
}
impl AdmissionLease {
pub(super) async fn release(mut self) {
self.renewal.take();
if let Err(error) = self
.port
.release(&self.client_api_key_id, &self.model_request_id)
.await
{
record_resource_failure(
self.diagnostics.as_ref(),
&self.model_request_id,
"release_client_admission",
"Client admission release failed; lease TTL remains active",
error,
)
.await;
}
self.armed = false;
}
}
impl Drop for AdmissionLease {
fn drop(&mut self) {
self.renewal.take();
if self.armed {
self.port
.abandon(&self.client_api_key_id, &self.model_request_id);
}
}
}
@@ -0,0 +1,480 @@
//! 执行请求、认证事实与对外会话端口
use crate::{
engine::{
CoordinatedEvent, EngineError, ModelRequestId,
authentication::ClientAuthenticationRequest,
continuation::PreviousResponseId,
extensions::ExtensionCallScope,
middleware::{
FrozenMiddlewarePlan, MiddlewareAuthority, MiddlewareContext, MiddlewareTarget,
},
nested::ExecutionEffects,
},
error::{GatewayError, GatewayErrorKind},
event::{ProviderEvent, ProviderResponseHeader},
identity::ProviderKind,
lifecycle::{CancellationToken, Deadline},
operation::Operation,
policy::{ClientApiKeyId, ClientPolicy},
routing::{
ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, RuntimeSnapshot,
UpstreamModelId, request_settings::RequestSettings,
},
};
use futures::future::BoxFuture;
use std::{
fmt,
net::IpAddr,
sync::Arc,
time::{Duration, SystemTime},
};
use uuid::Uuid;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientTransport {
HttpJson,
HttpSse,
WebSocket,
InternalProbe,
InternalPlugin,
}
impl ClientTransport {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::HttpJson => "http_json",
Self::HttpSse => "http_sse",
Self::WebSocket => "websocket",
Self::InternalProbe => "internal",
Self::InternalPlugin => "internal_plugin",
}
}
}
/// API 解码后交给 Core 的稳定请求元数据
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecutionRequestMetadata {
pub protocol: String,
pub endpoint: String,
pub transport: ClientTransport,
pub stream: bool,
pub client_ip: Option<IpAddr>,
pub user_agent: Option<String>,
pub previous_response_id: Option<PreviousResponseId>,
}
#[derive(Clone)]
pub struct AuthenticatedClient {
pub(super) settings: Option<RequestSettings>,
pub(super) snapshot: Arc<RuntimeSnapshot>,
pub(super) policy: ClientPolicy,
pub(super) authentication: Option<ClientAuthenticationRequest>,
}
impl AuthenticatedClient {
pub(super) fn execution_timeout(&self) -> Option<Duration> {
self.settings
.as_ref()
.and_then(|settings| settings.execution_timeout(self.policy.key_id()))
}
#[must_use]
pub fn request_settings(&self) -> Option<&RequestSettings> {
self.settings.as_ref()
}
/// 后续仍须重新认证;这里只更新下一次认证使用的冻结配置
#[must_use]
pub fn with_request_settings(mut self, settings: RequestSettings) -> Self {
self.settings = Some(settings);
self
}
#[must_use]
pub const fn snapshot(&self) -> &Arc<RuntimeSnapshot> {
&self.snapshot
}
#[must_use]
pub const fn policy(&self) -> &ClientPolicy {
&self.policy
}
}
impl fmt::Debug for AuthenticatedClient {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("AuthenticatedClient")
.field("key_id", &self.policy.key_id())
.field("revision", &self.snapshot.revision())
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum ClientAuthenticationError {
#[error("client API key is invalid")]
InvalidKey,
#[error("runtime snapshot is unavailable")]
SnapshotUnavailable,
#[error("frontend authentication provider is unavailable")]
ProviderUnavailable,
}
pub struct StartExecution {
pub client: AuthenticatedClient,
pub public_model: PublicModelId,
pub operation: Operation,
pub metadata: ExecutionRequestMetadata,
}
/// 入口中间件开始前冻结的根请求身份与生命周期
///
/// 中间件短路时本值直接释放;调用 `next` 后必须原样交回
/// [`ExecutionService::start_prepared`],不能重新认证为另一个 Key
pub struct PreparedRootExecution {
pub(super) extension_scope: ExtensionCallScope,
pub(super) response_control: crate::engine::response_control::ResponseControl,
pub(super) client: AuthenticatedClient,
pub(super) request_id: ModelRequestId,
pub(super) started_at: SystemTime,
pub(super) deadline_at: Deadline,
pub(super) cancellation: CancellationToken,
pub(super) execution_effects: Arc<ExecutionEffects>,
pub(super) execution_effects_baseline: usize,
}
impl PreparedRootExecution {
#[must_use]
pub fn request_settings(&self) -> RequestSettings {
self.client
.settings
.clone()
.unwrap_or_else(|| RequestSettings::new(self.client.snapshot.clone()))
.with_execution(
&self.client.policy,
self.deadline_at
.at()
.map(|at| duration_ms(at.duration_since(self.started_at).unwrap_or_default())),
)
}
/// 入口延续发起实例集合;保留已经交给调用方的取消句柄
#[must_use]
pub fn with_extension_scope(mut self, scope: ExtensionCallScope) -> Self {
self.extension_scope = scope;
self
}
#[must_use]
pub fn extension_scope(&self) -> ExtensionCallScope {
self.extension_scope.clone()
}
/// 设置已经在改写时完成编译;此处只将同一份结果应用到冻结身份和生命周期
pub fn apply_settings(&mut self, settings: &RequestSettings) -> Result<(), GatewayError> {
let deadline = settings
.execution_deadline(self.client.policy.key_id(), self.started_at)
.map_err(|_| {
GatewayError::new(
GatewayErrorKind::InvalidRequest,
"execution settings are invalid",
)
})?;
self.client.policy = settings.apply_policy(self.client.policy.clone());
self.client.snapshot = settings.snapshot();
self.client.settings = Some(settings.clone());
self.deadline_at = deadline;
Ok(())
}
#[must_use]
pub fn response_control(&self) -> crate::engine::response_control::ResponseControl {
self.response_control.clone()
}
pub(super) fn new(client: AuthenticatedClient) -> Result<Self, GatewayError> {
let started_at = SystemTime::now();
let deadline_at = Deadline::from_timeout(started_at, client.execution_timeout())
.ok_or_else(|| {
GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid")
})?;
let execution_effects = Arc::new(ExecutionEffects::default());
let execution_effects_baseline = execution_effects.epoch();
Ok(Self {
extension_scope: ExtensionCallScope::default(),
client,
response_control: crate::engine::response_control::ResponseControl::default(),
request_id: new_request_id()?,
started_at,
deadline_at,
cancellation: CancellationToken::new(),
execution_effects,
execution_effects_baseline,
})
}
#[must_use]
pub const fn request_id(&self) -> &ModelRequestId {
&self.request_id
}
#[must_use]
pub const fn client(&self) -> &AuthenticatedClient {
&self.client
}
#[must_use]
pub const fn started_at(&self) -> SystemTime {
self.started_at
}
#[must_use]
pub const fn deadline_at(&self) -> Deadline {
self.deadline_at
}
#[must_use]
pub fn cancellation(&self) -> CancellationToken {
self.cancellation.clone()
}
/// 以冻结的入口身份补齐 Runtime-only binding 事实;调用位置事实由 API/Core 提供
#[must_use]
pub fn middleware_context(
&self,
mut target: MiddlewareTarget,
extension_scope: ExtensionCallScope,
) -> MiddlewareContext {
target.request_id = self.request_id.clone();
let account_group_ids = self
.client
.policy
.account_scope()
.routing_snapshot()
.groups_snapshot()
.iter()
.map(|group| group.id().clone())
.collect::<Vec<_>>();
MiddlewareContext::new(
target,
MiddlewareAuthority {
client_key_id: self.client.policy.key_id().clone(),
account_group_ids: Arc::from(account_group_ids),
cancellation: self.cancellation.clone(),
deadline: self.deadline_at,
extension_scope,
execution_effects: Some(Arc::clone(&self.execution_effects)),
},
)
}
}
/// 入口中间件调用 `next` 后交给 Core 的已解码执行输入
pub struct PreparedExecutionRequest {
pub public_model: PublicModelId,
pub operation: Operation,
pub metadata: ExecutionRequestMetadata,
}
/// 启动一个由协议 adapter 明确绑定到 Provider 自有端点的请求
pub struct StartProviderExecution {
pub client: AuthenticatedClient,
pub provider: ProviderKind,
/// Provider 自有模型端点可固定上游模型;非模型端点保持 `None`
pub upstream_model: Option<UpstreamModelId>,
pub operation: Operation,
pub metadata: ExecutionRequestMetadata,
}
pub struct StartedExecution {
pub request_id: ModelRequestId,
pub created_at: SystemTime,
pub stream: bool,
pub session: Box<dyn ExecutionSession>,
}
pub trait ExecutionSession: Send {
fn trace(&self) -> crate::diagnostics::TraceContext {
crate::diagnostics::TraceContext::default()
}
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>>;
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>>;
fn response_headers(&self) -> &[ProviderResponseHeader];
fn response_status_code(&self) -> Option<u16> {
None
}
fn discard_pending_delivery(&mut self) -> Result<(), EngineError> {
Err(EngineError::InvalidDeliveryState)
}
fn commit_downstream(
&mut self,
client_status_code: Option<u16>,
) -> BoxFuture<'_, Result<(), EngineError>>;
fn record_client_status(
&mut self,
client_status_code: u16,
) -> BoxFuture<'_, Result<(), EngineError>>;
/// 执行已终结且结算、准入释放均已返回,协议层才可以放弃清理责任
/// 结算失败仍由 Store 保留费用重试,释放失败仍按租约 TTL 收敛
fn is_finalized(&self) -> bool;
fn cancel(&self);
/// 将会话交给宿主持续驱动清理;取消请求事件的等待不会丢弃已开始的结算
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()>;
}
pub trait ExecutionService: Send + Sync {
/// 未就绪时管理与健康接口仍可运行,模型认证按既有路径报告不可用
fn request_settings(&self) -> Option<RequestSettings> {
None
}
fn authenticate(
&self,
plaintext: &str,
) -> Result<AuthenticatedClient, ClientAuthenticationError>;
fn authenticate_request(
&self,
request: ClientAuthenticationRequest,
) -> BoxFuture<'_, Result<AuthenticatedClient, ClientAuthenticationError>> {
Box::pin(async move {
let plaintext = request
.native_bearer()
.ok_or(ClientAuthenticationError::InvalidKey)?;
self.authenticate(plaintext.expose_for_auth())
})
}
/// 只验证入口认证信封,不记录 Key 使用事实或执行推理准入
fn verify_request(
&self,
_request: ClientAuthenticationRequest,
) -> BoxFuture<'_, Result<AuthenticatedClient, ClientAuthenticationError>> {
Box::pin(async { Err(ClientAuthenticationError::InvalidKey) })
}
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId>;
fn client_model_catalog<'a>(
&'a self,
_client: &'a AuthenticatedClient,
_protocol: &'a str,
_client_version: &'a str,
) -> BoxFuture<'a, Result<Vec<PublicModelDescriptor>, ProviderCatalogUnavailable>> {
Box::pin(async { Err(ProviderCatalogUnavailable) })
}
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool;
/// 重新鉴权并冻结本次根请求身份、请求 ID、deadline 与取消域
fn prepare_execution(
&self,
client: AuthenticatedClient,
) -> BoxFuture<'_, Result<PreparedRootExecution, GatewayError>> {
Box::pin(async move { PreparedRootExecution::new(client) })
}
/// 由已经验证管理会话与插件 models 域的宿主入口选择 Key
/// 不向页面提供明文,也不跳过普通请求的扩展、准入和计量
fn prepare_plugin_execution(
&self,
_client_key_id: &ClientApiKeyId,
) -> BoxFuture<'_, Result<PreparedRootExecution, GatewayError>> {
Box::pin(async {
Err(GatewayError::new(
GatewayErrorKind::Unauthorized,
"plugin model entry is unavailable",
))
})
}
/// 为已经通过 `verify_request` 的只读入口建立中间件生命周期,不重复写 Key 使用事实
fn prepare_verified_execution(
&self,
client: AuthenticatedClient,
) -> Result<PreparedRootExecution, GatewayError> {
PreparedRootExecution::new(client)
}
/// 解析与准备阶段冻结的发布代次一致的中间件计划
fn middleware_plan(&self, _prepared: &PreparedRootExecution) -> Option<FrozenMiddlewarePlan> {
None
}
/// 消费一次已准备身份并进入原有 Core 路由、准入、attempt 与结算路径
fn start_prepared(
&self,
prepared: PreparedRootExecution,
request: PreparedExecutionRequest,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
self.start(StartExecution {
client: prepared.client,
public_model: request.public_model,
operation: request.operation,
metadata: request.metadata,
})
.await
})
}
/// 消费一次已准备身份,并进入 Provider 自有端点的原有准入与结算路径
fn start_prepared_provider_endpoint(
&self,
prepared: PreparedRootExecution,
provider: ProviderKind,
upstream_model: Option<UpstreamModelId>,
operation: Operation,
metadata: ExecutionRequestMetadata,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
self.start_provider_endpoint(StartProviderExecution {
client: prepared.client,
provider,
upstream_model,
operation,
metadata,
})
.await
})
}
fn start(
&self,
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>>;
fn start_provider_endpoint(
&self,
request: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>>;
/// 返回 Provider 注册表中首个声明 Live sideband 能力的网关。
/// `None` 表示当前组合不含语音 sideband;协议 adapter 据此回退到稳定 501。
fn live_gateway(&self) -> Option<Arc<dyn crate::live::LiveGateway>> {
None
}
}
/// 只验证 Client API Key 并返回稳定 Key ID,不产生请求使用事实
///
/// 管理侧登录等非数据面场景必须使用该端口,避免把认证本身误记为一次 Key 使用
pub trait ClientKeyVerifier: Send + Sync {
fn verify_client_key(
&self,
plaintext: &str,
) -> Result<ClientApiKeyId, ClientAuthenticationError>;
}
/// 成功认证后的 API Key 使用事实接收器
///
/// 认证仍是同步快照读取;实现必须自行异步、去重地持久化,不得阻塞客户端请求
pub trait ClientApiKeyUsageSink: Send + Sync {
fn record_used(&self, key_id: &ClientApiKeyId);
}
pub(super) fn duration_ms(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
pub(super) fn new_request_id() -> Result<ModelRequestId, GatewayError> {
ModelRequestId::new(format!("req_{}", Uuid::now_v7().simple()))
.map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID"))
}
@@ -0,0 +1,100 @@
//! 执行边界的错误分类与资源释放诊断
use super::contract::ClientAuthenticationError;
use crate::{
engine::{EngineError, ModelRequestId},
error::{GatewayError, GatewayErrorKind},
};
pub(super) fn authentication_gateway_error(error: ClientAuthenticationError) -> GatewayError {
match error {
ClientAuthenticationError::InvalidKey => {
GatewayError::new(GatewayErrorKind::Unauthorized, "client API key is invalid")
}
ClientAuthenticationError::SnapshotUnavailable => GatewayError::new(
GatewayErrorKind::Internal,
"runtime snapshot is unavailable",
),
ClientAuthenticationError::ProviderUnavailable => GatewayError::new(
GatewayErrorKind::Internal,
"frontend authentication provider is unavailable",
),
}
}
pub(super) fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError {
match error {
crate::validation::RoutingError::ModelNotFound {
model,
mapped_model,
} => GatewayError::new(
GatewayErrorKind::ModelNotFound,
if model == mapped_model {
"the requested model was not found in the provider catalogs available to this API key; check the model name"
} else {
"the requested model maps to an upstream model that was not found in the provider catalogs available to this API key; check the configured model mapping"
},
),
crate::validation::RoutingError::NoCapableProvider { .. }
| crate::validation::RoutingError::NoCapableProviderEndpoint { .. }
| crate::validation::RoutingError::EmptyAccountScope => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider can execute this request",
),
crate::validation::RoutingError::UnsupportedProviderEndpoint { .. } => GatewayError::new(
GatewayErrorKind::Unsupported,
"the selected provider does not support this operation",
),
_ => GatewayError::new(
GatewayErrorKind::Internal,
"runtime routing configuration is invalid",
),
}
}
pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
match error {
EngineError::Cancelled => {
GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled")
}
EngineError::Deadline => {
GatewayError::new(GatewayErrorKind::Timeout, "request deadline elapsed")
}
EngineError::Provider(provider) => GatewayError::from_provider(provider),
EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => {
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider is available",
)
}
EngineError::Store(source) => {
GatewayError::new(GatewayErrorKind::Internal, "request execution failed")
.with_source(source.clone())
}
_ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"),
}
}
pub(super) async fn record_resource_failure(
diagnostics: &dyn crate::diagnostics::OperationalDiagnostics,
request_id: &ModelRequestId,
operation: &'static str,
message: &'static str,
source: impl Into<crate::error::ErrorSource>,
) {
let mut failure = crate::diagnostics::OperationalFailure::new(
"core",
operation,
"resource_unavailable",
message,
);
failure.correlation_id = Some(request_id.as_str().to_owned());
failure.details = crate::error::ErrorDetails::capture(Some(&source.into()), None, false);
if diagnostics.record_failure(failure).await.is_err() {
tracing::warn!(
request_id = request_id.as_str(),
operation,
"resource diagnostic could not be recorded"
);
}
}
@@ -0,0 +1,24 @@
//! 数据面执行入口与请求、作用域、会话的单一所有权
mod admission;
mod contract;
mod failure;
pub use contract::{
AuthenticatedClient, ClientApiKeyUsageSink, ClientAuthenticationError, ClientKeyVerifier,
ClientTransport, ExecutionRequestMetadata, ExecutionService, ExecutionSession,
PreparedExecutionRequest, PreparedRootExecution, StartExecution, StartProviderExecution,
StartedExecution,
};
mod scope;
pub use scope::BoundModelExecutionContext;
mod service;
pub use failure::gateway_error_from_engine;
pub use service::DefaultExecutionService;
mod probe_store;
mod session;
use std::time::Duration;
const DIAGNOSTIC_TIMEOUT: Duration = Duration::from_secs(10 * 60);
const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100);
const MAX_NESTED_EXECUTIONS: usize = 16;
const MAX_CONCURRENT_NESTED_EXECUTIONS: usize = 4;
@@ -0,0 +1,64 @@
//! 诊断探测使用的非持久执行记录与观测事实
use crate::{
engine::{
AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
ModelRequestId, NewModelRequest, ProviderAccountId, RecoveryReport, UpstreamSendState,
},
error::StoreError,
identity::ProviderKind,
routing::UpstreamModelId,
};
use std::time::SystemTime;
pub(super) struct ProbeObservation {
pub(super) provider_kind: ProviderKind,
pub(super) account_id: ProviderAccountId,
pub(super) upstream_model: UpstreamModelId,
}
pub(super) struct TransientExecutionStore;
#[async_trait::async_trait]
impl ExecutionStore for TransientExecutionStore {
async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> {
Ok(())
}
async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> {
Ok(())
}
async fn mark_send_state(
&self,
_: &ModelRequestId,
_: UpstreamSendState,
) -> Result<(), StoreError> {
Ok(())
}
async fn mark_downstream_committed(
&self,
_: &ModelRequestId,
_: SystemTime,
_: Option<u16>,
) -> Result<(), StoreError> {
Ok(())
}
async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> {
Ok(())
}
async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> {
Ok(())
}
async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> {
Ok(())
}
async fn recover_expired(&self, _: SystemTime) -> Result<RecoveryReport, StoreError> {
Ok(RecoveryReport::default())
}
}
@@ -0,0 +1,301 @@
//! 请求嵌套图、活动权限与作用域 guard
use super::admission::{AdmissionLease, settle_budget};
use super::contract::{AuthenticatedClient, ExecutionRequestMetadata};
use super::{MAX_CONCURRENT_NESTED_EXECUTIONS, MAX_NESTED_EXECUTIONS};
use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort};
use crate::{
concurrency::ConcurrencyWaitBudget,
engine::{
ModelRequestId, ProviderAccountId, coordinator::CoordinationExtensions,
extensions::ExtensionCallScope, nested::ExecutionEffects,
},
error::{GatewayError, GatewayErrorKind},
identity::ProviderKind,
lifecycle::{CancellationToken, Deadline},
operation::Operation,
policy::ClientApiKeyId,
routing::{
FrozenAccountScope, PublicModelId, UpstreamModelId, request_settings::RequestSettings,
},
};
use std::{
collections::BTreeMap,
fmt,
sync::{
Arc, Mutex, Weak,
atomic::{AtomicUsize, Ordering},
},
time::SystemTime,
};
pub(super) enum ExecutionTarget {
Model(PublicModelId),
ProviderEndpoint {
provider: ProviderKind,
upstream_model: Option<UpstreamModelId>,
},
}
impl ExecutionTarget {
pub(super) fn public_model(&self) -> Option<&PublicModelId> {
match self {
Self::Model(model) => Some(model),
Self::ProviderEndpoint { .. } => None,
}
}
pub(super) fn into_public_model(self) -> Option<PublicModelId> {
match self {
Self::Model(model) => Some(model),
Self::ProviderEndpoint { .. } => None,
}
}
}
pub(super) struct PendingStartExecution {
pub(super) client: AuthenticatedClient,
pub(super) target: ExecutionTarget,
pub(super) operation: Operation,
pub(super) metadata: ExecutionRequestMetadata,
}
pub(super) struct AuthorizedExecution {
pub(super) response_control: Option<crate::engine::response_control::ResponseControl>,
pub(super) account_scope: Arc<FrozenAccountScope>,
pub(super) deadline_at: Deadline,
pub(super) cancellation: CancellationToken,
pub(super) extension_scope: ExtensionCallScope,
pub(super) required_provider: Option<ProviderKind>,
pub(super) required_account: Option<ProviderAccountId>,
pub(super) nested: Option<NestedExecutionFacts>,
pub(super) bound: Option<BoundExecutionFacts>,
pub(super) graph: Option<Arc<NestedExecutionGraph>>,
pub(super) execution_effects_baseline: usize,
pub(super) nested_permit: Option<NestedExecutionPermit>,
}
pub(super) struct ExecutionStartGuard {
pub(super) admission: ExecutionAdmission,
pub(super) active_request: Option<ActiveRequestLease>,
pub(super) concurrency_wait_budget: ConcurrencyWaitBudget,
pub(super) admission_decision_ms: Option<u64>,
}
pub(super) struct PreparedExecutionStart {
pub(super) request_id: ModelRequestId,
pub(super) started_at: SystemTime,
pub(super) plan: crate::routing::RoutingPlan,
pub(super) extensions: CoordinationExtensions,
pub(super) authorization: AuthorizedExecution,
pub(super) guard: ExecutionStartGuard,
}
pub(super) struct NestedExecutionFacts {
pub(super) parent_request_id: ModelRequestId,
pub(super) initiating_plugin_instance_id: String,
}
pub(super) struct BoundExecutionFacts {
pub(super) initiating_plugin_instance_id: String,
}
pub(super) struct NestedExecutionGraph {
pub(super) total_started: AtomicUsize,
pub(super) active: AtomicUsize,
pub(super) effects: Arc<ExecutionEffects>,
}
impl NestedExecutionGraph {
pub(super) fn new(effects: Arc<ExecutionEffects>) -> Self {
Self {
total_started: AtomicUsize::new(0),
active: AtomicUsize::new(0),
effects,
}
}
pub(super) fn acquire(self: &Arc<Self>) -> Result<NestedExecutionPermit, GatewayError> {
if self
.active
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |active| {
(active < MAX_CONCURRENT_NESTED_EXECUTIONS).then_some(active + 1)
})
.is_err()
{
return Err(GatewayError::new(
GatewayErrorKind::ConcurrencyQueueFull,
"nested model execution concurrency is exhausted",
));
}
if self
.total_started
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |total| {
(total < MAX_NESTED_EXECUTIONS).then_some(total + 1)
})
.is_err()
{
self.active.fetch_sub(1, Ordering::AcqRel);
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"nested model execution budget is exhausted",
));
}
Ok(NestedExecutionPermit {
graph: Arc::clone(self),
armed: true,
})
}
}
pub(super) struct NestedExecutionPermit {
pub(super) graph: Arc<NestedExecutionGraph>,
pub(super) armed: bool,
}
impl NestedExecutionPermit {
pub(super) fn release(mut self) {
if self.armed {
self.graph.active.fetch_sub(1, Ordering::AcqRel);
self.armed = false;
}
}
}
impl Drop for NestedExecutionPermit {
fn drop(&mut self) {
if self.armed {
self.graph.active.fetch_sub(1, Ordering::AcqRel);
}
}
}
pub(super) struct ActiveRequestAuthority {
pub(super) client: AuthenticatedClient,
pub(super) account_scope: Arc<FrozenAccountScope>,
pub(super) deadline_at: Deadline,
pub(super) cancellation: CancellationToken,
pub(super) extension_scope: ExtensionCallScope,
pub(super) graph: Arc<NestedExecutionGraph>,
}
pub(super) type ActiveRequestRegistry =
Arc<Mutex<BTreeMap<ModelRequestId, Weak<ActiveRequestAuthority>>>>;
pub(super) struct ActiveRequestLease {
pub(super) registry: ActiveRequestRegistry,
pub(super) request_id: ModelRequestId,
pub(super) authority: Arc<ActiveRequestAuthority>,
}
/// 管理页或 CLI 一次调用持有的显式模型执行身份
///
/// 结构不暴露 Key 明文;Runtime 只能把它原样交回 Core 发起模型请求
#[derive(Clone)]
pub struct BoundModelExecutionContext {
pub(super) authority: Arc<BoundModelExecutionAuthority>,
}
impl BoundModelExecutionContext {
#[must_use]
pub fn request_settings(&self) -> Option<&RequestSettings> {
self.authority.client.request_settings()
}
}
pub(super) struct BoundModelExecutionAuthority {
pub(super) client: AuthenticatedClient,
pub(super) account_scope: Arc<FrozenAccountScope>,
pub(super) cancellation: CancellationToken,
pub(super) extension_scope: ExtensionCallScope,
pub(super) initiating_plugin_instance_id: String,
}
impl fmt::Debug for BoundModelExecutionContext {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("BoundModelExecutionContext")
.field("key_id", &self.authority.client.policy.key_id())
.field(
"initiating_plugin_instance_id",
&self.authority.initiating_plugin_instance_id,
)
.finish_non_exhaustive()
}
}
impl Drop for BoundModelExecutionAuthority {
fn drop(&mut self) {
self.cancellation.cancel();
}
}
impl ActiveRequestLease {
pub(super) fn release(self) {
drop(self);
}
}
impl Drop for ActiveRequestLease {
fn drop(&mut self) {
self.authority.cancellation.cancel();
let mut active = self
.registry
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if active
.get(&self.request_id)
.and_then(Weak::upgrade)
.is_some_and(|authority| Arc::ptr_eq(&authority, &self.authority))
{
active.remove(&self.request_id);
}
}
}
impl ExecutionStartGuard {
pub(super) async fn release_failed(
self,
budget: Option<&dyn ClientBudgetPort>,
request_id: ModelRequestId,
key_id: ClientApiKeyId,
diagnostics: &dyn crate::diagnostics::OperationalDiagnostics,
) {
if let Some(active_request) = self.active_request {
active_request.release();
}
if let Some(budget) = budget {
settle_budget(
budget,
ClientBudgetCharge {
key_id,
request_id,
amount_usd: crate::metering::Decimal::ZERO,
completed_at: SystemTime::now(),
},
diagnostics,
)
.await;
}
self.admission.release().await;
}
}
pub(super) enum ExecutionAdmission {
Client(AdmissionLease),
Nested(NestedExecutionPermit),
Bound(AdmissionLease, NestedExecutionPermit),
}
impl ExecutionAdmission {
pub(super) async fn release(self) {
match self {
Self::Client(admission) => admission.release().await,
Self::Nested(permit) => permit.release(),
Self::Bound(admission, permit) => {
admission.release().await;
permit.release();
}
}
}
}
@@ -0,0 +1,213 @@
//! 执行会话的准入释放、费用结算与终态收敛
use super::COORDINATION_TIMEOUT;
use super::admission::settle_budget;
use super::contract::ExecutionSession;
use super::scope::{ActiveRequestLease, ExecutionAdmission};
use crate::{
engine::{
CoordinatedEvent, EngineError, ExecutionStore,
budget::ClientBudgetPort,
continuation::{NativeContinuationPin, NativeContinuationPort},
coordinator::ResponseExecutionSession,
},
event::{ProviderEvent, ProviderResponseHeader},
operation::ProviderSessionState,
};
use futures::{FutureExt as _, future::BoxFuture, pin_mut, select_biased};
use futures_timer::Delay;
use std::sync::Arc;
pub(super) struct DefaultExecutionSession {
pub(super) core: ResponseExecutionSession<dyn ExecutionStore>,
pub(super) admission: Option<ExecutionAdmission>,
pub(super) active_request: Option<ActiveRequestLease>,
pub(super) cleanup: Option<BoxFuture<'static, ()>>,
pub(super) continuation: Arc<dyn NativeContinuationPort>,
pub(super) continuation_recorded: bool,
pub(super) budget: Option<Arc<dyn ClientBudgetPort>>,
pub(super) diagnostics: Arc<dyn crate::diagnostics::OperationalDiagnostics>,
}
impl DefaultExecutionSession {
pub(super) fn new(
core: ResponseExecutionSession<dyn ExecutionStore>,
admission: ExecutionAdmission,
active_request: Option<ActiveRequestLease>,
continuation: Arc<dyn NativeContinuationPort>,
budget: Option<Arc<dyn ClientBudgetPort>>,
diagnostics: Arc<dyn crate::diagnostics::OperationalDiagnostics>,
) -> Self {
Self {
core,
admission: Some(admission),
active_request,
cleanup: None,
continuation,
continuation_recorded: false,
budget,
diagnostics,
}
}
pub(super) async fn settle_if_finalized(&mut self) {
if self.core.is_finalized()
&& let Some(admission) = self.admission.take()
{
if let Some(active_request) = self.active_request.take() {
active_request.release();
}
let budget = self.budget.take();
let charge = self.core.budget_charge();
let diagnostics = self.diagnostics.clone();
// 在首次 await 前把完整清理责任留在会话内
// 事件等待被取消后,后续 poll
// 或 detach 继续同一个 future,既不丢失费用,也不重启已完成的结算
self.cleanup = Some(Box::pin(async move {
if let Some(budget) = budget {
settle_budget(budget.as_ref(), charge, diagnostics.as_ref()).await;
}
admission.release().await;
}));
}
if let Some(cleanup) = self.cleanup.as_mut() {
cleanup.await;
self.cleanup = None;
}
}
pub(super) async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) {
if self.continuation_recorded {
return;
}
let Some(state) = state else {
return;
};
let Some(pin) = self.core.native_continuation_pin(state) else {
return;
};
self.continuation_recorded = true;
record_native_continuation(self.continuation.as_ref(), pin).await;
}
pub(super) async fn finalize_detached(&mut self) {
if let Err(error) = self.core.cancel_and_finalize().await {
tracing::warn!(request_id = self.core.request_id().as_str(), operation = "finalize_detached_execution", %error, "Detached execution 终态收敛失败");
}
self.settle_if_finalized().await;
}
}
impl Drop for DefaultExecutionSession {
fn drop(&mut self) {
self.core.cancel();
drop(self.active_request.take());
}
}
impl ExecutionSession for DefaultExecutionSession {
fn trace(&self) -> crate::diagnostics::TraceContext {
self.core.trace()
}
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
let result = self.core.next_event().await;
if let Ok(Some(event)) = result.as_ref() {
self.record_continuation(event.session_update()).await;
}
self.settle_if_finalized().await;
result
})
}
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
Box::pin(async move {
let result = self.core.collect_uncommitted().await;
if let Ok(events) = result.as_ref() {
let state = events.iter().find_map(ProviderEvent::session_update);
self.record_continuation(state).await;
}
self.settle_if_finalized().await;
result
})
}
fn response_headers(&self) -> &[ProviderResponseHeader] {
self.core.response_headers()
}
fn response_status_code(&self) -> Option<u16> {
self.core.response_status_code()
}
fn discard_pending_delivery(&mut self) -> Result<(), EngineError> {
self.core.discard_pending_delivery()
}
fn commit_downstream(
&mut self,
client_status_code: Option<u16>,
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
let result = self.core.commit_downstream(client_status_code).await;
self.settle_if_finalized().await;
result
})
}
fn record_client_status(
&mut self,
client_status_code: u16,
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
let result = self.core.record_client_status(client_status_code).await;
self.settle_if_finalized().await;
result
})
}
fn is_finalized(&self) -> bool {
self.core.is_finalized()
&& self.admission.is_none()
&& self.active_request.is_none()
&& self.cleanup.is_none()
}
fn cancel(&self) {
self.core.cancel();
}
fn detach_finalize(mut self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move { self.finalize_detached().await })
}
}
pub(super) async fn record_native_continuation(
continuation: &dyn NativeContinuationPort,
pin: NativeContinuationPin,
) {
let provider = pin.provider().as_str().to_owned();
let account = pin.account().as_str().to_owned();
let record = continuation.record(pin).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(record, timeout);
select_biased! {
result = record => {
if let Err(error) = result {
tracing::warn!(
provider = %provider,
account = %account,
%error,
"Continuation affinity 写入失败,后续请求将退化为外部续接"
);
}
},
_ = timeout => {
tracing::warn!(
provider = %provider,
account = %account,
"Continuation affinity 后台写入超时,已丢弃本次亲和记录"
);
},
}
}
@@ -1,4 +1,4 @@
//! 插件嵌套调用的递归作用域
//! 执行扩展计划的弱引用索引、冻结解析与插件递归作用域
use std::{collections::BTreeSet, sync::Arc};
@@ -43,3 +43,82 @@ impl ExtensionCallScope {
})
}
}
/// 同一代次的具体执行计划;不拥有发布状态,运行资源由对应发布集合保活
#[derive(Debug, Default)]
pub struct ExecutionExtensionPlans {
pub middleware: Option<Arc<dyn super::middleware::MiddlewarePlan>>,
pub upstream_adapters: Option<Arc<dyn super::upstream_adapter::UpstreamAdapterPlan>>,
}
/// 同一发布代次不能被静默替换
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("execution extension generation is already registered")]
pub struct ExtensionRegistrationError;
/// 仅以弱引用按冻结身份解析执行计划,不延长旧代次寿命或决定何时发布
#[derive(Clone, Default)]
pub struct ExecutionExtensionIndex {
sets: Arc<
std::sync::RwLock<
std::collections::BTreeMap<
crate::routing::extensions::ExtensionSetId,
std::sync::Weak<ExecutionExtensionPlans>,
>,
>,
>,
}
impl ExecutionExtensionIndex {
pub fn register(
&self,
id: crate::routing::extensions::ExtensionSetId,
plans: Arc<ExecutionExtensionPlans>,
) -> Result<Arc<ExecutionExtensionPlans>, ExtensionRegistrationError> {
let mut sets = self
.sets
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
sets.retain(|_, plans| plans.strong_count() > 0);
if sets.contains_key(&id) {
return Err(ExtensionRegistrationError);
}
sets.insert(id, Arc::downgrade(&plans));
Ok(plans)
}
fn resolve(
&self,
generation: &crate::routing::extensions::ExtensionSetReference,
) -> Option<Arc<ExecutionExtensionPlans>> {
self.sets
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(generation.id())
.and_then(std::sync::Weak::upgrade)
}
#[must_use]
pub fn middleware(
&self,
generation: &crate::routing::extensions::ExtensionSetReference,
) -> Option<super::middleware::FrozenMiddlewarePlan> {
let plan = self.resolve(generation)?.middleware.clone()?;
Some(super::middleware::FrozenMiddlewarePlan::new(
plan,
generation.clone(),
))
}
#[must_use]
pub fn upstream_adapters(
&self,
generation: &crate::routing::extensions::ExtensionSetReference,
) -> Option<super::upstream_adapter::FrozenUpstreamAdapterPlan> {
let plan = self.resolve(generation)?.upstream_adapters.clone()?;
Some(super::upstream_adapter::FrozenUpstreamAdapterPlan::new(
plan,
generation.clone(),
))
}
}
@@ -5,10 +5,11 @@
//! Runtime/SDK 负责;身份、路由、账号租约、发送事实、重试、计量与结算仍由 Core
//! 持有,不能通过本端口交给插件解释
use std::collections::{BTreeMap, BTreeSet};
use super::FrozenMiddlewarePlan;
use std::collections::BTreeSet;
use std::fmt;
use std::num::NonZeroU32;
use std::sync::{Arc, RwLock, Weak};
use std::sync::Arc;
use bytes::Bytes;
use futures::future::BoxFuture;
@@ -27,8 +28,7 @@ use crate::operation::{
CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload,
};
use crate::policy::ClientApiKeyId;
use crate::runtime::extensions::{ExtensionSetId, ExtensionSetReference};
use crate::settings::RequestSettings;
use crate::routing::request_settings::RequestSettings;
/// 中间件在一次逻辑请求中的挂载位置
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
@@ -543,7 +543,7 @@ impl fmt::Debug for MiddlewareResponseEnvelope {
/// Core 冻结的调用事实;完整业务身份可由插件读取
#[derive(Clone)]
pub struct MiddlewareContext {
plan: Option<FrozenMiddlewarePlan>,
pub(super) plan: Option<FrozenMiddlewarePlan>,
request_id: ModelRequestId,
mount: MiddlewareMount,
attempt_index: Option<NonZeroU32>,
@@ -799,175 +799,3 @@ impl From<ProviderError> for 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: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>>;
}
/// 调用期间同时保活发布代次与中间件计划
#[derive(Clone)]
pub struct FrozenMiddlewarePlan {
plan: Arc<dyn MiddlewarePlan>,
_generation: ExtensionSetReference,
}
impl FrozenMiddlewarePlan {
#[must_use]
pub fn has_service(&self) -> bool {
self.plan.has_service()
}
pub fn handle_service(
&self,
context: crate::middleware::service::Context,
input: crate::middleware::service::Value,
next: crate::middleware::service::Next,
) -> BoxFuture<
'static,
Result<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 {
plan,
_generation: generation,
}
}
pub fn handle(
&self,
mut context: MiddlewareContext,
request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
context.plan = Some(self.clone());
self.plan.handle(context, request, next)
}
}
impl fmt::Debug for FrozenMiddlewarePlan {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("FrozenMiddlewarePlan")
.finish_non_exhaustive()
}
}
/// 同一发布代次不能被静默替换
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("middleware generation is already registered")]
pub struct MiddlewareRegistrationError;
/// 按请求冻结的发布集合解析中间件计划;索引不延长旧代次寿命
#[derive(Clone, Default)]
pub struct MiddlewareExtensionIndex {
sets: Arc<RwLock<BTreeMap<ExtensionSetId, Weak<dyn MiddlewarePlan>>>>,
}
impl MiddlewareExtensionIndex {
pub fn register(
&self,
id: ExtensionSetId,
plan: Arc<dyn MiddlewarePlan>,
) -> Result<Arc<dyn MiddlewarePlan>, MiddlewareRegistrationError> {
let mut sets = self
.sets
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
sets.retain(|_, plan| plan.strong_count() > 0);
if sets.contains_key(&id) {
return Err(MiddlewareRegistrationError);
}
sets.insert(id, Arc::downgrade(&plan));
Ok(plan)
}
#[must_use]
pub fn resolve(&self, generation: &ExtensionSetReference) -> Option<FrozenMiddlewarePlan> {
self.sets
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(generation.id())
.and_then(Weak::upgrade)
.map(|plan| FrozenMiddlewarePlan::new(plan, generation.clone()))
}
}

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