Compare commits

...
35 Commits
Author SHA1 Message Date
ZyyandCodex e91ec88a1f chore: bump version to 3.3.24
Co-authored-by: Codex <noreply@openai.com>
2026-09-07 03:40:53 +08:00
ZyyandCodex 6f7a50b1b6 fix(scheduling): rotate smart selections among similarly scored accounts
Score load, remaining quota and first-output latency on fixed scales so
small differences do not monopolize unbound requests. Rotate accounts
within a narrow score tolerance in stable ID order using the existing
cursor, preserving session affinity and availability checks.

Cover small signal differences, changing score order, material advantages
and quota invalidation after credential rotation.

Validation: workspace tests (1844 passed, none ignored), rustfmt, and strict
workspace Clippy with all targets and features.

Co-authored-by: Codex <noreply@openai.com>
2026-09-07 03:40:34 +08:00
ZyyandCodex db0c294e83 fix(openai): prioritize session affinity and isolate child failover
Prefer client-scoped session bindings over ordinary account weights while
keeping existing availability checks and failover behavior. Let child
threads inherit the root account on first selection and update only their
own binding when they migrate.

Share session and thread identity extraction across Responses, Search and
image endpoints. Renew existing Redis bindings for 24 hours and preserve
account-bound state isolation without adding a table or scheduler.

Validation: workspace tests (1838 passed, none ignored), rustfmt, and strict
workspace Clippy with all targets and features.

Co-authored-by: Codex <noreply@openai.com>
2026-09-07 03:16:45 +08:00
ZyyandCodex d8181e8e91 docs(readme): install published images with Docker Compose
Download Compose and configuration files directly for installation, and
wait for healthy services before the first health check. Align deployment
paths and image upgrade instructions with the installation directory,
while documenting source checkout requirements for development and builds.

Validation: checked 16 shell examples and 17 local documentation links;
downloaded the deployment files into a temporary directory and verified
Compose configuration and a startup dry run without a source checkout.

Co-authored-by: Codex <noreply@openai.com>
2026-09-07 03:10:55 +08:00
ZyyandCodex af134cf488 fix(docker): apply system updates to the runtime base image
Upgrade installed Debian packages before adding runtime dependencies so
digest-pinned base images receive available fixes, including the PCRE2
security update.

Validation: rebuilt the amd64 runtime-base image without cache; Trivy
reported zero fixable OS vulnerabilities across all severities; curl,
pg_dump, and pg_restore started successfully as the cpr user.

Co-authored-by: Codex <noreply@openai.com>
2026-09-07 03:05:59 +08:00
Zyy d918149c2f fix(openai): recover exhausted quota windows independently
Persist per-window recovery baselines separately from current usage. Identify windows by upstream limit ID and actual duration, using protocol slots to distinguish equal or unspecified durations. Apply the same reset advancement and usage-below-10% rule to all periods without requiring allowed=true, retaining partial recovery across refreshes and restarts.

Use the exhausted window for reset tracking, preserve access during failure follow-ups, and return the updated account status after an admin quota refresh. Remove the ignored live Desktop artifact download test.

Validation: workspace tests (1828 passed), final OpenAI suite (516 passed, none ignored), fmt, clippy, and build. Real-account checks cover quota refresh and monthly recovery; seven controlled replay scenarios cover independent and partial recovery through the backend and PostgreSQL.
2026-09-07 00:59:48 +08:00
Zyy aa6bf6497f refactor(openai): unify WebSocket recovery and session fallback
Consolidate retry decisions in the provider and share the failure budget with session HTTP fallback. Preserve WebSocket priority and the 800ms cold connection budget.

Remove transport-local retries, fresh recovery probes, and metadata-based replay exceptions; use a 300s stream idle timeout.

Validation: backend tests (1816 passed, 1 ignored), fmt, clippy, and real Codex CLI 10-turn WebSocket continuation run.
2026-09-07 00:45:08 +08:00
ZyyandCodex 074c07b5ca refactor(store): tighten persisted facts and trim redundant observability
Remove unused event occurrence counts across storage, API and UI; count event
rows directly and keep provider metadata focused on facts absent from typed
observations. Preserve actual upstream-reported model and handshake status.

Narrow four live-reference indexes, drop eight low-value indexes, enforce
complete grouped facts and normalize administrator audit identities. Make
backup storage changes invalidate verification and pause schedules atomically,
while unchanged saves preserve state and revision.

Validation: 1,816 backend tests passed with isolated PostgreSQL and Redis,
including the official artifact download test. Strict Clippy, rustfmt, migration
freeze checks, frontend lint/typecheck/build and actionlint passed. Synthetic
10,000-row query plans verified live-reference and historical pagination indexes.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 17:27:29 +08:00
ZyyandCodex f70f46e87e fix(observability): exclude tokenless endpoints from usage
Define completed usage from delivery plus model, image, token, or cost evidence so standalone search remains operationally visible without appearing in usage records.

Exclude model-less provider endpoints from model diagnostics and keep attempt cost coverage on the shared usage fact.

Validation: 169 gateway-store tests passed against PostgreSQL; Clippy and rustfmt passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 17:02:41 +08:00
ZyyandCodex c21633655e fix(api): keep SSE responses alive during upstream silence
Send SSE keepalive comments after 15 seconds without output while preserving
the pending execution future. Disable response buffering and transformation
so long compaction requests can maintain the downstream connection.

Cover four minutes of upstream silence, cancellation, first-event commitment,
and conflicting upstream headers. Document transport errors and the separate
Codex event idle timeout.

Validation: 259 gateway-api tests and Clippy passed with serde/rc enabled;
cargo fmt and git diff --check passed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 16:49:35 +08:00
Zyy 1bafbb3864 fix(openai): support Auth0 profile avatars 2026-09-06 16:38:21 +08:00
Zyy 353dccedc4 feat: add request diagnostics and align Codex transport contracts 2026-09-06 16:24:14 +08:00
ZyyandCodex 697f1870a0 docs: refresh user-facing setup and existing guides
Refocus the README on installation, client setup and common questions. Consolidate native image generation, auth.json compatibility and obsolete Codex settings in the existing deployment guide.

Correct account import formats, API references, backup recovery and password rotation guidance. Synchronize architecture, theme, migration and repository skill documentation without adding new documents or changing temporary investigation notes.

Validated 53 local links, 22 shell examples, generated configuration parity and TOML syntax, Compose configuration, skill metadata and whitespace checks.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 15:28:32 +08:00
ZyyandCodex a404c2e09e feat(codex): enable native image generation with managed auth
Generate server-managed authentication settings for native image generation while preserving API-key auth.json and CCSwitch compatibility. Remove obsolete client settings and retain provider-level WebSocket control.

Keep gateway bearer authentication mandatory, strip client actor authorization headers before forwarding upstream, and cover the boundary with backend regression tests. Document setup and simplify the configuration dialog guidance.

Validated with backend tests, live Codex CLI image generation, WebSocket and legacy-auth checks, frontend build, and strict CLI configuration loading.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 15:09:05 +08:00
ZyyandCodex a6cfd124b0 fix(theme): balance popover header surfaces
Derive a dedicated popover header component token so light themes keep clear separation from striped table rows while dark themes retain the softer tertiary fill.

Apply the token consistently to structured popover headers and the health timeline arrow, and expose it through the theme editor.

Validated across all preset light and dark palettes, token overrides, frontend lint, typecheck and production build.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 14:18:04 +08:00
ZyyandCodex 4441b577c3 fix(theme): strengthen popover header separation
Promote structured popover headers from the tertiary fill to the secondary fill so they remain visually distinct from striped and hovered table rows in both light and dark themes.

Keep the health timeline arrow matched to its header surface when the popover placement flips.

Validated across all preset light and dark palettes, frontend lint, typecheck and production build.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 14:14:53 +08:00
ZyyandCodex a704fbc3f3 fix(websocket): serialize queued response requests
Queue incoming response.create frames until the active response finishes instead of closing the connection, aligning execution with the Codex client serial request flow.

Separate pump exit notifications from the bounded business queue so disconnects and shutdown remain prompt. Preserve expiry handling with a full queue, retain protocol validation, and document the serial execution contract.

Validated with 253 API tests, 12 architecture tests, strict workspace Clippy and Rust formatting.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 13:55:14 +08:00
ZyyandCodex dab221cdb9 refactor(theme): standardize container color roles
Rename primary, semantic and preset background roles to container and their paired foreground roles to on-container across derivation, CSS variables, utilities, consumers and documentation.

Preserve existing palette values and contrast relationships while making text on neutral surfaces distinct from content inside semantic containers.

Validated with token equivalence across presets and custom seeds, frontend typecheck, lint and production build.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 12:48:19 +08:00
ZyyandCodex 4f434ec525 fix(theme): restore semantic color roles
Separate semantic text on neutral surfaces from high-contrast text on tinted backgrounds, derive each role against its actual surfaces, and migrate badges and alerts to explicit on-background tokens.

Restore vivid usage-table accents and align input hover, active, and error outlines across light and dark themes.

Validated with semantic contrast checks across presets and custom seeds, frontend typecheck, lint and production build.

Co-authored-by: Codex <noreply@openai.com>
2026-09-06 12:08:34 +08:00
ZyyandCodex 2585d4cd7d refactor(architecture): consolidate usage and provider contracts
Move Dashboard quota projection to the backend, share pagination request state, and fix independent usage loads while keeping trend switches free of loading overlays.

Stream billing facts within the query budget, share pending authorization and usage validation, and separate Core value ownership and response observations. Remove the unused theme helper and strengthen dependency checks.

Validated with workspace integration tests against isolated PostgreSQL and Redis, strict Clippy, Rust formatting, and frontend typecheck, lint and build.

Co-authored-by: Codex <noreply@openai.com>
2026-09-05 05:29:12 +08:00
ZyyandCodex 93a34b0f56 fix(openai): add gpt-6-astra billing
Add standard, Flex, Fast, long-context and cache-write rates, and classify Astra as a reasoning model for preview web-search pricing. Cover tier aliases, the context threshold and cost components with backend regression tests.

Co-authored-by: Codex <noreply@openai.com>
2026-09-05 05:27:32 +08:00
ZyyandCodex aa3a07209a fix(theme): improve token contrast and color derivation
Guard text and primary button contrast across surfaces and interaction states, derive dark semantic text from its seed, and reuse focus tokens. Document the palette and component rules.

Co-authored-by: Codex <noreply@openai.com>
2026-09-05 05:11:17 +08:00
ZyyandCodex c2332bc8d0 chore: bump version to 3.3.23
Co-authored-by: Codex <noreply@openai.com>
2026-09-04 19:16:46 +08:00
ZyyandCodex 8f242e2875 feat(logging): add opt-in request dumps
Capture complete OpenAI Responses headers and payloads in a dedicated rotated log, keep the target out of normal sinks, and default the sensitive switch to off.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 19:15:28 +08:00
ZyyandCodex 51e9981bee fix(openai): resume metadata-close websocket requests
Treat metadata turn state as a same-request checkpoint, retry on the original account over WebSocket, and keep ambiguous requests without a checkpoint fail-closed.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 18:48:36 +08:00
ZyyandCodex 65fbbcb3ff chore: bump version to 3.3.22
Co-authored-by: Codex <noreply@openai.com>
2026-09-04 16:52:27 +08:00
ZyyandCodex bd086a7e3c fix(observability): show reasoning token usage
Expose aggregate reasoning tokens on account responses and use them in the token structure panel while retaining cache-write fields for API compatibility.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 16:51:18 +08:00
ZyyandCodex 3554ba13db fix(openai): preserve websocket turn state on failure
Bridge both metadata event names, project sticky state per response.create, and return captured state on pre-terminal failures without replaying ambiguous requests.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 16:51:01 +08:00
ZyyandCodex bcccebc9fa chore: bump version to 3.3.21
Co-authored-by: Codex <noreply@openai.com>
2026-09-04 11:52:25 +08:00
ZyyandCodex d67b46a9e2 fix(frontend): default Codex websocket config off
Keep inline switch text legible in the off state and omit WebSocket settings from generated Codex and CCSwitch configuration unless explicitly enabled.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 11:50:47 +08:00
ZyyandCodex 20c6d33cf4 fix(observability): retain recovered transport failures
Link successful same-session HTTP fallbacks to pending WebSocket failures, keep recovered records visible in Ops diagnostics, and label their recovery state.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 11:50:41 +08:00
ZyyandCodex cf0e84c548 fix(openai): isolate websocket continuation lanes
Include the downstream connection identity in pooled upstream WebSocket keys so concurrent clients cannot contend for connection-local response state.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 11:50:34 +08:00
ZyyandCodex 2332af9a08 chore: bump version to 3.3.20
Co-authored-by: Codex <noreply@openai.com>
2026-09-04 00:47:35 +08:00
ZyyandCodex 024fc9fcf0 fix(openai): fall back after websocket fast-path misses
Treat the 800 ms fast-path budget as a routing decision instead of a provider error, while preserving background preconnect for the next request.

Co-authored-by: Codex <noreply@openai.com>
2026-09-04 00:36:50 +08:00
ZyyandCodex 4220f4e095 fix(openai): recover websocket sessions after ambiguous failures
Co-authored-by: Codex <noreply@openai.com>
2026-09-03 23:33:26 +08:00
288 changed files with 11572 additions and 4095 deletions
+2
View File
@@ -33,6 +33,8 @@ jobs:
timeout-minutes: 30
env:
CARGO_TERM_COLOR: always
# Rust 1.97 编译 AWS SDK 调试信息需要更大的 codegen 线程栈。
RUST_MIN_STACK: 16777216
CPR_TEST_DATABASE_URL: postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy
CPR_TEST_REDIS_URL: redis://127.0.0.1:6379
services:
+92 -123
View File
@@ -5,192 +5,161 @@
# Codex Proxy RS
**基于 Rust 的多 Provider 透明 AI 网关**
集中管理 OpenAI、xAI 账号,为 Codex 和兼容客户端提供统一的访问地址与密钥。
[![CI](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
[![Release](https://img.shields.io/github/v/release/zyycn/codex-proxy-rs?display_name=tag&sort=semver&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/releases)
[![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
[![MIT](https://img.shields.io/badge/License-MIT-blue.svg?style=flat-square)](https://opensource.org/license/mit)
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运行维护](#运行维护) · [部署](deploy/README.md) · [接口](docs/api.md) · [架构](docs/architecture.md)
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [常见问题](#常见问题) · [部署说明](deploy/README.md)
</div>
> [!NOTE]
> 数据面支持 OpenAI Responses、Images 和模型目录协议,不提供 `/v1/chat/completions`。
## 功能
| 领域 | 当前能力 |
| --- | --- |
| 数据面 | Responses JSON、SSE、WebSocket(含 Codex 子代理语义)、Images 生成/编辑与模型目录 |
| Provider | 固定 `openai`、`xai` 两个 Provider,各自拥有 OAuth、credential、额度、目录与 transport |
| 路由 | Client Key 限定账号分组;按模型能力编译 Provider 候选;支持全局精确模型映射、会话亲和和安全 fallback |
| 透明边界 | OpenAI 保留未知请求字段及字段顺序,并原样转发 SSE、WebSocket 与 Images 业务字节;xAI 在 Provider 内转换协议 |
| 账号 | OAuth、OpenAI AT/RT 与账号文件导入,credential/额度刷新、账号恢复及主动额度重置卡 |
| 管理 | 账号与分组、Client Key、运行设置、用量/错误观测、S3/R2 备份、系统更新与回滚 |
| 计量 | Token、图片、费用、延迟、账号与 Provider 归因,并记录上游响应确认的实际 `service_tier` |
- 在网页中添加账号,查看额度、请求记录和费用,刷新或重新授权过期账号。
- 用账号分组控制不同密钥可使用的账号,设置并发和请求频率。
- 接入 Codex CLI、桌面端和支持 Responses API 的客户端。
- 支持流式对话、WebSocket、图片生成与编辑、Codex 独立搜索。
- 支持 S3/R2 数据库备份、同大版本在线更新和回滚。
架构边界、状态所有权和请求不变量见 [架构文档](docs/architecture.md)。
支持 OpenAI、xAI 账号,但图片与独立搜索目前仅通过 OpenAI 提供。模型和功能是否可用取决于上游账号权限。
**不支持 `/v1/chat/completions`**;仅支持该接口的客户端不能直接接入。
## 快速开始
需要 Docker Engine 与 Docker Compose Plugin。
使用 Docker Compose 部署发布镜像 `ghcr.io/zyycn/codex-proxy-rs:latest`,同时启动 PostgreSQL 和 Redis。
以下命令适用于 Linux amd64/arm64,需要 Docker Engine、Docker Compose Plugin、curl 和 OpenSSL。已有部署请先看
[升级说明](deploy/README.md#镜像升级与源码构建),不要覆盖原配置。
### 1. 准备配置和目录
### 1. 下载部署文件并配置
```bash
git clone https://github.com/zyycn/codex-proxy-rs.git
mkdir -p codex-proxy-rs/deploy
cd codex-proxy-rs
curl -fsSL https://raw.githubusercontent.com/zyycn/codex-proxy-rs/main/deploy/compose.yaml \
-o deploy/compose.yaml
curl -fsSL https://raw.githubusercontent.com/zyycn/codex-proxy-rs/main/deploy/config.example.yaml \
-o deploy/config.example.yaml
mkdir -p .runtime/data .runtime/logs
install -d -m 0750 .runtime/postgres .runtime/redis
cp deploy/config.example.yaml deploy/config.yaml
sudo chown "$(id -u):10001" deploy/config.yaml
chmod 0640 deploy/config.yaml
```
分别生成 PostgreSQL 与 Redis 密码:
```bash
openssl rand -hex 24
openssl rand -hex 24
```
把两个结果分别写入 `deploy/config.yaml`,再设置管理员初始密码:
| 配置 | 约束 |
| --- | --- |
| `store.database.password` | 48 位十六进制 |
| `store.redis.password` | 48 位十六进制 |
| `admin.default_password` | 至少 12 位,不能包含 `$` |
PostgreSQL、Redis 和应用从同一份 `config.yaml` 读取基础设施密码,不需要额外导出环境变量。
管理员密码只在首次创建管理员时生效。
Linux 需要允许容器组写入运行目录:
```bash
sudo chown -R "$(id -u):10001" .runtime/data .runtime/logs
chmod 0770 .runtime/data .runtime/logs
```
### 2. 启动
分别生成数据库和 Redis 密码:
```bash
openssl rand -hex 24
openssl rand -hex 24
```
编辑 `deploy/config.yaml`,填好以下三项:
| 配置项 | 填写内容 |
| --- | --- |
| `store.database.password` | 第一个生成的 48 位十六进制密码 |
| `store.redis.password` | 第二个生成的 48 位十六进制密码 |
| `admin.default_password` | 管理员初始密码,至少 12 位,不能包含 `$` |
不需要创建 `.env`。管理员初始密码只在首次创建账号时使用。
### 2. 启动服务
```bash
docker compose -f deploy/compose.yaml config --quiet
docker compose -f deploy/compose.yaml pull
docker compose -f deploy/compose.yaml up -d --no-build
docker compose -f deploy/compose.yaml up -d --no-build --wait
curl -i http://127.0.0.1:8080/healthz
```
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址为
`http://127.0.0.1:8080`。
健康检查返回 `204 No Content` 后,打开 `http://127.0.0.1:8080`,
使用 `admin@cpr.local` 和刚设置的管理员密码登录。
### 3. 初始化网关
默认地址只能在服务器本机访问。从其他设备使用时,需要配置
[HTTPS 反向代理](deploy/README.md#公网访问)。
1. 使用 `admin@cpr.local` 与初始密码登录。
2. 导入或授权 Provider 账号:OpenAI 支持 OAuth、逐行 AT/RT 与 OAuth JSON;
xAI 支持 OAuth 和账号 JSON。
3. 创建账号分组并按需配置全局模型映射。未命中的模型名原样交给候选 Provider。
4. 创建 `sk_...` Client Key,选择可用账号分组,并设置速率和并发限制;空分组范围表示全部账号。
### 3. 添加账号和密钥
OpenAI RT 导入会先向官方 token endpoint 换取 AT;仅 AT 的账号不具备自动续期能力,
能否参与调度仍取决于 token 中可验证的身份与后续上游事实。账号导入和首次 OAuth 创建会尽力读取一次额度,
额度上游暂时失败不会回滚已提交的账号。
1. 在「账号」中添加 OpenAI 或 xAI 账号。OpenAI 支持 OAuth 授权、AT/RT 和账号 JSON;xAI 支持 OAuth 授权和账号 JSON。
2. 按需创建账号分组,并把账号加入分组。
3. 创建客户端密钥,选择可用分组。**不选分组表示可使用全部账号**。
4. 点击该密钥的「使用密钥」,复制客户端配置。
OpenAI 账号的主动重置卡只在用户打开相应弹窗或手工刷新时查询。消费成功后管理端重新查询卡片与额度;
系统不会通过修改本地 `reset_at` 伪造重置。完整账号合同见 [账号接口](docs/api.md#5-账号)。
> [!IMPORTANT]
> xAI 使用 OAuth session,不支持把 xAI API Key 作为上游 credential。
完整的权限、密码轮换、备份和升级规则见 [部署文档](deploy/README.md)。
只导入 AT 的账号不能自动续期。xAI 暂不支持用 API Key 导入上游账号。
## 客户端接入
### Codex CLI 和桌面端
在「使用密钥」中按系统复制 `config.toml` 和 `auth.json`,也可以通过 CCSwitch 导入。
已有文件先备份,再合并配置并重启 Codex。客户端使用代理密钥,OpenAI 账号登录状态由服务端管理。
新配置已启用原生生图。可以直接提出需求,例如:
> 帮我做一个咖啡店首页,生成一张咖啡主题配图并用到页面中。
是否调用生图由 Codex 根据任务决定,也可以明确要求生成图片。需要支持生图的客户端、模型和 OpenAI 账号;
普通对话可用不代表账号一定有生图额度。
完整配置和旧版配置迁移见 [客户端配置](deploy/README.md#客户端配置)。
### 其他客户端
选择支持 **Responses API** 的客户端,填写:
| 配置 | 值 |
| --- | --- |
| Base URL | `http://127.0.0.1:8080/v1` |
| API Key | 管理端创建的 `sk_...` |
| 鉴权 | `Authorization: Bearer <client-api-key>` |
| Base URL | `http://127.0.0.1:8080/v1`,远程使用时换成服务器的 HTTPS 地址 |
| API Key | 管理端创建的 `sk_...` 密钥 |
可先检查密钥是否能读取模型列表:
```bash
curl http://127.0.0.1:8080/v1/models \
-H 'Authorization: Bearer <client-api-key>'
curl http://127.0.0.1:8080/v1/responses \
-H 'Authorization: Bearer <client-api-key>' \
-H 'Content-Type: application/json' \
-d '{
"model": "<model-id>",
"input": "Reply with pong.",
"stream": false
}'
```
<details>
<summary>公开数据面路由</summary>
接口列表与请求示例见 [接口文档](docs/api.md)。
| 路由 | 用途 |
| --- | --- |
| `POST /v1/responses` | JSON 响应或 SSE Responses 流 |
| `GET /v1/responses` | Responses WebSocket 升级 |
| `POST /v1/images/generations` | OpenAI Images 生成 JSON 透明转发 |
| `POST /v1/images/edits` | OpenAI Images 编辑 JSON 透明转发 |
| `GET /v1/models` | 当前 Client Key 账号范围内的模型列表 |
| `GET /v1/models/{model_id}` | OpenAI 兼容模型详情 |
## 常见问题
所有 `/v1/*` 路由都需要 Client Key。Images 是 OpenAI Provider 自有端点,不读取模型字段,也不参与
文本模型映射。
### 仍然要求登录,或无法生成图片
</details>
重新从管理端复制配置,确认修改的是客户端实际读取的用户配置目录,然后完全退出并重启客户端。
保留仅含代理密钥的 `auth.json` 不会使新配置失效。详细检查步骤见
[登录与生图排查](deploy/README.md#登录与生图排查)。
完整的客户端与管理端路由、鉴权、请求字段和 mutation 语义见 [接口文档](docs/api.md)。
### 没有可用账号或模型
## 运行维护
检查账号是否启用、凭据是否过期、额度是否耗尽,以及密钥绑定的分组是否包含可用账号。
「恢复账号」只清除本地错误状态,不能恢复上游权限或增加额度。
### 如何更新或查看日志
同一大版本内拉取发布镜像并重建应用容器:
```bash
# 同一大版本内升级发布镜像
docker compose -f deploy/compose.yaml pull codex-proxy-rs
docker compose -f deploy/compose.yaml up -d --no-build
# 查看日志
docker compose -f deploy/compose.yaml logs -f codex-proxy-rs
# 从当前源码构建
docker compose -f deploy/compose.yaml build codex-proxy-rs
docker compose -f deploy/compose.yaml up -d
docker compose -f deploy/compose.yaml up -d --no-build codex-proxy-rs
docker compose -f deploy/compose.yaml logs --tail=100 -f codex-proxy-rs
```
> [!IMPORTANT]
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点和在线更新状态。删除该目录会永久清除
> 本地运行状态。
> [!WARNING]
> 不支持跨大版本在线升级。升级到新的大版本时请使用全新的 `.runtime/` 数据目录,并重新导入或授权
> Provider 账号与 Client Key。
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket upgrade 与真实客户端
IP;不要暴露 PostgreSQL 或 Redis。当前运行拓扑为单副本,不能通过复制应用容器扩容。
## 维护与发布
仓库版本只能从干净、已同步上游的发布分支发布:
```bash
release/publish <version>
```
脚本依赖已登录的 GitHub CLI(`gh auth login`),会更新 `release/version.yaml`、创建约定提交和带注释的
`v<version>` tag,原子推送分支与 tag,再触发 Release 构建。源码提交、仓库发版和运行实例升级是三个
独立状态;脚本不会登录服务器或改变任何运行实例。
> 升级前先备份。跨大版本需使用新的数据目录重新部署,不支持直接沿用旧库在线升级。
> `.runtime/` 包含数据库和运行数据,请勿删除;账号文件、密钥和日志也不要公开分享。
## 文档
- [部署与运维](deploy/README.md)
- [客户端与管理端接口](docs/api.md)
- [客户端接入与生图](deploy/README.md#客户端配置)
- [部署、备份与恢复](deploy/README.md)
- [API 参考](docs/api.md)
- [系统架构](docs/architecture.md)
- [数据库迁移规则](backend/migrations/README.md)
- [Release](https://github.com/zyycn/codex-proxy-rs/releases)
- [容器镜像](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
- [管理端主题](docs/theme.md)
- [数据库迁移](backend/migrations/README.md)
+50 -44
View File
@@ -154,9 +154,9 @@ dependencies = [
[[package]]
name = "aws-lc-rs"
version = "1.17.1"
version = "1.16.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4342d8937fc7e5dd9b1c60292261c0670c882a2cd1719cfc11b1af41731e32ad"
checksum = "a054912289d18629dc78375ba2c3726a3afe3ff71b4edba9dedfca0e3446d1fc"
dependencies = [
"aws-lc-sys",
"untrusted 0.7.1",
@@ -165,15 +165,14 @@ dependencies = [
[[package]]
name = "aws-lc-sys"
version = "0.42.0"
version = "0.39.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d9ceb1da931507a12f4fccea479dccd00da1943e1b4ae72d8e502d707361444"
checksum = "1fa7e52a4c5c547c741610a2c6f123f3881e409b714cd27e6798ef020c514f0a"
dependencies = [
"cc",
"cmake",
"dunce",
"fs_extra",
"pkg-config",
]
[[package]]
@@ -195,7 +194,7 @@ dependencies = [
"bytes-utils",
"fastrand",
"http 0.2.12",
"http 1.4.2",
"http 1.4.0",
"http-body 0.4.6",
"http-body 1.0.1",
"percent-encoding",
@@ -231,7 +230,7 @@ dependencies = [
"hex",
"hmac",
"http 0.2.12",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"lru",
"percent-encoding",
@@ -257,7 +256,7 @@ dependencies = [
"hex",
"hmac",
"http 0.2.12",
"http 1.4.2",
"http 1.4.0",
"percent-encoding",
"sha2 0.11.0",
"time",
@@ -286,7 +285,7 @@ dependencies = [
"bytes",
"crc-fast",
"hex",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"md-5",
@@ -320,7 +319,7 @@ dependencies = [
"bytes-utils",
"futures-core",
"futures-util",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"percent-encoding",
@@ -339,7 +338,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-types",
"h2",
"http 1.4.2",
"http 1.4.0",
"hyper",
"hyper-rustls",
"hyper-util",
@@ -389,7 +388,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.4.2",
"http 1.4.0",
"http-body 0.4.6",
"http-body 1.0.1",
"http-body-util",
@@ -410,7 +409,7 @@ dependencies = [
"aws-smithy-types",
"bytes",
"http 0.2.12",
"http 1.4.2",
"http 1.4.0",
"pin-project-lite",
"tokio",
"tracing",
@@ -436,7 +435,7 @@ checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910"
dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-types",
"http 1.4.2",
"http 1.4.0",
]
[[package]]
@@ -450,7 +449,7 @@ dependencies = [
"bytes-utils",
"futures-core",
"http 0.2.12",
"http 1.4.2",
"http 1.4.0",
"http-body 0.4.6",
"http-body 1.0.1",
"http-body-util",
@@ -503,7 +502,7 @@ dependencies = [
"bytes",
"form_urlencoded",
"futures-util",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"hyper",
@@ -535,7 +534,7 @@ checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
dependencies = [
"bytes",
"futures-core",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"mime",
@@ -1478,12 +1477,14 @@ name = "gateway-core"
version = "0.1.0"
dependencies = [
"async-trait",
"base64 0.23.1",
"bytes",
"futures",
"futures-timer",
"semver",
"serde",
"serde_json",
"sha2 0.11.0",
"thiserror",
"tracing",
"uuid",
@@ -1533,6 +1534,7 @@ dependencies = [
name = "gateway-store"
version = "0.1.0"
dependencies = [
"async-stream",
"async-trait",
"aws-credential-types",
"aws-sdk-s3",
@@ -1638,7 +1640,7 @@ dependencies = [
"fnv",
"futures-core",
"futures-sink",
"http 1.4.2",
"http 1.4.0",
"indexmap",
"slab",
"tokio",
@@ -1686,7 +1688,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"headers-core",
"http 1.4.2",
"http 1.4.0",
"httpdate",
"mime",
"sha1 0.10.7",
@@ -1698,7 +1700,7 @@ version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4"
dependencies = [
"http 1.4.2",
"http 1.4.0",
]
[[package]]
@@ -1760,9 +1762,9 @@ dependencies = [
[[package]]
name = "http"
version = "1.4.2"
version = "1.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425"
checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a"
dependencies = [
"bytes",
"itoa",
@@ -1786,7 +1788,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184"
dependencies = [
"bytes",
"http 1.4.2",
"http 1.4.0",
]
[[package]]
@@ -1797,7 +1799,7 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a"
dependencies = [
"bytes",
"futures-core",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"pin-project-lite",
]
@@ -1833,21 +1835,22 @@ dependencies = [
[[package]]
name = "hyper"
version = "1.10.1"
version = "1.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498"
checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11"
dependencies = [
"atomic-waker",
"bytes",
"futures-channel",
"futures-core",
"h2",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"httparse",
"httpdate",
"itoa",
"pin-project-lite",
"pin-utils",
"smallvec",
"tokio",
"want",
@@ -1859,7 +1862,7 @@ version = "0.27.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58"
dependencies = [
"http 1.4.2",
"http 1.4.0",
"hyper",
"hyper-util",
"log",
@@ -1881,7 +1884,7 @@ dependencies = [
"bytes",
"futures-channel",
"futures-util",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"hyper",
"ipnet",
@@ -2593,6 +2596,7 @@ version = "0.0.0"
dependencies = [
"async-stream",
"async-trait",
"aws-lc-rs",
"base64 0.23.1",
"bytes",
"chrono",
@@ -2604,6 +2608,7 @@ dependencies = [
"gateway-protocol",
"getrandom 0.4.3",
"hex",
"hyper",
"hyper-rustls",
"reqwest",
"roxmltree",
@@ -2619,6 +2624,7 @@ dependencies = [
"tempfile",
"thiserror",
"tokio",
"tokio-rustls",
"tokio-tungstenite 0.28.0",
"tokio-util",
"tracing",
@@ -2910,7 +2916,7 @@ dependencies = [
"futures-core",
"futures-util",
"h2",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"hyper",
@@ -3020,9 +3026,9 @@ dependencies = [
[[package]]
name = "rustls-native-certs"
version = "0.8.4"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d"
checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63"
dependencies = [
"openssl-probe",
"rustls-pki-types",
@@ -3032,9 +3038,9 @@ dependencies = [
[[package]]
name = "rustls-pki-types"
version = "1.15.1"
version = "1.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96"
checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd"
dependencies = [
"web-time",
"zeroize",
@@ -3698,7 +3704,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
@@ -3799,9 +3805,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokio"
version = "1.53.1"
version = "1.52.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe"
dependencies = [
"bytes",
"libc",
@@ -3911,7 +3917,7 @@ dependencies = [
"bitflags",
"bytes",
"futures-util",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"pin-project-lite",
"tower",
@@ -3930,7 +3936,7 @@ dependencies = [
"bytes",
"futures-core",
"futures-util",
"http 1.4.2",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"http-range-header",
@@ -4062,7 +4068,7 @@ dependencies = [
"data-encoding",
"flate2",
"headers",
"http 1.4.2",
"http 1.4.0",
"httparse",
"log",
"rand 0.9.4",
@@ -4081,7 +4087,7 @@ checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8"
dependencies = [
"bytes",
"data-encoding",
"http 1.4.2",
"http 1.4.0",
"httparse",
"log",
"rand 0.9.4",
@@ -4331,9 +4337,9 @@ dependencies = [
[[package]]
name = "webpki-roots"
version = "1.0.8"
version = "1.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf"
checksum = "12bed680863276c63889429bfd6cab3b99943659923822de1c8a39c49e4d722c"
dependencies = [
"rustls-pki-types",
]
@@ -4535,7 +4541,7 @@ dependencies = [
"base64 0.22.1",
"deadpool",
"futures",
"http 1.4.2",
"http 1.4.0",
"http-body-util",
"hyper",
"hyper-util",
+1 -1
View File
@@ -26,7 +26,7 @@ tokio = { version = "1.52.3", default-features = false, features = ["rt-multi-th
[dev-dependencies]
config = { version = "0.15.23", default-features = false, features = ["yaml"] }
serde_json = "1"
syn = { version = "3.0.3", features = ["full"] }
syn = { version = "3.0.3", features = ["full", "visit"] }
tempfile = "3.27.0"
[lints]
+119 -20
View File
@@ -6,7 +6,7 @@ use std::{
path::{Path, PathBuf},
};
use syn::Item;
use syn::{Item, visit::Visit};
/// workspace 成员冻结清单;新增 crate 必须同步扩展本文件的依赖规则。
pub(super) const WORKSPACE_MEMBERS: &[&str] = &[
@@ -61,6 +61,82 @@ fn gateway_core_depends_on_no_http_db_redis_or_provider_crate() {
);
}
#[test]
fn core_value_owners_do_not_depend_on_execution_or_routing() {
let root = backend_root().join("crates/gateway-core/src");
for relative in super::rust_files(&root) {
let allowed: Option<&[&str]> = match relative.to_str().expect("source path") {
"validation.rs" => Some(&[]),
"identity.rs" => Some(&["validation"]),
"upstream.rs" => Some(&["validation"]),
"event.rs" => Some(&["metering", "operation", "upstream", "validation"]),
"account/store.rs" => Some(&["account", "error", "identity", "validation"]),
path if path.starts_with("policy/") => Some(&["account", "policy", "validation"]),
path if path.starts_with("account/") => Some(&["account", "identity", "validation"]),
_ => None,
};
let Some(allowed) = allowed else { continue };
let source = fs::read_to_string(root.join(&relative)).expect("read core source");
let syntax = syn::parse_file(&source).expect("parse core source");
let mut references = CrateReferences::default();
references.visit_file(&syntax);
for dependency in references.0 {
assert!(
allowed.contains(&dependency.as_str()),
"{} must not depend on crate::{dependency}",
relative.display()
);
}
}
}
#[derive(Default)]
struct CrateReferences(BTreeSet<String>);
impl<'ast> Visit<'ast> for CrateReferences {
fn visit_path(&mut self, path: &'ast syn::Path) {
let mut segments = path.segments.iter();
if segments
.next()
.is_some_and(|segment| segment.ident == "crate")
&& let Some(owner) = segments.next()
{
self.0.insert(owner.ident.to_string());
}
syn::visit::visit_path(self, path);
}
fn visit_item_use(&mut self, item: &'ast syn::ItemUse) {
self.use_tree(&item.tree, false);
}
}
impl CrateReferences {
fn use_tree(&mut self, tree: &syn::UseTree, crate_root: bool) {
match tree {
syn::UseTree::Path(path) if crate_root => {
self.0.insert(path.ident.to_string());
}
syn::UseTree::Name(name) if crate_root => {
self.0.insert(name.ident.to_string());
}
syn::UseTree::Rename(rename) if crate_root => {
self.0.insert(rename.ident.to_string());
}
syn::UseTree::Path(path) => self.use_tree(&path.tree, path.ident == "crate"),
syn::UseTree::Group(group) => {
for item in &group.items {
self.use_tree(item, crate_root);
}
}
syn::UseTree::Glob(_) if crate_root => {
self.0.insert("*".to_owned());
}
_ => {}
}
}
}
#[test]
fn gateway_protocol_has_no_workspace_dependencies() {
for name in dependency_names("crates/gateway-protocol") {
@@ -156,6 +232,46 @@ const ALLOWED_INTERNAL_EDGES: &[(&str, &str)] = &[
#[test]
fn workspace_internal_dependency_edges_match_frozen_dag() {
let metadata = cargo_metadata_json();
let actual = workspace_runtime_dependency_edges(&metadata);
let mut expected: Vec<(String, String)> = ALLOWED_INTERNAL_EDGES
.iter()
.map(|(from, to)| ((*from).to_owned(), (*to).to_owned()))
.collect();
expected.sort();
assert_eq!(
actual, expected,
"workspace internal dependency edges diverged from the frozen DAG"
);
}
#[test]
fn runtime_dependency_edges_exclude_dev_build_and_external_dependencies() {
let metadata = serde_json::json!({
"packages": [
{
"name": "gateway-api",
"dependencies": [
{ "name": "gateway-core", "kind": null },
{ "name": "gateway-store", "kind": "dev" },
{ "name": "gateway-host", "kind": "build" },
{ "name": "serde", "kind": null }
]
},
{
"name": "external-package",
"dependencies": [{ "name": "gateway-core", "kind": null }]
}
]
});
assert_eq!(
workspace_runtime_dependency_edges(&metadata),
vec![("gateway-api".to_owned(), "gateway-core".to_owned())]
);
}
fn workspace_runtime_dependency_edges(metadata: &serde_json::Value) -> Vec<(String, String)> {
let packages = metadata["packages"]
.as_array()
.expect("cargo metadata packages");
@@ -173,31 +289,14 @@ fn workspace_internal_dependency_edges_match_frozen_dag() {
if !internal_names.contains(dep_name) {
continue;
}
let is_runtime_edge = dependency["dep_kinds"]
.as_array()
.map(|kinds| {
kinds
.iter()
.any(|kind| kind["kind"].as_str() == Some("normal"))
})
.unwrap_or(true);
if is_runtime_edge {
if dependency.get("kind").expect("dependency kind").is_null() {
actual.push((name.to_owned(), dep_name.to_owned()));
}
}
}
let mut expected: Vec<(String, String)> = ALLOWED_INTERNAL_EDGES
.iter()
.map(|(from, to)| ((*from).to_owned(), (*to).to_owned()))
.collect();
actual.sort();
expected.sort();
assert_eq!(
actual, expected,
"workspace internal dependency edges diverged from the frozen DAG"
);
actual
}
#[test]
+2 -5
View File
@@ -143,10 +143,7 @@ fn config_loader_should_reject_removed_fingerprint_section() {
#[test]
fn config_loader_should_reject_invalid_codex_cli_version() {
assert_rejected(valid_config().replace(
"codex_version: '0.151.0-alpha.7.2'",
"codex_version: 'latest'",
));
assert_rejected(valid_config().replace("codex_version: '0.153.4'", "codex_version: 'latest'"));
}
#[test]
@@ -165,7 +162,7 @@ fn config_loader_should_reject_zero_server_port() {
#[test]
fn config_loader_should_reject_invalid_desktop_profile_fields() {
assert_rejected(valid_config().replace("desktop_build: '7377'", "desktop_build: 'build'"));
assert_rejected(valid_config().replace("desktop_build: '8109'", "desktop_build: 'build'"));
}
fn assert_rejected(config: String) {
+1 -1
View File
@@ -22,6 +22,7 @@ gateway-core = { path = "../gateway-core" }
hex = "0.4.3"
rand_core = { version = "0.6.4", features = ["getrandom"] }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
secrecy = { version = "0.10.3", features = ["serde"] }
subtle = "2.6.1"
thiserror = "2.0.18"
@@ -32,7 +33,6 @@ chrono-tz = "0.10"
tokio = { version = "1.52", features = ["rt", "time", "macros", "sync"] }
[dev-dependencies]
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
sha2 = "0.11"
tempfile = "3"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] }
+1
View File
@@ -304,6 +304,7 @@ pub async fn initialize(
client_distribution: Arc::new(DefaultClientDistributionService::new(client_distribution)),
observability: Arc::new(DefaultObservabilityService::new(
store.observability(),
store.accounts(),
store.settings(),
registry,
)),
@@ -4,6 +4,12 @@ use chrono::{DateTime, Utc};
use super::{MutationActor, MutationContext, Revision};
/// 会话审计使用的稳定管理员标识。
#[must_use]
pub fn admin_session_actor_ref(admin_user_id: &str) -> String {
format!("admin:{admin_user_id}")
}
/// 已认证的管理主体。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AdminPrincipal {
@@ -491,6 +491,8 @@ pub struct DashboardAccountUsage {
///
/// `None` 表示上游未提供可比较的百分比,不应伪造为零。
pub quota_used_percent: Option<f64>,
/// Provider 返回的代表窗口;缺失时保留未知语义。
pub quota_window: Option<super::provider_credentials::ProviderQuotaWindow>,
pub models: Vec<DashboardAccountModelUsage>,
}
@@ -702,6 +704,8 @@ pub struct UsageAttempt {
/// 一条请求及其全部尝试。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageDetail {
pub trace: Option<serde_json::Value>,
pub related_requests: Vec<serde_json::Value>,
pub request: UsageRecord,
pub attempts: Vec<UsageAttempt>,
}
@@ -783,7 +787,6 @@ pub struct OpsError {
pub latency_ms: Option<u64>,
pub message: String,
pub raw_upstream_error: Option<String>,
pub occurrence_count: u32,
pub client_ip: Option<String>,
pub user_agent: Option<String>,
pub reasoning_effort: Option<String>,
@@ -10,6 +10,8 @@ use gateway_core::{
account::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
routing::{ProviderKind, UpstreamModelId},
};
use serde::Deserialize;
use serde_json::{Map, Value};
use uuid::Uuid;
use super::{
@@ -279,6 +281,101 @@ impl PendingAuthorizationMutation {
}
}
impl PendingAuthorizationMutation {
/// 编码 v1 中立事务字段;版本号的位置和 Provider 私有外层文档由 adapter 保持。
#[must_use]
pub fn to_storage_v1(&self) -> Map<String, Value> {
let target = match &self.target {
AuthorizationMutationTarget::Create { name } => {
serde_json::json!({"kind": "create", "name": name})
}
AuthorizationMutationTarget::Reauthorize { account_id } => {
serde_json::json!({"kind": "reauthorize", "account_id": account_id.to_string()})
}
};
let owner = match self.owner_binding.owner() {
AuthorizationOwner::AdminSession { admin_user_id } => {
serde_json::json!({"kind": "admin_session", "admin_user_id": admin_user_id})
}
AuthorizationOwner::AdminApiKey => serde_json::json!({"kind": "admin_api_key"}),
AuthorizationOwner::System => serde_json::json!({"kind": "system"}),
};
Map::from_iter([
(
"provider_kind".to_owned(),
Value::String(self.provider_kind.as_str().to_owned()),
),
("target".to_owned(), target),
("owner".to_owned(), owner),
(
"started_request_id".to_owned(),
Value::String(self.owner_binding.started_request_id().to_owned()),
),
])
}
/// 恢复 adapter 已确认版本为 v1 的中立事务字段。
///
/// # Errors
///
/// 文档字段、Provider kind 或账号身份非法时返回错误,不携带原始 pending 内容。
pub fn from_storage_v1(value: Value) -> Result<Self, AdminError> {
let invalid = || AdminError::invalid("invalid pending authorization mutation");
let document: StoredAuthorizationMutationV1 =
serde_json::from_value(value).map_err(|_| invalid())?;
let provider_kind = ProviderKind::new(document.provider_kind).map_err(|_| invalid())?;
let target = match document.target {
StoredAuthorizationTargetV1::Create { name } => {
AuthorizationMutationTarget::Create { name }
}
StoredAuthorizationTargetV1::Reauthorize { account_id } => {
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new(account_id).map_err(|_| invalid())?,
}
}
};
let owner = match document.owner {
StoredAuthorizationOwnerV1::AdminSession { admin_user_id } => {
AuthorizationOwner::AdminSession { admin_user_id }
}
StoredAuthorizationOwnerV1::AdminApiKey => AuthorizationOwner::AdminApiKey,
StoredAuthorizationOwnerV1::System => AuthorizationOwner::System,
};
Ok(Self::new(
provider_kind,
target,
AuthorizationOwnerBinding {
owner,
started_request_id: document.started_request_id,
},
))
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct StoredAuthorizationMutationV1 {
provider_kind: String,
target: StoredAuthorizationTargetV1,
owner: StoredAuthorizationOwnerV1,
started_request_id: String,
}
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum StoredAuthorizationTargetV1 {
Create { name: String },
Reauthorize { account_id: String },
}
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum StoredAuthorizationOwnerV1 {
AdminSession { admin_user_id: String },
AdminApiKey,
System,
}
/// 启动 Provider OAuth Authorization Code 流程。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StartAuthorization {
@@ -776,6 +873,20 @@ impl ProviderQuota {
.map(|(_, used_percent)| used_percent)
}
/// 选择 Dashboard 展示的真实窗口;没有百分比时仍保留 Provider 的滚动窗口语义。
#[must_use]
pub fn representative_window(&self) -> Option<&ProviderQuotaWindow> {
self.representative_used_window()
.map(|(index, _)| &self.windows[index])
.or_else(|| {
self.windows
.iter()
.enumerate()
.min_by_key(|(index, window)| (quota_usage_priority(window), *index))
.map(|(_, window)| window)
})
}
/// 将已确认的账号级额度耗尽事实投影到展示窗口。
///
/// Provider 已指出具体触顶窗口时只归一化这些窗口;否则归一化 Dashboard
@@ -6,6 +6,7 @@ use std::sync::Arc;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::stream::BoxStream;
use super::backup::BackupStorePorts;
use crate::model::{
@@ -289,6 +290,10 @@ pub trait AccountGroupStore: Send + Sync {
) -> AdminStoreResult<AccountGroupMutation>;
}
/// 逐条读取的已计算费用事实;消费结束或丢弃时释放查询资源。
pub type UsageCalculatedBillingStream<'a> =
BoxStream<'a, AdminStoreResult<UsageCalculatedBillingFact>>;
/// 用量、趋势、诊断与运维错误的只读能力。
#[async_trait]
pub trait ObservabilityStore: Send + Sync {
@@ -318,12 +323,13 @@ pub trait ObservabilityStore: Send + Sync {
filter: UsageFilter,
) -> AdminStoreResult<Vec<RequestMetricPoint>>;
/// 返回可由 Provider 重新校验的已计算费用事实,用于恢复标准费用趋势。
async fn usage_calculated_billing_facts(
/// 流式返回可由 Provider 重新校验的已计算费用事实,不保证顺序。
/// 查询及解码错误由流返回;调用方应逐条聚合,避免收集整个区间。
fn usage_calculated_billing_facts(
&self,
range: TimeRange,
filter: UsageFilter,
) -> AdminStoreResult<Vec<UsageCalculatedBillingFact>>;
) -> UsageCalculatedBillingStream<'_>;
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage>;
@@ -269,6 +269,11 @@ impl DefaultAccountsService {
})
.await
.map_err(|error| map_provider_error(error, "provider quota"))?;
let stored = if refresh_quota {
self.load_account(account_id).await?
} else {
stored
};
self.attach_quota_local_usage(
std::slice::from_ref(&stored),
std::slice::from_mut(&mut quota),
@@ -69,7 +69,7 @@ impl DefaultAuthService {
id: format!("audit_{}", Uuid::now_v7().simple()),
actor_kind: AuditActorKind::AdminSession,
actor_admin_user_id: Some(self.default_admin_user_id.clone()),
actor_ref: format!("admin:{}", self.default_admin_user_id),
actor_ref: crate::model::auth::admin_session_actor_ref(&self.default_admin_user_id),
request_id: None,
action: action.to_owned(),
entity_kind: "admin_session".to_owned(),
@@ -183,7 +183,7 @@ impl AuthService for DefaultAuthService {
if let Some(session) = session {
let mut event = self.auth_audit("admin.logout", Utc::now());
event.actor_admin_user_id = Some(session.admin_user_id.clone());
event.actor_ref = format!("admin:{}", session.admin_user_id);
event.actor_ref = crate::model::auth::admin_session_actor_ref(&session.admin_user_id);
event.entity_ref = session.admin_user_id;
self.store
.append_audit_event(event)
@@ -326,7 +326,9 @@ impl DefaultBackupService {
crate::model::MutationActor::System => AuditActorKind::System,
};
let actor_ref = match &context.actor {
crate::model::MutationActor::AdminSession { admin_user_id } => admin_user_id.clone(),
crate::model::MutationActor::AdminSession { admin_user_id } => {
crate::model::auth::admin_session_actor_ref(admin_user_id)
}
crate::model::MutationActor::AdminApiKey => "admin_api_key".to_owned(),
crate::model::MutationActor::System => "system".to_owned(),
};
@@ -9,6 +9,7 @@ use std::{
use async_trait::async_trait;
use chrono::{DateTime, Duration, Timelike as _, Utc};
use futures::TryStreamExt;
use gateway_core::{account::ProviderAccountId, routing::ProviderKind};
use crate::{
@@ -29,7 +30,7 @@ use crate::{
},
ports::{
provider::{ProviderAdminErrorKind, ProviderAdminRegistry},
store::{ObservabilityStore, SettingsStore},
store::{AccountStore, ObservabilityStore, SettingsStore, UsageCalculatedBillingStream},
},
};
@@ -126,6 +127,7 @@ pub trait ObservabilityService: Send + Sync {
pub(crate) struct DefaultObservabilityService {
store: Arc<dyn ObservabilityStore>,
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
dashboard_queries: DashboardQueryCache,
@@ -135,11 +137,13 @@ impl DefaultObservabilityService {
#[must_use]
pub(crate) fn new(
store: Arc<dyn ObservabilityStore>,
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
) -> Self {
Self {
store,
accounts,
settings,
providers,
dashboard_queries: DashboardQueryCache::default(),
@@ -301,13 +305,21 @@ impl ObservabilityService for DefaultObservabilityService {
range: TimeRange,
filter: UsageFilter,
) -> Result<UsageInsights, AdminError> {
let (overview, trend, billing_facts) = futures::try_join!(
self.store.usage_summary(range, filter.clone()),
self.store.usage_trend(range, filter.clone()),
self.store.usage_calculated_billing_facts(range, filter),
)
.map_err(|error| map_store_error(error, "usage insights"))?;
let standard_costs = recover_standard_costs(&self.providers, billing_facts)?;
let aggregates = async {
futures::try_join!(
self.store.usage_summary(range, filter.clone()),
self.store.usage_trend(range, filter.clone()),
)
.map_err(|error| map_store_error(error, "usage insights"))
};
let ((overview, trend), standard_costs) = futures::try_join!(
aggregates,
recover_standard_costs(
&self.providers,
self.store
.usage_calculated_billing_facts(range, filter.clone()),
),
)?;
build_usage_insights(overview, trend, standard_costs)
}
@@ -590,12 +602,16 @@ struct UsageCostScenarios {
no_cache: ScenarioCosts,
}
fn recover_standard_costs(
async fn recover_standard_costs(
providers: &ProviderAdminRegistry,
facts: Vec<UsageCalculatedBillingFact>,
mut facts: UsageCalculatedBillingStream<'_>,
) -> Result<UsageCostScenarios, AdminError> {
let mut scenarios = UsageCostScenarios::default();
for fact in facts {
while let Some(fact) = facts
.try_next()
.await
.map_err(|error| map_store_error(error, "usage billing facts"))?
{
let Ok(provider_kind) = ProviderKind::new(fact.provider_kind.clone()) else {
continue;
};
@@ -708,8 +724,31 @@ fn usd_cost(costs: &[CurrencyCost]) -> Option<DecimalAmount> {
impl DefaultObservabilityService {
async fn enrich_dashboard_quotas(&self, accounts: &mut [DashboardAccountUsage]) {
if accounts.is_empty() {
return;
}
let now = Utc::now();
let ids = accounts
.iter()
.map(|account| account.account_id.clone())
.collect::<Vec<_>>();
let mut rolling_usage = self
.accounts
.load_account_usage(
TimeRange {
start: now - Duration::hours(24),
end: now,
},
&ids,
)
.await
.unwrap_or_default()
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<HashMap<_, _>>();
let providers = self.providers.clone();
let quota_reads = accounts.iter().map(|account| {
let rolling_usage = rolling_usage.remove(&account.account_id);
let account_id = ProviderAccountId::new(account.account_id.clone()).ok();
let provider = ProviderKind::new(account.provider_kind.clone())
.ok()
@@ -722,16 +761,19 @@ impl DefaultObservabilityService {
.quota(ProviderQuotaRequest {
account_id,
refresh: false,
rolling_usage: None,
rolling_usage,
})
.await
.ok()
.and_then(|quota| quota.representative_used_percent())
}
});
let quota_used = futures::future::join_all(quota_reads).await;
for (account, used_percent) in accounts.iter_mut().zip(quota_used) {
account.quota_used_percent = used_percent;
let quotas = futures::future::join_all(quota_reads).await;
for (account, quota) in accounts.iter_mut().zip(quotas) {
if let Some(mut quota) = quota {
quota.apply_limit_reached_display();
account.quota_used_percent = quota.representative_used_percent();
account.quota_window = quota.representative_window().cloned();
}
}
}
@@ -1,3 +1,4 @@
//! 模型模块测试:镜像 src/model/ 中带逻辑的模块。
mod backup;
mod provider_credentials;
@@ -0,0 +1,134 @@
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwnerBinding, PendingAuthorizationMutation,
ProviderQuota, ProviderQuotaWindow, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
};
use gateway_admin::model::{MutationActor, MutationContext};
use gateway_core::account::ProviderAccountId;
use gateway_core::routing::ProviderKind;
use serde_json::{Value, json};
#[test]
fn dashboard_quota_should_preserve_unknown_and_actual_window_facts() {
let mut quota = ProviderQuota {
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
limit_reached: false,
provider_data: None,
};
assert!(quota.representative_window().is_none());
assert!(quota.representative_used_percent().is_none());
let rolling = ProviderQuotaWindow {
key: "provider_rolling".to_owned(),
group: "requests".to_owned(),
label: "Provider rolling window".to_owned(),
limit_id: None,
limit_name: None,
role: Some(ProviderQuotaWindowRole::Primary),
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(86_400),
used_percent: None,
reset_at: None,
limit_reached: false,
local_usage: None,
provider_data: None,
};
quota.windows.push(rolling.clone());
assert_eq!(quota.representative_window(), Some(&rolling));
assert!(quota.representative_used_percent().is_none());
quota.windows[0].used_percent = Some(99.6);
quota.apply_limit_reached_display();
let actual = quota.representative_window().expect("actual window");
assert_eq!(actual.used_percent, Some(99.6));
assert!(
!actual.limit_reached,
"display rounding cannot exhaust quota"
);
assert_eq!(actual.window_seconds, Some(86_400));
assert_eq!(actual.reset_at, None);
quota.windows.insert(
0,
ProviderQuotaWindow {
key: "model_specific".to_owned(),
used_percent: Some(100.0),
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
..rolling
},
);
assert_eq!(
quota
.representative_window()
.map(|window| window.key.as_str()),
Some("provider_rolling")
);
}
#[test]
fn pending_mutation_v1_should_round_trip_all_targets_and_owners() {
for provider in ["openai", "xai"] {
for target in [
AuthorizationMutationTarget::Create {
name: "account".to_owned(),
},
AuthorizationMutationTarget::Reauthorize {
account_id: ProviderAccountId::new("acct_1").expect("account id"),
},
] {
for actor in [
MutationActor::AdminSession {
admin_user_id: "admin-1".to_owned(),
},
MutationActor::AdminApiKey,
MutationActor::System,
] {
let expected = PendingAuthorizationMutation::new(
ProviderKind::new(provider).expect("provider"),
target.clone(),
AuthorizationOwnerBinding::from_context(&MutationContext {
actor,
request_id: "request-1".to_owned(),
}),
);
assert_eq!(
PendingAuthorizationMutation::from_storage_v1(Value::Object(
expected.to_storage_v1()
)),
Ok(expected),
);
}
}
}
}
#[test]
fn pending_mutation_v1_should_keep_existing_storage_fields_and_reject_invalid_documents() {
let legacy = json!({
"provider_kind": "openai",
"target": { "kind": "reauthorize", "account_id": "acct_1" },
"owner": { "kind": "admin_session", "admin_user_id": "admin-1" },
"started_request_id": "request-1"
});
let restored =
PendingAuthorizationMutation::from_storage_v1(legacy.clone()).expect("v1 document");
assert_eq!(Value::Object(restored.to_storage_v1()), legacy);
for (pointer, invalid) in [
("/provider_kind", json!("")),
("/target/account_id", json!("")),
("/target/kind", json!("other")),
("/owner/kind", json!("other")),
("/started_request_id", json!(null)),
] {
let mut document = legacy.clone();
*document.pointer_mut(pointer).expect("field") = invalid;
assert!(
PendingAuthorizationMutation::from_storage_v1(document).is_err(),
"{pointer}"
);
}
let mut document = legacy;
document["schema_version"] = json!(2);
assert!(PendingAuthorizationMutation::from_storage_v1(document).is_err());
}
@@ -72,6 +72,7 @@ pub(super) struct FakeProviderAdmin {
import_account_ids: Mutex<Vec<String>>,
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
quota: Mutex<ProviderQuota>,
quota_refresh_account: Mutex<Option<(Arc<FakeAccountStore>, AccountRecord)>>,
current_credential_revision: Mutex<Revision>,
reset_credit_commands: Mutex<Vec<ConsumeProviderResetCredit>>,
}
@@ -89,6 +90,7 @@ impl FakeProviderAdmin {
import_account_ids: Mutex::new(vec!["acct_prepared".to_owned()]),
quota_requests: Mutex::new(Vec::new()),
quota: Mutex::new(empty_quota()),
quota_refresh_account: Mutex::new(None),
current_credential_revision: Mutex::new(revision(1)),
reset_credit_commands: Mutex::new(Vec::new()),
})
@@ -346,6 +348,14 @@ impl ProviderAdmin for FakeProviderAdmin {
) -> Result<ProviderQuota, ProviderAdminError> {
if request.refresh {
self.record("provider.quota");
if let Some((store, account)) = self
.quota_refresh_account
.lock()
.expect("quota refresh account")
.take()
{
store.set_accounts(vec![account]);
}
}
self.quota_requests
.lock()
@@ -1098,6 +1108,47 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
assert_eq!(store.audit_requests(), ["refresh-request"]);
}
#[tokio::test]
async fn accounts_quota_refresh_should_return_the_updated_account_status() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let mut account = account_record("openai");
account.quota = QuotaState::exhausted(
QuotaEvidence::UsageLimitReached,
std::time::SystemTime::now(),
None,
);
let store = FakeAccountStore::with_account(account.clone(), events.clone());
account.quota = QuotaState::allowed(std::time::SystemTime::now());
*provider
.quota_refresh_account
.lock()
.expect("quota refresh account") = Some((store.clone(), account));
let result = accounts_service(provider, store)
.await
.accounts()
.quota(
&ProviderAccountId::new("acct_test").expect("account ID"),
true,
)
.await
.expect("refresh quota");
assert_eq!(
result.account.quota.access(),
gateway_core::account::QuotaAccessState::Allowed
);
assert_eq!(
result.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
assert_eq!(
recorded(&events),
["store.load_account", "provider.quota", "store.load_account"]
);
}
#[tokio::test]
async fn accounts_recover_should_commit_facts_then_return_normal_account() {
let events = events();
@@ -536,13 +536,14 @@ impl ObservabilityStore for UnavailableStore {
Err(unavailable("usage trend"))
}
async fn usage_calculated_billing_facts(
fn usage_calculated_billing_facts(
&self,
_: TimeRange,
_: UsageFilter,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::UsageCalculatedBillingFact>>
{
Err(unavailable("usage billing facts"))
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
Box::pin(futures::stream::once(async {
Err(unavailable("usage billing facts"))
}))
}
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
@@ -2,7 +2,7 @@ use std::{
str::FromStr as _,
sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
atomic::{AtomicBool, AtomicUsize, Ordering},
},
time::Duration as StdDuration,
};
@@ -537,10 +537,41 @@ async fn usage_records_should_enrich_provider_reported_totals_when_pricing_match
));
}
#[tokio::test]
async fn usage_insights_should_reject_partial_costs_when_billing_stream_fails() {
let now = Utc::now();
let range = observation_range(now);
let store = Arc::new(FixtureObservabilityStore::new(range));
store.replace_calculated_billing_facts(vec![UsageCalculatedBillingFact {
bucket_start: quarter_hour_start(now),
provider_kind: "openai".to_owned(),
upstream_model_id: "gpt-5.5".to_owned(),
service_tier: None,
input_tokens: Some(800),
output_tokens: Some(200),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
total: CurrencyCost {
currency: "USD".to_owned(),
amount: "1.25".parse().expect("cost"),
},
}]);
store.billing_stream_fails.store(true, Ordering::SeqCst);
let services = observability_services_with_calculated_billing(store).await;
assert!(
services
.observability()
.usage_insights(range, UsageFilter::default())
.await
.is_err()
);
}
struct FixtureObservabilityStore {
trend: Mutex<Vec<RequestMetricPoint>>,
overview: Mutex<UsageOverview>,
calculated_billing_facts: Mutex<Vec<UsageCalculatedBillingFact>>,
billing_stream_fails: AtomicBool,
diagnostics: Mutex<Vec<DiagnosticObservation>>,
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
summary_observed_at: Mutex<Option<DateTime<Utc>>>,
@@ -561,6 +592,7 @@ impl FixtureObservabilityStore {
providers: Vec::new(),
}),
calculated_billing_facts: Mutex::new(Vec::new()),
billing_stream_fails: AtomicBool::new(false),
diagnostics: Mutex::new(Vec::new()),
runtime_slots: Mutex::new(None),
summary_observed_at: Mutex::new(None),
@@ -663,16 +695,23 @@ impl ObservabilityStore for FixtureObservabilityStore {
Ok(self.trend.lock().expect("trend").clone())
}
async fn usage_calculated_billing_facts(
fn usage_calculated_billing_facts(
&self,
_: TimeRange,
_: UsageFilter,
) -> AdminStoreResult<Vec<UsageCalculatedBillingFact>> {
Ok(self
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
let facts = self
.calculated_billing_facts
.lock()
.expect("calculated billing facts")
.clone())
.clone();
let error = self
.billing_stream_fails
.load(Ordering::SeqCst)
.then(|| Err(super::unavailable("billing stream failed")));
Box::pin(futures::stream::iter(
facts.into_iter().map(Ok).chain(error),
))
}
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
@@ -70,6 +70,7 @@ mod wire;
pub use credentials::*;
pub use handlers::{profile_avatar_response, router};
pub(super) use presenter::quota_window_view;
pub use wire::*;
use credentials::{
@@ -140,7 +140,7 @@ pub(super) fn account_quota_view(
)
}
pub(super) fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
pub(crate) fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
let ProviderQuotaWindow {
key,
group,
@@ -230,6 +230,8 @@ pub(super) fn account_usage_view(
output_tokens_display: display_optional_tokens(usage.output_tokens),
cached_tokens: usage.cached_tokens,
cached_tokens_display: display_optional_tokens(usage.cached_tokens),
reasoning_tokens: usage.reasoning_tokens,
reasoning_tokens_display: display_optional_tokens(usage.reasoning_tokens),
image_input_tokens: usage.image_input_tokens,
image_input_tokens_display: display_optional_tokens(usage.image_input_tokens),
image_output_tokens: usage.image_output_tokens,
@@ -347,6 +349,8 @@ pub(super) fn empty_account_usage() -> AccountUsageView {
output_tokens_display: "—".to_owned(),
cached_tokens: None,
cached_tokens_display: "—".to_owned(),
reasoning_tokens: None,
reasoning_tokens_display: "—".to_owned(),
image_input_tokens: None,
image_input_tokens_display: "—".to_owned(),
image_output_tokens: None,
@@ -276,6 +276,8 @@ pub struct AccountUsageView {
pub output_tokens_display: String,
pub cached_tokens: Option<u64>,
pub cached_tokens_display: String,
pub reasoning_tokens: Option<u64>,
pub reasoning_tokens_display: String,
pub image_input_tokens: Option<u64>,
pub image_input_tokens_display: String,
pub image_output_tokens: Option<u64>,
@@ -505,6 +505,8 @@ pub(crate) fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptV
pub(crate) fn usage_detail_view(detail: domain::UsageDetail) -> UsageRecordDetailView {
UsageRecordDetailView {
trace: detail.trace,
related_requests: detail.related_requests,
request: usage_record_view(detail.request),
attempts: detail
.attempts
@@ -560,7 +562,6 @@ pub(crate) fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
failure_class: error.failure_kind,
upstream_send_state: error.upstream_send_state,
provider_error_code: error.provider_error_code,
occurrence_count: error.occurrence_count,
response_id: error.client_response_id,
upstream_request_id: error.upstream_request_id,
latency_ms: error.latency_ms,
@@ -864,6 +865,28 @@ pub(crate) fn dashboard_view(
} = observation;
let mut account_usage_views = Vec::with_capacity(account_usage.len());
for credential in account_usage {
let local_requests = credential
.quota_window
.as_ref()
.filter(|window| window.used_percent.is_none())
.and_then(|window| window.local_usage.as_ref())
.map(|usage| usage.request_count);
let (metric_label, metric_value) = match local_requests {
Some(count) => (
"次数".to_owned(),
if count > 0 {
format_compact_number(count)
} else {
"—".to_owned()
},
),
None => (
"今日 Token".to_owned(),
credential
.total_tokens
.map_or_else(|| "—".to_owned(), format_compact_number),
),
};
account_usage_views.push(DashboardAccountUsageView {
id: credential.account_id.clone(),
provider: credential.provider_kind.clone(),
@@ -886,6 +909,11 @@ pub(crate) fn dashboard_view(
})
.collect(),
quota_used_percent: credential.quota_used_percent,
usage_window: credential
.quota_window
.map(crate::admin::accounts::quota_window_view),
metric_label,
metric_value,
last_used: relative_time(credential.last_used_at, range.end),
});
}
@@ -247,6 +247,8 @@ pub struct UsageAttemptView {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageRecordDetailView {
pub trace: Option<serde_json::Value>,
pub related_requests: Vec<serde_json::Value>,
#[serde(flatten)]
pub request: UsageRecordView,
pub attempts: Vec<UsageAttemptView>,
@@ -375,6 +377,9 @@ pub struct DashboardAccountUsageView {
pub request_count: u64,
pub request_buckets: Vec<DashboardAccountRequestBucketView>,
pub quota_used_percent: Option<f64>,
pub usage_window: Option<crate::admin::accounts::AccountQuotaWindowView>,
pub metric_label: String,
pub metric_value: String,
pub last_used: String,
}
@@ -778,7 +783,6 @@ pub struct OpsErrorView {
pub failure_class: String,
pub upstream_send_state: Option<String>,
pub provider_error_code: Option<String>,
pub occurrence_count: u32,
pub response_id: Option<String>,
pub upstream_request_id: Option<String>,
pub latency_ms: Option<u64>,
+8 -5
View File
@@ -180,13 +180,16 @@ pub fn initialize(
"request",
request_id = %request_id,
method = %request.method().as_str(),
uri = %request
.uri()
.path_and_query()
.map(axum::http::uri::PathAndQuery::as_str)
.unwrap_or_default(),
uri = %request.uri().path(),
)
}
})
.on_request(|_request: &Request<_>, _span: &tracing::Span| {
tracing::info!(target: "request_trace", stage = "http.received", "HTTP request received");
})
.on_response(|response: &axum::http::Response<_>, latency: Duration, _span: &tracing::Span| {
tracing::info!(target: "request_trace", stage = "http.response", status = response.status().as_u16(),
headers_ms = latency.as_millis(), "HTTP response headers ready");
});
let router = router
.layer(PropagateRequestIdLayer::new(request_id_header.clone()))
@@ -10,14 +10,14 @@ use axum::{
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{ImageRequest, ImageRequestKind, Operation, RawJsonPayload};
use serde_json::{Map, Value};
use serde_json::Value;
use crate::ApiState;
use crate::openai::{
auth::{authenticate_client, client_access_error_response},
error::gateway_error_response,
provider_endpoint::collect_raw_json_response,
responses::request_client_context,
responses::{OpenAiRequestHeaders, request_client_context},
};
const OPENAI_PROTOCOL: &str = "openai";
@@ -95,7 +95,7 @@ fn image_operation(
headers: &HeaderMap,
kind: ImageRequestKind,
) -> Result<Operation, GatewayError> {
let mut context = Map::new();
let mut context = OpenAiRequestHeaders::from_headers(headers).session_context();
if let Some(turn_id) = headers
.get("x-codex-image-turn-id")
.and_then(|value| value.to_str().ok())
@@ -3,6 +3,7 @@
use std::collections::VecDeque;
use std::convert::Infallible;
use std::net::{IpAddr, SocketAddr};
use std::time::Duration;
use axum::{
body::{Body, Bytes},
@@ -13,9 +14,10 @@ use axum::{
},
response::{IntoResponse, Response},
};
use futures::stream;
use gateway_core::engine::CommitRequirement;
use futures::{StreamExt, stream};
use gateway_core::diagnostics::TraceContext;
use gateway_core::engine::execution::{ClientTransport, ExecutionSession, StartedExecution};
use gateway_core::engine::{CommitRequirement, EngineError};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
use gateway_core::lifecycle::ConnectionGuard;
@@ -35,10 +37,13 @@ use super::{
request::decode_request_with_headers,
};
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
/// `POST /v1/responses`。
pub(crate) async fn responses(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
ingress_id: Option<Extension<tower_http::request_id::RequestId>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
@@ -83,15 +88,25 @@ pub(crate) async fn responses(
Ok(started) => started,
Err(error) => return gateway_error_response(&error),
};
let trace = started.session.trace();
trace.headers("client.request", serde_json::json!({
"ingressRequestId": ingress_id.as_ref().and_then(|Extension(id)| id.header_value().to_str().ok()),
}), headers.iter().map(|(name, value)| (name.as_str(), value.as_bytes())));
trace.capture("client.request.body", &body);
let request_id_header = HeaderValue::from_str(started.request_id.as_str()).ok();
let StartedExecution {
stream, session, ..
} = started;
if stream {
let mut response = if stream {
stream_execution_response(session, connection_guard).await
} else {
drop(connection_guard);
collect_execution_response(session).await
};
if let Some(id) = request_id_header {
response.headers_mut().insert("x-gateway-request-id", id);
}
response
}
/// 从 socket 与标准转发头提取旧 Usage 页面使用的诊断事实。
@@ -152,7 +167,8 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
let events = match session.collect_uncommitted().await {
Ok(events) => events,
Err(error) => {
let response = engine_error_response(&error);
let response_headers = session.response_headers().to_vec();
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
};
@@ -168,6 +184,11 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
}
};
let response_headers = session.response_headers().to_vec();
session.trace().record(
"downstream.encoded",
serde_json::json!({"transport": "http_json", "bytes": encoded.len()}),
);
session.trace().dump("downstream.body", &encoded);
let response = json_body_response(encoded, &response_headers);
let Some(session) = execution.session_mut() else {
@@ -178,7 +199,7 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
.await
{
execution.cancel_and_finalize().await;
let response = engine_error_response(&error);
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
if !session.is_finalized() {
@@ -246,6 +267,13 @@ fn apply_response_headers(
response
}
fn engine_error_response_with_headers(
error: &EngineError,
response_headers: &[ProviderResponseHeader],
) -> Response {
apply_response_headers(engine_error_response(error), response_headers)
}
/// 编码首个 SSE frame 后提交下游,再持续驱动同一执行会话。
pub async fn stream_execution_response(
session: Box<dyn ExecutionSession>,
@@ -267,7 +295,8 @@ pub async fn stream_execution_response(
return response;
}
Err(error) => {
let response = engine_error_response(&error);
let response_headers = session.response_headers().to_vec();
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
};
@@ -299,7 +328,7 @@ pub async fn stream_execution_response(
.await
{
execution.cancel_and_finalize().await;
let response = engine_error_response(&error);
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
let Some(session) = execution.into_session() else {
@@ -309,9 +338,11 @@ pub async fn stream_execution_response(
if state.encoder.is_completed() {
state.finish_completed(Vec::new()).await;
}
let body = Body::from_stream(stream::unfold(state, |mut state| async move {
let output = Box::pin(stream::unfold(state, |mut state| async move {
loop {
if let Some(chunk) = state.pending.pop_front() {
state.trace.dump("downstream.chunk", &chunk);
state.handed_off_bytes += chunk.len() as u64;
return Some((Ok::<Bytes, Infallible>(chunk), state));
}
if state.output_finished {
@@ -320,19 +351,35 @@ pub async fn stream_execution_response(
state.advance().await;
}
}));
// 保活只轮询已固定的输出 stream;不能取消并重建 advance/next_event future,
// 否则一次心跳就可能丢失正在等待的上游事件或执行终态清理。
let body = Body::from_stream(stream::unfold(output, |mut output| async move {
let chunk = tokio::select! {
biased;
chunk = output.next() => chunk?,
() = tokio::time::sleep(SSE_KEEPALIVE_INTERVAL) => {
Ok(Bytes::from_static(b": keep-alive\n\n"))
}
};
Some((chunk, output))
}));
event_stream_response(body, &response_headers)
}
fn event_stream_response(body: Body, response_headers: &[ProviderResponseHeader]) -> Response {
let mut response = Response::new(body);
let mut response = apply_response_headers(Response::new(body), response_headers);
*response.status_mut() = StatusCode::OK;
response
.headers_mut()
.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
response.headers_mut().insert(
CACHE_CONTROL,
HeaderValue::from_static("no-cache, no-transform"),
);
response
.headers_mut()
.insert(CACHE_CONTROL, HeaderValue::from_static("no-cache"));
apply_response_headers(response, response_headers)
.insert("x-accel-buffering", HeaderValue::from_static("no"));
response
}
fn internal_gateway_response(message: &'static str) -> Response {
@@ -394,12 +441,18 @@ impl Drop for PendingExecution {
if session.is_finalized() {
return;
}
session.trace().record(
"downstream.cancelled",
serde_json::json!({"reason": "response_guard_dropped"}),
);
session.cancel();
detach_finalize(session);
}
}
struct ResponsesStreamState {
trace: TraceContext,
handed_off_bytes: u64,
session: Option<Box<dyn ExecutionSession>>,
encoder: OpenAiResponsesEncoder,
pending: VecDeque<Bytes>,
@@ -416,6 +469,8 @@ impl ResponsesStreamState {
connection_guard: Option<Box<dyn ConnectionGuard>>,
) -> Self {
Self {
trace: session.trace(),
handed_off_bytes: 0,
session: Some(session),
encoder,
pending: initial_frames.into_iter().collect(),
@@ -571,6 +626,11 @@ impl ResponsesStreamState {
impl Drop for ResponsesStreamState {
fn drop(&mut self) {
self.trace.record("downstream.body.closed", serde_json::json!({
"handedOffBytes": self.handed_off_bytes,
"pendingChunks": self.pending.len(),
"outputFinished": self.output_finished, "executionTerminal": self.execution_terminal,
}));
if self.execution_terminal {
return;
}
@@ -16,6 +16,7 @@ const OPENAI_PROTOCOL: &str = "openai";
const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent";
const CODEX_TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
const DOWNSTREAM_WEBSOCKET_CONNECTION_ID_CONTEXT_KEY: &str = "downstream_websocket_connection_id";
/// Responses 请求进入共享解码内核时的下游传输来源。
#[derive(Clone, Copy)]
@@ -35,6 +36,7 @@ pub struct OpenAiRequestHeaders {
include_timing_metrics: Option<String>,
codex_window_id: Option<String>,
parent_thread_id: Option<String>,
downstream_websocket_connection_id: Option<String>,
conversation_id: Option<String>,
session_id: Option<String>,
@@ -60,6 +62,7 @@ impl OpenAiRequestHeaders {
include_timing_metrics: header_string(headers, "x-responsesapi-include-timing-metrics"),
codex_window_id: header_string(headers, "x-codex-window-id"),
parent_thread_id: header_string(headers, "x-codex-parent-thread-id"),
downstream_websocket_connection_id: None,
conversation_id: header_string(headers, "conversation-id")
.or_else(|| header_string(headers, "conversation_id")),
session_id: header_string(headers, "session-id")
@@ -74,12 +77,31 @@ impl OpenAiRequestHeaders {
}
}
/// 绑定代理为当前下游 WebSocket 分配的连接身份。
#[must_use]
pub(super) fn with_downstream_websocket_connection_id(
mut self,
connection_id: impl Into<String>,
) -> Self {
self.downstream_websocket_connection_id = Some(connection_id.into());
self
}
fn apply_subagent(&self, body: &mut Map<String, Value>) {
if let Some(subagent) = self.subagent.as_deref() {
inject_subagent_metadata(body, subagent);
}
}
/// 独立端点只消费会话身份与 turn metadata,不携带 Responses 的连接状态。
pub(crate) fn session_context(&self) -> Map<String, Value> {
let mut context = Map::new();
insert_protocol_context(&mut context, "session_id", self.session_id.as_ref());
insert_protocol_context(&mut context, "thread_id", self.thread_id.as_ref());
insert_protocol_context(&mut context, "turn_metadata", self.turn_metadata.as_ref());
context
}
fn protocol_context(
&self,
use_websocket: Option<bool>,
@@ -109,6 +131,11 @@ impl OpenAiRequestHeaders {
"parent_thread_id",
self.parent_thread_id.as_ref(),
);
insert_protocol_context(
&mut context,
DOWNSTREAM_WEBSOCKET_CONNECTION_ID_CONTEXT_KEY,
self.downstream_websocket_connection_id.as_ref(),
);
insert_protocol_context(
&mut context,
"conversation_id",
@@ -465,6 +492,8 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
// 下游鉴权和账号 cookie 绝不能成为上游账号身份。
| "authorization"
| "x-api-key"
// Codex 的服务端托管认证标记只用于客户端能力判断,不代表上游身份。
| "x-openai-actor-authorization"
| "cookie"
| "cookie2"
| "chatgpt-account-id"
@@ -195,6 +195,7 @@ pub struct ResponsesWebSocketConnection {
expired: Arc<AtomicBool>,
commands: Option<mpsc::Sender<ConnectionCommand>>,
incoming: mpsc::Receiver<ConnectionEvent>,
exited: oneshot::Receiver<PumpExitReason>,
pump_task: Option<JoinHandle<()>>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
@@ -231,7 +232,16 @@ impl ResponsesWebSocketConnection {
/// 等待下一个需要业务层处理的客户端事件。
pub async fn next_event(&mut self) -> Option<ConnectionEvent> {
let event = self.incoming.recv().await;
if let Some(reason) = self.exit_reason {
return Some(ConnectionEvent::Exited(reason));
}
let event = tokio::select! {
biased;
reason = &mut self.exited => {
Some(ConnectionEvent::Exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
}
event = self.incoming.recv() => event,
};
if let Some(ConnectionEvent::Exited(reason)) = event.as_ref() {
self.exit_reason.get_or_insert(*reason);
} else if event.is_none() {
@@ -240,6 +250,18 @@ impl ResponsesWebSocketConnection {
event
}
/// 等待连接退出,不消费留给后续串行请求的业务帧。
pub async fn wait_for_exit(&mut self) -> PumpExitReason {
if let Some(reason) = self.exit_reason {
return reason;
}
let reason = (&mut self.exited)
.await
.unwrap_or(PumpExitReason::PumpStopped);
self.exit_reason = Some(reason);
reason
}
/// 串行写入文本帧,并等待 pump 确认 transport 写入结果。
///
/// # Errors
@@ -425,10 +447,12 @@ where
let stats = Arc::new(ConnectionStats::default());
let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER);
let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER);
let (exit_tx, exit_rx) = oneshot::channel();
let pump_task = tokio::spawn(run_pump(
socket,
command_rx,
incoming_tx,
exit_tx,
Arc::clone(&connection_id),
cancellation,
opened_at,
@@ -442,6 +466,7 @@ where
expired,
commands: Some(command_tx),
incoming: incoming_rx,
exited: exit_rx,
pump_task: Some(pump_task),
stats,
config,
@@ -454,6 +479,7 @@ async fn run_pump<S, E>(
mut socket: S,
mut commands: mpsc::Receiver<ConnectionCommand>,
incoming: mpsc::Sender<ConnectionEvent>,
exited: oneshot::Sender<PumpExitReason>,
connection_id: Arc<str>,
cancellation: CancellationToken,
opened_at: Instant,
@@ -474,8 +500,13 @@ async fn run_pump<S, E>(
() = &mut deadline, if !deadline_elapsed => {
deadline_elapsed = true;
expired.store(true, Ordering::Release);
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Expired) {
break reason;
// 满队列本身已能唤醒空闲协调层;到期标志仍阻止启动下一轮。
// 不能因无法追加 Expired 通知而中断正在收尾的响应。
match incoming.try_send(ConnectionEvent::Expired) {
Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
Err(mpsc::error::TrySendError::Closed(_)) => {
break PumpExitReason::CoordinatorDropped;
}
}
}
command = commands.recv() => {
@@ -550,7 +581,8 @@ async fn run_pump<S, E>(
},
}
};
let _ = incoming.try_send(ConnectionEvent::Exited(reason));
// 退出不能排在待执行请求之后,也不能因业务队列已满而丢失。
let _ = exited.send(reason);
tracing::debug!(
websocket_connection_id = %connection_id,
pump_exit_reason = reason.as_str(),
@@ -13,12 +13,10 @@ use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
use super::{
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, PendingExecution, ProtocolErrorBody},
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
connection::{FramePhase, ResponsesWebSocketConnection, WriteContext},
protocol::{error_event, response_metadata_event},
};
const ACTIVE_RESPONSE_VIOLATION: &str = "Only one response.create may be active per connection";
#[derive(Clone, Copy, PartialEq, Eq)]
pub(super) enum ForwardOutcome {
Continue,
@@ -66,29 +64,20 @@ pub(super) async fn forward_execution(
);
return send_gateway_error(connection, &error, &request_id).await;
}
let first = loop {
match next_active_input(connection, &mut execution).await {
ActiveInput::Event(Ok(Some(event))) => break event,
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response ended before its first event",
);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Control => continue,
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
ActiveInput::ProtocolViolation => {
connection
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(&request_id))
.await;
return ForwardOutcome::Disconnect;
}
let first = match next_active_input(connection, &mut execution).await {
ActiveInput::Event(Ok(Some(event))) => event,
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response ended before its first event",
);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
return send_initial_engine_error(connection, &mut execution, &error, &request_id)
.await;
}
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
};
let requirement = first.commit_requirement();
let mut first = first.into_provider_events();
@@ -259,18 +248,35 @@ pub(super) async fn forward_execution(
let error = gateway_error_from_engine(&error);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Control => {}
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
ActiveInput::ProtocolViolation => {
connection
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(&request_id))
.await;
return ForwardOutcome::Disconnect;
}
}
}
}
async fn send_initial_engine_error(
connection: &mut ResponsesWebSocketConnection,
execution: &mut PendingExecution,
error: &EngineError,
request_id: &Arc<str>,
) -> ForwardOutcome {
let response_headers = execution
.session_mut()
.map(|session| session.response_headers().to_vec())
.unwrap_or_default();
if !response_headers.is_empty()
&& !send_metadata(
connection,
request_id,
response_metadata_event(request_id, &response_headers),
)
.await
{
return ForwardOutcome::Disconnect;
}
let error = gateway_error_from_engine(error);
send_gateway_error(connection, &error, request_id).await
}
fn commit_connection_replay(
replay: &mut ConnectionReplaySnapshot,
encoder: &OpenAiResponsesEncoder,
@@ -286,50 +292,39 @@ async fn confirm_completed_execution(
execution: &mut PendingExecution,
request_id: &Arc<str>,
) -> Result<(), ForwardOutcome> {
loop {
match next_active_input(connection, execution).await {
ActiveInput::Event(Ok(None))
if execution
.session_mut()
.is_some_and(|session| session.is_finalized()) =>
{
return Ok(());
}
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response was not finalized after its terminal event",
);
return Err(send_gateway_error(connection, &error, request_id).await);
}
ActiveInput::Event(Ok(Some(_))) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response continued after its terminal event",
);
return Err(send_gateway_error(connection, &error, request_id).await);
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return Err(send_gateway_error(connection, &error, request_id).await);
}
ActiveInput::Control => {}
ActiveInput::Disconnect => return Err(ForwardOutcome::Disconnect),
ActiveInput::ProtocolViolation => {
connection
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(request_id))
.await;
return Err(ForwardOutcome::Disconnect);
}
match next_active_input(connection, execution).await {
ActiveInput::Event(Ok(None))
if execution
.session_mut()
.is_some_and(|session| session.is_finalized()) =>
{
Ok(())
}
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response was not finalized after its terminal event",
);
Err(send_gateway_error(connection, &error, request_id).await)
}
ActiveInput::Event(Ok(Some(_))) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response continued after its terminal event",
);
Err(send_gateway_error(connection, &error, request_id).await)
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
Err(send_gateway_error(connection, &error, request_id).await)
}
ActiveInput::Disconnect => Err(ForwardOutcome::Disconnect),
}
}
enum ActiveInput {
Event(Result<Option<CoordinatedEvent>, EngineError>),
Control,
Disconnect,
ProtocolViolation,
}
async fn next_active_input(
@@ -339,15 +334,17 @@ async fn next_active_input(
let Some(session) = execution.session_mut() else {
return ActiveInput::Disconnect;
};
// 与 Codex stream_request 的连接锁一致:本轮结束前不消费下一条请求。
// pump 继续接收有界业务帧和处理 Ping/Pong;退出通过独立通知取消本轮。
tokio::select! {
event = session.next_event() => ActiveInput::Event(event),
message = connection.next_event() => match message {
Some(ConnectionEvent::Expired) => ActiveInput::Control,
Some(ConnectionEvent::Text(_) | ConnectionEvent::Binary) => {
ActiveInput::ProtocolViolation
}
Some(ConnectionEvent::Exited(_)) | None => ActiveInput::Disconnect,
biased;
_ = connection.wait_for_exit() => {
session.trace().record("downstream.cancelled", serde_json::json!({
"reason": "websocket_connection_exit", "connectionId": connection.id(),
}));
ActiveInput::Disconnect
},
event = session.next_event() => ActiveInput::Event(event),
}
}
@@ -459,6 +456,7 @@ fn log_terminal_write_success(
provider_terminal_to_write: std::time::Duration,
) {
tracing::info!(
target: "request_trace", stage = "downstream.terminal.written",
websocket_connection_id = connection.id(),
request_id = %request_id,
provider_terminal_to_terminal_write_ms = provider_terminal_to_write.as_millis(),
@@ -473,6 +471,7 @@ fn log_terminal_write_failure(
provider_terminal_to_write: std::time::Duration,
) {
tracing::info!(
target: "request_trace", stage = "downstream.terminal.write_failed",
websocket_connection_id = connection.id(),
request_id = %request_id,
provider_terminal_to_terminal_write_ms = provider_terminal_to_write.as_millis(),
@@ -19,6 +19,7 @@ use axum::{
response::{IntoResponse, Response},
};
use gateway_core::{
diagnostics::{TraceContext, diagnostic_json},
engine::execution::{AuthenticatedClient, ClientTransport},
lifecycle::{ConnectionGuard, ConnectionLifecycle},
};
@@ -60,12 +61,14 @@ pub(crate) async fn responses_websocket(
&headers,
connect_info.map(|Extension(ConnectInfo(address))| address),
);
let request_headers = OpenAiRequestHeaders::from_headers(&headers);
ResponsesWebSocketAdapter::new(service).upgrade_with_client_context(
websocket,
client,
client_ip,
user_agent,
OpenAiRequestHeaders::from_headers(&headers),
request_headers,
headers,
)
}
@@ -89,15 +92,24 @@ impl ResponsesWebSocketAdapter {
client_ip: Option<IpAddr>,
user_agent: Option<String>,
request_headers: OpenAiRequestHeaders,
raw_headers: HeaderMap,
) -> Response {
let connection_guard = match self.service.try_register_connection() {
Ok(guard) => guard,
Err(_) => return runtime_unavailable_response().into_response(),
};
let connection_id = self.service.next_request_id().replacen("req_", "ws_", 1);
TraceContext::new(&connection_id).headers(
"client.connection.headers",
serde_json::json!({"transport": "websocket", "method": "GET", "path": "/v1/responses"}),
raw_headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_bytes())),
);
let session = ResponsesWebSocketSession {
service: self.service.clone(),
client,
connection_id: self.service.next_request_id().replacen("req_", "ws_", 1),
connection_id,
client_ip,
user_agent,
request_headers,
@@ -143,6 +155,8 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
);
let _connection_guard = connection_guard;
let cancellation = lifecycle.cancellation();
let request_headers =
request_headers.with_downstream_websocket_connection_id(connection_id.clone());
let mut connection = ResponsesWebSocketConnection::new(socket, connection_id, cancellation);
let mut request_count = 0_u64;
let mut replay = ConnectionReplaySnapshot::default();
@@ -172,11 +186,12 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
let decoded = match decode_response_create_with_context(&payload, &request_headers) {
Ok(decoded) => decoded.with_client_context(client_ip, user_agent.clone()),
Err(error) => {
tracing::info!(
websocket_connection_id = connection.id(),
request_id = %correlation_id,
error = %error,
"Responses WebSocket request rejected"
trace_rejected_request(
&correlation_id,
connection.id(),
&payload,
"decode",
&error.to_string(),
);
if send_protocol_error(
&mut connection,
@@ -195,6 +210,13 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
let decoded = replay.prepare(decoded);
// deadline 与 Text 可能同时就绪;在任何上游执行开始前再次封住该竞争窗口。
if connection.is_expired() {
trace_rejected_request(
&correlation_id,
connection.id(),
&payload,
"connection_expired",
"",
);
expire_connection(&mut connection).await;
break;
}
@@ -209,6 +231,13 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
{
Ok(started) => started,
Err(error) => {
trace_rejected_request(
&correlation_id,
connection.id(),
&payload,
"start",
&error.to_string(),
);
if send_gateway_error(&mut connection, &error, &correlation_id).await
== ForwardOutcome::Disconnect
{
@@ -217,6 +246,14 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
continue;
}
};
started.session.trace().record("client.connection", serde_json::json!({
"transport": "websocket", "connectionId": connection.id(), "correlationId": correlation_id,
"requestIndex": request_count,
}));
started
.session
.trace()
.capture("client.request.body", payload.as_bytes());
if forward_execution(&mut connection, started, &mut replay).await
== ForwardOutcome::Disconnect
@@ -232,6 +269,28 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
connection.log_summary(request_count);
}
// 尚未建立模型请求的拒绝也保留原文入口,用返回给客户端的 correlation ID 检索。
fn trace_rejected_request(
correlation_id: &str,
connection_id: &str,
payload: &str,
phase: &'static str,
error: &str,
) {
let trace = TraceContext::new(correlation_id);
trace.record(
"client.connection",
serde_json::json!({"connectionId": connection_id}),
);
trace.capture("client.request.body", payload.as_bytes());
trace.record(
"client.request.rejected",
serde_json::json!({
"phase": phase, "detail": diagnostic_json(&serde_json::json!({"message": error})),
}),
);
}
async fn expire_connection(connection: &mut ResponsesWebSocketConnection) {
if connection
.send_text(
@@ -10,18 +10,16 @@ use axum::{
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{Operation, RawJsonPayload, StandaloneSearchRequest};
use serde_json::{Map, Value};
use crate::ApiState;
use crate::openai::{
auth::{authenticate_client, client_access_error_response},
error::gateway_error_response,
provider_endpoint::collect_raw_json_response,
responses::request_client_context,
responses::{OpenAiRequestHeaders, request_client_context},
};
const OPENAI_PROTOCOL: &str = "openai";
const TURN_METADATA_CONTEXT_KEY: &str = "turn_metadata";
/// `POST /v1/alpha/search`。
pub(crate) async fn standalone_search(
@@ -54,16 +52,7 @@ pub(crate) async fn standalone_search(
}
fn search_operation(body: Bytes, headers: &HeaderMap) -> Result<Operation, GatewayError> {
let mut context = Map::new();
if let Some(turn_metadata) = headers
.get("x-codex-turn-metadata")
.and_then(|value| value.to_str().ok())
{
context.insert(
TURN_METADATA_CONTEXT_KEY.to_owned(),
Value::String(turn_metadata.to_owned()),
);
}
let context = OpenAiRequestHeaders::from_headers(headers).session_context();
let payload = RawJsonPayload::new(OPENAI_PROTOCOL, body)
.map_err(|_| {
GatewayError::new(
@@ -348,6 +348,8 @@ mod response {
output_tokens_display: "-".to_owned(),
cached_tokens: None,
cached_tokens_display: "-".to_owned(),
reasoning_tokens: None,
reasoning_tokens_display: "-".to_owned(),
image_input_tokens: None,
image_input_tokens_display: "-".to_owned(),
image_output_tokens: None,
@@ -374,6 +376,8 @@ mod response {
let value = serde_json::to_value(view).expect("serialize account usage");
assert!(value["inputTokens"].is_null());
assert!(value["totalTokens"].is_null());
assert!(value["reasoningTokens"].is_null());
assert_eq!(value["reasoningTokensDisplay"], "-");
assert_eq!(value["createdTokensDisplay"], "-");
}
@@ -875,13 +875,14 @@ impl ObservabilityStore for UnusedStore {
Err(unavailable("usage trend"))
}
async fn usage_calculated_billing_facts(
fn usage_calculated_billing_facts(
&self,
_: TimeRange,
_: UsageFilter,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::UsageCalculatedBillingFact>>
{
Err(unavailable("usage billing facts"))
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
Box::pin(futures::stream::once(async {
Err(unavailable("usage billing facts"))
}))
}
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
@@ -71,6 +71,9 @@ fn dashboard_account_usage_should_keep_daily_request_timeline() {
request_count: 3,
}],
quota_used_percent: None,
usage_window: None,
metric_label: "次数".to_owned(),
metric_value: "3".to_owned(),
last_used: "刚刚".to_owned(),
})
.expect("dashboard account usage");
@@ -426,6 +429,8 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
.lock()
.expect("usage detail")
.replace(UsageDetail {
trace: None,
related_requests: Vec::new(),
request: usage_record_with_account(
"req_detail",
"acct_snap_a",
@@ -563,7 +568,6 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
raw_upstream_error: Some(
r#"{"error":{"code":"upstream","message":"raw upstream marker"}}"#.to_owned(),
),
occurrence_count: 1,
client_ip: Some("203.0.113.8".to_owned()),
user_agent: Some("codex-cli/0.144.0".to_owned()),
reasoning_effort: Some("medium".to_owned()),
@@ -13,6 +13,32 @@ fn bearer_client_api_key_should_reject_missing_authorization() {
);
}
#[test]
fn actor_authorization_marker_should_not_replace_client_api_key() {
let mut headers = HeaderMap::new();
headers.insert(
"x-openai-actor-authorization",
HeaderValue::from_static("proxy-managed"),
);
assert_eq!(
bearer_client_api_key(&headers),
Err(ClientApiKeyAuthError::MissingAuthorization)
);
}
#[test]
fn actor_authorization_marker_should_preserve_existing_bearer_authentication() {
let mut headers = HeaderMap::new();
headers.insert(
"x-openai-actor-authorization",
HeaderValue::from_static("proxy-managed"),
);
headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer sk_client"));
assert_eq!(bearer_client_api_key(&headers), Ok("sk_client"));
}
#[test]
fn bearer_client_api_key_should_reject_non_utf8_authorization() {
let mut headers = HeaderMap::new();
@@ -260,6 +260,8 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio
.header(AUTHORIZATION, "Bearer sk_images_test")
.header("content-type", "application/json")
.header("x-codex-image-turn-id", "turn_image_route")
.header("session-id", "root-image-session")
.header("thread-id", "child-image-thread")
.body(Body::from(body.to_vec()))
.expect("image request"),
)
@@ -286,7 +288,7 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio
assert_eq!(captured.body.as_ref(), *body);
assert_eq!(
captured.context,
json!({"image_turn_id": "turn_image_route"})
json!({"image_turn_id": "turn_image_route", "session_id": "root-image-session", "thread_id": "child-image-thread"})
);
}
assert_eq!(
@@ -4,6 +4,7 @@ use std::sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
};
use std::time::Duration;
use axum::{
body::{Body, to_bytes},
@@ -14,6 +15,7 @@ use axum::{
},
};
use bytes::Bytes;
use futures::StreamExt;
use futures::future::{BoxFuture, pending};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
@@ -79,6 +81,7 @@ impl Trace {
enum NextStep {
Event(CoordinatedEvent),
DelayedEvent(Duration, CoordinatedEvent),
Error(EngineError),
FinalizeCancelled,
FinalizeSuccess,
@@ -280,6 +283,12 @@ impl ExecutionSession for FakeSession {
self.trace.push("next_event");
Ok(Some(event))
}
NextStep::DelayedEvent(delay, event) => {
self.trace.push("wait_event");
tokio::time::sleep(delay).await;
self.trace.push("next_event");
Ok(Some(event))
}
NextStep::Error(error) => {
self.trace.push("next_error");
self.finalized = true;
@@ -709,6 +718,154 @@ async fn streaming_encodes_first_frame_before_commit_and_http_delivery() {
std::mem::forget(response);
}
#[tokio::test(start_paused = true)]
async fn streaming_should_keep_alive_during_silent_compaction_without_restarting_execution() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::Event(delivery(started(), CommitRequirement::CommitBeforeDelivery)),
NextStep::DelayedEvent(
Duration::from_secs(240),
delivery(completed(), CommitRequirement::AlreadyCommitted),
),
NextStep::FinalizeSuccess,
],
);
let response = stream_execution_response(Box::new(session), None).await;
let mut body = response.into_body().into_data_stream();
let first = body.next().await.expect("first frame").expect("body bytes");
assert!(String::from_utf8_lossy(&first).contains("response.created"));
let start = tokio::time::Instant::now();
for heartbeat in 1..16 {
let chunk = tokio::time::timeout(Duration::from_secs(20), body.next())
.await
.expect("silent upstream must not leave the downstream idle")
.expect("heartbeat frame")
.expect("heartbeat bytes");
assert_eq!(chunk, Bytes::from_static(b": keep-alive\n\n"));
assert_eq!(start.elapsed(), Duration::from_secs(15 * heartbeat));
}
let terminal = body
.next()
.await
.expect("terminal frame")
.expect("body bytes");
assert!(String::from_utf8_lossy(&terminal).contains("response.completed"));
assert_eq!(start.elapsed(), Duration::from_secs(240));
assert_eq!(
body.next().await.expect("done frame").expect("body bytes"),
Bytes::from_static(b"data: [DONE]\n\n")
);
assert!(body.next().await.is_none());
assert_eq!(
trace.snapshot(),
vec![
"next_event",
"commit",
"wait_event",
"next_event",
"next_end"
]
);
assert!(!trace.is_cancelled());
}
#[tokio::test(start_paused = true)]
async fn dropping_stream_during_keepalive_should_cancel_and_finalize_execution() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::Event(delivery(started(), CommitRequirement::CommitBeforeDelivery)),
NextStep::DelayedEvent(
Duration::from_secs(240),
delivery(completed(), CommitRequirement::AlreadyCommitted),
),
NextStep::FinalizeCancelled,
],
);
let response = stream_execution_response(Box::new(session), None).await;
let mut body = response.into_body().into_data_stream();
let _ = body.next().await.expect("first frame").expect("body bytes");
let heartbeat = tokio::time::timeout(Duration::from_secs(20), body.next())
.await
.expect("heartbeat before client disconnect")
.expect("heartbeat frame")
.expect("heartbeat bytes");
assert_eq!(heartbeat, Bytes::from_static(b": keep-alive\n\n"));
drop(body);
tokio::task::yield_now().await;
assert!(trace.is_cancelled());
assert_eq!(
trace.snapshot(),
vec!["next_event", "commit", "wait_event", "cancel_finalize"]
);
}
#[tokio::test(start_paused = true)]
async fn streaming_keepalive_should_wait_for_first_event_before_committing_http_status() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::DelayedEvent(
Duration::from_secs(45),
delivery(started(), CommitRequirement::CommitBeforeDelivery),
),
NextStep::FinalizeCancelled,
],
);
let response = tokio::spawn(stream_execution_response(Box::new(session), None));
tokio::task::yield_now().await;
tokio::time::advance(Duration::from_secs(30)).await;
assert!(!response.is_finished());
assert!(trace.client_statuses().is_empty());
let response = response.await.expect("response task");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(trace.snapshot(), vec!["wait_event", "next_event", "commit"]);
}
#[tokio::test]
async fn streaming_response_should_disable_proxy_buffering_and_cache_transformation() {
let session = FakeSession::streaming(
Arc::new(Trace::default()),
vec![
NextStep::Event(delivery(started(), CommitRequirement::CommitBeforeDelivery)),
NextStep::FinalizeCancelled,
],
)
.with_response_headers(vec![
ProviderResponseHeader::new("x-accel-buffering", Bytes::from_static(b"yes")),
ProviderResponseHeader::new("cache-control", Bytes::from_static(b"public, max-age=60")),
]);
let response = stream_execution_response(Box::new(session), None).await;
assert_eq!(response.headers()["x-accel-buffering"], "no");
assert_eq!(
response.headers()["cache-control"],
"no-cache, no-transform"
);
assert_eq!(
response
.headers()
.get_all("x-accel-buffering")
.iter()
.count(),
1
);
assert_eq!(
response.headers().get_all("cache-control").iter().count(),
1
);
}
#[tokio::test]
async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresentable_or_sensitive_ones()
{
@@ -893,6 +1050,41 @@ async fn streaming_rate_limit_before_first_frame_should_persist_the_returned_429
assert_eq!(trace.client_statuses(), vec![429]);
}
#[tokio::test]
async fn streaming_failure_before_first_frame_should_return_safe_observed_headers() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![NextStep::Error(EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::Ambiguous,
)))],
)
.with_response_headers(vec![
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-before-close"),
),
ProviderResponseHeader::new(
"authorization",
Bytes::from_static(b"must-not-cross-boundary"),
),
]);
let response = stream_execution_response(Box::new(session), None).await;
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
assert_eq!(
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok()),
Some("turn-state-before-close")
);
assert!(response.headers().get("authorization").is_none());
assert_eq!(trace.client_statuses(), vec![502]);
}
#[tokio::test]
async fn streaming_upstream_http_failure_before_first_frame_should_preserve_raw_response() {
let trace = Arc::new(Trace::default());
@@ -979,11 +1171,29 @@ async fn buffered_upstream_http_failure_should_preserve_raw_response() {
raw_body.clone(),
));
let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()])
.with_collect_error(EngineError::Provider(error));
.with_collect_error(EngineError::Provider(error))
.with_response_headers(vec![
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-before-close"),
),
ProviderResponseHeader::new(
"authorization",
Bytes::from_static(b"must-not-cross-boundary"),
),
]);
let response = collect_execution_response(Box::new(session)).await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok()),
Some("turn-state-before-close")
);
assert!(response.headers().get("authorization").is_none());
assert_eq!(
to_bytes(response.into_body(), usize::MAX)
.await
@@ -281,6 +281,10 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
headers.insert("content-length", HeaderValue::from_static("999"));
headers.insert("authorization", HeaderValue::from_static("Bearer client"));
headers.insert("x-api-key", HeaderValue::from_static("client-key"));
headers.insert(
"x-openai-actor-authorization",
HeaderValue::from_static("proxy-managed"),
);
headers.insert("cookie", HeaderValue::from_static("client=cookie"));
headers.insert(
"chatgpt-account-id",
@@ -348,6 +352,7 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
"content-length",
"authorization",
"x-api-key",
"x-openai-actor-authorization",
"cookie",
"chatgpt-account-id",
"chatgpt-project-id",
@@ -2,8 +2,8 @@ use std::{
fmt,
pin::Pin,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
Arc, Mutex,
atomic::{AtomicBool, AtomicUsize, Ordering},
},
task::{Context, Poll},
time::{Duration, SystemTime},
@@ -25,6 +25,7 @@ use gateway_core::{
error::GatewayError,
event::{GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta},
lifecycle::CancellationToken,
operation::Operation,
routing::PublicModelId,
};
use serde_json::{Value, json};
@@ -32,7 +33,7 @@ use tokio::sync::mpsc::{
UnboundedReceiver, UnboundedSender, error::TryRecvError, unbounded_channel,
};
use tokio_tungstenite::{
connect_async,
MaybeTlsStream, WebSocketStream, connect_async,
tungstenite::{
Message as ClientMessage, client::IntoClientRequest, protocol::frame::coding::CloseCode,
},
@@ -44,27 +45,33 @@ use crate::openai::{api_router, authenticated_client};
#[derive(Default)]
struct CrossingLimitTrace {
release_terminal: tokio::sync::Notify,
started_emitted: AtomicBool,
terminal_emitted: AtomicBool,
starts: AtomicUsize,
inputs: Mutex<Vec<Value>>,
cancelled: AtomicBool,
cancellation_observed: tokio::sync::Notify,
finalized: AtomicBool,
}
struct CrossingLimitSession {
trace: Arc<CrossingLimitTrace>,
started_emitted: bool,
terminal_emitted: bool,
finalized: bool,
}
impl ExecutionSession for CrossingLimitSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
if !self.trace.started_emitted.swap(true, Ordering::AcqRel) {
if !self.started_emitted {
self.started_emitted = true;
return Ok(Some(crossing_limit_started_batch()));
}
if !self.trace.terminal_emitted.load(Ordering::Acquire) {
if !self.terminal_emitted {
self.trace.release_terminal.notified().await;
if !self.trace.terminal_emitted.swap(true, Ordering::AcqRel) {
return Ok(Some(crossing_limit_terminal_batch()));
}
self.terminal_emitted = true;
return Ok(Some(crossing_limit_terminal_batch()));
}
self.finalized = true;
self.trace.finalized.store(true, Ordering::Release);
Ok(None)
})
@@ -87,10 +94,13 @@ impl ExecutionSession for CrossingLimitSession {
}
fn is_finalized(&self) -> bool {
self.trace.finalized.load(Ordering::Acquire)
self.finalized
}
fn cancel(&self) {}
fn cancel(&self) {
self.trace.cancelled.store(true, Ordering::Release);
self.trace.cancellation_observed.notify_one();
}
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move {
@@ -127,6 +137,14 @@ impl ExecutionService for CrossingLimitExecution {
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
self.trace.starts.fetch_add(1, Ordering::AcqRel);
if let Operation::Generate(generate) = &request.operation {
self.trace
.inputs
.lock()
.expect("input trace lock")
.push(generate.protocol_payload().body()["input"].clone());
}
Ok(StartedExecution {
request_id: gateway_core::engine::ModelRequestId::new("req_ws_crossing_limit")
.expect("request id"),
@@ -134,6 +152,9 @@ impl ExecutionService for CrossingLimitExecution {
stream: request.metadata.stream,
session: Box::new(CrossingLimitSession {
trace: Arc::clone(&self.trace),
started_emitted: false,
terminal_emitted: false,
finalized: false,
}),
})
})
@@ -192,7 +213,7 @@ fn crossing_limit_terminal_batch() -> CoordinatedEvent {
}
use gateway_api::openai::responses::websocket::connection::{
ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase,
ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase, PumpExitReason,
ResponsesWebSocketConnection, WriteContext, spawn_connection,
};
@@ -255,12 +276,14 @@ struct PumpHarness {
incoming: UnboundedSender<Result<Message, TestSocketError>>,
written: UnboundedReceiver<Message>,
dropped: Arc<AtomicBool>,
cancellation: CancellationToken,
}
fn test_connection(stall_writes: bool) -> PumpHarness {
let (incoming_tx, incoming_rx) = unbounded_channel();
let (written_tx, written_rx) = unbounded_channel();
let dropped = Arc::new(AtomicBool::new(false));
let cancellation = CancellationToken::new();
let socket = TestSocket {
incoming: incoming_rx,
written: written_tx,
@@ -270,7 +293,7 @@ fn test_connection(stall_writes: bool) -> PumpHarness {
let connection = spawn_connection(
socket,
Arc::from("ws_test"),
CancellationToken::new(),
cancellation.clone(),
ConnectionConfig::PRODUCTION,
);
PumpHarness {
@@ -278,6 +301,7 @@ fn test_connection(stall_writes: bool) -> PumpHarness {
incoming: incoming_tx,
written: written_rx,
dropped,
cancellation,
}
}
@@ -316,6 +340,110 @@ async fn pump_consumes_pong_without_forwarding_a_business_event() {
assert!(result.is_err(), "Pong unexpectedly reached the coordinator");
}
#[tokio::test(start_paused = true)]
async fn waiting_for_exit_should_preserve_queued_business_frames() {
let PumpHarness {
mut connection,
incoming,
..
} = test_connection(false);
for payload in ["first", "second"] {
incoming
.send(Ok(Message::Text(payload.into())))
.expect("queue text");
}
assert!(
tokio::time::timeout(Duration::from_secs(1), connection.wait_for_exit())
.await
.is_err()
);
for expected in ["first", "second"] {
assert!(
matches!(connection.next_event().await, Some(ConnectionEvent::Text(payload)) if payload == expected)
);
}
}
#[tokio::test]
async fn client_close_should_take_priority_over_queued_business_frames() {
let PumpHarness {
mut connection,
incoming,
..
} = test_connection(false);
incoming
.send(Ok(Message::Text("queued".into())))
.expect("queue text");
incoming
.send(Ok(Message::Close(None)))
.expect("close socket");
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), connection.wait_for_exit())
.await
.expect("observe close"),
PumpExitReason::ClientClose
);
assert!(matches!(
connection.next_event().await,
Some(ConnectionEvent::Exited(PumpExitReason::ClientClose))
));
}
#[tokio::test]
async fn full_business_queue_should_not_block_lifecycle_shutdown() {
let PumpHarness {
mut connection,
incoming,
mut written,
cancellation,
..
} = test_connection(false);
for _ in 0..32 {
incoming
.send(Ok(Message::Text("queued".into())))
.expect("queue text");
}
incoming
.send(Ok(Message::Ping(vec![1].into())))
.expect("send Ping after queue fills");
assert!(matches!(written.recv().await, Some(Message::Pong(_))));
cancellation.cancel();
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), connection.wait_for_exit())
.await
.expect("observe shutdown"),
PumpExitReason::LifecycleShutdown
);
}
#[tokio::test]
async fn inbound_overload_should_close_without_executing_queued_frames() {
let PumpHarness {
mut connection,
incoming,
..
} = test_connection(false);
for _ in 0..33 {
incoming
.send(Ok(Message::Text("queued".into())))
.expect("queue text");
}
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), connection.wait_for_exit())
.await
.expect("observe overload"),
PumpExitReason::InboundOverload
);
assert!(matches!(
connection.next_event().await,
Some(ConnectionEvent::Exited(PumpExitReason::InboundOverload))
));
}
#[tokio::test(start_paused = true)]
async fn pump_does_not_emit_an_active_ping() {
let PumpHarness {
@@ -371,6 +499,7 @@ async fn stalled_write_times_out_at_three_hundred_seconds_and_aborts_the_pump()
incoming: _incoming,
written: _written,
dropped,
..
} = test_connection(true);
let request_id = Arc::<str>::from("req_test");
@@ -413,6 +542,7 @@ async fn dropping_the_connection_aborts_and_drops_the_socket_owner() {
incoming: _incoming,
written: _written,
dropped,
..
} = test_connection(false);
drop(connection);
@@ -428,6 +558,7 @@ async fn idle_connection_reaches_the_official_limit_without_starting_an_executio
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: Vec::new(),
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -539,6 +670,27 @@ async fn active_response_crossing_the_limit_writes_terminal_before_the_limit_err
}
}
for _ in 0..32 {
socket
.send(ClientMessage::Text(
json!({"type": "response.create", "model": "model-a", "input": "queued after limit"})
.to_string()
.into(),
))
.await
.expect("fill request queue before expiry");
}
socket
.send(ClientMessage::Ping(vec![1].into()))
.await
.expect("send Ping after queued request");
let pong = tokio::time::timeout(Duration::from_secs(1), socket.next())
.await
.expect("Pong timeout")
.expect("connection remains open")
.expect("valid frame");
assert!(matches!(pong, ClientMessage::Pong(_)));
tokio::time::pause();
tokio::time::advance(Duration::from_secs(60 * 60)).await;
tokio::task::yield_now().await;
@@ -567,15 +719,21 @@ async fn active_response_crossing_the_limit_writes_terminal_before_the_limit_err
assert_eq!(event_types, ["response.completed", "error"]);
assert!(trace.finalized.load(Ordering::Acquire));
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
server.abort();
}
#[tokio::test]
async fn second_text_frame_during_an_active_response_keeps_the_single_response_contract() {
type TestClientSocket = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
async fn start_active_response() -> (
Arc<CrossingLimitTrace>,
TestClientSocket,
tokio::task::JoinHandle<()>,
) {
let trace = Arc::new(CrossingLimitTrace::default());
let execution = Arc::new(CrossingLimitExecution {
client: authenticated_client("sk_ws_crossing_limit"),
trace,
trace: Arc::clone(&trace),
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -607,7 +765,7 @@ async fn second_text_frame_during_an_active_response_keeps_the_single_response_c
.into(),
);
socket
.send(response_create.clone())
.send(response_create)
.await
.expect("send first response.create");
loop {
@@ -625,19 +783,169 @@ async fn second_text_frame_during_an_active_response_keeps_the_single_response_c
}
}
(trace, socket, server)
}
#[tokio::test]
async fn requests_received_during_an_active_response_should_execute_in_order() {
let (trace, mut socket, server) = start_active_response().await;
for input in ["second", "third"] {
socket
.send(ClientMessage::Text(
json!({"type": "response.create", "model": "model-a", "input": input})
.to_string()
.into(),
))
.await
.expect("send queued response.create");
}
socket
.send(response_create)
.send(ClientMessage::Ping(vec![1, 2, 3].into()))
.await
.expect("send overlapping response.create");
.expect("send Ping after queued requests");
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("Pong timeout while response is active")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
assert!(matches!(message, ClientMessage::Pong(_)));
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
for completed_count in 1..=3 {
trace.release_terminal.notify_one();
let expected_types = if completed_count < 3 {
vec![
"response.completed",
"response.metadata",
"response.created",
]
} else {
vec!["response.completed"]
};
let mut event_types = Vec::new();
for _ in &expected_types {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("serial response timeout")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
panic!("unexpected frame during serial execution: {message:?}");
};
let event: Value = serde_json::from_str(&text).expect("response event JSON");
event_types.push(event["type"].as_str().expect("event type").to_owned());
}
assert_eq!(event_types, expected_types);
assert_eq!(
trace.starts.load(Ordering::Acquire),
(completed_count + 1).min(3)
);
}
assert_eq!(
*trace.inputs.lock().expect("input trace lock"),
vec![json!("hello"), json!("second"), json!("third")]
);
assert!(!trace.cancelled.load(Ordering::Acquire));
socket.close(None).await.expect("close WebSocket");
server.abort();
}
#[tokio::test]
async fn queued_unknown_message_should_be_rejected_before_the_next_valid_request() {
let (trace, mut socket, server) = start_active_response().await;
for payload in [
json!({"type": "response.future"}),
json!({"type": "response.create", "model": "model-a", "input": "next"}),
] {
socket
.send(ClientMessage::Text(payload.to_string().into()))
.await
.expect("queue frame");
}
trace.release_terminal.notify_one();
let mut event_types = Vec::new();
for _ in 0..4 {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("response timeout")
.expect("connection remains open")
.expect("valid frame");
let ClientMessage::Text(text) = message else {
panic!("unexpected frame: {message:?}");
};
let event: Value = serde_json::from_str(&text).expect("event JSON");
event_types.push(event["type"].as_str().expect("event type").to_owned());
}
assert_eq!(
event_types,
[
"response.completed",
"error",
"response.metadata",
"response.created"
]
);
assert_eq!(trace.starts.load(Ordering::Acquire), 2);
assert_eq!(
*trace.inputs.lock().expect("input trace lock"),
vec![json!("hello"), json!("next")]
);
socket.close(None).await.expect("close WebSocket");
server.abort();
}
#[tokio::test]
async fn queued_binary_frame_should_be_rejected_after_the_active_response_finishes() {
let (trace, mut socket, server) = start_active_response().await;
socket
.send(ClientMessage::Binary(b"{}".to_vec().into()))
.await
.expect("queue binary frame");
trace.release_terminal.notify_one();
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("terminal timeout")
.expect("connection remains open")
.expect("valid frame");
let ClientMessage::Text(text) = message else {
panic!("active response did not finish before binary rejection: {message:?}");
};
let event: Value = serde_json::from_str(&text).expect("event JSON");
assert_eq!(event["type"], "response.completed");
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("policy close timeout")
.expect("WebSocket returns a policy close")
.expect("valid WebSocket frame");
let ClientMessage::Close(Some(frame)) = message else {
panic!("overlapping response.create must close the connection");
};
.expect("connection closes")
.expect("valid frame");
assert!(
matches!(message, ClientMessage::Close(Some(frame)) if frame.code == CloseCode::Policy)
);
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
server.abort();
}
assert_eq!(frame.code, CloseCode::Policy);
#[tokio::test]
async fn client_close_should_cancel_active_execution_without_starting_queued_requests() {
let (trace, mut socket, server) = start_active_response().await;
socket
.send(ClientMessage::Text(
json!({"type": "response.create", "model": "model-a", "input": "queued"})
.to_string()
.into(),
))
.await
.expect("queue request");
socket.close(None).await.expect("close WebSocket");
tokio::time::timeout(
Duration::from_secs(2),
trace.cancellation_observed.notified(),
)
.await
.expect("cancel active execution promptly");
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
assert!(trace.cancelled.load(Ordering::Acquire));
server.abort();
}
@@ -2,7 +2,7 @@ mod connection;
mod protocol;
use std::sync::{
Arc,
Arc, Mutex,
atomic::{AtomicBool, AtomicUsize, Ordering},
};
use std::time::{Duration, SystemTime};
@@ -280,17 +280,26 @@ struct AtomicFailureTrace {
next_calls: AtomicUsize,
committed: AtomicBool,
finalized: AtomicBool,
downstream_websocket_connection_ids: Mutex<Vec<String>>,
}
struct AtomicFailureSession {
trace: Arc<AtomicFailureTrace>,
response_headers: Vec<ProviderResponseHeader>,
fail_before_first_event: bool,
}
impl ExecutionSession for AtomicFailureSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
match self.trace.next_calls.fetch_add(1, Ordering::AcqRel) {
let next_call = self.trace.next_calls.fetch_add(1, Ordering::AcqRel);
if self.fail_before_first_event && next_call == 0 {
return Err(EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::Ambiguous,
)));
}
match next_call {
0 => Ok(Some(atomic_failure_batch())),
1 => {
self.trace.finalized.store(true, Ordering::Release);
@@ -340,6 +349,7 @@ struct AtomicFailureExecution {
client: AuthenticatedClient,
trace: Arc<AtomicFailureTrace>,
response_headers: Vec<ProviderResponseHeader>,
fail_before_first_event: bool,
}
impl ExecutionService for AtomicFailureExecution {
@@ -365,6 +375,19 @@ impl ExecutionService for AtomicFailureExecution {
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
self.trace.starts.fetch_add(1, Ordering::AcqRel);
if let Operation::Generate(generate) = &request.operation
&& let Some(connection_id) = generate
.protocol_payload()
.context()
.get("downstream_websocket_connection_id")
.and_then(Value::as_str)
{
self.trace
.downstream_websocket_connection_ids
.lock()
.expect("downstream connection trace lock")
.push(connection_id.to_owned());
}
Box::pin(async move {
Ok(StartedExecution {
request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic")
@@ -374,6 +397,7 @@ impl ExecutionService for AtomicFailureExecution {
session: Box::new(AtomicFailureSession {
trace: Arc::clone(&self.trace),
response_headers: self.response_headers.clone(),
fail_before_first_event: self.fail_before_first_event,
}),
})
})
@@ -439,6 +463,7 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: Vec::new(),
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -510,6 +535,14 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() {
assert!(trace.committed.load(Ordering::Acquire));
assert_eq!(trace.next_calls.load(Ordering::Acquire), 2);
assert!(trace.finalized.load(Ordering::Acquire));
{
let downstream_connection_ids = trace
.downstream_websocket_connection_ids
.lock()
.expect("downstream connection trace lock");
assert_eq!(downstream_connection_ids.len(), 1);
assert!(downstream_connection_ids[0].starts_with("ws_"));
}
socket.close(None).await.expect("close WebSocket");
server.abort();
@@ -540,6 +573,7 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
ProviderResponseHeader::new("content-type", Bytes::from_static(b"application/private")),
ProviderResponseHeader::new("x-request-id", Bytes::from_static(b"req_upstream")),
],
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -624,6 +658,91 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
server.abort();
}
#[tokio::test]
async fn websocket_failure_before_first_event_should_return_observed_metadata_then_error() {
let trace = Arc::new(AtomicFailureTrace::default());
let execution = Arc::new(AtomicFailureExecution {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: vec![
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-before-close"),
),
ProviderResponseHeader::new(
"authorization",
Bytes::from_static(b"must-not-cross-boundary"),
),
],
fail_before_first_event: true,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_atomic".parse().expect("authorization"),
);
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
socket
.send(ClientMessage::Text(
json!({
"type": "response.create",
"model": "model-a",
"input": "hello"
})
.to_string()
.into(),
))
.await
.expect("send response.create");
let mut messages = Vec::<Value>::new();
while messages.len() < 2 {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("WebSocket response timeout")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
continue;
};
messages.push(serde_json::from_str(&text).expect("JSON WebSocket event"));
}
assert_eq!(
messages
.iter()
.filter_map(|message| message.get("type").and_then(Value::as_str))
.collect::<Vec<_>>(),
vec!["response.metadata", "error"]
);
let headers = messages[0]
.get("headers")
.and_then(Value::as_object)
.expect("response metadata headers");
assert_eq!(
headers.get("x-codex-turn-state"),
Some(&json!("turn-state-before-close"))
);
assert!(!headers.contains_key("authorization"));
assert!(!trace.committed.load(Ordering::Acquire));
socket.close(None).await.expect("close WebSocket");
server.abort();
}
#[tokio::test]
async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit() {
let trace = Arc::new(AtomicFailureTrace::default());
@@ -631,6 +750,7 @@ async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit(
client: authenticated_client("sk_ws_atomic"),
trace,
response_headers: Vec::new(),
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -162,6 +162,8 @@ async fn search_route_should_preserve_request_and_success_response_bytes() {
.header(AUTHORIZATION, "Bearer sk_search_test")
.header("content-type", "application/json")
.header("x-codex-turn-metadata", turn_metadata)
.header("session-id", "root-session")
.header("thread-id", "child-thread")
.body(Body::from(request_body.to_vec()))
.expect("search request"),
)
@@ -184,7 +186,10 @@ async fn search_route_should_preserve_request_and_success_response_bytes() {
assert_eq!(captured[0].endpoint, "/v1/alpha/search");
assert_eq!(captured[0].transport, ClientTransport::HttpJson);
assert_eq!(captured[0].body.as_ref(), request_body);
assert_eq!(captured[0].context, json!({"turn_metadata": turn_metadata}));
assert_eq!(
captured[0].context,
json!({"turn_metadata": turn_metadata, "session_id": "root-session", "thread_id": "child-thread"})
);
assert_eq!(
*execution
.committed_statuses
+2
View File
@@ -13,11 +13,13 @@ path = "tests/main.rs"
[dependencies]
async-trait = "0.1.89"
base64 = "0.23.0"
bytes = "1.11.1"
futures = "0.3.32"
futures-timer = "3.0.3"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
sha2 = "0.11.0"
semver = "1.0.27"
thiserror = "2.0.18"
tracing = "0.1.44"
@@ -2,6 +2,7 @@
mod error;
mod model;
pub mod scope;
mod selection;
mod store;
@@ -6,8 +6,8 @@ use std::time::{Duration, SystemTime};
use serde_json::{Map, Value};
use crate::error::{IdentifierError, validate_text};
use crate::routing::ProviderKind;
use crate::identity::ProviderKind;
use crate::validation::{IdentifierError, validate_text};
use super::CredentialError;
@@ -0,0 +1,291 @@
//! 请求认证时冻结的账号范围与目录,不依赖路由选择器。
use super::ProviderAccountId;
use crate::identity::ProviderKind;
use crate::validation::{IdentifierError, RoutingError};
use std::{
collections::{BTreeMap, BTreeSet},
fmt,
sync::Arc,
};
/// `account_groups.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountGroupId(String);
impl AccountGroupId {
/// 校验并创建账号分组 ID。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
let Some(suffix) = value.strip_prefix("grp_") else {
return Err(IdentifierError::MissingPrefix { expected: "grp_" });
};
if suffix.len() != 32
|| !suffix
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
{
return Err(IdentifierError::InvalidFormat);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for AccountGroupId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 快照中一个账号的 Provider 与分组归属。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeAccount {
provider_kind: ProviderKind,
group_ids: Arc<BTreeSet<AccountGroupId>>,
}
impl RuntimeAccount {
#[must_use]
pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet<AccountGroupId>) -> Self {
Self {
provider_kind,
group_ids: Arc::new(group_ids),
}
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
pub fn group_ids(&self) -> &BTreeSet<AccountGroupId> {
&self.group_ids
}
}
/// 全快照共享的账号、Provider 与分组反向索引。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuntimeAccountDirectory {
accounts: BTreeMap<ProviderAccountId, RuntimeAccount>,
providers_with_accounts: BTreeSet<ProviderKind>,
providers_by_group: BTreeMap<AccountGroupId, BTreeSet<ProviderKind>>,
}
impl RuntimeAccountDirectory {
#[must_use]
pub fn new(accounts: BTreeMap<ProviderAccountId, RuntimeAccount>) -> Self {
let mut providers_with_accounts = BTreeSet::new();
let mut providers_by_group = BTreeMap::<AccountGroupId, BTreeSet<ProviderKind>>::new();
for account in accounts.values() {
providers_with_accounts.insert(account.provider_kind.clone());
for group_id in account.group_ids.iter() {
providers_by_group
.entry(group_id.clone())
.or_default()
.insert(account.provider_kind.clone());
}
}
Self {
accounts,
providers_with_accounts,
providers_by_group,
}
}
#[must_use]
pub fn account(&self, account_id: &ProviderAccountId) -> Option<&RuntimeAccount> {
self.accounts.get(account_id)
}
#[must_use]
pub fn providers_with_accounts(&self) -> &BTreeSet<ProviderKind> {
&self.providers_with_accounts
}
#[must_use]
pub fn providers_for_groups<'a>(
&self,
group_ids: impl IntoIterator<Item = &'a AccountGroupId>,
) -> BTreeSet<ProviderKind> {
group_ids
.into_iter()
.filter_map(|group_id| self.providers_by_group.get(group_id))
.flatten()
.cloned()
.collect()
}
}
/// 历史请求保存的账号范围种类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountRoutingScopeKind {
All,
Groups,
}
impl AccountRoutingScopeKind {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::All => "all",
Self::Groups => "groups",
}
}
}
/// 请求开始时冻结的分组 ID 与名称。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingGroupSnapshot {
id: AccountGroupId,
name: String,
}
impl RoutingGroupSnapshot {
#[must_use]
pub fn new(id: AccountGroupId, name: String) -> Self {
Self { id, name }
}
#[must_use]
pub const fn id(&self) -> &AccountGroupId {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
/// 请求历史所需的完整、稳定账号范围快照。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRoutingSnapshot {
kind: AccountRoutingScopeKind,
groups: Arc<[RoutingGroupSnapshot]>,
}
impl AccountRoutingSnapshot {
#[must_use]
pub fn all() -> Self {
Self {
kind: AccountRoutingScopeKind::All,
groups: Arc::from([]),
}
}
#[must_use]
pub fn groups(groups: Vec<RoutingGroupSnapshot>) -> Self {
Self {
kind: AccountRoutingScopeKind::Groups,
groups: Arc::from(groups),
}
}
#[must_use]
pub const fn kind(&self) -> AccountRoutingScopeKind {
self.kind
}
#[must_use]
pub fn groups_snapshot(&self) -> &[RoutingGroupSnapshot] {
&self.groups
}
}
/// Key 持久 binding 编译出的账号权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientRoutingScope {
AllAccounts,
Restricted {
bound_groups: Arc<[RoutingGroupSnapshot]>,
enabled_group_ids: Arc<BTreeSet<AccountGroupId>>,
provider_kinds: Arc<BTreeSet<ProviderKind>>,
},
}
impl ClientRoutingScope {
#[must_use]
pub fn all_accounts() -> Self {
Self::AllAccounts
}
pub fn restricted(
bound_groups: Vec<RoutingGroupSnapshot>,
enabled_group_ids: BTreeSet<AccountGroupId>,
provider_kinds: BTreeSet<ProviderKind>,
) -> Result<Self, RoutingError> {
if bound_groups.is_empty() {
return Err(RoutingError::InvalidAccountScope);
}
Ok(Self::Restricted {
bound_groups: Arc::from(bound_groups),
enabled_group_ids: Arc::new(enabled_group_ids),
provider_kinds: Arc::new(provider_kinds),
})
}
}
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
}
impl FrozenAccountScope {
#[must_use]
pub const fn new(
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
) -> Self {
Self {
directory,
client_scope,
}
}
#[must_use]
pub fn allows(&self, account_id: &ProviderAccountId) -> bool {
let Some(account) = self.directory.account(account_id) else {
return false;
};
match &self.client_scope {
ClientRoutingScope::AllAccounts => true,
ClientRoutingScope::Restricted {
enabled_group_ids, ..
} => account
.group_ids()
.iter()
.any(|group_id| enabled_group_ids.contains(group_id)),
}
}
#[must_use]
pub fn provider_kinds(&self) -> &BTreeSet<ProviderKind> {
match &self.client_scope {
ClientRoutingScope::AllAccounts => self.directory.providers_with_accounts(),
ClientRoutingScope::Restricted { provider_kinds, .. } => provider_kinds,
}
}
#[must_use]
pub fn routing_snapshot(&self) -> AccountRoutingSnapshot {
match &self.client_scope {
ClientRoutingScope::AllAccounts => AccountRoutingSnapshot::all(),
ClientRoutingScope::Restricted { bound_groups, .. } => {
AccountRoutingSnapshot::groups(bound_groups.to_vec())
}
}
}
#[must_use]
pub const fn directory(&self) -> &Arc<RuntimeAccountDirectory> {
&self.directory
}
}
@@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, RwLock};
use std::time::{Duration, Instant, SystemTime};
use crate::routing::ProviderKind;
use crate::identity::ProviderKind;
use super::{AccountStatus, ProviderAccount, ProviderAccountId};
@@ -87,6 +87,7 @@ pub struct AccountRuntimeSignals {
pub in_flight: u32,
pub last_started_at: Option<SystemTime>,
pub quota_reset_at: Option<SystemTime>,
/// Provider 归一化的剩余额度基点:0 耗尽,10_000 全部可用。
pub quota_remaining_rank: Option<u64>,
pub rate_limited_until: Option<SystemTime>,
pub failure_rate_basis_points: Option<u16>,
@@ -323,6 +324,7 @@ pub struct AccountQuotaSignals {
}
impl AccountQuotaSignals {
/// `remaining_rank` 使用剩余额度基点,范围为 0..=10_000。
#[must_use]
pub const fn new(reset_at: Option<SystemTime>, remaining_rank: Option<u64>) -> Self {
Self {
@@ -386,9 +388,11 @@ pub struct AccountSelectionContext {
pub now: SystemTime,
pub excluded_accounts: BTreeSet<ProviderAccountId>,
pub preferred_account: Option<ProviderAccountId>,
/// 优先账号可调度时,是否优先于其它账号的持久权重。
pub preferred_account_overrides_weight: bool,
pub round_robin_cursor: u64,
pub eligibility: AccountEligibilityPolicy,
pub account_scope: Option<std::sync::Arc<crate::routing::FrozenAccountScope>>,
pub account_scope: Option<std::sync::Arc<crate::account::scope::FrozenAccountScope>>,
}
/// 选择账号时是否执行本地调度资格投影。
@@ -516,7 +520,9 @@ impl AccountSelector {
{
Some(candidate) => match self.scheduling_blocker(candidate, context) {
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
None if candidate.account.weight() < highest_weight => {
None if !context.preferred_account_overrides_weight
&& candidate.account.weight() < highest_weight =>
{
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
}
None => {
@@ -636,6 +642,10 @@ const SMART_LOAD_WEIGHT: f64 = 1.0;
const SMART_QUOTA_WEIGHT: f64 = 0.8;
const SMART_FAILURE_WEIGHT: f64 = 1.0;
const SMART_LATENCY_WEIGHT: f64 = 0.5;
// 容忍 5 个百分点的单项负载/失败率差异,避免微小信号波动独占新会话。
const SMART_SCORE_TOLERANCE: f64 = 0.05;
// 首输出 10 秒时延迟得分减半;固定尺度不随其他候选账号变化。
const SMART_LATENCY_HALF_SCORE_MS: f64 = 10_000.0;
fn capacity_utilization(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 {
f64::from(candidate.signals.in_flight)
@@ -647,94 +657,32 @@ fn capacity_utilization(candidate: &AccountCandidate, default_concurrency: NonZe
)
}
struct SmartNormalization {
max_capacity_utilization: f64,
min_quota: Option<u64>,
max_quota: Option<u64>,
min_latency_ms: Option<u64>,
max_latency_ms: Option<u64>,
}
impl SmartNormalization {
fn from_candidates(candidates: &[&AccountCandidate], default_concurrency: NonZeroU32) -> Self {
let max_capacity_utilization = candidates
.iter()
.map(|candidate| capacity_utilization(candidate, default_concurrency))
.reduce(f64::max)
.unwrap_or_default();
let min_quota = candidates
.iter()
.filter_map(|candidate| candidate.signals.quota_remaining_rank)
.min();
let max_quota = candidates
.iter()
.filter_map(|candidate| candidate.signals.quota_remaining_rank)
.max();
let min_latency_ms = candidates
.iter()
.filter_map(|candidate| candidate.signals.first_output_latency_ms)
.filter(|latency| *latency > 0)
.min();
let max_latency_ms = candidates
.iter()
.filter_map(|candidate| candidate.signals.first_output_latency_ms)
.filter(|latency| *latency > 0)
.max();
Self {
max_capacity_utilization,
min_quota,
max_quota,
min_latency_ms,
max_latency_ms,
}
}
}
fn select_smart_candidate<'a>(
candidates: &[&'a AccountCandidate],
default_concurrency: NonZeroU32,
cursor: u64,
) -> Option<&'a AccountCandidate> {
let normalization = SmartNormalization::from_candidates(candidates, default_concurrency);
let mut ranked = candidates
.iter()
.map(|candidate| {
(
*candidate,
smart_score(candidate, default_concurrency, &normalization),
)
})
.map(|candidate| (*candidate, smart_score(candidate, default_concurrency)))
.collect::<Vec<_>>();
ranked.sort_by(|(left, left_score), (right, right_score)| {
right_score
.total_cmp(left_score)
.then_with(|| left.account.id().cmp(right.account.id()))
});
let best_score = ranked.first()?.1.to_bits();
let tied = ranked
let best_score = ranked
.iter()
.take_while(|(_, score)| score.to_bits() == best_score)
.count();
let index = cursor as usize % tied;
.map(|(_, score)| *score)
.max_by(f64::total_cmp)?;
ranked.retain(|(_, score)| best_score - score <= SMART_SCORE_TOLERANCE);
// 轮换顺序保持稳定,避免分数轻微交错与 cursor 同步后仍反复命中同一账号。
ranked.sort_unstable_by(|(left, _), (right, _)| left.account.id().cmp(right.account.id()));
let index = (cursor % ranked.len() as u64) as usize;
Some(ranked[index].0)
}
fn smart_score(
candidate: &AccountCandidate,
default_concurrency: NonZeroU32,
normalization: &SmartNormalization,
) -> f64 {
let load = lower_ratio_is_better(
capacity_utilization(candidate, default_concurrency),
normalization.max_capacity_utilization,
);
let quota = candidate.signals.quota_remaining_rank.map_or(0.5, |quota| {
higher_is_better(
quota,
normalization.min_quota.unwrap_or(quota),
normalization.max_quota.unwrap_or(quota),
)
});
fn smart_score(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 {
let load = 1.0 - capacity_utilization(candidate, default_concurrency).clamp(0.0, 1.0);
let quota = candidate
.signals
.quota_remaining_rank
.map_or(0.5, |quota| quota.min(10_000) as f64 / 10_000.0);
let failure = 1.0
- f64::from(
candidate
@@ -748,11 +696,7 @@ fn smart_score(
.first_output_latency_ms
.filter(|latency| *latency > 0)
.map_or(1.0, |latency| {
lower_is_better(
latency,
normalization.min_latency_ms.unwrap_or(latency),
normalization.max_latency_ms.unwrap_or(latency),
)
SMART_LATENCY_HALF_SCORE_MS / (SMART_LATENCY_HALF_SCORE_MS + latency as f64)
});
SMART_LOAD_WEIGHT * load
@@ -760,24 +704,3 @@ fn smart_score(
+ SMART_FAILURE_WEIGHT * failure
+ SMART_LATENCY_WEIGHT * latency
}
fn lower_ratio_is_better(value: f64, maximum: f64) -> f64 {
if maximum <= 0.0 {
return 1.0;
}
1.0 - (value / maximum).clamp(0.0, 1.0)
}
fn lower_is_better(value: u64, minimum: u64, maximum: u64) -> f64 {
if maximum <= minimum {
return 1.0;
}
1.0 - (value.saturating_sub(minimum) as f64 / (maximum - minimum) as f64).clamp(0.0, 1.0)
}
fn higher_is_better(value: u64, minimum: u64, maximum: u64) -> f64 {
if maximum <= minimum {
return 1.0;
}
(value.saturating_sub(minimum) as f64 / (maximum - minimum) as f64).clamp(0.0, 1.0)
}
@@ -3,7 +3,7 @@
use async_trait::async_trait;
use crate::error::StoreError;
use crate::routing::ProviderKind;
use crate::identity::ProviderKind;
use super::{
AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
@@ -0,0 +1,157 @@
//! 诊断内容的统一筛选边界。未知字符串只保留长度和摘要。
use serde_json::{Map, Value, json};
use sha2::{Digest, Sha256};
/// 原始字节的长度与 SHA-256,不复制或保留正文。
#[must_use]
pub fn body_fingerprint(bytes: &[u8]) -> Value {
let digest = Sha256::digest(bytes);
let hex: String = digest
.iter()
.flat_map(|byte| {
let alphabet = b"0123456789abcdef";
[
char::from(alphabet[usize::from(byte >> 4)]),
char::from(alphabet[usize::from(byte & 15)]),
]
})
.collect();
json!({"bytes": bytes.len(), "sha256": hex})
}
/// 保留可诊断的协议结构;凭据永不进入普通日志或数据库。
#[must_use]
pub fn diagnostic_json(value: &Value) -> Value {
capture(value, "", 0, &mut 96)
}
fn capture(value: &Value, key: &str, depth: usize, budget: &mut usize) -> Value {
if *budget == 0 || depth > 6 {
return json!({"omitted": true});
}
*budget -= 1;
let key = key.to_ascii_lowercase();
if secret_key(&key) {
return Value::String("<redacted>".to_owned());
}
match value {
Value::Object(object) => {
let mut result = Map::new();
for (name, value) in object.iter().take(48) {
if *budget == 0 {
break;
}
result.insert(bounded(name, 96), capture(value, name, depth + 1, budget));
}
if result.len() < object.len() {
result.insert(
"_omittedFields".to_owned(),
json!(object.len() - result.len()),
);
}
Value::Object(result)
}
Value::Array(items) => json!({
"length": items.len(),
"sample": items.iter().take(4)
.map(|item| capture(item, key.as_str(), depth + 1, budget))
.collect::<Vec<_>>()
}),
Value::String(text)
if safe_header(&key) || safe_value_key(&key) && protocol_label(text) =>
{
Value::String(bounded(text, 256))
}
Value::String(text) => body_fingerprint(text.as_bytes()),
other => other.clone(),
}
}
/// 请求/响应头保持多值;安全 trace ID 原样保存,未知值只保存摘要。
#[must_use]
pub fn diagnostic_headers<'a>(headers: impl IntoIterator<Item = (&'a str, &'a str)>) -> Value {
let mut result = Map::new();
for (name, value) in headers.into_iter().take(64) {
let name = bounded(&name.to_ascii_lowercase(), 96);
let value = if secret_key(&name) {
Value::String("<redacted>".to_owned())
} else if safe_header(&name) {
Value::String(bounded(value, 256))
} else {
body_fingerprint(value.as_bytes())
};
let values = result
.entry(name)
.or_insert_with(|| Value::Array(Vec::new()));
if let Value::Array(values) = values {
values.push(value);
}
}
Value::Object(result)
}
fn secret_key(key: &str) -> bool {
key.contains("authorization")
|| key.contains("cookie")
|| key.contains("api_key")
|| key.contains("api-key")
|| key.contains("token") && !key.ends_with("tokens")
|| key.contains("password")
|| key.contains("secret")
|| key.contains("attestation")
|| matches!(key, "x-oai-is" | "x-oai-is-update")
}
fn safe_header(key: &str) -> bool {
matches!(
key,
"x-oai-request-id"
| "x-client-request-id"
| "x-request-id"
| "request-id"
| "openai-request-id"
| "x-openai-request-id"
| "cf-ray"
| "traceparent"
| "tracestate"
| "content-type"
| "content-length"
| "content-encoding"
| "retry-after"
| "openai-processing-ms"
| "x-processing-ms"
| "date"
| "server"
| "sec-websocket-extensions"
| "x-codex-allowed"
| "x-codex-limit-reached"
| "x-codex-active-limit"
| "x-models-etag"
)
}
fn safe_value_key(key: &str) -> bool {
safe_header(key)
|| matches!(
key,
"type" | "event" | "status" | "code" | "model" | "role" | "object" | "service_tier"
)
}
pub(super) fn bounded(value: &str, limit: usize) -> String {
value
.chars()
.take(limit)
.map(|c| if c.is_control() { ' ' } else { c })
.collect()
}
/// Event names/codes are protocol labels, never arbitrary upstream text.
pub(super) fn protocol_label(value: &str) -> bool {
!value.is_empty()
&& value.len() <= 128
&& value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-' | b':'))
}
@@ -0,0 +1,9 @@
//! 有界、请求局部的诊断证据;不参与路由、重试或客户端协议决策。
mod capture;
mod stream;
mod trace;
pub use capture::{body_fingerprint, diagnostic_headers, diagnostic_json};
pub use stream::{StreamCapture, StreamFormat};
pub use trace::TraceContext;
@@ -0,0 +1,165 @@
//! 仅用于观测的有界流分帧;超过限额时明确记缺口,不影响实际协议解析。
use serde_json::json;
use super::TraceContext;
const MAX_FRAME_BYTES: usize = 128 * 1024;
#[derive(Clone, Copy)]
pub enum StreamFormat {
Sse,
JsonLines,
}
/// 同一 exchange 持有一个 observer,跨 HTTP chunk 拼接完整事件。
pub struct StreamCapture {
trace: TraceContext,
format: StreamFormat,
line: Vec<u8>,
data: Vec<u8>,
event_name: Option<String>,
previous_cr: bool,
first_line: bool,
skipping_line: bool,
skipping_event: bool,
finished: bool,
}
impl StreamCapture {
#[must_use]
pub fn new(trace: TraceContext, format: StreamFormat) -> Self {
Self {
trace,
format,
line: Vec::new(),
data: Vec::new(),
event_name: None,
previous_cr: false,
first_line: true,
skipping_line: false,
skipping_event: false,
finished: false,
}
}
pub fn push(&mut self, bytes: &[u8]) {
self.trace.dump("upstream.chunk", bytes);
for byte in bytes {
if self.previous_cr && *byte == b'\n' {
self.previous_cr = false;
continue;
}
self.previous_cr = *byte == b'\r';
if matches!(*byte, b'\n' | b'\r') {
if !self.skipping_line {
self.line();
}
self.line.clear();
self.skipping_line = false;
} else if !self.skipping_line {
if self.line.len() < MAX_FRAME_BYTES {
self.line.push(*byte);
} else {
self.trace.record(
"capture.gap",
json!({"reason": "line_size_limit", "limitBytes": MAX_FRAME_BYTES}),
);
self.line.clear();
self.skipping_line = true;
self.skipping_event = matches!(self.format, StreamFormat::Sse);
self.data.clear();
}
}
}
}
pub fn finish(&mut self) {
if self.finished {
return;
}
let partial_frame = !self.line.is_empty()
|| !self.data.is_empty()
|| self.skipping_event
|| self.skipping_line;
if !self.line.is_empty() && !self.skipping_line {
self.line();
}
if !self.data.is_empty() && !self.skipping_event {
self.trace.capture_named_event(
"upstream.event",
&self.data,
self.event_name.as_deref(),
);
}
self.trace
.record("upstream.eof", json!({"partialFrame": partial_frame}));
self.finished = true;
}
fn line(&mut self) {
let line = if self.first_line {
self.first_line = false;
self.line
.strip_prefix(b"\xef\xbb\xbf")
.unwrap_or(&self.line)
} else {
&self.line
};
if matches!(self.format, StreamFormat::JsonLines) {
if !line.is_empty() {
self.trace.capture_event("upstream.event", line);
}
return;
}
if line.is_empty() {
if !self.data.is_empty() && !self.skipping_event {
self.trace.capture_named_event(
"upstream.event",
&self.data,
self.event_name.as_deref(),
);
}
self.data.clear();
self.event_name = None;
self.skipping_event = false;
} else if let Some(name) = line.strip_prefix(b"event:") {
self.event_name = Some(super::capture::bounded(
String::from_utf8_lossy(name).trim(),
128,
));
} else if let Some(data) = line.strip_prefix(b"data:") {
let data = data.strip_prefix(b" ").unwrap_or(data);
if self.skipping_event {
return;
}
if self.data.len() + data.len() + 1 > MAX_FRAME_BYTES {
self.trace.record(
"capture.gap",
json!({"reason": "event_size_limit", "limitBytes": MAX_FRAME_BYTES}),
);
self.data.clear();
self.skipping_event = true;
} else {
if !self.data.is_empty() {
self.data.push(b'\n');
}
self.data.extend_from_slice(data);
}
}
}
}
impl Drop for StreamCapture {
fn drop(&mut self) {
if !self.finished {
self.trace.record(
"upstream.reader.dropped",
json!({
"partialBytes": self.line.len() + self.data.len(),
"captureIncomplete": self.skipping_event || self.skipping_line,
}),
);
}
}
}
@@ -0,0 +1,338 @@
//! 首部与最近事件同时保留,流式 delta 合并,所有淘汰均有计数。
use std::{
collections::VecDeque,
sync::{Arc, Mutex},
time::{Instant, SystemTime, UNIX_EPOCH},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use serde::Serialize;
use serde_json::{Value, json};
use super::capture::bounded;
use super::{body_fingerprint, diagnostic_headers, diagnostic_json};
const MAX_EVENTS: usize = 128;
const HEAD_EVENTS: usize = 8;
const MAX_DATA_BYTES: usize = 4096;
const MAX_BUFFER_BYTES: usize = 64 * 1024;
/// 显式传给异步任务的关联上下文;默认值禁用捕获,避免全局请求映射。
#[derive(Clone, Default)]
pub struct TraceContext {
state: Option<Arc<Mutex<TraceState>>>,
attempt_index: u32,
exchange_id: Option<u64>,
}
impl std::fmt::Debug for TraceContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TraceContext")
.field("attempt_index", &self.attempt_index)
.field("exchange_id", &self.exchange_id)
.finish_non_exhaustive()
}
}
struct TraceState {
request_id: String,
started_at_ms: u128,
started: Instant,
sequence: u64,
exchanges: u64,
wire_frames: u64,
wire_bytes: u64,
dropped_events: u64,
bytes: usize,
events: VecDeque<TraceEvent>,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct TraceEvent {
sequence: u64,
last_sequence: u64,
elapsed_ms: u64,
last_elapsed_ms: u64,
attempt_index: u32,
exchange_id: Option<u64>,
stage: &'static str,
count: u64,
data: Value,
#[serde(skip)]
bytes: usize,
}
impl TraceContext {
/// 一个模型请求只创建一次,所有 attempt 共享同一有界时间线。
#[must_use]
pub fn new(request_id: &str) -> Self {
Self {
state: Some(Arc::new(Mutex::new(TraceState {
request_id: bounded(request_id, 128),
started_at_ms: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis(),
started: Instant::now(),
sequence: 0,
exchanges: 0,
wire_frames: 0,
wire_bytes: 0,
dropped_events: 0,
bytes: 0,
events: VecDeque::new(),
}))),
..Self::default()
}
}
#[must_use]
pub fn attempt(&self, attempt_index: u32) -> Self {
Self {
state: self.state.clone(),
attempt_index,
exchange_id: None,
}
}
/// 一个 attempt 中的每次实际 transport exchange 也有独立编号。
#[must_use]
pub fn exchange(&self, transport: &'static str) -> Self {
let mut context = self.clone();
if let Some(state) = &self.state {
let mut state = state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.exchanges += 1;
context.exchange_id = Some(state.exchanges);
}
context.record("upstream.exchange.started", json!({"transport": transport}));
context
}
/// details 只允许调用方构造的诊断事实;不接收原始请求/响应正文。
pub fn record(&self, stage: &'static str, data: Value) {
self.push(stage, data, false);
}
/// 同一个头部边界只采集一次:默认脱敏,显式开启 dump 时另存完整字节。
/// facts 必须为调用方构造的安全对象;原始头值只进入独立 dump。
pub fn headers<'a>(
&self,
stage: &'static str,
mut facts: Value,
headers: impl IntoIterator<Item = (&'a str, &'a [u8])>,
) {
if self.state.is_none() {
return;
}
let Some(facts_object) = facts.as_object_mut() else {
return;
};
let headers: Vec<_> = headers.into_iter().collect();
let summary = diagnostic_headers(
headers
.iter()
.map(|(name, value)| (*name, std::str::from_utf8(value).unwrap_or("<binary>"))),
);
facts_object.insert("headers".to_owned(), summary);
self.record(stage, facts);
if tracing::enabled!(target: "request_dump", tracing::Level::INFO) {
let values: Vec<_> = headers
.iter()
.map(|(name, value)| json!({"name": name, "valueBase64": STANDARD.encode(value)}))
.collect();
if let Ok(bytes) = serde_json::to_vec(&json!({"headers": values})) {
self.dump(stage, &bytes);
}
}
}
/// 入站 JSON 先记录再解析业务语义,包括未知事件和 metadata。
/// 原文仅在现有 request_dump 开关打开时写入其独立文件。
pub fn capture(&self, stage: &'static str, bytes: &[u8]) {
self.dump(stage, bytes);
self.capture_event(stage, bytes);
}
/// StreamCapture 已转储原始 chunk 时,只提取完整事件的安全摘要。
pub fn capture_event(&self, stage: &'static str, bytes: &[u8]) {
self.capture_named_event(stage, bytes, None);
}
pub fn capture_named_event(&self, stage: &'static str, bytes: &[u8], name: Option<&str>) {
if self.state.is_none() {
return;
}
let fingerprint = body_fingerprint(bytes);
if bytes.len() > 128 * 1024 {
self.record("capture.gap", json!({"reason": "json_inspection_size_limit", "stage": stage, "body": fingerprint}));
return;
}
let parsed = serde_json::from_slice::<Value>(bytes).ok();
let event_type = parsed
.as_ref()
.and_then(|v| v.get("type"))
.and_then(Value::as_str)
.or(name)
.filter(|kind| super::capture::protocol_label(kind));
let delta = event_type.is_some_and(|kind| kind.ends_with(".delta"));
let data = json!({
"body": fingerprint,
"eventType": event_type.map(|kind| bounded(kind, 128)),
"jsonValid": parsed.is_some(),
"metadata": parsed.as_ref().filter(|_| !delta).map(diagnostic_json),
});
self.push(stage, data, delta);
}
/// 只记录业务解析边界;metadata 已由 transport 在解析前捕获。
pub fn wire_event(&self, protocol: &str, event_type: Option<&str>) {
let event_type = event_type.filter(|kind| super::capture::protocol_label(kind));
let delta = event_type.is_some_and(|kind| kind.ends_with(".delta"));
self.push(
"provider.event",
json!({
"protocol": protocol,
"eventType": event_type.map(|kind| bounded(kind, 128)),
}),
delta,
);
}
/// 完整报文捕获与摘要使用相同 request / attempt / exchange ID。
pub fn dump(&self, stage: &'static str, bytes: &[u8]) {
if !tracing::enabled!(target: "request_dump", tracing::Level::INFO) {
return;
}
if let Some(state) = &self.state {
let mut state = state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.wire_frames += 1;
state.wire_bytes = state.wire_bytes.saturating_add(bytes.len() as u64);
let wire_sequence = state.wire_frames;
let chunk_count = bytes.len().div_ceil(32 * 1024).max(1);
// Each log entry stays bounded even for images or very large WebSocket frames.
for chunk_index in 0..chunk_count {
let start = chunk_index * 32 * 1024;
let chunk = &bytes[start..bytes.len().min(start + 32 * 1024)];
tracing::info!(target: "request_dump", request_id = %state.request_id,
attempt_index = self.attempt_index, exchange_id = self.exchange_id,
stage, wire_sequence, chunk_index, chunk_count,
body_bytes = bytes.len(), body_base64 = %STANDARD.encode(chunk),
contains_sensitive_data = true, "diagnostic wire capture");
}
}
}
/// 与请求终态原子持久化;旧记录没有此字段,不伪造历史时间线。
#[must_use]
pub fn snapshot(&self) -> Option<Value> {
let state = self
.state
.as_ref()?
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
Some(json!({
"schemaVersion": 1, "requestId": state.request_id,
"startedAtUnixMs": state.started_at_ms, "totalEvents": state.sequence,
"droppedEvents": state.dropped_events, "maxEvents": MAX_EVENTS,
"captureMode": "sanitized", "events": state.events,
"wireDumpEnabled": tracing::enabled!(target: "request_dump", tracing::Level::INFO),
"wireFrames": state.wire_frames, "wireBytes": state.wire_bytes,
"maxBufferBytes": MAX_BUFFER_BYTES, "headEvents": HEAD_EVENTS,
"snapshotBoundary": "execution_finalization",
}))
}
fn push(&self, stage: &'static str, mut data: Value, coalesce: bool) {
let Some(state) = &self.state else { return };
let mut encoded = serde_json::to_vec(&data).unwrap_or_default();
if encoded.len() > MAX_DATA_BYTES {
data = json!({"truncated": true, "summary": body_fingerprint(&encoded),
"eventType": data.get("eventType"), "body": data.get("body"),
"jsonValid": data.get("jsonValid"),
});
encoded = serde_json::to_vec(&data).unwrap_or_default();
}
let mut state = state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.sequence += 1;
let sequence = state.sequence;
let elapsed_ms = u64::try_from(state.started.elapsed().as_millis()).unwrap_or(u64::MAX);
if coalesce
&& let Some(last) = state.events.iter_mut().rev().find(|event| {
event.stage == stage
&& event.attempt_index == self.attempt_index
&& event.exchange_id == self.exchange_id
})
&& last.stage == stage
&& last.attempt_index == self.attempt_index
&& last.exchange_id == self.exchange_id
&& last.data.get("eventType") == data.get("eventType")
{
last.count += 1;
last.last_sequence = sequence;
last.last_elapsed_ms = elapsed_ms;
return;
}
// 普通日志只输出有界安全事实;进程意外退出时仍可按 ID 检索已发生阶段。
tracing::info!(target: "request_trace", request_id = %state.request_id,
attempt_index = self.attempt_index, exchange_id = self.exchange_id,
sequence, elapsed_ms, stage, data = %data, "request trace");
let bytes = encoded.len() + 256;
state.bytes += bytes;
state.events.push_back(TraceEvent {
sequence,
last_sequence: sequence,
elapsed_ms,
last_elapsed_ms: elapsed_ms,
attempt_index: self.attempt_index,
exchange_id: self.exchange_id,
stage,
count: 1,
data,
bytes,
});
while state.events.len() > MAX_EVENTS || state.bytes > MAX_BUFFER_BYTES {
// Preserve failures and retry decisions ahead of routine middle-of-stream events.
// Head/tail and important stages compete under the same hard memory limit.
let index = (HEAD_EVENTS..state.events.len().saturating_sub(8))
.find(|index| {
!matches!(
state.events[*index].stage,
"attempt.started"
| "account.selected"
| "attempt.failed"
| "retry.decided"
| "upstream.close"
| "upstream.read.failed"
| "capture.gap"
| "request.finished"
)
})
.unwrap_or(if state.events.len() > HEAD_EVENTS {
HEAD_EVENTS
} else {
0
});
if let Some(event) = state.events.remove(index) {
state.bytes = state.bytes.saturating_sub(event.bytes);
state.dropped_events += event.count;
}
}
}
}
impl Drop for TraceState {
fn drop(&mut self) {
tracing::info!(target: "request_trace", request_id = %self.request_id,
stage = "request.trace.closed", total_events = self.sequence,
dropped_events = self.dropped_events, wire_frames = self.wire_frames,
wire_bytes = self.wire_bytes, "request diagnostic capture closed");
}
}
@@ -8,9 +8,9 @@ use std::fmt;
use futures::future::BoxFuture;
use crate::account::ProviderAccountId;
use crate::identity::ProviderKind;
use crate::operation::ProviderSessionState;
use crate::policy::ClientApiKeyId;
use crate::routing::ProviderKind;
/// 客户端或 Provider 传递的 opaque response handle。
///
@@ -1,21 +1,26 @@
//! 唯一的账号重试、发送与下游 commit barrier owner。
use crate::diagnostics::TraceContext;
use serde_json::json;
use std::collections::BTreeSet;
use std::future::Future;
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use super::observation::ResponseObservation;
use crate::engine::continuation::{
ContinuationBinding, NativeContinuationPin, NativeContinuationScope, PreviousResponseId,
};
use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream};
use crate::engine::{
AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTransport, AttemptTrigger,
CommitRequirement, ContinuationAttempt, CoordinatedEvent, EngineError, ExecutionOutcome,
ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFailureObservation,
ModelRequestFinalization, ModelRequestId, ModelRequestTimings, NewModelRequest,
ProviderAccountStateOwner, ProviderAttemptOutcome, RequestAttemptContext, UpstreamSendState,
ModelRequestFinalization, ModelRequestId, NewModelRequest, ProviderAccountStateOwner,
ProviderAttemptOutcome, RequestAttemptContext, UpstreamSendState,
};
use crate::error::{
ContinuationRecoveryDisposition, GatewayError, GatewayErrorKind, ProviderError,
@@ -25,7 +30,6 @@ use crate::event::{
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
};
use crate::lifecycle::CancellationToken;
use crate::metering::{CostEstimate, CostSource, Usage};
use crate::operation::{Operation, ProviderSessionState};
use crate::routing::RoutingPlan;
use futures::future::Fuse;
@@ -143,12 +147,26 @@ where
.position(|candidate| candidate.provider() == pin.provider())
.ok_or(EngineError::ContinuationPinMismatch)
})?;
let trace = TraceContext::new(request_id.as_str());
trace.record(
"request.routed",
json!({
"endpoint": request.endpoint, "protocol": request.protocol,
"clientTransport": request.client_transport,
"model": request.requested_model.as_ref().map(|model| model.as_str()),
"configRevision": format!("{:?}", request.config_revision),
"admissionMs": request.admission_decision_ms,
"continuation": format!("{continuation_attempt:?}"),
"candidateCount": plan.candidates().len(),
}),
);
let image_generation_requested = operation.image_generation_requested();
let mut session = ResponseExecutionSession {
engine: Arc::clone(&self.engine),
request_id,
client_api_key_ref,
timing_started_at,
observation: ResponseObservation::new(timing_started_at),
trace,
deadline,
deadline_timer: Delay::new(
deadline
@@ -179,12 +197,7 @@ where
delivery_pending: false,
upstream_complete: false,
finalized: false,
usage: Usage::new(),
image_generation_requested,
cost: CostEstimate::unavailable(),
timings: ModelRequestTimings::default(),
client_response_id: None,
upstream_response_id: None,
last_retryable_failure: None,
last_retryable_failure_events: Vec::new(),
provider_attempt_outcomes: Vec::new(),
@@ -239,7 +252,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
engine: Arc<GatewayEngine<S>>,
request_id: ModelRequestId,
client_api_key_ref: crate::policy::ClientApiKeyId,
timing_started_at: Instant,
observation: ResponseObservation,
trace: TraceContext,
deadline: SystemTime,
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
deadline_timer: Fuse<Delay>,
@@ -274,12 +288,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
delivery_pending: bool,
upstream_complete: bool,
finalized: bool,
usage: Usage,
image_generation_requested: bool,
cost: CostEstimate,
timings: ModelRequestTimings,
client_response_id: Option<String>,
upstream_response_id: Option<String>,
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
last_retryable_failure: Option<ProviderError>,
/// 与 `last_retryable_failure` 同属一个 attempt 的原始失败批次。
@@ -298,6 +307,11 @@ impl<S: ?Sized> ResponseExecutionSession<S>
where
S: ExecutionStore,
{
/// 当前请求共享的诊断上下文。
pub fn trace(&self) -> TraceContext {
self.trace.clone()
}
/// 读取下一条 canonical event;首条未提交事件会携带 commit 要求。
///
/// # Errors
@@ -405,6 +419,10 @@ where
return Err(EngineError::InvalidDeliveryState);
}
let committed_at = SystemTime::now();
self.trace.record(
"downstream.committed",
json!({"status": client_status_code}),
);
if self.request_persisted {
best_effort_store_write(
"mark_downstream_committed",
@@ -498,7 +516,7 @@ where
return None;
}
let account = current.metadata.provider_account_id().clone();
let response_id = self.upstream_response_id.as_deref()?;
let response_id = self.observation.upstream_response_id.as_deref()?;
let previous_response_id = PreviousResponseId::new(response_id.to_owned());
let upstream_response_id = PreviousResponseId::new(response_id.to_owned());
Some(
@@ -570,11 +588,19 @@ where
return Err(EngineError::Deadline);
}
PollBoundary::Item(Some(Ok(mut event))) => {
if let Some(wire) = event.wire_event() {
self.trace
.attempt(self.attempts)
.wire_event(wire.protocol(), wire.event_type());
}
if let Some(state) = event.session_update() {
self.observe_session_update(state);
}
if let Some(observation) = event.take_observation() {
self.observe_response(observation);
}
for fact in event.canonical_facts() {
self.observe_response_identity(fact);
self.observation.observe_identity(fact);
}
for fact in event.canonical_facts() {
self.observe_event(fact).await;
@@ -686,7 +712,8 @@ where
}));
let context = AttemptContext::new(
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
.with_timing_started_at(self.timing_started_at),
.with_timing_started_at(self.observation.timing_started_at)
.with_trace(self.trace.clone()),
next_attempt,
self.deadline,
self.plan.account_selection_policy(),
@@ -709,6 +736,18 @@ where
.ok_or_else(|| EngineError::ProviderNotRegistered {
provider: candidate.provider().as_str().to_owned(),
})?;
let attempt_trace = self.trace.attempt(next_attempt.get());
attempt_trace.record(
"attempt.started",
json!({
"provider": candidate.provider().as_str(),
"model": candidate.upstream_model().map(|model| model.as_str()),
"transportPolicy": format!("{attempt_transport:?}"),
"continuation": format!("{:?}", self.continuation_attempt),
"pinnedAccount": pinned_account.as_ref().map(|id| id.as_str()),
"transportRecovery": is_transport_recovery,
}),
);
let provider_request = ProviderRequest::new(self.operation.clone(), candidate.clone());
let stream = match poll_provider(
provider,
@@ -733,6 +772,7 @@ where
ProviderBoundary::Result(result) => match *result {
Ok(stream) => stream,
Err(error) => {
record_trace_error(&attempt_trace, &error);
if self.prepare_unavailable_native_continuation_replay(&error) {
return Ok(Some(PullOutcome::AttemptDiscarded));
}
@@ -804,6 +844,11 @@ where
}
let metadata = stream.metadata().clone();
attempt_trace.record("account.selected", json!({
"provider": metadata.provider().as_str(), "accountId": metadata.provider_account_id().as_str(),
"transport": metadata.transport().as_str(),
"selectionMs": metadata.selection_observation().map(|o| o.account_selection_wait_ms()),
}));
let selection_observation = metadata.selection_observation();
let capacity = selection_observation.and_then(|observation| observation.capacity());
if !self.account_selection.is_diagnostic()
@@ -948,30 +993,24 @@ where
}
async fn observe_event(&mut self, event: &GatewayEvent) {
self.observe_event_facts(event);
self.observation.observe_event(event);
self.mark_send_observed().await;
}
fn observe_event_facts(&mut self, event: &GatewayEvent) {
let elapsed = elapsed_ms(self.timing_started_at);
observe_event_timing(&mut self.timings, event, elapsed);
if let GatewayEvent::Usage(observed) = event {
self.usage.merge(observed);
}
if let GatewayEvent::CalculatedCost(observed) = event
&& self.cost.source() != CostSource::ProviderReported
{
self.cost = observed.into_estimate();
}
if let GatewayEvent::ProviderCost(observed) = event {
self.cost = observed.into_estimate();
}
}
async fn observe_wire_event(&mut self) {
self.mark_send_observed().await;
}
fn observe_session_update(&mut self, state: &ProviderSessionState) {
let Some(current) = self.current.as_mut() else {
return;
};
if state.provider() != current.metadata.provider().as_str() {
return;
}
self.operation.set_provider_session_state(state.clone());
}
async fn mark_send_observed(&mut self) {
let Some(current) = self.current.as_mut() else {
return;
@@ -1003,44 +1042,10 @@ where
{
return;
}
let observed = observation.timings();
if let Some(value) = observed.transport_decision_wait_ms {
self.timings.transport_decision_wait_ms = Some(value);
}
if let Some(value) = observed.connect_ms {
self.timings.connect_ms = Some(value);
}
if let Some(value) = observed.headers_ms {
self.timings.headers_ms = Some(value);
}
if let Some(value) = observed.first_event_ms {
self.timings.first_event_ms = Some(value);
}
if let Some(value) = observed.first_reasoning_ms {
self.timings.first_reasoning_ms = Some(value);
}
if let Some(value) = observed.first_text_ms {
self.timings.first_text_ms = Some(value);
}
if let Some(value) = observed.first_token_ms {
self.timings.first_token_ms = Some(value);
}
if let Some(value) = observed.provider_processing_ms {
self.timings.provider_processing_ms = Some(value);
}
self.observation.observe_response(&observation);
current.response_observation = Some(observation);
}
fn observe_response_identity(&mut self, event: &GatewayEvent) {
let metadata = match event {
GatewayEvent::Started(metadata) | GatewayEvent::Completed(metadata) => metadata,
_ => return,
};
let response_id = metadata.response_id().to_owned();
self.client_response_id = Some(response_id.clone());
self.upstream_response_id = Some(response_id);
}
/// 失败可重试时要求丢弃本 attempt;预算耗尽时可返回最后一批原始失败事件。
async fn handle_stream_error(
&mut self,
@@ -1048,6 +1053,7 @@ where
) -> Result<StreamErrorOutcome, EngineError> {
// 原始 wire 只活在 request-local 决策状态;clone、attempt 记录与持久化终态
// 均只接触已剥离的稳定错误字段。
record_trace_error(&self.trace.attempt(self.attempts), &error);
let mut atomic_client_events = error.take_atomic_client_events();
let current = self.current.take().ok_or(EngineError::NoActiveAttempt)?;
self.record_provider_failure(current.metadata.provider().clone(), error.kind());
@@ -1128,6 +1134,14 @@ where
|| account_rotation_retry
|| transport_recovery.is_some();
self.trace.attempt(current.index.get()).record("retry.decided", json!({
"retryable": retryable, "continuationRetry": continuation_retry,
"sameAccountRetry": same_account_retry, "accountRotationRetry": account_rotation_retry,
"ordinaryRetry": ordinary_retry, "transportRecovery": transport_recovery.is_some(),
"delayMs": transport_recovery.map(|(_, delay)| duration_ms(delay)),
"downstreamCommitted": self.downstream_committed_at.is_some(),
"sendState": format!("{attempt_send_state:?}"),
}));
if retryable {
// 普通 clone 只保留稳定事实;原始 wire/HTTP response 由 request-local
// 所有权保留到下一次 attempt 成功,或最终空选路时返回客户端。
@@ -1203,8 +1217,8 @@ where
self.observe_response(observation);
}
for fact in event.canonical_facts() {
self.observe_response_identity(fact);
self.observe_event_facts(fact);
self.observation.observe_identity(fact);
self.observation.observe_event(fact);
}
}
}
@@ -1303,18 +1317,7 @@ where
}
fn reset_uncommitted_observations(&mut self) {
self.usage = Usage::new();
self.cost = CostEstimate::unavailable();
self.client_response_id = None;
self.upstream_response_id = None;
self.timings.transport_decision_wait_ms = None;
self.timings.connect_ms = None;
self.timings.headers_ms = None;
self.timings.first_event_ms = None;
self.timings.first_reasoning_ms = None;
self.timings.first_text_ms = None;
self.timings.first_token_ms = None;
self.timings.provider_processing_ms = None;
self.observation.reset_for_attempt();
self.upstream_complete = false;
}
@@ -1323,7 +1326,7 @@ where
return Ok(());
}
let completed_at = SystemTime::now();
self.timings.latency_ms = Some(elapsed_ms(self.timing_started_at));
self.observation.finish();
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
.response_observation
@@ -1349,6 +1352,10 @@ where
self.current_transport_observation();
let service_tier = self.current_service_tier();
let provider_metadata_json = self.current_provider_metadata_json();
self.trace.record(
"request.finished",
json!({"outcome": "succeeded", "attempts": self.attempts}),
);
if self.request_persisted {
best_effort_store_write(
"finalize_model_request",
@@ -1363,23 +1370,24 @@ where
downstream_committed_at: self.downstream_committed_at,
client_status_code: self.client_status_code,
upstream_status_code,
client_response_id: self.client_response_id.clone(),
client_response_id: self.observation.client_response_id.clone(),
upstream_request_id,
upstream_response_id: self.upstream_response_id.clone(),
upstream_response_id: self.observation.upstream_response_id.clone(),
upstream_transport,
http_version,
websocket_pool,
service_tier,
provider_metadata_json,
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
error: None,
provider_error_code: None,
raw_upstream_error: None,
failure_observation: ModelRequestFailureObservation::default(),
retry_after_ms: None,
usage: self.usage.clone(),
usage: self.observation.usage.clone(),
image_generation_succeeded: self.image_generation_succeeded(),
cost: self.cost.clone(),
timings: self.timings.clone(),
cost: self.observation.cost.clone(),
timings: self.observation.timings.clone(),
completed_at,
}),
)
@@ -1489,12 +1497,20 @@ where
if self.finalized {
return Ok(());
}
self.trace.record(
"request.finished",
json!({
"outcome": format!("{:?}", finalization.outcome),
"errorKind": finalization.error.kind().as_str(),
"sendState": format!("{:?}", finalization.send_state), "attempts": self.attempts,
}),
);
if !self.request_persisted {
self.finalized = true;
return Ok(());
}
let completed_at = SystemTime::now();
self.timings.latency_ms = Some(elapsed_ms(self.timing_started_at));
self.observation.finish();
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
.response_observation
@@ -1535,23 +1551,24 @@ where
upstream_status_code: finalization
.upstream_status_code
.or(observed_status_code),
client_response_id: self.client_response_id.clone(),
client_response_id: self.observation.client_response_id.clone(),
upstream_request_id,
upstream_response_id: self.upstream_response_id.clone(),
upstream_response_id: self.observation.upstream_response_id.clone(),
upstream_transport,
http_version,
websocket_pool,
service_tier,
provider_metadata_json,
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
error: Some(finalization.error),
provider_error_code: finalization.provider_error_code,
raw_upstream_error: finalization.raw_upstream_error,
failure_observation: finalization.observation,
retry_after_ms: finalization.retry_after_ms,
usage: self.usage.clone(),
usage: self.observation.usage.clone(),
image_generation_succeeded: self.image_generation_succeeded(),
cost: self.cost.clone(),
timings: self.timings.clone(),
cost: self.observation.cost.clone(),
timings: self.observation.timings.clone(),
completed_at,
}),
)
@@ -1599,7 +1616,11 @@ where
if matches!(self.operation, Operation::GenerateImage(_)) {
self.upstream_complete
} else {
self.usage.image_output_tokens.unwrap_or_default() > 0
self.observation
.usage
.image_output_tokens
.unwrap_or_default()
> 0
}
})
}
@@ -1637,7 +1658,7 @@ where
fn record_provider_failure(
&mut self,
provider_kind: crate::routing::ProviderKind,
provider_kind: crate::identity::ProviderKind,
error_kind: ProviderErrorKind,
) {
self.provider_attempt_outcomes
@@ -1826,31 +1847,18 @@ const fn escalate_send_state(a: UpstreamSendState, b: UpstreamSendState) -> Upst
}
}
fn observe_event_timing(timings: &mut ModelRequestTimings, event: &GatewayEvent, elapsed_ms: u64) {
timings.first_event_ms.get_or_insert(elapsed_ms);
match event {
GatewayEvent::ReasoningDelta(_) => {
timings.first_reasoning_ms.get_or_insert(elapsed_ms);
timings.first_token_ms.get_or_insert(elapsed_ms);
}
GatewayEvent::TextDelta(_) => {
timings.first_text_ms.get_or_insert(elapsed_ms);
timings.first_token_ms.get_or_insert(elapsed_ms);
}
// `response.output_item.added` 会先投影一个空参数的 tool delta;它只是结构帧,
// 不能抢在真实工具参数之前成为首个可消费 token。
GatewayEvent::ToolCallDelta(delta) if !delta.arguments_delta.is_empty() => {
timings.first_token_ms.get_or_insert(elapsed_ms);
}
GatewayEvent::CalculatedCost(_) | GatewayEvent::ProviderCost(_) => {}
_ => {}
}
}
fn elapsed_ms(started_at: Instant) -> u64 {
duration_ms(started_at.elapsed())
}
fn duration_ms(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
fn record_trace_error(trace: &TraceContext, error: &ProviderError) {
trace.record("attempt.failed", json!({
"kind": error.kind().as_str(), "sendState": format!("{:?}", error.send_state()),
"upstreamStatus": error.upstream_status(),
"upstreamCode": error.upstream_code().map(|code| code.as_str()),
"rawError": error.raw_upstream_error().map(|raw| {
let value = serde_json::from_str(raw.as_str()).unwrap_or_else(|_| json!(raw.as_str()));
crate::diagnostics::diagnostic_json(&value)
}),
}));
}
@@ -31,12 +31,12 @@ use crate::engine::{
};
use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind, StoreError};
use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader};
use crate::identity::ProviderKind;
use crate::lifecycle::CancellationToken;
use crate::operation::{Operation, ProviderSessionState};
use crate::policy::{ClientApiKeyId, ClientPolicy};
use crate::routing::{
ProviderKind, PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot,
UpstreamModelId,
PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot, UpstreamModelId,
};
use crate::runtime::RuntimeSnapshotHandle;
@@ -155,6 +155,9 @@ pub struct StartedExecution {
}
pub trait ExecutionSession: Send {
fn trace(&self) -> crate::diagnostics::TraceContext {
crate::diagnostics::TraceContext::default()
}
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>>;
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>>;
fn response_headers(&self) -> &[ProviderResponseHeader];
@@ -1035,6 +1038,12 @@ impl Drop for DefaultExecutionSession {
}
impl ExecutionSession for DefaultExecutionSession {
fn trace(&self) -> crate::diagnostics::TraceContext {
self.core
.as_ref()
.map(ResponseExecutionSession::trace)
.unwrap_or_default()
}
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
let result = self.core_mut()?.next_event().await;
@@ -1202,11 +1211,11 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
.map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID"))
}
fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
fn map_routing_error(error: crate::validation::RoutingError) -> GatewayError {
match error {
crate::error::RoutingError::NoCapableProvider { .. }
| crate::error::RoutingError::NoCapableProviderEndpoint { .. }
| crate::error::RoutingError::EmptyAccountScope => GatewayError::new(
crate::validation::RoutingError::NoCapableProvider { .. }
| crate::validation::RoutingError::NoCapableProviderEndpoint { .. }
| crate::validation::RoutingError::EmptyAccountScope => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider can execute this request",
),
+28 -7
View File
@@ -4,6 +4,7 @@ pub mod admission;
pub mod continuation;
pub mod coordinator;
pub mod execution;
mod observation;
pub mod probe;
pub mod provider;
@@ -22,17 +23,18 @@ use thiserror::Error;
use crate::account::{AccountSelectionPolicy, ProviderAccountId};
use crate::engine::continuation::{ContinuationBinding, NativeContinuationPin};
use crate::error::{
GatewayError, IdentifierError, ProviderConnectionObservation, ProviderError, ProviderErrorKind,
StoreError, validate_text,
GatewayError, ProviderConnectionObservation, ProviderError, ProviderErrorKind, StoreError,
};
use crate::event::ProviderEvent;
use crate::identity::ProviderKind;
use crate::lifecycle::CancellationToken;
use crate::metering::{CostEstimate, Usage};
use crate::operation::OperationKind;
use crate::operation::ProviderSessionState;
use crate::policy::ClientApiKeyId;
use crate::routing::{ConfigRevision, ProviderKind, PublicModelId, UpstreamModelId};
use crate::routing::{ConfigRevision, PublicModelId, UpstreamModelId};
use crate::upstream::UpstreamSendState;
use crate::validation::{IdentifierError, validate_text};
/// `model_requests.id`。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -153,7 +155,7 @@ pub struct AccountAttemptContext {
state_owner: Option<ProviderAccountStateOwner>,
credential_recovery_attempted: bool,
diagnostic_required_account: bool,
account_scope: Option<Arc<crate::routing::FrozenAccountScope>>,
account_scope: Option<Arc<crate::account::scope::FrozenAccountScope>>,
}
impl AccountAttemptContext {
@@ -194,7 +196,10 @@ impl AccountAttemptContext {
/// 附着普通请求认证时冻结的账号范围。
#[must_use]
pub fn with_account_scope(mut self, scope: Arc<crate::routing::FrozenAccountScope>) -> Self {
pub fn with_account_scope(
mut self,
scope: Arc<crate::account::scope::FrozenAccountScope>,
) -> Self {
self.account_scope = Some(scope);
self
}
@@ -233,7 +238,7 @@ impl AccountAttemptContext {
}
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
pub const fn account_scope(&self) -> Option<&Arc<crate::account::scope::FrozenAccountScope>> {
self.account_scope.as_ref()
}
}
@@ -311,6 +316,7 @@ pub struct RequestAttemptContext {
request_id: ModelRequestId,
client_api_key_ref: ClientApiKeyId,
timing_started_at: Instant,
trace: crate::diagnostics::TraceContext,
}
impl RequestAttemptContext {
@@ -320,9 +326,16 @@ impl RequestAttemptContext {
request_id,
client_api_key_ref,
timing_started_at: Instant::now(),
trace: crate::diagnostics::TraceContext::default(),
}
}
#[must_use]
pub fn with_trace(mut self, trace: crate::diagnostics::TraceContext) -> Self {
self.trace = trace;
self
}
/// 覆盖本次请求的单调计时原点。
#[must_use]
pub fn with_timing_started_at(mut self, timing_started_at: Instant) -> Self {
@@ -362,6 +375,12 @@ pub struct AttemptContext {
}
impl AttemptContext {
/// 当前 attempt 的诊断关联;克隆后可传给后台 transport 任务。
#[must_use]
pub fn trace(&self) -> crate::diagnostics::TraceContext {
self.request.trace.attempt(self.attempt_index.get())
}
#[must_use]
pub const fn new(
request: RequestAttemptContext,
@@ -469,7 +488,7 @@ impl AttemptContext {
/// 普通请求认证时冻结的账号范围;管理端诊断为 `None`。
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
pub const fn account_scope(&self) -> Option<&Arc<crate::account::scope::FrozenAccountScope>> {
self.account.account_scope()
}
@@ -607,6 +626,8 @@ pub struct ModelRequestFinalization {
pub service_tier: Option<String>,
/// Provider 已筛选的专有观测 JSON;Core 不解释字段。
pub provider_metadata_json: Option<String>,
/// 请求全程的有界诊断快照,跨 Provider 与重试保留。
pub diagnostic_trace_json: Option<String>,
pub error: Option<GatewayError>,
pub provider_error_code: Option<String>,
/// Provider 返回的原始错误正文或 WebSocket close/error frame。
@@ -0,0 +1,127 @@
//! 单次响应的计量、时间与响应身份事实;重试丢弃时统一清理。
use super::ModelRequestTimings;
use crate::event::GatewayEvent;
use crate::event::ProviderResponseObservation;
use crate::metering::{CostEstimate, CostSource, Usage};
use std::time::Instant;
pub(super) struct ResponseObservation {
pub(super) timing_started_at: Instant,
pub(super) usage: Usage,
pub(super) cost: CostEstimate,
pub(super) timings: ModelRequestTimings,
pub(super) client_response_id: Option<String>,
pub(super) upstream_response_id: Option<String>,
}
impl ResponseObservation {
pub(super) fn new(timing_started_at: Instant) -> Self {
Self {
timing_started_at,
usage: Usage::new(),
cost: CostEstimate::unavailable(),
timings: ModelRequestTimings::default(),
client_response_id: None,
upstream_response_id: None,
}
}
fn elapsed_ms(&self) -> u64 {
u64::try_from(self.timing_started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
}
pub(super) fn finish(&mut self) {
self.timings.latency_ms = Some(self.elapsed_ms());
}
pub(super) fn observe_event(&mut self, event: &GatewayEvent) {
let elapsed = self.elapsed_ms();
observe_event_timing(&mut self.timings, event, elapsed);
if let GatewayEvent::Usage(observed) = event {
self.usage.merge(observed);
}
if let GatewayEvent::CalculatedCost(observed) = event
&& self.cost.source() != CostSource::ProviderReported
{
self.cost = observed.into_estimate();
}
if let GatewayEvent::ProviderCost(observed) = event {
self.cost = observed.into_estimate();
}
}
pub(super) fn observe_identity(&mut self, event: &GatewayEvent) {
let metadata = match event {
GatewayEvent::Started(metadata) | GatewayEvent::Completed(metadata) => metadata,
_ => return,
};
let response_id = metadata.response_id().to_owned();
self.client_response_id = Some(response_id.clone());
self.upstream_response_id = Some(response_id);
}
pub(super) fn reset_for_attempt(&mut self) {
self.usage = Usage::new();
self.cost = CostEstimate::unavailable();
self.client_response_id = None;
self.upstream_response_id = None;
self.timings.transport_decision_wait_ms = None;
self.timings.connect_ms = None;
self.timings.headers_ms = None;
self.timings.first_event_ms = None;
self.timings.first_reasoning_ms = None;
self.timings.first_text_ms = None;
self.timings.first_token_ms = None;
self.timings.provider_processing_ms = None;
}
pub(super) fn observe_response(&mut self, observation: &ProviderResponseObservation) {
let observed = observation.timings();
if let Some(value) = observed.transport_decision_wait_ms {
self.timings.transport_decision_wait_ms = Some(value);
}
if let Some(value) = observed.connect_ms {
self.timings.connect_ms = Some(value);
}
if let Some(value) = observed.headers_ms {
self.timings.headers_ms = Some(value);
}
if let Some(value) = observed.first_event_ms {
self.timings.first_event_ms = Some(value);
}
if let Some(value) = observed.first_reasoning_ms {
self.timings.first_reasoning_ms = Some(value);
}
if let Some(value) = observed.first_text_ms {
self.timings.first_text_ms = Some(value);
}
if let Some(value) = observed.first_token_ms {
self.timings.first_token_ms = Some(value);
}
if let Some(value) = observed.provider_processing_ms {
self.timings.provider_processing_ms = Some(value);
}
}
}
fn observe_event_timing(timings: &mut ModelRequestTimings, event: &GatewayEvent, elapsed_ms: u64) {
timings.first_event_ms.get_or_insert(elapsed_ms);
match event {
GatewayEvent::ReasoningDelta(_) => {
timings.first_reasoning_ms.get_or_insert(elapsed_ms);
timings.first_token_ms.get_or_insert(elapsed_ms);
}
GatewayEvent::TextDelta(_) => {
timings.first_text_ms.get_or_insert(elapsed_ms);
timings.first_token_ms.get_or_insert(elapsed_ms);
}
// `response.output_item.added` 会先投影一个空参数的 tool delta;它只是结构帧,
// 不能抢在真实工具参数之前成为首个可消费 token。
GatewayEvent::ToolCallDelta(delta) if !delta.arguments_delta.is_empty() => {
timings.first_token_ms.get_or_insert(elapsed_ms);
}
GatewayEvent::CalculatedCost(_) | GatewayEvent::ProviderCost(_) => {}
_ => {}
}
}
@@ -6,7 +6,8 @@ use bytes::Bytes;
use futures::future::BoxFuture;
use crate::error::{ClientVisibleUpstreamResponse, GatewayError, GatewayErrorKind};
use crate::routing::{ProviderKind, UpstreamModelId};
use crate::identity::ProviderKind;
use crate::routing::UpstreamModelId;
use crate::{account::ProviderAccountId, operation::Operation, upstream::UpstreamSendState};
#[derive(Debug, Clone, PartialEq)]
@@ -15,13 +15,13 @@ use crate::account::{
AccountAttemptFeedback, AccountCapacitySnapshot, AccountFeedbackStats, ProviderAccountId,
};
use crate::engine::AttemptContext;
use crate::error::{OpaqueUpstreamValue, PreDeliveryRetry, ProviderError, ProviderErrorKind};
use crate::error::{PreDeliveryRetry, ProviderError, ProviderErrorKind};
use crate::event::{EventSequenceValidator, ProviderEvent};
use crate::identity::ProviderKind;
use crate::operation::Operation;
use crate::policy::ClientApiKeyId;
use crate::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId,
};
use crate::routing::{ModelCapabilities, ModelPresentation, ProviderCandidate, UpstreamModelId};
use crate::upstream::OpaqueUpstreamValue;
use crate::upstream::{UpstreamSendState, UpstreamTransport};
/// Box 只出现在 Provider Registry 的统一 event envelope 边界。
+13 -173
View File
@@ -8,126 +8,11 @@ use bytes::Bytes;
use thiserror::Error;
use crate::event::{ProviderEvent, ProviderResponseHeader};
pub use crate::upstream::OpaqueUpstreamValue;
use crate::upstream::UpstreamSendState;
/// 应用层标识不满足约束。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum IdentifierError {
/// 标识为空。
#[error("identifier must not be empty")]
Empty,
/// 标识超过核心允许的字节数。
#[error("identifier exceeds {max_bytes} bytes")]
TooLong {
/// 最大字节数。
max_bytes: usize,
},
/// 标识使用了保留的系统前缀。
#[error("identifier uses the reserved system prefix")]
ReservedPrefix,
/// 标识包含控制字符。
#[error("identifier contains control characters")]
ControlCharacter,
/// 标识缺少规定的语义前缀。
#[error("identifier must start with `{expected}`")]
MissingPrefix {
/// 规定前缀。
expected: &'static str,
},
/// 标识不满足该领域值对象的固定格式。
#[error("identifier has an invalid format")]
InvalidFormat,
}
/// Operation 构造或校验失败。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OperationError {
/// 必填文本为空。
#[error("`{field}` must not be empty")]
EmptyField {
/// 字段名。
field: &'static str,
},
/// 数量字段为零。
#[error("`{field}` must be greater than zero")]
ZeroValue {
/// 字段名。
field: &'static str,
},
/// JSON 字段必须是 object。
#[error("`{field}` must be a JSON object")]
JsonObjectRequired {
/// 字段名。
field: &'static str,
},
}
/// 路由快照或 Route Plan 不满足不变量。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum RoutingError {
/// 动态 Provider/model 标识无法构造。
#[error("routing identifier is invalid")]
InvalidIdentifier,
/// 配置 revision 必须为正数。
#[error("config revision must be greater than zero")]
InvalidRevision,
/// Restricted scope 必须保留至少一个持久化分组 binding。
#[error("account routing scope is invalid")]
InvalidAccountScope,
/// Key 的账号范围中当前没有任何可参与路由的账号。
#[error("account routing scope is empty")]
EmptyAccountScope,
/// 快照中存在重复实体。
#[error("duplicate {entity} `{id}`")]
DuplicateEntity {
/// 实体类型。
entity: &'static str,
/// 实体 ID。
id: String,
},
/// 实体引用不存在。
#[error("{entity} `{id}` was not found")]
NotFound {
/// 实体类型。
entity: &'static str,
/// 实体 ID。
id: String,
},
/// 固定平台内没有可执行本次请求的 Provider。
#[error("no provider can execute model `{model}`")]
NoCapableProvider {
/// 客户端提交的模型名称。
model: String,
},
/// 固定 Provider 的原生端点当前不可执行。
#[error("provider endpoint `{provider}` is unavailable")]
NoCapableProviderEndpoint {
/// adapter 已绑定的 Provider。
provider: String,
},
}
/// 调用方策略不满足约束或拒绝请求。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum PolicyError {
/// 请求超过调用方策略。
#[error("request was denied by caller policy: {reason}")]
Denied {
/// 稳定拒绝原因。
reason: &'static str,
},
}
/// 用量或价格估算不满足事实约束。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum MeteringError {
/// 十进制定点值格式无效或超过 `numeric(20, 10)`。
#[error("decimal value must fit unsigned numeric(20, 10)")]
InvalidDecimal,
/// 货币代码无效。
#[error("currency must be a three-letter uppercase ASCII code")]
InvalidCurrency,
}
pub use crate::validation::{
IdentifierError, MeteringError, OperationError, PolicyError, RoutingError,
};
/// 跨 Provider 稳定的上游失败分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -231,33 +116,6 @@ pub enum PreDeliveryRetry {
SameAccountTransportFallback,
}
/// Adapter 已明确分类为非 bearer 的不透明上游值。
///
/// Core 不解释或校验内容,只通过自定义 [`Debug`] 避免它意外进入日志。协议 adapter
/// 可在原客户端响应中读取原值。
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct OpaqueUpstreamValue(String);
impl OpaqueUpstreamValue {
/// 原样保存上游值。
#[must_use]
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
/// 返回已经由 adapter 安全分类的原值。
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Debug for OpaqueUpstreamValue {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("OpaqueUpstreamValue(<redacted-from-Debug>)")
}
}
/// 已由 Provider 从结构化上游错误中提取、仅供原客户端协议展示的错误详情。
///
/// 它不是诊断日志或持久化错误信息:[`Debug`] 会隐藏全部值,调用方应只在
@@ -623,13 +481,21 @@ impl ProviderError {
.get_or_insert_with(|| Box::new(ProviderErrorFailureObservation::default()))
}
/// 标记 Provider 已证明本次拒绝没有执行生成,可在下游提交前重放。
/// 标记 Provider 已证明本次失败可在下游提交前安全重放。
///
/// 证明可以来自明确的“未执行”拒绝,也可以来自同 attempt 已交付、可恢复
/// 同一业务请求的状态检查点;后者仍须由 Core 的同账号恢复门禁单独校验。
#[must_use]
pub const fn with_replay_safe(mut self) -> Self {
self.replay_safe = true;
self
}
/// 原地标记 Provider 已证明本次失败可从已交付的状态检查点安全恢复。
pub const fn set_replay_safe(&mut self) {
self.replay_safe = true;
}
/// 允许 Core 仅在客户端尚未收到任何事件时执行一次受预算约束的换号恢复。
///
/// 该标记不证明上游未执行,也不等价于 [`Self::with_replay_safe`];它只表达
@@ -1234,29 +1100,3 @@ impl StoreError {
self.kind
}
}
pub(crate) fn validate_text(
value: &str,
max_bytes: usize,
reject_reserved_prefix: bool,
required_prefix: Option<&'static str>,
) -> Result<(), IdentifierError> {
if value.is_empty() {
return Err(IdentifierError::Empty);
}
if value.len() > max_bytes {
return Err(IdentifierError::TooLong { max_bytes });
}
if value.chars().any(char::is_control) {
return Err(IdentifierError::ControlCharacter);
}
if reject_reserved_prefix && value.starts_with("__") {
return Err(IdentifierError::ReservedPrefix);
}
if let Some(prefix) = required_prefix
&& !value.starts_with(prefix)
{
return Err(IdentifierError::MissingPrefix { expected: prefix });
}
Ok(())
}
+4 -2
View File
@@ -7,10 +7,11 @@ use bytes::Bytes;
use serde_json::Value;
use thiserror::Error;
use crate::error::{IdentifierError, OpaqueUpstreamValue, validate_text};
use crate::metering::{CalculatedCost, ProviderReportedCost, Usage};
use crate::operation::ProviderSessionState;
use crate::upstream::OpaqueUpstreamValue;
use crate::upstream::UpstreamTransport;
use crate::validation::{IdentifierError, validate_text};
/// 一次响应的稳定元数据。
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -763,7 +764,8 @@ impl ProviderEvent {
}
}
/// 把本事件标记为 Provider 连接内状态的提交边界。
/// 附加 Provider 私有状态检查点;Core 可将其用于同请求恢复,协议连接可在
/// terminal 事件上持有它供下一轮使用。
pub fn attach_session_update(&mut self, state: ProviderSessionState) {
self.session_update = Some(Box::new(state));
}
@@ -0,0 +1,28 @@
//! 编译期 Provider 身份;账号与路由共同使用,不依赖路由计划。
use crate::validation::{IdentifierError, validate_text};
use std::fmt;
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderKind(String);
impl ProviderKind {
/// 校验 Provider slug。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
validate_text(&value, 64, true, None)?;
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for ProviderKind {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
+3
View File
@@ -4,10 +4,12 @@
//! 具体客户端协议和具体 Provider 都通过外层 adapter 接入。
pub mod account;
pub mod diagnostics;
pub mod engine;
pub mod error;
pub mod event;
pub mod health;
pub mod identity;
pub mod lifecycle;
pub mod metering;
pub mod operation;
@@ -17,6 +19,7 @@ pub mod routing;
pub mod runtime;
pub mod task;
pub mod upstream;
pub mod validation;
use std::sync::Arc;
use std::time::SystemTime;
@@ -3,7 +3,7 @@
use std::fmt;
use std::str::FromStr;
use crate::error::MeteringError;
use crate::validation::MeteringError;
const DECIMAL_SCALE: u128 = 10_000_000_000;
const MAX_SCALED_DECIMAL: u128 = 99_999_999_999_999_999_999;
@@ -12,7 +12,7 @@ use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::error::{OperationError, validate_text};
use crate::validation::{OperationError, validate_text};
/// 网关支持的稳定 operation 分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -11,8 +11,8 @@ pub use client_version::{
use std::fmt;
use std::sync::Arc;
use crate::error::{IdentifierError, PolicyError, validate_text};
use crate::routing::FrozenAccountScope;
use crate::account::scope::FrozenAccountScope;
use crate::validation::{IdentifierError, PolicyError, validate_text};
/// `client_api_keys.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -12,9 +12,10 @@ use crate::account::{
AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, CredentialState,
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
};
use crate::error::{IdentifierError, validate_text};
use crate::identity::ProviderKind;
use crate::policy::ClientApiKeyId;
use crate::routing::{ProviderKind, UpstreamModelId};
use crate::routing::UpstreamModelId;
use crate::validation::{IdentifierError, validate_text};
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
+7 -307
View File
@@ -2,6 +2,11 @@
pub mod snapshot;
pub use crate::account::scope::{
AccountGroupId, AccountRoutingScopeKind, AccountRoutingSnapshot, ClientRoutingScope,
FrozenAccountScope, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory,
};
pub use crate::identity::ProviderKind;
pub use snapshot::RuntimeSnapshot;
use std::collections::{BTreeMap, BTreeSet};
@@ -9,317 +14,12 @@ use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::sync::Arc;
use crate::account::{AccountSelectionPolicy, ProviderAccountId};
use crate::error::{IdentifierError, RoutingError, validate_text};
use crate::account::AccountSelectionPolicy;
use crate::operation::{CapabilityRequirements, Feature, OperationKind};
use crate::validation::{IdentifierError, RoutingError, validate_text};
const MAX_REQUEST_ATTEMPTS: u32 = 32;
/// `account_groups.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountGroupId(String);
impl AccountGroupId {
/// 校验并创建账号分组 ID。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
let Some(suffix) = value.strip_prefix("grp_") else {
return Err(IdentifierError::MissingPrefix { expected: "grp_" });
};
if suffix.len() != 32
|| !suffix
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
{
return Err(IdentifierError::InvalidFormat);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for AccountGroupId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 快照中一个账号的 Provider 与分组归属。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeAccount {
provider_kind: ProviderKind,
group_ids: Arc<BTreeSet<AccountGroupId>>,
}
impl RuntimeAccount {
#[must_use]
pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet<AccountGroupId>) -> Self {
Self {
provider_kind,
group_ids: Arc::new(group_ids),
}
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
pub fn group_ids(&self) -> &BTreeSet<AccountGroupId> {
&self.group_ids
}
}
/// 全快照共享的账号、Provider 与分组反向索引。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuntimeAccountDirectory {
accounts: BTreeMap<ProviderAccountId, RuntimeAccount>,
providers_with_accounts: BTreeSet<ProviderKind>,
providers_by_group: BTreeMap<AccountGroupId, BTreeSet<ProviderKind>>,
}
impl RuntimeAccountDirectory {
#[must_use]
pub fn new(accounts: BTreeMap<ProviderAccountId, RuntimeAccount>) -> Self {
let mut providers_with_accounts = BTreeSet::new();
let mut providers_by_group = BTreeMap::<AccountGroupId, BTreeSet<ProviderKind>>::new();
for account in accounts.values() {
providers_with_accounts.insert(account.provider_kind.clone());
for group_id in account.group_ids.iter() {
providers_by_group
.entry(group_id.clone())
.or_default()
.insert(account.provider_kind.clone());
}
}
Self {
accounts,
providers_with_accounts,
providers_by_group,
}
}
#[must_use]
pub fn account(&self, account_id: &ProviderAccountId) -> Option<&RuntimeAccount> {
self.accounts.get(account_id)
}
#[must_use]
pub fn providers_with_accounts(&self) -> &BTreeSet<ProviderKind> {
&self.providers_with_accounts
}
#[must_use]
pub fn providers_for_groups<'a>(
&self,
group_ids: impl IntoIterator<Item = &'a AccountGroupId>,
) -> BTreeSet<ProviderKind> {
group_ids
.into_iter()
.filter_map(|group_id| self.providers_by_group.get(group_id))
.flatten()
.cloned()
.collect()
}
}
/// 历史请求保存的账号范围种类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountRoutingScopeKind {
All,
Groups,
}
impl AccountRoutingScopeKind {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::All => "all",
Self::Groups => "groups",
}
}
}
/// 请求开始时冻结的分组 ID 与名称。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingGroupSnapshot {
id: AccountGroupId,
name: String,
}
impl RoutingGroupSnapshot {
#[must_use]
pub fn new(id: AccountGroupId, name: String) -> Self {
Self { id, name }
}
#[must_use]
pub const fn id(&self) -> &AccountGroupId {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
/// 请求历史所需的完整、稳定账号范围快照。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRoutingSnapshot {
kind: AccountRoutingScopeKind,
groups: Arc<[RoutingGroupSnapshot]>,
}
impl AccountRoutingSnapshot {
#[must_use]
pub fn all() -> Self {
Self {
kind: AccountRoutingScopeKind::All,
groups: Arc::from([]),
}
}
#[must_use]
pub fn groups(groups: Vec<RoutingGroupSnapshot>) -> Self {
Self {
kind: AccountRoutingScopeKind::Groups,
groups: Arc::from(groups),
}
}
#[must_use]
pub const fn kind(&self) -> AccountRoutingScopeKind {
self.kind
}
#[must_use]
pub fn groups_snapshot(&self) -> &[RoutingGroupSnapshot] {
&self.groups
}
}
/// Key 持久 binding 编译出的账号权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientRoutingScope {
AllAccounts,
Restricted {
bound_groups: Arc<[RoutingGroupSnapshot]>,
enabled_group_ids: Arc<BTreeSet<AccountGroupId>>,
provider_kinds: Arc<BTreeSet<ProviderKind>>,
},
}
impl ClientRoutingScope {
#[must_use]
pub fn all_accounts() -> Self {
Self::AllAccounts
}
pub fn restricted(
bound_groups: Vec<RoutingGroupSnapshot>,
enabled_group_ids: BTreeSet<AccountGroupId>,
provider_kinds: BTreeSet<ProviderKind>,
) -> Result<Self, RoutingError> {
if bound_groups.is_empty() {
return Err(RoutingError::InvalidAccountScope);
}
Ok(Self::Restricted {
bound_groups: Arc::from(bound_groups),
enabled_group_ids: Arc::new(enabled_group_ids),
provider_kinds: Arc::new(provider_kinds),
})
}
}
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
}
impl FrozenAccountScope {
#[must_use]
pub const fn new(
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
) -> Self {
Self {
directory,
client_scope,
}
}
#[must_use]
pub fn allows(&self, account_id: &ProviderAccountId) -> bool {
let Some(account) = self.directory.account(account_id) else {
return false;
};
match &self.client_scope {
ClientRoutingScope::AllAccounts => true,
ClientRoutingScope::Restricted {
enabled_group_ids, ..
} => account
.group_ids()
.iter()
.any(|group_id| enabled_group_ids.contains(group_id)),
}
}
#[must_use]
pub fn provider_kinds(&self) -> &BTreeSet<ProviderKind> {
match &self.client_scope {
ClientRoutingScope::AllAccounts => self.directory.providers_with_accounts(),
ClientRoutingScope::Restricted { provider_kinds, .. } => provider_kinds,
}
}
#[must_use]
pub fn routing_snapshot(&self) -> AccountRoutingSnapshot {
match &self.client_scope {
ClientRoutingScope::AllAccounts => AccountRoutingSnapshot::all(),
ClientRoutingScope::Restricted { bound_groups, .. } => {
AccountRoutingSnapshot::groups(bound_groups.to_vec())
}
}
}
#[must_use]
pub const fn directory(&self) -> &Arc<RuntimeAccountDirectory> {
&self.directory
}
}
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderKind(String);
impl ProviderKind {
/// 校验 Provider slug。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
validate_text(&value, 64, true, None)?;
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for ProviderKind {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 客户端请求中的模型名称。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PublicModelId(String);
@@ -9,12 +9,12 @@ use futures::future::BoxFuture;
use crate::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
use crate::error::RoutingError;
use crate::operation::Operation;
use crate::policy::{
ClientApiKeyId, ClientPolicy, CodexClientMinVersions, CodexClientVersion,
PlaintextClientApiKey, RateLimits,
};
use crate::validation::RoutingError;
use super::{
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
@@ -11,11 +11,12 @@ use futures_timer::Delay;
use crate::engine::provider::ProviderCatalogGeneration;
use crate::health::{HealthProbe, HealthState};
use crate::identity::ProviderKind;
use crate::lifecycle::CancellationToken;
use crate::routing::ConfigRevision;
use crate::routing::snapshot::{
RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotStorePort,
};
use crate::routing::{ConfigRevision, ProviderKind};
use crate::task::{
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable,
+30 -1
View File
@@ -1,6 +1,8 @@
//! 上游调用跨 Engine、Event、Error 与 Provider 共享的中立边界事实。
use crate::error::{IdentifierError, validate_text};
use std::fmt;
use crate::validation::{IdentifierError, validate_text};
/// 上游是否可能已经收到业务 payload。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -43,3 +45,30 @@ impl UpstreamTransport {
&self.0
}
}
/// Adapter 已明确分类为非 bearer 的不透明上游值。
///
/// Core 不解释或校验内容,只通过自定义 [`Debug`] 避免它意外进入日志。协议 adapter
/// 可在原客户端响应中读取原值。
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct OpaqueUpstreamValue(String);
impl OpaqueUpstreamValue {
/// 原样保存上游值。
#[must_use]
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
/// 返回已经由 adapter 安全分类的原值。
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Debug for OpaqueUpstreamValue {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("OpaqueUpstreamValue(<redacted-from-Debug>)")
}
}
@@ -0,0 +1,148 @@
//! 领域值对象的纯校验错误与文本约束,不依赖执行事件或路由状态。
use thiserror::Error;
/// 应用层标识不满足约束。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum IdentifierError {
/// 标识为空。
#[error("identifier must not be empty")]
Empty,
/// 标识超过核心允许的字节数。
#[error("identifier exceeds {max_bytes} bytes")]
TooLong {
/// 最大字节数。
max_bytes: usize,
},
/// 标识使用了保留的系统前缀。
#[error("identifier uses the reserved system prefix")]
ReservedPrefix,
/// 标识包含控制字符。
#[error("identifier contains control characters")]
ControlCharacter,
/// 标识缺少规定的语义前缀。
#[error("identifier must start with `{expected}`")]
MissingPrefix {
/// 规定前缀。
expected: &'static str,
},
/// 标识不满足该领域值对象的固定格式。
#[error("identifier has an invalid format")]
InvalidFormat,
}
/// Operation 构造或校验失败。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OperationError {
/// 必填文本为空。
#[error("`{field}` must not be empty")]
EmptyField {
/// 字段名。
field: &'static str,
},
/// 数量字段为零。
#[error("`{field}` must be greater than zero")]
ZeroValue {
/// 字段名。
field: &'static str,
},
/// JSON 字段必须是 object。
#[error("`{field}` must be a JSON object")]
JsonObjectRequired {
/// 字段名。
field: &'static str,
},
}
/// 路由快照或 Route Plan 不满足不变量。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum RoutingError {
/// 动态 Provider/model 标识无法构造。
#[error("routing identifier is invalid")]
InvalidIdentifier,
/// 配置 revision 必须为正数。
#[error("config revision must be greater than zero")]
InvalidRevision,
/// Restricted scope 必须保留至少一个持久化分组 binding。
#[error("account routing scope is invalid")]
InvalidAccountScope,
/// Key 的账号范围中当前没有任何可参与路由的账号。
#[error("account routing scope is empty")]
EmptyAccountScope,
/// 快照中存在重复实体。
#[error("duplicate {entity} `{id}`")]
DuplicateEntity {
/// 实体类型。
entity: &'static str,
/// 实体 ID。
id: String,
},
/// 实体引用不存在。
#[error("{entity} `{id}` was not found")]
NotFound {
/// 实体类型。
entity: &'static str,
/// 实体 ID。
id: String,
},
/// 固定平台内没有可执行本次请求的 Provider。
#[error("no provider can execute model `{model}`")]
NoCapableProvider {
/// 客户端提交的模型名称。
model: String,
},
/// 固定 Provider 的原生端点当前不可执行。
#[error("provider endpoint `{provider}` is unavailable")]
NoCapableProviderEndpoint {
/// adapter 已绑定的 Provider。
provider: String,
},
}
/// 调用方策略不满足约束或拒绝请求。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum PolicyError {
/// 请求超过调用方策略。
#[error("request was denied by caller policy: {reason}")]
Denied {
/// 稳定拒绝原因。
reason: &'static str,
},
}
/// 用量或价格估算不满足事实约束。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum MeteringError {
/// 十进制定点值格式无效或超过 `numeric(20, 10)`。
#[error("decimal value must fit unsigned numeric(20, 10)")]
InvalidDecimal,
/// 货币代码无效。
#[error("currency must be a three-letter uppercase ASCII code")]
InvalidCurrency,
}
pub(crate) fn validate_text(
value: &str,
max_bytes: usize,
reject_reserved_prefix: bool,
required_prefix: Option<&'static str>,
) -> Result<(), IdentifierError> {
if value.is_empty() {
return Err(IdentifierError::Empty);
}
if value.len() > max_bytes {
return Err(IdentifierError::TooLong { max_bytes });
}
if value.chars().any(char::is_control) {
return Err(IdentifierError::ControlCharacter);
}
if reject_reserved_prefix && value.starts_with("__") {
return Err(IdentifierError::ReservedPrefix);
}
if let Some(prefix) = required_prefix
&& !value.starts_with(prefix)
{
return Err(IdentifierError::MissingPrefix { expected: prefix });
}
Ok(())
}
@@ -100,6 +100,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
preferred_account: None,
preferred_account_overrides_weight: false,
round_robin_cursor: 0,
eligibility: AccountEligibilityPolicy::Enforce,
account_scope: None,
@@ -417,6 +418,84 @@ fn selector_should_not_let_soft_affinity_bypass_a_higher_weight_account() {
);
}
#[test]
fn selector_should_prioritize_affinity_over_weight_when_requested() {
let candidates = [
weighted_candidate("acct_preferred", 1, 0),
weighted_candidate("acct_higher", 100, 0),
];
for strategy in [
RotationStrategy::Smart,
RotationStrategy::QuotaResetPriority,
RotationStrategy::RoundRobin,
RotationStrategy::Sticky,
] {
let mut selection = context(strategy);
selection.preferred_account = Some(candidates[0].account.id().clone());
selection.preferred_account_overrides_weight = true;
let selected = AccountSelector
.select(&candidates, &selection)
.expect("preferred account");
assert_eq!(
(
selected.candidate().account.id().as_str(),
selected.preferred()
),
("acct_preferred", PreferredAccountSelection::Hit),
"{strategy:?}"
);
}
}
#[test]
fn highest_priority_affinity_should_still_obey_existing_scheduling_constraints() {
for blocker in [
AccountSchedulingBlocker::Excluded,
AccountSchedulingBlocker::ConcurrencyLimit,
AccountSchedulingBlocker::RequestInterval,
AccountSchedulingBlocker::LocalAvailability,
] {
let mut candidates = [
weighted_candidate("acct_preferred", 1, 0),
weighted_candidate("acct_higher", 100, 0),
];
let mut selection = context(RotationStrategy::Smart);
selection.preferred_account = Some(candidates[0].account.id().clone());
selection.preferred_account_overrides_weight = true;
match blocker {
AccountSchedulingBlocker::Excluded => {
selection
.excluded_accounts
.insert(candidates[0].account.id().clone());
}
AccountSchedulingBlocker::ConcurrencyLimit => candidates[0].signals.in_flight = 3,
AccountSchedulingBlocker::RequestInterval => {
candidates[0].signals.last_started_at = Some(selection.now);
selection.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(3).expect("limit"),
Duration::from_secs(1),
);
}
AccountSchedulingBlocker::LocalAvailability => {
candidates[0].signals.rate_limited_until =
Some(selection.now + Duration::from_secs(60))
}
_ => unreachable!("test cases"),
}
let selected = AccountSelector
.select(&candidates, &selection)
.expect("fallback account");
assert_eq!(
(
selected.candidate().account.id().as_str(),
selected.preferred()
),
("acct_higher", PreferredAccountSelection::Blocked(blocker))
);
}
}
#[test]
fn selector_should_fall_back_when_a_higher_weight_account_reaches_its_override() {
let mut saturated = weighted_candidate("acct_high_weight", 100, 1);
@@ -618,8 +697,8 @@ fn account_feedback_should_decay_failure_rate_after_success() {
#[test]
fn smart_selector_should_use_provider_quota_rank_after_load_is_equal() {
let candidates = vec![
candidate("acct_low_quota", 0, Some(20)),
candidate("acct_high_quota", 0, Some(80)),
candidate("acct_low_quota", 0, Some(2_000)),
candidate("acct_high_quota", 0, Some(8_000)),
];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
@@ -664,10 +743,10 @@ fn quota_reset_selector_should_use_capacity_utilization_as_load_tiebreaker() {
}
#[test]
fn smart_selector_should_prefer_known_quota_over_unknown_after_load_is_equal() {
fn smart_selector_should_prefer_sufficient_known_quota_over_unknown_after_load_is_equal() {
let candidates = vec![
candidate("acct_unknown", 0, None),
candidate("acct_known", 0, Some(1)),
candidate("acct_known", 0, Some(8_000)),
];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
@@ -1,8 +1,13 @@
use std::time::{Duration, Instant};
use gateway_core::account::{AccountAttemptFeedback, AccountFeedbackStats, ProviderAccountId};
use gateway_core::account::{
AccountAttemptFeedback, AccountCandidate, AccountFeedbackStats, AccountSelector,
ProviderAccountId, RotationStrategy,
};
use gateway_core::routing::ProviderKind;
use super::{candidate, candidate_with_concurrency, context};
const FAILURE_RATE_HALF_LIFE: Duration = Duration::from_secs(15 * 60);
fn feedback_subject() -> (AccountFeedbackStats, ProviderKind, ProviderAccountId) {
@@ -96,3 +101,139 @@ fn concurrent_account_failures_should_not_lose_samples() {
assert_eq!(failure_rate, Some(8_322));
}
fn smart_selection_ids(candidates: &[AccountCandidate]) -> Vec<&str> {
let mut selection = context(RotationStrategy::Smart);
(0..20)
.map(|cursor| {
selection.round_robin_cursor = cursor;
AccountSelector
.select(candidates, &selection)
.expect("candidate available")
.candidate()
.account
.id()
.as_str()
})
.collect()
}
#[test]
fn smart_selector_should_rotate_despite_small_signal_differences() {
for signal in ["quota", "latency", "load", "failure"] {
let mut candidates = [
candidate_with_concurrency("acct_a", 0, 100),
candidate_with_concurrency("acct_b", 0, 100),
];
match signal {
"quota" => {
candidates[0].signals.quota_remaining_rank = Some(600);
candidates[1].signals.quota_remaining_rank = Some(601);
}
"latency" => {
candidates[0].signals.first_output_latency_ms = Some(2_500);
candidates[1].signals.first_output_latency_ms = Some(2_501);
}
"load" => candidates[0].signals.in_flight = 1,
"failure" => candidates[0].signals.failure_rate_basis_points = Some(1),
_ => unreachable!(),
}
let selected = smart_selection_ids(&candidates);
assert_eq!(selected, ["acct_a", "acct_b"].repeat(10), "{signal}");
}
}
#[test]
fn smart_selector_should_keep_rotation_order_when_nearby_scores_cross() {
let mut candidates = [
candidate("acct_a", 0, Some(8_000)),
candidate("acct_b", 0, Some(8_001)),
candidate("acct_worse", 0, Some(2_000)),
];
let mut selection = context(RotationStrategy::Smart);
let mut selected = Vec::new();
for cursor in 0..20 {
selection.round_robin_cursor = cursor;
for candidate in &mut candidates {
match candidate.account.id().as_str() {
"acct_a" => candidate.signals.quota_remaining_rank = Some(8_000 + cursor % 2),
"acct_b" => candidate.signals.quota_remaining_rank = Some(8_001 - cursor % 2),
_ => {}
}
}
candidates.reverse();
selected.push(
AccountSelector
.select(&candidates, &selection)
.expect("candidate available")
.candidate()
.account
.id()
.as_str()
.to_owned(),
);
}
assert_eq!(selected, ["acct_a", "acct_b"].repeat(10));
}
#[test]
fn smart_selector_should_only_rotate_among_candidates_close_to_the_best() {
let candidates = [
candidate("acct_a", 0, Some(10_000)),
candidate("acct_b", 0, Some(9_500)),
candidate("acct_c", 0, Some(9_000)),
];
assert_eq!(
smart_selection_ids(&candidates),
["acct_a", "acct_b"].repeat(10)
);
}
#[test]
fn smart_selector_should_preserve_material_signal_advantages() {
for signal in ["quota", "latency", "load", "failure"] {
let mut candidates = [
candidate_with_concurrency("acct_a", 0, 10),
candidate_with_concurrency("acct_b", 0, 10),
];
match signal {
"quota" => {
candidates[0].signals.quota_remaining_rank = Some(600);
candidates[1].signals.quota_remaining_rank = Some(2_600);
}
"latency" => {
candidates[0].signals.first_output_latency_ms = Some(20_000);
candidates[1].signals.first_output_latency_ms = Some(2_500);
}
"load" => candidates[0].signals.in_flight = 1,
"failure" => candidates[0].signals.failure_rate_basis_points = Some(2_000),
_ => unreachable!(),
}
assert_eq!(smart_selection_ids(&candidates), ["acct_b"; 20], "{signal}");
}
}
#[test]
fn smart_selector_should_balance_actual_load_against_remaining_quota() {
let mut candidates = [
candidate_with_concurrency("acct_a", 0, 10),
candidate_with_concurrency("acct_b", 1, 10),
];
candidates[0].signals.quota_remaining_rank = Some(600);
candidates[1].signals.quota_remaining_rank = Some(2_600);
assert_eq!(smart_selection_ids(&candidates), ["acct_b"; 20]);
}
#[test]
fn smart_selector_should_treat_unknown_quota_as_neutral() {
let candidates = [
candidate("acct_known_low", 0, Some(600)),
candidate("acct_unknown", 0, None),
];
assert_eq!(smart_selection_ids(&candidates), ["acct_unknown"; 20]);
}
@@ -0,0 +1,49 @@
use gateway_core::diagnostics::{diagnostic_headers, diagnostic_json};
use serde_json::json;
#[test]
fn credentials_and_content_never_enter_diagnostics_but_trace_ids_survive() {
let headers = diagnostic_headers([
("Authorization", "Bearer never-persist-this"),
("Set-Cookie", "session=never-persist-this"),
("X-OAI-Request-ID", "upstream-123"),
("x-codex-turn-state", "opaque-never-persist-this"),
("cf-ray", "ray-one"),
("cf-ray", "ray-two"),
]);
assert_eq!(headers["x-oai-request-id"][0], "upstream-123");
assert_eq!(headers["cf-ray"].as_array().unwrap().len(), 2);
assert!(!headers.to_string().contains("never-persist-this"));
let content = diagnostic_json(&json!({
"type": "codex.response.metadata", "headers": headers,
"access_token": "never-persist-this", "refresh_token": "never-persist-this",
"response": {"output": [{"text": "private model output"}]},
"input": "private user prompt", "unknown": "private extension",
}));
assert_eq!(content["type"], "codex.response.metadata");
let text = content.to_string();
for secret in [
"never-persist-this",
"private model output",
"private user prompt",
"private extension",
] {
assert!(!text.contains(secret), "leaked {secret}");
}
}
#[test]
fn arbitrary_event_type_text_is_fingerprinted_instead_of_logged() {
let trace = gateway_core::diagnostics::TraceContext::new("req_type");
trace.capture(
"upstream.event",
br#"{"type":"secret@example.com","message":"private-content"}"#,
);
let snapshot = trace.snapshot().unwrap();
assert_eq!(
snapshot["events"][0]["data"]["eventType"],
serde_json::Value::Null
);
assert!(!snapshot.to_string().contains("secret@example.com"));
assert!(!snapshot.to_string().contains("private-content"));
}
@@ -0,0 +1,3 @@
mod capture;
mod stream;
mod trace;
@@ -0,0 +1,70 @@
use gateway_core::diagnostics::{StreamCapture, StreamFormat, TraceContext};
#[test]
fn sse_capture_reassembles_every_chunk_boundary_and_preserves_unknown_events() {
let bytes = "event: future.metadata\r\ndata: {\"headers\":{\"x-request-id\":\"up-1\"},\"unknown\":\"中文\"}\r\n\r\ndata: {\"type\":\"response.completed\"}\n\n".as_bytes();
for boundary in 0..=bytes.len() {
let trace = TraceContext::new("req_sse");
let mut capture = StreamCapture::new(trace.clone(), StreamFormat::Sse);
capture.push(&bytes[..boundary]);
capture.push(&bytes[boundary..]);
capture.finish();
let snapshot = trace.snapshot().unwrap();
let events = snapshot["events"].as_array().unwrap();
assert_eq!(events.len(), 3, "boundary {boundary}");
assert_eq!(events[0]["data"]["eventType"], "future.metadata");
assert_eq!(
events[0]["data"]["metadata"]["headers"]["x-request-id"],
"up-1"
);
assert_eq!(events[1]["data"]["eventType"], "response.completed");
assert_eq!(events[2]["data"]["partialFrame"], false);
}
}
#[test]
fn oversize_frames_report_a_gap_and_resume_at_the_next_frame() {
let trace = TraceContext::new("req_sse");
let mut capture = StreamCapture::new(trace.clone(), StreamFormat::Sse);
capture.push(b"data: ");
capture.push(&vec![b'x'; 130 * 1024]);
capture.push(b"\n\ndata: {\"type\":\"response.completed\"}\n\n");
capture.finish();
let snapshot = trace.snapshot().unwrap();
let events = snapshot["events"].as_array().unwrap();
assert_eq!(events[0]["stage"], "capture.gap");
assert_eq!(events[1]["data"]["eventType"], "response.completed");
}
#[test]
fn dropping_a_partial_stream_is_distinct_from_upstream_eof() {
let trace = TraceContext::new("req_partial");
{
let mut capture = StreamCapture::new(trace.clone(), StreamFormat::Sse);
capture.push(b"data: {\"type\":");
}
let snapshot = trace.snapshot().unwrap();
assert_eq!(snapshot["events"][0]["stage"], "upstream.reader.dropped");
assert!(
snapshot["events"][0]["data"]["partialBytes"]
.as_u64()
.unwrap()
> 0
);
}
#[test]
fn sse_accepts_bom_and_bare_carriage_return_even_when_split_into_single_bytes() {
let trace = TraceContext::new("req_cr");
let mut capture = StreamCapture::new(trace.clone(), StreamFormat::Sse);
for byte in b"\xef\xbb\xbfdata: {\"type\":\"codex.response.metadata\"}\r\r" {
capture.push(&[*byte]);
}
capture.finish();
let snapshot = trace.snapshot().unwrap();
assert_eq!(
snapshot["events"][0]["data"]["eventType"],
"codex.response.metadata"
);
assert_eq!(snapshot["events"][1]["data"]["partialFrame"], false);
}
@@ -0,0 +1,142 @@
use gateway_core::diagnostics::TraceContext;
use serde_json::json;
#[test]
fn bounded_history_preserves_start_failure_and_final_result() {
let trace = TraceContext::new("req_trace");
trace.record("request.started", json!({}));
for index in 0..400 {
trace.record(
"test.event",
json!({"index": index, "detail": "x".repeat(2000)}),
);
}
trace
.attempt(1)
.record("attempt.failed", json!({"code": "websocket_close_1000"}));
trace
.attempt(2)
.record("request.finished", json!({"outcome": "succeeded"}));
let snapshot = trace.snapshot().unwrap();
let events = snapshot["events"].as_array().unwrap();
assert_eq!(events[0]["stage"], "request.started");
assert_eq!(events[events.len() - 2]["attemptIndex"], 1);
assert_eq!(events.last().unwrap()["attemptIndex"], 2);
assert!(snapshot["droppedEvents"].as_u64().unwrap() > 0);
assert!(events.len() <= 128);
assert!(snapshot.to_string().len() < 64 * 1024);
let kept: u64 = events
.iter()
.map(|event| event["count"].as_u64().unwrap())
.sum();
assert_eq!(
kept + snapshot["droppedEvents"].as_u64().unwrap(),
snapshot["totalEvents"]
);
}
#[test]
fn shared_attempts_and_exchanges_coalesce_delta_without_losing_metadata() {
let trace = TraceContext::new("req_trace");
let ws = trace.attempt(1).exchange("websocket");
ws.capture(
"upstream.event",
br#"{"type":"codex.response.metadata","headers":{"x-request-id":"up-1"}}"#,
);
trace
.attempt(1)
.wire_event("openai", Some("codex.response.metadata"));
for _ in 0..1000 {
ws.capture(
"upstream.event",
br#"{"type":"response.output_text.delta","delta":"secret"}"#,
);
trace
.attempt(1)
.wire_event("openai", Some("response.output_text.delta"));
}
let http = trace.attempt(1).exchange("http_sse");
http.capture("upstream.event", br#"{"type":"response.completed"}"#);
let snapshot = trace.snapshot().unwrap();
let events = snapshot["events"].as_array().unwrap();
assert_eq!(snapshot["droppedEvents"], 0);
assert_eq!(
events
.iter()
.find(|e| e["data"]["eventType"] == "response.output_text.delta")
.unwrap()["count"],
1000
);
assert_eq!(events.last().unwrap()["exchangeId"], 2);
assert!(
events
.iter()
.any(|e| e["data"]["metadata"]["headers"]["x-request-id"] == "up-1")
);
assert!(!snapshot.to_string().contains("secret"));
assert!(
events
.iter()
.filter(|event| event["stage"] == "provider.event")
.all(|event| event["data"].get("metadata").is_none())
);
}
#[test]
fn header_capture_keeps_duplicates_but_excludes_credentials_from_timeline() {
let trace = TraceContext::new("req_headers");
trace.headers(
"upstream.response.headers",
json!({"status": 101}),
[
("authorization", b"Bearer SECRET".as_slice()),
("x-request-id", b"upstream-one".as_slice()),
("x-request-id", b"upstream-two".as_slice()),
("set-cookie", b"PRIVATE_COOKIE".as_slice()),
],
);
let snapshot = trace.snapshot().unwrap();
assert_eq!(
snapshot["events"][0]["data"]["headers"]["x-request-id"],
json!(["upstream-one", "upstream-two"])
);
assert!(!snapshot.to_string().contains("SECRET"));
assert!(!snapshot.to_string().contains("PRIVATE_COOKIE"));
}
#[test]
fn disabled_context_does_not_allocate_a_history() {
let trace = TraceContext::default().attempt(2).exchange("http");
trace.capture("upstream.event", b"not json");
assert!(trace.snapshot().is_none());
}
#[test]
fn a_long_successful_retry_keeps_the_earlier_failure_ahead_of_routine_events() {
let trace = TraceContext::new("req_retry");
for _ in 0..12 {
trace.record("setup", json!({}));
}
trace
.attempt(1)
.record("attempt.failed", json!({"kind": "unavailable"}));
trace
.attempt(1)
.record("retry.decided", json!({"retryable": true}));
for index in 0..500 {
trace
.attempt(2)
.record("upstream.event", json!({"index": index}));
}
trace.record("request.finished", json!({"outcome": "succeeded"}));
let snapshot = trace.snapshot().unwrap();
let events = snapshot["events"].as_array().unwrap();
assert!(
events
.iter()
.any(|event| event["stage"] == "attempt.failed" && event["attemptIndex"] == 1)
);
assert!(events.iter().any(|event| event["stage"] == "retry.decided"));
assert_eq!(events.last().unwrap()["stage"], "request.finished");
assert!(events.len() <= 128);
}
@@ -49,6 +49,7 @@ use serde_json::{Map, Value, json};
#[derive(Debug, Clone, PartialEq, Eq)]
struct FinalState {
diagnostic_trace_json: Option<String>,
outcome: ExecutionOutcome,
send_state: UpstreamSendState,
attempt_count: u32,
@@ -207,6 +208,7 @@ impl ExecutionStore for FakeStore {
.expect("store lock")
.finalizations
.push(FinalState {
diagnostic_trace_json: finalization.diagnostic_trace_json,
outcome: finalization.outcome,
send_state: finalization.send_state,
attempt_count: finalization.attempt_count,
@@ -270,6 +272,7 @@ enum Script {
struct ScriptedProvider {
scripts: Mutex<VecDeque<Script>>,
contexts: Mutex<Vec<AttemptContext>>,
operations: Mutex<Vec<Operation>>,
}
impl ScriptedProvider {
@@ -277,6 +280,7 @@ impl ScriptedProvider {
Self {
scripts: Mutex::new(scripts.into()),
contexts: Mutex::new(Vec::new()),
operations: Mutex::new(Vec::new()),
}
}
}
@@ -302,6 +306,10 @@ impl Provider for ScriptedProvider {
request: ProviderRequest,
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
self.operations
.lock()
.expect("operations lock")
.push(request.operation().clone());
self.contexts.lock().expect("contexts lock").push(context);
let script = self
.scripts
@@ -1308,6 +1316,21 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
let state = store.state.lock().expect("store lock");
let finalization = &state.finalizations[0];
let trace: Value =
serde_json::from_str(finalization.diagnostic_trace_json.as_deref().unwrap()).unwrap();
let events = trace["events"].as_array().unwrap();
assert!(
events
.iter()
.any(|event| event["stage"] == "attempt.failed" && event["attemptIndex"] == 1)
);
assert!(
events
.iter()
.any(|event| event["stage"] == "attempt.started" && event["attemptIndex"] == 2)
);
assert_eq!(events.last().unwrap()["stage"], "request.finished");
assert_eq!(finalization.upstream_transport.as_deref(), Some("http_sse"));
assert_eq!(finalization.http_version.as_deref(), Some("HTTP/2"));
assert_eq!(
@@ -2838,14 +2861,25 @@ fn ambiguous_send_state_stops_retry() {
fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() {
let operation = generate_operation();
let route_plan = plan(&operation);
let checkpoint = ProviderSessionState::new(
"openai",
Map::from_iter([("checkpoint".to_owned(), json!("opaque-state"))]),
)
.expect("provider checkpoint");
let mut checkpoint_event = ProviderEvent::observation(ProviderResponseObservation::new(
UpstreamTransport::new("websocket").expect("transport"),
));
checkpoint_event.attach_session_update(checkpoint);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
Script::ObservedStream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::Ambiguous,
)
.with_pre_delivery_retry())],
items: vec![
Ok(checkpoint_event),
Err(
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
.with_pre_delivery_retry(),
),
],
},
Script::Stream {
account_id: "acct_second",
@@ -2874,6 +2908,57 @@ fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() {
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn ambiguous_transport_retry_is_rejected_even_with_provider_session_state() {
let operation = generate_operation();
let route_plan = plan(&operation);
let checkpoint = ProviderSessionState::new(
"openai",
Map::from_iter([("checkpoint".to_owned(), json!("opaque-state"))]),
)
.expect("provider checkpoint");
let mut checkpoint_event = ProviderEvent::observation(ProviderResponseObservation::new(
UpstreamTransport::new("websocket").expect("transport"),
));
checkpoint_event.attach_session_update(checkpoint);
let retry_index = NonZeroU32::new(1).expect("retry index");
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_first",
items: vec![
Ok(checkpoint_event),
Err(
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
.with_replay_safe()
.with_pre_delivery_transport_retry(retry_index, Duration::ZERO),
),
],
},
Script::Stream {
account_id: "acct_first",
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(),
))
.expect("start execution");
let error =
block_on(session.collect_uncommitted()).expect_err("ambiguous send cannot be replayed");
assert!(matches!(error, EngineError::Provider(_)));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 0);
assert_eq!(state.finalizations[0].attempt_count, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn pre_delivery_transport_fallback_retries_the_same_account_with_fallback_transport() {
let operation = generate_operation();
@@ -1,4 +1,5 @@
mod account;
mod diagnostics;
mod engine;
mod error;
mod event;
@@ -118,6 +118,9 @@ pub struct LoggingConfig {
pub level: String,
pub stdout: bool,
pub file: FileLoggingConfig,
/// 将完整请求原文写入独立诊断日志;默认关闭,因为内容包含凭据。
#[serde(default)]
pub request_dump: bool,
}
impl LoggingConfig {
+68 -26
View File
@@ -6,16 +6,19 @@ use std::io::{self, Write};
use std::path::{Path, PathBuf};
use chrono::Utc;
use tracing_appender::non_blocking::{NonBlockingBuilder, WorkerGuard};
use tracing_appender::non_blocking::{NonBlocking, NonBlockingBuilder, WorkerGuard};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::{Layer as _, layer::SubscriberExt as _, util::SubscriberInitExt as _};
use crate::config::LoggingConfig;
use crate::config::{FileLoggingConfig, LoggingConfig};
const LOG_FILE_PREFIX: &str = "codex-proxy-rs";
const OAUTH_RECOVERY_LOG_FILE_PREFIX: &str = "codex-proxy-rs-oath";
/// OAuth 恢复记录含原始 AT/RT,文件日志开启时必须落到独立文件。
const OAUTH_RECOVERY_LOG_TARGET: &str = "oauth_recovery";
const REQUEST_DUMP_LOG_FILE_PREFIX: &str = "codex-proxy-rs-request-dump";
/// 请求转储含完整凭据与正文,只允许在显式开启时写入独立文件。
const REQUEST_DUMP_LOG_TARGET: &str = "request_dump";
/// non-blocking 日志 writer 的进程级守卫。
pub struct LogGuard {
@@ -39,6 +42,7 @@ pub fn initialize_logging(config: &LoggingConfig) -> Result<LogGuard, LogError>
let directive = env::var("RUST_LOG").unwrap_or_else(|_| config.level.clone());
let file_filter = application_file_filter(&directive)?;
let recovery_file_filter = oauth_recovery_file_filter();
let request_dump_file_filter = request_dump_file_filter();
let stdout_filter =
EnvFilter::try_new(stdout_filter_directive(&directive, config.file.enabled))
.map_err(|_| LogError::InvalidFilter)?;
@@ -52,37 +56,30 @@ pub fn initialize_logging(config: &LoggingConfig) -> Result<LogGuard, LogError>
writer
});
let (file_writer, recovery_file_writer) = if config.file.enabled {
let maximum_bytes = config
.file
.max_file_size_mb
.checked_mul(1024 * 1024)
.ok_or(LogError::SizeOverflow)?;
let file_writer = RotatingLogWriter::open(
config.file.directory.clone(),
LOG_FILE_PREFIX,
maximum_bytes,
config.file.retention_days,
config.file.max_files,
)?;
let recovery_file_writer = RotatingLogWriter::open(
config.file.directory.clone(),
let (file_writer, file_guard) =
rotating_file_writer(&config.file, LOG_FILE_PREFIX, "gateway-log-file")?;
let (recovery_file_writer, recovery_file_guard) = rotating_file_writer(
&config.file,
OAUTH_RECOVERY_LOG_FILE_PREFIX,
maximum_bytes,
config.file.retention_days,
config.file.max_files,
"gateway-log-oauth-recovery-file",
)?;
let (file_writer, file_guard) = NonBlockingBuilder::default()
.thread_name("gateway-log-file")
.finish(file_writer);
let (recovery_file_writer, recovery_file_guard) = NonBlockingBuilder::default()
.thread_name("gateway-log-oauth-recovery-file")
.finish(recovery_file_writer);
guards.push(file_guard);
guards.push(recovery_file_guard);
(Some(file_writer), Some(recovery_file_writer))
} else {
(None, None)
};
let request_dump_file_writer = if config.request_dump {
let (writer, guard) = rotating_file_writer(
&config.file,
REQUEST_DUMP_LOG_FILE_PREFIX,
"gateway-log-request-dump-file",
)?;
guards.push(guard);
Some(writer)
} else {
None
};
let stdout_layer = stdout_writer.map(|writer| {
tracing_subscriber::fmt::layer()
@@ -118,10 +115,24 @@ pub fn initialize_logging(config: &LoggingConfig) -> Result<LogGuard, LogError>
.with_span_list(true)
.with_filter(recovery_file_filter)
});
let request_dump_file_layer = request_dump_file_writer.map(|writer| {
tracing_subscriber::fmt::layer()
.json()
.with_writer(writer)
.with_target(true)
.with_file(true)
.with_line_number(true)
.with_thread_ids(true)
.with_thread_names(true)
.with_current_span(true)
.with_span_list(true)
.with_filter(request_dump_file_filter)
});
tracing_subscriber::registry()
.with(stdout_layer)
.with(file_layer)
.with(recovery_file_layer)
.with(request_dump_file_layer)
.try_init()
.map_err(|_| LogError::AlreadyInitialized)?;
Ok(LogGuard { _writers: guards })
@@ -131,8 +142,12 @@ fn application_file_filter(directive: &str) -> Result<EnvFilter, LogError> {
let recovery_directive = format!("{OAUTH_RECOVERY_LOG_TARGET}=off")
.parse()
.expect("static OAuth recovery log directive is valid");
let request_dump_directive = format!("{REQUEST_DUMP_LOG_TARGET}=off")
.parse()
.expect("static request dump log directive is valid");
EnvFilter::try_new(directive)
.map(|filter| filter.add_directive(recovery_directive))
.map(|filter| filter.add_directive(request_dump_directive))
.map_err(|_| LogError::InvalidFilter)
}
@@ -142,14 +157,41 @@ fn oauth_recovery_file_filter() -> EnvFilter {
.expect("static OAuth recovery log filter is valid")
}
fn request_dump_file_filter() -> EnvFilter {
format!("off,{REQUEST_DUMP_LOG_TARGET}=info")
.parse()
.expect("static request dump log filter is valid")
}
fn stdout_filter_directive(directive: &str, persistent_log_enabled: bool) -> String {
if !persistent_log_enabled {
directive.to_owned()
format!("{directive},{REQUEST_DUMP_LOG_TARGET}=off")
} else {
"off,gateway_startup=info".to_owned()
}
}
fn rotating_file_writer(
config: &FileLoggingConfig,
file_prefix: &'static str,
thread_name: &'static str,
) -> Result<(NonBlocking, WorkerGuard), LogError> {
let maximum_bytes = config
.max_file_size_mb
.checked_mul(1024 * 1024)
.ok_or(LogError::SizeOverflow)?;
let writer = RotatingLogWriter::open(
config.directory.clone(),
file_prefix,
maximum_bytes,
config.retention_days,
config.max_files,
)?;
Ok(NonBlockingBuilder::default()
.thread_name(thread_name)
.finish(writer))
}
struct RotatingLogWriter {
directory: PathBuf,
file_prefix: &'static str,
@@ -103,6 +103,7 @@ fn valid_config() -> HostConfig {
max_file_size_mb: 100,
max_files: 30,
},
request_dump: false,
},
system_update,
drain_timeout_seconds: 30,
+76 -7
View File
@@ -8,13 +8,18 @@ use gateway_host::system_update::SystemUpdateConfig;
const LOG_DIRECTORY_ENV: &str = "CPR_LOGGING_TEST_DIRECTORY";
const CHILD_PROCESS_ENV: &str = "CPR_LOGGING_TEST_CHILD";
const REQUEST_DUMP_ENABLED_ENV: &str = "CPR_LOGGING_TEST_REQUEST_DUMP_ENABLED";
const APPLICATION_LOG_FILE_PREFIX: &str = "codex-proxy-rs.";
const OAUTH_RECOVERY_LOG_FILE_PREFIX: &str = "codex-proxy-rs-oath.";
const REQUEST_DUMP_LOG_FILE_PREFIX: &str = "codex-proxy-rs-request-dump.";
const APPLICATION_LOG_TARGET: &str = "logging_test_application";
const APPLICATION_LOG_MARKER: &str = "application-file-filter-test";
const OAUTH_RECOVERY_LOG_TARGET: &str = "oauth_recovery";
const OAUTH_RECOVERY_LOG_MARKER: &str = "oauth-recovery-file-filter-test";
const OAUTH_RECOVERY_PROVIDER: &str = "openai";
const REQUEST_DUMP_LOG_TARGET: &str = "request_dump";
const REQUEST_DUMP_LOG_MARKER: &str = "request-dump-file-filter-test";
const REQUEST_DUMP_SECRET: &str = "unredacted-authorization-value";
#[test]
fn logging_requires_at_least_one_sink() {
@@ -34,6 +39,7 @@ fn logging_requires_at_least_one_sink() {
max_file_size_mb: 100,
max_files: 30,
},
request_dump: false,
},
system_update: SystemUpdateConfig::default(),
drain_timeout_seconds: 30,
@@ -48,11 +54,12 @@ fn logging_requires_at_least_one_sink() {
}
#[test]
fn oauth_recovery_file_logging_is_separate_and_overrides_global_log_level() {
fn sensitive_file_logging_is_separate_and_overrides_global_log_level() {
if env::var_os(CHILD_PROCESS_ENV).is_some() {
write_oauth_recovery_log(PathBuf::from(
env::var_os(LOG_DIRECTORY_ENV).expect("child log directory"),
));
write_sensitive_logs(
PathBuf::from(env::var_os(LOG_DIRECTORY_ENV).expect("child log directory")),
env::var(REQUEST_DUMP_ENABLED_ENV).is_ok_and(|value| value == "true"),
);
return;
}
@@ -60,14 +67,15 @@ fn oauth_recovery_file_logging_is_separate_and_overrides_global_log_level() {
let output = Command::new(env::current_exe().expect("current test executable"))
.args([
"--exact",
"logging::oauth_recovery_file_logging_is_separate_and_overrides_global_log_level",
"logging::sensitive_file_logging_is_separate_and_overrides_global_log_level",
"--nocapture",
])
.env(CHILD_PROCESS_ENV, "1")
.env(LOG_DIRECTORY_ENV, directory.path())
.env(REQUEST_DUMP_ENABLED_ENV, "true")
.env(
"RUST_LOG",
"off,logging_test_application=info,oauth_recovery=off",
"off,logging_test_application=info,oauth_recovery=off,request_dump=trace",
)
.output()
.expect("run logging child process");
@@ -84,6 +92,8 @@ fn oauth_recovery_file_logging_is_separate_and_overrides_global_log_level() {
assert!(application_log.contains(APPLICATION_LOG_MARKER));
assert!(!application_log.contains(&json_target_field(OAUTH_RECOVERY_LOG_TARGET)));
assert!(!application_log.contains(OAUTH_RECOVERY_LOG_MARKER));
assert!(!application_log.contains(&json_target_field(REQUEST_DUMP_LOG_TARGET)));
assert!(!application_log.contains(REQUEST_DUMP_SECRET));
let recovery_log = read_log_file_set(directory.path(), OAUTH_RECOVERY_LOG_FILE_PREFIX);
assert!(recovery_log.contains(&json_target_field(OAUTH_RECOVERY_LOG_TARGET)));
@@ -91,9 +101,49 @@ fn oauth_recovery_file_logging_is_separate_and_overrides_global_log_level() {
assert!(!recovery_log.contains(&json_target_field(APPLICATION_LOG_TARGET)));
assert!(!recovery_log.contains(APPLICATION_LOG_MARKER));
assert!(recovery_log.contains(&format!(r#""provider":"{OAUTH_RECOVERY_PROVIDER}""#)));
assert!(!recovery_log.contains(&json_target_field(REQUEST_DUMP_LOG_TARGET)));
assert!(!recovery_log.contains(REQUEST_DUMP_SECRET));
let request_dump_log = read_log_file_set(directory.path(), REQUEST_DUMP_LOG_FILE_PREFIX);
assert!(request_dump_log.contains(&json_target_field(REQUEST_DUMP_LOG_TARGET)));
assert!(request_dump_log.contains(REQUEST_DUMP_LOG_MARKER));
assert!(request_dump_log.contains(REQUEST_DUMP_SECRET));
assert!(!request_dump_log.contains(&json_target_field(APPLICATION_LOG_TARGET)));
assert!(!request_dump_log.contains(APPLICATION_LOG_MARKER));
assert!(!request_dump_log.contains(&json_target_field(OAUTH_RECOVERY_LOG_TARGET)));
assert!(!request_dump_log.contains(OAUTH_RECOVERY_LOG_MARKER));
let disabled_directory = tempfile::tempdir().expect("create disabled log directory");
let output = Command::new(env::current_exe().expect("current test executable"))
.args([
"--exact",
"logging::sensitive_file_logging_is_separate_and_overrides_global_log_level",
"--nocapture",
])
.env(CHILD_PROCESS_ENV, "1")
.env(LOG_DIRECTORY_ENV, disabled_directory.path())
.env(REQUEST_DUMP_ENABLED_ENV, "false")
.env(
"RUST_LOG",
"off,logging_test_application=info,oauth_recovery=off,request_dump=trace",
)
.output()
.expect("run logging child process with request dump disabled");
assert!(
output.status.success(),
"logging child process failed:\nstdout:\n{}\nstderr:\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr),
);
assert!(!log_file_set_exists(
disabled_directory.path(),
REQUEST_DUMP_LOG_FILE_PREFIX
));
let application_log = read_log_file_set(disabled_directory.path(), APPLICATION_LOG_FILE_PREFIX);
assert!(!application_log.contains(REQUEST_DUMP_SECRET));
}
fn write_oauth_recovery_log(directory: PathBuf) {
fn write_sensitive_logs(directory: PathBuf, request_dump: bool) {
let config = HostConfig {
listen: ListenConfig {
host: "127.0.0.1".to_owned(),
@@ -110,6 +160,7 @@ fn write_oauth_recovery_log(directory: PathBuf) {
max_file_size_mb: 1,
max_files: 1,
},
request_dump,
},
system_update: SystemUpdateConfig::default(),
drain_timeout_seconds: 30,
@@ -134,6 +185,12 @@ fn write_oauth_recovery_log(directory: PathBuf) {
marker = OAUTH_RECOVERY_LOG_MARKER,
"OAuth recovery test record"
);
tracing::info!(
target: REQUEST_DUMP_LOG_TARGET,
marker = REQUEST_DUMP_LOG_MARKER,
authorization = REQUEST_DUMP_SECRET,
"request dump test record"
);
drop(bundle);
}
@@ -156,6 +213,18 @@ fn read_log_file_set(directory: &Path, file_prefix: &str) -> String {
.collect()
}
fn log_file_set_exists(directory: &Path, file_prefix: &str) -> bool {
fs::read_dir(directory)
.expect("read log directory")
.filter_map(Result::ok)
.any(|entry| {
entry
.file_name()
.to_str()
.is_some_and(|name| name.starts_with(file_prefix))
})
}
fn json_target_field(target: &str) -> String {
format!(r#""target":"{target}""#)
}
@@ -4,6 +4,41 @@ const MULTI_AGENT_MODE_OPEN_TAG: &str = "<multi_agent_mode>";
const MULTI_AGENT_MODE_CLOSE_TAG: &str = "</multi_agent_mode>";
const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegation is active.";
/// 各 Codex 端点共用的显式根会话身份。连接边界的 session_id 优先,
/// 正文及官方 metadata 只做回退;不从 turn、window 或请求内容猜测归属。
#[must_use]
pub fn codex_session_id(body: &Map<String, Value>, context: &Map<String, Value>) -> Option<String> {
codex_identity_field(body, context, "session_id")
}
/// 各 Codex 端点共用的显式线程身份,与根会话采用相同的来源优先级。
#[must_use]
pub fn codex_thread_id(body: &Map<String, Value>, context: &Map<String, Value>) -> Option<String> {
codex_identity_field(body, context, "thread_id")
}
fn codex_identity_field(
body: &Map<String, Value>,
context: &Map<String, Value>,
field: &str,
) -> Option<String> {
non_empty_string(context.get(field))
.or_else(|| non_empty_string(body.get(field)))
.or_else(|| client_metadata_string(body, field))
.map(str::to_owned)
.or_else(|| {
let metadata = non_empty_string(context.get("turn_metadata"))
.or_else(|| request_turn_metadata(body))?;
let metadata: Value = serde_json::from_str(metadata).ok()?;
non_empty_string(metadata.get(field)).map(str::to_owned)
})
.or_else(|| {
body.get("context")
.and_then(|context| non_empty_string(context.get(field)))
.map(str::to_owned)
})
}
/// 从 OpenAI Responses 请求中提取的稳定 Codex 请求语义。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodexResponsesRequestSemantics {
@@ -98,6 +98,37 @@ pub fn extract_usage(body: &Value) -> Option<TokenUsage> {
})
}
/// 校验计价所需的输入、输出和缓存读取字段,拒绝缺失、非法或不一致的数值。
///
/// 用量提取允许缺失字段补零;计价必须额外确认这些零值有明确的上游事实。
#[must_use]
pub fn billable_usage_is_complete(response: &Value, usage: TokenUsage) -> bool {
let Some(raw) = response.get("usage").filter(|value| value.is_object()) else {
return false;
};
let input = raw
.get("input_tokens")
.or_else(|| raw.get("prompt_tokens"))
.and_then(Value::as_u64);
let output = raw
.get("output_tokens")
.or_else(|| raw.get("completion_tokens"))
.and_then(Value::as_u64);
let cached = raw
.pointer("/input_tokens_details/cached_tokens")
.or_else(|| raw.pointer("/prompt_tokens_details/cached_tokens"))
.or_else(|| raw.get("cached_tokens"));
let cached = match cached {
Some(value) => value.as_u64(),
None => Some(0),
};
input == Some(usage.input_tokens)
&& output == Some(usage.output_tokens)
&& cached == Some(usage.cached_tokens)
&& usage.cached_tokens <= usage.input_tokens
}
/// 从完整 SSE 文本中提取最终可见用量。
///
/// 如果存在 `response.completed` 的 usage,则优先返回它;否则回退到最后一条
@@ -18,5 +18,5 @@ pub mod sse;
pub use codex::{
CodexResponsesRequestSemantics, codex_responses_request_semantics,
codex_responses_request_semantics_with_turn_metadata,
codex_responses_request_semantics_with_turn_metadata, codex_session_id, codex_thread_id,
};
@@ -1,6 +1,101 @@
use gateway_protocol::openai::codex_responses_request_semantics;
use gateway_protocol::openai::{
codex_responses_request_semantics, codex_session_id, codex_thread_id,
};
use serde_json::{Map, Value, json};
#[test]
fn codex_session_should_prefer_the_connection_identity_over_body_metadata() {
let body = object(
json!({"session_id":"body-session", "client_metadata":{"session_id":"metadata-session"}}),
);
let context = object(
json!({"session_id":" root-session ", "turn_metadata":"{\"session_id\":\"turn-session\"}"}),
);
assert_eq!(
codex_session_id(&body, &context).as_deref(),
Some("root-session")
);
}
#[test]
fn codex_session_should_accept_explicit_metadata_and_ignore_unrelated_ids() {
for (body, context) in [
(json!({"session_id":"root"}), json!({})),
(json!({"client_metadata":{"session_id":"root"}}), json!({})),
(
json!({"client_metadata":{"x-codex-turn-metadata":"{\"session_id\":\"root\"}"}}),
json!({}),
),
(
json!({}),
json!({"turn_metadata":"{\"session_id\":\"root\"}"}),
),
(json!({"context":{"session_id":"root"}}), json!({})),
] {
assert_eq!(
codex_session_id(&object(body), &object(context)).as_deref(),
Some("root")
);
}
for body in [
json!({"session_id":" ", "thread_id":"child", "turn_id":"turn", "id":"request"}),
json!({"client_metadata":{"x-codex-turn-metadata":"invalid JSON"}}),
] {
assert!(codex_session_id(&object(body), &Map::new()).is_none());
}
}
#[test]
fn codex_thread_should_use_explicit_identity_and_never_guess_from_other_ids() {
for (body, context, expected) in [
(
json!({"thread_id":"body", "client_metadata":{"thread_id":"metadata"}}),
json!({"thread_id":" child "}),
Some("child"),
),
(
json!({"thread_id":"body", "client_metadata":{"thread_id":"metadata"}}),
json!({}),
Some("body"),
),
(
json!({"client_metadata":{"thread_id":"metadata"}}),
json!({}),
Some("metadata"),
),
(
json!({"client_metadata":{"x-codex-turn-metadata":"{\"thread_id\":\"metadata-thread\"}"}}),
json!({}),
Some("metadata-thread"),
),
(
json!({"turnMetadata":"{\"thread_id\":\"body-turn\"}"}),
json!({"turn_metadata":"{\"thread_id\":\"header-turn\"}"}),
Some("header-turn"),
),
(
json!({"context":{"thread_id":"context-thread"}}),
json!({}),
Some("context-thread"),
),
(
json!({"thread_id":" ", "session_id":"root", "turn_id":"turn", "agent_name":"agent"}),
json!({}),
None,
),
(
json!({"client_metadata":{"x-codex-turn-metadata":"invalid JSON"}}),
json!({}),
None,
),
] {
assert_eq!(
codex_thread_id(&object(body), &object(context)).as_deref(),
expected
);
}
}
#[test]
fn codex_semantics_should_use_protocol_context_before_body_metadata() {
let body = json!({
@@ -1,10 +1,55 @@
use gateway_protocol::openai::events::{
RateLimitKeySource, RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage,
is_codex_quota_header_name, is_rate_limit_header_name, parse_rate_limit_headers,
RateLimitKeySource, RateLimitWindow, TokenUsage, billable_usage_is_complete, extract_sse_usage,
extract_usage, is_codex_quota_header_name, is_rate_limit_header_name, parse_rate_limit_headers,
parse_rate_limits_event, rate_limits_to_header_pairs, retry_after_seconds_from_body,
};
use serde_json::json;
#[test]
fn billable_usage_should_require_valid_input_output_and_cache_facts() {
for (raw, expected) in [
(json!({}), false),
(json!({"input_tokens": 0}), false),
(json!({"input_tokens": 0, "output_tokens": 0}), true),
(
json!({"prompt_tokens": 8, "completion_tokens": 4, "cached_tokens": 2}),
true,
),
(json!({"input_tokens": -1, "output_tokens": 1}), false),
(
json!({"input_tokens": "8", "prompt_tokens": 8, "output_tokens": 1}),
false,
),
(json!({"input_tokens": 8, "output_tokens": 1.5}), false),
(
json!({"input_tokens": 8, "output_tokens": 1, "cached_tokens": 9}),
false,
),
(
json!({"input_tokens": 8, "output_tokens": 1, "cached_tokens": null}),
false,
),
(
json!({"input_tokens": 8, "output_tokens": 1, "input_tokens_details": {"cached_tokens": "2"}, "cached_tokens": 2}),
false,
),
(
json!({"input_tokens": 8, "prompt_tokens": 9, "output_tokens": 1}),
true,
),
] {
let body = json!({"usage": raw});
let usage = extract_usage(&body).unwrap_or_default();
assert_eq!(billable_usage_is_complete(&body, usage), expected, "{body}");
}
let body = json!({"usage": {"input_tokens": 8, "output_tokens": 4}});
assert!(!billable_usage_is_complete(&body, TokenUsage::default()));
assert!(!billable_usage_is_complete(
&json!({}),
TokenUsage::default()
));
}
#[test]
fn extract_usage_should_read_codex_usage_shape() {
let body = json!({
+1
View File
@@ -12,6 +12,7 @@ name = "main"
path = "tests/main.rs"
[dependencies]
async-stream = "0.3.6"
async-trait = "0.1.89"
aws-credential-types = "1.2"
aws-sdk-s3 = { version = "1.140", default-features = false, features = ["rt-tokio", "default-https-client"] }
@@ -59,12 +59,25 @@ impl BackupRepository for PgBackupRepository {
.map_err(|_| store_unavailable("begin backup storage update"))?;
let result = async {
let current = lock_settings_in_transaction(&mut transaction).await?;
let mut changed = storage_changed_fields(&current, &command);
if changed.is_empty() {
let revision: i64 =
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
.fetch_one(&mut *transaction)
.await
.map_err(|_| store_unavailable("load unchanged backup revision"))?;
let revision = AdminRevision::new(
u64::try_from(revision).map_err(|_| store_invalid("config revision"))?,
)
.map_err(|_| store_invalid("config revision"))?;
return Ok((current, revision));
}
ensure_storage_identity_stable(&current, &command, &mut transaction).await?;
let secret = command
.secret_access_key
.as_ref()
.map(|secret| secret.expose_secret().to_owned());
.map(|secret| secret.expose_secret());
sqlx::query(
"update backup_settings
set endpoint = $1,
@@ -76,6 +89,8 @@ impl BackupRepository for PgBackupRepository {
force_path_style = $7,
storage_revision = storage_revision + 1,
last_verified_at = null,
schedule_enabled = false,
next_run_at = null,
updated_at = now()
where id = 1",
)
@@ -83,7 +98,7 @@ impl BackupRepository for PgBackupRepository {
.bind(&command.region)
.bind(&command.bucket)
.bind(&command.access_key_id)
.bind(secret.as_deref())
.bind(secret)
.bind(&command.prefix)
.bind(command.force_path_style)
.execute(&mut *transaction)
@@ -94,7 +109,13 @@ impl BackupRepository for PgBackupRepository {
.await
.map_err(map_admin_error)?;
let revision = admin_revision(store_revision)?;
let changed = storage_changed_fields(&current, &command);
if current.last_verified_at.is_some() {
changed.push("last_verified_at".to_owned());
}
if current.schedule_enabled {
changed.push("schedule_enabled".to_owned());
changed.push("next_run_at".to_owned());
}
let audit = crate::mutation_audit(
context,
"backup.s3_config_updated",
@@ -719,7 +740,13 @@ fn storage_changed_fields(
if current.access_key_id.as_deref() != Some(command.access_key_id.as_str()) {
fields.push("access_key_id".to_owned());
}
if command.secret_access_key.is_some() {
if command.secret_access_key.as_ref().is_some_and(|secret| {
current
.secret_access_key
.as_ref()
.map(|value| value.expose_secret())
!= Some(secret.expose_secret())
}) {
fields.push("secret_access_key".to_owned());
}
if current.prefix.as_deref() != Some(command.prefix.as_str()) {
@@ -283,6 +283,7 @@ pub struct ModelRequestFinalization {
pub websocket_pool: Option<String>,
pub service_tier: Option<String>,
pub provider_metadata_json: Option<Value>,
pub diagnostic_trace_json: Option<Value>,
pub error_kind: Option<String>,
pub provider_error_code: Option<String>,
pub error_message: Option<String>,
@@ -353,6 +354,15 @@ impl ModelRequestFinalization {
{
return Err(invalid("provider observation must be a JSON object"));
}
if self
.diagnostic_trace_json
.as_ref()
.is_some_and(|trace| !trace.is_object() || trace.to_string().len() > 64 * 1024)
{
return Err(invalid(
"diagnostic trace must be a JSON object within 64 KiB",
));
}
validate_optional_stable_reason(
self.continuation_unavailable_reason.as_deref(),
"continuation unavailable reason",
@@ -725,10 +735,11 @@ impl ModelRequestRepository for PgExecutionStore {
upstream_connection_id = $44,
upstream_connection_exit_reason = $45,
upstream_connection_age_ms = $46,
upstream_connection_idle_ms = $47
upstream_connection_idle_ms = $47, diagnostic_trace_json = $48
where id = $1 and outcome = 'running'
returning id, client_api_key_ref, continuation_affinity_hash,
continuation_requested, outcome, started_at, completed_at
continuation_requested, provider_kind, upstream_transport,
outcome, started_at, completed_at
), recovery_target as (
select prior.id
from model_requests prior
@@ -775,9 +786,47 @@ impl ModelRequestRepository for PgExecutionStore {
where prior.id = target.id
and prior.recovered_at is null
returning prior.id
), session_transport_recovery_target as (
select prior.id
from model_requests prior
cross join finalized current
where current.outcome = 'succeeded'
and current.provider_kind = 'openai'
and current.upstream_transport = 'http_sse'
and current.continuation_affinity_hash is not null
and prior.client_api_key_ref = current.client_api_key_ref
and prior.continuation_affinity_hash = current.continuation_affinity_hash
and prior.provider_kind = current.provider_kind
and prior.outcome = 'failed'
and prior.error_kind = 'upstream_unavailable'
and prior.upstream_transport = 'websocket'
and prior.upstream_send_state = 'ambiguous'
and prior.downstream_committed_at is null
and prior.recovered_at is null
and current.started_at >= prior.completed_at
and current.started_at <= prior.completed_at + interval '30 seconds'
), session_transport_recovery_update as (
update model_requests prior
set recovery_attempt_count = prior.recovery_attempt_count + 1,
recovery_request_id = current.id,
recovered_at = current.completed_at,
recovery_retry_delay_ms = greatest(
0,
floor(extract(epoch from (current.started_at - prior.completed_at)) * 1000)
)::bigint,
recovery_total_latency_ms = greatest(
0,
floor(extract(epoch from (current.completed_at - prior.completed_at)) * 1000)
)::bigint
from finalized current
join session_transport_recovery_target target on true
where prior.id = target.id
and prior.recovered_at is null
returning prior.id
)
select (select count(*) from finalized)::bigint
+ (select count(*) * 0 from recovery_update)::bigint",
+ (select count(*) * 0 from recovery_update)::bigint
+ (select count(*) * 0 from session_transport_recovery_update)::bigint",
)
.bind(&finalization.model_request_id)
.bind(finalization.outcome.as_str())
@@ -877,6 +926,7 @@ impl ModelRequestRepository for PgExecutionStore {
finalization.upstream_connection_idle_ms,
"upstream_connection_idle_ms",
)?)
.bind(finalization.diagnostic_trace_json.map(sqlx::types::Json))
.fetch_one(&self.pool)
.await
.map_err(|_| postgres_unavailable("finalize model request"))?;
@@ -1039,7 +1089,6 @@ impl ExecutionStore for PgExecutionStore {
|| "intermediate upstream failure".to_owned(),
|diagnostic| diagnostic.as_str().to_owned(),
),
occurrence_count: 1,
created_at: Utc::now(),
},
)
@@ -1084,7 +1133,6 @@ impl ExecutionStore for PgExecutionStore {
|| "account connection test failed".to_owned(),
|diagnostic| diagnostic.as_str().to_owned(),
),
occurrence_count: 1,
created_at: Utc::now(),
},
)
@@ -1176,6 +1224,12 @@ impl ExecutionStore for PgExecutionStore {
websocket_pool: finalization.websocket_pool,
service_tier: finalization.service_tier,
provider_metadata_json,
diagnostic_trace_json: finalization
.diagnostic_trace_json
.as_deref()
.map(serde_json::from_str)
.transpose()
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
error_kind,
provider_error_code: finalization.provider_error_code,
error_message,
@@ -583,6 +583,7 @@ impl ExecutionObservationWrite {
finalization.websocket_pool.as_deref(),
finalization.service_tier.as_deref(),
finalization.provider_metadata_json.as_deref(),
finalization.diagnostic_trace_json.as_deref(),
finalization.provider_error_code.as_deref(),
])
.saturating_add(error_bytes)
@@ -1,5 +1,7 @@
//! Pg 观测 adapter:实现 `ObservabilityRepository` 与 `AdminObservabilityStore`。
use futures::{StreamExt, TryStreamExt, stream::BoxStream};
use gateway_admin::ports::store::UsageCalculatedBillingStream;
use gateway_core::provider_ports::ProviderCooldownPort;
use std::sync::Arc;
@@ -185,17 +187,15 @@ impl ObservabilityRepository for PgObservabilityRepository {
.await
}
async fn usage_calculated_billing_facts(
fn usage_calculated_billing_facts(
&self,
range: ObservabilityRange,
filter: UsageRecordFilter,
) -> StoreResult<Vec<CalculatedUsageBillingFact>> {
self.query_budget
.run(
"load calculated usage billing facts",
calculated_usage_billing_facts(&self.pool, range, &filter),
)
.await
) -> BoxStream<'_, StoreResult<CalculatedUsageBillingFact>> {
self.query_budget.run_stream(
"load calculated usage billing facts",
calculated_usage_billing_facts(&self.pool, range, filter),
)
}
async fn provider_account_usage(
@@ -387,18 +387,20 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
.collect()
}
async fn usage_calculated_billing_facts(
fn usage_calculated_billing_facts(
&self,
range: admin_observability::TimeRange,
filter: admin_observability::UsageFilter,
) -> AdminStoreResult<Vec<admin_observability::UsageCalculatedBillingFact>> {
) -> UsageCalculatedBillingStream<'_> {
let range = match store_range(range) {
Ok(range) => range,
Err(error) => return Box::pin(futures::stream::once(async { Err(error) })),
};
self.repository
.usage_calculated_billing_facts(store_range(range)?, store_usage_filter(filter))
.await
.map_err(observability_error)?
.into_iter()
.map(admin_calculated_usage_billing_fact)
.collect()
.usage_calculated_billing_facts(range, store_usage_filter(filter))
.map_err(observability_error)
.map(|fact| fact.and_then(admin_calculated_usage_billing_fact))
.boxed()
}
async fn list_usage_records(
@@ -319,6 +319,7 @@ pub(crate) fn admin_dashboard_account_usage(
)
.collect(),
quota_used_percent: None,
quota_window: None,
models: usage
.models
.into_iter()
@@ -583,6 +584,8 @@ pub(crate) fn admin_usage_detail(
detail: UsageRecordDetail,
) -> AdminStoreResult<admin_observability::UsageDetail> {
Ok(admin_observability::UsageDetail {
trace: detail.trace,
related_requests: detail.related_requests,
request: admin_usage_record(detail.request)?,
attempts: detail
.attempts
@@ -721,7 +724,6 @@ pub(crate) fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::Ops
latency_ms: error.latency_ms,
message: error.message,
raw_upstream_error: error.raw_upstream_error,
occurrence_count: error.occurrence_count,
client_ip: error.client_ip,
user_agent: error.user_agent,
reasoning_effort: error.reasoning_effort,
@@ -912,7 +914,6 @@ pub(crate) fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<Ops
latency_ms: optional_unsigned(row, "latency_ms")?,
message: get(row, "message")?,
raw_upstream_error: get(row, "raw_upstream_error")?,
occurrence_count: to_u32(get(row, "occurrence_count")?)?,
client_ip: get(row, "client_ip")?,
user_agent: get(row, "user_agent")?,
reasoning_effort: get(row, "reasoning_effort")?,
@@ -1,6 +1,7 @@
//! 查询模型、校验与观测端口契约。
use super::*;
use futures::stream::BoxStream;
pub(crate) const MAX_PAGE_SIZE: u16 = 100;
pub(crate) const MAX_FILTER_BYTES: usize = 256;
@@ -644,6 +645,8 @@ pub struct UsageAttemptObservation {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageRecordDetail {
pub trace: Option<serde_json::Value>,
pub related_requests: Vec<serde_json::Value>,
pub request: UsageRecord,
pub attempts: Vec<UsageAttemptObservation>,
}
@@ -714,7 +717,6 @@ pub struct OpsErrorRecord {
pub latency_ms: Option<u64>,
pub message: String,
pub raw_upstream_error: Option<String>,
pub occurrence_count: u32,
pub client_ip: Option<String>,
pub user_agent: Option<String>,
pub reasoning_effort: Option<String>,
@@ -762,11 +764,11 @@ pub trait ObservabilityRepository: Send + Sync {
range: ObservabilityRange,
filter: UsageRecordFilter,
) -> StoreResult<Vec<RequestMetricPoint>>;
async fn usage_calculated_billing_facts(
fn usage_calculated_billing_facts(
&self,
range: ObservabilityRange,
filter: UsageRecordFilter,
) -> StoreResult<Vec<CalculatedUsageBillingFact>>;
) -> BoxStream<'_, StoreResult<CalculatedUsageBillingFact>>;
async fn provider_account_usage(
&self,
query: ProviderAccountUsageQuery,
@@ -283,42 +283,44 @@ async fn request_metric_series_inner(
fill_metric_gaps(range, granularity, points)
}
pub(crate) async fn calculated_usage_billing_facts(
pub(crate) fn calculated_usage_billing_facts(
pool: &PgPool,
range: ObservabilityRange,
filter: &UsageRecordFilter,
) -> StoreResult<Vec<CalculatedUsageBillingFact>> {
filter.validate()?;
let granularity = granularity_for(range);
let mut query = QueryBuilder::<Postgres>::new("select date_bin(");
query.push_bind(granularity.sql_interval());
query.push(
"::interval, mr.started_at, timestamptz '1970-01-01 00:00:00+00') as bucket_start,
mr.provider_kind, mr.upstream_model_id, mr.service_tier,
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.cache_write_tokens,
mr.cost_currency, mr.cost_amount::text as amount
from model_requests mr where mr.started_at >= ",
);
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
query.push(
" and mr.cost_source = 'calculated'
and mr.cost_amount is not null and mr.cost_currency is not null
and nullif(mr.provider_kind, '') is not null
and nullif(mr.upstream_model_id, '') is not null",
);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" order by bucket_start, mr.id");
query
.build()
.fetch_all(pool)
.await
.map_err(|_| postgres_unavailable("load calculated usage billing facts"))?
.iter()
.map(calculated_usage_billing_fact_from_row)
.collect()
filter: UsageRecordFilter,
) -> futures::stream::BoxStream<'_, StoreResult<CalculatedUsageBillingFact>> {
use futures::TryStreamExt;
Box::pin(async_stream::try_stream! {
filter.validate()?;
let granularity = granularity_for(range);
let mut query = QueryBuilder::<Postgres>::new("select date_bin(");
query.push_bind(granularity.sql_interval());
query.push(
"::interval, mr.started_at, timestamptz '1970-01-01 00:00:00+00') as bucket_start,
mr.provider_kind, mr.upstream_model_id, mr.service_tier,
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.cache_write_tokens,
mr.cost_currency, mr.cost_amount::text as amount
from model_requests mr where mr.started_at >= ",
);
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
query.push(
" and mr.cost_source = 'calculated'
and mr.cost_amount is not null and mr.cost_currency is not null
and nullif(mr.provider_kind, '') is not null
and nullif(mr.upstream_model_id, '') is not null",
);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, &filter, "mr");
// 费用逐条累加与行顺序无关,避免全量物化和不必要的数据库排序。
let mut rows = query.build().fetch(pool);
while let Some(row) = rows.try_next().await
.map_err(|_| postgres_unavailable("load calculated usage billing facts"))?
{
yield calculated_usage_billing_fact_from_row(&row)?;
}
})
}
pub(crate) async fn request_costs_by_bucket(
@@ -363,22 +365,22 @@ pub(crate) async fn attempt_metrics(
filter: &UsageRecordFilter,
) -> StoreResult<AttemptMetrics> {
filter.validate()?;
let fact = completed_usage_fact_predicate("sr");
let completed_usage = completed_usage_fact_predicate("mr");
// PERF: `model_requests` 包含宽请求元数据。这里只物化两个聚合族会复用的列,
// 避免生产基数下 `mr.*` 与多次标量子查询把 CTE 反复写入临时文件。
let mut query = QueryBuilder::<Postgres>::new(
let mut query = QueryBuilder::<Postgres>::new(format!(
"with selected_requests as materialized (
select mr.id, mr.attempt_count, mr.outcome, mr.error_kind,
mr.upstream_status_code, mr.client_status_code,
mr.downstream_committed_at, mr.client_transport, mr.cost_source
from model_requests mr where mr.started_at >= ",
);
mr.cost_source, ({completed_usage}) as is_completed_usage
from model_requests mr where mr.started_at >= "
));
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
push_unrecovered_request_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(format!(
query.push(
"), request_aggregate as (
select coalesce(sum(sr.attempt_count), 0)::bigint as attempt_count,
count(*) filter (
@@ -413,30 +415,30 @@ pub(crate) async fn attempt_metrics(
sr.client_status_code) between 500 and 599
)::bigint as request_provider_5xx_count,
count(*) filter (
where {fact} and sr.cost_source = 'provider_reported'
where sr.is_completed_usage and sr.cost_source = 'provider_reported'
)::bigint as provider_reported_count,
count(*) filter (
where {fact} and sr.cost_source = 'calculated'
where sr.is_completed_usage and sr.cost_source = 'calculated'
)::bigint as calculated_count,
count(*) filter (
where {fact} and sr.cost_source = 'unavailable'
where sr.is_completed_usage and sr.cost_source = 'unavailable'
)::bigint as unavailable_count
from selected_requests sr
), ops_aggregate as (
select coalesce(sum(coalesce(oe.occurrence_count, 1)), 0)::bigint
select count(*)::bigint
as ops_failure_count,
coalesce(sum(coalesce(oe.occurrence_count, 1)) filter (
count(*) filter (
where oe.failure_kind in ('rate_limited', 'quota_exhausted')
or oe.status_code = 429
), 0)::bigint as ops_rate_limited_count,
coalesce(sum(coalesce(oe.occurrence_count, 1)) filter (
)::bigint as ops_rate_limited_count,
count(*) filter (
where oe.failure_kind in
('authentication', 'authorization', 'invalid_credential')
or oe.status_code in (401, 403)
), 0)::bigint as ops_auth_failure_count,
coalesce(sum(coalesce(oe.occurrence_count, 1)) filter (
)::bigint as ops_auth_failure_count,
count(*) filter (
where oe.status_code between 500 and 599
), 0)::bigint as ops_provider_5xx_count
)::bigint as ops_provider_5xx_count
from ops_events oe
join selected_requests sr on sr.id = oe.model_request_id
)
@@ -455,8 +457,8 @@ pub(crate) async fn attempt_metrics(
requests.calculated_count,
requests.unavailable_count
from request_aggregate requests
cross join ops_aggregate ops"
));
cross join ops_aggregate ops",
);
let row = query
.build()
.fetch_one(pool)
@@ -20,7 +20,6 @@ const REQUEST_ERROR_SELECT: &str = "select 'model_request'::text as source,
mr.latency_ms,
coalesce(mr.error_message, mr.error_kind, 'request failed') as message,
mr.raw_upstream_error,
1::integer as occurrence_count,
host(mr.client_ip) as client_ip, mr.user_agent,
mr.reasoning_effort, mr.reasoning_preset, mr.request_kind,
mr.subagent_kind, mr.compact,
@@ -48,7 +47,7 @@ const OPS_EVENT_SELECT: &str = "select 'ops_event'::text as source,
oe.upstream_send_state,
null::integer as client_status_code, oe.status_code as upstream_status_code,
oe.provider_error_code, mr.client_response_id, oe.upstream_request_id,
oe.latency_ms, oe.message, oe.raw_upstream_error, oe.occurrence_count,
oe.latency_ms, oe.message, oe.raw_upstream_error,
host(mr.client_ip) as client_ip, mr.user_agent,
mr.reasoning_effort, mr.reasoning_preset, mr.request_kind,
mr.subagent_kind, mr.compact,
@@ -128,7 +127,6 @@ fn push_request_error_predicates(
range: ObservabilityRange,
filter: &OpsErrorFilter,
) {
statement.push(" and mr.recovered_at is null");
push_range(statement, "mr.completed_at", range);
for (column, value) in [
("mr.client_api_key_ref", &filter.client_api_key_ref),
@@ -171,7 +169,6 @@ fn push_ops_event_predicates(
range: ObservabilityRange,
filter: &OpsErrorFilter,
) {
statement.push(" and (mr.id is null or mr.recovered_at is null)");
push_range(statement, "oe.created_at", range);
for (column, value) in [
("mr.client_api_key_ref", &filter.client_api_key_ref),

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