mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a2922b410e | ||
|
|
d303efc6f7 | ||
|
|
dd17332c31 | ||
|
|
47d08786ad | ||
|
|
3434ad5e72 | ||
|
|
e91ec88a1f | ||
|
|
6f7a50b1b6 | ||
|
|
db0c294e83 | ||
|
|
d8181e8e91 | ||
|
|
af134cf488 | ||
|
|
d918149c2f | ||
|
|
aa6bf6497f | ||
|
|
074c07b5ca | ||
|
|
f70f46e87e | ||
|
|
c21633655e | ||
|
|
1bafbb3864 | ||
|
|
353dccedc4 | ||
|
|
697f1870a0 | ||
|
|
a404c2e09e | ||
|
|
a6cfd124b0 | ||
|
|
4441b577c3 | ||
|
|
a704fbc3f3 | ||
|
|
dab221cdb9 | ||
|
|
4f434ec525 | ||
|
|
2585d4cd7d | ||
|
|
93a34b0f56 | ||
|
|
aa3a07209a | ||
|
|
c2332bc8d0 | ||
|
|
8f242e2875 | ||
|
|
51e9981bee | ||
|
|
65fbbcb3ff | ||
|
|
bd086a7e3c | ||
|
|
3554ba13db | ||
|
|
bcccebc9fa | ||
|
|
d67b46a9e2 | ||
|
|
20c6d33cf4 | ||
|
|
cf0e84c548 |
@@ -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:
|
||||
|
||||
@@ -5,192 +5,161 @@
|
||||
|
||||
# Codex Proxy RS
|
||||
|
||||
**基于 Rust 的多 Provider 透明 AI 网关**
|
||||
集中管理 OpenAI、xAI 账号,为 Codex 和兼容客户端提供统一的访问地址与密钥。
|
||||
|
||||
[](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](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)
|
||||
|
||||
Generated
+50
-44
@@ -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",
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,10 +2,12 @@
|
||||
|
||||
mod error;
|
||||
mod model;
|
||||
pub mod scope;
|
||||
mod selection;
|
||||
mod store;
|
||||
|
||||
pub use error::CredentialError;
|
||||
pub use model::*;
|
||||
pub use selection::*;
|
||||
pub(crate) use selection::{SMART_SCORE_TOLERANCE, smart_score};
|
||||
pub use store::ProviderAccountStore;
|
||||
|
||||
@@ -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 => {
|
||||
@@ -584,7 +590,7 @@ impl AccountSelector {
|
||||
})
|
||||
}
|
||||
|
||||
fn scheduling_blocker(
|
||||
pub(crate) fn scheduling_blocker(
|
||||
&self,
|
||||
candidate: &AccountCandidate,
|
||||
context: &AccountSelectionContext,
|
||||
@@ -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 个百分点的单项负载/失败率差异,避免微小信号波动独占新会话。
|
||||
pub(crate) 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),
|
||||
)
|
||||
});
|
||||
pub(crate) 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,10 @@
|
||||
//! 有界、请求局部的诊断证据;不参与路由、重试或客户端协议决策。
|
||||
|
||||
mod capture;
|
||||
mod selection;
|
||||
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,73 @@
|
||||
//! 复用选号时的输入,保留有界的候选评分与阻断原因。
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::TraceContext;
|
||||
use crate::account::{
|
||||
AccountCandidate, AccountSelection, AccountSelectionContext, AccountSelector, RotationStrategy,
|
||||
SMART_SCORE_TOLERANCE, smart_score,
|
||||
};
|
||||
|
||||
// 单个 trace 事件最多 4 KiB;优先保留实际选中的账号。
|
||||
const MAX_CANDIDATES: usize = 6;
|
||||
|
||||
impl TraceContext {
|
||||
/// 记录实际选号输入与结果;只复用评分规则,不参与调度决策。
|
||||
pub fn account_selection(
|
||||
&self,
|
||||
candidates: &[AccountCandidate],
|
||||
context: &AccountSelectionContext,
|
||||
selection: Option<&AccountSelection<'_>>,
|
||||
) {
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
let selected = selection.as_ref().map(|s| s.candidate());
|
||||
let selected_id = selected.map(|c| c.account.id());
|
||||
let smart = context.policy.strategy() == RotationStrategy::Smart;
|
||||
let observations = selected
|
||||
.into_iter()
|
||||
.chain(
|
||||
candidates
|
||||
.iter()
|
||||
.filter(|c| Some(c.account.id()) != selected_id),
|
||||
)
|
||||
.take(MAX_CANDIDATES)
|
||||
.map(|candidate| {
|
||||
let signals = &candidate.signals;
|
||||
json!({
|
||||
"accountId": candidate.account.id().as_str(),
|
||||
"weight": candidate.account.weight().get(),
|
||||
"blocker": AccountSelector.scheduling_blocker(candidate, context)
|
||||
.map(|blocker| format!("{blocker:?}")),
|
||||
"inFlight": signals.in_flight,
|
||||
"concurrencyLimit": candidate.account
|
||||
.effective_concurrency(context.policy.max_concurrent_per_account()).get(),
|
||||
"quotaRemainingBasisPoints": signals.quota_remaining_rank,
|
||||
"quotaResetAtUnixMs": signals.quota_reset_at.and_then(|reset|
|
||||
reset.duration_since(std::time::UNIX_EPOCH).ok()
|
||||
).map(|duration| duration.as_millis()),
|
||||
"failureRateBasisPoints": signals.failure_rate_basis_points,
|
||||
"firstOutputLatencyMs": signals.first_output_latency_ms,
|
||||
"smartScore": smart.then(|| smart_score(
|
||||
candidate, context.policy.max_concurrent_per_account()
|
||||
)),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
self.record(
|
||||
"account.selection",
|
||||
json!({
|
||||
"strategy": context.policy.strategy().as_str(),
|
||||
"selectedAccountId": selected_id.map(|id| id.as_str()),
|
||||
"preferredAccountId": context.preferred_account.as_ref().map(|id| id.as_str()),
|
||||
"preferredResult": selection.as_ref().map(|s| format!("{:?}", s.preferred())),
|
||||
"roundRobinCursor": context.round_robin_cursor,
|
||||
"smartScoreTolerance": smart.then_some(SMART_SCORE_TOLERANCE),
|
||||
"candidateCount": candidates.len(),
|
||||
"omittedCandidates": candidates.len().saturating_sub(observations.len()),
|
||||
"candidates": observations,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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,345 @@
|
||||
//! 首部与最近事件同时保留,流式 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 {
|
||||
/// 是否已启用请求诊断,供调用方避免构造不会保存的事实。
|
||||
#[must_use]
|
||||
pub const fn is_enabled(&self) -> bool {
|
||||
self.state.is_some()
|
||||
}
|
||||
|
||||
/// 一个模型请求只创建一次,所有 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.selection"
|
||||
| "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",
|
||||
),
|
||||
|
||||
@@ -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 边界。
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
mod selection;
|
||||
mod selection_diagnostics;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
@@ -100,6 +101,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 +419,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 +698,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 +744,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,92 @@
|
||||
use gateway_core::account::{AccountSelector, RotationStrategy};
|
||||
use gateway_core::diagnostics::TraceContext;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{candidate, candidate_with_concurrency, context};
|
||||
|
||||
#[test]
|
||||
fn selection_trace_should_distinguish_load_quota_and_affinity() {
|
||||
let mut candidates = [
|
||||
candidate_with_concurrency("acct_74", 1, 10),
|
||||
candidate_with_concurrency("acct_96", 0, 10),
|
||||
];
|
||||
candidates[0].signals.quota_remaining_rank = Some(2600);
|
||||
candidates[1].signals.quota_remaining_rank = Some(400);
|
||||
let mut context = context(RotationStrategy::Smart);
|
||||
let trace = TraceContext::new("req_smart_trace");
|
||||
let selection = AccountSelector.select(&candidates, &context);
|
||||
trace.account_selection(&candidates, &context, selection.as_ref());
|
||||
context.preferred_account = Some(candidates[1].account.id().clone());
|
||||
let selection = AccountSelector.select(&candidates, &context);
|
||||
trace.account_selection(&candidates, &context, selection.as_ref());
|
||||
|
||||
let snapshot = trace.snapshot().expect("trace");
|
||||
let first = &snapshot["events"][0]["data"];
|
||||
let second = &snapshot["events"][1]["data"];
|
||||
assert_eq!(first["selectedAccountId"], "acct_74");
|
||||
assert_eq!(first["preferredResult"], "NotRequested");
|
||||
let busy = &first["candidates"][0];
|
||||
assert_eq!(busy["inFlight"], 1);
|
||||
assert_eq!(busy["concurrencyLimit"], 10);
|
||||
assert_eq!(busy["quotaRemainingBasisPoints"], 2600);
|
||||
assert!(busy["blocker"].is_null());
|
||||
assert!(
|
||||
busy["smartScore"].as_f64().expect("score")
|
||||
> first["candidates"][1]["smartScore"]
|
||||
.as_f64()
|
||||
.expect("score")
|
||||
);
|
||||
assert_eq!(second["selectedAccountId"], "acct_96");
|
||||
assert_eq!(second["preferredResult"], "Hit");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selection_trace_should_bound_candidates_and_preserve_the_selected_account() {
|
||||
let mut candidates = (0..30)
|
||||
.map(|index| candidate(&format!("acct_{index}"), 0, None))
|
||||
.collect::<Vec<_>>();
|
||||
candidates.push(candidate("acct_best", 0, Some(10_000)));
|
||||
let trace = TraceContext::new("req_many_candidates");
|
||||
// 选号位于首部保留区之外,也应在长流的常规事件淘汰后保留。
|
||||
for _ in 0..10 {
|
||||
trace.record("test.event", json!({}));
|
||||
}
|
||||
let context = context(RotationStrategy::Smart);
|
||||
let selection = AccountSelector.select(&candidates, &context);
|
||||
trace.account_selection(&candidates, &context, selection.as_ref());
|
||||
for index in 0..400 {
|
||||
trace.record("test.event", json!({"index": index}));
|
||||
}
|
||||
|
||||
let snapshot = trace.snapshot().expect("trace");
|
||||
let event = snapshot["events"]
|
||||
.as_array()
|
||||
.expect("events")
|
||||
.iter()
|
||||
.find(|event| event["stage"] == "account.selection")
|
||||
.expect("selection survives event eviction");
|
||||
let data = &event["data"];
|
||||
let recorded = data["candidates"].as_array().expect("bounded candidates");
|
||||
assert!(recorded.len() < candidates.len());
|
||||
assert_eq!(recorded[0]["accountId"], "acct_best");
|
||||
assert!(recorded[1]["quotaRemainingBasisPoints"].is_null());
|
||||
assert_eq!(
|
||||
recorded.len() as u64 + data["omittedCandidates"].as_u64().expect("omitted"),
|
||||
candidates.len() as u64
|
||||
);
|
||||
assert!(data.to_string().len() < 4096);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selection_trace_should_explain_when_all_accounts_are_busy() {
|
||||
let candidates = [candidate_with_concurrency("acct_busy", 10, 10)];
|
||||
let trace = TraceContext::new("req_no_candidate");
|
||||
let context = context(RotationStrategy::Smart);
|
||||
let selection = AccountSelector.select(&candidates, &context);
|
||||
assert!(selection.is_none());
|
||||
trace.account_selection(&candidates, &context, selection.as_ref());
|
||||
let snapshot = trace.snapshot().expect("trace");
|
||||
let data = &snapshot["events"][0]["data"];
|
||||
assert!(data["selectedAccountId"].is_null());
|
||||
assert_eq!(data["candidates"][0]["blocker"], "ConcurrencyLimit");
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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!({
|
||||
|
||||
@@ -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(¤t, &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(¤t, &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(¤t, &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,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user