Compare commits

...
20 Commits
Author SHA1 Message Date
ZyyandCodex 78bc57c5cf chore(release): v3.29.7
Co-authored-by: Codex <noreply@openai.com>
2026-10-21 11:23:31 +08:00
ef11b72d8a fix(workers): recover daemon health after successful work (#499)
Report per-run daemon progress to reset health failures and restart backoff
after successful work, with stale-run isolation and recovery regression tests.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-21 11:08:15 +08:00
ZyyandCodex 52fb1e7180 fix(deps): update vulnerable redirect and TLS dependencies
Update follow-redirects, rustls-webpki, and webpki-roots to patched releases so frontend and backend security audits pass.

Co-authored-by: Codex <noreply@openai.com>
2026-10-21 10:47:57 +08:00
ZyyandCodex 7a7a7ed6ad fix(openai): avoid full MSIX scans during version checks
Use the known MSIX archive offset to keep Windows Desktop version checks within the range download budget. Cover both architectures with a bounded-read regression test.

Co-authored-by: Codex <noreply@openai.com>
2026-10-21 10:33:18 +08:00
b8479b59db fix(plugins): preserve request origin and raw upstream observations (#497)
* fix(plugins): preserve request origin and raw upstream observations

Co-authored-by: Codex <noreply@openai.com>

* perf(plugins): share immutable response observations

Avoid copying response metadata and header collections for every stream event. Update Workbench byte headers and virtual-time timeout coverage, and align the large-request fixture with RPC v5.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-21 01:09:58 +08:00
2480f07f2d fix(plugins): honor configured model and HTTP request body limits (#496)
* fix(plugins): honor configured request body limits

* test(plugin-sdk): mirror upstream adapter module layout

Co-authored-by: Codex <noreply@openai.com>

* fix(plugins): preserve frozen request body budgets across callbacks

Carry the execution body budget through policy and attempt callbacks, and validate explicit-key model requests against their bound snapshot. Add regression coverage for configuration publication and document budget ownership and lifetime.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
2026-10-21 00:51:13 +08:00
niceSprite b2dc9f493d fix(plugins): retain request middleware resource streams (#492) 2026-10-20 23:09:06 +08:00
d92d310e45 fix(accounts): select the bound outbound proxy by default (#493)
Expose the stored proxy ID through account queries and use it to display
the bound proxy name without matching redacted endpoints. Preserve
unchanged bindings and keep the preserve option for batch editing.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-20 23:02:51 +08:00
ZyyandCodex 3caf83131a docs: refine GitHub reply tone
Co-authored-by: Codex <noreply@openai.com>
2026-10-20 20:48:35 +08:00
ZyyandCodex af35cde374 chore(release): v3.29.6
Co-authored-by: Codex <noreply@openai.com>
2026-10-20 18:44:26 +08:00
5c115fb38f fix(core): reserve rotation budget during connection recovery (#487)
* fix(retry): rotate accounts after connection failures

Reliable NotSent HTTP connection failures stay on the same account and
egress until the shared 30-second opening window is exhausted. When an
account's proxy does not respond, each opening waits for the full
15-second connect timeout on the same route, so requests that selected
that account fail after about 30 seconds while other accounts stay idle.

After a connection failure, exclude the account and reselect through
the existing rotation path, consuming routing and rotation budget and
sharing the same opening window. Rotation happens after the first
failure because a second same-account opening on an unresponsive
egress would leave too little of the window to rotate. If no other
account is eligible, restore the exclusion and continue the original
same-account HTTP recovery, so single-account deployments keep their
current behavior. Required accounts and continuation attempts are
unchanged.

Refs #485

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(retry): retain same-account connection recovery

Keep existing account rotation rules and cover rotation after a recovered
connection receives a replay-safe rate-limit rejection.

Co-authored-by: Codex <noreply@openai.com>

* fix(core): reserve connection budget for account rotation

Retry the original account before rotating on unsent connection failures.
Split the retry window and retain an opening for HTTP recovery on another
account, with bounded fallback when no alternative can be selected.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-20 17:42:42 +08:00
ed558e5396 fix(plugins): decouple upstream streams from initial deadlines (#489)
Transfer upstream stream lifetime to the model request after the initial RPC reply and give managed HTTP body reads an operation budget. Cover heartbeat continuity, cancellation cleanup, and per-read timeouts.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-20 17:17:28 +08:00
ZyyandCodex 929f781a63 perf(store): slim business backups and accelerate history cleanup
Export business data and published hourly usage from a shared snapshot as
bounded gzip SQL archives, excluding request history and diagnostic data.
Run bounded retention cycles every minute, deleting oldest unlocked rows
with per-batch lock and statement deadlines.

Add real PostgreSQL export/restore and lock-contention coverage, document
the backup scope, and move detailed UI guidance into a help popover.

Validation: 4011 backend tests passed in partitions; Rustfmt, Clippy,
frontend lint/type checks, and actionlint passed. Resource-limited load
checks covered current-adapter million-row backups and a read-only policy
check on the existing 31-million-row database. Final popover interaction
checks were stopped at the user's request. Online restore is deferred.

Co-authored-by: Codex <noreply@openai.com>
2026-10-20 14:42:29 +08:00
e0326ecba4 fix(openai): base pooled WS breaker on opening results (#484)
* fix(openai): base pooled WS breaker on opening results

Keep foreground fallback independent of pooled opening health and fence late results by circuit generation.

Refs #482

Co-authored-by: Codex <noreply@openai.com>

* fix(openai): preserve concurrent breaker results while closed

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
2026-10-20 14:06:23 +08:00
ZyyandCodex 6057088300 chore(release): v3.29.5
Co-authored-by: Codex <noreply@openai.com>
2026-10-19 17:42:17 +08:00
e3a49a6217 fix(openai): supplement declarations for hosted search history (#478)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-19 17:07:27 +08:00
b32f87fa8a fix(store): bound extrema rebuilds to their analytics bucket (#477)
Group extrema rebuilds by collection and analytics bucket, with explicit started_at bounds so large fact histories can use the time-range index. Preserve per-group build prefixes and aggregate results.

Keep the scan regression test on a single PostgreSQL connection so forced statistics flushes include the worker's reads. The updated test passes with the bounded query and rejects the previous query after reading 6,000,040 rows.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-19 16:50:01 +08:00
ZyyandCodex 91e8105e5d fix(openai): measure first-token latency from first response data
Record first-token timing from nonempty HTTP SSE data or WebSocket messages before JSON and event classification. Preserve semantic output, reasoning and text timing boundaries.

Co-authored-by: Codex <noreply@openai.com>
2026-10-19 13:59:05 +08:00
dependabot[bot]andCodex c4d9576639 chore(deps): update GitHub Actions and composite action pins (#472)
Align pinned action versions across CI, release workflows and composite actions.

Co-authored-by: Codex <noreply@openai.com>
2026-10-19 13:52:12 +08:00
ZyyandCodex 015cd60be1 fix(ci): allow scoped Dependabot updates through rules guard
Co-authored-by: Codex <noreply@openai.com>
2026-10-19 13:46:36 +08:00
175 changed files with 5310 additions and 865 deletions
+2 -2
View File
@@ -18,7 +18,7 @@ runs:
run: ruby .github/scripts/toolchains.rb >> "$GITHUB_OUTPUT"
- name: Setup Node
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
uses: actions/setup-node@949feb2413d6458794dcd2491c4babbbce0c15c1 # v7.1.0
with:
node-version: ${{ steps.toolchains.outputs.node }}
cache: pnpm
@@ -38,7 +38,7 @@ runs:
- name: Upload frontend artifact
if: inputs.upload-artifact == 'true'
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: frontend-dist
path: frontend/dist
+1 -1
View File
@@ -242,7 +242,7 @@ runs:
- name: Upload container diagnostics
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: ${{ inputs.artifact-name }}
path: ${{ runner.temp }}/codex-proxy-rs-diagnostics
+15 -1
View File
@@ -122,7 +122,7 @@ jobs:
run: cargo clippy --all-targets --all-features --locked -- -D warnings
- name: Setup nextest
uses: taiki-e/install-action@4cef1412cce204788f482e778a0b9187f9626a29 # v2.87.21
uses: taiki-e/install-action@8cca0918f8471c33515202b73f56ab27f2bed283 # v2.87.29
with:
tool: cargo-nextest@0.9.143
fallback: none
@@ -156,6 +156,20 @@ jobs:
echo "CPR_TEST_REDIS_URL=redis://:${redis_password}@127.0.0.1:6379"
} >> "$GITHUB_ENV"
- name: Expose matching PostgreSQL backup and restore tools
shell: bash
env:
CPR_TEST_POSTGRES_CONTAINER: ${{ job.services.postgres.id }}
run: |
set -euo pipefail
client_dir="$RUNNER_TEMP/pg-client"
mkdir -p "$client_dir"
for tool in pg_dump psql; do
printf '#!/bin/sh\nexec docker exec -i --env PGPASSWORD "%s" %s "$@"\n' "$CPR_TEST_POSTGRES_CONTAINER" "$tool" > "$client_dir/$tool"
chmod +x "$client_dir/$tool"
done
echo "$client_dir" >> "$GITHUB_PATH"
- name: Rust tests
working-directory: backend
run: |
+1 -1
View File
@@ -119,7 +119,7 @@ jobs:
run: ruby .github/scripts/toolchains.rb >> "$GITHUB_OUTPUT"
- name: Setup Node
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
uses: actions/setup-node@949feb2413d6458794dcd2491c4babbbce0c15c1 # v7.1.0
with:
node-version: ${{ steps.toolchains.outputs.node }}
- name: Setup Rust
+13 -13
View File
@@ -57,7 +57,7 @@ jobs:
run: ruby .github/scripts/toolchains.rb >> "$GITHUB_OUTPUT"
- name: Upload sealed plugin assets
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: plugin-release-assets
path: dist/plugin-release-assets
@@ -143,7 +143,7 @@ jobs:
CPR_BUILD_TYPE: ${{ needs.release-info.outputs.build_type }}
- name: Upload binary artifact
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: codex-proxy-rs-${{ matrix.os }}-${{ matrix.arch }}
path: dist/binaries/${{ matrix.os }}-${{ matrix.arch }}/
@@ -196,7 +196,7 @@ jobs:
CPR_BUILD_TYPE: ${{ needs.release-info.outputs.build_type }}
- name: Upload binary artifact
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: codex-proxy-rs-${{ matrix.os }}-${{ matrix.arch }}
path: dist/binaries/${{ matrix.os }}-${{ matrix.arch }}/
@@ -219,19 +219,19 @@ jobs:
ref: ${{ needs.release-info.outputs.git_sha }}
- name: Download frontend artifact
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
name: frontend-dist
path: dist/frontend
- name: Download binary artifacts
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
pattern: codex-proxy-rs-*
path: dist/binaries
- name: Download sealed plugin assets
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
name: plugin-release-assets
path: dist/plugin-release-assets
@@ -250,7 +250,7 @@ jobs:
subject-path: dist/release/*
- name: Upload release assets
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: release-assets
path: dist/release
@@ -283,19 +283,19 @@ jobs:
ref: ${{ needs.release-info.outputs.git_sha }}
- name: Download frontend artifact
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
name: frontend-dist
path: dist/docker/web/dist
- name: Download Docker binary artifacts
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
name: codex-proxy-rs-${{ matrix.os }}-${{ matrix.arch }}
path: dist/docker-artifacts/codex-proxy-rs-${{ matrix.os }}-${{ matrix.arch }}
- name: Download sealed plugin assets
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
name: plugin-release-assets
path: dist/plugin-release-assets
@@ -362,7 +362,7 @@ jobs:
printf '%s\n' "$IMAGE_DIGEST" > "dist/digests/linux-${IMAGE_ARCH}.txt"
- name: Upload verified image digest
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
with:
name: image-digest-${{ matrix.os }}-${{ matrix.arch }}
path: dist/digests/*.txt
@@ -390,13 +390,13 @@ jobs:
ref: ${{ needs.release-info.outputs.git_sha }}
- name: Download release assets
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
name: release-assets
path: dist/release
- name: Download verified image digests
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
with:
pattern: image-digest-*
path: dist/digests
+11 -2
View File
@@ -33,14 +33,23 @@ jobs:
...params,
per_page: 100,
});
// 按 PR 作者识别官方机器人,维护者重开或编辑 PR 时仍使用同一判断
const isDependabot = pr.user.login === 'dependabot[bot]'
&& pr.user.type === 'Bot'
&& pr.head.repo?.full_name === `${owner}/${repo}`;
const isDependencyUpdate = file => isDependabot
&& file.status === 'modified'
&& (/^\.github\/workflows\/[^/]+\.ya?ml$/.test(file.filename)
|| /^\.github\/actions\/.+\/(?:action\.ya?ml|Dockerfile)$/.test(file.filename));
const isProtectedRule = path => typeof path === 'string'
&& (path === 'AGENTS.md' || path.endsWith('/AGENTS.md')
|| path === 'CONTRIBUTING.md'
|| path === '.agents/skills' || path.startsWith('.agents/skills/')
|| path === '.github' || path.startsWith('.github/'));
// 同时检查重命名前后的路径,避免通过移走规则文件绕过检查。
if (!files.some(file => isProtectedRule(file.filename) || isProtectedRule(file.previous_filename))) {
// 依赖更新只豁免已有文件的修改,重命名前后的受保护路径仍参与检查
if (!files.some(file => !isDependencyUpdate(file)
&& (isProtectedRule(file.filename) || isProtectedRule(file.previous_filename)))) {
core.info('No protected contribution rule changes detected.');
return;
}
+2 -2
View File
@@ -28,7 +28,7 @@ jobs:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
- name: Install cargo-audit
uses: taiki-e/install-action@4cef1412cce204788f482e778a0b9187f9626a29 # v2.87.21
uses: taiki-e/install-action@8cca0918f8471c33515202b73f56ab27f2bed283 # v2.87.29
with:
tool: cargo-audit@0.22.2
fallback: none
@@ -56,7 +56,7 @@ jobs:
run: ruby .github/scripts/toolchains.rb >> "$GITHUB_OUTPUT"
- name: Setup Node
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
uses: actions/setup-node@949feb2413d6458794dcd2491c4babbbce0c15c1 # v7.1.0
with:
node-version: ${{ steps.toolchains.outputs.node }}
cache: pnpm
+7 -4
View File
@@ -23,11 +23,12 @@
AI 草拟或发布 Issue、PR 讨论及审查意见的回复时:
- 用第一人称“我”,不用“我们”“本助手”代称自己
- 围绕问题、事实和结论展开,能自然省略主语时省略;需要说明个人判断或行动归属时用“我”,避免反复以“我”起句或强调自己,不用“我们”“本助手”代称自己
- 普通回复直接回应问题,不加问候语或称呼;向对方提问时称呼“佬”,同一条回复不重复称呼
- 客气、委婉地回应具体问题,补充信息或提出不同意见时说明缘由,避免命令式、责备式表达
- 使用自然、客气的日常交流口吻,适度用“看了下”“方便补充一下吗”等缓和表达,例如“我看了截图”可写为“我看了下截图”,不机械堆叠客套或语气词
- 补充信息或提出不同意见时先说明观察和依据,再温和表达判断,避免生硬纠正、命令或责备;已确认的事实保持明确,不为委婉写成猜测
- 使用自然短句,简单回复用短段,并列信息较多时再分点;少用句号,避免客服套话、机械模板和空泛保证
- 只把实际完成的排查、复现、修改和验证写成“我已……”;未采纳的意见说明依据,未确认、待处理或未验证的内容如实交代,说清结论、限制和下一步
- 排查、复现、修改和验证的完成状态与实际一致;未采纳的意见说明依据,未确认、待处理或未验证的内容如实交代,说清结论、限制和下一步
## 问题与方案依据
@@ -85,7 +86,9 @@ PR 不接受以下协作规则文件的新增、修改、删除或重命名:
- `.agents/skills/` 中的所有文件,包括技能入口、引用文档、配置和脚本
- `.github/` 中的所有文件,包括工作流、Actions、Issue/PR 模板和依赖更新配置
调整这些文件时先提交 Issue 说明原因和建议,由维护者处理;[规则保护工作流](.github/workflows/rules.yml) 会提示并关闭包含此类变更的 PR。撤销相关变更后可以重新打开 PR,继续正常检查
官方 `dependabot[bot]` 从本仓库分支发起的依赖更新可修改 `.github/workflows/` 中已有的 YAML 工作流,以及 `.github/actions/` 下已有的 `action.yml`、`action.yaml` 和 `Dockerfile`,继续执行正常审查与 CI;其他受保护文件以及新增、删除、重命名不在例外范围内
除此之外,调整这些文件时先提交 Issue 说明原因和建议,由维护者处理;[规则保护工作流](.github/workflows/rules.yml) 会提示并关闭包含此类变更的 PR。撤销相关变更后可以重新打开 PR,继续正常检查
源码变更遵守现有架构边界与迁移冻结规则,由对应 CI 检查执行;不额外设立整块业务目录的 PR 禁改清单。构建、测试和安全扫描复用现有工作流
+5 -4
View File
@@ -1903,6 +1903,7 @@ dependencies = [
name = "gateway-store"
version = "0.1.0"
dependencies = [
"async-compression",
"async-stream",
"async-trait",
"aws-credential-types",
@@ -3832,9 +3833,9 @@ dependencies = [
[[package]]
name = "rustls-webpki"
version = "0.103.15"
version = "0.103.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2"
checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281"
dependencies = [
"aws-lc-rs",
"ring",
@@ -5222,9 +5223,9 @@ dependencies = [
[[package]]
name = "webpki-roots"
version = "1.0.5"
version = "1.0.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "12bed680863276c63889429bfd6cab3b99943659923822de1c8a39c49e4d722c"
checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9"
dependencies = [
"rustls-pki-types",
]
@@ -96,6 +96,10 @@ async fn execute(
Some(plugin_runtime.policy_registry()),
Some(plugin_runtime.execution_registry()),
);
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_snapshots(&core.snapshots())
);
let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control());
or_shutdown!(plugin_runtime, management.bind(&plugin_runtime));
let core = or_shutdown!(plugin_runtime, core.activate().await);
@@ -31,6 +31,10 @@ pub async fn run() -> Result<(), BootstrapError> {
Some(plugin_runtime.execution_registry()),
Some(plugin_runtime.frontend_authentication_registry()),
);
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_snapshots(&core.snapshots())
);
let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control());
or_shutdown!(plugin_runtime, management.bind(&plugin_runtime));
let mut core = or_shutdown!(plugin_runtime, core.activate().await);
@@ -11,7 +11,7 @@ use chrono::{DateTime, Utc};
use futures::future::BoxFuture;
use gateway_core::{
lifecycle::CancellationToken,
task::{DaemonTask, WorkerTaskError},
task::{DaemonTask, WorkerDaemonContext, WorkerTaskError},
};
use tracing::{info, warn};
@@ -75,18 +75,20 @@ impl BackupTask {
}
impl DaemonTask for BackupTask {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move { self.run_daemon(cancellation).await })
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move { self.run_daemon(context).await })
}
}
impl BackupTask {
async fn run_daemon(&self, cancellation: CancellationToken) -> Result<(), WorkerTaskError> {
async fn run_daemon(&self, context: WorkerDaemonContext) -> Result<(), WorkerTaskError> {
let cancellation = context.cancellation();
loop {
if cancellation.is_cancelled() {
return Ok(());
}
self.run_cycle(&cancellation).await?;
self.run_cycle(cancellation).await?;
context.succeeded();
tokio::select! {
_ = cancellation.cancelled() => return Ok(()),
@@ -706,6 +708,10 @@ impl BackupTask {
fn classify_dump_error(error: BackupError) -> (&'static str, &'static str) {
match error.code() {
code::CANCELLED => (code::CANCELLED, "备份已取消"),
code::ANALYTICS_NOT_READY => (
code::ANALYTICS_NOT_READY,
"用量统计正在重建,请在统计就绪后重试备份",
),
code::STAGING_SPACE_EXHAUSTED => (code::STAGING_SPACE_EXHAUSTED, "暂存磁盘空间不足"),
_ => (code::PG_DUMP_FAILED, "数据库导出失败"),
}
@@ -123,6 +123,7 @@ pub struct AccountRecord {
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub model_access: gateway_core::account::AccountModelAccess,
pub outbound_proxy_id: Option<String>,
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub outbound_proxy_enabled: bool,
pub credential_state: CredentialState,
@@ -15,6 +15,8 @@ pub const STATE_CONFLICT: &str = "backup.state_conflict";
pub const STORAGE_IDENTITY_LOCKED: &str = "backup.storage_identity_locked";
/// `pg_dump` 非零退出或失败
pub const PG_DUMP_FAILED: &str = "backup.pg_dump_failed";
/// 用量汇总尚未完整发布
pub const ANALYTICS_NOT_READY: &str = "backup.analytics_not_ready";
/// 暂存磁盘空间不足
pub const STAGING_SPACE_EXHAUSTED: &str = "backup.staging_space_exhausted";
/// S3 认证失败
@@ -499,7 +499,7 @@ impl BackupError {
/// 构造稳定错误码常量
pub mod code;
/// 构建对象 key:`{prefix}/YYYY/MM/DD/codex-proxy-rs_{UTC seconds}_{backup id}.dump`
/// 构建对象 key:`{prefix}/YYYY/MM/DD/codex-proxy-rs_{UTC seconds}_{backup id}.sql.gz`
///
/// 规范化规则:prefix 不得为空、不得以前导 `/` 开头、不得包含 `..`、反斜杠或控制字符;
/// 输出统一使用 `/` 分隔且不含空段
@@ -510,7 +510,7 @@ pub fn build_object_key(
) -> Result<String, BackupError> {
let prefix = normalize_prefix(prefix)?;
Ok(format!(
"{prefix}/{}/{:02}/{:02}/codex-proxy-rs_{}_{}.dump",
"{prefix}/{}/{:02}/{:02}/codex-proxy-rs_{}_{}.sql.gz",
at.format("%Y"),
at.format("%m"),
at.format("%d"),
@@ -519,10 +519,10 @@ pub fn build_object_key(
))
}
/// 下载文件名:`codex-proxy-rs_{短 id}.dump`,与对象 key 的文件名一致
/// 下载文件名:`codex-proxy-rs_{短 id}.sql.gz`,与对象 key 的文件名一致
#[must_use]
pub fn build_download_file_name(backup_id: &str) -> String {
format!("codex-proxy-rs_{}.dump", short_backup_id(backup_id))
format!("codex-proxy-rs_{}.sql.gz", short_backup_id(backup_id))
}
/// 截取备份 id 的可读短标识:去掉 `backup_` 前缀,最多保留 8 位 hex
@@ -9,10 +9,7 @@ use std::{
use chrono::Utc;
use futures::{FutureExt as _, StreamExt as _, future::BoxFuture, stream::FuturesUnordered};
use gateway_core::{
lifecycle::CancellationToken,
task::{DaemonTask, WorkerTaskError},
};
use gateway_core::task::{DaemonTask, WorkerDaemonContext, WorkerTaskError};
use tokio::{sync::Notify, time::Instant};
use uuid::Uuid;
@@ -332,8 +329,9 @@ impl ImportTasksService for DefaultImportTasksService {
pub(crate) struct ImportTaskWorker(pub Arc<DefaultImportTasksService>);
impl DaemonTask for ImportTaskWorker {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let cancellation = context.cancellation();
let mut active = FuturesUnordered::new();
// 固定一个取消 future,避免循环反复注册中立 CancellationToken 的 waiter
let cancelled = cancellation.cancelled();
@@ -354,6 +352,7 @@ impl DaemonTask for ImportTaskWorker {
};
active.push(self.0.execute(id, index, input));
}
context.succeeded();
tokio::select! {
() = &mut cancelled => {},
_ = active.next(), if !active.is_empty() => {},
@@ -5,6 +5,7 @@
//! 供 backup::task 与 use_case::backup 的测试复用
use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, Ordering};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
@@ -92,6 +93,7 @@ pub(crate) fn queued_record(seed: &BackupRecordSeed) -> BackupRecord {
/// 内存版备份仓储
pub(crate) struct FakeBackupRepository {
settings: Mutex<BackupSettings>,
settings_unavailable: AtomicBool,
records: Mutex<Vec<BackupRecord>>,
}
@@ -99,10 +101,16 @@ impl FakeBackupRepository {
pub(crate) fn new(settings: BackupSettings) -> Self {
Self {
settings: Mutex::new(settings),
settings_unavailable: AtomicBool::new(false),
records: Mutex::new(Vec::new()),
}
}
pub(crate) fn set_settings_unavailable(&self, unavailable: bool) {
self.settings_unavailable
.store(unavailable, Ordering::SeqCst);
}
/// 测试辅助:读取当前全部记录
pub(crate) fn all_records(&self) -> Vec<BackupRecord> {
self.records.lock().expect("records").clone()
@@ -134,6 +142,13 @@ impl FakeBackupRepository {
#[async_trait]
impl BackupRepository for FakeBackupRepository {
async fn load_settings(&self) -> AdminStoreResult<BackupSettings> {
if self.settings_unavailable.load(Ordering::SeqCst) {
return Err(AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
"backup",
"test store unavailable",
));
}
Ok(self.settings.lock().expect("settings").clone())
}
@@ -11,6 +11,7 @@ use gateway_admin::{
ports::backup::{BackupRepository, DatabaseDumpPort, DumpRequest},
};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::task::{DaemonTask, WorkerDaemonContext};
use super::support::{
DUMP_CONTENT, FakeBackupRepository, FakeDumpPort, FakeObjectStore, backup_id,
@@ -379,3 +380,40 @@ async fn manual_backup_preserves_primary_and_final_delete_failures() {
);
}
}
#[tokio::test(start_paused = true)]
async fn daemon_reports_success_only_after_a_complete_cycle_and_recovers_after_store_failure() {
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
let task = Arc::new(BackupTask::new(
repository.clone(),
Arc::new(FakeDumpPort::new()),
Arc::new(FakeObjectStore::new()),
));
let context = WorkerDaemonContext::new(CancellationToken::new());
repository.set_settings_unavailable(true);
let error = task.run(context.clone()).await.unwrap_err();
assert_eq!(error.as_safe_str(), "backup store failed");
assert_eq!(context.completed_cycles(), 0);
repository.set_settings_unavailable(false);
let recovered = WorkerDaemonContext::new(CancellationToken::new());
let worker = tokio::spawn({
let task = task.clone();
let context = recovered.clone();
async move { task.run(context).await }
});
recovered.progress_since(0).await;
assert_eq!(recovered.completed_cycles(), 1);
repository.set_settings_unavailable(true);
tokio::time::advance(std::time::Duration::from_secs(30)).await;
assert_eq!(
worker.await.unwrap().unwrap_err().as_safe_str(),
"backup store failed"
);
assert_eq!(recovered.completed_cycles(), 1);
let cancelled = WorkerDaemonContext::new(CancellationToken::new());
cancelled.cancellation().cancel();
task.run(cancelled.clone()).await.unwrap();
assert_eq!(cancelled.completed_cycles(), 0);
}
@@ -58,7 +58,7 @@ fn object_key_shape() {
let key = build_object_key("production/backups", "backup_abc", at).unwrap();
assert_eq!(
key,
"production/backups/2026/08/01/codex-proxy-rs_20260801_123000_abc.dump"
"production/backups/2026/08/01/codex-proxy-rs_20260801_123000_abc.sql.gz"
);
}
@@ -3194,6 +3194,7 @@ fn quota_local_usage(account_id: &str, total_tokens: u64) -> AccountUsage {
pub(super) fn account_record(kind: &str) -> AccountRecord {
let now = Utc::now();
AccountRecord {
outbound_proxy_id: None,
outbound_proxy_enabled: true,
notes: None,
model_access: Default::default(),
@@ -14,7 +14,7 @@ use gateway_admin::{
use gateway_core::{
lifecycle::CancellationToken,
routing::ProviderKind,
task::{WorkerContribution, WorkerKind, WorkerRunnable},
task::{WorkerContribution, WorkerDaemonContext, WorkerKind, WorkerRunnable},
};
use uuid::Uuid;
@@ -75,7 +75,9 @@ fn start(bundle: &mut AdminBundle) -> (CancellationToken, tokio::task::JoinHandl
let cancellation = CancellationToken::new();
let token = cancellation.clone();
let handle = tokio::spawn(async move {
task.run(token).await.expect("import daemon exits cleanly");
task.run(WorkerDaemonContext::new(token))
.await
.expect("import daemon exits cleanly");
});
(cancellation, handle)
}
@@ -120,6 +120,7 @@ pub(super) fn account_view(
},
weight: account.weight.get(),
model_access: account.model_access,
outbound_proxy_id: account.outbound_proxy_id,
outbound_proxy_endpoint: account
.outbound_proxy
.as_ref()
@@ -267,6 +267,7 @@ pub struct AccountSummaryView {
#[serde(rename_all = "camelCase")]
pub struct AccountView {
pub capabilities: AccountCapabilitiesView,
pub outbound_proxy_id: Option<String>,
pub outbound_proxy_endpoint: Option<String>,
pub id: String,
pub name: String,
@@ -9,6 +9,67 @@ use tower::ServiceExt as _;
use super::super::{AdminTestFixture, AdminTestState};
#[tokio::test]
async fn account_proxy_binding_is_distinct_from_the_redacted_endpoint() {
use chrono::Utc;
use gateway_core::account::{AccountStatusFacts, CredentialState, OutboundProxy, QuotaState};
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
for (proxy_id, proxy_url, enabled) in [
(
Some("proxy-a"),
Some("socks5h://alice:secret@127.0.0.1:50101"),
true,
),
(
Some("proxy-b"),
Some("socks5h://bob:secret@127.0.0.1:50101"),
false,
),
(None, Some("socks5h://alice:secret@127.0.0.1:50101"), true),
(None, None, true),
] {
let now = Utc::now();
let facts = AccountStatusFacts {
enabled: true,
credential_state: CredentialState::Ready,
access_token_expires_at: None,
quota: QuotaState::unknown(),
cooldown: None,
last_error_reason: None,
last_error_message: None,
};
let mut item = account_page_item(facts, now, "oauth");
item.account.outbound_proxy_id = proxy_id.map(str::to_owned);
item.account.outbound_proxy = proxy_url.map(|url| OutboundProxy::parse(url).unwrap());
item.account.outbound_proxy_enabled = enabled;
*fixture.account.lock().unwrap() = Some(item);
let response = admin::accounts::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
Request::builder()
.uri("/api/admin/accounts")
.header("x-request-id", "req_proxy_binding")
.header(header::COOKIE, "cpr_session=valid-session")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), 32768).await.unwrap();
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
let account = &value["data"]["items"][0];
assert_eq!(account["outboundProxyId"], serde_json::json!(proxy_id));
assert_eq!(
account["outboundProxyEndpoint"],
serde_json::json!(proxy_url.map(|_| "socks5h://127.0.0.1:50101"))
);
assert!(!String::from_utf8(body.to_vec()).unwrap().contains("secret"));
}
}
#[tokio::test]
async fn account_cooldown_returns_an_instant_and_server_formatted_recovery() {
use chrono::{DateTime, Duration, Utc};
@@ -112,6 +173,7 @@ fn account_page_item(
AccountPageItem {
account: AccountRecord {
outbound_proxy_id: None,
outbound_proxy_enabled: true,
id: "acct_cooldown".to_owned(),
provider_kind: gateway_core::routing::ProviderKind::new("openai").unwrap(),
@@ -511,6 +511,7 @@ mod provider {
total_slots: None,
},
account: AccountRecord {
outbound_proxy_id: None,
outbound_proxy_enabled: true,
notes: None,
model_access: Default::default(),
@@ -18,6 +18,7 @@ struct State {
pub struct ConnectionBudget {
state: Arc<Mutex<State>>,
now: fn() -> Instant,
reserve_next_opening: bool,
}
impl Default for ConnectionBudget {
@@ -37,9 +38,26 @@ impl ConnectionBudget {
Self {
state: Arc::default(),
now,
reserve_next_opening: false,
}
}
/// 仅本次恢复使用一半剩余时间并留下一次 opening;已启动的后台连接不受影响
#[must_use]
pub fn reserve_next_opening(mut self) -> Self {
self.reserve_next_opening = true;
self
}
#[must_use]
pub fn remaining_openings(&self) -> u32 {
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
MAX_OPENINGS.saturating_sub(state.openings)
}
/// 在可能建立连接前预约一次 opening;返回本次连接允许使用的剩余时间
/// HTTP 连接池隐藏了是否冷建连,因此一次 HTTP opening 也计入次数上限
pub fn begin(&self) -> Result<Option<Duration>, ConnectionBudgetExhausted> {
@@ -53,11 +71,16 @@ impl ConnectionBudget {
let started = *state.started.get_or_insert_with(self.now);
let remaining =
STARTUP_TIMEOUT.saturating_sub((self.now)().saturating_duration_since(started));
if state.openings >= MAX_OPENINGS || remaining.is_zero() {
let opening_limit = MAX_OPENINGS - u32::from(self.reserve_next_opening);
if state.openings >= opening_limit || remaining.is_zero() {
return Err(ConnectionBudgetExhausted);
}
state.openings += 1;
Ok(Some(remaining))
Ok(Some(if self.reserve_next_opening {
remaining / 2
} else {
remaining
}))
}
/// 仅连接成功不足以清零;确认业务发送后退出新增的 NotSent 恢复策略
@@ -318,6 +318,7 @@ where
concurrency_wait_budget: ConcurrencyWaitBudget::default(),
connection_budget: super::connection::ConnectionBudget::default(),
connection_retries: 0,
connection_rotation: None,
observation: ResponseObservation::new(timing_started_at),
request_observation,
budget_prior_attempts_usd: Decimal::ZERO,
@@ -424,6 +425,7 @@ struct CurrentAttempt {
metadata: ProviderCallMetadata,
trigger: AttemptTrigger,
transport: AttemptTransport,
connection_rotation_reserved: bool,
index: NonZeroU32,
started_at: SystemTime,
send_observed: bool,
@@ -458,6 +460,7 @@ struct PendingAttemptRetry {
transport: AttemptTransport,
delay: Duration,
transport_recovery: bool,
reserve_connection_rotation: bool,
}
/// API 可逐事件消费的 Core 执行会话
@@ -471,6 +474,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
concurrency_wait_budget: ConcurrencyWaitBudget,
connection_budget: super::connection::ConnectionBudget,
connection_retries: u32,
/// 原账号恢复仍建连失败后尝试重选;选号不可用时消费该状态,恢复原账号剩余预算
connection_rotation: Option<PendingAttemptRetry>,
observation: ResponseObservation,
request_observation: Option<RequestObservationDispatch>,
budget_prior_attempts_usd: Decimal,
@@ -999,6 +1004,9 @@ where
let is_transport_recovery = pending_retry
.as_ref()
.is_some_and(|retry| retry.transport_recovery);
let connection_rotation_reserved = pending_retry
.as_ref()
.is_some_and(|retry| retry.reserve_connection_rotation);
if !is_transport_recovery && self.routing_attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
@@ -1060,7 +1068,9 @@ where
self.recovery_account
.take()
.or_else(|| self.account_selection.required_account().cloned()),
AttemptTransport::Default,
self.connection_rotation
.as_ref()
.map_or(AttemptTransport::Default, |recovery| recovery.transport),
),
}
};
@@ -1122,8 +1132,13 @@ where
.with_pricing(self.plan.pricing())
.with_request_location(self.plan.request_location().cloned())
.with_privacy(self.plan.privacy())
.with_request_body_limit(self.plan.request_body_limit())
.with_concurrency_wait_budget(self.concurrency_wait_budget.clone())
.with_connection_budget(self.connection_budget.clone())
.with_connection_budget(if connection_rotation_reserved {
self.connection_budget.clone().reserve_next_opening()
} else {
self.connection_budget.clone()
})
.with_timing_started_at(self.observation.timing_started_at)
.with_request_policy(self.request_policy.clone())
.with_execution_effects(self.execution_effects.as_ref().map(Arc::clone))
@@ -1186,9 +1201,21 @@ where
Ok(stream) => stream,
Err(error) => {
attempt_trace.record_provider_failure(&error);
let connection_fallback = self.connection_rotation.is_some()
&& !self.connection_budget.exhausted()
&& !error.retry_is_prohibited()
&& error.send_state() == UpstreamSendState::NotSent
&& matches!(
error.kind(),
ProviderErrorKind::NoEligibleAccount
| ProviderErrorKind::AccountCapacityUnavailable
| ProviderErrorKind::QuotaExhausted
| ProviderErrorKind::ConcurrencyQueueFull
| ProviderErrorKind::ConcurrencyQueueTimeout
);
let continuation_retry = !error.retry_is_prohibited()
&& self.prepare_unavailable_native_continuation_replay(&error);
let candidate_retry = !error.retry_is_prohibited() && !continuation_retry
let candidate_retry = !error.retry_is_prohibited() && !connection_fallback && !continuation_retry
&& matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
@@ -1224,12 +1251,15 @@ where
.get()
.saturating_sub(self.routing_attempts),
remaining_deadline: Duration::ZERO,
retry_allowed: continuation_retry || candidate_retry,
retry_allowed: connection_fallback
|| continuation_retry
|| candidate_retry,
})
.await?;
attempt_trace.record("retry.decided", json!({
"phase": "prepare", "retryable": retryable,
"continuationRetry": continuation_retry, "candidateRetry": candidate_retry,
"connectionFallback": connection_fallback,
"retryProhibited": error.retry_is_prohibited(),
"sendState": error.send_state().as_str(),
"requestSendState": self.current_send_state().as_str(),
@@ -1239,6 +1269,17 @@ where
"remainingRoutingAttempts": self.plan.max_attempts().get().saturating_sub(self.routing_attempts),
}));
if retryable {
if connection_fallback {
let recovery = self
.connection_rotation
.take()
.expect("pending connection rotation");
self.excluded_accounts.remove(&recovery.account);
self.pending_retry = Some(recovery);
} else if candidate_retry {
// 原账号只属于原 Provider,跨候选后不能拿旧钉选作恢复兜底
self.connection_rotation = None;
}
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
@@ -1306,6 +1347,7 @@ where
}
let metadata = stream.metadata().clone();
self.connection_rotation = None;
attempt_trace.record("account.selected", json!({
"provider": metadata.provider().as_str(), "accountId": metadata.provider_account_id().as_str(),
"transport": metadata.transport().as_str(),
@@ -1429,6 +1471,7 @@ where
metadata,
trigger,
transport: attempt_transport,
connection_rotation_reserved,
index: next_attempt,
started_at: SystemTime::now(),
send_observed: false,
@@ -1539,6 +1582,12 @@ where
.attempt(self.attempts)
.record_provider_failure(&error);
let mut atomic_client_events = error.take_atomic_client_events();
// 失败随附的观测也属于当前 attempt,换号记录不能等到最终交付才读取
for event in &atomic_client_events {
if let Some(observation) = event.response_observation() {
self.observe_response(observation.clone());
}
}
let current = self.current.take().ok_or(EngineError::NoActiveAttempt)?;
if self.request_observation.is_some() {
let websocket_attempt = websocket_observation_attempt(&current);
@@ -1642,6 +1691,18 @@ where
}
_ => None,
};
let can_rotate_connection = connection_retry_requested
&& self.account_selection.required_account().is_none()
&& self.routing_attempts < self.plan.max_attempts().get()
&& self.account_rotations
< self.plan.account_selection_policy().max_account_rotations();
// A 先恢复一次,再利用预留预算重选;重选继续走已有降级传输,避免 WS 抢占最后一次 opening
let connection_rotation = transport_recovery
.filter(|_| can_rotate_connection && current.connection_rotation_reserved);
let reserve_connection_rotation = can_rotate_connection
&& !current.connection_rotation_reserved
&& self.connection_budget.remaining_openings() >= 2;
let transport_recovery = transport_recovery.filter(|_| connection_rotation.is_none());
let ordinary_retry = !self.connection_budget.exhausted()
&& !connection_retry_requested
&& !execution_effect_observed
@@ -1714,6 +1775,7 @@ where
|| same_account_retry
|| transient_retry.is_some()
|| transport_recovery.is_some()
|| connection_rotation.is_some()
|| rotation_retry);
let retryable = match self
@@ -1750,6 +1812,8 @@ where
"retryable": retryable, "continuationRetry": continuation_retry,
"sameAccountRetry": same_account_retry, "accountRotationRetry": account_rotation_retry,
"ordinaryRetry": ordinary_retry, "transportRecovery": transport_recovery.is_some(),
"connectionRotation": connection_rotation.is_some(),
"reserveConnectionRotation": reserve_connection_rotation,
"transientRetry": transient_retry.is_some(),
"clientRetrySuppressed": suppress_transient_retry,
"rotationRetry": rotation_retry, "accountRotations": self.account_rotations,
@@ -1795,6 +1859,7 @@ where
transport: current.transport,
delay,
transport_recovery: false,
reserve_connection_rotation: false,
});
} else if let Some((transport, delay)) = transport_recovery {
self.pending_retry = Some(PendingAttemptRetry {
@@ -1802,6 +1867,17 @@ where
transport,
delay,
transport_recovery: true,
reserve_connection_rotation,
});
} else if let Some((transport, delay)) = connection_rotation {
let account = current.metadata.provider_account_id().clone();
self.excluded_accounts.insert(account.clone());
self.connection_rotation = Some(PendingAttemptRetry {
account,
transport,
delay,
transport_recovery: true,
reserve_connection_rotation: false,
});
} else if !continuation_retry {
self.excluded_accounts
@@ -65,6 +65,22 @@ pub struct ExecutionRequestMetadata {
pub previous_response_id: Option<PreviousResponseId>,
}
impl ExecutionRequestMetadata {
pub(super) fn inherit_request_origin(&mut self, scope: &ExtensionCallScope) {
let Some(request) = scope.request() else {
return;
};
// 执行入口可以是内部调用,客户端来源仍属于原始入口快照
self.client_ip = request.client_ip;
self.user_agent = request
.headers
.iter()
.find(|header| header.name().eq_ignore_ascii_case("user-agent"))
.and_then(|header| std::str::from_utf8(header.value()).ok())
.map(str::to_owned);
}
}
#[derive(Clone)]
pub struct AuthenticatedClient {
pub(super) settings: Option<RequestSettings>,
@@ -197,6 +197,15 @@ pub struct BoundModelExecutionContext {
}
impl BoundModelExecutionContext {
/// 与实际执行共用绑定时的快照,不能回退到后续发布的配置
#[must_use]
pub fn request_body_limit(&self) -> usize {
self.authority
.client
.snapshot
.responses_max_decompressed_body_bytes()
}
#[must_use]
pub fn request_settings(&self) -> Option<&RequestSettings> {
self.authority.client.request_settings()
@@ -448,6 +448,12 @@ impl DefaultExecutionService {
account_group_ids.clone(),
)
.with_extension_scope(authorization.extension_scope.clone())
.with_request_body_limit(
request
.client
.snapshot
.responses_max_decompressed_body_bytes(),
)
.with_execution_effects(Arc::clone(&authorization.graph.as_ref()?.effects)),
)
});
@@ -995,8 +1001,7 @@ impl DefaultExecutionService {
let execution_effects_baseline = parent.graph.effects.epoch();
metadata.endpoint = "host.model".to_owned();
metadata.transport = ClientTransport::InternalPlugin;
metadata.client_ip = None;
metadata.user_agent = None;
metadata.inherit_request_origin(&extension_scope);
let authorization = AuthorizedExecution {
account_scope,
response_control: None,
@@ -1149,8 +1154,7 @@ impl DefaultExecutionService {
let execution_effects_baseline = graph.effects.epoch();
metadata.endpoint = "host.model".to_owned();
metadata.transport = ClientTransport::InternalPlugin;
metadata.client_ip = None;
metadata.user_agent = None;
metadata.inherit_request_origin(&authority.extension_scope);
self.client_api_key_usage
.record_used(authority.client.policy.key_id());
self.start_authorized(
@@ -559,9 +559,22 @@ pub struct MiddlewareContext {
deadline: crate::lifecycle::Deadline,
extension_scope: ExtensionCallScope,
execution_effects: Option<Arc<ExecutionEffects>>,
request_body_limit: Option<usize>,
}
impl MiddlewareContext {
/// attempt 只继承执行预算,不开放 request 设置改写
#[must_use]
pub const fn with_request_body_limit(mut self, bytes: Option<usize>) -> Self {
self.request_body_limit = bytes;
self
}
#[must_use]
pub const fn request_body_limit(&self) -> Option<usize> {
self.request_body_limit
}
#[must_use]
pub fn plan(&self) -> Option<FrozenMiddlewarePlan> {
self.plan.clone()
@@ -590,6 +603,7 @@ impl MiddlewareContext {
} = authority;
Self {
plan: None,
request_body_limit: None,
request_id,
mount,
attempt_index,
+15 -1
View File
@@ -286,6 +286,7 @@ impl ContinuationAttempt {
/// Provider 每次执行可见的 request-local context
#[derive(Debug, Clone)]
pub struct RequestAttemptContext {
request_body_limit: Option<usize>,
lease_metadata: Option<Arc<crate::provider_ports::ProviderLeaseMetadata>>,
response_control: Option<response_control::ResponseControl>,
pricing: Arc<crate::metering::PricingOverrides>,
@@ -311,6 +312,12 @@ pub struct RequestAttemptContext {
}
impl RequestAttemptContext {
#[must_use]
pub const fn with_request_body_limit(mut self, bytes: usize) -> Self {
self.request_body_limit = Some(bytes);
self
}
#[must_use]
pub fn with_lease_metadata(
mut self,
@@ -387,6 +394,7 @@ impl RequestAttemptContext {
pub fn new(request_id: ModelRequestId, client_api_key_ref: ClientApiKeyId) -> Self {
Self {
response_control: None,
request_body_limit: None,
lease_metadata: None,
request_id,
client_api_key_ref,
@@ -542,6 +550,11 @@ pub struct AttemptContext {
}
impl AttemptContext {
#[must_use]
pub const fn request_body_limit(&self) -> Option<usize> {
self.request.request_body_limit
}
#[must_use]
pub fn lease_metadata(&self) -> Option<Arc<crate::provider_ports::ProviderLeaseMetadata>> {
self.request.lease_metadata.clone()
@@ -747,7 +760,8 @@ impl AttemptContext {
extension_scope: self.request.extension_scope.clone(),
execution_effects: self.request.execution_effects.as_ref().map(Arc::clone),
},
);
)
.with_request_body_limit(self.request.request_body_limit);
provider::execute_attempt_middleware(
self.request.middleware.as_ref(),
context,
@@ -34,9 +34,15 @@ pub struct ModelRouteInput {
requested_model: PublicModelId,
available_providers: BTreeSet<ProviderKind>,
execution_effects: Arc<ExecutionEffects>,
request_body_limit: Option<usize>,
}
impl ModelRouteInput {
#[must_use]
pub const fn request_body_limit(&self) -> Option<usize> {
self.request_body_limit
}
#[must_use]
pub const fn request_id(&self) -> &super::ModelRequestId {
&self.request_id
@@ -188,9 +194,15 @@ pub struct AccountScheduleInput {
model: Option<String>,
candidates: Vec<AccountScheduleCandidate>,
execution_effects: Arc<ExecutionEffects>,
request_body_limit: Option<usize>,
}
impl AccountScheduleInput {
#[must_use]
pub const fn request_body_limit(&self) -> Option<usize> {
self.request_body_limit
}
#[must_use]
pub const fn request_id(&self) -> &super::ModelRequestId {
&self.request_id
@@ -282,9 +294,15 @@ pub struct RetryInput {
pub extension_scope: ExtensionCallScope,
pub facts: RetryFacts,
execution_effects: Arc<ExecutionEffects>,
request_body_limit: Option<usize>,
}
impl RetryInput {
#[must_use]
pub const fn request_body_limit(&self) -> Option<usize> {
self.request_body_limit
}
#[must_use]
pub fn execution_effects(&self) -> Arc<ExecutionEffects> {
Arc::clone(&self.execution_effects)
@@ -339,6 +357,7 @@ pub struct RequestPolicyContext {
account_group_ids: Arc<[AccountGroupId]>,
extension_scope: ExtensionCallScope,
execution_effects: Arc<ExecutionEffects>,
request_body_limit: Option<usize>,
}
impl fmt::Debug for RequestPolicyContext {
@@ -353,6 +372,13 @@ impl fmt::Debug for RequestPolicyContext {
}
impl RequestPolicyContext {
/// 路由、调度和重试沿用同一次模型执行的冻结预算
#[must_use]
pub const fn with_request_body_limit(mut self, bytes: usize) -> Self {
self.request_body_limit = Some(bytes);
self
}
pub async fn retry_decision(&self, facts: RetryFacts) -> RetryDecision {
self.plan
.retry_decision(RetryInput {
@@ -362,6 +388,7 @@ impl RequestPolicyContext {
extension_scope: self.extension_scope.clone(),
facts,
execution_effects: Arc::clone(&self.execution_effects),
request_body_limit: self.request_body_limit,
})
.await
.unwrap_or_else(|_| {
@@ -391,6 +418,7 @@ impl RequestPolicyContext {
account_group_ids: account_group_ids.into(),
extension_scope: ExtensionCallScope::default(),
execution_effects: Arc::new(ExecutionEffects::default()),
request_body_limit: None,
}
}
@@ -428,6 +456,7 @@ impl RequestPolicyContext {
requested_model,
available_providers,
execution_effects: Arc::clone(&self.execution_effects),
request_body_limit: self.request_body_limit,
})
.await
}
@@ -474,6 +503,7 @@ impl RequestPolicyContext {
model: model.map(str::to_owned),
candidates: projected,
execution_effects: Arc::clone(&self.execution_effects),
request_body_limit: self.request_body_limit,
})
.await
.map_err(|_| AccountPolicyError::Fault)?;
+6 -5
View File
@@ -547,7 +547,7 @@ pub struct ProviderResponseTimings {
/// 这样 Provider 专有协议不会渗入路由
/// 或管理领域
#[derive(Clone, PartialEq, Eq)]
pub struct ProviderResponseMetadata(String);
pub struct ProviderResponseMetadata(Arc<str>);
impl ProviderResponseMetadata {
const MAX_BYTES: usize = 32 * 1024;
@@ -561,7 +561,7 @@ impl ProviderResponseMetadata {
serde_json::from_str::<Value>(&json)
.ok()
.filter(Value::is_object)
.map(|_| Self(json))
.map(|_| Self(json.into()))
}
/// 返回不透明的 JSON 文本,供持久化 adapter 原样验证和写入
@@ -626,7 +626,8 @@ pub struct ProviderResponseObservation {
service_tier: Option<String>,
upstream_response_model: Option<String>,
timings: ProviderResponseTimings,
client_headers: Vec<ProviderResponseHeader>,
// 流式事件分别持有快照,共享不可变头集合,避免逐事件复制头名和集合
client_headers: Arc<[ProviderResponseHeader]>,
provider_metadata: Option<ProviderResponseMetadata>,
}
@@ -642,7 +643,7 @@ impl ProviderResponseObservation {
service_tier: None,
upstream_response_model: None,
timings: ProviderResponseTimings::default(),
client_headers: Vec::new(),
client_headers: Arc::default(),
provider_metadata: None,
}
}
@@ -719,7 +720,7 @@ impl ProviderResponseObservation {
#[must_use]
pub fn with_client_headers(mut self, client_headers: Vec<ProviderResponseHeader>) -> Self {
self.client_headers = client_headers;
self.client_headers = client_headers.into();
self
}
@@ -88,6 +88,15 @@ pub struct ExtensionPreparationError;
/// 准备阶段读取与请求快照相同的持久 revision;不得自行发布当前代次
pub trait ExtensionPreparationPort: Send + Sync {
/// 初始发布尚无可读快照,注册阶段沿用本次编译的正文预算
fn prepare_with_request_body_limit(
&self,
revision: ConfigRevision,
_request_body_limit: usize,
) -> BoxFuture<'_, Result<ExtensionSetReference, ExtensionPreparationError>> {
self.prepare(revision)
}
fn prepare(
&self,
revision: ConfigRevision,
@@ -561,6 +561,7 @@ impl ProviderCandidate {
/// 一次请求冻结的 Provider 尝试顺序
#[derive(Debug, Clone)]
pub struct RoutingPlan {
request_body_limit: usize,
pricing: Arc<crate::metering::PricingOverrides>,
request_location: Option<crate::account::RequestLocation>,
privacy: Option<Arc<dyn crate::settings::privacy::CompiledPrivacyPolicy>>,
@@ -573,6 +574,11 @@ pub struct RoutingPlan {
}
impl RoutingPlan {
#[must_use]
pub const fn request_body_limit(&self) -> usize {
self.request_body_limit
}
#[must_use]
pub fn pricing(&self) -> Arc<crate::metering::PricingOverrides> {
Arc::clone(&self.pricing)
@@ -331,7 +331,13 @@ impl RuntimeSnapshotCompiler {
});
if must_prepare {
let prepared = preparer
.prepare(revision)
.prepare_with_request_body_limit(
revision,
usize::try_from(
facts.settings.responses_max_decompressed_body_bytes,
)
.unwrap_or(usize::MAX),
)
.await
.map_err(|_| RuntimeSnapshotCompileError::ExtensionsUnavailable)?;
prepared_extensions = Some((revision, prepared));
@@ -1201,6 +1207,7 @@ impl RuntimeSnapshot {
}
Ok(RoutingPlan {
request_body_limit: self.responses_max_decompressed_body_bytes(),
config_revision: self.revision,
pricing: Arc::clone(&self.settings.values.pricing),
request_location: self.settings.request_location.clone(),
@@ -1282,6 +1289,7 @@ impl RuntimeSnapshot {
account_scope: Arc::clone(&account_scope),
};
Ok(RoutingPlan {
request_body_limit: self.responses_max_decompressed_body_bytes(),
config_revision: self.revision,
pricing: Arc::clone(&self.settings.values.pricing),
request_location: self.settings.request_location.clone(),
@@ -3,7 +3,9 @@
mod publications;
mod publisher;
pub use publications::{RuntimeSnapshotHandle, RuntimeSnapshotUnavailable};
pub use publications::{
RuntimeSnapshotHandle, RuntimeSnapshotUnavailable, WeakRuntimeSnapshotHandle,
};
pub use publisher::{
RuntimeSnapshotPublisher, SnapshotControl, SnapshotRevisionStream, SnapshotSubscriptionError,
SnapshotSubscriptionPort, runtime_revision_needs_refresh,
@@ -8,7 +8,7 @@ use crate::{
use futures::future::BoxFuture;
use std::{
collections::BTreeMap,
sync::{Arc, Mutex as SyncMutex, RwLock},
sync::{Arc, Mutex as SyncMutex, RwLock, Weak},
};
/// 请求级冻结失败;此状态必须 fail closed
@@ -23,7 +23,27 @@ pub struct RuntimeSnapshotHandle {
publications: Arc<SyncMutex<Vec<futures::channel::mpsc::Sender<()>>>>,
}
/// 不延长发布集合生命周期的只读句柄,供集合持有的扩展读取后续配置
#[derive(Clone, Default)]
pub struct WeakRuntimeSnapshotHandle {
current: Weak<RwLock<Option<Arc<RuntimeSnapshot>>>>,
}
impl WeakRuntimeSnapshotHandle {
pub fn acquire(&self) -> Result<Arc<RuntimeSnapshot>, RuntimeSnapshotUnavailable> {
let current = self.current.upgrade().ok_or(RuntimeSnapshotUnavailable)?;
acquire_current(&current)
}
}
impl RuntimeSnapshotHandle {
#[must_use]
pub fn downgrade(&self) -> WeakRuntimeSnapshotHandle {
WeakRuntimeSnapshotHandle {
current: Arc::downgrade(&self.current),
}
}
#[must_use]
pub fn new(initial: RuntimeSnapshot) -> Self {
Self {
@@ -86,10 +106,7 @@ impl RuntimeSnapshotHandle {
/// 冻结当前 Arc;后续发布不改变已经开始的请求
pub fn acquire(&self) -> Result<Arc<RuntimeSnapshot>, RuntimeSnapshotUnavailable> {
read_unpoisoned(&self.current)
.clone()
.filter(|snapshot| snapshot.extensions().is_none_or(|set| set.can_serve()))
.ok_or(RuntimeSnapshotUnavailable)
acquire_current(&self.current)
}
}
@@ -109,6 +126,15 @@ impl HealthProbe for RuntimeSnapshotHandle {
}
}
fn acquire_current(
current: &RwLock<Option<Arc<RuntimeSnapshot>>>,
) -> Result<Arc<RuntimeSnapshot>, RuntimeSnapshotUnavailable> {
read_unpoisoned(current)
.clone()
.filter(|snapshot| snapshot.extensions().is_none_or(|set| set.can_serve()))
.ok_or(RuntimeSnapshotUnavailable)
}
fn read_unpoisoned<T>(lock: &RwLock<T>) -> std::sync::RwLockReadGuard<'_, T> {
lock.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
@@ -372,8 +372,12 @@ struct RuntimeSnapshotSubscriptionTask {
}
impl DaemonTask for RuntimeSnapshotSubscriptionTask {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
fn run(
&self,
context: crate::task::WorkerDaemonContext,
) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let cancellation = context.cancellation();
let mut retry_delay = INITIAL_BACKOFF;
loop {
if cancellation.is_cancelled() {
@@ -383,10 +387,11 @@ impl DaemonTask for RuntimeSnapshotSubscriptionTask {
let mut subscription = match subscription {
Ok(subscription) => {
retry_delay = INITIAL_BACKOFF;
context.succeeded();
subscription
}
Err(_) => {
wait_or_cancel(&cancellation, retry_delay).await;
wait_or_cancel(cancellation, retry_delay).await;
retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF);
continue;
}
+53 -1
View File
@@ -6,8 +6,11 @@
use std::fmt;
use std::num::NonZeroU64;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use event_listener::Event;
use futures::future::BoxFuture;
use crate::lifecycle::CancellationToken;
@@ -341,11 +344,60 @@ pub trait ScheduledTask: Send + Sync {
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>>;
}
/// 一次 daemon 运行的取消与成功进度,重启时必须创建新上下文
#[derive(Clone)]
pub struct WorkerDaemonContext {
cancellation: CancellationToken,
progress: Arc<DaemonProgress>,
}
#[derive(Default)]
struct DaemonProgress {
completed: AtomicU64,
changed: Event,
}
impl WorkerDaemonContext {
#[must_use]
pub fn new(cancellation: CancellationToken) -> Self {
Self {
cancellation,
progress: Arc::default(),
}
}
#[must_use]
pub const fn cancellation(&self) -> &CancellationToken {
&self.cancellation
}
/// 仅在实际工作或依赖检查成功后通知 Host;进入 run 不代表恢复
pub fn succeeded(&self) {
if !self.cancellation.is_cancelled() {
self.progress.completed.fetch_add(1, Ordering::Release);
self.progress.changed.notify(usize::MAX);
}
}
#[must_use]
pub fn completed_cycles(&self) -> u64 {
self.progress.completed.load(Ordering::Acquire)
}
/// 先注册再检查计数,避免报告与等待交错时丢失成功通知
pub async fn progress_since(&self, completed: u64) {
let changed = self.progress.changed.listen();
if self.completed_cycles() == completed {
changed.await;
}
}
}
/// 长驻任务
///
/// 实现负责连接级重连退避;Host 负责任务退出或 panic 后的重启退避
pub trait DaemonTask: Send + Sync {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>>;
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>>;
}
/// Host 可执行的两种任务形态,无法组合出“守护任务 + 周期调度”等非法状态
@@ -27,6 +27,40 @@ fn elapsed_recovery_window_is_shared_without_waiting_for_real_time() {
assert_eq!(fallback.startup_remaining(), Some(Duration::ZERO));
}
#[test]
fn same_account_retry_preserves_time_and_an_opening_for_rotation() {
use std::{
cell::Cell,
time::{Duration, Instant},
};
thread_local! {
static NOW: Cell<Instant> = Cell::new(Instant::now());
}
fn now() -> Instant {
NOW.get()
}
let request = ConnectionBudget::with_clock(now);
// 后台 WS 先占一次,HTTP 首次建连再占一次,模拟代理持续不响应
let background = request.clone();
assert_eq!(background.begin().unwrap(), Some(Duration::from_secs(30)));
NOW.set(now() + Duration::from_secs(1));
assert_eq!(request.begin().unwrap(), Some(Duration::from_secs(29)));
NOW.set(now() + Duration::from_secs(15));
assert_eq!(request.remaining_openings(), 2);
let retry = request.clone().reserve_next_opening();
assert_eq!(retry.begin().unwrap(), Some(Duration::from_secs(7)));
assert_eq!(request.remaining_openings(), 1);
assert!(retry.begin().is_err());
assert!(!request.exhausted());
NOW.set(now() + Duration::from_secs(7));
// 重选使用原预算,既不继承 A 的局部时间限制,也不能重置总次数
assert_eq!(request.begin().unwrap(), Some(Duration::from_secs(7)));
assert!(request.exhausted());
assert!(background.begin().is_err());
request.complete();
assert_eq!(retry.begin().unwrap(), None);
}
#[test]
fn all_transports_share_one_opening_budget_but_other_requests_do_not() {
let request = ConnectionBudget::default();
@@ -3467,13 +3467,16 @@ fn explicit_429_after_structural_event_should_retry_before_commit() {
fn atomic_response_failed_should_rotate_account_before_stream_commit() {
let operation = generate_operation();
let route_plan = plan(&operation);
let mut failure = atomic_response_failed("response-first-failed", "first account failure");
let mut events = failure.take_atomic_client_events();
events[0].attach_observation(
ProviderResponseObservation::new(UpstreamTransport::new("http_sse").unwrap())
.with_request_id(OpaqueUpstreamValue::new("upstream-failed-attempt")),
);
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_first",
items: vec![Err(atomic_response_failed(
"response-first-failed",
"first account failure",
))],
items: vec![Err(failure.with_atomic_client_events(events))],
},
Script::Stream {
account_id: "acct_second",
@@ -3516,6 +3519,10 @@ fn atomic_response_failed_should_rotate_account_before_stream_commit() {
let state = store.state.lock().expect("store lock");
assert_eq!(state.commits, 1);
assert_eq!(state.intermediate_failures, 1);
assert_eq!(
state.intermediate_request_ids,
[Some("upstream-failed-attempt".to_owned())]
);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
@@ -5831,11 +5838,11 @@ fn provider_default_profile_is_applied_once_and_frozen_across_account_retries()
}
#[test]
fn connection_recovery_keeps_one_account_and_http_transport_until_success() {
fn connection_recovery_rotates_only_after_retrying_the_first_account() {
let operation = generate_operation();
let route_plan = plan(&operation);
let failure = || Script::Stream {
account_id: "acct_only",
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
@@ -5843,11 +5850,10 @@ fn connection_recovery_keeps_one_account_and_http_transport_until_success() {
.with_connection_retry(AttemptTransport::Fallback))],
};
let (coordinator, store, provider) = coordinator(vec![
failure(),
failure(),
failure(),
Script::Stream {
account_id: "acct_only",
account_id: "acct_second",
items: complete_stream(None),
},
]);
@@ -5863,37 +5869,250 @@ fn connection_recovery_keeps_one_account_and_http_transport_until_success() {
block_on(session.collect_uncommitted()).unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 4);
let account = ProviderAccountId::new("acct_only").unwrap();
assert_eq!(contexts.len(), 3);
let first = ProviderAccountId::new("acct_first").unwrap();
assert_eq!(contexts[1].required_account(), Some(&first));
assert!(!contexts[1].excluded_accounts().contains(&first));
assert!(contexts[2].required_account().is_none());
assert!(contexts[2].excluded_accounts().contains(&first));
for context in contexts.iter().skip(1) {
assert_eq!(context.required_account(), Some(&account));
assert_eq!(context.transport(), AttemptTransport::Fallback);
assert!(!context.excluded_accounts().contains(&account));
}
let state = store.state.lock().unwrap();
assert_eq!(state.intermediate_failures, 3);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 3);
assert_eq!(store.state.lock().unwrap().intermediate_failures, 2);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 2);
drop(session);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 4);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 3);
}
#[test]
fn connection_recovery_keeps_one_account_and_http_transport_until_success() {
for unavailable in [
ProviderErrorKind::NoEligibleAccount,
ProviderErrorKind::AccountCapacityUnavailable,
ProviderErrorKind::QuotaExhausted,
ProviderErrorKind::ConcurrencyQueueFull,
ProviderErrorKind::ConcurrencyQueueTimeout,
] {
let operation = generate_operation();
let route_plan = plan(&operation);
let failure = || Script::Stream {
account_id: "acct_only",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
};
let (coordinator, store, provider) = coordinator(vec![
failure(),
failure(),
Script::Error(ProviderError::new(unavailable, UpstreamSendState::NotSent)),
failure(),
Script::Stream {
account_id: "acct_only",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 5);
let account = ProviderAccountId::new("acct_only").unwrap();
assert!(contexts[2].required_account().is_none());
assert!(contexts[2].excluded_accounts().contains(&account));
assert_eq!(contexts[2].transport(), AttemptTransport::Fallback);
for index in [1, 3, 4] {
let context = &contexts[index];
assert_eq!(context.required_account(), Some(&account));
assert_eq!(context.transport(), AttemptTransport::Fallback);
assert!(!context.excluded_accounts().contains(&account));
}
let state = store.state.lock().unwrap();
assert_eq!(state.intermediate_failures, 3);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 3);
drop(session);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 4);
}
}
#[test]
fn connection_rotation_does_not_bypass_selection_replay_prohibition() {
for send_state in [
UpstreamSendState::NotSent,
UpstreamSendState::Sent,
UpstreamSendState::Ambiguous,
] {
let operation = generate_operation();
let route_plan = plan(&operation);
let failure = || Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
};
let selection_error = ProviderError::new(ProviderErrorKind::NoEligibleAccount, send_state);
let selection_error = if send_state == UpstreamSendState::NotSent {
selection_error.with_retry_prohibited()
} else {
selection_error
};
let (coordinator, _, provider) =
coordinator(vec![failure(), failure(), Script::Error(selection_error)]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
assert!(block_on(session.collect_uncommitted()).is_err());
assert_eq!(provider.contexts.lock().unwrap().len(), 3);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 2);
}
}
#[test]
fn connection_recovery_does_not_delay_existing_account_rotation() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
},
Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::RateLimited,
UpstreamSendState::Sent,
)
.with_status(429)
.with_replay_safe())],
},
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 3);
let first = ProviderAccountId::new("acct_first").unwrap();
// 临时建连失败先在原账号恢复,后续明确可重放的限流拒绝仍按已有规则换号
assert_eq!(contexts[1].required_account(), Some(&first));
assert_eq!(contexts[1].transport(), AttemptTransport::Fallback);
assert!(!contexts[1].excluded_accounts().contains(&first));
assert!(contexts[2].required_account().is_none());
assert_eq!(contexts[2].transport(), AttemptTransport::Default);
assert!(contexts[2].excluded_accounts().contains(&first));
let state = store.state.lock().unwrap();
assert_eq!(state.intermediate_failures, 2);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 2);
drop(session);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 3);
}
#[test]
fn connection_recovery_stays_on_account_when_rotation_is_not_allowed() {
for required in [true, false] {
let operation = generate_operation();
let route_plan = if required {
plan(&operation)
} else {
plan_with_policy(
&operation,
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
.with_max_account_rotations(0),
)
};
let failure = || Script::Stream {
account_id: "acct_required",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
};
let (coordinator, _, provider) = coordinator(vec![
failure(),
failure(),
Script::Stream {
account_id: "acct_required",
items: complete_stream(None),
},
]);
let account = ProviderAccountId::new("acct_required").unwrap();
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
required.then_some(account.clone()),
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 3);
assert_eq!(contexts[0].required_account(), required.then_some(&account));
for context in contexts.iter().skip(1) {
assert_eq!(context.required_account(), Some(&account));
assert!(!context.excluded_accounts().contains(&account));
}
}
}
#[test]
fn connection_recovery_exhaustion_preserves_transport_failure() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(
(0..4)
.map(|_| Script::Stream {
account_id: "acct_only",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
})
.collect(),
let mut scripts: Vec<_> = (0..4)
.map(|_| Script::Stream {
account_id: "acct_only",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
})
.collect();
scripts.insert(
2,
Script::Error(ProviderError::new(
ProviderErrorKind::NoEligibleAccount,
UpstreamSendState::NotSent,
)),
);
let (coordinator, store, provider) = coordinator(scripts);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
@@ -5907,7 +6126,7 @@ fn connection_recovery_exhaustion_preserves_transport_failure() {
assert!(
matches!(error, EngineError::Provider(ref error) if error.kind() == ProviderErrorKind::Transport)
);
assert_eq!(provider.contexts.lock().unwrap().len(), 4);
assert_eq!(provider.contexts.lock().unwrap().len(), 5);
assert_eq!(
store.state.lock().unwrap().finalizations[0].attempt_count,
4
@@ -1195,6 +1195,114 @@ fn request_policy_snapshot(
.with_extensions(Some(generation))
}
fn plugin_origin_scope() -> gateway_core::engine::extensions::ExtensionCallScope {
gateway_core::engine::extensions::ExtensionCallScope::default().with_request(Arc::new(
gateway_core::engine::extensions::RequestContext {
method: "POST".into(),
uri: "/plugin/proxy".into(),
http_version: "HTTP/1.1".into(),
headers: vec![gateway_core::engine::middleware::MiddlewareHeader::new(
"User-Agent",
Bytes::from_static(b"fixture-client/494"),
)],
peer_addr: None,
client_ip: Some("192.0.2.49".parse().unwrap()),
},
))
}
#[test]
fn explicit_key_model_execution_preserves_http_origin_and_leaves_background_origin_unknown() {
use gateway_core::engine::nested::{
BoundModelExecutionBinding, BoundModelExecutionRequest, NestedModelExecutionPort,
};
block_on(async {
for has_origin in [true, false] {
let store = Arc::new(TrackingExecutionStore::default());
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(request_policy_snapshot(
super::extensions::reference("origin-fixture"),
&[ProviderKind::new("openai").unwrap()],
)),
store.clone(),
ProviderRegistry::new([Arc::new(ProcessingProvider) as Arc<dyn Provider>]).unwrap(),
Arc::new(Admissions::default()),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
Arc::new(crate::RecordingDiagnostics::default()),
);
let input = request(&service, ClientTransport::InternalPlugin);
let context = NestedModelExecutionPort::bind(
&service,
BoundModelExecutionBinding {
settings: None,
client_key_id: ClientApiKeyId::new("key_start_test").unwrap(),
initiating_plugin_instance_id: "origin-fixture".into(),
cancellation: Default::default(),
extension_scope: if has_origin {
plugin_origin_scope()
} else {
Default::default()
},
},
)
.await
.unwrap();
let mut started = NestedModelExecutionPort::start_bound(
&service,
BoundModelExecutionRequest {
context: context.clone(),
public_model: input.public_model,
operation: input.operation,
metadata: input.metadata,
provider: None,
account: None,
},
)
.await
.unwrap();
let mut child = NestedModelExecutionPort::start(
&service,
NestedModelExecutionRequest {
parent_request_id: started.request_id.clone(),
initiating_plugin_instance_id: "nested-origin-fixture".into(),
public_model: PublicModelId::new("gpt-start").unwrap(),
operation: start_operation(),
metadata: request(&service, ClientTransport::InternalPlugin).metadata,
provider: None,
account: None,
parent_account: None,
},
)
.await
.unwrap();
child.session.collect_uncommitted().await.unwrap();
child.session.commit_downstream(Some(200)).await.unwrap();
child.session.detach_finalize().await;
started.session.collect_uncommitted().await.unwrap();
{
let requests = store.requests.lock().unwrap();
assert_eq!(requests.len(), 2);
for request in requests.iter() {
assert_eq!(
request.client_ip,
has_origin.then(|| "192.0.2.49".parse().unwrap())
);
assert_eq!(
request.user_agent.as_deref(),
has_origin.then_some("fixture-client/494")
);
assert_eq!(
request.client_transport,
ClientTransport::InternalPlugin.as_str()
);
}
}
started.session.detach_finalize().await;
}
});
}
#[test]
fn completed_parent_rejects_new_nested_execution_and_cancels_an_active_child() {
block_on(async {
+61 -2
View File
@@ -3,8 +3,9 @@
use bytes::Bytes;
use gateway_core::event::{
ContentItem, ContentKind, EventSequenceError, EventSequenceValidator, GatewayEvent,
ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta,
ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseMetadata,
ProviderResponseObservation, ProviderResponseTimings, ReasoningDelta, ResponseMeta, TextDelta,
ToolCallDelta,
};
use gateway_core::upstream::UpstreamTransport;
use serde_json::json;
@@ -335,3 +336,61 @@ fn validator_should_accept_reasoning_and_tool_call_stream() {
assert_eq!(validator.finish(), Ok(()));
}
#[test]
fn response_observation_snapshots_preserve_values_without_copying_large_payloads() {
let metadata = format!("{{\"raw\":\"{}\"}}", "a".repeat(32 * 1024 - 10));
let observation =
ProviderResponseObservation::new(UpstreamTransport::new("http_sse").expect("transport"))
.with_provider_metadata(
ProviderResponseMetadata::new(metadata.clone()).expect("metadata"),
)
.with_client_headers(
(0..32)
.map(|_| {
ProviderResponseHeader::new("set-cookie", Bytes::from_static(b"raw-\xff"))
})
.collect(),
)
.with_timings(ProviderResponseTimings {
first_event_ms: Some(17),
..ProviderResponseTimings::default()
});
let allocations = allocation_counter::measure(|| {
for _ in 0..1_000 {
std::hint::black_box(observation.clone());
}
});
// 允许小标量字段的复制,但流式快照不能随大头集合和 metadata 重复分配
assert!(allocations.bytes_total < 256 * 1_000, "{allocations:?}");
assert_eq!(allocations.bytes_current, 0);
let updated = observation
.clone()
.with_client_headers(Vec::new())
.with_provider_metadata(
ProviderResponseMetadata::new("{\"updated\":true}".to_owned()).expect("metadata"),
)
.with_timings(ProviderResponseTimings {
first_event_ms: Some(23),
..ProviderResponseTimings::default()
});
assert_eq!(
observation.provider_metadata().expect("metadata").as_json(),
metadata
);
assert_eq!(observation.client_headers().len(), 32);
assert!(
observation.client_headers().iter().all(|header| {
header.name() == "set-cookie" && header.value().as_ref() == b"raw-\xff"
})
);
assert_eq!(observation.timings().first_event_ms, Some(17));
assert_eq!(
updated.provider_metadata().expect("metadata").as_json(),
"{\"updated\":true}"
);
assert!(updated.client_headers().is_empty());
assert_eq!(updated.timings().first_event_ms, Some(23));
}
@@ -1253,7 +1253,8 @@ fn decompression_setting_should_validate_and_remain_frozen_across_publication()
));
}
let handle = RuntimeSnapshotHandle::new(compile(1, 64 * 1024 * 1024).unwrap());
let frozen = handle.acquire().unwrap();
let reader = handle.downgrade();
let frozen = reader.acquire().unwrap();
handle.publish(compile(2, 128 * 1024 * 1024).unwrap());
assert_eq!(
handle
@@ -1266,6 +1267,13 @@ fn decompression_setting_should_validate_and_remain_frozen_across_publication()
frozen.responses_max_decompressed_body_bytes(),
64 * 1024 * 1024
);
assert_eq!(
reader
.acquire()
.unwrap()
.responses_max_decompressed_body_bytes(),
128 * 1024 * 1024
);
handle.publish(compile(3, 1024 * 1024).unwrap());
assert_eq!(
handle
@@ -1274,6 +1282,10 @@ fn decompression_setting_should_validate_and_remain_frozen_across_publication()
.responses_max_decompressed_body_bytes(),
1024 * 1024
);
handle.suspend();
assert!(reader.acquire().is_err());
drop(handle);
assert!(reader.acquire().is_err());
}
#[test]
@@ -28,7 +28,8 @@ use gateway_core::runtime::{
};
use gateway_core::settings::SettingsValues;
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerKind, WorkerRunnable,
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDaemonContext, WorkerKind,
WorkerRunnable,
};
type ReadGate<T> = Arc<Mutex<Option<oneshot::Receiver<Result<T, SnapshotStoreError>>>>>;
@@ -490,7 +491,7 @@ fn subscription_refresh_should_wait_for_inflight_compile_and_reload_authoritativ
})
.expect("subscription task");
let cancellation = CancellationToken::new();
let mut daemon = task.run(cancellation.clone());
let mut daemon = task.run(WorkerDaemonContext::new(cancellation.clone()));
let mut old = Box::pin(publisher.refresh());
assert!(futures::poll!(old.as_mut()).is_pending());
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
@@ -40,10 +40,13 @@ pub struct HttpRequest {
pub url: String,
pub headers: Vec<(String, Vec<u8>)>,
pub body: Vec<u8>,
/// 调用方从有效配置派生的请求体上限;零只允许空正文
pub maximum_body_bytes: usize,
}
pub struct HttpResponse {
pub status: u16,
pub version: http::Version,
pub headers: Vec<(String, Vec<u8>)>,
pub body: HttpBody,
}
@@ -51,6 +54,7 @@ pub struct HttpResponse {
pub struct HttpBody {
response: hyper::body::Incoming,
pending: Bytes,
trailers: Vec<(String, Vec<u8>)>,
received: usize,
deadline: Instant,
_permit: OwnedSemaphorePermit,
@@ -64,7 +68,8 @@ pub struct HttpError {
pub send_state: UpstreamSendState,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum HttpErrorKind {
InvalidRequest,
AddressDenied,
@@ -191,7 +196,7 @@ impl HttpClient {
scope: Option<&str>,
websocket_key: Option<&str>,
) -> Result<PendingResponse, HttpError> {
if request.body.len() > 16 * 1024 * 1024
if request.body.len() > request.maximum_body_bytes
|| request.headers.len() > 128
|| timeout.is_zero()
|| timeout > Duration::from_secs(120)
@@ -364,10 +369,12 @@ impl PendingResponse {
pub(super) fn into_http(self) -> HttpResponse {
HttpResponse {
status: self.status,
version: self.response.version(),
headers: self.headers,
body: HttpBody {
response: self.response.into_body(),
pending: Bytes::new(),
trailers: Vec::new(),
received: 0,
deadline: self.deadline,
_permit: self.permit,
@@ -377,23 +384,55 @@ impl PendingResponse {
}
impl HttpBody {
/// 取出截至当前读取位置收到的尾部头,保留多值与原始字节
pub fn take_trailers(&mut self) -> Vec<(String, Vec<u8>)> {
std::mem::take(&mut self.trailers)
}
pub async fn read(&mut self, maximum: usize) -> Result<Option<Bytes>, HttpError> {
self.read_until(maximum, self.deadline).await
}
/// 资源 owner 已接管总寿命时,每次读取使用独立预算,取消仍由 owner 回收正文
pub async fn read_with_timeout(
&mut self,
maximum: usize,
timeout: Duration,
) -> Result<Option<Bytes>, HttpError> {
let deadline = Instant::now()
.checked_add(timeout)
.ok_or_else(|| HttpError::sent("read timeout"))?;
self.read_until(maximum, deadline).await
}
async fn read_until(
&mut self,
maximum: usize,
deadline: Instant,
) -> Result<Option<Bytes>, HttpError> {
if maximum == 0 || maximum > 64 * 1024 {
return Err(HttpError::sent("read size"));
}
loop {
// 已缓冲的数据和就绪帧也受同一期限约束,慢读者不能无限续用 HTTP 资源
if Instant::now() >= self.deadline {
// 已缓冲的数据和就绪帧也受本次期限约束,不能绕过 owner 的读取预算
if Instant::now() >= deadline {
return Err(HttpError::sent("response deadline").with_kind(HttpErrorKind::Timeout));
}
if !self.pending.is_empty() {
return Ok(Some(self.pending.split_to(maximum.min(self.pending.len()))));
}
self.pending = match tokio::time::timeout_at(self.deadline, self.response.frame()).await
{
self.pending = match tokio::time::timeout_at(deadline, self.response.frame()).await {
Ok(Some(Ok(frame))) => match frame.into_data() {
Ok(chunk) => chunk,
Err(_) => continue,
Err(frame) => {
let trailers = frame
.into_trailers()
.map_err(|_| HttpError::sent("response frame"))?;
self.trailers.extend(trailers.iter().map(|(name, value)| {
(name.as_str().to_owned(), value.as_bytes().to_vec())
}));
continue;
}
},
Ok(None) => return Ok(None),
Ok(Some(Err(_))) => return Err(HttpError::sent("response body")),
@@ -27,6 +27,8 @@ type Socket = WebSocketStream<TokioIo<hyper::upgrade::Upgraded>>;
/// 读写分别串行,允许等待输出时发送控制消息;关闭会唤醒在途操作并释放容量
pub struct ManagedWebSocket {
handshake_version: http::Version,
handshake_headers: Vec<(String, Vec<u8>)>,
writer: Mutex<Option<(SplitSink<Socket, Message>, OwnedSemaphorePermit)>>,
reader: Mutex<Option<SplitStream<Socket>>>,
closed: watch::Sender<bool>,
@@ -39,6 +41,7 @@ pub enum WebSocketMessage {
pub struct WebSocketResponse {
pub status: u16,
pub version: http::Version,
pub headers: Vec<(String, Vec<u8>)>,
pub connection: Option<ManagedWebSocket>,
pub body: Option<HttpBody>,
@@ -70,6 +73,7 @@ impl HttpClient {
let response = response.into_http();
return Ok(WebSocketResponse {
status: response.status,
version: response.version,
headers: response.headers,
connection: None,
body: Some(response.body),
@@ -112,8 +116,11 @@ impl HttpClient {
let (writer, reader) = socket.split();
Ok(WebSocketResponse {
status: response.status,
headers: response.headers,
version: response.response.version(),
headers: response.headers.clone(),
connection: Some(ManagedWebSocket {
handshake_version: response.response.version(),
handshake_headers: response.headers,
writer: Mutex::new(Some((writer, response.permit))),
reader: Mutex::new(Some(reader)),
closed: watch::channel(false).0,
@@ -124,6 +131,17 @@ impl HttpClient {
}
impl ManagedWebSocket {
/// 连接复用仍可读取最初握手的原始元数据,不伪造一次新握手
#[must_use]
pub const fn handshake_version(&self) -> http::Version {
self.handshake_version
}
#[must_use]
pub fn handshake_headers(&self) -> &[(String, Vec<u8>)] {
&self.handshake_headers
}
pub async fn send(
&self,
message: WebSocketMessage,
@@ -96,6 +96,7 @@ impl Downloads {
url: url.to_string(),
headers,
body: vec![],
maximum_body_bytes: 0,
},
egress.map(|egress| &egress.proxy),
&self.network,
+2 -1
View File
@@ -33,7 +33,8 @@ pub fn worker(store: Arc<dyn RetentionStore>) -> Result<WorkerContribution, Work
// 保持稳定的租约身份,滚动升级期间仍与已有进程互斥
let id = WorkerId::try_new(WorkerKind::Retention, "postgres")?;
let schedule = WorkerSchedule::try_new(
Duration::from_secs(60 * 60),
// 有界单轮结束后一分钟继续检查,避免日级写入量超过清理容量
Duration::from_secs(60),
Duration::from_secs(1),
Duration::from_secs(60),
Duration::from_secs(15 * 60),
+46 -24
View File
@@ -14,9 +14,9 @@ use gateway_core::health::{
use gateway_core::lifecycle::CancellationToken;
use gateway_core::task::{
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
WorkerDefinitionError, WorkerDisabledReason, WorkerId, WorkerKind, WorkerLeaderLeaseGuard,
WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseRequest, WorkerRegistration,
WorkerRunnable, WorkerSchedule, WorkerTaskError,
WorkerDaemonContext, WorkerDefinitionError, WorkerDisabledReason, WorkerId, WorkerKind,
WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseRequest,
WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError,
};
use tokio::task::JoinHandle;
@@ -235,7 +235,7 @@ pub enum WorkerStartError {
#[derive(Clone)]
struct HealthEntry {
snapshot: WorkerHealthSnapshot,
success_freshness: Duration,
success_freshness: Option<Duration>,
}
#[derive(Default)]
@@ -253,11 +253,14 @@ impl WorkerHealthRegistry {
.values()
.map(|registration| {
let freshness = match &registration.runnable {
WorkerRunnable::Scheduled { schedule, .. } => (*schedule)
.interval()
.saturating_mul(2)
.saturating_add(schedule.maximum_backoff()),
WorkerRunnable::Daemon { restart, .. } => restart.maximum_backoff(),
WorkerRunnable::Scheduled { schedule, .. } => Some(
(*schedule)
.interval()
.saturating_mul(2)
.saturating_add(schedule.maximum_backoff()),
),
// daemon 可长期空闲或执行备份,重启退避不约束工作耗时
WorkerRunnable::Daemon { .. } => None,
};
let key = WorkerHealthKey::Task(registration.id.clone());
(
@@ -275,7 +278,7 @@ impl WorkerHealthRegistry {
key.clone(),
HealthEntry {
snapshot: empty_snapshot(key, WorkerRuntimeState::Disabled),
success_freshness: Duration::ZERO,
success_freshness: None,
},
)
}));
@@ -315,11 +318,11 @@ impl WorkerHealthRegistry {
});
}
fn succeeded(&self, id: &WorkerId) {
fn succeeded(&self, id: &WorkerId, runtime_state: WorkerRuntimeState, cycles: u64) {
self.update(id, |state| {
state.state = WorkerRuntimeState::Idle;
state.state = runtime_state;
state.consecutive_failures = 0;
state.completed_cycles = state.completed_cycles.saturating_add(1);
state.completed_cycles = state.completed_cycles.saturating_add(cycles);
state.last_success_at = Some(SystemTime::now());
state.last_error = None;
});
@@ -361,8 +364,11 @@ impl WorkerHealthSource for WorkerHealthRegistry {
snapshot.state,
WorkerRuntimeState::Running | WorkerRuntimeState::Idle
) && snapshot.last_success_at.is_some_and(|last_success| {
now.duration_since(last_success)
.is_ok_and(|age| age > entry.success_freshness)
now.duration_since(last_success).is_ok_and(|age| {
entry
.success_freshness
.is_some_and(|freshness| age > freshness)
})
});
if stale {
snapshot.state = WorkerRuntimeState::BackingOff;
@@ -407,7 +413,7 @@ async fn supervise_scheduled(
let delay = match outcome {
ScheduledOutcome::Succeeded => {
health.succeeded(&id);
health.succeeded(&id, WorkerRuntimeState::Idle, 1);
backoff = schedule.initial_backoff();
schedule.interval()
}
@@ -595,17 +601,33 @@ async fn supervise_daemon(
break;
}
health.running(&id, None);
let daemon_cancel = CancellationToken::new();
let run = AssertUnwindSafe(task.run(daemon_cancel.clone())).catch_unwind();
let context = WorkerDaemonContext::new(CancellationToken::new());
let run = AssertUnwindSafe(task.run(context.clone())).catch_unwind();
tokio::pin!(run);
let result = tokio::select! {
result = &mut run => Some(result),
() = cancellation.cancelled() => {
daemon_cancel.cancel();
let _ = run.await;
None
let mut completed = 0;
let result = loop {
let outcome = tokio::select! {
biased;
() = cancellation.cancelled() => {
context.cancellation().cancel();
let _ = run.await;
break None;
}
result = &mut run => Some(result),
() = context.progress_since(completed) => None,
};
// 成功与退出可在同一 poll 内发生,先消费进度再记录本次退出故障
let cycles = context.completed_cycles();
if cycles > completed {
health.succeeded(&id, WorkerRuntimeState::Running, cycles - completed);
completed = cycles;
backoff = restart.initial_backoff();
}
if let Some(result) = outcome {
break Some(result);
}
};
context.cancellation().cancel();
let Some(result) = result else { break };
let error = match result {
Ok(Ok(())) => WorkerTaskError::safe("daemon exited unexpectedly"),
@@ -16,6 +16,7 @@ fn request(port: u16, secure: bool) -> HttpRequest {
),
headers: vec![],
body: vec![],
maximum_body_bytes: 0,
}
}
@@ -12,6 +12,7 @@ fn request(url: String) -> HttpRequest {
url,
headers: vec![],
body: vec![],
maximum_body_bytes: 0,
}
}
@@ -221,6 +222,63 @@ async fn buffered_body_cannot_extend_the_request_deadline() {
assert_eq!(error.kind(), HttpErrorKind::Timeout);
}
#[tokio::test]
async fn resource_body_reads_outlive_opening_but_keep_a_per_read_timeout() {
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}", listener.local_addr().unwrap());
let (send_next, next) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut head = Vec::new();
let mut buffer = [0; 1024];
while !head.ends_with(b"\r\n\r\n") {
let count = socket.read(&mut buffer).await.unwrap();
assert!(count > 0 && head.len() < 16 * 1024);
head.extend_from_slice(&buffer[..count]);
}
socket
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 3\r\nConnection: close\r\n\r\na")
.await
.unwrap();
next.await.unwrap();
socket.write_all(b"b").await.unwrap();
assert_eq!(socket.read(&mut buffer).await.unwrap(), 0);
});
let mut response = HttpClient::new()
.unwrap()
.open(request(url), None, &local_network(), Duration::from_secs(1))
.await
.unwrap();
assert_eq!(response.body.read(1).await.unwrap().unwrap(), "a");
tokio::time::pause();
tokio::time::advance(Duration::from_secs(2)).await;
tokio::time::resume();
send_next.send(()).unwrap();
assert_eq!(
response
.body
.read_with_timeout(1, Duration::from_secs(1))
.await
.unwrap()
.unwrap(),
"b"
);
let error = response
.body
.read_with_timeout(1, Duration::from_millis(100))
.await
.unwrap_err();
assert_eq!(error.kind(), HttpErrorKind::Timeout);
assert_eq!(error.send_state, UpstreamSendState::Sent);
drop(response);
tokio::time::timeout(Duration::from_secs(1), server)
.await
.unwrap()
.unwrap();
}
#[tokio::test]
async fn cancelling_dns_releases_the_request_slot_without_starting_http() {
use std::sync::Arc;
@@ -259,3 +317,40 @@ async fn cancelling_dns_releases_the_request_slot_without_starting_http() {
);
}
}
#[tokio::test]
async fn managed_http_enforces_the_callers_body_limit_before_sending() {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(ResponseTemplate::new(200))
.expect(2)
.mount(&server)
.await;
let client = HttpClient::new().unwrap();
// 覆盖旧固定上限之外的正文,并验证同一个客户端逐次使用调用方的预算
let maximum = 17 * 1024 * 1024;
for (size, limit, accepted) in [
(maximum, maximum, true),
(maximum + 1, maximum, false),
(maximum + 1, maximum + 1, true),
] {
let mut request = request(server.uri());
request.method = "POST".into();
request.body = vec![b'x'; size];
request.maximum_body_bytes = limit;
let result = client
.open(request, None, &local_network(), Duration::from_secs(10))
.await;
if accepted {
assert_eq!(result.unwrap().status, 200);
} else {
let error = result.err().unwrap();
assert_eq!(error.kind(), HttpErrorKind::InvalidRequest);
assert_eq!(error.send_state, UpstreamSendState::NotSent);
}
}
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 2);
assert_eq!(requests[0].body.len(), maximum);
assert_eq!(requests[1].body.len(), maximum + 1);
}
@@ -98,6 +98,7 @@ fn request(url: String) -> HttpRequest {
url,
headers: vec![],
body: vec![],
maximum_body_bytes: 0,
}
}
@@ -16,6 +16,7 @@ fn request(url: String) -> HttpRequest {
url,
headers: vec![],
body: vec![],
maximum_body_bytes: 0,
}
}
fn network() -> NetworkPolicy {
+5 -3
View File
@@ -12,7 +12,7 @@ use gateway_admin::ports::system::{
use gateway_core::diagnostics::{OperationalDiagnostics, OperationalFailure};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::task::{
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerId, WorkerKind,
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDaemonContext, WorkerId, WorkerKind,
WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest,
WorkerRegistration, WorkerRunnable, WorkerTaskError,
};
@@ -365,8 +365,9 @@ impl Drop for TaskLifetime {
}
impl DaemonTask for WriterTask {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let cancellation = context.cancellation();
let _lifetime = TaskLifetime(self.probe.dropped.clone());
let mut received = self.received.lock().await;
self.probe.started.cancel();
@@ -394,8 +395,9 @@ impl DaemonTask for WriterTask {
struct IdleTask;
impl DaemonTask for IdleTask {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let cancellation = context.cancellation();
cancellation.cancelled().await;
Ok(())
})
+188 -4
View File
@@ -11,10 +11,10 @@ use futures::future::BoxFuture;
use gateway_core::health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerRuntimeState};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDisabledReason,
WorkerFencingToken, WorkerId, WorkerKind, WorkerLeaderLeaseGuard, WorkerLeaderLeasePort,
WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest, WorkerRegistration,
WorkerRunnable, WorkerSchedule, WorkerTaskError,
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
WorkerDaemonContext, WorkerDisabledReason, WorkerFencingToken, WorkerId, WorkerKind,
WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseError,
WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError,
};
use gateway_host::workers::{WorkerStartError, WorkerSupervisor};
use tokio::sync::Notify;
@@ -1023,3 +1023,187 @@ async fn worker_boundary_records_one_native_failure_and_recovers_health() {
.contains("PRIVATE_WORKER_NATIVE_CAUSE")
);
}
// 每次运行保存自己的上下文,测试可模拟已退出任务的迟到通知
struct ControlledDaemon {
contexts: Arc<Mutex<Vec<WorkerDaemonContext>>>,
commands: tokio::sync::Mutex<tokio::sync::mpsc::UnboundedReceiver<DaemonCommand>>,
}
enum DaemonCommand {
Succeed,
Fail,
SucceedThenFail,
Panic,
Exit,
}
impl DaemonTask for ControlledDaemon {
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
lock_unpoisoned(&self.contexts).push(context.clone());
let mut commands = self.commands.lock().await;
loop {
let command = tokio::select! {
() = context.cancellation().cancelled() => return Ok(()),
command = commands.recv() => command.expect("daemon command"),
};
match command {
DaemonCommand::Succeed => context.succeeded(),
DaemonCommand::SucceedThenFail => {
context.succeeded();
return Err(WorkerTaskError::safe("test daemon failure"));
}
DaemonCommand::Fail => {
return Err(WorkerTaskError::safe("test daemon failure"));
}
DaemonCommand::Panic => panic!("test daemon panic"),
DaemonCommand::Exit => return Ok(()),
}
}
})
}
}
fn daemon_plan(task: impl DaemonTask + 'static) -> Vec<WorkerContribution> {
complete_plan(vec![WorkerContribution::Registration(
WorkerRegistration::try_new(
WorkerId::try_new(WorkerKind::Backup, "backup").unwrap(),
WorkerRunnable::Daemon {
restart: DaemonRestartPolicy::try_new(
Duration::from_millis(10),
Duration::from_millis(20),
)
.unwrap(),
task: Box::new(task),
},
)
.unwrap(),
)])
}
#[tokio::test(start_paused = true)]
async fn daemon_recovers_only_after_current_run_succeeds_and_resets_backoff() {
let (commands, receiver) = tokio::sync::mpsc::unbounded_channel();
let contexts = Arc::new(Mutex::new(Vec::new()));
let supervisor = supervisor();
let id = WorkerId::try_new(WorkerKind::Backup, "backup").unwrap();
supervisor
.start(
daemon_plan(ControlledDaemon {
contexts: contexts.clone(),
commands: tokio::sync::Mutex::new(receiver),
}),
Arc::new(FakeLeasePort::default()),
Arc::new(RecordingDiagnostics::default()),
)
.unwrap();
yield_until(|| lock_unpoisoned(&contexts).len() == 1).await;
// 连续失败逐步退避;仅重新进入 run 不能恢复健康
for (attempt, delay) in [(1, 10), (2, 20)] {
commands.send(DaemonCommand::Fail).unwrap();
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
assert_eq!(task_health(&supervisor, &id).consecutive_failures, attempt);
tokio::time::advance(Duration::from_millis(delay - 1)).await;
assert_eq!(lock_unpoisoned(&contexts).len(), attempt as usize);
tokio::time::advance(Duration::from_millis(1)).await;
yield_until(|| lock_unpoisoned(&contexts).len() == attempt as usize + 1).await;
assert_eq!(task_health(&supervisor, &id).consecutive_failures, attempt);
}
lock_unpoisoned(&contexts)[0].succeeded();
tokio::task::yield_now().await;
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 2);
commands.send(DaemonCommand::Succeed).unwrap();
yield_until(|| task_health(&supervisor, &id).consecutive_failures == 0).await;
let recovered = task_health(&supervisor, &id);
assert_eq!(recovered.state, WorkerRuntimeState::Running);
assert_eq!(recovered.completed_cycles, 1);
assert!(recovered.last_error.is_none());
assert!(recovered.last_success_at.is_some());
assert!(recovered.last_failure_at.is_some());
commands.send(DaemonCommand::Fail).unwrap();
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 1);
tokio::time::advance(Duration::from_millis(10)).await;
yield_until(|| lock_unpoisoned(&contexts).len() == 4).await;
commands.send(DaemonCommand::Fail).unwrap();
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
tokio::time::advance(Duration::from_millis(20)).await;
yield_until(|| lock_unpoisoned(&contexts).len() == 5).await;
// 同一 poll 中先成功再退出,必须保留最新故障并重新从初始退避开始
commands.send(DaemonCommand::SucceedThenFail).unwrap();
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 1);
assert_eq!(task_health(&supervisor, &id).completed_cycles, 2);
tokio::time::advance(Duration::from_millis(10)).await;
yield_until(|| lock_unpoisoned(&contexts).len() == 6).await;
supervisor.shutdown(Duration::from_secs(1)).await;
lock_unpoisoned(&contexts)[5].succeeded();
assert_eq!(
task_health(&supervisor, &id).state,
WorkerRuntimeState::Stopped
);
}
#[tokio::test]
async fn daemon_success_remains_healthy_while_idle_beyond_restart_backoff() {
let (commands, receiver) = tokio::sync::mpsc::unbounded_channel();
let supervisor = supervisor();
let id = WorkerId::try_new(WorkerKind::Backup, "backup").unwrap();
supervisor
.start(
daemon_plan(ControlledDaemon {
contexts: Arc::default(),
commands: tokio::sync::Mutex::new(receiver),
}),
Arc::new(FakeLeasePort::default()),
Arc::new(RecordingDiagnostics::default()),
)
.unwrap();
commands.send(DaemonCommand::Succeed).unwrap();
wait_until(|| task_health(&supervisor, &id).completed_cycles == 1).await;
tokio::time::sleep(Duration::from_millis(50)).await;
let health = task_health(&supervisor, &id);
assert_eq!(health.state, WorkerRuntimeState::Running);
assert_eq!(health.consecutive_failures, 0);
supervisor.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test(start_paused = true)]
async fn daemon_panic_and_unexpected_exit_remain_unhealthy_until_success() {
for (command, error) in [
(DaemonCommand::Panic, "daemon panicked"),
(DaemonCommand::Exit, "daemon exited unexpectedly"),
] {
let (commands, receiver) = tokio::sync::mpsc::unbounded_channel();
let contexts = Arc::new(Mutex::new(Vec::new()));
let supervisor = supervisor();
let id = WorkerId::try_new(WorkerKind::Backup, "backup").unwrap();
supervisor
.start(
daemon_plan(ControlledDaemon {
contexts: contexts.clone(),
commands: tokio::sync::Mutex::new(receiver),
}),
Arc::new(FakeLeasePort::default()),
Arc::new(RecordingDiagnostics::default()),
)
.unwrap();
commands.send(command).unwrap();
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
assert_eq!(
task_health(&supervisor, &id).last_error.as_deref(),
Some(error)
);
tokio::time::advance(Duration::from_millis(10)).await;
yield_until(|| lock_unpoisoned(&contexts).len() == 2).await;
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 1);
commands.send(DaemonCommand::Succeed).unwrap();
yield_until(|| task_health(&supervisor, &id).consecutive_failures == 0).await;
supervisor.shutdown(Duration::from_secs(1)).await;
}
}
@@ -1,7 +1,7 @@
{
"schema_version": 2,
"manifest_schema_versions": [2],
"protocol_versions": [4],
"protocol_versions": [5],
"capabilities": [
{ "capability": "frontend_authentication", "versions": [2] },
{ "capability": "scheduler", "versions": [3] },
@@ -9,7 +9,7 @@
{ "capability": "model_catalog", "versions": [1] },
{ "capability": "retry_policy", "versions": [2] },
{ "capability": "middleware", "versions": [5] },
{ "capability": "upstream_adapter", "versions": [3] },
{ "capability": "upstream_adapter", "versions": [4] },
{ "capability": "observer", "versions": [2] },
{ "capability": "command_line", "versions": [1] },
{ "capability": "management", "versions": [2] },
@@ -69,6 +69,7 @@ pub(super) async fn invoke(
},
maximum_timeout,
);
call_context.resource_stream = true;
call_context.request = crate::adapter::request_context(context.extension_scope().request());
call_context.request_id = Some(context.request_id().as_str().to_owned());
call_context.attempt_id = context
@@ -85,6 +86,7 @@ pub(super) async fn invoke(
context.extension_scope().clone(),
context.plan(),
context.cancellation().clone(),
context.request_body_limit(),
)
.map_err(|_| MiddlewareError::Fault)?;
let _binding = invocation_ports.callbacks.bind_middleware(
@@ -91,6 +91,7 @@ async fn call(
input.extension_scope.clone(),
None,
gateway_core::lifecycle::CancellationToken::new(),
input.request_body_limit(),
)
.map_err(|_| ())?;
let reply = invocation
@@ -137,6 +137,7 @@ async fn route_with_entry(
input.extension_scope().clone(),
None,
gateway_core::lifecycle::CancellationToken::new(),
input.request_body_limit(),
)
.map_err(|_| ())?;
let reply = invocation
@@ -212,6 +213,7 @@ async fn schedule_with_entry(
input.extension_scope().clone(),
None,
gateway_core::lifecycle::CancellationToken::new(),
input.request_body_limit(),
)
.map_err(|_| ())?;
let reply = invocation
@@ -24,7 +24,7 @@ use gateway_protocol::openai::sse::SseEventDecoder;
pub(super) struct EventDecoder {
sequence: EventSequenceValidator,
response_id: Option<String>,
service_tier: Option<String>,
observation: Option<ProviderResponseObservation>,
usage: Usage,
contents: usize,
pub(super) completed: bool,
@@ -54,6 +54,7 @@ impl EventDecoder {
transport: &str,
account: &dyn UpstreamAccountConnection,
sent: UpstreamSendState,
started_at: std::time::Instant,
) -> Result<DecodedEvent, EventDecodeError> {
let message = UpstreamAdapterEvent::decode(payload)
.map_err(|source| invalid(sent).with_source(source))?;
@@ -65,22 +66,31 @@ impl EventDecoder {
.wire
.map(|wire| decode_wire(wire, protocol, sent))
.transpose()?;
let has_observation = self.observation.is_some() || message.observation.is_some();
let mut observation = match self.observation.take() {
Some(observation) => observation,
None => ProviderResponseObservation::new(
UpstreamTransport::new(transport)
.map_err(|source| invalid(sent).with_source(source))?,
),
};
if let Some(update) = message.observation {
observation = super::observation::merge(observation, update, sent)?;
}
if let Some(failure) = message.failure {
if !message.event.facts.is_empty() || message.continuation.is_some() {
return Err(invalid(sent).into());
}
let mut error = failure_error(failure, sent)?;
let mut events = Vec::new();
if has_observation {
events.push(ProviderEvent::observation(observation));
}
if let Some(wire) = wire {
error = error.with_atomic_client_events(vec![ProviderEvent::wire(wire)]);
events.push(ProviderEvent::wire(wire));
}
return Err(EventDecodeError::Upstream(error));
}
if let Some(tier) = message.service_tier {
if tier.is_empty() || tier.len() > 64 || tier.chars().any(char::is_control) {
return Err(invalid(sent).into());
}
// 上游可能在完成时才把 auto 解析为实际档位;与原生观测一样保留最新值
self.service_tier = Some(tier);
return Err(EventDecodeError::Upstream(
failure_error(failure, sent)?.with_atomic_client_events(events),
));
}
let mut facts = Vec::new();
for fact in message.event.facts {
@@ -181,7 +191,7 @@ impl EventDecoder {
return Err(invalid(sent).into());
}
if let Some(cost) =
account.calculate_cost(self.service_tier.as_deref(), &self.usage)
account.calculate_cost(observation.service_tier(), &self.usage)
{
let cost = GatewayEvent::CalculatedCost(cost);
self.sequence
@@ -213,22 +223,26 @@ impl EventDecoder {
if message.continuation.is_some() && !self.completed {
return Err(invalid(sent).into());
}
observation =
super::observation::observe_output(observation, &facts, wire.is_some(), started_at);
self.observation = Some(observation.clone());
let mut event = if let Some(wire) = wire {
ProviderEvent::canonical_with_wire(facts, wire)
} else {
let mut facts = facts.into_iter();
let first = facts.next().ok_or_else(|| invalid(sent))?;
facts.fold(ProviderEvent::canonical(first), ProviderEvent::with_fact)
match facts.next() {
Some(first) => {
facts.fold(ProviderEvent::canonical(first), ProviderEvent::with_fact)
}
None => {
return Ok(DecodedEvent {
event: ProviderEvent::observation(observation),
continuation: message.continuation,
});
}
}
};
if let Some(tier) = &self.service_tier {
let observation = ProviderResponseObservation::new(
UpstreamTransport::new(transport)
.map_err(|source| invalid(sent).with_source(source))?,
)
.try_with_service_tier(tier.clone())
.map_err(|source| invalid(sent).with_source(source))?;
event.attach_observation(observation);
}
event.attach_observation(observation);
Ok(DecodedEvent {
event,
continuation: message.continuation,
@@ -308,6 +322,20 @@ fn decode_wire(
ProtocolWireEvent::json_with_sse_metadata(protocol, event, data, id, retry)
}
}
WirePayload::WebSocket {
event,
data,
message,
} => {
let message =
String::from_utf8(message).map_err(|source| invalid(sent).with_source(source))?;
if let Some(data) = data {
ProtocolWireEvent::json(protocol, event, data)
.map(|wire| wire.with_raw_websocket_message(message))
} else {
ProtocolWireEvent::raw_websocket(protocol, message)
}
}
WirePayload::RawSse { frame } => {
single_sse(&frame, sent)?;
ProtocolWireEvent::raw_sse(protocol, Bytes::from(frame))
@@ -106,6 +106,7 @@ impl Execution {
self.adapter.declaration.transport.as_str(),
self.invocation.account.as_ref(),
active.managed.send_state.get(),
self.invocation.context.timing_started_at(),
)
.map_err(|error| match error {
EventDecodeError::Invalid(error) => {
@@ -131,7 +132,7 @@ impl Execution {
}.await;
if let Err(error) = terminal {
// 尾部故障不得交付成功 wire,但已经解析的计量事实仍归 Core 结算
let metering = decoded
let mut metering: Vec<_> = decoded
.event
.canonical_facts()
.iter()
@@ -145,6 +146,11 @@ impl Execution {
.cloned()
.map(gateway_core::event::ProviderEvent::canonical)
.collect();
if let Some(observation) = decoded.event.response_observation() {
metering.push(gateway_core::event::ProviderEvent::observation(
observation.clone(),
));
}
return Err(error.with_atomic_client_events(metering));
}
self.ended = true;
@@ -177,6 +183,8 @@ impl Execution {
));
}
let mut context = self.adapter.session.context(Stage::Upstream, timeout);
// 初始回复后由模型执行持有流的期限与取消,RPC 预算只约束建立阶段
context.resource_stream = true;
context.request =
crate::adapter::request_context(invocation.context.extension_scope().request());
context.request_id = Some(invocation.context.request_id().as_str().to_owned());
@@ -210,10 +218,21 @@ impl Execution {
.extension_scope()
.extending(self.adapter.instance_id.clone())
.ok_or_else(invalid)?;
let request_body_limit = invocation
.context
.request_body_limit()
.map(Ok)
.unwrap_or_else(|| self.adapter.callbacks.request_body_limit())
.map_err(|error| fault_error(error, UpstreamSendState::NotSent))?;
let scope = self
.adapter
.callbacks
.prepare_upstream(&context, Arc::clone(&managed), extension_scope)
.prepare_upstream(
&context,
Arc::clone(&managed),
extension_scope,
request_body_limit,
)
.map_err(|error| fault_error(error, UpstreamSendState::NotSent))?;
let request = UpstreamAdapterRequest {
adapter_id: declaration.id.clone(),
@@ -230,6 +249,10 @@ impl Execution {
protocol: invocation.operation.protocol().to_owned(),
client_transport: invocation.context.client_transport().as_str().to_owned(),
fast_mode: invocation.context.fast_mode().as_str().to_owned(),
request_elapsed_ms: u64::try_from(
invocation.context.timing_started_at().elapsed().as_millis(),
)
.unwrap_or(u64::MAX),
headers: invocation
.headers
.iter()
@@ -237,14 +260,18 @@ impl Execution {
.collect(),
continuation: previous.map(|state| state.continuation),
};
let body = invocation
.operation
.middleware_body()
.map_err(|source| invalid().with_source(source))?;
if body.len() > gateway_plugin_sdk::call::model::request_payload_limit(request_body_limit) {
return Err(fault_error(
crate::callback::request_limits::too_large(),
UpstreamSendState::NotSent,
));
}
let payload = request
.encode(
invocation
.operation
.middleware_body()
.map_err(|source| invalid().with_source(source))?
.to_vec(),
)
.encode(body.to_vec())
.map_err(|source| invalid().with_source(source))?;
invocation.context.trace().record("plugin.upstream", serde_json::json!({ "instanceId": self.adapter.instance_id, "adapterId": declaration.id, "generation": context.generation, "transport": declaration.transport.as_str() }));
let mut stream = tokio::select! {
@@ -3,6 +3,7 @@
mod event;
mod execution;
mod failure;
mod observation;
mod registration;
pub(crate) mod target;
@@ -0,0 +1,129 @@
//! 合并适配器的原始上游观测,并记录宿主实际收到规范化输出的时间
use gateway_core::event::{
GatewayEvent, ProviderResponseHeader, ProviderResponseMetadata, ProviderResponseObservation,
ProviderResponseTimings, UpstreamHttpVersion, WebSocketPoolKind,
};
use gateway_core::upstream::{OpaqueUpstreamValue, UpstreamSendState};
use gateway_plugin_sdk::call::model::facts::ResponseObservation;
use super::failure::invalid;
pub(super) fn merge(
mut current: ProviderResponseObservation,
update: ResponseObservation,
sent: UpstreamSendState,
) -> Result<ProviderResponseObservation, gateway_core::error::ProviderError> {
// 传输类型由已选适配器固定,不能让不一致的报告绕过 Core 的观测归属
if update.transport != current.transport().as_str() {
return Err(invalid(sent));
}
if let Some(value) = update.http_version {
let version = if value == "unknown" {
UpstreamHttpVersion::Unknown
} else {
UpstreamHttpVersion::parse(&value).ok_or_else(|| invalid(sent))?
};
current = current.with_http_version(version);
}
if let Some(value) = update.websocket_pool {
current = current.with_websocket_pool(match value.as_str() {
"new" => WebSocketPoolKind::New,
"reuse" => WebSocketPoolKind::Reuse,
_ => return Err(invalid(sent)),
});
}
if let Some(value) = update.status_code {
if !(100..=599).contains(&value) {
return Err(invalid(sent));
}
current = current.with_status_code(value);
}
if let Some(value) = update.request_id {
current = current.with_request_id(OpaqueUpstreamValue::new(value));
}
if let Some(value) = update.service_tier {
current = current
.try_with_service_tier(value)
.map_err(|source| invalid(sent).with_source(source))?;
}
if let Some(value) = update.upstream_response_model {
current = current.with_upstream_response_model_if_valid(&value);
if current.upstream_response_model() != Some(value.trim()) {
return Err(invalid(sent));
}
}
if !update.client_headers.is_empty() {
current = current.with_client_headers(
update
.client_headers
.into_iter()
.map(|header| ProviderResponseHeader::new(header.name, header.value.into()))
.collect(),
);
}
if let Some(value) = update.provider_metadata {
current = current.with_provider_metadata(
ProviderResponseMetadata::new(value).ok_or_else(|| invalid(sent))?,
);
}
let previous = current.timings();
let update = update.timings;
let timings = ProviderResponseTimings {
transport_decision_wait_ms: update
.transport_decision_wait_ms
.or(previous.transport_decision_wait_ms),
connect_ms: update.connect_ms.or(previous.connect_ms),
headers_ms: update.headers_ms.or(previous.headers_ms),
first_event_ms: update.first_event_ms.or(previous.first_event_ms),
first_reasoning_ms: update.first_reasoning_ms.or(previous.first_reasoning_ms),
first_text_ms: update.first_text_ms.or(previous.first_text_ms),
first_token_ms: update.first_token_ms.or(previous.first_token_ms),
provider_processing_ms: update
.provider_processing_ms
.or(previous.provider_processing_ms),
upstream_response_ms: update
.upstream_response_ms
.or(previous.upstream_response_ms),
upstream_api_overhead_ms: update
.upstream_api_overhead_ms
.or(previous.upstream_api_overhead_ms),
upstream_engine_ms: update.upstream_engine_ms.or(previous.upstream_engine_ms),
};
Ok(current.with_timings(timings))
}
pub(super) fn observe_output(
current: ProviderResponseObservation,
facts: &[GatewayEvent],
has_wire: bool,
started_at: std::time::Instant,
) -> ProviderResponseObservation {
let elapsed = u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX);
let mut timings = current.timings();
if has_wire || !facts.is_empty() {
timings.first_event_ms.get_or_insert(elapsed);
}
// 内容边界来自插件报告的 canonical 事实,不按上游私有事件名猜测
for fact in facts {
let has_output = match fact {
GatewayEvent::TextDelta(delta) if !delta.text.is_empty() => {
timings.first_text_ms.get_or_insert(elapsed);
true
}
GatewayEvent::ReasoningDelta(delta) if !delta.text.is_empty() => {
timings.first_reasoning_ms.get_or_insert(elapsed);
true
}
GatewayEvent::ToolCallDelta(delta) => {
!delta.arguments_delta.is_empty()
|| delta.name.as_ref().is_some_and(|name| !name.is_empty())
}
_ => false,
};
if has_output {
timings.first_token_ms.get_or_insert(elapsed);
}
}
current.with_timings(timings)
}
@@ -18,7 +18,7 @@ use gateway_plugin_sdk::{
FinishReason as WireFinishReason, Usage as WireUsage, WireEvent, WirePayload,
facts::{self, ExecutionCost, ExecutionFacts},
},
observation::{RequestMoney, RequestTimings},
observation::RequestMoney,
},
};
@@ -120,7 +120,7 @@ fn observation(value: &ProviderResponseObservation) -> facts::ResponseObservatio
request_id: value.request_id().map(|value| value.as_str().to_owned()),
service_tier: value.service_tier().map(str::to_owned),
upstream_response_model: value.upstream_response_model().map(str::to_owned),
timings: RequestTimings {
timings: facts::ResponseTimings {
transport_decision_wait_ms: timings.transport_decision_wait_ms,
connect_ms: timings.connect_ms,
headers_ms: timings.headers_ms,
@@ -132,7 +132,6 @@ fn observation(value: &ProviderResponseObservation) -> facts::ResponseObservatio
upstream_response_ms: timings.upstream_response_ms,
upstream_api_overhead_ms: timings.upstream_api_overhead_ms,
upstream_engine_ms: timings.upstream_engine_ms,
latency_ms: None,
},
client_headers: value
.client_headers()
@@ -271,6 +270,12 @@ fn project_wire(wire: &ProtocolWireEvent) -> Result<WireEvent, PluginFault> {
WirePayload::RawJson {
body: body.to_vec(),
}
} else if let Some(message) = wire.raw_websocket_message() {
WirePayload::WebSocket {
event: wire.event_type().map(str::to_owned),
data: wire.has_json_data().then(|| wire.data().clone()),
message: message.as_bytes().to_vec(),
}
} else if wire.has_json_data() {
WirePayload::Json {
event: wire.event_type().map(str::to_owned),
@@ -283,15 +288,6 @@ fn project_wire(wire: &ProtocolWireEvent) -> Result<WireEvent, PluginFault> {
WirePayload::RawSse {
frame: frame.to_vec(),
}
} else if let Some(message) = wire.raw_websocket_message() {
// 子模型回调采用 SDK 的事件流表达,不把 WebSocket 文本冒充 HTTP body
WirePayload::RawSse {
frame: gateway_protocol::openai::sse::encode_sse_event(
wire.event_type().unwrap_or_default(),
message,
)
.into_bytes(),
}
} else {
return Err(PluginFault::new(
ErrorCode::Fault,
@@ -17,9 +17,18 @@ pub(super) struct HttpCallbacks<'a> {
pub(super) scope: &'a Arc<CallbackScope>,
pub(super) call: &'a CallResources,
pub(super) maximum_payload: usize,
pub(super) request_limits: &'a super::request_limits::RequestLimits,
}
impl HttpCallbacks<'_> {
pub(super) fn body_limit(&self, body: &[u8]) -> Result<usize, PluginFault> {
// 空正文不占请求预算,也允许尚未发布设置的注册阶段执行无正文查询
if body.is_empty() {
return Ok(0);
}
self.request_limits.body_bytes(self.call)
}
pub(super) async fn send(
&self,
request: HttpRequest,
@@ -32,7 +41,11 @@ impl HttpCallbacks<'_> {
scope,
call,
maximum_payload,
..
} = *self;
if request.body.len() > request.maximum_body_bytes {
return Err(super::request_limits::too_large());
}
let timeout = call.timeout()?;
let attempt = call.start_upstream(upstream_purpose)?;
let response = client
@@ -45,16 +58,12 @@ impl HttpCallbacks<'_> {
);
let response = response.map_err(http_error)?;
let mut body = response.body;
let headers = response
.headers
.into_iter()
.map(|(name, value)| String::from_utf8(value).map(|value| (name, value)))
.collect::<Result<Vec<_>, _>>()
.map_err(|_| sent_fault(ErrorCode::Fault, "HTTP header cannot be represented"))?;
let mut result = wire::HttpResponse {
status: response.status,
headers,
http_version: format!("{:?}", response.version),
headers: response.headers,
stream: None,
trailers: Vec::new(),
};
let mut payload = Vec::new();
if streaming {
@@ -90,6 +99,7 @@ impl HttpCallbacks<'_> {
}
payload.extend_from_slice(&chunk);
}
result.trailers = body.take_trailers();
}
Ok(RpcReply {
result: serde_json::to_value(result).map_err(|_| invalid())?,
@@ -111,11 +121,8 @@ impl HttpCallbacks<'_> {
HttpRequest {
method: request.method,
url: request.url,
headers: request
.headers
.into_iter()
.map(|(name, value)| (name, value.into_bytes()))
.collect(),
headers: request.headers,
maximum_body_bytes: self.body_limit(&payload)?,
body: payload,
},
method == "host.http.do_stream",
@@ -142,12 +149,16 @@ impl HttpCallbacks<'_> {
let mut body = stream.body.try_lock().map_err(|_| {
PluginFault::new(ErrorCode::Capacity, "HTTP stream already has a reader")
})?;
let timeout = call.timeout()?;
let mut closed = stream.closed.subscribe();
let result = tokio::select! {
biased;
_ = closed.wait_for(|closed| *closed) => Err(denied()),
result = body.as_mut().ok_or_else(denied)?.read(request.maximum_bytes as usize) => result.map_err(http_error),
result = body.as_mut().ok_or_else(denied)?.read_with_timeout(request.maximum_bytes as usize, timeout) => result.map_err(http_error),
};
let trailers = body
.as_mut()
.map_or_else(Vec::new, |body| body.take_trailers());
if !matches!(&result, Ok(Some(_))) {
body.take();
stream.closed.send_replace(true);
@@ -159,7 +170,7 @@ impl HttpCallbacks<'_> {
}
let chunk = result?;
Ok(RpcReply {
result: serde_json::json!({"eof":chunk.is_none()}),
result: serde_json::json!({"eof":chunk.is_none(), "trailers": trailers}),
payload: chunk.map_or_else(Vec::new, |chunk| chunk.to_vec()),
})
}
@@ -216,6 +227,10 @@ pub(super) fn http_error(error: HttpError) -> PluginFault {
HttpErrorKind::Response => (ErrorCode::Upstream, "目标网站响应读取失败"),
};
let mut fault = PluginFault::new(code, message);
fault.details = Some(serde_json::json!({
"kind": error.kind(),
"reason": error.reason(),
}));
fault.send_state = match error.send_state {
UpstreamSendState::NotSent => SendState::NotSent,
UpstreamSendState::Sent => SendState::Sent,
@@ -15,6 +15,7 @@ mod log;
mod middleware;
mod model;
pub(crate) mod private_state;
pub(crate) mod request_limits;
mod resources;
mod scope;
pub(crate) mod services;
@@ -80,6 +81,7 @@ impl PluginCallbackPorts {
}
pub(crate) struct PluginCallbacks {
pub(crate) registration_body_limit: Option<usize>,
services: Arc<services::ServicePorts>,
service_actor: gateway_admin::model::usage_data::UsageActor,
accounts: Arc<accounts::PluginAccounts>,
@@ -96,6 +98,7 @@ pub(crate) struct PluginCallbacks {
pending_middleware: Mutex<BTreeMap<String, Arc<dyn MiddlewareCallback>>>,
calls: Mutex<BTreeMap<u64, Arc<CallResources>>>,
maximum_payload: usize,
request_limits: Arc<request_limits::RequestLimits>,
}
impl PluginCallbacks {
@@ -106,8 +109,10 @@ impl PluginCallbacks {
log_slots: Arc<tokio::sync::Semaphore>,
private_state: Arc<private_state::PluginPrivateState>,
ports: PluginCallbackPorts,
request_limits: Arc<request_limits::RequestLimits>,
) -> Result<Self, AdminError> {
Ok(Self {
registration_body_limit: None,
services: ports.services,
service_actor: gateway_admin::model::usage_data::UsageActor::Plugin {
publisher: manifest.publisher.clone(),
@@ -126,7 +131,11 @@ impl PluginCallbacks {
log::PluginLog::new(manifest, log_slots)
.map_err(|_| AdminError::invalid("插件清单身份不合法"))?,
),
models: Arc::new(model::PluginModels::new(ports.models, maximum_payload)),
models: Arc::new(model::PluginModels::new(
ports.models,
request_limits.clone(),
maximum_payload,
)),
private_state,
http: ports.http,
network: Arc::new(ports.network),
@@ -134,6 +143,7 @@ impl PluginCallbacks {
pending_middleware: Mutex::new(BTreeMap::new()),
calls: Mutex::new(BTreeMap::new()),
maximum_payload,
request_limits,
})
}
@@ -181,6 +191,7 @@ impl PluginCallbacks {
extensions: gateway_core::engine::extensions::ExtensionCallScope,
plan: Option<gateway_core::engine::middleware::FrozenMiddlewarePlan>,
cancellation: CancellationToken,
request_body_limit: Option<usize>,
) -> Result<Arc<CallbackScope>, PluginFault> {
if !matches!(
context.stage,
@@ -197,6 +208,7 @@ impl PluginCallbacks {
scope = scope.with_execution_effects(effects);
}
scope.extension_scope = extensions;
scope.request_body_limit = request_body_limit;
scope.origin = plan.map(|plan| InvocationContext {
request_id: context
.request_id
@@ -212,11 +224,16 @@ impl PluginCallbacks {
self.register_scope(context, scope)
}
pub(crate) fn request_body_limit(&self) -> Result<usize, PluginFault> {
self.request_limits.current_body_bytes()
}
pub(crate) fn prepare_upstream(
&self,
context: &CallContext,
managed: Arc<upstream::ManagedUpstream>,
extension_scope: gateway_core::engine::extensions::ExtensionCallScope,
request_body_limit: usize,
) -> Result<Arc<CallbackScope>, PluginFault> {
if context.stage != gateway_plugin_sdk::Stage::Upstream
|| context.account_id.as_deref() != Some(managed.account.account_id().as_str())
@@ -227,6 +244,7 @@ impl PluginCallbacks {
let mut scope = CallbackScope::new(context, managed.account.outbound_proxy().cloned())
.with_upstream(managed);
scope.extension_scope = extension_scope;
scope.request_body_limit = Some(request_body_limit);
self.register_scope(context, scope)
}
@@ -352,7 +370,13 @@ impl CallbackHandler for PluginCallbacks {
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&context.resource_scope_id)
.and_then(Weak::upgrade)
.unwrap_or_else(|| Arc::new(CallbackScope::for_call(context)));
.unwrap_or_else(|| {
let mut scope = CallbackScope::for_call(context);
if context.stage == gateway_plugin_sdk::Stage::Registration {
scope.request_body_limit = self.registration_body_limit;
}
Arc::new(scope)
});
let middleware = self
.pending_middleware
.lock()
@@ -401,6 +425,7 @@ impl CallbackHandler for PluginCallbacks {
let data = self.data.clone();
let keys = self.keys.clone();
let models = self.models.clone();
let request_limits = self.request_limits.clone();
let affinity = self.affinity.clone();
let services = self.services.clone();
let service_actor = self.service_actor.clone();
@@ -514,6 +539,7 @@ impl CallbackHandler for PluginCallbacks {
scope,
call: &call,
maximum_payload,
request_limits: &request_limits,
};
if method.starts_with("host.upstream.") {
return upstream::dispatch(&http, &method, params, payload).await;
@@ -63,13 +63,19 @@ impl PluginModelPortSlot {
pub(super) struct PluginModels {
slot: Arc<PluginModelPortSlot>,
request_limits: Arc<super::request_limits::RequestLimits>,
maximum_payload: usize,
}
impl PluginModels {
pub(super) fn new(slot: Arc<PluginModelPortSlot>, maximum_payload: usize) -> Self {
pub(super) fn new(
slot: Arc<PluginModelPortSlot>,
request_limits: Arc<super::request_limits::RequestLimits>,
maximum_payload: usize,
) -> Self {
Self {
slot,
request_limits,
maximum_payload,
}
}
@@ -121,17 +127,27 @@ impl PluginModels {
"host.model.execute" | "host.model.execute_stream" => {
let mut request: ModelExecuteRequest =
serde_json::from_value(params).map_err(|_| invalid())?;
if payload.is_empty() || payload.len() > self.maximum_payload {
if payload.is_empty() {
return Err(invalid());
}
let check_size = |limit| {
if payload.len() > gateway_plugin_sdk::call::model::request_payload_limit(limit)
{
Err(super::request_limits::too_large())
} else {
Ok(())
}
};
let port = self.slot.upgrade()?;
let stream = method == "host.model.execute_stream";
let started = if let Some(client_key_id) = request.client_key_id.take() {
let bound = self.binding(&port, context, call, client_key_id).await?;
check_size(bound.request_body_limit())?;
port.start_bound(self.bound_request(request, payload, stream, bound)?)
.await
.map_err(super::error::gateway)?
} else {
check_size(self.request_limits.body_bytes(call)?)?;
let request = self.request(context, request, payload, stream)?;
port.start(request).await.map_err(super::error::gateway)?
};
@@ -0,0 +1,49 @@
//! 从当前请求或已发布设置读取正文预算,不持有发布集合的强引用
use super::CallResources;
use gateway_core::runtime::{RuntimeSnapshotHandle, WeakRuntimeSnapshotHandle};
use gateway_plugin_sdk::{ErrorCode, PluginFault};
use std::sync::OnceLock;
#[derive(Default)]
pub(crate) struct RequestLimits {
snapshots: OnceLock<WeakRuntimeSnapshotHandle>,
}
impl RequestLimits {
pub(crate) fn bind(
&self,
snapshots: &RuntimeSnapshotHandle,
) -> Result<(), gateway_admin::model::AdminError> {
self.snapshots
.set(snapshots.downgrade())
.map_err(|_| gateway_admin::model::AdminError::conflict("插件设置读取端口已经绑定"))
}
pub(super) fn body_bytes(&self, call: &CallResources) -> Result<usize, PluginFault> {
if let Some(settings) = call.request_settings() {
return Ok(settings.snapshot().responses_max_decompressed_body_bytes());
}
if let Some(bytes) = call.scope.request_body_limit {
return Ok(bytes);
}
self.current_body_bytes()
}
pub(crate) fn current_body_bytes(&self) -> Result<usize, PluginFault> {
self.snapshots
.get()
.and_then(|snapshots| snapshots.acquire().ok())
.map(|snapshot| snapshot.responses_max_decompressed_body_bytes())
.ok_or_else(|| PluginFault::new(ErrorCode::Fault, "request body limit is unavailable"))
}
}
pub(crate) fn too_large() -> PluginFault {
let mut fault = PluginFault::new(
ErrorCode::InvalidInput,
"request body exceeds the configured limit",
);
fault.http_status = Some(413);
fault
}
@@ -19,6 +19,7 @@ pub(crate) struct InvocationContext {
/// 一次回调的调用链、账号与网络事实,由操作持有到完成
pub(crate) struct CallbackScope {
pub(super) origin: Option<InvocationContext>,
pub(super) request_body_limit: Option<usize>,
stage: Stage,
account_id: Option<String>,
credential_revision: Option<u64>,
@@ -61,6 +62,7 @@ impl CallbackScope {
pub(super) fn new(context: &CallContext, proxy: Option<OutboundProxy>) -> Self {
Self {
origin: None,
request_body_limit: None,
stage: context.stage,
account_id: context.account_id.clone(),
credential_revision: context.credential_revision,
@@ -133,6 +133,7 @@ impl ManagedUpstream {
&self,
request: adapter::UpstreamHttpRequest,
body: Vec<u8>,
maximum_body_bytes: usize,
) -> Result<(HttpRequest, UpstreamPathPurpose), PluginFault> {
let (target, purpose) = self.target.resolve(&request.path, &request.query)?;
if request.headers.len() > 128 {
@@ -158,18 +159,14 @@ impl ManagedUpstream {
(name, value.to_vec())
})
.collect();
headers.extend(
request
.headers
.into_iter()
.map(|(name, value)| (name, value.into_bytes())),
);
headers.extend(request.headers);
Ok((
HttpRequest {
method: request.method,
url: target.into(),
headers,
body,
maximum_body_bytes,
},
purpose,
))
@@ -188,6 +185,7 @@ pub(super) async fn dispatch(
scope,
call,
maximum_payload,
..
} = *http;
let managed = scope.upstream.as_ref().ok_or_else(denied)?;
let timeout = || {
@@ -197,7 +195,8 @@ pub(super) async fn dispatch(
match method {
"host.upstream.http.do" | "host.upstream.http.do_stream" => {
let request = serde_json::from_value(params).map_err(|_| invalid())?;
let (request, purpose) = managed.request(request, payload)?;
let maximum_body_bytes = http.body_limit(&payload)?;
let (request, purpose) = managed.request(request, payload, maximum_body_bytes)?;
http.send(request, method.ends_with("do_stream"), Some(purpose))
.await
}
@@ -227,12 +226,13 @@ pub(super) async fn dispatch(
headers: request.headers,
},
vec![],
0,
)?;
if purpose != UpstreamPathPurpose::Inference {
return Err(denied());
}
let mut socket = managed.socket.try_lock().map_err(|_| denied())?;
if let Some((url, _)) = socket
if let Some((url, connection)) = socket
.as_ref()
.filter(|(_, connection)| !connection.is_closed())
{
@@ -240,7 +240,13 @@ pub(super) async fn dispatch(
return Err(denied());
}
return Ok(RpcReply {
result: serde_json::json!({"status":101,"headers":[],"stream":null}),
result: serde_json::json!({
"status": 101,
"http_version": format!("{:?}", connection.handshake_version()),
"headers": connection.handshake_headers(),
"stream": null,
"trailers": [],
}),
payload: vec![],
});
}
@@ -261,12 +267,8 @@ pub(super) async fn dispatch(
.map_or_else(|error| error.send_state, |_| UpstreamSendState::Sent),
);
let mut response = response.map_err(super::http::http_error)?;
let headers = response
.headers
.into_iter()
.map(|(name, value)| String::from_utf8(value).map(|value| (name, value)))
.collect::<Result<Vec<_>, _>>()
.map_err(|_| invalid())?;
let headers = response.headers;
let mut trailers = Vec::new();
let mut payload = Vec::new();
if let Some(body) = response.body.as_mut() {
while let Some(chunk) = body
@@ -279,6 +281,7 @@ pub(super) async fn dispatch(
}
payload.extend_from_slice(&chunk);
}
trailers = body.take_trailers();
}
if let Some(connection) = response.connection {
*socket = Some((url, Arc::new(connection)));
@@ -286,8 +289,10 @@ pub(super) async fn dispatch(
Ok(RpcReply {
result: serde_json::to_value(wire::HttpResponse {
status: response.status,
http_version: format!("{:?}", response.version),
headers,
stream: None,
trailers,
})
.map_err(|_| invalid())?,
payload,
@@ -21,6 +21,7 @@ impl PluginRuntime {
&self,
instance: PluginInstance,
target_revision: Revision,
request_body_limit: Option<usize>,
) -> Result<Arc<PreparedInstance>, AdminError> {
let mut observer_entries = Vec::new();
let mut commands = Vec::new();
@@ -151,7 +152,7 @@ impl PluginRuntime {
configuration,
contributes: manifest.contributes.clone(),
};
let callbacks = Arc::new(PluginCallbacks::new(
let mut callbacks = PluginCallbacks::new(
&instance,
self.config.rpc_limits.maximum_buffered_body_bytes,
&manifest,
@@ -166,7 +167,10 @@ impl PluginRuntime {
self.affinity_ports.clone(),
self.resource_ports.clone(),
),
)?);
self.request_limits.clone(),
)?;
callbacks.registration_body_limit = request_body_limit;
let callbacks = Arc::new(callbacks);
let session = Arc::new(
RpcSession::start_supervised(
package.clone(),
@@ -7,8 +7,9 @@ use gateway_core::{
lifecycle::CancellationToken,
runtime::RuntimeSnapshotHandle,
task::{
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDefinitionError, WorkerId,
WorkerKind, WorkerRegistration, WorkerRunnable, WorkerTaskError,
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDaemonContext,
WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable,
WorkerTaskError,
},
};
use gateway_plugin_sdk::Stage;
@@ -52,8 +53,9 @@ impl PluginRuntime {
}
impl DaemonTask for MaintenanceTask {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let cancellation = context.cancellation();
let mut publications = self.snapshots.subscribe_publications();
let mut active: Option<Arc<PreparedSet>> = None;
let mut revision = None;
@@ -91,6 +93,9 @@ impl DaemonTask for MaintenanceTask {
}
}
revision = snapshot.as_ref().map(|s| s.revision());
if snapshot.is_some() {
context.succeeded();
}
// 发布变更优先于下一轮维护;不在快照发布锁中等待插件或数据库
tokio::select! {
biased;
@@ -13,7 +13,7 @@ impl PluginRuntime {
.load_instances()
.await
.map_err(|_| AdminError::unavailable("无法读取已安装插件"))?;
let reference = self.prepare_snapshot(snapshot, None).await?;
let reference = self.prepare_snapshot(snapshot, None, None).await?;
let set = self.prepared_set(&reference).await?;
Ok(crate::PluginCommandSession::new(
reference,
@@ -52,6 +52,7 @@ pub struct PluginRuntimeConfig {
/// 只索引有主人的候选;当前代次仅由 Core 的发布视图持有
pub struct PluginRuntime {
pub(super) request_limits: Arc<crate::callback::request_limits::RequestLimits>,
pub(super) service_ports: Arc<crate::callback::services::ServicePorts>,
pub(super) store: Arc<dyn PluginStore>,
pub(super) state: Arc<dyn PluginStateStore>,
@@ -109,6 +110,7 @@ impl PluginRuntime {
processes,
validators: Arc::new(Semaphore::new(MAXIMUM_PARALLEL_PREPARATIONS)),
configuration_schemas: Arc::default(),
request_limits: Arc::default(),
service_ports: Arc::new(crate::callback::services::ServicePorts::default()),
account_ports: Arc::new(PluginAccountPortSlot::new()),
client_key_ports: Arc::new(PluginClientKeyPortSlot::new()),
@@ -138,6 +140,14 @@ impl PluginRuntime {
self.service_ports.bind_http(dispatcher)
}
/// 读取当前发布设置,同时避免回调反向保活自身所在的插件集合
pub fn bind_snapshots(
&self,
snapshots: &gateway_core::runtime::RuntimeSnapshotHandle,
) -> Result<(), AdminError> {
self.request_limits.bind(snapshots)
}
pub fn bind_services(
&self,
registry: &Arc<gateway_admin::public_service::Registry>,
@@ -35,10 +35,13 @@ impl PluginRuntime {
&self,
snapshot: PluginInstanceSnapshot,
required_revision: Option<Revision>,
request_body_limit: Option<usize>,
) -> Result<ExtensionSetReference, AdminError> {
let target_revision = snapshot.config_revision.get();
self.record_preparation(target_revision, PreparationDiagnostic::Preparing);
let result = self.prepare_inner(snapshot, required_revision).await;
let result = self
.prepare_inner(snapshot, required_revision, request_body_limit)
.await;
match &result {
Ok(reference) => self.record_preparation(
target_revision,
@@ -58,6 +61,7 @@ impl PluginRuntime {
&self,
snapshot: PluginInstanceSnapshot,
required_revision: Option<Revision>,
request_body_limit: Option<usize>,
) -> Result<ExtensionSetReference, AdminError> {
if self.shutting_down.load(Ordering::Acquire) {
return Err(AdminError::unavailable("插件运行时正在关闭"));
@@ -124,7 +128,11 @@ impl PluginRuntime {
tokio::time::timeout(
// 沿用准备阶段原有的 30 秒预算;打包校验与能力恢复也在预算内
Duration::from_secs(30),
self.prepare_instance(instance.clone(), snapshot.config_revision),
self.prepare_instance(
instance.clone(),
snapshot.config_revision,
request_body_limit,
),
)
.await
.unwrap_or_else(|_| Err(AdminError::unavailable("插件启动超时")))
@@ -363,15 +371,33 @@ impl PluginPreparation for PluginRuntime {
snapshot: PluginInstanceSnapshot,
) -> Result<ExtensionSetReference, AdminError> {
let required_revision = snapshot.config_revision;
self.prepare_snapshot(snapshot, Some(required_revision))
self.prepare_snapshot(snapshot, Some(required_revision), None)
.await
}
}
impl ExtensionPreparationPort for PluginRuntime {
fn prepare_with_request_body_limit(
&self,
revision: ConfigRevision,
request_body_limit: usize,
) -> BoxFuture<'_, Result<ExtensionSetReference, ExtensionPreparationError>> {
self.prepare_revision(revision, Some(request_body_limit))
}
fn prepare(
&self,
revision: ConfigRevision,
) -> BoxFuture<'_, Result<ExtensionSetReference, ExtensionPreparationError>> {
self.prepare_revision(revision, None)
}
}
impl PluginRuntime {
fn prepare_revision(
&self,
revision: ConfigRevision,
request_body_limit: Option<usize>,
) -> BoxFuture<'_, Result<ExtensionSetReference, ExtensionPreparationError>> {
Box::pin(async move {
let snapshot = self
@@ -382,7 +408,7 @@ impl ExtensionPreparationPort for PluginRuntime {
if snapshot.config_revision.get() != revision.get() {
return Err(ExtensionPreparationError);
}
self.prepare_snapshot(snapshot, None)
self.prepare_snapshot(snapshot, None, request_body_limit)
.await
.map_err(|_| ExtensionPreparationError)
})
@@ -6,6 +6,116 @@ use serde_json::json;
use crate::support::environment::Environment;
#[tokio::test]
async fn model_requests_have_a_body_limit_independent_of_the_response_buffer() {
use gateway_plugin_sdk::call::model::request_payload_limit;
for stream in [false, true] {
let Some(mut environment) = Environment::create_command().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let account = environment.account(None).await;
let key_id = format!("key_{}", uuid::Uuid::new_v4().simple());
environment
.client_key(&key_id, "sk-large-body-fixture")
.await;
let body_path = environment.directory.path().join("model-body.json");
let mut body = serde_json::to_vec(&json!({
"model":crate::support::native::MODEL,
"input":"x".repeat(2 * 1024 * 1024),
}))
.unwrap();
let raw_limit = 2 * 1024 * 1024;
body.resize(request_payload_limit(raw_limit), b' ');
std::fs::write(&body_path, &body).unwrap();
let marker = environment.directory.path().join("model-results.jsonl");
environment
.install_plugin(json!({
"plugin_id":"test.large-model-request",
"command_registration":{"commands":[{"name":"run","description":"执行模型"}]},
"command_nested_model_fixture":{
"request":{
"client_key_id":key_id,
"model":crate::support::native::MODEL,
"protocol":"openai",
"operation":"generate",
"provider":"openai",
"account_id":account.as_str(),
},
"body_file":body_path,
"stream":stream,
"expect_error":true,
},
"command_nested_model_marker":marker,
"command_result":{"stdout":"ok\n","stderr":"","exit_code":0},
}))
.await;
let (runtime, core) = environment.runtime().await;
let snapshots = core.snapshots();
let configure = |bytes: usize| {
let snapshot = snapshots.acquire().unwrap();
let settings = snapshot
.settings()
.clone()
.with_responses_max_decompressed_body_bytes(bytes as u64);
let updated = gateway_core::routing::request_settings::RequestSettings::new(snapshot)
.replace(settings, "fixture")
.unwrap();
snapshots.publish((*updated.snapshot()).clone());
};
configure(raw_limit);
environment.store.start_command_line_writes().unwrap();
let instances = environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap();
let commands = runtime.prepare_command_line().await.unwrap();
let instance = &instances.instances[0];
commands.execute(&instance.id, "run", &[]).await.unwrap();
environment.wait_for_bound_model_requests(&key_id, 1).await;
let first: serde_json::Value = serde_json::from_str(
std::fs::read_to_string(&marker)
.unwrap()
.lines()
.next()
.unwrap(),
)
.unwrap();
assert!(first["request_id"].is_string(), "{first}");
assert!(first["events"].as_u64().unwrap() > 0);
// 合法 JSON 仅多一个空白字节,超限必须在进入 Core 和账本前拒绝
let mut oversized = body.clone();
oversized.push(b'\n');
std::fs::write(&body_path, oversized).unwrap();
commands.execute(&instance.id, "run", &[]).await.unwrap();
let records = std::fs::read_to_string(&marker).unwrap();
let rejected: serde_json::Value =
serde_json::from_str(records.lines().last().unwrap()).unwrap();
assert_eq!(rejected["error"], "invalid_input");
assert_eq!(rejected["http_status"], 413);
assert_eq!(environment.bound_model_requests(&key_id).await.len(), 1);
// 不重新准备插件实例:发布更大的配置后,同一个请求必须恢复成功
configure(raw_limit + 1);
commands.execute(&instance.id, "run", &[]).await.unwrap();
environment.wait_for_bound_model_requests(&key_id, 2).await;
commands.shutdown().await;
environment
.store
.shutdown_command_line_writes()
.await
.unwrap();
drop(core);
drop(runtime);
environment.close().await;
}
}
#[tokio::test]
async fn command_selects_key_at_call_time_and_uses_normal_model_ledger() {
let Some(mut environment) = Environment::create_command().await else {
@@ -277,3 +387,131 @@ async fn command_accounts_domain_does_not_need_separate_login_or_write_grants()
drop(runtime);
environment.close().await;
}
#[tokio::test]
async fn bound_model_body_budget_stays_frozen_until_the_parent_call_ends() {
use gateway_plugin_sdk::call::model::request_payload_limit;
for stream in [false, true] {
for growing in [false, true] {
let Some(mut environment) = Environment::create_command().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let account = environment.account(None).await;
let key_id = format!("key_{}", uuid::Uuid::new_v4().simple());
environment
.client_key(&key_id, "sk-frozen-body-fixture")
.await;
let old_limit = if growing { 64 * 1024 } else { 128 * 1024 };
let new_limit = if growing { 128 * 1024 } else { 64 * 1024 };
let body_path = environment.directory.path().join("model-body.json");
let body = json!({"model":crate::support::native::MODEL,"input":"hello"});
let mut large_body = serde_json::to_vec(&body).unwrap();
large_body.resize(
request_payload_limit(old_limit) + usize::from(growing),
b' ',
);
std::fs::write(&body_path, large_body).unwrap();
let marker = environment.directory.path().join("model-results.jsonl");
let gate = environment.directory.path().join("resume-model");
let request = json!({
"client_key_id":key_id,"model":crate::support::native::MODEL,
"protocol":"openai","operation":"generate","provider":"openai",
"account_id":account.as_str(),
});
environment
.install_plugin(json!({
"plugin_id":"test.frozen-model-request",
"command_registration":{"commands":[{"name":"run","description":"执行模型"}]},
"command_nested_model_fixture":[
{"request":request,"body":body,"stream":stream},
{"request":request,"body_file":body_path,"stream":stream,
"expect_error":true,"wait_for_file":gate},
],
"command_nested_model_marker":marker,
"command_result":{"stdout":"ok\n","stderr":"","exit_code":0},
}))
.await;
let (runtime, core) = environment.runtime().await;
let snapshots = core.snapshots();
let configure = |bytes: usize| {
let snapshot = snapshots.acquire().unwrap();
let settings = snapshot
.settings()
.clone()
.with_responses_max_decompressed_body_bytes(bytes as u64);
let updated =
gateway_core::routing::request_settings::RequestSettings::new(snapshot)
.replace(settings, "fixture")
.unwrap();
snapshots.publish((*updated.snapshot()).clone());
};
configure(old_limit);
environment.store.start_command_line_writes().unwrap();
let instances = environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap();
let commands = runtime.prepare_command_line().await.unwrap();
let instance = &instances.instances[0];
let update = async {
tokio::time::timeout(std::time::Duration::from_secs(10), async {
while std::fs::read_to_string(&marker)
.unwrap_or_default()
.lines()
.count()
< 1
{
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
})
.await
.unwrap();
configure(new_limit);
std::fs::write(&gate, b"resume").unwrap();
};
let (result, ()) = tokio::join!(commands.execute(&instance.id, "run", &[]), update);
result.unwrap();
let expected = if growing { 1 } else { 2 };
environment
.wait_for_bound_model_requests(&key_id, expected)
.await;
assert_eq!(
environment.bound_model_requests(&key_id).await.len(),
expected
);
// 父调用结束后重新绑定,下一次调用才使用新发布的预算
commands.execute(&instance.id, "run", &[]).await.unwrap();
environment.wait_for_bound_model_requests(&key_id, 3).await;
let records = std::fs::read_to_string(&marker)
.unwrap()
.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.collect::<Vec<_>>();
assert_eq!(records.len(), 4);
for (index, rejected) in [(1, growing), (3, !growing)] {
if rejected {
assert_eq!(records[index]["http_status"], 413, "{records:?}");
assert_eq!(records[index]["error"], "invalid_input");
} else {
assert!(records[index]["request_id"].is_string(), "{records:?}");
}
}
assert_eq!(environment.bound_model_requests(&key_id).await.len(), 3);
commands.shutdown().await;
environment
.store
.shutdown_command_line_writes()
.await
.unwrap();
drop(core);
drop(runtime);
environment.close().await;
}
}
}
@@ -16,7 +16,7 @@ use gateway_admin::{
use gateway_core::{
lifecycle::CancellationToken,
policy::ClientApiKeyId,
task::{WorkerContribution, WorkerRunnable},
task::{WorkerContribution, WorkerDaemonContext, WorkerRunnable},
};
use serde_json::{Value, json};
use tokio::{sync::Notify, time::timeout};
@@ -277,7 +277,11 @@ async fn published_maintenance_resets_budget_with_only_budget_permission() {
};
let stop = CancellationToken::new();
let cancellation = stop.clone();
let worker = tokio::spawn(async move { task.run(cancellation).await.unwrap() });
let worker = tokio::spawn(async move {
task.run(WorkerDaemonContext::new(cancellation))
.await
.unwrap()
});
let result: Value = timeout(Duration::from_secs(15), async {
loop {
if let Ok(content) = std::fs::read_to_string(&marker)
@@ -27,6 +27,7 @@ struct Facts(AtomicUsize, AtomicUsize, Option<String>);
fn account() -> AccountRecord {
let at = "2026-01-01T00:00:00Z".parse().unwrap();
AccountRecord {
outbound_proxy_id: None,
outbound_proxy_enabled: true,
id: "acct_facts".into(),
provider_kind: ProviderKind::new("openai").unwrap(),
@@ -9,6 +9,7 @@ use gateway_admin::{
},
ports::plugin_resources::PluginResourceAccess,
};
use gateway_core::task::WorkerDaemonContext;
use gateway_core::{
account::FastMode,
lifecycle::CancellationToken,
@@ -165,7 +166,9 @@ async fn start(
let stop = CancellationToken::new();
let cancellation = stop.clone();
let task = tokio::spawn(async move {
task.run(cancellation).await.unwrap();
task.run(WorkerDaemonContext::new(cancellation))
.await
.unwrap();
});
(runtime, core, resources, stop, task)
}
@@ -48,6 +48,14 @@ async fn setup_with_restart_circuit(
async fn setup_with_contributions_and_restart_circuit(
contributes: Contributions,
restart_circuit: PluginRestartCircuitConfig,
) -> (tempfile::TempDir, Arc<Store>, PluginRuntime) {
setup_with_limits(contributes, restart_circuit, RpcLimits::default()).await
}
async fn setup_with_limits(
contributes: Contributions,
restart_circuit: PluginRestartCircuitConfig,
rpc_limits: RpcLimits,
) -> (tempfile::TempDir, Arc<Store>, PluginRuntime) {
let cache = tempfile::tempdir().unwrap();
let artifact = PackageInspector::new(PackageLimits::default(), "1.0.0".parse().unwrap())
@@ -83,7 +91,7 @@ async fn setup_with_contributions_and_restart_circuit(
cache_directory: cache.path().to_owned(),
host_version: "1.0.0".parse().unwrap(),
package_limits: PackageLimits::default(),
rpc_limits: RpcLimits::default(),
rpc_limits,
restart_circuit,
},
Arc::new(gateway_host::outbound::HttpClient::new().unwrap()),
@@ -102,7 +102,21 @@ async fn reading_websocket_facts_preserves_original_text_and_opaque_messages() {
);
}
assert_eq!(delivered, messages);
assert_eq!(marker_lines(&marker, 4).await.len(), 4);
let snapshots = marker_lines(&marker, 4).await;
for (snapshot, original) in snapshots[1..].iter().zip(&messages) {
let event: gateway_plugin_sdk::call::model::ExecutionEvent =
serde_json::from_value(snapshot.clone()).unwrap();
let gateway_plugin_sdk::call::model::WirePayload::WebSocket { data, message, .. } =
event.wire.unwrap().payload
else {
panic!("WebSocket facts must retain their original transport");
};
assert_eq!(message, original.as_bytes());
assert_eq!(
data.is_some(),
serde_json::from_str::<serde_json::Value>(original).is_ok()
);
}
drop(stream);
runtime.shutdown().await;
}
@@ -214,6 +214,8 @@ async fn nested_sdk_middleware_waits_beyond_budget_and_can_read_response_callbac
.expect("host next must not consume either plugin's processing budget");
let (_, status, _, mut body, _) = response.into_parts();
assert_eq!(status, 200);
// 正文消费可晚于初始处理预算,惰性 body_read 仍须归属原资源流
tokio::time::sleep(Duration::from_millis(900)).await;
// 正文改写还会发起新的 body_read,证明 callback 有效期也随暂停恢复
let frame = body.next_frame().await.unwrap().unwrap();
assert_eq!(
@@ -187,3 +187,164 @@ async fn mixed_mounts_keep_independent_order_and_fresh_continuations() {
drop(generation);
runtime.shutdown().await;
}
#[tokio::test]
async fn data_plane_http_callbacks_keep_the_frozen_body_budget_across_publications() {
use gateway_core::{
engine::policy::RetryFacts,
error::ProviderErrorKind,
policy::ClientPolicy,
routing::{RuntimeSnapshot, request_settings::RequestSettings},
runtime::RuntimeSnapshotHandle,
settings::SettingsValues,
upstream::UpstreamSendState,
};
let snapshot = |limit: u64, revision| {
RuntimeSnapshot::new(
ConfigRevision::new(revision).unwrap(),
SettingsValues::new(1, "smart", Default::default(), None, None)
.with_responses_max_decompressed_body_bytes(limit),
vec![],
vec![],
vec![],
)
.unwrap()
};
for body_bytes in [128, 129] {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(ResponseTemplate::new(200))
.mount(&server)
.await;
let records = tempfile::tempdir().unwrap();
let marker = records.path().join("http-results.jsonl");
let (cache, runtime) = setup(vec![InstanceFixture {
id: "frozen-limits",
configuration: serde_json::json!({
"http_probe_methods":["middleware.handle","policy.route_model","policy.schedule_account","policy.retry_decision"],
"http_probe_url":server.uri(),"http_probe_body_bytes":body_bytes,"http_probe_marker":marker,
"retry_decision":{"decision":"delegate"},
"route_decision":{"decision":"unhandled"},
"schedule_decision":{"decision":"delegate"},
}),
bindings: vec![
binding(MIDDLEWARE_CONTRIBUTION,"request",0,PluginFailurePolicy::Reject),
binding(MIDDLEWARE_CONTRIBUTION,"attempt",0,PluginFailurePolicy::Reject),
binding(MODEL_ROUTER_CONTRIBUTION,"routing",0,PluginFailurePolicy::Reject),
binding(SCHEDULER_CONTRIBUTION,"scheduling",0,PluginFailurePolicy::Reject),
binding("test.example.retryPolicy","retry",0,PluginFailurePolicy::Delegate),
],
}]).await;
let snapshots = RuntimeSnapshotHandle::new(snapshot(128, 1));
runtime.bind_snapshots(&snapshots).unwrap();
let settings = RequestSettings::new(snapshots.acquire().unwrap()).with_execution(
&ClientPolicy::new(
ClientApiKeyId::new("key-middleware").unwrap(),
gateway_core::policy::PlaintextClientApiKey::new("frozen-budget-fixture").unwrap(),
Arc::new(gateway_core::routing::FrozenAccountScope::new(
Arc::default(),
gateway_core::routing::ClientRoutingScope::all_accounts(),
)),
true,
Default::default(),
),
Some(60_000),
);
let generation = prepare(&runtime).await;
let plan = runtime
.execution_registry()
.middleware(&generation)
.unwrap();
let policy = policy_context(&runtime, generation.clone(), "req_frozen_limits")
.with_request_body_limit(128);
for (revision, current_limit) in [(2, 64), (3, 256)] {
snapshots.publish(snapshot(current_limit, revision));
for attempt in [false, true] {
let mut request = MiddlewareRequest::new(
"openai",
middleware_headers(),
Bytes::from_static(br#"{"input":"hello"}"#),
);
let context = if attempt {
attempt_middleware_context(Arc::new(ExecutionEffects::default()))
.with_request_body_limit(Some(128))
} else {
request = request.with_settings(settings.clone());
middleware_context(ClientTransport::HttpJson)
};
let response = plan
.handle(
context,
request,
Downstream {
calls: Arc::default(),
closes: Arc::default(),
reads: Arc::default(),
error: false,
}
.into_next(),
)
.await
.unwrap();
let (_, _, _, mut body, _) = response.into_parts();
while body.next_frame().await.unwrap().is_some() {}
body.close().await;
}
policy
.route_model(
operation(),
PublicModelId::new("public-a").unwrap(),
BTreeSet::from([ProviderKind::new("example").unwrap()]),
)
.await
.unwrap();
policy
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("example").unwrap(),
None,
&[candidate("acct_high", 100)],
&mut selection_context(),
)
.await
.unwrap();
policy
.retry_decision(RetryFacts {
attempt_index: NonZeroU32::MIN,
provider: ProviderKind::new("example").unwrap(),
model: None,
error_kind: ProviderErrorKind::RateLimited,
upstream_status: Some(429),
send_state: UpstreamSendState::NotSent,
remaining_routing_attempts: 1,
remaining_deadline: Duration::from_secs(3),
retry_allowed: true,
})
.await;
}
let records = std::fs::read_to_string(marker)
.unwrap()
.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.collect::<Vec<_>>();
assert_eq!(records.len(), 10, "{records:?}");
let expected = if body_bytes == 128 { 200 } else { 413 };
assert!(
records.iter().all(|record| record["status"] == expected),
"{records:?}"
);
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), if body_bytes == 128 { 10 } else { 0 });
assert!(
received
.iter()
.all(|request| request.body == vec![b'x'; body_bytes])
);
drop(policy);
drop(plan);
drop(generation);
runtime.shutdown().await;
super::super::wait_until_empty(cache.path()).await;
}
}
@@ -9,6 +9,44 @@ use gateway_plugin_sdk::Capability;
use gateway_plugin_sdk::Contributions;
use gateway_plugin_sdk::Stage;
#[tokio::test]
async fn initial_registration_http_uses_the_prepared_settings_without_a_publication() {
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
for limit in [63, 64] {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200))
.expect(u64::from(limit == 64))
.mount(&server)
.await;
let (cache, store, runtime) =
super::super::setup_with_contributions(Contributions::from([
crate::support::contribution(
Capability::ModelCatalog,
vec![Stage::Registration],
vec![],
vec![],
),
]))
.await;
store.snapshot.lock().unwrap().instances[0].configuration = serde_json::json!({
"registration_http": {"params": {"method":"POST", "url":server.uri(), "headers":[]}, "body":"x".repeat(64)},
"model_catalog":{"models":[{"id":"fixture-alias","provider":"openai","model":"fixture-upstream"}]}
});
let prepared = ExtensionPreparationPort::prepare_with_request_body_limit(
&runtime,
ConfigRevision::new(1).unwrap(),
limit,
)
.await
.unwrap();
assert_eq!(prepared.model_aliases().len(), usize::from(limit == 64));
drop(prepared);
runtime.shutdown().await;
super::super::wait_until_empty(cache.path()).await;
}
}
#[tokio::test]
async fn model_catalog_is_frozen_without_entering_the_request_chain() {
let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([
@@ -87,7 +87,7 @@ impl UpstreamAccountConnection for Account {
fn configuration(base: &str, streaming: bool) -> Value {
json!({
"upstream_registration":{"adapters":[{"id":"example","provider":"openai","base_url":format!("{base}/adapter/"),"paths":[{"path":"responses","purpose":"inference"}],"authentication_kinds":["oauth"],"transport":if streaming {"http_sse"} else {"http_json"},"protocol":"openai"}]},
"upstream_callbacks":[{"method":if streaming {"host.upstream.http.do_stream"} else {"host.upstream.http.do"},"params":{"method":"POST","path":"responses","headers":[["content-type","application/json"]]},"body":"request"}],
"upstream_callbacks":[{"method":if streaming {"host.upstream.http.do_stream"} else {"host.upstream.http.do"},"params":{"method":"POST","path":"responses","headers":[["content-type",b"application/json".to_vec()]]},"body":"request"}],
"upstream_events":[
{"event":{"facts":[{"type":"started","id":"response-one","model":"native-model"}]}},
{"event":{"facts":[{"type":"usage","usage":{"input_tokens":7,"output_tokens":2}},{"type":"completed","id":"response-one","model":"native-model","reason":"stop"}]},"wire":if streaming {"sse"} else {"http"},"continuation":{"scope":"persisted","upstream_response_id":"private-upstream-id","state":{"cursor":"opaque"}}}
@@ -101,6 +101,17 @@ async fn setup(
tempfile::TempDir,
Arc<super::Store>,
gateway_plugin_runtime::PluginRuntime,
) {
setup_with_limits(config, gateway_plugin_runtime::RpcLimits::default()).await
}
async fn setup_with_limits(
config: Value,
limits: gateway_plugin_runtime::RpcLimits,
) -> (
tempfile::TempDir,
Arc<super::Store>,
gateway_plugin_runtime::PluginRuntime,
) {
let contributes = Contributions::from([crate::support::contribution(
Capability::UpstreamAdapter,
@@ -109,7 +120,7 @@ async fn setup(
vec!["openai".into()],
)]);
let (cache, store, runtime) =
super::setup_with_contributions_and_restart_circuit(contributes, Default::default()).await;
super::setup_with_limits(contributes, Default::default(), limits).await;
{
let mut snapshot = store.snapshot.lock().unwrap();
let instance = &mut snapshot.instances[0];
@@ -148,6 +159,16 @@ fn context_with_fast_mode(
generation: &ExtensionSetReference,
key: &str,
mode: gateway_core::account::FastMode,
) -> AttemptContext {
context_with_body_limit(runtime, generation, key, mode, 64 * 1024 * 1024)
}
fn context_with_body_limit(
runtime: &gateway_plugin_runtime::PluginRuntime,
generation: &ExtensionSetReference,
key: &str,
mode: gateway_core::account::FastMode,
bytes: usize,
) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
@@ -155,7 +176,8 @@ fn context_with_fast_mode(
ClientApiKeyId::new(key).unwrap(),
)
.with_upstream_adapters(runtime.execution_registry().upstream_adapters(generation))
.with_fast_mode(mode),
.with_fast_mode(mode)
.with_request_body_limit(bytes),
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(5),
AccountSelectionPolicy::new(
@@ -309,6 +331,319 @@ async fn adapter_http_and_sse_are_cold_authenticated_and_settled_once() {
}
}
#[tokio::test]
async fn adapter_sse_outlives_rpc_budget_and_cancellation_closes_its_resources() {
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
for cancel in [false, true] {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let (started, receiving) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = Vec::new();
let mut buffer = [0; 1024];
let head_end = loop {
if let Some(end) = request.windows(4).position(|bytes| bytes == b"\r\n\r\n") {
break end + 4;
}
let count = socket.read(&mut buffer).await.unwrap();
assert!(count > 0 && request.len() < 16 * 1024);
request.extend_from_slice(&buffer[..count]);
};
let length: usize = std::str::from_utf8(&request[..head_end])
.unwrap()
.lines()
.filter_map(|line| line.split_once(':'))
.find(|(name, _)| name.eq_ignore_ascii_case("content-length"))
.unwrap()
.1
.trim()
.parse()
.unwrap();
while request.len() < head_end + length {
let count = socket.read(&mut buffer).await.unwrap();
assert!(count > 0);
request.extend_from_slice(&buffer[..count]);
}
socket
.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n: heartbeat\n\n")
.await
.unwrap();
started.send(()).unwrap();
for _ in 0..6 {
tokio::select! {
result = socket.read(&mut buffer) => {
assert_eq!(result.unwrap(), 0);
return false;
}
() = tokio::time::sleep(Duration::from_millis(200)) => {}
}
if socket.write_all(b": heartbeat\n\n").await.is_err() {
return false;
}
}
socket.shutdown().await.unwrap();
true
});
let mut config = configuration(&base, true);
// 此夹具汇总 HTTP 正文后产出规范事件,不把多个心跳伪装成一个 SSE 帧
config["upstream_events"][1]
.as_object_mut()
.unwrap()
.remove("wire");
let (cache, _, runtime) = setup_with_limits(
config,
gateway_plugin_runtime::RpcLimits {
maximum_call_timeout: Duration::from_millis(500),
..Default::default()
},
)
.await;
let generation =
ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let account = Account::new();
let context = context(&runtime, &generation, "key-one");
let cancellation = context.cancellation().clone();
let mut stream = execute_with_context(account.clone(), context, None);
let consume = async {
let mut completed = 0;
while let Some(event) = stream.next().await {
completed += event?
.canonical_facts()
.iter()
.filter(|fact| matches!(fact, GatewayEvent::Completed(_)))
.count();
}
Ok::<_, ProviderError>(completed)
};
let interrupt = async {
receiving.await.unwrap();
tokio::time::sleep(Duration::from_millis(800)).await;
if cancel {
cancellation.cancel();
}
};
let (result, ()) = tokio::time::timeout(Duration::from_secs(5), async {
tokio::join!(consume, interrupt)
})
.await
.unwrap();
if cancel {
assert_eq!(result.unwrap_err().kind(), ProviderErrorKind::Cancelled);
} else {
assert_eq!(result.unwrap(), 1);
}
drop(stream);
assert_eq!(
tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap(),
!cancel
);
tokio::time::timeout(Duration::from_secs(2), async {
while Arc::strong_count(&account) != 1 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
}
#[tokio::test]
async fn adapter_preserves_parsed_and_opaque_websocket_messages() {
use gateway_plugin_sdk::call::model::{WireEvent, WirePayload};
let messages = [
"{ \"type\": \"future.event\", \"value\": 1e3 }".to_owned(),
"opaque future message".to_owned(),
format!("{{\"value\":{}0{}}}", "[".repeat(140), "]".repeat(140)),
];
let mut config = configuration("http://127.0.0.1:1", false);
config["upstream_callbacks"] = json!([]);
config["upstream_registration"]["adapters"][0]["transport"] = json!("websocket");
let mut events = vec![config["upstream_events"][0].clone()];
for message in &messages {
events.push(json!({"event":{"facts":[],"wire":WireEvent {
protocol: "openai".into(),
payload: WirePayload::WebSocket {
event: None, data: serde_json::from_str(message).ok(), message: message.as_bytes().to_vec(),
},
}}}));
}
events.push(json!({"event":config["upstream_events"][1]["event"].clone()}));
config["upstream_events"] = json!(events);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None);
let mut received = Vec::new();
while let Some(event) = stream.next().await {
if let Some(wire) = event.unwrap().wire_event() {
assert!(wire.raw_sse_frame().is_none());
received.push(wire.raw_websocket_message().unwrap().to_owned());
}
}
assert_eq!(received, messages);
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn managed_http_failure_keeps_host_reason_and_error_kind() {
let mut config = configuration("http://127.0.0.1:1", false);
config["upstream_callbacks"][0]["params"]["method"] = json!("invalid method");
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None);
let error = stream.next().await.unwrap().unwrap_err();
let raw: Value = serde_json::from_str(error.raw_upstream_error().unwrap().as_str()).unwrap();
assert_eq!(raw["fault"]["details"]["kind"], "invalid_request");
assert_eq!(raw["fault"]["details"]["reason"], "method");
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn adapter_preserves_observation_updates_and_measures_only_real_output() {
let mut config = configuration("http://127.0.0.1:1", true);
config["upstream_callbacks"] = json!([]);
config["upstream_events"] = json!([
{"event":{"facts":[{"type":"started","id":"response-one","model":"native-model"},
{"type":"content_added","index":0,"kind":"reasoning"}]},
"observation":{"transport":"http_sse","http_version":"HTTP/2","status_code":200,
"request_id":"upstream-494","upstream_response_model":"real-model",
"client_headers":[{"name":"x-bytes","value":[255]},{"name":"x-bytes","value":[254]}],
"provider_metadata":"{ \"future\": 1e3 }","timings":{"connect_ms":5,"headers_ms":8}}},
{"event":{"facts":[{"type":"reasoning_delta","index":0,"text":""}]}},
{"event":{"facts":[{"type":"reasoning_delta","index":0,"text":"thinking"}]}},
{"event":{"facts":[{"type":"content_added","index":1,"kind":"text"},{"type":"text_delta","index":1,"text":"answer"}]}},
{"event":{"facts":[]},"observation":{"transport":"http_sse","service_tier":"priority",
"timings":{"upstream_response_ms":1000,"upstream_api_overhead_ms":1.25,"upstream_engine_ms":998.75}}},
{"event":{"facts":[{"type":"usage","usage":{"input_tokens":7,"output_tokens":2}},
{"type":"completed","id":"response-one","model":"native-model","reason":"stop"}]},
"observation":{"transport":"http_sse","timings":{"first_token_ms":41}}}
]);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let context = context(&runtime, &generation, "key-one");
tokio::time::sleep(Duration::from_millis(50)).await;
let mut stream = execute_with_context(Account::new(), context, None);
let mut observations = Vec::new();
while let Some(event) = stream.next().await {
observations.push(event.unwrap().response_observation().unwrap().clone());
}
assert_eq!(observations.len(), 6);
assert_eq!(observations[0].timings().first_token_ms, None);
assert_eq!(observations[1].timings().first_token_ms, None);
let first_output = observations[2].timings().first_token_ms.unwrap();
assert!(first_output >= 50, "首字计时必须包括调用前的等待");
assert_eq!(
observations[2].timings().first_reasoning_ms,
Some(first_output)
);
assert!(observations[3].timings().first_text_ms.unwrap() >= first_output);
assert_eq!(observations[3].timings().upstream_response_ms, None);
let last = observations.last().unwrap();
assert_eq!(last.timings().first_token_ms, Some(41));
assert_eq!(last.timings().connect_ms, Some(5));
assert_eq!(last.timings().headers_ms, Some(8));
assert_eq!(last.timings().upstream_response_ms, Some(1000));
assert_eq!(last.timings().upstream_api_overhead_ms, Some(1.25));
assert_eq!(last.timings().upstream_engine_ms, Some(998.75));
assert_eq!(last.status_code(), Some(200));
assert_eq!(last.http_version().unwrap().as_str(), "HTTP/2");
assert_eq!(last.request_id().unwrap().as_str(), "upstream-494");
assert_eq!(last.upstream_response_model(), Some("real-model"));
assert_eq!(last.service_tier(), Some("priority"));
assert_eq!(
last.provider_metadata().unwrap().as_json(),
"{ \"future\": 1e3 }"
);
assert_eq!(last.client_headers().len(), 2);
assert_eq!(last.client_headers()[0].value().as_ref(), [255]);
assert_eq!(last.client_headers()[1].value().as_ref(), [254]);
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn managed_http_preserves_byte_headers_trailers_and_version_for_plugins() {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
for streaming in [false, true] {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let upstream = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = Vec::new();
while !request.windows(4).any(|part| part == b"\r\n\r\n") {
assert!(socket.read_buf(&mut request).await.unwrap() > 0);
}
socket.write_all(b"HTTP/1.1 200 OK\r\nX-Bytes: \xff\r\nX-Bytes: other\r\nTransfer-Encoding: chunked\r\nTrailer: x-tail\r\n\r\n").await.unwrap();
let body = if streaming { "data: {}\n\n" } else { "{}" };
socket
.write_all(format!("{:x}\r\n{body}\r\n", body.len()).as_bytes())
.await
.unwrap();
socket
.write_all(b"0\r\nx-tail: \xfe\r\nx-tail: final\r\n\r\n")
.await
.unwrap();
});
let records = tempfile::tempdir().unwrap();
let marker = records.path().join("http.jsonl");
let mut config = configuration(&base, streaming);
config["upstream_callbacks"][0]["body"] = json!("");
config["upstream_http_marker"] = json!(marker);
let (cache, _, runtime) = setup(config).await;
let generation =
ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None);
while let Some(event) = stream.next().await {
event.unwrap();
}
upstream.await.unwrap();
let observations: Vec<Value> = std::fs::read_to_string(&marker)
.unwrap()
.lines()
.map(|line| serde_json::from_str(line).unwrap())
.collect();
assert_eq!(observations[0]["http_version"], "HTTP/1.1");
let headers = observations[0]["headers"].as_array().unwrap();
assert!(headers.contains(&json!(["x-bytes", [255]])));
assert!(headers.contains(&json!(["x-bytes", b"other".to_vec()])));
let trailers = &observations.last().unwrap()["trailers"];
assert_eq!(
trailers,
&json!([["x-tail", [254]], ["x-tail", b"final".to_vec()]])
);
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
}
#[tokio::test]
async fn adapter_accepts_service_tier_resolved_by_the_terminal_response() {
let server = MockServer::start().await;
@@ -318,8 +653,10 @@ async fn adapter_accepts_service_tier_resolved_by_the_terminal_response() {
.mount(&server)
.await;
let mut config = configuration(&server.uri(), false);
config["upstream_events"][0]["service_tier"] = json!("auto");
config["upstream_events"][1]["service_tier"] = json!("default");
config["upstream_events"][0]["observation"] =
json!({"transport":"http_json","service_tier":"auto"});
config["upstream_events"][1]["observation"] =
json!({"transport":"http_json","service_tier":"default"});
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
@@ -361,6 +698,7 @@ async fn adapter_failure_preserves_upstream_diagnostics_and_native_feedback() {
let mut config = configuration(&server.uri(), false);
config["upstream_events"] = json!([{
"event": {"facts": []},
"observation":{"transport":"http_json","status_code":429,"timings":{"upstream_response_ms":230}},
"failure": {
"kind": "rate_limited", "status": 429, "retry_after_ms": 1500,
"message": "fixture rate limit", "code": "rate_limit_exceeded"
@@ -372,7 +710,7 @@ async fn adapter_failure_preserves_upstream_diagnostics_and_native_feedback() {
.unwrap();
let account = Account::new();
let mut stream = execute(&runtime, &generation, Arc::clone(&account), "key-one", None);
let error = stream.next().await.unwrap().unwrap_err();
let mut error = stream.next().await.unwrap().unwrap_err();
assert_eq!(error.kind(), ProviderErrorKind::RateLimited);
assert_eq!(error.send_state(), UpstreamSendState::Sent);
assert_eq!(error.upstream_status(), Some(429));
@@ -381,6 +719,10 @@ async fn adapter_failure_preserves_upstream_diagnostics_and_native_feedback() {
"rate_limit_exceeded"
);
assert_eq!(error.retry_after(), Some(Duration::from_millis(1500)));
let observations = error.take_atomic_client_events();
let observation = observations[0].response_observation().unwrap();
assert_eq!(observation.status_code(), Some(429));
assert_eq!(observation.timings().upstream_response_ms, Some(230));
let snapshot = error.stable_snapshot();
let raw: Value = serde_json::from_str(snapshot.raw_upstream_error().unwrap().as_str()).unwrap();
assert_eq!(raw["message"], "fixture rate limit");
@@ -436,7 +778,8 @@ async fn adapter_overrides_auth_headers_and_rejects_stale_continuation() {
super::wait_until_empty(cache.path()).await;
let mut config = configuration(&server.uri(), false);
config["upstream_callbacks"][0]["params"]["headers"] = json!([["Authorization", "injected"]]);
config["upstream_callbacks"][0]["params"]["headers"] =
json!([["Authorization", b"injected".to_vec()]]);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
@@ -488,8 +831,13 @@ async fn adapter_does_not_publish_completed_when_rpc_fails_after_terminal_frame(
.contains("fixture terminal fault")
);
assert!(!format!("{error:?}").contains("fixture terminal fault"));
let facts = error
.take_atomic_client_events()
let atomic = error.take_atomic_client_events();
assert!(
atomic
.iter()
.any(|event| event.response_observation().is_some())
);
let facts = atomic
.into_iter()
.flat_map(|event| event.into_parts().0)
.collect::<Vec<_>>();
@@ -578,11 +926,12 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle
let (socket, _) = listener.accept().await.unwrap();
let mut ws = tokio_tungstenite::accept_hdr_async(
socket,
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| {
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, mut response: tokio_tungstenite::tungstenite::handshake::server::Response| {
assert_eq!(
request.headers()["authorization"],
"Bearer fixture-native-token"
);
response.headers_mut().append("x-bytes", http::HeaderValue::from_static("handshake-value"));
Ok(response)
},
)
@@ -597,7 +946,10 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle
}
assert!(ws.next().await.is_none_or(|result| result.is_err()));
});
let records = tempfile::tempdir().unwrap();
let marker = records.path().join("handshake.jsonl");
let mut config = configuration(&base, false);
config["upstream_http_marker"] = json!(marker);
config["upstream_registration"]["adapters"][0]["transport"] = json!("websocket");
config["upstream_events"][1]["continuation"]["scope"] = json!("connection_local");
config["upstream_callbacks"] = json!([
@@ -641,6 +993,22 @@ async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle
event.unwrap();
}
drop(second);
let handshakes: Vec<Value> = std::fs::read_to_string(&marker)
.unwrap()
.lines()
.map(|line| serde_json::from_str::<Value>(line).unwrap())
.filter(|value| value["status"] == 101)
.collect();
assert_eq!(handshakes.len(), 2);
assert_eq!(handshakes[0], handshakes[1]);
assert_eq!(handshakes[0]["http_version"], "HTTP/1.1");
assert!(
handshakes[0]["headers"]
.as_array()
.unwrap()
.contains(&json!(["x-bytes", b"handshake-value".to_vec()]))
);
let error = execute(
&runtime,
&generation,
@@ -1088,3 +1456,120 @@ async fn malformed_adapter_events_fail_only_the_current_stream() {
super::wait_until_empty(cache.path()).await;
}
}
#[tokio::test]
async fn large_adapter_requests_use_the_frozen_budget_for_both_http_capabilities() {
for streaming in [false, true] {
for upstream in [false, true] {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_string(if streaming {
"data: {}\n\n"
} else {
"{}"
}))
.expect(if upstream { 2 } else { 4 })
.mount(&server)
.await;
let mut config = configuration(&server.uri(), streaming);
if !upstream {
let mut inference = config["upstream_callbacks"][0].clone();
inference["body"] = json!("{}");
config["upstream_callbacks"]
.as_array_mut()
.unwrap()
.push(inference);
config["upstream_callbacks"][0]["method"] = json!(if streaming {
"host.http.do_stream"
} else {
"host.http.do"
});
config["upstream_callbacks"][0]["params"] = json!({"method":"POST", "url":format!("{}/auxiliary", server.uri()), "headers":[]});
}
let (cache, _, runtime) = setup(config).await;
let generation =
ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let account = Account::new();
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
json!({"model":"native-model","input":"x".repeat(17 * 1024 * 1024)})
.as_object()
.unwrap()
.clone(),
)
.unwrap(),
));
let expected = operation.middleware_body().unwrap();
for accepted in [true, false, true] {
let context = context_with_body_limit(
&runtime,
&generation,
"key-one",
Default::default(),
expected.len() - usize::from(!accepted),
);
let provider = ProviderKind::new("openai").unwrap();
let model = UpstreamModelId::new("native-model").unwrap();
let adapter = context
.upstream_adapter(&provider, &model)
.unwrap()
.unwrap();
let metadata = ProviderCallMetadata::new(
provider,
model,
account.id.clone(),
UpstreamTransport::new(adapter.transport()).unwrap(),
);
let mut stream = adapter.execute(UpstreamAdapterInvocation {
operation: operation.clone(),
headers: vec![],
context,
metadata,
account: account.clone(),
});
let mut completed = false;
let mut failure = None;
while let Some(event) = stream.next().await {
match event {
Ok(event) => {
completed |= event
.canonical_facts()
.iter()
.any(|fact| matches!(fact, GatewayEvent::Completed(_)))
}
Err(error) => {
failure = Some(error);
break;
}
}
}
if accepted {
assert!(
failure.is_none(),
"streaming={streaming} upstream={upstream}: {failure:?}"
);
assert!(completed);
} else {
assert_eq!(failure.unwrap().send_state(), UpstreamSendState::NotSent);
}
}
let requests = server
.received_requests()
.await
.unwrap()
.into_iter()
.filter(|request| request.body.len() > 1024 * 1024)
.collect::<Vec<_>>();
assert_eq!(requests.len(), 2);
for request in requests {
assert_eq!(request.body.as_slice(), expected.as_ref());
}
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
}
}
@@ -48,7 +48,8 @@ fn host_compatibility_requires_the_trusted_middleware_contract() {
assert!(compatibility.supports_capability("middleware", 5));
assert!(!compatibility.supports_capability("upstream_adapter", 1));
assert!(!compatibility.supports_capability("upstream_adapter", 2));
assert!(compatibility.supports_capability("upstream_adapter", 3));
assert!(!compatibility.supports_capability("upstream_adapter", 3));
assert!(compatibility.supports_capability("upstream_adapter", 4));
assert!(!compatibility.supports_capability("openai", 1));
}
@@ -474,6 +474,7 @@ impl Environment {
Some(runtime.execution_registry()),
Some(runtime.frontend_authentication_registry()),
);
runtime.bind_snapshots(&core.snapshots()).unwrap();
let access = gateway_admin::initialize_plugin_accounts(
super::native::admin_registry(),
self.store.admin_ports().accounts(),
@@ -509,6 +510,7 @@ impl Environment {
.activate()
.await
.unwrap();
runtime.bind_snapshots(&core.snapshots()).unwrap();
runtime
.bind_model_ports(
&core.nested_model_execution_port(),
@@ -284,15 +284,51 @@ impl Peer {
let Some(fixture) = self.configuration.get(fixture_field) else {
return Ok(());
};
for fixture in fixture
.as_array()
.map_or(std::slice::from_ref(fixture), Vec::as_slice)
{
if let Some(gate) = fixture["wait_for_file"].as_str() {
while !std::path::Path::new(gate).exists() {
tokio::time::sleep(Duration::from_millis(5)).await;
}
}
self.execute_nested_model_fixture(parent, fixture, marker_field)
.await?;
}
Ok(())
}
async fn execute_nested_model_fixture(
&self,
parent: u64,
fixture: &Value,
marker_field: &str,
) -> Result<(), PluginFault> {
let body = if let Some(path) = fixture["body_file"].as_str() {
std::fs::read(path).unwrap()
} else {
serde_json::to_vec(&fixture["body"]).unwrap()
};
if fixture["stream"] == true {
let (result, payload) = self
let response = self
.callback_payload(
parent,
"host.model.execute_stream",
fixture["request"].clone(),
serde_json::to_vec(&fixture["body"]).unwrap(),
body,
)
.await?;
.await;
let (result, payload) = match response {
Err(error) if fixture["expect_error"] == true => {
self.append_observation_marker(
marker_field,
&json!({"error":error.code,"http_status":error.http_status}),
);
return Ok(());
}
result => result?,
};
assert!(payload.is_empty());
let stream: gateway_plugin_sdk::call::host::ModelStreamResult =
serde_json::from_value(result).unwrap();
@@ -348,7 +384,7 @@ impl Peer {
parent,
"host.model.execute",
fixture["request"].clone(),
serde_json::to_vec(&fixture["body"]).unwrap(),
body,
)
.await;
match callback {
@@ -366,7 +402,10 @@ impl Peer {
Ok(())
}
Err(error) if fixture["expect_error"] == true => {
self.append_observation_marker(marker_field, &json!({"error":error.code}));
self.append_observation_marker(
marker_field,
&json!({"error":error.code,"http_status":error.http_status}),
);
Ok(())
}
Err(error) => {
@@ -439,6 +478,37 @@ impl Peer {
self.send(Message::End { id, error: None }, vec![]).await;
return;
}
if self.configuration["http_probe_methods"]
.as_array()
.is_some_and(|methods| {
methods
.iter()
.any(|probe| probe.as_str() == Some(method.as_str()))
})
{
let body = vec![
b'x';
self.configuration["http_probe_body_bytes"]
.as_u64()
.unwrap() as usize
];
let response = self
.callback_payload(
id,
"host.http.do",
json!({"method":"POST", "url":self.configuration["http_probe_url"]}),
body,
)
.await;
let status = match response {
Ok((result, _)) => result["status"].as_u64().unwrap(),
Err(error) => u64::from(error.http_status.unwrap()),
};
self.append_observation_marker(
"http_probe_marker",
&json!({"method":method,"status":status}),
);
}
match method.as_str() {
"upstream_adapter.register" => {
self.send(
@@ -1304,6 +1374,20 @@ impl Peer {
return;
}
"plugin.register" => {
if let Some(request) = self.configuration.get("registration_http") {
let result = self
.callback_payload(
id,
"host.http.do",
request["params"].clone(),
request["body"].as_str().unwrap().as_bytes().to_vec(),
)
.await;
if let Err(error) = result {
self.send(Message::Error { id, error }, vec![]).await;
return;
}
}
if self.configuration["exit_during_registration"] == true {
std::process::exit(31);
}
@@ -1598,17 +1682,25 @@ impl Peer {
)
}))
.await?;
for (result, payload) in results {
for (callback, (result, payload)) in callbacks.iter().zip(results) {
self.append_observation_marker("upstream_http_marker", &result);
if let Some(stream) = result["stream"].as_str() {
response_body.clear();
let read_method = if callback["method"] == "host.http.do_stream" {
"host.http.stream_read"
} else {
"host.upstream.http.stream_read"
};
loop {
let (result, chunk) = self
.callback_payload(
id,
"host.upstream.http.stream_read",
read_method,
json!({"stream":stream,"maximum_bytes":65536}),
vec![],
)
.await?;
self.append_observation_marker("upstream_http_marker", &result);
response_body.extend(chunk);
if result["eof"] == true {
break;
@@ -1625,7 +1717,9 @@ impl Peer {
let mut message = UpstreamAdapterEvent::new(
serde_json::from_value::<ExecutionEvent>(event["event"].clone()).unwrap(),
);
message.service_tier = event["service_tier"].as_str().map(str::to_owned);
message.observation = event
.get("observation")
.map(|value| serde_json::from_value(value.clone()).unwrap());
message.continuation = event
.get("continuation")
.map(|value| serde_json::from_value(value.clone()).unwrap());
+3 -3
View File
@@ -8,8 +8,8 @@ SDK 处于实验阶段;能力是否可用取决于宿主支持、清单声明
| --- | --- | --- |
| SDK 包 | `0.1.0` | Rust 开发依赖,独立于网关版本 |
| `manifestVersion` | `2` | 安装清单格式 |
| `package.protocolVersion` | `4` | 宿主与插件的进程通信 |
| `contributes.*.version` | `middleware` 为 `5`,`scheduler` / `upstream_adapter` 为 `3`,其余见声明 | 单项能力合同,见[扩展项声明](docs/manifest.md#扩展项简写) |
| `package.protocolVersion` | `5` | 宿主与插件的进程通信 |
| `contributes.*.version` | `middleware` 为 `5`,`scheduler` 为 `3`,`upstream_adapter` 为 `4`,其余见声明 | 单项能力合同,见[扩展项声明](docs/manifest.md#扩展项简写) |
开发顺序:编写[清单](docs/manifest.md) → 实现[能力处理器](docs/capabilities.md) → 用 [CLI](../../../apps/plugin-cli/README.md)
打包 → [安装、配置与启用](../../../../docs/plugins.md)。完整可运行示例位于独立仓库
@@ -92,7 +92,7 @@ sequenceDiagram
首个流式结果成功返回后,正文、会话及关联回调由资源取消与终态回收,不再沿用初始期限。
普通调用继续使用原期限;资源模式仍保留并发容量、Credit、父调用关联及取消约束
协议 v4 的帧头由大端 `u32` 元数据长度和大端 `u64` 载荷长度组成,随后依次传输 JSON 元数据和二进制载荷。
协议 v5 的帧头由大端 `u32` 元数据长度和大端 `u64` 载荷长度组成,随后依次传输 JSON 元数据和二进制载荷。
元数据保留 64 KiB 上限;正文总量不由插件传输层额外限制。传输内部以 256 KiB 为单次 I/O 块,
不逐块发起 RPC,也不要求插件配置与宿主同步。读取按进度扩展缓冲区;截断正文必须报错,不能作为完整消息交付
@@ -209,9 +209,21 @@ SDK 不自动重试。超时或断连不能证明写入未提交;重试上限
- `host.model.execute[_stream]`:交给 Core 的路由、准入、租约、重试、账本和计费链路执行
- `host.model.stream_read/close`:读取或关闭当前父调用创建的有界流
`ModelExecuteRequest.client_key_id` 可显式选择执行 Key,不与触发插件的父请求 Key 求交。省略时继承现有模型父请求身份;没有模型父请求的调用必须提供 Key。所选 Key 的启用状态、模型规则、账号组、并发和预算继续由 Core 执行
模型请求正文使用独立的二进制载荷,不共用响应缓冲预算。宿主从有效设置
`responses_max_decompressed_body_bytes` 取得原始正文预算 `L`,通过
`call::model::request_payload_limit(L)` 为内部请求预留 JSON 字符串转义和元数据空间
(`2 × L + 64 KiB`,受请求编码的 32 位分段长度约束)。超限返回带 HTTP 413 的
`invalid_input`,不启动模型执行。上游适配器的请求 `GPAQ` 也使用请求预算,响应事件仍受
原有 8 MiB 编码上限、Runtime 缓冲预算及流控窗口限制
模型事件使用 `call::model::ExecutionEvent` 的 `GPE2` 封套,将 canonical 事实、宿主事实和原始 wire 分段编码。`host` 包含费用及定价明细、完整上游观测、会话状态和改写来源;当前 wire 与原始 wire 都保留二进制字节。读取或回传快照不触发第二次计费、续接提交或发送状态更新。
HTTP 和 request 中间件中的模型回调沿用有效请求设置;attempt、路由、调度、重试和上游适配器
沿用该次模型请求冻结的正文预算。显式 Key 模型调用的预检与执行使用同一绑定快照,
同一父调用内复用 Key 时不受后续配置发布影响;父调用结束后重新绑定才读取新配置,无需重新启用插件
`ModelExecuteRequest.client_key_id` 可显式选择执行 Key,不与触发插件的父请求 Key 求交。省略时继承现有模型父请求身份;没有模型父请求的调用必须提供 Key。所选 Key 的启用状态、模型规则、账号组、并发和预算继续由 Core 执行
嵌套调用与显式 Key 调用继承入口快照中的客户端 IP 和 User-Agent;没有 HTTP 入口的后台调用保持来源未知
模型事件使用 `call::model::ExecutionEvent` 的 `GPE2` 封套,将 canonical 事实、宿主事实和原始 wire 分段编码。`host` 包含费用及定价明细、完整上游观测、会话状态和改写来源;当前 wire 与原始 wire 都保留二进制字节;`WirePayload::WebSocket` 独立保留完整文本消息和可选 JSON 视图,不把未知文本改造成 SSE。读取或回传快照不触发第二次计费、续接提交或发送状态更新。
非流式模型结果通过 `ModelEventBatch` 的 `HME1` 有界二进制批次返回;流式句柄只在创建它的父调用内有效。
发起回调的插件实例在子调用中跳过,Core 在单次执行的子调用图内限制递归深度、总数和并发,避免插件调用自身形成环。
显式 Key 的身份快照随父调用保留,每个模型执行独立计时并建立调用图;父调用结束会取消尚未完成的执行,连接空闲不消耗后续执行的期限
@@ -226,10 +238,21 @@ SDK 不自动重试。超时或断连不能证明写入未提交;重试上限
父调用取消或结束时宿主回收未关闭的流,SDK 不后台预读。取消一次 `read()` 等待后,再次读取会继续同一次操作,不丢弃已经返回的分块。
账号上游使用同一正文对象,使用方式见[受管上游](upstream-adapters.md)
请求和响应 `headers` 使用 `(String, Vec<u8>)`,保留重复头与非 UTF-8 值;响应另有 `http_version`。
读取至 EOF 后通过 `body.trailers()` 取得完整尾部头,读取中只包含已收到的部分。
`collect` 只返回正文,需要尾部头时使用 `read`;底层非流式回复携带 `trailers`,流式回复在 `stream_read` 中交付尾部头
`host.http.do/do_stream` 和流读取/关闭方法提供受管 HTTP。正文使用独立二进制载荷,Host 统一执行代理、超时和背压,不按插件身份限制目标地址段。DNS 由宿主解析,连接保留目标 Host 与 TLS 域名校验;重定向不会自动跟随。
受管 HTTP(包括 `host.upstream.http.do[_stream]`)的请求正文按有效设置
`responses_max_decompressed_body_bytes` 检查。HTTP/request 中间件使用有效请求设置;
attempt、路由、调度、重试和上游适配器使用执行冻结预算,其余调用读取当前发布配置。
超限返回带 HTTP 413 的 `invalid_input` 和 `not_sent`,不向目标发送请求。
注册阶段沿用正在编译的配置预算,首次发布前也可发送非空正文。
非空正文无法取得有效配置时拒绝调用;无正文查询不消耗正文预算。
这里检查实际出站正文,内部模型请求的转义余量不增加出站预算。响应缓冲及流窗口不随请求配置扩大
普通调用使用父调用剩余期限;已建立的受管 HTTP 流或 WebSocket 会话按每次网络操作计时,仍受父资源取消约束
期限耗尽返回 `timeout`,解析、连接或响应读取失败返回 `upstream`。错误消息提供分类提示,响应正文仍由插件读取;协议故障不应直接展示内部消息
期限耗尽返回 `timeout`,解析、连接或响应读取失败返回 `upstream`。错误 `details.kind` 和 `details.reason` 保留 Host 的故障分类与原因标记,响应正文仍由插件读取;协议故障不应直接展示内部消息
## 请求扩展
@@ -34,7 +34,7 @@
清单与宿主兼容声明的格式版本为 `2`,进程协议版本为 `4`;清单不接受 `permissions` 字段
RPC v4 的帧布局、编码与会话生命周期见[通信边界](../README.md#通信边界)。插件使用 SDK 构建,并由 CLI 生成包含协议版本的安装包
RPC v5 的帧布局、编码与会话生命周期见[通信边界](../README.md#通信边界)。插件使用 SDK 构建,并由 CLI 生成包含协议版本的安装包
安装包清单校验与已安装元数据读取是两个边界:清单拒绝未知字段,并检查必需字段、类型和版本;
数据库元数据的读取规则见[制品与发布](../../../../../docs/architecture.md#制品配置与发布),不改变清单或 RPC 合同。
@@ -49,7 +49,7 @@ RPC v4 的帧布局、编码与会话生命周期见[通信边界](../README.md#
## 扩展项简写
作者声明可以省略 `id` 和固定阶段。版本字段省略时为 `1`,适用于 `model_catalog`、`command_line` 和 `maintenance`;其余能力显式声明当前版本:`middleware` 为 `5`,`scheduler` 和 `upstream_adapter` 为 `3`,认证、路由、重试、观察和管理能力为 `2`:
作者声明可以省略 `id` 和固定阶段。版本字段省略时为 `1`,适用于 `model_catalog`、`command_line` 和 `maintenance`;其余能力显式声明当前版本:`middleware` 为 `5`,`scheduler` 为 `3`,`upstream_adapter` 为 `4`,认证、路由、重试、观察和管理能力为 `2`:
```json
{
@@ -1,12 +1,12 @@
# 受管上游适配器
`upstream_adapter` v3 是现有生成请求洋葱链的终端,只挂到内置 `openai` / `xai`,不替换独立图像或压缩端点。
`upstream_adapter` v4 是现有生成请求洋葱链的终端,只挂到内置 `openai` / `xai`,不替换独立图像或压缩端点。
插件定义 Base URL、业务路径、请求编码和响应解析;宿主继续持有身份、ACL、选号、租约、凭据、代理、重试和账本。
普通请求或响应加工使用 [middleware](capabilities.md#洋葱中间件)
## 声明与绑定
清单声明 `upstream_adapter`、版本 `3`,以及匹配的 `inputFormats` / `outputFormats`;固定阶段 `upstream` 由作者清单规范化入口生成。
清单声明 `upstream_adapter`、版本 `4`,以及匹配的 `inputFormats` / `outputFormats`;固定阶段 `upstream` 由作者清单规范化入口生成。
使用 `PluginBuilder::on` 注册
`methods::UPSTREAM_ADAPTER_REGISTER` 与 `methods::UPSTREAM_ADAPTER_EXECUTE`
@@ -55,15 +55,15 @@ while let Some(bytes) = body.read().await? {
正文与普通 `HostClient::http` 共用 `read`、`collect(maximum_bytes)` 和 `close`;每次读取至多 64 KiB,
不后台预读,不向作者暴露流 ID。提前关闭及收集超限会释放正文,父调用结束或取消时宿主兜底回收。
JSON 可有界收集,SSE 应增量解析,不能把读取块当成完整事件
JSON 可有界收集,SSE 应增量解析,不能把读取块当成完整事件。HTTP 版本、字节头及 trailers 的读取见[网络](capabilities.md#网络)
WebSocket 使用 `call.host.upstream_websocket(request).await?`,返回
`UpstreamWebSocketUpgrade::Connected { headers, connection }` 或 `Rejected { status, headers, body }`。
`UpstreamWebSocketUpgrade::Connected { http_version, headers, connection }` 或 `Rejected { status, http_version, headers, body, trailers }`。
连接提供 `send(kind, body)`、`read()`、`close()`,支持完整文本与二进制消息,Ping/Pong 由宿主处理。
读写可同时进行,等待上游输出时可以发送控制消息;有先后依赖的消息须顺序等待 `send`。
同一时刻只允许一个读取者;取消一次 `read()` 等待后,再次读取继续同一次接收,不跳过消息。
`close` 会唤醒在途操作并释放连接,并发关闭等待同一次完成;读取超时或连接错误也会关闭连接。
握手拒绝保留 HTTP 状态和正文。底层 `host.upstream.*` 是 SDK 使用的资源协议
握手拒绝保留 HTTP 状态、正文和 trailers;连接复用返回该连接最初的握手版本和响应头。底层 `host.upstream.*` 是 SDK 使用的资源协议
```rust,ignore
let (message, ()) = tokio::try_join!(
@@ -80,12 +80,28 @@ let (message, ()) = tokio::try_join!(
执行返回 `TypedReply<Empty>` 的有界 `ResponseStream`。每个 chunk 使用 `UpstreamAdapterEvent::encode()`,
包含标准事实与可选原始 wire;事件需符合开始、内容、工具、用量和唯一结束的顺序。
原始 JSON / SSE 字节与标准事实分开传递,错误使用 `UpstreamFailure`;宿主网络观测决定发送状态,插件不能将其降为未发送
原始 JSON / SSE / WebSocket 字节与标准事实分开传递,错误使用 `UpstreamFailure`;宿主网络观测决定发送状态,插件不能将其降为未发送
初始 RPC 回复受插件处理预算约束,默认上限为 120 秒;持续接收与解析应放在返回的流中。
初始流式回复后,流的寿命遵循模型请求期限与取消,受管 HTTP 正文每次读取单独计时。
持续有数据的流不会因初始 RPC 预算耗尽而结束,单次读取超时、客户端取消和流结束仍会回收资源
用量提交标准 `Usage`,费用由内置 Provider 按已冻结的上游模型与价格计算,响应回显模型不切换计价。
`service_tier` 记录上游回显的最新档位,允许上游在完成时确定实际值;OpenAI 计价沿用请求策略,回显档位仅作观测。
`observation.service_tier` 记录上游回显的最新档位,允许上游在完成时确定实际值;OpenAI 计价沿用请求策略,回显档位仅作观测。
未知用量或价格保留未知。宿主确认 RPC 正常结束后才发布完成事实,插件不能独立结算或组织换号重试
`UpstreamAdapterEvent.observation` 接收 `ResponseObservation` 原始观测:HTTP 版本、状态、请求 ID、回显模型、
服务档位、字节响应头、Provider JSON 元数据及 `ResponseTimings`。`transport` 与注册传输类型一致;
已提供字段覆盖同次执行的旧值,缺失字段或空 headers 保留已有观测,允许独立的观测事件及失败事件携带观测。
`ExecutionEvent.host` 是宿主结果的只读快照,不能用于回写费用、发送状态或续接
`first_*_ms` 相对 Core 请求计时起点,包含本次请求的排队和重试等待。
`request_elapsed_ms` 提供宿主发送本次适配器调用前已耗的时间,插件自行计时须对齐该基线;
没有报告时 Runtime 在收到首事件、非空正文/推理/工具增量时记录相应时间。
开始、心跳和空增量不构成首个输出,显式报告的计时优先于宿主缺省观测。
`upstream_response_ms`、`upstream_api_overhead_ms`、`upstream_engine_ms` 只能来自上游原始计时,
没有来源时保持缺失。观测合同不包含总耗时、速率或计算费用,这些由宿主业务层消费原始值计算
成功终态可附带 `UpstreamContinuation`,包含上游响应 ID、不超过 32 KiB 的私有续接材料和作用域:
- `persisted`:续接材料随宿主会话保存,仍受实例、代次、Key、账号和凭据版本限制
@@ -63,15 +63,17 @@ pub struct HttpRequest {
pub method: String,
pub url: String,
#[serde(default)]
pub headers: Vec<(String, String)>,
pub headers: Vec<(String, Vec<u8>)>,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct HttpResponse {
pub status: u16,
pub headers: Vec<(String, String)>,
pub http_version: String,
pub headers: Vec<(String, Vec<u8>)>,
pub stream: Option<String>,
pub trailers: Vec<(String, Vec<u8>)>,
}
/// 宿主账号列表查询;cursor 是上次结果返回的稳定账号 ID
@@ -45,29 +45,44 @@ pub(in crate::call) fn pack<T: Serialize, const N: usize>(
prefix: [u8; 4],
metadata: &T,
payloads: [Vec<u8>; N],
) -> Result<Vec<u8>, ExecutionEncodingError> {
pack_with_limit(prefix, metadata, payloads, MAX_EXECUTION_PAYLOAD_BYTES)
}
pub(in crate::call) fn pack_with_limit<T: Serialize, const N: usize>(
prefix: [u8; 4],
metadata: &T,
payloads: [Vec<u8>; N],
maximum: usize,
) -> Result<Vec<u8>, ExecutionEncodingError> {
let header_bytes = 4 + 4 * (N + 1);
let raw_bytes = payloads
.iter()
.fold(0usize, |total, bytes| total.saturating_add(bytes.len()));
if raw_bytes > MAX_EXECUTION_PAYLOAD_BYTES {
.try_fold(0usize, |total, bytes| total.checked_add(bytes.len()))
.ok_or(ExecutionEncodingError)?;
if raw_bytes > maximum {
return Err(ExecutionEncodingError);
}
let metadata = serde_json::to_vec(metadata).map_err(|_| ExecutionEncodingError)?;
let total = header_bytes
.saturating_add(metadata.len())
.saturating_add(raw_bytes);
if total > MAX_EXECUTION_PAYLOAD_BYTES {
.checked_add(metadata.len())
.and_then(|total| total.checked_add(raw_bytes))
.ok_or(ExecutionEncodingError)?;
let lengths = std::iter::once(metadata.len())
.chain(payloads.iter().map(Vec::len))
.map(u32::try_from)
.collect::<Result<Vec<_>, _>>()
.map_err(|_| ExecutionEncodingError)?;
if total > maximum {
return Err(ExecutionEncodingError);
}
let mut result = Vec::with_capacity(total);
let mut result = Vec::new();
result
.try_reserve_exact(total)
.map_err(|_| ExecutionEncodingError)?;
result.extend_from_slice(&prefix);
for length in std::iter::once(metadata.len()).chain(payloads.iter().map(Vec::len)) {
result.extend_from_slice(
&u32::try_from(length)
.map_err(|_| ExecutionEncodingError)?
.to_be_bytes(),
);
for length in lengths {
result.extend_from_slice(&length.to_be_bytes());
}
result.extend_from_slice(&metadata);
for bytes in payloads {
@@ -79,12 +94,17 @@ pub(in crate::call) fn pack<T: Serialize, const N: usize>(
pub(in crate::call) fn unpack<T: DeserializeOwned, const N: usize>(
prefix: [u8; 4],
bytes: &[u8],
) -> Result<(T, [&[u8]; N]), ExecutionEncodingError> {
unpack_with_limit(prefix, bytes, MAX_EXECUTION_PAYLOAD_BYTES)
}
pub(in crate::call) fn unpack_with_limit<T: DeserializeOwned, const N: usize>(
prefix: [u8; 4],
bytes: &[u8],
maximum: usize,
) -> Result<(T, [&[u8]; N]), ExecutionEncodingError> {
let header_bytes = 4 + 4 * (N + 1);
if bytes.len() < header_bytes
|| bytes.len() > MAX_EXECUTION_PAYLOAD_BYTES
|| bytes[..4] != prefix
{
if bytes.len() < header_bytes || bytes.len() > maximum || bytes[..4] != prefix {
return Err(ExecutionEncodingError);
}
let mut cursor = header_bytes;
@@ -122,6 +142,7 @@ fn attach(target: Option<&mut Vec<u8>>, bytes: &[u8]) -> Result<(), ExecutionEnc
fn raw_wire(wire: Option<&mut super::WireEvent>) -> Option<&mut Vec<u8>> {
match &mut wire?.payload {
WirePayload::Json { raw_sse, .. } => raw_sse.as_mut(),
WirePayload::WebSocket { message, .. } => Some(message),
WirePayload::RawSse { frame } => Some(frame),
WirePayload::RawJson { body } | WirePayload::RawBody { body } => Some(body),
}
@@ -4,10 +4,7 @@
use serde::{Deserialize, Serialize};
use super::WireEvent;
use crate::call::{
middleware::MiddlewareHeader,
observation::{RequestMoney, RequestTimings},
};
use crate::call::{middleware::MiddlewareHeader, observation::RequestMoney};
#[derive(Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
@@ -76,12 +73,60 @@ pub struct ResponseObservation {
pub request_id: Option<String>,
pub service_tier: Option<String>,
pub upstream_response_model: Option<String>,
pub timings: RequestTimings,
#[serde(default)]
pub timings: ResponseTimings,
#[serde(default)]
pub client_headers: Vec<MiddlewareHeader>,
/// Provider 已产生的完整观测 JSON 文本,保留原始表示
pub provider_metadata: Option<String>,
}
/// 原始上游计时;不包含最终总耗时、速率或费用等业务结果
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ResponseTimings {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub transport_decision_wait_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub connect_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub headers_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_event_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_reasoning_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_text_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_token_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_processing_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub upstream_response_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub upstream_api_overhead_ms: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub upstream_engine_ms: Option<f64>,
}
impl ResponseObservation {
#[must_use]
pub fn new(transport: impl Into<String>) -> Self {
Self {
transport: transport.into(),
http_version: None,
websocket_pool: None,
status_code: None,
request_id: None,
service_tier: None,
upstream_response_model: None,
timings: ResponseTimings::default(),
client_headers: Vec::new(),
provider_metadata: None,
}
}
}
#[derive(Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SessionState {
@@ -6,6 +6,15 @@ pub(in crate::call) mod codec;
pub mod facts;
pub use codec::{ExecutionEncodingError, MAX_EXECUTION_PAYLOAD_BYTES};
/// 内部请求允许 JSON 字符串转义与元数据开销,出站正文仍按原始配置上限检查
#[must_use]
pub fn request_payload_limit(body_bytes: usize) -> usize {
body_bytes
.saturating_mul(2)
.saturating_add(64 * 1024)
.min(u32::MAX as usize)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContentKind {
@@ -149,6 +158,12 @@ pub enum WirePayload {
#[serde(default, skip_serializing_if = "Option::is_none")]
raw_sse: Option<Vec<u8>>,
},
/// 完整 WebSocket 文本消息及可用的解析视图,原文独立放入二进制段
WebSocket {
event: Option<String>,
data: Option<serde_json::Value>,
message: Vec<u8>,
},
/// 注释、非 JSON data 等原生 SSE 表达,不能夹带第二个帧
RawSse { frame: Vec<u8> },
/// 一个完整的非流式 JSON 响应正文,保留空白与字段顺序
@@ -5,19 +5,19 @@ use serde::{Deserialize, Serialize};
use super::{UpstreamAdapterEvent, UpstreamAdapterRequest, UpstreamContinuation, UpstreamFailure};
use crate::call::model::{
ExecutionEncodingError, ExecutionEvent,
codec::{pack, unpack},
codec::{pack, pack_with_limit, unpack, unpack_with_limit},
};
#[derive(Serialize, Deserialize)]
struct Metadata {
service_tier: Option<String>,
observation: Option<super::super::model::facts::ResponseObservation>,
continuation: Option<UpstreamContinuation>,
failure: Option<UpstreamFailure>,
}
pub(super) fn encode(value: UpstreamAdapterEvent) -> Result<Vec<u8>, ExecutionEncodingError> {
let metadata = Metadata {
service_tier: value.service_tier,
observation: value.observation,
continuation: value.continuation,
failure: value.failure,
};
@@ -28,7 +28,7 @@ pub(super) fn decode(bytes: &[u8]) -> Result<UpstreamAdapterEvent, ExecutionEnco
let (metadata, [event]): (Metadata, _) = unpack(*b"GPA1", bytes)?;
Ok(UpstreamAdapterEvent {
event: ExecutionEvent::decode(event)?,
service_tier: metadata.service_tier,
observation: metadata.observation,
continuation: metadata.continuation,
failure: metadata.failure,
})
@@ -38,12 +38,13 @@ pub(super) fn encode_request(
request: UpstreamAdapterRequest,
body: Vec<u8>,
) -> Result<Vec<u8>, ExecutionEncodingError> {
pack(*b"GPAQ", &request, [body])
// 请求预算由宿主配置检查,线协议只约束 u32 长度;响应事件仍受独立预算限制
pack_with_limit(*b"GPAQ", &request, [body], usize::MAX)
}
pub(super) fn decode_request(
bytes: &[u8],
) -> Result<(UpstreamAdapterRequest, Vec<u8>), ExecutionEncodingError> {
let (request, [body]) = unpack(*b"GPAQ", bytes)?;
let (request, [body]) = unpack_with_limit(*b"GPAQ", bytes, usize::MAX)?;
Ok((request, body.to_vec()))
}

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