Compare commits

..
23 Commits
Author SHA1 Message Date
zyycnandCodex 40a2b70bf2 chore: bump version to 2.3.2
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:07:00 +08:00
zyycnandCodex 1e6ac77d3c fix(frontend): scope usage record refresh
Refresh only the request detail table with an independent current-time snapshot, while preserving latest-request-wins ordering and leaving analytics cards untouched.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:06:38 +08:00
zyycnandCodex 3340934944 fix(transport): preserve websocket openings after fallback
Continue pool-owned WebSocket handshakes after the foreground budget expires, retain late successes for reuse, and cancel them on account eviction or shutdown. Keep breaker accounting and transport metrics aligned with the foreground decision.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:06:26 +08:00
zyycnandCodex 37a368de32 fix(ci): avoid duplicate pull request checks
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 11:33:11 +08:00
zyycnandCodex 79cd8ac965 chore(deps): delay fresh npm updates
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 11:13:50 +08:00
zyycnandCodex 57c6b95596 fix: report total dashboard billing
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 10:28:07 +08:00
zyycnandCodex a4f1843876 fix: unify account failures and usage diagnostics
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 10:15:10 +08:00
zyycnandCodex 096e48621a chore: bump version to 2.3.1
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 03:09:52 +08:00
zyycnandCodex 36828b1a7e fix: harden websocket liveness and usage reporting
Enforce pump-owned Pong deadlines without adding request-path probes, and pause watchdogs during inbound backpressure.

Refresh usage snapshots against the current time, expose account identities, preserve full plan badges, and source dashboard totals from global request buckets.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 03:09:39 +08:00
zyycnandCodex 05ec1a4d38 chore: bump version to 2.3.0
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 01:42:15 +08:00
zyycnandCodex 9808e5939c chore(deps): update frontend and container toolchain
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 01:36:13 +08:00
zyycnandCodex a58911c127 feat(transport): implement low-latency dual transport
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 01:11:43 +08:00
zyycnandCodex 0b9e080f1b fix(frontend): correct scrollbar and table overflow
Hide native scrollbars behind the custom modal scrollbar and truncate overflowing account cells with accessible hover titles.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 23:15:19 +08:00
zyycnandCodex f218e31c72 refactor(dispatch): harden account failover boundaries
Invalidate unavailable accounts in memory before asynchronous persistence, evict account WebSockets, and classify explicit auth, quota, ban, and workspace failures for immediate rotation.

Separate account-scoped installation identity from client session semantics, preserve logical IDs across failover, clear opaque account state, remove forced HTTP/1.1 configuration, and document the lifecycle rules.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 23:15:05 +08:00
zyycnandCodex 83f2f55e4a refactor(dispatch): enforce single-owner failure decisions
Move Cyber classification and retry decisions behind its controller-owned typed result.

Share feature failure owner ordering across attempt and committed stream paths, with architecture guards and documentation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 18:35:06 +08:00
zyycnandCodex 0384e39289 refactor(dispatch): implement v1 lifecycle onion
Co-authored-by: Codex <noreply@openai.com>
2026-07-14 17:55:20 +08:00
zyycnandCodex 9036288d69 docs: define v1 lifecycle controller boundaries
Scope the onion lifecycle model to v1 proxy APIs and establish single-feature ownership across controllers, lifecycle, adapters, API, and storage.

Document the target dispatch directory and mechanical boundary checks for future controller migrations.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 11:01:51 +08:00
zyycnandCodex 77b18eec68 fix(dispatch): handle deactivated and cyber policy failures
Classify deactivated_workspace as banned while preserving ordinary 402 quota behavior.

Track cyber_policy per client session so the next request avoids failed accounts, cap rotation at three accounts, guard success cleanup with ABA-safe token CAS, and document the Request/Attempt/Stream lifecycle controller target.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 05:13:48 +08:00
zyycnandCodex 3f7d473bea fix(ci): allow database bind initialization
Prepare ephemeral CI bind roots as sticky hand-off directories so image-specific database users can initialize storage on user-namespaced runners.

Keep PostgreSQL data permissions image-managed and verify PostgreSQL and Redis persistence across container recreation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 01:27:14 +08:00
zyycnandCodex 99cf39c76d refactor: consolidate runtime configuration
Use deploy/config.yaml as the structured source for application settings and service credentials, with strict validation and runtime-relative paths.

Consolidate Compose assets under deploy, remove obsolete environment and smoke files, keep terminal logs readable while retaining JSON file output, and refresh the Codex Desktop fingerprint.

Co-authored-by: Codex <noreply@openai.com>
2026-07-14 00:46:57 +08:00
Zyy f4d0576e6b docs: update README 2026-07-13 19:09:22 +08:00
Zyy 2d9510d9e3 feat: improve usage insights and dashboard trends 2026-07-13 19:09:22 +08:00
Zyy 36d6e26291 fix(frontend): polish modal and responsive styles 2026-07-13 12:00:08 +08:00
220 changed files with 19424 additions and 8657 deletions
+2
View File
@@ -34,6 +34,8 @@ updates:
day: monday
time: '09:20'
timezone: Asia/Shanghai
cooldown:
default-days: 1
groups:
minor-and-patch:
update-types:
+52 -34
View File
@@ -19,21 +19,24 @@ jobs:
env:
COMPOSE_PROJECT_NAME: cpr-ci-${{ github.run_id }}-${{ github.run_attempt }}
CPR_IMAGE: codex-proxy-rs:ci-${{ inputs.ref }}
CPR_ADMIN_DEFAULT_PASSWORD: 'admin @:/?#%2F ci password'
CPR_POSTGRES_PASSWORD: 'pg @:/?#%2F ci'
CPR_REDIS_PASSWORD: 'redis @:/?#%2F ci'
CPR_SMOKE_PORT: '18080'
TRIVY_QUIET: 'true'
TRIVY_SKIP_VERSION_CHECK: 'true'
steps:
- name: Define isolated runtime paths
- name: Define isolated runtime values
shell: bash
run: |
set -euo pipefail
admin_password="$(openssl rand -hex 24)"
postgres_password="$(openssl rand -hex 24)"
redis_password="$(openssl rand -hex 24)"
echo "::add-mask::$admin_password"
echo "::add-mask::$postgres_password"
echo "::add-mask::$redis_password"
{
echo "CPR_CONFIG_FILE=$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT/config.yaml"
echo "CPR_DATA_DIR=$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT/data"
echo "CPR_LOG_DIR=$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT/logs"
echo "DIAGNOSTICS_DIR=$RUNNER_TEMP/codex-proxy-rs-diagnostics"
echo "CI_ADMIN_PASSWORD=$admin_password"
echo "CI_POSTGRES_PASSWORD=$postgres_password"
echo "CI_REDIS_PASSWORD=$redis_password"
} >> "$GITHUB_ENV"
- name: Checkout
@@ -94,19 +97,35 @@ jobs:
shell: bash
run: |
set -euo pipefail
runtime_dir="$(dirname "$CPR_CONFIG_FILE")"
mkdir -p "$runtime_dir" "$DIAGNOSTICS_DIR"
sudo install -d -m 0770 -o "$(id -u)" -g 10001 "$CPR_DATA_DIR" "$CPR_LOG_DIR"
cp deploy/config.example.yaml "$CPR_CONFIG_FILE"
sudo chown "$(id -u):10001" "$CPR_CONFIG_FILE"
chmod 0640 "$CPR_CONFIG_FILE"
umask 0077
# Bind roots hand off ownership to the image-specific database users.
# Sticky directories keep that initialization compatible with user-namespaced runners.
install -d -m 1777 .runtime/postgres .runtime/redis
mkdir -p "$DIAGNOSTICS_DIR"
sudo install -d -m 0770 -o "$(id -u)" -g 10001 \
.runtime/data .runtime/logs
ruby <<'RUBY'
path = "deploy/config.example.yaml"
content = File.read(path)
replacements = {
"password: &postgres_password ''" => "password: &postgres_password '#{ENV.fetch("CI_POSTGRES_PASSWORD")}'",
"password: &redis_password ''" => "password: &redis_password '#{ENV.fetch("CI_REDIS_PASSWORD")}'",
"default_password: ''" => "default_password: '#{ENV.fetch("CI_ADMIN_PASSWORD")}'",
}
replacements.each do |placeholder, configured|
raise "missing config placeholder: #{placeholder}" unless content.sub!(placeholder, configured)
end
File.write("deploy/config.yaml", content)
RUBY
sudo chown "$(id -u):10001" deploy/config.yaml
chmod 0640 deploy/config.yaml
- name: Validate Docker Compose
shell: bash
run: |
set -euo pipefail
docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml config --quiet
-f deploy/compose.yaml config --quiet
- name: Configure Redis host memory overcommit
shell: bash
@@ -120,25 +139,25 @@ jobs:
run: |
set -euo pipefail
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
-f deploy/compose.yaml)
"${compose[@]}" up -d --no-build --wait postgres redis
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
psql -At -U codex_proxy -d codex_proxy -c "SHOW TimeZone;" | tr -d '\r')" = Asia/Shanghai
"${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
"${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
psql -v ON_ERROR_STOP=1 -U codex_proxy -d codex_proxy -c \
"CREATE TABLE cpr_ci_persistence_probe (value text PRIMARY KEY); INSERT INTO cpr_ci_persistence_probe VALUES ('persisted');"
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CPR_REDIS_PASSWORD" redis \
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CI_REDIS_PASSWORD" redis \
redis-cli SET cpr:ci:persistence-probe persisted
"${compose[@]}" up -d --no-build --wait --force-recreate postgres redis
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
psql -At -U codex_proxy -d codex_proxy -c \
"SELECT value FROM cpr_ci_persistence_probe;")" = persisted
test "$("${compose[@]}" exec -T -e REDISCLI_AUTH="$CPR_REDIS_PASSWORD" redis \
test "$("${compose[@]}" exec -T -e REDISCLI_AUTH="$CI_REDIS_PASSWORD" redis \
redis-cli --raw GET cpr:ci:persistence-probe | tr -d '\r')" = persisted
"${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
"${compose[@]}" exec -T -e PGPASSWORD="$CI_POSTGRES_PASSWORD" postgres \
psql -v ON_ERROR_STOP=1 -U codex_proxy -d codex_proxy -c \
"DROP TABLE cpr_ci_persistence_probe;"
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CPR_REDIS_PASSWORD" redis \
"${compose[@]}" exec -T -e REDISCLI_AUTH="$CI_REDIS_PASSWORD" redis \
redis-cli DEL cpr:ci:persistence-probe
- name: Start and verify application
@@ -146,24 +165,24 @@ jobs:
run: |
set -euo pipefail
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
-f deploy/compose.yaml)
"${compose[@]}" up -d --no-build --wait codex-proxy-rs
status="$(curl --silent --show-error --output /dev/null --write-out '%{http_code}' \
"http://127.0.0.1:${CPR_SMOKE_PORT}/healthz")"
"http://127.0.0.1:8080/healthz")"
test "$status" = 204
test "$("${compose[@]}" exec -T codex-proxy-rs id -u | tr -d '\r')" = 10001
test "$("${compose[@]}" exec -T codex-proxy-rs id -g | tr -d '\r')" = 10001
test "$("${compose[@]}" exec -T codex-proxy-rs date +%z | tr -d '\r')" = +0800
test "$("${compose[@]}" exec -T codex-proxy-rs cat /etc/timezone | tr -d '\r')" = Asia/Shanghai
"${compose[@]}" exec -T codex-proxy-rs sh -c \
'touch /app/data/.ci-write /app/logs/.ci-write && rm /app/data/.ci-write /app/logs/.ci-write'
'touch /app/.runtime/data/.ci-write /app/.runtime/logs/.ci-write && rm /app/.runtime/data/.ci-write /app/.runtime/logs/.ci-write'
- name: Verify clean logs and graceful shutdown
shell: bash
run: |
set -euo pipefail
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
-f deploy/compose.yaml)
"${compose[@]}" stop --timeout 30 codex-proxy-rs
container_id="$("${compose[@]}" ps -a -q codex-proxy-rs)"
test -n "$container_id"
@@ -195,7 +214,7 @@ jobs:
mkdir -p "$DIAGNOSTICS_DIR"
raw_dir="$(mktemp -d)"
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
-f deploy/compose.yaml)
"${compose[@]}" config > "$raw_dir/compose-config.yaml" 2>&1
"${compose[@]}" ps -a > "$raw_dir/compose-ps.txt" 2>&1
"${compose[@]}" logs --no-color --timestamps > "$raw_dir/compose.log" 2>&1
@@ -207,9 +226,9 @@ jobs:
raw_dir = ENV.fetch("RAW_DIR")
output_dir = ENV.fetch("OUTPUT_DIR")
secrets = [
ENV.fetch("CPR_ADMIN_DEFAULT_PASSWORD"),
ENV.fetch("CPR_POSTGRES_PASSWORD"),
ENV.fetch("CPR_REDIS_PASSWORD"),
ENV.fetch("CI_ADMIN_PASSWORD"),
ENV.fetch("CI_POSTGRES_PASSWORD"),
ENV.fetch("CI_REDIS_PASSWORD"),
]
encoded_secrets = secrets.map do |secret|
secret.encode("UTF-8").bytes.map do |byte|
@@ -243,6 +262,5 @@ jobs:
shell: bash
run: |
docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml \
down --volumes --remove-orphans || true
sudo rm -rf "$GITHUB_WORKSPACE/.runtime/ci-$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT"
-f deploy/compose.yaml down --remove-orphans || true
sudo rm -rf "$GITHUB_WORKSPACE/.runtime" "$GITHUB_WORKSPACE/deploy/config.yaml"
+3 -1
View File
@@ -3,8 +3,10 @@ name: CI
on:
push:
branches:
- '**'
- main
pull_request:
branches:
- main
permissions:
contents: read
+3 -1
View File
@@ -3,8 +3,10 @@ name: Security Scan
on:
push:
branches:
- '**'
- main
pull_request:
branches:
- main
schedule:
- cron: '0 3 * * 1'
+1 -7
View File
@@ -19,13 +19,7 @@
*.log
# Local configuration and secrets
.env
.env.*
!.env.example
!**/.env.example
/.envrc
/config.yaml
/backend/config.yaml
/deploy/config.yaml
# Editors and operating systems
/.idea/
+146 -157
View File
@@ -1,122 +1,119 @@
<div align="center">
<h1>Codex Proxy RS</h1>
<p>OpenAI Responses 兼容的 Codex 多账号网关</p>
<p>
<a href="https://www.rust-lang.org/"><img alt="Rust 1.97" src="https://img.shields.io/badge/Rust-1.97-000?style=flat-square&amp;logo=rust" /></a>
<a href="https://github.com/tokio-rs/axum"><img alt="Axum 0.8" src="https://img.shields.io/badge/Axum-0.8-2f7d95?style=flat-square" /></a>
<a href="https://vuejs.org/"><img alt="Vue 3.5" src="https://img.shields.io/badge/Vue-3.5-42b883?style=flat-square&amp;logo=vuedotjs&amp;logoColor=white" /></a>
<a href="https://vite.dev/"><img alt="Vite 8" src="https://img.shields.io/badge/Vite-8-646cff?style=flat-square&amp;logo=vite&amp;logoColor=white" /></a>
<a href="https://www.postgresql.org/"><img alt="PostgreSQL 18" src="https://img.shields.io/badge/PostgreSQL-18-4169e1?style=flat-square&amp;logo=postgresql&amp;logoColor=white" /></a>
<a href="https://redis.io/"><img alt="Redis 8" src="https://img.shields.io/badge/Redis-8-ff4438?style=flat-square&amp;logo=redis&amp;logoColor=white" /></a>
<a href="https://opensource.org/license/mit"><img alt="MIT License" src="https://img.shields.io/badge/License-MIT-blue?style=flat-square" /></a>
</p>
<p>
<a href="#快速部署">快速部署</a>
·
<a href="#客户端接入">客户端接入</a>
·
<a href="docs/architecture.md">架构说明</a>
</p>
<img src="frontend/public/favicon.svg" alt="Codex Proxy RS" width="88" height="88" />
# Codex Proxy RS
**基于 Rust 的高性能 Codex 多账号网关代理**
[![CI](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
[![Release](https://img.shields.io/github/v/release/zyycn/codex-proxy-rs?display_name=tag&sort=semver)](https://github.com/zyycn/codex-proxy-rs/releases)
[![Rust 1.97](https://img.shields.io/badge/Rust-1.97-000?logo=rust&logoColor=white)](https://www.rust-lang.org/)
[![Vue 3.5](https://img.shields.io/badge/Vue-3.5-42b883?logo=vuedotjs&logoColor=white)](https://vuejs.org/)
[![Vite 8](https://img.shields.io/badge/Vite-8-646cff?logo=vite&logoColor=white)](https://vite.dev/)
[![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
[![MIT](https://img.shields.io/badge/License-MIT-blue.svg)](#开源许可)
[快速开始](#快速开始) · [首次使用](#首次使用) · [客户端接入](#客户端接入) · [常见问题](#常见问题) · [部署文档](deploy/README.md)
</div>
Codex Proxy RS 为多个 Codex 账号提供统一的 Responses API。客户端使用同一个服务地址和 API Key;账号调度、并发限制、额度检查、故障换号和会话续接都在服务端完成。
> [!NOTE]
> 当前兼容 OpenAI Responses API,不提供 `/v1/chat/completions`。
## 功能
- 兼容 `/v1/responses`、SSE、官方 Responses WebSocket 和模型查询。
- 客户端与上游均可使用 HTTP/SSE 或 WebSocket。
- 提供 `smart`、额度重置优先、轮询和粘性四种账号调度策略。
- 单个账号发生额度、认证、模型或传输故障时,继续使用账号池中的其他可用账号。
- 支持跨账号会话续接,并隔离不同账号的会话状态。
- 提供账号、API Key、模型映射、运行参数、用量、错误记录和系统更新管理界面。
- 使用 PostgreSQL 保存权威业务数据,使用 Redis 保存管理会话、租约、模型快照和响应归属。
- 内置日志轮转、遥测保留期清理和 PostgreSQL/Redis 健康检查。
- 支持 Responses HTTP、SSE 和官方 WebSocket 请求。
- 提供智能调度、额度重置优先、轮询和粘滞四种账号策略。
- 某个账号遇到额度、认证、模型或传输错误时,继续尝试其他账号。
- 在网页中管理账号、客户端 Key、模型别名和运行参数。
- 按请求、模型、账号或客户端 Key 查看 Token、费用、缓存、TTFT、延迟和错误。
- 支持账号 OAuth、Token 导入以及 CPR、Sub2API、CLIProxyAPI JSON 导入。
- 可直接生成 Codex CLI 配置,或通过 deeplink 导入 CCSwitch。
除账号调度和会话恢复外,请求与响应尽量保持原有语义。
## 快速开始
## 快速部署
Docker Compose 是最简单的启动方式。开始前请安装 Docker Engine 和 Docker Compose Plugin。
需要 Docker Engine 和 Docker Compose Plugin。下面的 Compose 默认使用:
- 应用镜像:`ghcr.io/zyycn/codex-proxy-rs:latest`
- PostgreSQL 18
- Redis 8
- 管理端地址:`http://127.0.0.1:8080`
### 1. 获取项目
### 1. 准备配置
```bash
git clone https://github.com/zyycn/codex-proxy-rs.git
cd codex-proxy-rs
```
### 2. 创建配置
```bash
mkdir -p .runtime/data .runtime/logs
cp deploy/config.example.yaml .runtime/config.yaml
cp deploy/.env.example deploy/.env
chmod 0600 deploy/.env
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
```
编辑 `deploy/.env`:
```dotenv
CPR_ADMIN_DEFAULT_PASSWORD='<管理员密码,至少 12 个字符>'
CPR_POSTGRES_PASSWORD='<PostgreSQL 密码>'
CPR_REDIS_PASSWORD='<Redis 密码>'
```
`deploy/.env` 已被 Git 忽略。应用会安全地把 PostgreSQL 和 Redis 密码写入运行时连接 URL,不需要把真实密码填入 `.runtime/config.yaml`。
Linux 宿主机需要保证容器用户 `10001:10001` 可以读写数据和日志目录:
分别执行三次以下命令生成密码:
```bash
sudo chown -R "$(id -u):10001" .runtime/data .runtime/logs
chmod 0770 .runtime/data .runtime/logs
sudo chown "$(id -u):10001" .runtime/config.yaml
chmod 0640 .runtime/config.yaml
openssl rand -hex 24
```
### 3. 启动
把结果填入 `deploy/config.yaml` 的 `database.password`、`redis.password` 和
`admin.default_password`。PostgreSQL 与 Redis 密码必须是 48 位十六进制字符;管理员密码
至少 12 位且不能包含 `$`。
> [!IMPORTANT]
> 应用容器以非 root 用户 `10001:10001` 运行。Linux Bind Mount 会保留宿主机权限,因此需要允许容器用户组读取配置并写入数据和日志目录:
>
> ```bash
> sudo chown -R "$(id -u):10001" .runtime/data .runtime/logs
> chmod 0770 .runtime/data .runtime/logs
> ```
>
> 配置文件通过 Compose `configs` 只读挂载;普通 Compose 会保留宿主机权限,因此由
> 当前用户持有、仅向容器组开放读取权限。
> macOS 和 Windows 的 Docker Desktop 通常也不需要调整数据目录权限。
### 2. 启动服务
```bash
docker compose --env-file deploy/.env -f deploy/docker-compose.yml config
docker compose --env-file deploy/.env -f deploy/docker-compose.yml pull
docker compose --env-file deploy/.env -f deploy/docker-compose.yml up -d --no-build
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
```
检查服务:
启动后检查:
```bash
docker compose --env-file deploy/.env -f deploy/docker-compose.yml ps
docker compose -f deploy/compose.yaml ps
curl -i http://127.0.0.1:8080/healthz
```
`/healthz` 正常时返回 `204 No Content`。随后打开 <http://127.0.0.1:8080>:
`/healthz` 返回 `204 No Content` 即表示服务正常。管理端地址为 <http://127.0.0.1:8080>。
- 用户名:`admin@cpr.local`
- 密码:`deploy/.env` 中的 `CPR_ADMIN_DEFAULT_PASSWORD`
## 首次使用
首次登录后导入账号,在“API 密钥”页面创建客户端 Key。
1. 使用 `admin@cpr.local` 和 `x-cpr.admin.default_password` 中设置的密码登录。
2. 打开“账号”,用 OAuth、Access/Refresh Token 或 JSON 文件添加账号。
3. 点一次连接测试,能看到模型和额度即可。
4. 打开“API 密钥”,创建一个 `sk_...` 客户端 Key。
5. 点击“使用密钥”复制 Codex CLI 配置,或一键导入 CCSwitch。
### 本地构建镜像
需要验证当前源码时,不拉取 GHCR 镜像,直接构建:
```bash
docker compose --env-file deploy/.env -f deploy/docker-compose.yml build codex-proxy-rs
docker compose --env-file deploy/.env -f deploy/docker-compose.yml up -d
```
> [!TIP]
> `admin.default_password` 只用于首次创建管理员。已有管理员后修改该字段不会重置登录密码。
## 客户端接入
OpenAI 兼容客户端填写:
| 配置 | 值 |
| -------- | ---------------------------------------- |
| Base URL | `http://127.0.0.1:8080/v1` |
| API Key | 管理端创建的 `sk_...` 客户端 Key |
| 鉴权 | `Authorization: Bearer <client-api-key>` |
- Base URL:`http://127.0.0.1:8080/v1`
- API Key:管理端创建的 `sk_...` 客户端 Key
先查询当前可用模型:
Responses 请求示例:
```bash
curl http://127.0.0.1:8080/v1/models \
-H 'Authorization: Bearer <client-api-key>'
```
发送一个非流式 Responses 请求:
```bash
curl http://127.0.0.1:8080/v1/responses \
@@ -129,98 +126,90 @@ curl http://127.0.0.1:8080/v1/responses \
}'
```
可用接口:
<details>
<summary>查看全部客户端 API 路由</summary>
| 路由 | 说明 |
| --- | --- |
| `POST /v1/responses` | Responses 非流式或 SSE 流式请求 |
| `GET /v1/responses` | 官方 Responses WebSocket |
| `POST /v1/responses/review` | Review 请求 |
| `GET /v1/models` | 模型列表 |
| `GET /v1/models/catalog` | 模型目录 |
| `GET /v1/models/{model_id}` | 模型详情 |
| `GET /v1/models/{model_id}/info` | 模型运行信息 |
| 路由 | 说明 |
| -------------------------------- | ------------------------------- |
| `POST /v1/responses` | Responses 非流式或 SSE 流式请求 |
| `GET /v1/responses` | 官方 Responses WebSocket |
| `POST /v1/responses/review` | Review 请求 |
| `GET /v1/models` | OpenAI 格式的模型列表 |
| `GET /v1/models/catalog` | 完整模型目录 |
| `GET /v1/models/{model_id}` | 模型详情 |
| `GET /v1/models/{model_id}/info` | 模型运行信息 |
Compose 默认只监听宿主机 `127.0.0.1`。需要从其他设备访问时,应在前面配置 HTTPS 反向代理,不要直接暴露 PostgreSQL 和 Redis 端口。
所有 `/v1/*` 路由都需要客户端 API Key。
</details>
## 管理端
管理端提供以下页面:
| 页面 | 内容 |
| -------- | ---------------------------------------------------------------------- |
| 概览 | 查看服务状态、请求趋势、Token、账号池容量、请求健康和最近调用 |
| 账号 | 添加、导入、导出、测试、启停账号,以及刷新 Token 和额度 |
| API 密钥 | 创建、复制、启停或删除 Key,生成 Codex CLI 配置并导入 CCSwitch |
| 使用统计 | 按时间、模型、账号和客户端 Key 查看用量、费用、缓存、延迟和错误 |
| 系统设置 | 配置模型别名、调度策略、单账号并发、请求间隔、刷新参数和管理员 API Key |
| 系统更新 | 查看版本和发布说明;Release 版本可以在线更新并重启 |
- Dashboard:请求趋势、服务状态、账号概览和最近请求。
- 账号:导入、导出、OAuth、连接测试、quota 和 token 刷新。
- API 密钥:创建、停用、删除和客户端接入信息。
- 用量:请求记录、token、延迟、模型分布和运维错误。
- 设置:模型别名、账号调度策略、单账号并发、请求间隔和刷新参数。
- 系统:版本检查、更新、重启和回滚。
## 常用操作
管理端登录会话保存在 Redis,并按 `admin.session_ttl_minutes` 自动过期。
## 配置
启动配置位于 `.runtime/config.yaml`,模板见 [deploy/config.example.yaml](deploy/config.example.yaml)。
| 配置段 | 用途 |
| --- | --- |
| `server` | HTTP 监听地址和端口 |
| `api` | Codex 上游地址 |
| `database`、`redis` | PostgreSQL 与 Redis 连接 |
| `quota` | quota 刷新周期和耗尽账号过滤 |
| `tls` | 上游 HTTP 协议偏好 |
| `ws_pool` | WebSocket 连接池与首事件超时 |
| `fingerprint` | Codex Desktop 请求指纹 |
| `admin` | 默认管理员和登录会话 TTL |
| `logging` | 日志级别、目录、大小和保留期 |
| `telemetry` | 请求事实记录开关 |
模型别名、调度策略、单账号并发、请求间隔和刷新参数由管理端保存到 PostgreSQL,可在运行中更新。请求事实、运维错误和聚合时间桶由后台任务按各自保留期自动清理。
## 数据与升级
| 数据 | 默认位置 |
| --- | --- |
| 账号、API Key、设置、用量和错误记录 | Compose 命名卷 `postgres-data` |
| 管理会话、租约、模型快照和响应归属 | Compose 命名卷 `redis-data` |
| 身份派生密钥和更新状态 | `.runtime/data` |
| 文件日志 | `.runtime/logs` |
| 启动配置 | `.runtime/config.yaml` |
| Docker 密钥 | `deploy/.env` |
普通升级只替换应用镜像:
### 升级
```bash
docker compose --env-file deploy/.env -f deploy/docker-compose.yml pull codex-proxy-rs
docker compose --env-file deploy/.env -f deploy/docker-compose.yml up -d --no-build
docker compose -f deploy/compose.yaml pull codex-proxy-rs
docker compose -f deploy/compose.yaml up -d --no-build
```
更新镜像不会删除数据库卷,但仍应在升级前备份 PostgreSQL 和 `.runtime/data`。
不要使用 `docker compose down -v` 做普通升级,该命令会删除 PostgreSQL 与 Redis 命名卷。
`.runtime/data/identity_hmac_secret` 必须长期保留。丢失该文件会改变全部账号的 installation ID 和账号作用域身份。
## 查看状态
### 查看日志
```bash
# 容器状态
docker compose --env-file deploy/.env -f deploy/docker-compose.yml ps
# 应用日志
docker compose --env-file deploy/.env -f deploy/docker-compose.yml logs -f codex-proxy-rs
# PostgreSQL 和 Redis 日志
docker compose --env-file deploy/.env -f deploy/docker-compose.yml logs postgres redis
# 展开并检查最终 Compose 配置
docker compose --env-file deploy/.env -f deploy/docker-compose.yml config
docker compose -f deploy/compose.yaml logs -f codex-proxy-rs
```
`/healthz` 返回 `503` 时,先检查 PostgreSQL 和 Redis 的健康状态。容器反复重启时,再检查 `.runtime/config.yaml`、`deploy/.env` 和 `.runtime/data`、`.runtime/logs` 的权限。
### 从当前源码构建
## 架构
```bash
docker compose -f deploy/compose.yaml build codex-proxy-rs
docker compose -f deploy/compose.yaml up -d
```
模块边界、请求链路、存储归属和后台任务见 [架构说明](docs/architecture.md)。
> [!WARNING]
> PostgreSQL、Redis、应用数据和文件日志都保存在仓库根目录 `.runtime/`。普通
> `docker compose down` 不会删除这些绑定目录;删除 `.runtime/` 才会清除数据。升级前请备份
> 整个目录,数据库备份中包含客户端凭据,应按敏感数据保管。
## License
## 常见问题
MIT
### `/healthz` 返回 `503`
检查 PostgreSQL 和 Redis 容器是否健康:
```bash
docker compose -f deploy/compose.yaml ps
docker compose -f deploy/compose.yaml logs postgres redis
```
### 容器反复重启
检查 `deploy/config.yaml` 的三个密码是否填写完整,以及 `.runtime/data`、`.runtime/logs`
的权限是否正确。
### 需要从其他设备访问
Compose 默认仅绑定 `127.0.0.1`。请在服务前配置 HTTPS 反向代理,并正确转发 WebSocket 升级和真实客户端 IP 请求头。不要直接暴露 PostgreSQL 或 Redis。
## 相关链接
- [发布版本](https://github.com/zyycn/codex-proxy-rs/releases)
- [容器镜像](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
- [详细部署说明](deploy/README.md)
- [架构说明](docs/architecture.md)
Docker 镜像支持 Linux amd64/arm64。Release 压缩包另有 macOS arm64 和 Windows amd64 版本。
## 开源许可
Codex Proxy RS 基于 [MIT License](https://opensource.org/license/mit) 开源。
+1 -49
View File
@@ -359,7 +359,6 @@ dependencies = [
"bytes",
"chrono",
"config",
"dirs",
"flate2",
"futures",
"hex",
@@ -368,7 +367,6 @@ dependencies = [
"hyper-util",
"indexmap",
"ipnet",
"percent-encoding",
"rand 0.10.2",
"rand_core 0.6.4",
"redis",
@@ -653,27 +651,6 @@ dependencies = [
"ctutils",
]
[[package]]
name = "dirs"
version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e"
dependencies = [
"dirs-sys",
]
[[package]]
name = "dirs-sys"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab"
dependencies = [
"libc",
"option-ext",
"redox_users",
"windows-sys 0.61.2",
]
[[package]]
name = "displaydoc"
version = "0.2.6"
@@ -1388,15 +1365,6 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libredox"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652"
dependencies = [
"libc",
]
[[package]]
name = "libsqlite3-sys"
version = "0.37.0"
@@ -1560,12 +1528,6 @@ version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe"
[[package]]
name = "option-ext"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
[[package]]
name = "parking"
version = "2.2.1"
@@ -1856,17 +1818,6 @@ dependencies = [
"bitflags",
]
[[package]]
name = "redox_users"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac"
dependencies = [
"getrandom 0.2.17",
"libredox",
"thiserror",
]
[[package]]
name = "regex"
version = "1.13.0"
@@ -2053,6 +2004,7 @@ version = "0.10.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a"
dependencies = [
"serde",
"zeroize",
]
+1 -3
View File
@@ -13,7 +13,6 @@ base64 = "0.22.1"
bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["serde", "clock"] }
config = { version = "0.15.23", default-features = false, features = ["yaml"] }
dirs = "6.0.0"
flate2 = "1.1.9"
futures = "0.3.32"
hex = "0.4.3"
@@ -22,7 +21,6 @@ hyper = { version = "1.10.1", features = ["http1", "server"] }
hyper-util = { version = "0.1.20", features = ["tokio"] }
indexmap = { version = "2.0", features = ["serde"] }
ipnet = "2.12.0"
percent-encoding = "2.3.2"
rand = "0.10.1"
# Enables password_hash::rand_core::OsRng for argon2 0.5.
rand_core = { version = "0.6.4", features = ["getrandom"] }
@@ -32,7 +30,7 @@ reqwest = { version = "=0.12.28", default-features = false, features = ["json",
rustls = { version = "=0.23.36", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8.4"
rustls-pki-types = "1.14.1"
secrecy = "0.10.3"
secrecy = { version = "0.10.3", features = ["serde"] }
semver = "1.0.27"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["preserve_order"] }
+1 -1
View File
@@ -114,7 +114,7 @@ pub(crate) async fn logout(
.delete_session(&session_id)
.await
{
tracing::warn!(error = %error, "failed to revoke admin session during logout");
tracing::warn!(error = %error, "Failed to revoke admin session during logout");
}
let mut response = AdminResponse::new(
+5 -5
View File
@@ -205,12 +205,12 @@ pub(crate) async fn dashboard_summary(
let capacity = state.services.account_pool.capacity_summary_now().await;
let now = Utc::now();
let today_filter = today_usage_record_filter(now);
let lifetime_usage = state
let retained_usage = state
.services
.usage
.summary()
.retained_summary()
.await
.map_err(|error| dashboard_data_error("account usage summary", &error))?;
.map_err(|error| dashboard_data_error("retained usage summary", &error))?;
let account_usage_records = state
.services
.usage_records
@@ -256,7 +256,7 @@ pub(crate) async fn dashboard_summary(
cards: dashboard_cards(
dashboard_account_counts(&accounts),
&time_buckets,
&lifetime_usage,
&retained_usage,
),
trend,
health_timeline: dashboard_health_timeline_data(&time_buckets),
@@ -419,7 +419,7 @@ async fn account_quota_used_percent_by_id(
}
fn dashboard_data_error(source: &'static str, error: &impl std::fmt::Display) -> AdminError {
tracing::error!(source, error = %error, "failed to load dashboard data");
tracing::error!(source, error = %error, "Failed to load dashboard data");
AdminError::internal("Failed to load dashboard data")
}
+10 -1
View File
@@ -30,7 +30,8 @@ use super::{
},
usage_routes::{
clear_usage_records, usage_record_detail, usage_records,
usage_records_endpoint_distribution, usage_records_latency_trend,
usage_records_endpoint_distribution, usage_records_insights_diagnostics,
usage_records_insights_overview, usage_records_latency_trend,
usage_records_model_distribution, usage_records_summary, usage_records_token_trend,
},
};
@@ -101,6 +102,14 @@ pub fn router() -> Router<AppState> {
"/api/admin/usage/records/insights/latency-trend",
get(usage_records_latency_trend),
)
.route(
"/api/admin/usage/records/insights/overview",
get(usage_records_insights_overview),
)
.route(
"/api/admin/usage/records/insights/diagnostics",
get(usage_records_insights_diagnostics),
)
.route("/api/admin/usage/records/delete", post(clear_usage_records))
.route("/api/admin/usage/records/detail", get(usage_record_detail))
.route("/api/admin/ops/errors", get(ops_errors))
+72 -4
View File
@@ -1,4 +1,4 @@
//! 使用记录 HTTP 处理器。
//! 使用统计与请求记录 HTTP 处理器。
use std::collections::HashMap;
@@ -26,6 +26,7 @@ use crate::{
time::china_datetime,
},
telemetry::{
usage::insights::{UsageDiagnosticsDimension, default_time_range},
usage::query::{
UsageQueryError, UsageQueryFilter, UsageRecordBreakdown, UsageRecordModelSource,
UsageRecordSummary, UsageRecordTrendPoint, usage_record_billing,
@@ -64,6 +65,21 @@ pub(crate) struct UsageRecordDistributionQuery {
source: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UsageInsightsQuery {
start_time: Option<String>,
end_time: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UsageDiagnosticsQuery {
#[serde(flatten)]
range: UsageInsightsQuery,
dimension: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UsageRecordDetailQuery {
@@ -304,6 +320,39 @@ pub(crate) async fn usage_records_latency_trend(
usage_records_trend_response(state, query).await
}
/// `GET /api/admin/usage/records/insights/overview`
pub(crate) async fn usage_records_insights_overview(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<UsageInsightsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let (start, end) = insights_time_range(query)?;
state
.services
.usage_records
.insights_overview(start, end)
.await
.map(|insights| AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(insights)))
.map_err(|error| log_error(&error))
}
/// `GET /api/admin/usage/records/insights/diagnostics`
pub(crate) async fn usage_records_insights_diagnostics(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<UsageDiagnosticsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let dimension = diagnostics_dimension(query.dimension)?;
let (start, end) = insights_time_range(query.range)?;
state
.services
.usage_records
.insights_diagnostics(start, end, dimension)
.await
.map(|insights| AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(insights)))
.map_err(|error| log_error(&error))
}
async fn usage_records_trend_response(
state: AppState,
query: UsageRecordsQuery,
@@ -405,6 +454,25 @@ fn optional_datetime(
.map_err(|_| AdminError::invalid_time_range("Invalid time range"))
}
fn insights_time_range(
query: UsageInsightsQuery,
) -> Result<(chrono::DateTime<chrono::Utc>, chrono::DateTime<chrono::Utc>), AdminError> {
let end = optional_datetime(query.end_time)?.unwrap_or_else(chrono::Utc::now);
let start = optional_datetime(query.start_time)?.unwrap_or_else(|| default_time_range(end).0);
if start >= end {
return Err(AdminError::invalid_time_range(
"Start time must be earlier than end time",
));
}
Ok((start, end))
}
fn diagnostics_dimension(value: Option<String>) -> Result<UsageDiagnosticsDimension, AdminError> {
let value = non_empty(value).unwrap_or_else(|| "model".to_string());
UsageDiagnosticsDimension::parse(&value)
.ok_or_else(|| AdminError::bad_request("Invalid diagnostics dimension"))
}
fn log_error(error: &UsageQueryError) -> AdminError {
match error {
UsageQueryError::List | UsageQueryError::Get | UsageQueryError::Clear => {
@@ -616,9 +684,9 @@ fn usage_record_breakdown_items(items: Vec<UsageRecordBreakdown>) -> Vec<UsageRe
impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
fn from(point: UsageRecordTrendPoint) -> Self {
let prompt_tokens = point.input_tokens.saturating_add(point.cached_tokens);
let cache_hit_rate_value = if prompt_tokens > 0 {
(point.cached_tokens as f64 / prompt_tokens as f64 * 100.0).round()
let cache_hit_rate_value = if point.input_tokens > 0 {
(point.cached_tokens.min(point.input_tokens) as f64 / point.input_tokens as f64 * 100.0)
.round()
} else {
0.0
};
+2 -2
View File
@@ -93,13 +93,13 @@ fn log_client_api_key_auth_failure(error: &ClientApiKeyAuthError) {
tracing::warn!(
auth_failure = error.reason(),
error = %message,
"client api key authorization failed"
"Client API key authorization failed"
);
}
_ => {
tracing::info!(
auth_failure = error.reason(),
"client api key authorization failed"
"Client API key authorization failed"
);
}
}
+64 -1
View File
@@ -9,7 +9,10 @@ use serde_json::{Value, json};
use crate::{
dispatch::errors::{DispatchFailureClass, ResponseDispatchError},
upstream::openai::protocol::sse::response_failed_sse_event as encode_response_failed_sse_event,
upstream::openai::protocol::{
responses::ResponsesSseFailure,
sse::response_failed_sse_event as encode_response_failed_sse_event,
},
upstream::openai::transport::CodexClientError,
};
@@ -47,6 +50,13 @@ struct OpenAiErrorKind {
code: &'static str,
}
struct EncodedClientFailure<'a> {
status: StatusCode,
kind: OpenAiErrorKind,
code: &'a str,
message: &'a str,
}
const NO_AVAILABLE_ACCOUNTS_ERROR: OpenAiErrorKind = OpenAiErrorKind {
error_type: "server_error",
code: "no_available_accounts",
@@ -123,6 +133,16 @@ pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> R
return (*status, Json(body)).into_response();
}
if let Some(failure) = response_dispatch_client_failure(error) {
return openai_error_response(
failure.status,
failure.message,
failure.kind.error_type,
failure.code,
)
.into_response();
}
match error {
ResponseDispatchError::NoActiveAccount => {
responses_no_available_accounts_response().into_response()
@@ -132,6 +152,13 @@ pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> R
}
pub fn responses_stream_dispatch_failed_sse_event(error: &ResponseDispatchError) -> String {
if let Some(failure) = response_dispatch_client_failure(error) {
return encode_response_failed_sse_event(
failure.kind.error_type,
failure.code,
failure.message,
);
}
let http_error =
response_dispatch_http_error(error, ResponseDispatchMessageStyle::ResponsesStream);
let kind = response_dispatch_error_kind(error);
@@ -142,6 +169,15 @@ pub(super) fn responses_websocket_dispatch_error_event(
error: &ResponseDispatchError,
request_id: &str,
) -> String {
if let Some(failure) = response_dispatch_client_failure(error) {
return responses_websocket_error_event(
failure.status,
failure.kind.error_type,
failure.code,
failure.message,
request_id,
);
}
if let Some(event) = upstream_client_websocket_error_event(error, request_id) {
return event;
}
@@ -313,6 +349,33 @@ fn responses_error_kind_for_status(status: StatusCode) -> OpenAiErrorKind {
}
}
fn response_dispatch_client_failure(
error: &ResponseDispatchError,
) -> Option<EncodedClientFailure<'_>> {
let client_failure = match error {
ResponseDispatchError::Failed(failure) => failure,
_ => return None,
};
let failure = client_failure.exposed_upstream()?;
Some(encoded_client_failure(
failure,
status_from_u16(client_failure.status_code(), StatusCode::BAD_GATEWAY),
))
}
fn encoded_client_failure(
failure: &ResponsesSseFailure,
status: StatusCode,
) -> EncodedClientFailure<'_> {
let kind = responses_error_kind_for_status(status);
EncodedClientFailure {
status,
kind,
code: failure.upstream_code.as_deref().unwrap_or(kind.code),
message: &failure.message,
}
}
fn response_dispatch_message(
error: &ResponseDispatchError,
style: ResponseDispatchMessageStyle,
+15 -11
View File
@@ -103,32 +103,36 @@ pub fn build_codex_request(
None => {}
}
apply_context_header_fields(&mut request, headers);
apply_identity_fields(&mut request, headers);
apply_session_fields(&mut request, headers);
request.client_user_agent = user_agent_from_headers(headers);
request
}
fn apply_identity_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.client_conversation_id = request.prompt_cache_key().map(ToString::to_string);
request.client_session_id = identity_context_string(request.body(), "session_id")
.or_else(|| header_string(headers, "session-id"));
request.client_thread_id = identity_context_string(request.body(), "thread_id")
fn apply_session_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.client_conversation_id = session_context_string(request.body(), "conversation_id")
.or_else(|| header_string(headers, "conversation-id"))
.or_else(|| header_string(headers, "conversation_id"))
.or_else(|| request.prompt_cache_key().map(ToString::to_string));
request.client_session_id = session_context_string(request.body(), "session_id")
.or_else(|| header_string(headers, "session-id"))
.or_else(|| header_string(headers, "session_id"));
request.client_thread_id = session_context_string(request.body(), "thread_id")
.or_else(|| header_string(headers, "thread-id"));
request.client_request_id = identity_context_string(request.body(), "x-client-request-id")
request.client_request_id = session_context_string(request.body(), "x-client-request-id")
.or_else(|| header_string(headers, "x-client-request-id"));
request.client_turn_id = identity_context_string(request.body(), "turn_id")
request.client_turn_id = session_context_string(request.body(), "turn_id")
.or_else(|| header_string(headers, "x-codex-turn-id"));
request.codex_window_id = request.codex_window_id.take().or_else(|| {
identity_context_string(request.body(), "x-codex-window-id")
session_context_string(request.body(), "x-codex-window-id")
.or_else(|| header_string(headers, "x-codex-window-id"))
});
request.parent_thread_id = request.parent_thread_id.take().or_else(|| {
identity_context_string(request.body(), "x-codex-parent-thread-id")
session_context_string(request.body(), "x-codex-parent-thread-id")
.or_else(|| header_string(headers, "x-codex-parent-thread-id"))
});
}
fn identity_context_string(body: &Map<String, Value>, key: &str) -> Option<String> {
fn session_context_string(body: &Map<String, Value>, key: &str) -> Option<String> {
body_context_string(body, key).or_else(|| {
body.get("client_metadata")
.and_then(Value::as_object)
+49 -170
View File
@@ -1,7 +1,5 @@
//! 官方 Responses WebSocket 入站协议。
use std::sync::Arc;
use axum::{
body::Body,
extract::{Request, State},
@@ -24,11 +22,11 @@ use tokio_tungstenite::{
use crate::{
api::{AppState, middleware::request_id::RequestId},
dispatch::{recovery::history::sanitize_replay_items, service::ResponseDispatchStream},
upstream::openai::protocol::{
sse::{SseEvent, SseEventDecoder},
websocket::{is_terminal_websocket_event, websocket_event_type},
dispatch::{
service::{ConnectionTranscriptFacts, ResponseDispatchStream},
transport::canonical::CanonicalResponseEvent,
},
upstream::openai::protocol::websocket::is_terminal_websocket_event,
};
use super::{PreparedResponsesRequest, ResponsesRequestValidationError, prepare_responses_request};
@@ -72,7 +70,7 @@ pub async fn responses_websocket(State(state): State<AppState>, mut request: Req
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"invalid Responses WebSocket upgrade request"
"Invalid Responses WebSocket upgrade request"
);
return openai_error_response(
StatusCode::BAD_REQUEST,
@@ -133,129 +131,42 @@ struct ResponsesWebSocketSession {
connection_id: String,
}
struct ConnectionReplayState {
last_response_id: Option<String>,
input: Option<Arc<Vec<Value>>>,
}
enum PendingReplayUpdate {
Replace(Vec<Value>),
Append(Vec<Value>),
Unavailable,
}
struct CompletedTurn {
response_id: String,
output: Vec<Value>,
}
#[derive(Default)]
struct CompletedTurnCollector {
struct ConnectionTranscriptCollector {
completed_items: Vec<Value>,
completed: Option<CompletedTurn>,
completed: Option<ConnectionTranscriptFacts>,
}
enum ForwardDispatchOutcome {
Continue(Option<CompletedTurn>),
Continue(Option<ConnectionTranscriptFacts>),
Disconnect,
}
impl ConnectionReplayState {
fn new() -> Self {
Self {
last_response_id: None,
input: None,
}
}
fn prepare(
&self,
request: &mut crate::upstream::openai::protocol::responses::CodexResponsesRequest,
) -> PendingReplayUpdate {
let turn_input = request.input().to_vec();
match request.previous_response_id() {
None => {
request.local_replay_input = None;
PendingReplayUpdate::Replace(turn_input)
}
Some(previous_response_id)
if self.last_response_id.as_deref() == Some(previous_response_id)
&& self.input.is_some() =>
{
request.local_replay_input.clone_from(&self.input);
PendingReplayUpdate::Append(turn_input)
}
Some(_) => {
request.local_replay_input = None;
PendingReplayUpdate::Unavailable
}
}
}
fn commit(&mut self, pending: PendingReplayUpdate, completed: CompletedTurn) {
self.last_response_id = Some(completed.response_id);
match pending {
PendingReplayUpdate::Replace(input) => self.replace(input, completed.output),
PendingReplayUpdate::Append(input) => self.append(input, completed.output),
PendingReplayUpdate::Unavailable => self.clear_input(),
}
}
fn replace(&mut self, mut input: Vec<Value>, mut output: Vec<Value>) {
sanitize_replay_items(&mut input);
sanitize_replay_items(&mut output);
input.append(&mut output);
self.input = Some(Arc::new(input));
}
fn append(&mut self, mut input: Vec<Value>, mut output: Vec<Value>) {
sanitize_replay_items(&mut input);
sanitize_replay_items(&mut output);
input.append(&mut output);
let Some(history) = self.input.as_mut() else {
return;
};
Arc::make_mut(history).append(&mut input);
}
fn clear_input(&mut self) {
self.input = None;
}
}
impl CompletedTurnCollector {
fn observe(&mut self, event: &SseEvent) {
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
return;
};
let event_type = event
.event
.as_deref()
.or_else(|| value.get("type").and_then(Value::as_str));
match event_type {
impl ConnectionTranscriptCollector {
fn observe(&mut self, event: &CanonicalResponseEvent) {
match event.event_type() {
Some("response.output_item.done") => {
if let Some(item) = value.get("item") {
if let Some(item) = event.data().get("item") {
self.completed_items.push(item.clone());
}
}
Some("response.completed") => {
let Some(response_id) = value
let Some(response_id) = event
.data()
.pointer("/response/id")
.and_then(Value::as_str)
.map(ToString::to_string)
else {
return;
};
let output = value
let output = event
.data()
.pointer("/response/output")
.and_then(Value::as_array)
.filter(|output| !output.is_empty())
.cloned()
.unwrap_or_else(|| self.completed_items.clone());
self.completed = Some(CompletedTurn {
response_id,
output,
});
self.completed = Some(ConnectionTranscriptFacts::new(response_id, output));
}
_ => {}
}
@@ -271,7 +182,11 @@ async fn serve_responses_websocket(
"Responses WebSocket connected"
);
let mut request_count = 0u64;
let mut replay_state = ConnectionReplayState::new();
let mut replay_snapshot = session
.state
.services
.responses
.connection_replay_snapshot();
while let Some(message) = socket.next().await {
let payload = match message {
@@ -361,11 +276,19 @@ async fn serve_responses_websocket(
continue;
}
};
let pending_replay = replay_state.prepare(&mut prepared.request);
let replay_plan = session
.state
.services
.responses
.prepare_connection_replay(&replay_snapshot, &mut prepared.request);
match dispatch_response_create(&mut socket, &session, request_id, prepared).await {
ForwardDispatchOutcome::Continue(Some(completed)) => {
replay_state.commit(pending_replay, completed);
ForwardDispatchOutcome::Continue(Some(transcript)) => {
session.state.services.responses.commit_connection_replay(
&mut replay_snapshot,
replay_plan,
transcript,
);
}
ForwardDispatchOutcome::Continue(None) => {}
ForwardDispatchOutcome::Disconnect => break,
@@ -424,6 +347,7 @@ async fn forward_dispatch_stream(
) -> ForwardDispatchOutcome {
let ResponseDispatchStream {
body,
mut canonical_events,
response_headers,
} = stream;
if !send_text(
@@ -436,9 +360,8 @@ async fn forward_dispatch_stream(
}
let mut body = body;
let mut decoder = SseEventDecoder::default();
let mut terminal_sent = false;
let mut completed = CompletedTurnCollector::default();
let mut completed = ConnectionTranscriptCollector::default();
loop {
let next = if terminal_sent {
ActiveResponseInput::Upstream(body.next().await)
@@ -450,28 +373,15 @@ async fn forward_dispatch_stream(
};
match next {
ActiveResponseInput::Upstream(Some(Ok(chunk))) => {
let events = match decoder.push(&chunk) {
Ok(events) => events,
Err(error) => {
tracing::warn!(request_id, error = %error, "failed to decode Responses SSE for WebSocket forwarding");
let sent = send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"invalid_upstream_response",
INVALID_UPSTREAM_EVENT_MESSAGE,
request_id,
)
.await;
drop(body);
return if sent {
ForwardDispatchOutcome::Continue(None)
} else {
ForwardDispatchOutcome::Disconnect
};
}
let Some(events) = canonical_events.recv().await else {
tracing::warn!(
request_id,
"Canonical Responses event stream closed before body stream"
);
return ForwardDispatchOutcome::Disconnect;
};
if !forward_sse_events(
drop(chunk);
if !forward_canonical_events(
socket,
events,
request_id,
@@ -504,37 +414,6 @@ async fn forward_dispatch_stream(
};
}
ActiveResponseInput::Upstream(None) => {
let events = match decoder.finish() {
Ok(events) => events,
Err(error) => {
tracing::warn!(request_id, error = %error, "failed to finish Responses SSE decoding");
return if send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"invalid_upstream_response",
INVALID_UPSTREAM_EVENT_MESSAGE,
request_id,
)
.await
{
ForwardDispatchOutcome::Continue(None)
} else {
ForwardDispatchOutcome::Disconnect
};
}
};
if !forward_sse_events(
socket,
events,
request_id,
&mut terminal_sent,
&mut completed,
)
.await
{
return ForwardDispatchOutcome::Disconnect;
}
if terminal_sent {
return ForwardDispatchOutcome::Continue(completed.completed);
}
@@ -580,15 +459,15 @@ enum ActiveResponseInput {
Downstream(Option<Result<Message, WebSocketError>>),
}
async fn forward_sse_events(
async fn forward_canonical_events(
socket: &mut ResponsesWebSocket,
events: Vec<SseEvent>,
events: Vec<CanonicalResponseEvent>,
request_id: &str,
terminal_sent: &mut bool,
completed: &mut CompletedTurnCollector,
completed: &mut ConnectionTranscriptCollector,
) -> bool {
for event in events {
let Some(event_type) = websocket_event_type(&event.data) else {
let Some(event_type) = event.event_type() else {
tracing::warn!(
request_id,
"Responses SSE event does not contain a valid JSON type"
@@ -605,10 +484,10 @@ async fn forward_sse_events(
return false;
};
completed.observe(&event);
if !send_text(socket, event.data).await {
if !send_text(socket, event.data().to_string()).await {
return false;
}
*terminal_sent |= is_terminal_websocket_event(&event_type);
*terminal_sent |= is_terminal_websocket_event(event_type);
}
true
}
@@ -28,6 +28,11 @@ impl Default for ConnectionDrain {
}
impl ConnectionDrain {
/// 为下层长生命周期执行器提供与连接排空一致的取消信号。
pub(crate) fn cancellation_token(&self) -> CancellationToken {
self.cancellation.clone()
}
/// 注册需要随进程关闭而终止的长连接任务。
pub fn spawn(&self, future: impl Future<Output = ()> + Send + 'static) {
let cancellation = self.cancellation.clone();
+3 -3
View File
@@ -48,7 +48,7 @@ fn make_http_span(request: &Request<Body>) -> Span {
}
fn on_http_request(_request: &Request<Body>, _span: &Span) {
debug!("received HTTP request");
debug!("Received HTTP request");
}
fn on_http_response(response: &Response<Body>, latency: Duration, _span: &Span) {
@@ -58,7 +58,7 @@ fn on_http_response(response: &Response<Body>, latency: Duration, _span: &Span)
info!(
status = response.status().as_u16(),
latency_ms = latency.as_millis(),
"completed HTTP request"
"Completed HTTP request"
);
}
@@ -70,7 +70,7 @@ fn on_http_failure(failure_class: ServerErrorsFailureClass, latency: Duration, _
warn!(
failure = %failure_class,
latency_ms = latency.as_millis(),
"failed HTTP request"
"Failed HTTP request"
);
}
+1 -1
View File
@@ -96,7 +96,7 @@ pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
async fn healthz(State(state): State<AppState>) -> StatusCode {
if let Err(error) = state.services.health_probe.check().await {
tracing::warn!(error, "health check failed");
tracing::warn!(error, "Health check failed");
return StatusCode::SERVICE_UNAVAILABLE;
}
StatusCode::NO_CONTENT
+512 -275
View File
@@ -1,9 +1,23 @@
use std::{collections::BTreeMap, env, path::Path};
use std::{
collections::{BTreeMap, HashSet},
env,
path::{Path, PathBuf},
};
use percent_encoding::{NON_ALPHANUMERIC, utf8_percent_encode};
use serde::{Deserialize, Serialize};
use axum::http::HeaderName;
use secrecy::{ExposeSecret, SecretBox};
use serde::{Deserialize, de::IgnoredAny};
use tracing_subscriber::EnvFilter;
use url::Url;
const CONFIG_SCHEMA_VERSION: u32 = 1;
const CONFIG_RELATIVE_PATH: &str = "deploy/config.yaml";
const SERVER_HOST_ENV: &str = "CPR_SERVER_HOST";
const SERVER_PORT_ENV: &str = "CPR_SERVER_PORT";
const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL";
const REDIS_URL_ENV: &str = "CPR_REDIS_URL";
const SERVICE_PASSWORD_HEX_LENGTH: usize = 48;
const WEAK_ADMIN_PASSWORDS: &[&str] = &[
"",
"admin",
@@ -15,49 +29,230 @@ const WEAK_ADMIN_PASSWORDS: &[&str] = &[
"codex-proxy-rs",
];
/// 应用运行总配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
type SecretValue = SecretBox<String>;
/// 已加载并完成校验的启动配置。
///
/// 数据库、Redis 和管理员初始化密码只在启动阶段保留,并在释放时清零。
#[derive(Debug)]
pub struct BootstrapConfig {
app: AppConfig,
database_url: SecretValue,
redis_url: SecretValue,
admin_default_password: SecretValue,
}
impl BootstrapConfig {
/// 从当前目录或父目录中的 `deploy/config.yaml` 加载配置。
pub fn load() -> Result<Self, ConfigError> {
let current_directory = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?;
let config_path = discover_config_path(&current_directory)?;
let overrides = TopologyOverrides::from_environment()?;
Self::load_from_path_with_overrides(config_path, overrides)
}
/// 从指定 YAML 文件加载配置,不应用容器拓扑覆盖。
pub fn load_from_path(config_path: impl AsRef<Path>) -> Result<Self, ConfigError> {
Self::load_from_path_with_overrides(config_path, TopologyOverrides::default())
}
/// 从指定 YAML 文件加载配置,并应用明确的容器拓扑覆盖。
pub fn load_from_path_with_overrides(
config_path: impl AsRef<Path>,
overrides: TopologyOverrides,
) -> Result<Self, ConfigError> {
let config_path = absolute_path(config_path.as_ref())?;
let config_directory = config_path.parent().ok_or(ConfigError::InvalidConfigPath)?;
let document: ConfigDocument = ::config::Config::builder()
.add_source(::config::File::from(config_path.as_path()).required(true))
.build()
.and_then(::config::Config::try_deserialize)
.map_err(|_| ConfigError::InvalidDocument {
path: config_path.clone(),
})?;
let ConfigDocument {
cpr: mut file_config,
_services: _,
} = document;
if file_config.schema_version != CONFIG_SCHEMA_VERSION {
return Err(ConfigError::UnsupportedSchemaVersion);
}
overrides.apply(&mut file_config);
validate_file_config(&file_config)?;
resolve_relative_path(config_directory, &mut file_config.runtime.data_directory);
resolve_relative_path(config_directory, &mut file_config.logging.file.directory);
let database_url = connection_url_with_password(
&file_config.database.url,
&file_config.database.password,
"database.url",
)?;
let redis_url = connection_url_with_password(
&file_config.redis.url,
&file_config.redis.password,
"redis.url",
)?;
let FileAppConfig {
schema_version: _,
server,
api,
database,
redis,
runtime,
quota,
ws_pool,
fingerprint,
admin,
logging,
telemetry,
} = file_config;
let app = AppConfig {
server,
api,
model_aliases: BTreeMap::new(),
auth: AuthConfig::default(),
quota,
database: DatabaseConfig { url: database.url },
redis: RedisConfig { url: redis.url },
runtime,
ws_pool,
fingerprint,
admin: AdminConfig {
session_ttl_minutes: admin.session_ttl_minutes,
default_username: admin.default_username,
},
logging,
telemetry,
};
Ok(Self {
app,
database_url,
redis_url,
admin_default_password: admin.default_password,
})
}
/// 返回不包含启动密码的应用运行配置。
pub fn app(&self) -> &AppConfig {
&self.app
}
/// 返回已安全注入密码的 PostgreSQL 连接地址。
pub fn database_url(&self) -> &str {
self.database_url.expose_secret()
}
/// 返回已安全注入密码的 Redis 连接地址。
pub fn redis_url(&self) -> &str {
self.redis_url.expose_secret()
}
pub(crate) fn into_parts(self) -> BootstrapConfigParts {
BootstrapConfigParts {
app: self.app,
database_url: self.database_url,
redis_url: self.redis_url,
admin_default_password: self.admin_default_password,
}
}
}
pub(crate) struct BootstrapConfigParts {
pub(crate) app: AppConfig,
pub(crate) database_url: SecretValue,
pub(crate) redis_url: SecretValue,
pub(crate) admin_default_password: SecretValue,
}
/// Docker Compose 对容器内部网络拓扑的固定覆盖。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct TopologyOverrides {
/// 容器内监听地址。
pub server_host: Option<String>,
/// 容器内监听端口。
pub server_port: Option<u16>,
/// 容器内 PostgreSQL 地址,不包含密码。
pub database_url: Option<String>,
/// 容器内 Redis 地址,不包含密码。
pub redis_url: Option<String>,
}
impl TopologyOverrides {
fn from_environment() -> Result<Self, ConfigError> {
let server_port = optional_environment_value(SERVER_PORT_ENV)?
.map(|value| {
value
.parse::<u16>()
.ok()
.filter(|port| *port > 0)
.ok_or(ConfigError::InvalidTopologyOverride(SERVER_PORT_ENV))
})
.transpose()?;
Ok(Self {
server_host: optional_environment_value(SERVER_HOST_ENV)?,
server_port,
database_url: optional_environment_value(DATABASE_URL_ENV)?,
redis_url: optional_environment_value(REDIS_URL_ENV)?,
})
}
fn apply(self, config: &mut FileAppConfig) {
if let Some(host) = self.server_host {
config.server.host = host;
}
if let Some(port) = self.server_port {
config.server.port = port;
}
if let Some(url) = self.database_url {
config.database.url = url;
}
if let Some(url) = self.redis_url {
config.redis.url = url;
}
}
}
/// 应用运行总配置,不包含任何启动密码。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AppConfig {
/// HTTP 服务配置。
pub server: ServerConfig,
/// API 地址配置。
pub api: ApiConfig,
/// 运行时模型别名,由数据库设置加载,配置文件不承载该字段。
#[serde(skip)]
pub model_aliases: BTreeMap<String, String>,
/// 认证与刷新配置,由固定默认值和数据库运行时设置承载,配置文件不承载该字段。
#[serde(skip)]
pub auth: AuthConfig,
/// 配额配置。
#[serde(default)]
pub quota: QuotaConfig,
/// 数据库配置。
/// PostgreSQL 基础连接地址,不包含密码。
pub database: DatabaseConfig,
/// Redis 运行态存储配置。
/// Redis 基础连接地址,不包含密码。
pub redis: RedisConfig,
/// TLS 偏好配置。
#[serde(default)]
pub tls: TlsConfig,
/// 本地持久运行目录。
pub runtime: RuntimePathsConfig,
/// WebSocket 连接池启动设置。
#[serde(default)]
pub ws_pool: WebSocketPoolSettings,
/// 上游请求指纹默认配置。
#[serde(default)]
pub fingerprint: FingerprintConfig,
/// 管理员初始化配置。
#[serde(default)]
/// 管理员初始化配置,不包含密码。
pub admin: AdminConfig,
/// 日志配置。
#[serde(default)]
pub logging: LoggingConfig,
/// PostgreSQL 遥测事实记录配置。
#[serde(default)]
pub telemetry: TelemetryConfig,
}
/// HTTP 监听配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct ServerConfig {
/// 监听主机。
@@ -67,14 +262,15 @@ pub struct ServerConfig {
}
/// 上游 API 地址配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct ApiConfig {
/// 上游基础 URL。
pub base_url: String,
}
/// 认证、轮换与 token 续期配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
pub struct AuthConfig {
/// 提前刷新秒数。
pub refresh_margin_seconds: u64,
@@ -113,58 +309,39 @@ impl Default for AuthConfig {
}
/// 配额刷新与跳过配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct QuotaConfig {
/// 配额刷新周期。
#[serde(default = "default_quota_refresh_interval_minutes")]
pub refresh_interval_minutes: u64,
/// 是否跳过已耗尽配额的账号。
#[serde(default = "default_skip_exhausted")]
pub skip_exhausted: bool,
}
impl Default for QuotaConfig {
fn default() -> Self {
Self {
refresh_interval_minutes: default_quota_refresh_interval_minutes(),
skip_exhausted: default_skip_exhausted(),
}
}
}
fn default_quota_refresh_interval_minutes() -> u64 {
5
}
fn default_skip_exhausted() -> bool {
true
}
/// 数据库连接配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatabaseConfig {
/// PostgreSQL 连接 URL。
/// PostgreSQL 基础连接 URL,不包含密码。
pub url: String,
}
/// Redis 连接配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisConfig {
/// Redis 连接 URL。
/// Redis 基础连接 URL,不包含密码。
pub url: String,
}
/// TLS/HTTP 协议偏好配置。
#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
pub struct TlsConfig {
/// 是否强制 HTTP/1.1。
pub force_http11: bool,
/// 本地持久运行目录。
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct RuntimePathsConfig {
/// 持久身份密钥与更新状态目录。
pub data_directory: PathBuf,
}
/// WebSocket 连接池启动设置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct WebSocketPoolSettings {
/// 是否启用连接池。
@@ -174,31 +351,11 @@ pub struct WebSocketPoolSettings {
/// 单账号最大连接数。
pub max_per_account: usize,
/// 首个真实输出(首事件)到达前的绝对超时;`0` 表示禁用。
///
/// 从发出 `response.create` 起算,覆盖建连/发送/上游排队/思考的全程,直到收到
/// 首个真实内容帧(`response.created`/`response.in_progress` 等纯生命周期帧不计)。
/// 超时判定为连接落到病态上游后端,丢弃并换新连接重试。
#[serde(default = "default_ws_initial_event_timeout_ms")]
pub initial_event_timeout_ms: u64,
}
fn default_ws_initial_event_timeout_ms() -> u64 {
20_000
}
impl Default for WebSocketPoolSettings {
fn default() -> Self {
Self {
enabled: true,
max_age_ms: 3_300_000,
max_per_account: 8,
initial_event_timeout_ms: default_ws_initial_event_timeout_ms(),
}
}
}
/// 上游请求指纹默认配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct FingerprintConfig {
/// 客户端来源名。
@@ -221,53 +378,8 @@ pub struct FingerprintConfig {
pub header_order: Vec<String>,
}
impl Default for FingerprintConfig {
fn default() -> Self {
Self {
originator: "Codex Desktop".to_string(),
app_version: "26.519.81530".to_string(),
build_number: "3178".to_string(),
platform: "darwin".to_string(),
arch: "arm64".to_string(),
chromium_version: "146".to_string(),
user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(),
default_headers: vec![
FingerprintHeaderConfig::new("Accept-Encoding", "gzip, deflate, br, zstd"),
FingerprintHeaderConfig::new("Accept-Language", "en-US,en;q=0.9"),
FingerprintHeaderConfig::new("sec-ch-ua-mobile", "?0"),
FingerprintHeaderConfig::new("sec-ch-ua-platform", "\"macOS\""),
FingerprintHeaderConfig::new("sec-fetch-site", "same-origin"),
FingerprintHeaderConfig::new("sec-fetch-mode", "cors"),
FingerprintHeaderConfig::new("sec-fetch-dest", "empty"),
],
header_order: vec![
"authorization".to_string(),
"chatgpt-account-id".to_string(),
"originator".to_string(),
"x-openai-internal-codex-residency".to_string(),
"x-client-request-id".to_string(),
"x-codex-installation-id".to_string(),
"x-codex-turn-state".to_string(),
"openai-beta".to_string(),
"user-agent".to_string(),
"sec-ch-ua".to_string(),
"sec-ch-ua-mobile".to_string(),
"sec-ch-ua-platform".to_string(),
"accept-encoding".to_string(),
"accept-language".to_string(),
"sec-fetch-site".to_string(),
"sec-fetch-mode".to_string(),
"sec-fetch-dest".to_string(),
"content-type".to_string(),
"accept".to_string(),
"cookie".to_string(),
],
}
}
}
/// 指纹默认请求头配置项。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct FingerprintHeaderConfig {
/// 请求头名称。
@@ -276,76 +388,20 @@ pub struct FingerprintHeaderConfig {
pub value: String,
}
impl FingerprintHeaderConfig {
fn new(name: &str, value: &str) -> Self {
Self {
name: name.to_string(),
value: value.to_string(),
}
}
}
/// 管理员初始化与会话配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
/// 管理员初始化与会话配置,不包含初始化密码。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminConfig {
/// 会话有效期(分钟)。
pub session_ttl_minutes: u64,
/// 默认管理员用户名(首次启动时创建)
#[serde(default = "default_admin_username")]
/// 默认管理员用户名(首次启动时创建)。
pub default_username: String,
/// 默认管理员密码(首次启动时创建)
#[serde(default = "default_admin_password")]
pub default_password: String,
}
impl Default for AdminConfig {
fn default() -> Self {
Self {
session_ttl_minutes: 1440,
default_username: default_admin_username(),
default_password: default_admin_password(),
}
}
}
impl AdminConfig {
/// 校验首次启动创建管理员时使用的默认密码。
pub fn validate_default_password(&self) -> Result<(), AdminConfigError> {
let password = self.default_password.trim();
if password.len() < 12 {
return Err(AdminConfigError::WeakDefaultPassword);
}
if WEAK_ADMIN_PASSWORDS.contains(&password.to_ascii_lowercase().as_str()) {
return Err(AdminConfigError::WeakDefaultPassword);
}
Ok(())
}
}
/// 管理员配置校验错误。
#[derive(Debug, Clone, Copy, thiserror::Error, PartialEq, Eq)]
pub enum AdminConfigError {
/// 默认管理员密码为空、过短或属于常见弱口令。
#[error(
"admin.default_password must be set to a non-default password with at least 12 characters"
)]
WeakDefaultPassword,
}
fn default_admin_username() -> String {
"admin".to_string()
}
fn default_admin_password() -> String {
String::new()
}
/// 应用结构化日志配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct LoggingConfig {
/// 默认日志级别;`RUST_LOG` 存在时优先使用环境变量。
/// 默认日志级别;进程 `RUST_LOG` 可作临时覆盖。
pub level: String,
/// 是否写入标准输出。
pub stdout: bool,
@@ -353,24 +409,14 @@ pub struct LoggingConfig {
pub file: FileLoggingConfig,
}
impl Default for LoggingConfig {
fn default() -> Self {
Self {
level: "info".to_string(),
stdout: true,
file: FileLoggingConfig::default(),
}
}
}
/// 应用结构化文件日志配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct FileLoggingConfig {
/// 是否写入轮转文件。
pub enabled: bool,
/// 日志目录。
pub directory: String,
pub directory: PathBuf,
/// 按中国自然日计算的保留天数。
pub retention_days: usize,
/// 单文件大小上限,单位 MiB。
@@ -379,104 +425,295 @@ pub struct FileLoggingConfig {
pub max_files: usize,
}
impl Default for FileLoggingConfig {
fn default() -> Self {
Self {
enabled: true,
directory: ".runtime/logs".to_string(),
retention_days: 14,
max_file_size_mb: 20,
max_files: 20,
}
}
}
/// PostgreSQL 遥测事实记录配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct TelemetryConfig {
/// 是否记录成功与失败代理事实。
pub enabled: bool,
}
impl Default for TelemetryConfig {
fn default() -> Self {
Self { enabled: true }
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct ConfigDocument {
#[serde(rename = "x-cpr")]
cpr: FileAppConfig,
#[serde(rename = "services")]
_services: IgnoredAny,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct FileAppConfig {
schema_version: u32,
server: ServerConfig,
api: ApiConfig,
database: FileDatabaseConfig,
redis: FileRedisConfig,
runtime: RuntimePathsConfig,
quota: QuotaConfig,
ws_pool: WebSocketPoolSettings,
fingerprint: FingerprintConfig,
admin: FileAdminConfig,
logging: LoggingConfig,
telemetry: TelemetryConfig,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct FileDatabaseConfig {
url: String,
password: SecretValue,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct FileRedisConfig {
url: String,
password: SecretValue,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct FileAdminConfig {
session_ttl_minutes: u64,
default_username: String,
default_password: SecretValue,
}
fn discover_config_path(start: &Path) -> Result<PathBuf, ConfigError> {
start
.ancestors()
.map(|directory| directory.join(CONFIG_RELATIVE_PATH))
.find(|candidate| candidate.is_file())
.ok_or(ConfigError::ConfigFileNotFound)
}
fn absolute_path(path: &Path) -> Result<PathBuf, ConfigError> {
if path.is_absolute() {
return Ok(path.to_path_buf());
}
let current_directory = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?;
Ok(current_directory.join(path))
}
fn resolve_relative_path(base: &Path, path: &mut PathBuf) {
if path.is_relative() {
*path = base.join(&*path);
}
}
const CONFIG_FILE_ENV: &str = "CPR_CONFIG_FILE";
const ADMIN_DEFAULT_PASSWORD_ENV: &str = "CPR_ADMIN_DEFAULT_PASSWORD";
const POSTGRES_PASSWORD_ENV: &str = "CPR_POSTGRES_PASSWORD";
const REDIS_PASSWORD_ENV: &str = "CPR_REDIS_PASSWORD";
impl AppConfig {
/// 从 `CPR_CONFIG_FILE` 或当前目录 `config.yaml` 加载配置文件,并应用运行时密码。
pub fn load() -> Result<Self, ::config::ConfigError> {
let config_file = env::var(CONFIG_FILE_ENV)
.ok()
.filter(|value| !value.trim().is_empty());
match config_file {
Some(path) => load_file(path),
None => Self::load_from_dir("."),
}
}
/// 从指定目录加载配置文件,并应用运行时密码。
pub fn load_from_dir(config_dir: impl AsRef<Path>) -> Result<Self, ::config::ConfigError> {
load_file(config_dir.as_ref().join("config.yaml"))
fn optional_environment_value(name: &'static str) -> Result<Option<String>, ConfigError> {
match env::var(name) {
Ok(value) if value.trim().is_empty() => Err(ConfigError::InvalidTopologyOverride(name)),
Ok(value) => Ok(Some(value)),
Err(env::VarError::NotPresent) => Ok(None),
Err(env::VarError::NotUnicode(_)) => Err(ConfigError::InvalidTopologyOverride(name)),
}
}
fn load_file(config_file: impl AsRef<Path>) -> Result<AppConfig, ::config::ConfigError> {
let mut app_config: AppConfig = ::config::Config::builder()
.add_source(::config::File::from(config_file.as_ref()).required(true))
.build()?
.try_deserialize()?;
apply_password_override(
&mut app_config.database.url,
"database.url",
POSTGRES_PASSWORD_ENV,
fn validate_file_config(config: &FileAppConfig) -> Result<(), ConfigError> {
validate_nonempty("server.host", &config.server.host)?;
if config.server.port == 0 {
return Err(ConfigError::InvalidField("server.port"));
}
validate_url(
"api.base_url",
&config.api.base_url,
&["http", "https"],
false,
)?;
apply_password_override(&mut app_config.redis.url, "redis.url", REDIS_PASSWORD_ENV)?;
if let Some(password) = password_override(ADMIN_DEFAULT_PASSWORD_ENV)? {
app_config.admin.default_password = password;
validate_url(
"database.url",
&config.database.url,
&["postgres", "postgresql"],
true,
)?;
validate_url("redis.url", &config.redis.url, &["redis", "rediss"], true)?;
validate_service_password("database.password", &config.database.password)?;
validate_service_password("redis.password", &config.redis.password)?;
validate_admin_password(config.admin.default_password.expose_secret())?;
validate_nonempty("admin.default_username", &config.admin.default_username)?;
if config.admin.session_ttl_minutes == 0 {
return Err(ConfigError::InvalidField("admin.session_ttl_minutes"));
}
Ok(app_config)
}
fn apply_password_override(
raw_url: &mut String,
config_key: &str,
env_name: &str,
) -> Result<(), ::config::ConfigError> {
let Some(password) = password_override(env_name)? else {
return Ok(());
};
let mut url = Url::parse(raw_url).map_err(|_| {
::config::ConfigError::Message(format!(
"{config_key} must be a valid URL when {env_name} is set"
))
})?;
let encoded_password = utf8_percent_encode(&password, NON_ALPHANUMERIC).to_string();
url.set_password(Some(&encoded_password)).map_err(|()| {
::config::ConfigError::Message(format!(
"{config_key} does not support a password override from {env_name}"
))
})?;
*raw_url = url.into();
if config.runtime.data_directory.as_os_str().is_empty() {
return Err(ConfigError::InvalidField("runtime.data_directory"));
}
if config.quota.refresh_interval_minutes == 0 {
return Err(ConfigError::InvalidField("quota.refresh_interval_minutes"));
}
if config.ws_pool.max_age_ms == 0 {
return Err(ConfigError::InvalidField("ws_pool.max_age_ms"));
}
if config.ws_pool.max_per_account == 0 {
return Err(ConfigError::InvalidField("ws_pool.max_per_account"));
}
validate_fingerprint(&config.fingerprint)?;
validate_logging(&config.logging)?;
Ok(())
}
fn password_override(env_name: &str) -> Result<Option<String>, ::config::ConfigError> {
match env::var(env_name) {
Ok(password) if password.is_empty() => Err(::config::ConfigError::Message(format!(
"{env_name} must not be empty"
))),
Ok(password) => Ok(Some(password)),
Err(env::VarError::NotPresent) => Ok(None),
Err(env::VarError::NotUnicode(_)) => Err(::config::ConfigError::Message(format!(
"{env_name} must contain valid UTF-8"
))),
fn validate_url(
field: &'static str,
raw_url: &str,
allowed_schemes: &[&str],
reject_password: bool,
) -> Result<(), ConfigError> {
let url = Url::parse(raw_url).map_err(|_| ConfigError::InvalidField(field))?;
if !allowed_schemes.contains(&url.scheme()) || url.host_str().is_none() {
return Err(ConfigError::InvalidField(field));
}
if reject_password && url.password().is_some() {
return Err(ConfigError::PasswordInUrl(field));
}
Ok(())
}
fn validate_service_password(
field: &'static str,
password: &SecretValue,
) -> Result<(), ConfigError> {
let password = password.expose_secret();
if password.len() != SERVICE_PASSWORD_HEX_LENGTH
|| !password.bytes().all(|byte| byte.is_ascii_hexdigit())
{
return Err(ConfigError::InvalidServicePassword(field));
}
Ok(())
}
fn validate_admin_password(password: &str) -> Result<(), ConfigError> {
let password = password.trim();
if password.len() < 12
|| password.contains('$')
|| WEAK_ADMIN_PASSWORDS.contains(&password.to_ascii_lowercase().as_str())
{
return Err(ConfigError::WeakAdminPassword);
}
Ok(())
}
fn validate_fingerprint(config: &FingerprintConfig) -> Result<(), ConfigError> {
for (field, value) in [
("fingerprint.originator", config.originator.as_str()),
("fingerprint.app_version", config.app_version.as_str()),
("fingerprint.build_number", config.build_number.as_str()),
("fingerprint.platform", config.platform.as_str()),
("fingerprint.arch", config.arch.as_str()),
(
"fingerprint.chromium_version",
config.chromium_version.as_str(),
),
(
"fingerprint.user_agent_template",
config.user_agent_template.as_str(),
),
] {
validate_nonempty(field, value)?;
}
if config.default_headers.is_empty() || config.header_order.is_empty() {
return Err(ConfigError::InvalidField("fingerprint"));
}
let mut header_names = HashSet::new();
for header in &config.default_headers {
let name = header.name.trim().to_ascii_lowercase();
if name.is_empty()
|| header.value.is_empty()
|| HeaderName::from_bytes(name.as_bytes()).is_err()
|| !header_names.insert(name)
{
return Err(ConfigError::InvalidField("fingerprint.default_headers"));
}
}
let mut ordered_headers = HashSet::new();
for header in &config.header_order {
let header = header.trim().to_ascii_lowercase();
if header.is_empty()
|| HeaderName::from_bytes(header.as_bytes()).is_err()
|| !ordered_headers.insert(header)
{
return Err(ConfigError::InvalidField("fingerprint.header_order"));
}
}
Ok(())
}
fn validate_logging(config: &LoggingConfig) -> Result<(), ConfigError> {
EnvFilter::try_new(&config.level).map_err(|_| ConfigError::InvalidField("logging.level"))?;
if !config.stdout && !config.file.enabled {
return Err(ConfigError::InvalidField("logging"));
}
if config.file.directory.as_os_str().is_empty() {
return Err(ConfigError::InvalidField("logging.file.directory"));
}
if config.file.retention_days == 0 {
return Err(ConfigError::InvalidField("logging.file.retention_days"));
}
if config.file.max_file_size_mb == 0 {
return Err(ConfigError::InvalidField("logging.file.max_file_size_mb"));
}
if config.file.max_files == 0 {
return Err(ConfigError::InvalidField("logging.file.max_files"));
}
Ok(())
}
fn validate_nonempty(field: &'static str, value: &str) -> Result<(), ConfigError> {
if value.trim().is_empty() {
return Err(ConfigError::InvalidField(field));
}
Ok(())
}
fn connection_url_with_password(
raw_url: &str,
password: &SecretValue,
field: &'static str,
) -> Result<SecretValue, ConfigError> {
let mut url = Url::parse(raw_url).map_err(|_| ConfigError::InvalidField(field))?;
url.set_password(Some(password.expose_secret()))
.map_err(|()| ConfigError::InvalidField(field))?;
Ok(SecretBox::new(Box::new(url.into())))
}
/// 启动配置加载错误。错误文本不会包含配置值或连接地址。
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ConfigError {
/// 无法确定当前工作目录。
#[error("failed to determine the current working directory")]
CurrentDirectory,
/// 从当前目录及父目录找不到配置文件。
#[error("deploy/config.yaml was not found in the current directory or its parents")]
ConfigFileNotFound,
/// 配置路径没有父目录。
#[error("configuration file path is invalid")]
InvalidConfigPath,
/// YAML 无法读取、解析或不符合结构;不透传解析器文本以避免泄露密码。
#[error("failed to load configuration file `{path}`; check its permissions and YAML structure")]
InvalidDocument {
/// 配置文件路径。
path: PathBuf,
},
/// 配置格式版本不受支持。
#[error("x-cpr.schema_version must be 1")]
UnsupportedSchemaVersion,
/// 字段值非法。
#[error("configuration field `{0}` is invalid")]
InvalidField(&'static str),
/// 连接 URL 不允许内嵌密码。
#[error("configuration field `{0}` must not contain a password; use its password field")]
PasswordInUrl(&'static str),
/// PostgreSQL 或 Redis 密码格式非法。
#[error("configuration field `{0}` must be exactly 48 hexadecimal characters")]
InvalidServicePassword(&'static str),
/// 管理员初始化密码非法。
#[error("admin.default_password must be strong, at least 12 characters, and contain no `$`")]
WeakAdminPassword,
/// Compose 拓扑覆盖非法。
#[error("container topology override `{0}` is invalid")]
InvalidTopologyOverride(&'static str),
}
+29 -19
View File
@@ -3,13 +3,14 @@
use std::{env, error::Error, future::IntoFuture, sync::Arc, time::Duration};
use axum::Router;
use secrecy::ExposeSecret;
use crate::{
auth::{
service::SessionService,
store::{PgAdminUserStore, RedisAdminSessionStore},
},
bootstrap::config::AppConfig,
bootstrap::config::{AppConfig, BootstrapConfig},
dispatch::{
affinity::{RedisSessionAffinityStore, SessionAffinityService},
service::{ResponseDispatchService, ResponseDispatchServiceParts},
@@ -211,7 +212,7 @@ impl Services {
let websocket_initial_event_timeout = (config.ws_pool.initial_event_timeout_ms > 0)
.then(|| std::time::Duration::from_millis(config.ws_pool.initial_event_timeout_ms));
let client = CodexBackendClient::new(
build_reqwest_client(config.tls.force_http11)?,
build_reqwest_client()?,
config.api.base_url.clone(),
fingerprint.clone(),
)
@@ -296,8 +297,10 @@ impl Services {
let session_affinity =
Arc::new(SessionAffinityService::new(stores.session_affinity.clone()));
let cloudflare_recovery =
crate::dispatch::recovery::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
let connection_drain = crate::api::middleware::connection_drain::ConnectionDrain::default();
let cloudflare_recovery = crate::dispatch::controllers::cloudflare::CloudflareRecovery::new(
stores.cookies.clone(),
);
let responses = Arc::new(ResponseDispatchService::new(ResponseDispatchServiceParts {
account_pool: account_pool.clone(),
models: models.clone(),
@@ -306,9 +309,9 @@ impl Services {
recorder,
account_identity,
cloudflare: cloudflare_recovery,
shutdown: connection_drain.cancellation_token(),
}));
let system_update = Arc::new(SystemUpdateService::from_env());
let connection_drain = crate::api::middleware::connection_drain::ConnectionDrain::default();
Ok(Self {
database,
@@ -445,13 +448,12 @@ pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOption
.pool_options(&settings_snapshot_from_config(config))
}
pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + Sync>> {
wait_for_scheduled_restart().await;
let config = AppConfig::load()?;
let host = config.server.host.clone();
let port = config.server.port;
let log_guard = init_logging(&config)?;
let host = config.app().server.host.clone();
let port = config.app().server.port;
let log_guard = init_logging(config.app())?;
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
let (router, task_coordinator, connection_drain) = build_router(config).await?;
let mut task_coordinator = Some(task_coordinator);
@@ -501,7 +503,7 @@ pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
Ok(())
}
() = crate::bootstrap::shutdown::shutdown_signal() => {
tracing::warn!("received a second shutdown signal; cancelling remaining connections");
tracing::warn!("Received a second shutdown signal; cancelling remaining connections");
Ok(())
}
}
@@ -534,7 +536,7 @@ pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
const HTTP_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
async fn build_router(
mut config: AppConfig,
bootstrap_config: BootstrapConfig,
) -> Result<
(
Router,
@@ -543,8 +545,17 @@ async fn build_router(
),
Box<dyn Error + Send + Sync>,
> {
let pool = connect(&config.database.url).await?;
let redis = RedisConnection::connect(&config.redis.url, "cpr").await?;
let crate::bootstrap::config::BootstrapConfigParts {
app: mut config,
database_url,
redis_url,
admin_default_password,
} = bootstrap_config.into_parts();
let pool = connect(database_url.expose_secret()).await?;
drop(database_url);
let redis = RedisConnection::connect(redis_url.expose_secret(), "cpr").await?;
drop(redis_url);
let settings =
SettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool).await?;
apply_settings_to_config(&mut config, &settings);
@@ -567,16 +578,14 @@ async fn build_router(
request_buckets: PgRequestBucketStore::new(pool),
};
config.admin.validate_default_password()?;
let default_admin_password = config.admin.default_password.clone();
let data_dir = ensure_data_dir()?;
let data_dir = ensure_data_dir(&config.runtime.data_directory)?;
let identity_secret = load_or_create_identity_secret(&data_dir)?;
let default_fingerprint = fingerprint_from_config(&config.fingerprint);
let runtime_fingerprint = fingerprint_store
.ensure_current_seed(&default_fingerprint)
.await?;
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(config.clone());
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(&config);
let services = Services::try_with_identity_secret_and_usage_record_options(
&config,
stores,
@@ -588,8 +597,9 @@ async fn build_router(
let created_default_admin = services
.admin_sessions
.ensure_default_admin(&default_admin_password)
.ensure_default_admin(admin_default_password.expose_secret())
.await?;
drop(admin_default_password);
tracing::info!(
created = created_default_admin,
"默认管理员账号已完成初始化检查"
+3 -5
View File
@@ -9,7 +9,6 @@ use crate::{
/// 运行时配置镜像(从 AppConfig 衍生)。
#[derive(Debug, Clone)]
pub struct RuntimeConfig {
pub admin: crate::bootstrap::config::AdminConfig,
pub auth: crate::bootstrap::config::AuthConfig,
pub quota: crate::bootstrap::config::QuotaConfig,
}
@@ -40,12 +39,11 @@ impl HealthProbe for RuntimeHealthProbe {
}
}
impl From<AppConfig> for RuntimeConfig {
fn from(config: AppConfig) -> Self {
impl From<&AppConfig> for RuntimeConfig {
fn from(config: &AppConfig) -> Self {
Self {
admin: config.admin.clone(),
auth: config.auth.clone(),
quota: config.quota,
quota: config.quota.clone(),
}
}
}
+1 -1
View File
@@ -82,6 +82,7 @@ impl TaskCoordinator {
stores.accounts.clone(),
stores.account_usage.clone(),
services.codex.clone(),
services.account_pool.clone(),
services.account_pseudonymizer.clone(),
config
.quota
@@ -91,7 +92,6 @@ impl TaskCoordinator {
30 * 60,
)
.with_cookie_store(stores.cookies.clone())
.with_account_pool(services.account_pool.clone())
.start(),
);
coordinator.push("fingerprint_update", {
@@ -44,8 +44,8 @@ pub fn start_fingerprint_update_task(
},
PeriodicTaskConfig::new(
APPCAST_POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动",
"fingerprint 后台版本检查器已关闭",
"Fingerprint 后台版本检查器已启动",
"Fingerprint 后台版本检查器已关闭",
),
)
}
@@ -68,15 +68,15 @@ impl PeriodicTaskRunner for FingerprintUpdateTask {
info!(
app_version = %app_version,
build_number = %build_number,
"fingerprint 运行时版本已更新"
"Fingerprint 运行时版本已更新"
);
}
Ok(None) => {}
Err(error) => {
if first_tick {
warn!(error = %error, "fingerprint 首次版本检查失败");
warn!(error = %error, "Fingerprint 首次版本检查失败");
} else {
warn!(error = %error, "fingerprint 定期版本检查失败");
warn!(error = %error, "Fingerprint 定期版本检查失败");
}
}
}
+6 -10
View File
@@ -35,6 +35,7 @@ impl QuotaRefreshTask {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<AccountPseudonymizer>,
interval_secs: u64,
min_refresh_interval_secs: u64,
@@ -44,6 +45,7 @@ impl QuotaRefreshTask {
store,
usage_store,
codex,
account_pool,
account_pseudonymizer,
min_refresh_interval_secs,
),
@@ -58,18 +60,12 @@ impl QuotaRefreshTask {
self
}
/// 设置运行时账号池,用于刷新后同步内存调度状态。
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.service = self.service.with_account_pool(account_pool);
self
}
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let config = PeriodicTaskConfig::new(
self.interval_secs,
"quota 刷新任务已启动",
"quota 刷新任务已关闭",
"Quota 刷新任务已启动",
"Quota 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
@@ -96,7 +92,7 @@ impl PeriodicTaskRunner for QuotaRefreshTask {
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
warn!(error = %error, "Quota 刷新任务失败");
}
}
})
@@ -107,6 +103,6 @@ fn log_refreshed_summary(summary: QuotaRefreshSummary) {
info!(
refreshed = summary.refreshed,
failed = summary.failed,
"quota 刷新任务完成"
"Quota 刷新任务完成"
);
}
+4 -4
View File
@@ -54,8 +54,8 @@ where
pub fn start(self) -> SchedulerHandle {
let config = PeriodicTaskConfig::new(
self.interval_secs,
"token 刷新任务已启动",
"token 刷新任务已关闭",
"Token 刷新任务已启动",
"Token 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
@@ -73,7 +73,7 @@ where
scheduled = summary.scheduled,
immediate = summary.immediate,
replaced = summary.replaced,
"token 刷新定时器已调度"
"Token 刷新定时器已调度"
);
}
Ok(summary) => {
@@ -84,7 +84,7 @@ where
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
warn!(error = %error, "Token 刷新定时器调度失败");
}
}
})
+211 -85
View File
@@ -1,10 +1,15 @@
//! 客户端身份、本地会话索引与账号作用域 wire 身份隔离。
use serde_json::{Map, Value};
use crate::{
infra::identity::AccountPseudonymizer,
upstream::openai::protocol::responses::CodexResponsesRequest,
};
const TURN_STATE_METADATA_KEY: &str = "x-codex-turn-state";
const TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
/// 持久化密钥驱动的账号身份隔离服务。
#[derive(Clone)]
pub struct AccountIdentityService {
@@ -21,25 +26,51 @@ impl std::fmt::Debug for AccountIdentityService {
/// 单次账号 attempt 的 wire 身份。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountScopedIdentity {
pub prompt_cache_key: Option<String>,
pub session_id: Option<String>,
pub thread_id: Option<String>,
pub client_request_id: String,
pub turn_id: Option<String>,
pub window_id: Option<String>,
pub parent_thread_id: Option<String>,
pub(in crate::dispatch) struct AccountScopedIdentity {
pub installation_id: String,
}
struct IdentityValues<'a> {
prompt_cache_key: Option<&'a str>,
session_id: Option<&'a str>,
thread_id: Option<&'a str>,
client_request_id: Option<&'a str>,
turn_id: Option<&'a str>,
window_id: Option<&'a str>,
parent_thread_id: Option<&'a str>,
/// 单个请求内账号绑定 opaque 状态的 owner。
pub(in crate::dispatch) struct AccountIdentityScope {
turn_state_account_id: Option<String>,
}
impl AccountIdentityScope {
pub(in crate::dispatch) fn new(preferred_account_id: Option<String>) -> Self {
Self {
turn_state_account_id: preferred_account_id,
}
}
fn preserves_turn_state(&mut self, account_id: &str, has_turn_state: bool) -> bool {
if !has_turn_state {
return false;
}
let owner = self
.turn_state_account_id
.get_or_insert_with(|| account_id.to_string());
owner == account_id
}
}
/// 已按当前账号重建 wire 身份、可安全交给上游 transport 的请求。
pub(in crate::dispatch) struct AccountScopedRequest {
request: CodexResponsesRequest,
identity: AccountScopedIdentity,
}
impl AccountScopedRequest {
pub(in crate::dispatch) fn request(&self) -> &CodexResponsesRequest {
&self.request
}
pub(in crate::dispatch) fn identity(&self) -> &AccountScopedIdentity {
&self.identity
}
pub(in crate::dispatch) fn into_request(self) -> CodexResponsesRequest {
self.request
}
}
impl AccountIdentityService {
@@ -63,88 +94,183 @@ impl AccountIdentityService {
.map(|anchor| format!("lc_{}", self.pseudonym("local-conversation", None, anchor)));
}
pub fn scope(
/// 从请求语义为单次账号 attempt 重建完整 wire 身份。
pub(in crate::dispatch) fn scope_request(
&self,
request: &CodexResponsesRequest,
scope: &mut AccountIdentityScope,
mut request: CodexResponsesRequest,
account_id: &str,
proxy_request_id: &str,
) -> AccountScopedIdentity {
self.scope_values(
) -> AccountScopedRequest {
let identity = AccountScopedIdentity {
installation_id: self.pseudonymizer.installation_id(account_id),
};
let client_metadata_turn_state = client_metadata_string(&request, TURN_STATE_METADATA_KEY);
let preserve_turn_state = scope.preserves_turn_state(
account_id,
proxy_request_id,
IdentityValues {
prompt_cache_key: request.prompt_cache_key(),
session_id: request.client_session_id.as_deref(),
thread_id: request.client_thread_id.as_deref(),
client_request_id: request.client_request_id.as_deref(),
turn_id: request.client_turn_id.as_deref(),
window_id: request.codex_window_id.as_deref(),
parent_thread_id: request.parent_thread_id.as_deref(),
},
)
request.turn_state.is_some() || client_metadata_turn_state.is_some(),
);
let turn_state = if preserve_turn_state {
request.turn_state.clone()
} else {
None
};
let client_metadata_turn_state = if preserve_turn_state {
client_metadata_turn_state
} else {
None
};
let turn_metadata = request
.turn_metadata
.as_deref()
.map(|raw| Self::scope_turn_metadata(raw, &identity));
let client_metadata_turn_metadata = client_metadata_string(&request, TURN_METADATA_KEY)
.map(|raw| Self::scope_turn_metadata(&raw, &identity));
apply_identity_to_request(
&mut request,
&identity,
turn_state.as_deref(),
turn_metadata.as_deref(),
client_metadata_turn_state.as_deref(),
client_metadata_turn_metadata.as_deref(),
);
request.turn_state = turn_state;
request.turn_metadata = turn_metadata;
AccountScopedRequest { request, identity }
}
pub fn scope_auxiliary(
&self,
account_id: &str,
proxy_request_id: &str,
) -> AccountScopedIdentity {
pub(in crate::dispatch) fn scope_auxiliary(&self, account_id: &str) -> AccountScopedIdentity {
AccountScopedIdentity {
client_request_id: format!(
"cr_{}",
self.pseudonym("client-request-id", Some(account_id), proxy_request_id)
),
installation_id: self.pseudonymizer.installation_id(account_id),
..AccountScopedIdentity::default()
}
}
fn scope_values(
&self,
account_id: &str,
proxy_request_id: &str,
values: IdentityValues<'_>,
) -> AccountScopedIdentity {
AccountScopedIdentity {
prompt_cache_key: self.scope_optional(
"prompt-cache-key",
account_id,
values.prompt_cache_key,
),
session_id: self.scope_optional("session-id", account_id, values.session_id),
thread_id: self.scope_optional("thread-id", account_id, values.thread_id),
client_request_id: format!(
"cr_{}",
self.pseudonym(
"client-request-id",
Some(account_id),
values.client_request_id.unwrap_or(proxy_request_id),
)
),
turn_id: self.scope_optional("turn-id", account_id, values.turn_id),
window_id: self.scope_optional("window-id", account_id, values.window_id),
parent_thread_id: self.scope_optional(
"parent-thread-id",
account_id,
values.parent_thread_id,
),
installation_id: self.pseudonymizer.installation_id(account_id),
}
}
fn scope_optional(
&self,
domain: &str,
account_id: &str,
value: Option<&str>,
) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!("wi_{}", self.pseudonym(domain, Some(account_id), value)))
fn scope_turn_metadata(raw: &str, identity: &AccountScopedIdentity) -> String {
let Ok(Value::Object(mut metadata)) = serde_json::from_str::<Value>(raw) else {
return raw.to_string();
};
let contains_installation_identity = [
"installation_id",
"installationId",
"x-codex-installation-id",
]
.into_iter()
.any(|key| metadata.contains_key(key));
if !contains_installation_identity {
return raw.to_string();
}
for key in [
"installation_id",
"installationId",
"x-codex-installation-id",
] {
replace_existing_metadata_field(&mut metadata, key, Some(&identity.installation_id));
}
serde_json::to_string(&metadata).unwrap_or_else(|_| raw.to_string())
}
fn pseudonym(&self, domain: &str, account_id: Option<&str>, value: &str) -> String {
self.pseudonymizer.scoped(domain, account_id, value)
}
}
fn client_metadata_string(request: &CodexResponsesRequest, key: &str) -> Option<String> {
request
.client_metadata()?
.as_object()?
.get(key)?
.as_str()
.map(ToString::to_string)
}
fn apply_identity_to_request(
request: &mut CodexResponsesRequest,
identity: &AccountScopedIdentity,
turn_state: Option<&str>,
turn_metadata: Option<&str>,
client_metadata_turn_state: Option<&str>,
client_metadata_turn_metadata: Option<&str>,
) {
for (key, value) in [
(
"x-codex-installation-id",
Some(identity.installation_id.as_str()),
),
("installation_id", Some(identity.installation_id.as_str())),
("installationId", Some(identity.installation_id.as_str())),
("turnState", turn_state),
("x-codex-turn-state", turn_state),
("turnMetadata", turn_metadata),
("x-codex-turn-metadata", turn_metadata),
] {
request.replace_existing_identity_field(key, value);
}
let client_metadata = scoped_client_metadata(
request.client_metadata(),
identity,
client_metadata_turn_state,
client_metadata_turn_metadata,
);
request.set_client_metadata(client_metadata);
}
fn scoped_client_metadata(
client_metadata: Option<&Value>,
identity: &AccountScopedIdentity,
turn_state: Option<&str>,
turn_metadata: Option<&str>,
) -> Option<Value> {
let client_metadata = client_metadata?;
let Value::Object(mut metadata) = client_metadata.clone() else {
return Some(client_metadata.clone());
};
for (key, value) in [
(
"x-codex-installation-id",
Some(identity.installation_id.as_str()),
),
(TURN_STATE_METADATA_KEY, turn_state),
(TURN_METADATA_KEY, turn_metadata),
] {
replace_metadata_field(&mut metadata, key, value);
}
replace_existing_metadata_field(
&mut metadata,
"installation_id",
Some(&identity.installation_id),
);
replace_existing_metadata_field(
&mut metadata,
"installationId",
Some(&identity.installation_id),
);
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn replace_existing_metadata_field(
metadata: &mut Map<String, Value>,
key: &str,
value: Option<&str>,
) {
if metadata.contains_key(key) {
replace_metadata_field(metadata, key, value);
}
}
fn replace_metadata_field(metadata: &mut Map<String, Value>, key: &str, value: Option<&str>) {
match value.filter(|value| !value.trim().is_empty()) {
Some(value) => {
metadata.insert(key.to_string(), Value::String(value.to_string()));
}
None => {
metadata.remove(key);
}
}
}
+2 -1
View File
@@ -6,7 +6,8 @@ pub mod service;
pub mod store;
pub mod types;
pub use identity::{AccountIdentityService, AccountScopedIdentity};
pub use identity::AccountIdentityService;
pub(in crate::dispatch) use identity::{AccountIdentityScope, AccountScopedRequest};
pub use resolve::*;
pub use service::*;
pub use store::*;
+1 -101
View File
@@ -1,107 +1,7 @@
use chrono::{DateTime, Utc};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::{
dispatch::{
affinity::{SessionAffinityEntry, SessionAffinityService},
recovery::history::HistoryRecoveryPlan,
service::ResponseDispatchService,
},
upstream::openai::protocol::{
events::TokenUsage,
responses::{CodexResponsesRequest, completed_response_metadata},
},
};
impl ResponseDispatchService {
pub(in crate::dispatch) async fn preferred_account_id_for_request(
&self,
request: &CodexResponsesRequest,
history: &HistoryRecoveryPlan,
now: DateTime<Utc>,
) -> Option<String> {
if let Some(account_id) = history.preferred_account_id() {
return Some(account_id.to_string());
}
if request.previous_response_id().is_some() {
return None;
}
self.conversation_account_id(request, now).await
}
async fn conversation_account_id(
&self,
request: &CodexResponsesRequest,
now: DateTime<Utc>,
) -> Option<String> {
let conversation_id = request
.local_conversation_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?;
let variant_hash = compute_variant_hash(request);
self.session_affinity
.lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now)
.await
}
}
pub(in crate::dispatch) async fn record_response_affinity(
session_affinity: &SessionAffinityService,
history: &HistoryRecoveryPlan,
original_request: &CodexResponsesRequest,
account_id: &str,
body: &str,
turn_state: Option<String>,
usage: Option<TokenUsage>,
) {
let metadata = match completed_response_metadata(body) {
Ok(Some(metadata)) => metadata,
Ok(None) => return,
Err(error) => {
tracing::warn!(
error = %error,
"failed to parse completed response metadata for session affinity"
);
return;
}
};
let conversation_id = history
.conversation_id(original_request)
.unwrap_or(&metadata.response_id)
.to_string();
let variant_hash = compute_variant_hash(original_request);
let entry = SessionAffinityEntry {
account_id: account_id.to_string(),
conversation_id,
turn_state: turn_state
.filter(|value| !value.trim().is_empty())
.or_else(|| original_request.turn_state.clone()),
instructions_hash: Some(hash_instructions(Some(original_request.instructions()))),
input_tokens: usage.map(|usage| usage.input_tokens),
function_call_ids: metadata.function_call_ids,
variant_hash: Some(variant_hash),
continuation_scope: if original_request.store() {
crate::upstream::openai::protocol::responses::PreviousResponseScope::Persisted
} else {
crate::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal
},
created_at: Utc::now(),
};
if let Err(error) = session_affinity
.record(metadata.response_id.clone(), entry)
.await
{
tracing::warn!(
error = %error,
response_id = %metadata.response_id,
account_id = %account_id,
"failed to record session affinity"
);
}
}
use crate::upstream::openai::protocol::responses::CodexResponsesRequest;
// ====================================================================
// Variant identity helpers
+10 -4
View File
@@ -3,6 +3,8 @@
use chrono::{DateTime, Duration, Utc};
use thiserror::Error;
use crate::infra::redis::RedisConnection;
use super::{
store::{RedisSessionAffinityStore, RedisSessionAffinityStoreError},
types::SessionAffinityEntry,
@@ -65,7 +67,7 @@ impl SessionAffinityService {
tracing::warn!(
response_id,
error = %error,
"failed to remove Redis session affinity"
"Failed to remove Redis session affinity"
);
false
}
@@ -76,17 +78,21 @@ impl SessionAffinityService {
match self.store.forget_account(account_id).await {
Ok(forgotten) => forgotten > 0,
Err(error) => {
tracing::warn!(account_id, error = %error, "failed to remove account affinities");
tracing::warn!(account_id, error = %error, "Failed to remove account affinities");
false
}
}
}
pub(crate) fn redis_connection(&self) -> RedisConnection {
self.store.redis_connection()
}
async fn entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<SessionAffinityEntry> {
match self.store.get(response_id, now, self.ttl).await {
Ok(entry) => entry,
Err(error) => {
tracing::warn!(response_id, error = %error, "failed to read Redis session affinity");
tracing::warn!(response_id, error = %error, "Failed to read Redis session affinity");
None
}
}
@@ -106,7 +112,7 @@ impl SessionAffinityService {
{
Ok(entry) => entry,
Err(error) => {
tracing::warn!(conversation_id, error = %error, "failed to query Redis affinity index");
tracing::warn!(conversation_id, error = %error, "Failed to query Redis affinity index");
None
}
}
+4
View File
@@ -296,6 +296,10 @@ impl RedisSessionAffinityStore {
Self { redis }
}
pub(crate) fn redis_connection(&self) -> RedisConnection {
self.redis.clone()
}
pub async fn upsert(
&self,
response_id: &str,
@@ -0,0 +1,208 @@
//! 账号级失败分类、隔离副作用与换号决策 owner。
use crate::{
dispatch::{
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
},
},
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_response_failure, classify_upstream_failure,
},
pool::AccountPoolService,
},
upstream::openai::{
failure::{UpstreamFailureFacts, UpstreamFailureKind},
protocol::responses::ResponsesSseFailure,
transport::CodexBackendClient,
},
};
pub(super) struct AccountFailureController;
pub(super) struct ClassifiedFailure {
exhaustion: AccountExhaustionRecord,
effect: Option<AccountStateEffect>,
}
impl AccountFailureController {
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedFailure> {
let account_id = observation.account.as_ref()?.id.as_str();
classify(observation, account_id)
}
pub(super) fn classify_failure(
account_id: &str,
failure: ControllerFailureFact<'_>,
) -> Option<ClassifiedFailure> {
classify_failure(account_id, failure)
}
pub(super) async fn apply_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
classified: &ClassifiedFailure,
) {
let Some(account_id) = classified.exhaustion.account_id.as_deref() else {
return;
};
if let Some(effect) = &classified.effect {
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
}
}
pub(super) fn decision(
observation: &AttemptObservation,
classified: ClassifiedFailure,
) -> AttemptDecision {
if observation.routing.can_retry_next_candidate {
AttemptDecision::RetryNextCandidate {
exhaustion: Some(classified.exhaustion),
on_exhaustion: None,
}
} else {
AttemptDecision::Return(
crate::dispatch::lifecycle::contract::AttemptReturnKind::Observed,
)
}
}
pub(super) fn is_retryable_transport(observation: &AttemptObservation) -> bool {
let AttemptObservationKind::UpstreamFailure(facts) = &observation.kind else {
return false;
};
matches!(
facts.kind,
UpstreamFailureKind::HttpConnect
| UpstreamFailureKind::HttpTimeout
| UpstreamFailureKind::HttpTransport
| UpstreamFailureKind::StreamIdle
| UpstreamFailureKind::WebSocketTransport
| UpstreamFailureKind::WebSocketTimeout
) || facts
.status_code
.is_some_and(|status| (500..=599).contains(&status))
}
pub(super) fn client_failure(
classified: &ClassifiedFailure,
failure: &ResponsesSseFailure,
) -> ClientFailure {
let status_code =
classified
.exhaustion
.status_code
.unwrap_or(match classified.exhaustion.kind {
ExhaustedAccountKind::ModelUnsupported => 400,
ExhaustedAccountKind::Expired | ExhaustedAccountKind::Disabled => 401,
ExhaustedAccountKind::Banned => 403,
ExhaustedAccountKind::QuotaExhausted
| ExhaustedAccountKind::RateLimited
| ExhaustedAccountKind::CloudflareChallenge
| ExhaustedAccountKind::CloudflarePathBlocked
| ExhaustedAccountKind::UpstreamUnavailable => 502,
});
ClientFailure::new(failure.clone(), status_code, false)
}
pub(super) fn unowned_client_failure(failure: &ResponsesSseFailure) -> Option<ClientFailure> {
let code = failure
.upstream_code
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
if code.contains("invalid_request")
|| code.contains("not_found")
|| code.contains("context_window")
|| code.contains("invalid_prompt")
|| code.contains("bad_request")
{
return Some(ClientFailure::new(failure.clone(), 400, false));
}
if code.contains("server_overloaded") {
return Some(ClientFailure::new(failure.clone(), 503, true));
}
classify_response_failure(failure).map(|classified| {
let status = match classified.kind {
AccountFailureKind::ModelUnsupported => 400,
AccountFailureKind::Expired => 401,
AccountFailureKind::Disabled | AccountFailureKind::Banned => 403,
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => 429,
};
ClientFailure::new(failure.clone(), status, false)
})
}
}
fn classify(observation: &AttemptObservation, account_id: &str) -> Option<ClassifiedFailure> {
match &observation.kind {
AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Empty) => {
Some(ClassifiedFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::UpstreamUnavailable,
"upstream response did not include visible output",
),
effect: None,
})
}
_ => ControllerFailureFact::from_attempt(observation)
.and_then(|failure| classify_failure(account_id, failure)),
}
}
fn classify_failure(
account_id: &str,
failure: ControllerFailureFact<'_>,
) -> Option<ClassifiedFailure> {
match failure {
ControllerFailureFact::Upstream(facts) => classify_upstream(account_id, facts),
ControllerFailureFact::Response(failure) => classify_sse(account_id, failure),
}
}
fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<ClassifiedFailure> {
classified_failure(
account_id,
&facts.body,
facts.status_code,
classify_upstream_failure(facts)?,
)
}
fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<ClassifiedFailure> {
classified_failure(
account_id,
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
failure.explicit_status_code,
classify_response_failure(failure)?,
)
}
fn classified_failure(
account_id: &str,
message: &str,
status_code: Option<u16>,
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
) -> Option<ClassifiedFailure> {
let kind = match classified.kind {
AccountFailureKind::ModelUnsupported => ExhaustedAccountKind::ModelUnsupported,
AccountFailureKind::Expired => ExhaustedAccountKind::Expired,
AccountFailureKind::Disabled => ExhaustedAccountKind::Disabled,
AccountFailureKind::Banned => ExhaustedAccountKind::Banned,
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => return None,
};
let mut exhaustion = AccountExhaustionRecord::new(account_id, kind, message);
if let Some(status_code) = status_code {
exhaustion = exhaustion.with_status_code(status_code);
}
Some(ClassifiedFailure {
exhaustion,
effect: classified.effect,
})
}
@@ -0,0 +1,137 @@
//! 会话亲和性路由与写入的唯一 feature owner。
use std::time::Duration;
use chrono::{DateTime, Utc};
use tokio::time::timeout;
use crate::dispatch::affinity::{
SessionAffinityEntry, SessionAffinityService, compute_variant_hash, hash_instructions,
};
use crate::upstream::openai::protocol::{
events::TokenUsage,
responses::{CodexResponsesRequest, PreviousResponseScope, completed_response_metadata},
};
const REDIS_BEST_EFFORT_TIMEOUT: Duration = Duration::from_millis(100);
pub(super) struct AffinityController;
pub(super) struct ResponseExit<'a> {
pub affinity: &'a SessionAffinityService,
pub conversation_id: Option<&'a str>,
pub request: &'a CodexResponsesRequest,
pub account_id: &'a str,
pub body: &'a str,
pub turn_state: Option<String>,
pub usage: Option<TokenUsage>,
pub continuation_scope: PreviousResponseScope,
}
pub(super) struct StreamExit<'a> {
pub response: ResponseExit<'a>,
pub completed: bool,
}
impl AffinityController {
pub(super) async fn preferred_account_id(
affinity: &SessionAffinityService,
request: &CodexResponsesRequest,
now: DateTime<Utc>,
) -> Option<String> {
if request.previous_response_id().is_some() {
return None;
}
let conversation_id = request
.local_conversation_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?;
let variant_hash = compute_variant_hash(request);
match timeout(
REDIS_BEST_EFFORT_TIMEOUT,
affinity.lookup_latest_account_by_conversation(
conversation_id,
None,
Some(&variant_hash),
now,
),
)
.await
{
Ok(account_id) => account_id,
Err(_) => {
tracing::warn!(
conversation_id,
timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(),
"Timed out reading conversation affinity; continuing without preference"
);
None
}
}
}
pub(super) async fn leave_complete(exit: ResponseExit<'_>) {
record_response_affinity(exit).await;
}
pub(super) async fn leave_stream(exit: StreamExit<'_>) {
if !exit.completed {
return;
}
record_response_affinity(exit.response).await;
}
}
async fn record_response_affinity(exit: ResponseExit<'_>) {
let ResponseExit {
affinity,
conversation_id,
request,
account_id,
body,
turn_state,
usage,
continuation_scope,
} = exit;
let metadata = match completed_response_metadata(body) {
Ok(Some(metadata)) => metadata,
Ok(None) => return,
Err(error) => {
tracing::warn!(error = %error, "Failed to parse completed response metadata for session affinity");
return;
}
};
let entry = SessionAffinityEntry {
account_id: account_id.to_string(),
conversation_id: conversation_id.unwrap_or(&metadata.response_id).to_string(),
turn_state: turn_state
.filter(|value| !value.trim().is_empty())
.or_else(|| request.turn_state.clone()),
instructions_hash: Some(hash_instructions(Some(request.instructions()))),
input_tokens: usage.map(|usage| usage.input_tokens),
function_call_ids: metadata.function_call_ids,
variant_hash: Some(compute_variant_hash(request)),
continuation_scope,
created_at: Utc::now(),
};
match timeout(
REDIS_BEST_EFFORT_TIMEOUT,
affinity.record(metadata.response_id.clone(), entry),
)
.await
{
Ok(Ok(())) => {}
Ok(Err(error)) => {
tracing::warn!(error = %error, response_id = %metadata.response_id, account_id, "Failed to record session affinity");
}
Err(_) => {
tracing::warn!(
response_id = %metadata.response_id,
account_id,
timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(),
"Timed out recording session affinity"
);
}
}
}
@@ -0,0 +1,237 @@
//! Cloudflare challenge、路径封禁与 Cookie 生命周期的唯一 feature owner。
use chrono::Utc;
use crate::{
dispatch::{
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
lifecycle::stream::StreamTerminal,
},
fleet::{
account::AccountStatus,
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
pool::AccountPoolService,
},
upstream::openai::failure::UpstreamFailureFacts,
};
pub(super) struct CloudflareController;
pub(super) struct StreamExit<'a> {
pub recovery: &'a CloudflareRecovery,
pub account_id: &'a str,
pub set_cookie_headers: &'a [String],
pub terminal: &'a StreamTerminal,
}
#[derive(Clone, Copy)]
enum CloudflareFailureKind {
Challenge,
PathBlocked,
}
pub(super) struct CloudflareFailure {
kind: CloudflareFailureKind,
exhaustion: AccountExhaustionRecord,
}
impl CloudflareController {
pub(super) fn classify(observation: &AttemptObservation) -> Option<CloudflareFailure> {
classify(observation)
}
pub(super) async fn apply_effect(
recovery: &CloudflareRecovery,
account_pool: &AccountPoolService,
failure: &CloudflareFailure,
) {
match failure.kind {
CloudflareFailureKind::Challenge => {
let Some(account_id) = failure.exhaustion.account_id.as_deref() else {
return;
};
recovery.apply_challenge(account_pool, account_id).await;
}
CloudflareFailureKind::PathBlocked => {
let Some(account_id) = failure.exhaustion.account_id.as_deref() else {
return;
};
recovery.apply_path_block(account_pool, account_id).await;
}
}
}
pub(super) fn decision(
observation: &AttemptObservation,
failure: CloudflareFailure,
) -> AttemptDecision {
if observation.routing.can_retry_next_candidate {
return AttemptDecision::RetryNextCandidate {
exhaustion: Some(failure.exhaustion),
on_exhaustion: None,
};
}
AttemptDecision::Return(AttemptReturnKind::Observed)
}
pub(super) async fn leave_complete(recovery: &CloudflareRecovery, account_id: &str) {
recovery.reset_account_recovery(account_id).await;
}
pub(super) async fn leave_stream(exit: StreamExit<'_>) {
exit.recovery
.capture_set_cookie_headers(exit.account_id, exit.set_cookie_headers)
.await;
if matches!(
exit.terminal,
StreamTerminal::Completed { .. } | StreamTerminal::Incomplete { .. }
) {
exit.recovery.reset_account_recovery(exit.account_id).await;
}
}
}
fn classify(observation: &AttemptObservation) -> Option<CloudflareFailure> {
let account_id = observation.account.as_ref()?.id.as_str();
let AttemptObservationKind::UpstreamFailure(facts) = &observation.kind else {
return None;
};
classify_upstream(account_id, facts)
}
fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<CloudflareFailure> {
if facts.status_code == Some(403) && is_challenge(&facts.body) {
return Some(CloudflareFailure {
kind: CloudflareFailureKind::Challenge,
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::CloudflareChallenge,
"Upstream blocked the request (Cloudflare challenge)",
),
});
}
if facts.status_code == Some(404) && facts.body.trim().is_empty() {
return Some(CloudflareFailure {
kind: CloudflareFailureKind::PathBlocked,
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::CloudflarePathBlocked,
"Upstream blocked the request (Cloudflare path-block)",
),
});
}
None
}
fn is_challenge(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("cf-mitigated")
|| value.contains("cf-chl-bypass")
|| value.contains("_cf_chl")
|| value.contains("cf_chl")
|| value.contains("attention required")
|| value.contains("just a moment")
}
#[derive(Clone)]
pub struct CloudflareRecovery {
path_block_tracker: CloudflarePathBlockTracker,
challenge_tracker: CloudflareChallengeCooldownTracker,
cookie_store: PgCookieStore,
}
impl CloudflareRecovery {
pub fn new(cookie_store: PgCookieStore) -> Self {
Self {
path_block_tracker: CloudflarePathBlockTracker::default(),
challenge_tracker: CloudflareChallengeCooldownTracker::default(),
cookie_store,
}
}
pub async fn cookie_header_for_request(&self, account_id: &str, path: &str) -> Option<String> {
self.cookie_store
.cookie_header_for_request(account_id, "chatgpt.com", path)
.await
.ok()?
}
pub async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) {
for header in headers {
if let Err(error) = self
.cookie_store
.capture_set_cookie(account_id, header)
.await
{
tracing::warn!(
account_id,
error = %error,
"Failed to persist upstream set-cookie header"
);
}
}
}
pub async fn apply_challenge(&self, account_pool: &AccountPoolService, account_id: &str) {
let now = Utc::now();
let cooldown = self
.challenge_tracker
.record_challenge(account_id, now)
.await;
account_pool
.set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until)
.await;
if let Err(error) = self
.cookie_store
.expire_account_cookies_at(account_id, cooldown.cooldown_until)
.await
{
tracing::warn!(
account_id,
error = %error,
"Failed to persist Cloudflare challenge cookie cleanup deadline"
);
}
}
pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) {
self.delete_account_cookies(account_id, "Cloudflare path-block")
.await;
let now = Utc::now();
self.path_block_tracker
.record_path_block(account_id, now)
.await;
if self
.path_block_tracker
.should_disable(account_id, now)
.await
{
account_pool
.set_status(account_id, AccountStatus::Disabled)
.await;
}
}
pub async fn reset_account_recovery(&self, account_id: &str) {
self.path_block_tracker.reset(account_id).await;
self.challenge_tracker.reset(account_id).await;
}
async fn delete_account_cookies(&self, account_id: &str, reason: &str) {
delete_account_cookies(&self.cookie_store, account_id, reason).await;
}
}
async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) {
if let Err(error) = cookie_store.delete_account_cookies(account_id).await {
tracing::warn!(
account_id,
reason,
error = %error,
"Failed to delete account cookies after Cloudflare recovery signal"
);
}
}
@@ -0,0 +1,227 @@
//! `cyber_policy` 会话级生命周期控制器。
use std::{collections::BTreeSet, sync::Arc, time::Duration as StdDuration};
use chrono::Duration;
use sha2::{Digest, Sha256};
use tokio::time::timeout;
mod store;
mod types;
use crate::{
dispatch::{
affinity::SessionAffinityService,
controllers::ControllerFailureFact,
errors::ClientFailure,
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
},
};
use self::{store::Store, types::SessionState};
const SESSION_STATE_TTL: Duration = Duration::hours(1);
const STATE_IO_TIMEOUT: StdDuration = StdDuration::from_millis(100);
/// 单次请求使用的 `cyber_policy` 会话路由计划。
pub(in crate::dispatch) struct CyberPolicyScope {
session_key: Option<String>,
state: SessionState,
}
impl CyberPolicyScope {
pub(in crate::dispatch) fn excluded_account_ids(&self) -> BTreeSet<String> {
self.state.failed_account_ids.iter().cloned().collect()
}
fn session_key(&self) -> Option<&str> {
self.session_key.as_deref()
}
fn has_failures(&self) -> bool {
!self.state.failed_account_ids.is_empty()
}
}
/// 识别失败、维护会话状态,并为下一次请求生成账号排除集合。
#[derive(Clone)]
pub(in crate::dispatch) struct CyberPolicyController {
store: Store,
}
/// 已由 Cyber owner 认领的失败;调用方只能请求 owner 产生 effect 或 decision。
pub(super) struct ClassifiedCyberPolicyFailure<'a> {
fact: ControllerFailureFact<'a>,
}
impl CyberPolicyController {
pub(in crate::dispatch) fn new(session_affinity: Arc<SessionAffinityService>) -> Self {
Self {
store: Store::new(session_affinity.redis_connection()),
}
}
pub(super) fn classify(
fact: ControllerFailureFact<'_>,
) -> Option<ClassifiedCyberPolicyFailure<'_>> {
let owned = match fact {
ControllerFailureFact::Upstream(facts) => {
facts
.status_code
.is_some_and(|status| (400..=499).contains(&status))
&& is_cyber_policy_code(facts.code.as_deref())
}
ControllerFailureFact::Response(failure) => is_cyber_policy_failure(failure),
};
owned.then_some(ClassifiedCyberPolicyFailure { fact })
}
pub(super) fn client_failure(classified: &ClassifiedCyberPolicyFailure<'_>) -> ClientFailure {
match classified.fact {
ControllerFailureFact::Upstream(facts) => ClientFailure::new(
failure_from_upstream_facts(facts),
facts.status_code.unwrap_or(400),
true,
),
ControllerFailureFact::Response(failure) => {
ClientFailure::new(failure.clone(), 400, true)
}
}
}
pub(super) fn decision(
observation: &AttemptObservation,
classified: ClassifiedCyberPolicyFailure<'_>,
) -> AttemptDecision {
let client_failure = Self::client_failure(&classified);
if observation.routing.can_retry_next_candidate {
AttemptDecision::RetryNextCandidate {
exhaustion: None,
on_exhaustion: Some(client_failure),
}
} else {
AttemptDecision::Return(AttemptReturnKind::Failed(client_failure))
}
}
pub(in crate::dispatch) async fn prepare(
&self,
request: &CodexResponsesRequest,
) -> CyberPolicyScope {
let Some(session_key) = session_key(request) else {
return CyberPolicyScope {
session_key: None,
state: SessionState::default(),
};
};
let state = match timeout(STATE_IO_TIMEOUT, self.store.load(&session_key)).await {
Ok(Ok(state)) => state.unwrap_or_default(),
Ok(Err(error)) => {
tracing::warn!(error = %error, "Failed to read cyber policy session state");
SessionState::default()
}
Err(_) => {
tracing::warn!("Timed out reading cyber policy session state");
SessionState::default()
}
};
CyberPolicyScope {
session_key: Some(session_key),
state,
}
}
pub(super) async fn exclude_account(&self, plan: &CyberPolicyScope, account_id: &str) {
let Some(session_key) = plan.session_key() else {
return;
};
match timeout(
STATE_IO_TIMEOUT,
self.store
.record_failure(session_key, account_id, SESSION_STATE_TTL),
)
.await
{
Ok(Ok(state)) => tracing::warn!(
account_id,
excluded_account_count = state.failed_account_ids.len(),
"Excluded cyber policy account from this session"
),
Ok(Err(error)) => tracing::warn!(
account_id,
error = %error,
"Failed to record cyber policy session state"
),
Err(_) => tracing::warn!(account_id, "Timed out recording cyber policy session state"),
}
}
pub(in crate::dispatch) async fn observe_success(&self, plan: &CyberPolicyScope) {
if !plan.has_failures() {
return;
}
let Some(session_key) = plan.session_key() else {
return;
};
match timeout(
STATE_IO_TIMEOUT,
self.store.clear(session_key, &plan.state.revision),
)
.await
{
Ok(Ok(_)) => {}
Ok(Err(error)) => {
tracing::warn!(error = %error, "Failed to clear cyber policy session state");
}
Err(_) => tracing::warn!("Timed out clearing cyber policy session state"),
}
}
}
fn session_key(request: &CodexResponsesRequest) -> Option<String> {
if request.previous_response_id().is_some() {
return None;
}
let api_key_id = non_empty(request.client_api_key_id.as_deref())?;
let explicit_session_id = non_empty(request.client_session_id.as_deref())
.or_else(|| non_empty(request.client_conversation_id.as_deref()))
.or_else(|| {
request
.explicit_prompt_cache_key
.then(|| non_empty(request.prompt_cache_key()))
.flatten()
})?;
let mut hasher = Sha256::new();
hasher.update(b"cyber-policy-session\0");
hasher.update(api_key_id.as_bytes());
hasher.update(b"\0");
hasher.update(explicit_session_id.as_bytes());
Some(hex::encode(hasher.finalize()))
}
fn is_cyber_policy_code(code: Option<&str>) -> bool {
code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy"))
}
fn failure_from_upstream_facts(facts: &UpstreamFailureFacts) -> ResponsesSseFailure {
ResponsesSseFailure {
event: "error".to_string(),
message: facts.message.clone(),
upstream_code: facts.code.clone(),
upstream_type: None,
explicit_status_code: facts.status_code,
retry_after_seconds: facts.retry_after_seconds,
}
}
pub(super) fn is_cyber_policy_failure(failure: &ResponsesSseFailure) -> bool {
is_cyber_policy_code(failure.upstream_code.as_deref())
}
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
@@ -0,0 +1,119 @@
use chrono::Duration;
use redis::AsyncCommands;
use thiserror::Error;
use crate::infra::redis::RedisConnection;
use super::types::SessionState;
const STATE_VERSION: &str = "cyber-policy:v2";
const RECORD_FAILURE_SCRIPT: &str = r#"
local raw = redis.call('GET', KEYS[1])
local state = { failedAccountIds = {}, revision = '' }
if raw then
local ok, decoded = pcall(cjson.decode, raw)
if ok and type(decoded) == 'table' and type(decoded.failedAccountIds) == 'table' then
state = decoded
end
end
local account_id = ARGV[1]
local known = false
for _, failed_account_id in ipairs(state.failedAccountIds) do
if failed_account_id == account_id then
known = true
break
end
end
if not known then
table.insert(state.failedAccountIds, account_id)
end
state.revision = ARGV[3]
local encoded = cjson.encode(state)
redis.call('SET', KEYS[1], encoded, 'EX', tonumber(ARGV[2]))
return encoded
"#;
const CLEAR_STATE_SCRIPT: &str = r#"
local raw = redis.call('GET', KEYS[1])
if not raw then
return 0
end
local ok, state = pcall(cjson.decode, raw)
if not ok or type(state) ~= 'table' then
return 0
end
if tostring(state.revision or '') ~= ARGV[1] then
return 0
end
return redis.call('DEL', KEYS[1])
"#;
#[derive(Clone)]
pub struct Store {
redis: RedisConnection,
}
#[derive(Debug, Error)]
pub enum StoreError {
#[error("Redis cyber policy operation failed: {0}")]
Redis(#[from] redis::RedisError),
#[error("invalid Redis cyber policy value: {0}")]
Json(#[from] serde_json::Error),
}
impl Store {
pub fn new(redis: RedisConnection) -> Self {
Self { redis }
}
pub async fn load(&self, session_key: &str) -> Result<Option<SessionState>, StoreError> {
let mut connection = self.redis.manager();
let value: Option<String> = connection.get(self.key(session_key)).await?;
value
.map(|value| serde_json::from_str(&value))
.transpose()
.map_err(Into::into)
}
pub async fn record_failure(
&self,
session_key: &str,
account_id: &str,
ttl: Duration,
) -> Result<SessionState, StoreError> {
let mut connection = self.redis.manager();
let revision = uuid::Uuid::new_v4().to_string();
let value: String = redis::Script::new(RECORD_FAILURE_SCRIPT)
.key(self.key(session_key))
.arg(account_id)
.arg(ttl.num_seconds().max(1))
.arg(revision)
.invoke_async(&mut connection)
.await?;
Ok(serde_json::from_str(&value)?)
}
pub async fn clear(
&self,
session_key: &str,
expected_revision: &str,
) -> Result<bool, StoreError> {
let mut connection = self.redis.manager();
let removed: usize = redis::Script::new(CLEAR_STATE_SCRIPT)
.key(self.key(session_key))
.arg(expected_revision)
.invoke_async(&mut connection)
.await?;
Ok(removed > 0)
}
fn key(&self, session_key: &str) -> String {
self.redis
.key(&format!("{STATE_VERSION}:session:{session_key}"))
}
}
@@ -0,0 +1,9 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionState {
pub failed_account_ids: Vec<String>,
#[serde(default)]
pub revision: String,
}
+438
View File
@@ -0,0 +1,438 @@
//! previous response 所有权、重放与恢复规则的唯一 owner。
use std::{sync::Arc, time::Duration};
use chrono::Utc;
use serde_json::Value;
use tokio::time::timeout;
use crate::{
dispatch::{
affinity::{SessionAffinityEntry, SessionAffinityService},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
AttemptRoutingFacts, CompleteResponseFacts, PinnedCandidateAcquireFailureKind,
},
},
upstream::openai::{
failure::{UpstreamFailureFacts, UpstreamFailureKind},
protocol::responses::{CodexResponsesRequest, PreviousResponseScope, StreamCommitPolicy},
transport::CodexBackendTransport,
transport::websocket::PreviousResponseUnavailableReason,
},
};
const REDIS_BEST_EFFORT_TIMEOUT: Duration = Duration::from_millis(100);
const CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE: &str =
"previous response history cannot be sent to another account";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum HistorySource {
None,
Managed,
ConnectionLocal,
ExternalUnknown,
}
/// 单个请求内的 previous response 状态;策略只由 `HistoryController` 解释。
pub(super) struct HistoryState {
source: HistorySource,
managed_entry: Option<SessionAffinityEntry>,
local_replay_input: Option<Arc<Vec<Value>>>,
external_attempted_account_id: Option<String>,
replay_owner_account: bool,
}
#[derive(Default)]
pub(in crate::dispatch) struct ConnectionReplaySnapshot {
last_response_id: Option<String>,
input: Option<Arc<Vec<Value>>>,
}
pub(in crate::dispatch) struct ConnectionReplayPlan {
update: ConnectionReplayUpdate,
}
enum ConnectionReplayUpdate {
Replace(Vec<Value>),
Append {
history: Arc<Vec<Value>>,
input: Vec<Value>,
},
Unavailable,
}
pub(in crate::dispatch) struct HistoryController;
impl HistoryController {
pub(in crate::dispatch) fn new_connection_replay_snapshot() -> ConnectionReplaySnapshot {
ConnectionReplaySnapshot::default()
}
pub(in crate::dispatch) fn prepare_connection_replay(
snapshot: &ConnectionReplaySnapshot,
request: &mut CodexResponsesRequest,
) -> ConnectionReplayPlan {
let turn_input = request.input().to_vec();
let update = match request.previous_response_id() {
None => {
request.local_replay_input = None;
ConnectionReplayUpdate::Replace(turn_input)
}
Some(previous_response_id)
if snapshot.last_response_id.as_deref() == Some(previous_response_id) =>
{
if let Some(history) = snapshot.input.as_ref() {
request.local_replay_input = Some(Arc::clone(history));
ConnectionReplayUpdate::Append {
history: Arc::clone(history),
input: turn_input,
}
} else {
request.local_replay_input = None;
ConnectionReplayUpdate::Unavailable
}
}
Some(_) => {
request.local_replay_input = None;
ConnectionReplayUpdate::Unavailable
}
};
request.local_replay_available = !matches!(update, ConnectionReplayUpdate::Unavailable);
ConnectionReplayPlan { update }
}
pub(in crate::dispatch) fn commit_connection_replay(
snapshot: &mut ConnectionReplaySnapshot,
plan: ConnectionReplayPlan,
response_id: String,
output: Vec<Value>,
) {
snapshot.last_response_id = Some(response_id);
match plan.update {
ConnectionReplayUpdate::Replace(input) => {
snapshot.input = Some(Arc::new(sanitized_transcript(input, output)));
}
ConnectionReplayUpdate::Append { mut history, input } => {
let mut turn = sanitized_transcript(input, output);
Arc::make_mut(&mut history).append(&mut turn);
snapshot.input = Some(history);
}
ConnectionReplayUpdate::Unavailable => snapshot.input = None,
}
}
pub(super) async fn enter(
affinity: &SessionAffinityService,
request: &CodexResponsesRequest,
) -> HistoryState {
let Some(previous_response_id) = request.previous_response_id() else {
return HistoryState::empty();
};
let managed_entry = match timeout(
REDIS_BEST_EFFORT_TIMEOUT,
affinity.lookup(previous_response_id, Utc::now()),
)
.await
{
Ok(entry) => entry,
Err(_) => {
tracing::warn!(
previous_response_id,
timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(),
"Timed out reading previous response affinity; continuing without managed history"
);
None
}
};
let local_replay_input = request.local_replay_input.clone();
let source = match (&managed_entry, &local_replay_input) {
(Some(_), _) => HistorySource::Managed,
(None, Some(_)) => HistorySource::ConnectionLocal,
(None, None) => HistorySource::ExternalUnknown,
};
HistoryState {
source,
managed_entry,
local_replay_input,
external_attempted_account_id: None,
replay_owner_account: false,
}
}
pub(super) fn preferred_account_id(state: &HistoryState) -> Option<&str> {
state
.managed_entry
.as_ref()
.map(|entry| entry.account_id.as_str())
}
pub(super) fn prepare_attempt(
state: &mut HistoryState,
request: &CodexResponsesRequest,
account_id: &str,
) -> Result<CodexResponsesRequest, String> {
let prepared = match state.source {
HistorySource::None => Some(request.clone()),
HistorySource::ExternalUnknown => {
match state.external_attempted_account_id.as_deref() {
None => {
state.external_attempted_account_id = Some(account_id.to_string());
Some(external_history_request(request))
}
Some(first_account_id) if first_account_id == account_id => {
Some(external_history_request(request))
}
Some(_) => None,
}
}
HistorySource::ConnectionLocal => state
.local_replay_input
.as_deref()
.map(|replay_input| full_replay_request(request, replay_input)),
HistorySource::Managed => prepare_managed_attempt(state, request, account_id),
};
prepared.ok_or_else(|| CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE.to_string())
}
pub(super) fn prepare_same_account_retry(state: &mut HistoryState, account_id: &str) -> bool {
if !Self::can_retry_same_account(state, account_id) {
return false;
}
state.replay_owner_account = true;
true
}
pub(super) fn routing_facts(
state: &HistoryState,
account_id: Option<&str>,
) -> AttemptRoutingFacts {
AttemptRoutingFacts {
external_origin: state.source == HistorySource::ExternalUnknown,
can_retry_same_account: account_id
.is_some_and(|account_id| Self::can_retry_same_account(state, account_id)),
can_retry_next_candidate: match state.source {
HistorySource::None | HistorySource::ConnectionLocal => true,
HistorySource::Managed => state.local_replay_input.is_some(),
HistorySource::ExternalUnknown => false,
},
}
}
pub(super) fn conversation_id<'a>(
state: &'a HistoryState,
request: &'a CodexResponsesRequest,
) -> Option<&'a str> {
state
.managed_entry
.as_ref()
.map(|entry| entry.conversation_id.as_str())
.or_else(|| non_empty(request.local_conversation_id.as_deref()))
}
pub(super) fn decide(observation: &AttemptObservation) -> Option<AttemptDecision> {
if let AttemptObservationKind::PinnedCandidateUnavailable { kind, .. } = observation.kind {
return Some(match kind {
PinnedCandidateAcquireFailureKind::Busy => {
AttemptDecision::Return(AttemptReturnKind::ContinuationBusy)
}
PinnedCandidateAcquireFailureKind::Unavailable => {
AttemptDecision::Return(AttemptReturnKind::RouteUnavailable {
message: "the account that owns the previous response is unavailable for same-account recovery".to_string(),
})
}
});
}
let detail = match &observation.kind {
AttemptObservationKind::UpstreamFailure(facts) if is_continuation_busy(facts) => {
if observation.routing.can_retry_same_account {
return Some(AttemptDecision::RetrySameAccount);
}
return Some(AttemptDecision::Return(AttemptReturnKind::ContinuationBusy));
}
AttemptObservationKind::UpstreamFailure(facts) if is_history_failure(facts) => {
facts.body.clone()
}
AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Failed(failure))
| AttemptObservationKind::StreamFailure(failure)
if is_history_code(failure.upstream_code.as_deref()) =>
{
crate::dispatch::failure::sse::sse_failure_error_body(failure)
}
_ => return None,
};
if observation.routing.can_retry_same_account {
return Some(AttemptDecision::RetrySameAccount);
}
if observation.routing.external_origin {
return Some(AttemptDecision::Return(AttemptReturnKind::Observed));
}
Some(AttemptDecision::Return(
AttemptReturnKind::RouteUnavailable { message: detail },
))
}
pub(in crate::dispatch) fn continuation_scope(
request: &CodexResponsesRequest,
transport: CodexBackendTransport,
connection_local_continuation: bool,
) -> PreviousResponseScope {
if request.store() {
PreviousResponseScope::Persisted
} else if transport == CodexBackendTransport::WebSocket && connection_local_continuation {
PreviousResponseScope::ConnectionLocal
} else if request.local_replay_available {
PreviousResponseScope::ReplayRequired
} else {
PreviousResponseScope::Unavailable
}
}
fn can_retry_same_account(state: &HistoryState, account_id: &str) -> bool {
let Some(entry) = state.managed_entry.as_ref() else {
return false;
};
entry.account_id == account_id
&& state.local_replay_input.is_some()
&& !state.replay_owner_account
}
}
fn prepare_managed_attempt(
state: &HistoryState,
request: &CodexResponsesRequest,
account_id: &str,
) -> Option<CodexResponsesRequest> {
let entry = state.managed_entry.as_ref()?;
if entry.account_id != account_id || state.replay_owner_account {
return replay_for_conversation(state, request, &entry.conversation_id);
}
match entry.continuation_scope {
PreviousResponseScope::Persisted
| PreviousResponseScope::ConnectionLocal
| PreviousResponseScope::ExternalUnknown => {
let mut prepared = request.clone();
prepared.local_conversation_id = Some(entry.conversation_id.clone());
prepared.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
prepared.previous_response_scope = Some(entry.continuation_scope);
if prepared.turn_state.as_deref().is_none_or(str::is_empty) {
prepared.turn_state.clone_from(&entry.turn_state);
}
Some(prepared)
}
PreviousResponseScope::ReplayRequired => {
replay_for_conversation(state, request, &entry.conversation_id)
}
PreviousResponseScope::Unavailable => None,
}
}
fn replay_for_conversation(
state: &HistoryState,
request: &CodexResponsesRequest,
conversation_id: &str,
) -> Option<CodexResponsesRequest> {
state.local_replay_input.as_deref().map(|replay_input| {
let mut prepared = full_replay_request(request, replay_input);
prepared.local_conversation_id = Some(conversation_id.to_string());
prepared
})
}
impl HistoryState {
fn empty() -> Self {
Self {
source: HistorySource::None,
managed_entry: None,
local_replay_input: None,
external_attempted_account_id: None,
replay_owner_account: false,
}
}
}
fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesRequest {
let mut request = original.clone();
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
request.previous_response_scope =
Some(crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown);
request
}
fn full_replay_request(
original: &CodexResponsesRequest,
replay_input: &[Value],
) -> CodexResponsesRequest {
let mut request = original.clone();
let mut input = Vec::with_capacity(replay_input.len() + original.input().len());
input.extend_from_slice(replay_input);
input.extend_from_slice(original.input());
sanitize_replay_items(&mut input);
request.set_input(input);
request.set_previous_response_id(None);
request.turn_state = None;
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
request
}
fn sanitized_transcript(mut input: Vec<Value>, mut output: Vec<Value>) -> Vec<Value> {
sanitize_replay_items(&mut input);
sanitize_replay_items(&mut output);
input.append(&mut output);
input
}
fn sanitize_replay_items(values: &mut [Value]) {
for value in values {
match value {
Value::Array(values) => sanitize_replay_items(values),
Value::Object(object) => {
object.remove("id");
object.remove("encrypted_content");
for value in object.values_mut() {
sanitize_replay_items(std::slice::from_mut(value));
}
}
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
}
}
}
fn is_continuation_busy(facts: &UpstreamFailureFacts) -> bool {
matches!(
facts.kind,
UpstreamFailureKind::ContinuationUnavailable(
PreviousResponseUnavailableReason::ConnectionBusy
)
)
}
fn is_history_failure(facts: &UpstreamFailureFacts) -> bool {
facts
.code
.as_deref()
.is_some_and(|code| is_history_code(Some(code)))
|| matches!(
facts.kind,
UpstreamFailureKind::ContinuationUnavailable(reason)
if reason != PreviousResponseUnavailableReason::ConnectionBusy
)
}
fn is_history_code(code: Option<&str>) -> bool {
matches!(
code,
Some(
"previous_response_not_found"
| "invalid_encrypted_content"
| "missing_tool_output"
| "no_tool_output"
)
)
}
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
+998
View File
@@ -0,0 +1,998 @@
//! v1 请求生命周期控制器集合。
//!
//! 控制器集合是静态类型的 feature owner 边界:生命周期只依赖这里暴露的
//! typed hook,不直接知道具体功能名称或状态存储实现。
use std::{
collections::BTreeSet,
future::Future,
sync::Arc,
time::{Duration, Instant},
};
use chrono::{DateTime, Utc};
use tokio::time::timeout;
use crate::{
dispatch::{
affinity::{
AccountIdentityScope, AccountIdentityService, AccountScopedRequest,
SessionAffinityService,
},
errors::{ClientFailure, ResponseDispatchError, upstream_error_set_cookie_headers},
lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
fleet::{account::Account, pool::AccountLease},
models::service::ModelService,
telemetry::recorder::Recorder,
upstream::openai::protocol::events::TokenUsage,
upstream::openai::transport::{
CodexBackendClient, CodexBackendResponse, CodexBackendTransport, CodexClientError,
CodexRateLimitHeaderUpdates, CodexResponseMetadata, CodexTurnStateUpdate,
CodexUpstreamDiagnostics, WebSocketPoolDecision,
},
};
use crate::{
dispatch::{
controllers::cloudflare::CloudflareRecovery,
controllers::cyber_policy::{
ClassifiedCyberPolicyFailure, CyberPolicyController, CyberPolicyScope,
},
controllers::{
account_failure::AccountFailureController, history::HistoryController,
history::HistoryState, quota::QuotaController,
},
lifecycle::{
contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
AttemptRoutingFacts, FinalOutcome,
},
stream::StreamTerminal,
},
},
fleet::pool::AccountPoolService,
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
},
};
mod account_failure;
mod affinity;
pub(crate) mod cloudflare;
pub(crate) mod cyber_policy;
pub(in crate::dispatch) mod history;
mod quota;
pub(in crate::dispatch) mod telemetry;
mod usage;
const BEST_EFFORT_IO_TIMEOUT: Duration = Duration::from_millis(100);
pub(in crate::dispatch) struct CompleteExit<'a> {
pub request: &'a CodexResponsesRequest,
pub account_id: &'a str,
pub body: &'a str,
pub turn_state: Option<String>,
pub usage: Option<TokenUsage>,
pub image_generation_requested: bool,
pub completed: bool,
pub recorder: &'a Recorder,
pub models: &'a ModelService,
pub request_id: &'a str,
pub route: &'a str,
pub requested_model: &'a str,
pub display_model: &'a str,
pub started_at: Instant,
pub response_id: Option<&'a str>,
pub response: &'a CodexBackendResponse,
pub trace: &'a ResponseDispatchTrace,
pub attempt: &'a ResponseDispatchAttempt,
}
/// 已提交 stream 在 controller 洋葱退出阶段消费的稳定事实。
///
/// 该上下文不包含 live transport 执行器、账号 lease 或 shutdown 控制器;各 feature
/// owner 只能通过 `ControllerSet` 分发的窄 typed exit 读取所需字段。
pub(in crate::dispatch) struct StreamControllerContext {
pub account_id: String,
pub account_plan_type: Option<String>,
pub request_id: String,
pub route: String,
pub display_model: String,
pub requested_model: String,
pub request: CodexResponsesRequest,
pub transport: CodexBackendTransport,
pub set_cookie_headers: Vec<String>,
pub rate_limit_headers: Vec<(String, String)>,
pub rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
pub turn_state_update: Option<CodexTurnStateUpdate>,
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
pub turn_state: Option<String>,
pub diagnostics: CodexUpstreamDiagnostics,
pub response_metadata: CodexResponseMetadata,
pub transport_metrics: crate::upstream::openai::transport::CodexTransportMetrics,
pub connection_local_continuation: bool,
pub attempt: ResponseDispatchAttempt,
pub attempts: Vec<ResponseDispatchAttempt>,
pub started_at: Instant,
}
/// 单次 attempt 的原始上游失败上下文。
///
/// 生命周期只负责把事实交给控制器集合;Cloudflare 状态与遥测记录分别由各自
/// controller owner 消费,避免失败信息在最终 exhaustion 折叠时丢失。
pub(in crate::dispatch) struct AttemptUpstreamErrorContext<'a> {
pub request_id: &'a str,
pub account: &'a Account,
pub route: &'a str,
pub model: &'a str,
pub started_at: Instant,
pub stream: bool,
pub transport: crate::upstream::openai::transport::CodexBackendTransport,
pub request: &'a CodexResponsesRequest,
pub error: &'a CodexClientError,
pub trace: &'a ResponseDispatchTrace,
pub attempt: &'a ResponseDispatchAttempt,
}
/// 请求在产生最终 dispatch error 后交给 telemetry owner 的稳定事实。
pub(in crate::dispatch) struct DispatchErrorObservation<'a> {
pub request_id: &'a str,
pub client_api_key_id: Option<&'a str>,
pub account_id: Option<&'a str>,
pub route: &'a str,
pub model: &'a str,
pub started_at: Instant,
pub stream: bool,
pub compact: bool,
pub transport: Option<&'a str>,
}
/// commit 前 stream failure 的 telemetry 事实;client status 已包含在 typed error 中。
pub(in crate::dispatch) struct PrefetchedStreamFailureObservation<'a> {
pub request_id: &'a str,
pub account_id: &'a str,
pub route: &'a str,
pub model: &'a str,
pub requested_model: &'a str,
pub started_at: Instant,
pub transport: CodexBackendTransport,
pub request: &'a CodexResponsesRequest,
pub failure: &'a ResponsesSseFailure,
pub error: &'a ResponseDispatchError,
pub diagnostics: &'a CodexUpstreamDiagnostics,
pub rate_limit_headers: &'a [(String, String)],
pub prefetched: &'a [u8],
pub trace: &'a ResponseDispatchTrace,
pub attempt: &'a ResponseDispatchAttempt,
}
/// 所有 controller owner 共同消费的失败事实。
///
/// attempt 与 commit 后的 stream finalizer 都必须先归一为该类型,避免同一种
/// 上游失败因发生阶段不同而漏掉 feature 副作用。
#[derive(Clone, Copy)]
enum ControllerFailureFact<'a> {
Upstream(&'a UpstreamFailureFacts),
Response(&'a ResponsesSseFailure),
}
#[derive(Clone, Copy)]
enum FailureObservation<'a> {
Attempt(&'a AttemptObservation),
Stream {
account_id: &'a str,
failure: &'a ResponsesSseFailure,
},
}
impl<'a> FailureObservation<'a> {
fn fact(self) -> Option<ControllerFailureFact<'a>> {
match self {
Self::Attempt(observation) => ControllerFailureFact::from_attempt(observation),
Self::Stream { failure, .. } => Some(ControllerFailureFact::Response(failure)),
}
}
}
enum SharedFailureClassification<'a> {
CyberPolicy(ClassifiedCyberPolicyFailure<'a>),
AccountFailure(account_failure::ClassifiedFailure),
Quota(quota::ClassifiedQuotaFailure),
}
impl<'a> ControllerFailureFact<'a> {
fn from_attempt(observation: &'a AttemptObservation) -> Option<Self> {
match &observation.kind {
AttemptObservationKind::UpstreamFailure(facts) => Some(Self::Upstream(facts)),
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Failed(failure),
)
| AttemptObservationKind::StreamFailure(failure) => Some(Self::Response(failure)),
_ => None,
}
}
}
/// 单次 attempt 的唯一分类结果;副作用与 decision 必须共同消费这个值。
enum AttemptClassification<'a> {
Shared(SharedFailureClassification<'a>),
Cloudflare(cloudflare::CloudflareFailure),
History(AttemptDecision),
Default(AttemptDecision),
}
struct StreamFailureClassification<'a> {
failure: &'a ResponsesSseFailure,
owner: Option<SharedFailureClassification<'a>>,
}
impl StreamFailureClassification<'_> {
fn client_failure(&self) -> ClientFailure {
match &self.owner {
Some(SharedFailureClassification::CyberPolicy(classified)) => {
CyberPolicyController::client_failure(classified)
}
Some(SharedFailureClassification::AccountFailure(classified)) => {
AccountFailureController::client_failure(classified, self.failure)
}
Some(SharedFailureClassification::Quota(_)) => {
QuotaController::client_failure(self.failure)
}
None => default_response_client_failure(self.failure),
}
}
}
/// 请求级控制器集合。使用具体类型而非动态 middleware 链,保证调用顺序和所有权可见。
#[derive(Clone)]
pub(in crate::dispatch) struct ControllerSet {
cyber_policy: CyberPolicyController,
session_affinity: Arc<SessionAffinityService>,
account_pool: Arc<AccountPoolService>,
codex: Arc<CodexBackendClient>,
cloudflare: CloudflareRecovery,
recorder: Arc<Recorder>,
}
/// 控制器在单个请求中的不可变状态快照。
pub(in crate::dispatch) struct ControllerRequestScope {
cyber_policy: CyberPolicyScope,
history: HistoryState,
identity: AccountIdentityScope,
usage: usage::UsageScope,
}
/// 请求 enter 阶段的控制器结果。
pub(in crate::dispatch) struct ControllerEnter {
pub scope: ControllerRequestScope,
pub preferred_account_id: Option<String>,
pub excluded_account_ids: BTreeSet<String>,
}
pub(in crate::dispatch) struct AttemptAccountPreparationContext<'a> {
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
}
pub(in crate::dispatch) struct PreparedAttemptAccount {
pub lease: AccountLease,
pub cookie_header: Option<String>,
}
pub(in crate::dispatch) enum AttemptAccountPreparation {
Ready(Box<PreparedAttemptAccount>),
Rejected,
}
pub(in crate::dispatch) struct AttemptRoutePreparationContext<'a> {
pub account_identity: &'a AccountIdentityService,
pub account_id: &'a str,
}
pub(in crate::dispatch) enum AttemptRoutePreparation {
Ready(Box<AccountScopedRequest>),
Unavailable { message: String },
}
impl ControllerSet {
pub(in crate::dispatch) async fn prepare_attempt_account(
&self,
context: AttemptAccountPreparationContext<'_>,
acquired: AccountLease,
) -> AttemptAccountPreparation {
let usage_cookie = self
.cloudflare
.cookie_header_for_request(&acquired.account.id, "/codex/usage")
.await;
let acquired = match QuotaController::enter(
quota::QuotaEnterContext {
account_pool: &self.account_pool,
codex: context.codex,
cookie_header: usage_cookie.as_deref(),
account_identity: context.account_identity,
request_id: context.request_id,
},
acquired,
)
.await
{
quota::QuotaEnterOutcome::Ready(acquired) => *acquired,
quota::QuotaEnterOutcome::LimitReached => {
return AttemptAccountPreparation::Rejected;
}
};
let cookie_header = self
.cloudflare
.cookie_header_for_request(&acquired.account.id, "/codex/responses")
.await;
AttemptAccountPreparation::Ready(Box::new(PreparedAttemptAccount {
lease: acquired,
cookie_header,
}))
}
pub(in crate::dispatch) fn prepare_attempt_route(
&self,
scope: &mut ControllerRequestScope,
request: &CodexResponsesRequest,
context: AttemptRoutePreparationContext<'_>,
) -> AttemptRoutePreparation {
let request = match HistoryController::prepare_attempt(
&mut scope.history,
request,
context.account_id,
) {
Ok(request) => request,
Err(message) => return AttemptRoutePreparation::Unavailable { message },
};
AttemptRoutePreparation::Ready(Box::new(context.account_identity.scope_request(
&mut scope.identity,
request,
context.account_id,
)))
}
pub(in crate::dispatch) fn prepare_same_account_retry(
&self,
scope: &mut ControllerRequestScope,
account_id: &str,
) -> bool {
HistoryController::prepare_same_account_retry(&mut scope.history, account_id)
}
pub(in crate::dispatch) fn attempt_routing_facts(
&self,
scope: &ControllerRequestScope,
account_id: Option<&str>,
) -> AttemptRoutingFacts {
HistoryController::routing_facts(&scope.history, account_id)
}
pub(in crate::dispatch) async fn observe_complete_upstream(
&self,
account: &Account,
set_cookie_headers: &[String],
rate_limit_headers: &[(String, String)],
) {
let _ = best_effort_controller_io(
"cookie.complete",
self.cloudflare
.capture_set_cookie_headers(&account.id, set_cookie_headers),
)
.await;
let _ = best_effort_controller_io(
"quota.complete_headers",
self.account_pool
.sync_passive_rate_limit_headers(account, rate_limit_headers),
)
.await;
}
pub(in crate::dispatch) async fn observe_upstream_error(
&self,
context: AttemptUpstreamErrorContext<'_>,
) {
let _ = best_effort_controller_io(
"cookie.attempt_error",
self.cloudflare.capture_set_cookie_headers(
&context.account.id,
upstream_error_set_cookie_headers(context.error),
),
)
.await;
let _ = best_effort_controller_io(
"telemetry.attempt_error",
telemetry::TelemetryController::observe_upstream_error(&self.recorder, &context),
)
.await;
}
pub(in crate::dispatch) async fn observe_dispatch_error(
&self,
observation: DispatchErrorObservation<'_>,
error: &ResponseDispatchError,
) {
let _ = best_effort_controller_io(
"telemetry.dispatch_error",
telemetry::TelemetryController::observe_dispatch_error(
&self.recorder,
observation,
error,
),
)
.await;
}
pub(in crate::dispatch) async fn observe_prefetched_stream_failure(
&self,
observation: PrefetchedStreamFailureObservation<'_>,
) {
let _ = best_effort_controller_io(
"telemetry.prefetched_failure",
telemetry::TelemetryController::observe_prefetched_stream_failure(
&self.recorder,
observation,
),
)
.await;
}
pub(in crate::dispatch) async fn finalize_complete(
&self,
scope: ControllerRequestScope,
exit: CompleteExit<'_>,
outcome: FinalOutcome,
) {
let successful = matches!(outcome, FinalOutcome::Completed | FinalOutcome::Incomplete);
if successful {
let _ = best_effort_controller_io(
"cloudflare.complete",
cloudflare::CloudflareController::leave_complete(&self.cloudflare, exit.account_id),
)
.await;
}
if successful {
let _ = best_effort_controller_io(
"usage.complete",
usage::UsageController::leave_complete(
&self.account_pool,
exit.account_id,
exit.usage,
exit.image_generation_requested,
),
)
.await;
}
if exit.completed {
let conversation_id = HistoryController::conversation_id(&scope.history, exit.request);
let continuation_scope = HistoryController::continuation_scope(
exit.request,
exit.response.transport,
exit.response.connection_local_continuation,
);
let _ = best_effort_controller_io(
"affinity.complete",
affinity::AffinityController::leave_complete(affinity::ResponseExit {
affinity: &self.session_affinity,
conversation_id,
request: exit.request,
account_id: exit.account_id,
body: exit.body,
turn_state: exit.turn_state.clone(),
usage: exit.usage,
continuation_scope,
}),
)
.await;
}
if successful {
let _ = best_effort_controller_io(
"telemetry.complete",
telemetry::TelemetryController::leave_complete(&exit),
)
.await;
}
exit.models
.observe_models_etag(exit.response.response_metadata.models_etag.as_deref());
}
pub(in crate::dispatch) fn new(
session_affinity: Arc<crate::dispatch::affinity::SessionAffinityService>,
account_pool: Arc<AccountPoolService>,
codex: Arc<CodexBackendClient>,
cloudflare: CloudflareRecovery,
recorder: Arc<Recorder>,
) -> Self {
Self {
cyber_policy: CyberPolicyController::new(Arc::clone(&session_affinity)),
session_affinity,
account_pool,
codex,
cloudflare,
recorder,
}
}
pub(in crate::dispatch) async fn enter(
&self,
request: &CodexResponsesRequest,
now: DateTime<Utc>,
) -> ControllerEnter {
let usage = usage::UsageController::enter(request);
let (history, cyber_policy, affinity_preferred_account_id) = tokio::join!(
HistoryController::enter(&self.session_affinity, request),
self.cyber_policy.prepare(request),
affinity::AffinityController::preferred_account_id(
&self.session_affinity,
request,
now,
),
);
let preferred_account_id = HistoryController::preferred_account_id(&history)
.map(str::to_owned)
.or(affinity_preferred_account_id);
let excluded_account_ids = cyber_policy.excluded_account_ids();
ControllerEnter {
scope: ControllerRequestScope {
cyber_policy,
history,
identity: AccountIdentityScope::new(preferred_account_id.clone()),
usage,
},
preferred_account_id,
excluded_account_ids,
}
}
async fn finalize_stream_success(
&self,
scope: &ControllerRequestScope,
terminal: &StreamTerminal,
) {
match terminal {
StreamTerminal::Completed { response } | StreamTerminal::Incomplete { response }
if response.is_object() =>
{
self.cyber_policy.observe_success(&scope.cyber_policy).await
}
StreamTerminal::Completed { .. }
| StreamTerminal::Incomplete { .. }
| StreamTerminal::Failed { .. }
| StreamTerminal::UpstreamClosed
| StreamTerminal::UpstreamError { .. }
| StreamTerminal::ProtocolError { .. }
| StreamTerminal::CaptureLimitExceeded
| StreamTerminal::Cancelled
| StreamTerminal::DownstreamClosed
| StreamTerminal::Shutdown => {}
}
}
fn classify_stream_failure<'a>(
&self,
account_id: &'a str,
failure: &'a ResponsesSseFailure,
) -> StreamFailureClassification<'a> {
StreamFailureClassification {
failure,
owner: Self::classify_shared_failure(FailureObservation::Stream {
account_id,
failure,
}),
}
}
async fn apply_stream_account_state(&self, classification: &StreamFailureClassification<'_>) {
match classification.owner.as_ref() {
Some(SharedFailureClassification::AccountFailure(classified)) => {
AccountFailureController::apply_effect(&self.account_pool, &self.codex, classified)
.await;
}
Some(SharedFailureClassification::Quota(classified)) => {
QuotaController::apply_effect(&self.account_pool, &self.codex, classified).await;
}
Some(SharedFailureClassification::CyberPolicy(_)) | None => {}
}
}
async fn apply_stream_policy_failure(
&self,
scope: &ControllerRequestScope,
account_id: &str,
classification: &StreamFailureClassification<'_>,
) {
if matches!(
classification.owner,
Some(SharedFailureClassification::CyberPolicy(_))
) {
self.cyber_policy
.exclude_account(&scope.cyber_policy, account_id)
.await;
}
}
/// 按固定的反向 leave 顺序收敛 stream 终态;生命周期只调用这一 typed exit。
pub(in crate::dispatch) async fn leave_stream(
&self,
scope: ControllerRequestScope,
context: StreamControllerContext,
summary: &crate::dispatch::lifecycle::stream::StreamSummary,
body: &str,
) {
let stream_failure = match &summary.terminal {
StreamTerminal::Failed { failure } => {
Some(self.classify_stream_failure(&context.account_id, failure))
}
_ => None,
};
let client_failure_status = stream_failure
.as_ref()
.map(|failure| i64::from(failure.client_failure().status_code()));
let outcome = match &summary.terminal {
StreamTerminal::Completed { .. } => FinalOutcome::Completed,
StreamTerminal::Incomplete { .. } => FinalOutcome::Incomplete,
StreamTerminal::Failed { .. }
| StreamTerminal::UpstreamClosed
| StreamTerminal::UpstreamError { .. }
| StreamTerminal::ProtocolError { .. }
| StreamTerminal::CaptureLimitExceeded => FinalOutcome::Failed,
StreamTerminal::Cancelled | StreamTerminal::DownstreamClosed => FinalOutcome::Cancelled,
StreamTerminal::Shutdown => FinalOutcome::Shutdown,
};
if !matches!(outcome, FinalOutcome::Cancelled | FinalOutcome::Shutdown)
&& let Some(classification) = stream_failure.as_ref()
{
self.apply_stream_account_state(classification).await;
}
let _ = best_effort_controller_io(
"cloudflare",
cloudflare::CloudflareController::leave_stream(cloudflare::StreamExit {
recovery: &self.cloudflare,
account_id: &context.account_id,
set_cookie_headers: &context.set_cookie_headers,
terminal: &summary.terminal,
}),
)
.await;
let fallback = usage::StreamUsageExit {
rate_limit_headers: context.rate_limit_headers.clone(),
turn_state: context.turn_state.clone(),
};
let mut usage_exit = best_effort_controller_io(
"usage",
usage::UsageController::leave_stream(usage::StreamExit {
account_pool: &self.account_pool,
account_id: &context.account_id,
account_plan_type: context.account_plan_type.as_deref(),
rate_limit_headers: &context.rate_limit_headers,
rate_limit_header_updates: context.rate_limit_header_updates.as_ref(),
turn_state_update: context.turn_state_update.as_ref(),
turn_state: context.turn_state.as_deref(),
usage: summary.usage,
}),
)
.await
.unwrap_or(fallback);
let conversation_id = HistoryController::conversation_id(&scope.history, &context.request);
let continuation_scope = HistoryController::continuation_scope(
&context.request,
context.transport,
context.connection_local_continuation,
);
affinity::AffinityController::leave_stream(affinity::StreamExit {
response: affinity::ResponseExit {
affinity: &self.session_affinity,
conversation_id,
request: &context.request,
account_id: &context.account_id,
body,
turn_state: usage_exit.turn_state.take(),
usage: summary.usage,
continuation_scope,
},
completed: matches!(summary.terminal, StreamTerminal::Completed { .. }),
})
.await;
let _ = best_effort_controller_io(
"telemetry",
telemetry::TelemetryController::leave_stream(telemetry::StreamExit {
context: telemetry::StreamContext {
recorder: &self.recorder,
account_id: &context.account_id,
request_id: &context.request_id,
route: &context.route,
display_model: &context.display_model,
requested_model: &context.requested_model,
request: &context.request,
transport: context.transport,
websocket_pool_decision: context.websocket_pool_decision,
diagnostics: &context.diagnostics,
response_metadata: &context.response_metadata,
transport_metrics: &context.transport_metrics,
attempt: &context.attempt,
attempts: &context.attempts,
started_at: context.started_at,
},
summary,
rate_limit_headers: &usage_exit.rate_limit_headers,
failure_status: client_failure_status,
body,
}),
)
.await;
if !matches!(outcome, FinalOutcome::Cancelled | FinalOutcome::Shutdown) {
if let Some(classification) = stream_failure.as_ref() {
self.apply_stream_policy_failure(&scope, &context.account_id, classification)
.await;
} else {
self.finalize_stream_success(&scope, &summary.terminal)
.await;
}
}
}
/// 对一次 observation 只分类一次,并保证 feature effects 先于 decision 返回。
pub(in crate::dispatch) async fn handle_attempt(
&self,
scope: &ControllerRequestScope,
observation: &AttemptObservation,
) -> AttemptDecision {
let classification = self.classify_attempt(observation);
let _ = best_effort_controller_io(
"usage.attempt",
usage::UsageController::observe_attempt(&self.account_pool, &scope.usage, observation),
)
.await;
if matches!(
observation.kind,
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Completed
| crate::dispatch::lifecycle::contract::CompleteResponseFacts::Incomplete,
)
) {
self.cyber_policy.observe_success(&scope.cyber_policy).await;
}
match &classification {
AttemptClassification::Shared(shared) => match shared {
SharedFailureClassification::CyberPolicy(_) => {
if let Some(account_id) = observation
.account
.as_ref()
.map(|account| account.id.as_str())
{
self.cyber_policy
.exclude_account(&scope.cyber_policy, account_id)
.await;
}
}
SharedFailureClassification::AccountFailure(classified) => {
AccountFailureController::apply_effect(
&self.account_pool,
&self.codex,
classified,
)
.await;
}
SharedFailureClassification::Quota(classified) => {
QuotaController::apply_effect(&self.account_pool, &self.codex, classified)
.await;
}
},
AttemptClassification::Cloudflare(failure) => {
let _ = best_effort_controller_io(
"cloudflare.attempt",
cloudflare::CloudflareController::apply_effect(
&self.cloudflare,
&self.account_pool,
failure,
),
)
.await;
}
AttemptClassification::History(_) | AttemptClassification::Default(_) => {}
}
self.attempt_decision(observation, classification)
}
fn classify_attempt<'a>(
&self,
observation: &'a AttemptObservation,
) -> AttemptClassification<'a> {
let shared = Self::classify_shared_failure(FailureObservation::Attempt(observation));
let shared = match shared {
Some(shared @ SharedFailureClassification::CyberPolicy(_)) => {
return AttemptClassification::Shared(shared);
}
shared => shared,
};
if let Some(classified) = cloudflare::CloudflareController::classify(observation) {
return AttemptClassification::Cloudflare(classified);
}
if let Some(decision) = HistoryController::decide(observation) {
return AttemptClassification::History(decision);
}
if let Some(shared) = shared {
return AttemptClassification::Shared(shared);
}
AttemptClassification::Default(Self::default_attempt_decision(observation))
}
/// complete、prefetch 与 committed stream 共用这一份 feature owner 优先级。
fn classify_shared_failure<'a>(
observation: FailureObservation<'a>,
) -> Option<SharedFailureClassification<'a>> {
let fact = observation.fact();
if let Some(classified) = fact.and_then(CyberPolicyController::classify) {
return Some(SharedFailureClassification::CyberPolicy(classified));
}
let account_failure = match observation {
FailureObservation::Attempt(observation) => {
AccountFailureController::classify(observation)
}
FailureObservation::Stream { account_id, .. } => {
fact.and_then(|fact| AccountFailureController::classify_failure(account_id, fact))
}
};
if let Some(classified) = account_failure {
return Some(SharedFailureClassification::AccountFailure(classified));
}
let quota = match observation {
FailureObservation::Attempt(observation) => QuotaController::classify(observation),
FailureObservation::Stream { account_id, .. } => {
fact.and_then(|fact| QuotaController::classify_failure(account_id, fact))
}
};
quota.map(SharedFailureClassification::Quota)
}
fn attempt_decision(
&self,
observation: &AttemptObservation,
classification: AttemptClassification<'_>,
) -> AttemptDecision {
match classification {
AttemptClassification::Shared(shared) => match shared {
SharedFailureClassification::CyberPolicy(classified) => {
CyberPolicyController::decision(observation, classified)
}
SharedFailureClassification::AccountFailure(classified) => {
let client_failure = response_failure(observation).map(|failure| {
AccountFailureController::client_failure(&classified, failure)
});
with_response_client_failure(
AccountFailureController::decision(observation, classified),
client_failure,
)
}
SharedFailureClassification::Quota(classified) => {
let client_failure =
response_failure(observation).map(QuotaController::client_failure);
with_response_client_failure(
QuotaController::decision(observation, classified),
client_failure,
)
}
},
AttemptClassification::Cloudflare(classified) => {
cloudflare::CloudflareController::decision(observation, classified)
}
AttemptClassification::History(decision) => with_response_client_failure(
decision,
response_failure(observation)
.map(|failure| ClientFailure::new(failure.clone(), 400, false)),
),
AttemptClassification::Default(decision) => decision,
}
}
fn default_attempt_decision(observation: &AttemptObservation) -> AttemptDecision {
match &observation.kind {
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Completed
| crate::dispatch::lifecycle::contract::CompleteResponseFacts::Incomplete,
) => AttemptDecision::Accept,
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Failed(failure),
) => AttemptDecision::Return(AttemptReturnKind::Failed(
default_response_client_failure(failure),
)),
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::MissingCompleted
| crate::dispatch::lifecycle::contract::CompleteResponseFacts::Empty,
) => AttemptDecision::Return(AttemptReturnKind::Observed),
AttemptObservationKind::StreamFailure(failure) => AttemptDecision::Return(
AttemptReturnKind::Failed(default_response_client_failure(failure)),
),
AttemptObservationKind::NoCandidate { .. } => {
AttemptDecision::Return(AttemptReturnKind::Observed)
}
AttemptObservationKind::CandidatePreparationRejected => {
unreachable!("candidate preparation controller owns it")
}
AttemptObservationKind::RoutePreparationRejected { message } => {
AttemptDecision::Return(AttemptReturnKind::RouteUnavailable {
message: message.clone(),
})
}
AttemptObservationKind::PinnedCandidateUnavailable { .. } => {
unreachable!("route controller owns pinned-candidate acquisition failures")
}
AttemptObservationKind::ProtocolFailure { .. } => {
AttemptDecision::Return(AttemptReturnKind::Observed)
}
AttemptObservationKind::UpstreamFailure(facts) => {
if observation.routing.can_retry_next_candidate
&& AccountFailureController::is_retryable_transport(observation)
{
return AttemptDecision::RetryNextCandidate {
exhaustion: observation.account.as_ref().map(|account| {
crate::dispatch::failure::exhaustion::AccountExhaustionRecord::new(
account.id.clone(),
crate::dispatch::failure::exhaustion::ExhaustedAccountKind::UpstreamUnavailable,
facts.message.clone(),
)
}),
on_exhaustion: None,
};
}
AttemptDecision::Return(AttemptReturnKind::Observed)
}
}
}
}
fn response_failure(observation: &AttemptObservation) -> Option<&ResponsesSseFailure> {
match &observation.kind {
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Failed(failure),
)
| AttemptObservationKind::StreamFailure(failure) => Some(failure),
_ => None,
}
}
fn with_response_client_failure(
decision: AttemptDecision,
failure: Option<ClientFailure>,
) -> AttemptDecision {
match (decision, failure) {
(AttemptDecision::Return(AttemptReturnKind::Observed), Some(failure)) => {
AttemptDecision::Return(AttemptReturnKind::Failed(failure))
}
(decision, _) => decision,
}
}
fn default_response_client_failure(failure: &ResponsesSseFailure) -> ClientFailure {
AccountFailureController::unowned_client_failure(failure)
.unwrap_or_else(|| ClientFailure::new(failure.clone(), 502, false))
}
async fn best_effort_controller_io<T>(
controller: &'static str,
operation: impl Future<Output = T>,
) -> Option<T> {
match timeout(BEST_EFFORT_IO_TIMEOUT, operation).await {
Ok(output) => Some(output),
Err(_) => {
tracing::warn!(
controller,
timeout_ms = BEST_EFFORT_IO_TIMEOUT.as_millis(),
"Timed out running best-effort controller I/O"
);
None
}
}
}
+247
View File
@@ -0,0 +1,247 @@
//! 请求前配额验证规则 owner。
use serde_json::Value;
use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
},
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_response_failure, classify_upstream_failure,
},
pool::{AccountLease, AccountPoolService},
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
},
},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexRequestContext},
},
};
const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached";
pub(super) enum QuotaEnterOutcome {
Ready(Box<AccountLease>),
LimitReached,
}
pub(super) struct QuotaEnterContext<'a> {
pub account_pool: &'a AccountPoolService,
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
pub cookie_header: Option<&'a str>,
}
pub(super) struct QuotaController;
impl QuotaController {
pub(super) async fn enter(
context: QuotaEnterContext<'_>,
acquired: AccountLease,
) -> QuotaEnterOutcome {
if !acquired.account.quota_verify_required {
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
let account_id = acquired.account.id.clone();
let identity = context.account_identity.scope_auxiliary(&account_id);
let usage = context
.codex
.fetch_usage(CodexRequestContext {
access_token: &acquired.account.access_token,
account_id: acquired.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: None,
parent_thread_id: None,
cookie_header: context.cookie_header,
installation_id: Some(&identity.installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
})
.await;
let raw = match usage {
Ok(raw) => raw,
Err(error) => {
tracing::warn!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "upstream_error",
retry_with_another_account = false,
error = %error,
"Failed to verify stale quota state before upstream request"
);
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
};
let quota = quota_from_usage(&raw);
context
.account_pool
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
acquired.release_without_usage().await;
tracing::info!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "limit_reached",
retry_with_another_account = true,
"Quota verification reported exhausted account before upstream request"
);
return QuotaEnterOutcome::LimitReached;
}
QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, &quota)))
}
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
if matches!(
observation.kind,
AttemptObservationKind::CandidatePreparationRejected
) {
return observation
.account
.as_ref()
.map(|account| ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account.id.clone(),
ExhaustedAccountKind::RateLimited,
LIMIT_REACHED_MESSAGE,
),
effect: None,
});
}
classify(observation)
}
pub(super) fn classify_failure(
account_id: &str,
failure: ControllerFailureFact<'_>,
) -> Option<ClassifiedQuotaFailure> {
classify_failure(account_id, failure)
}
pub(super) fn decision(
observation: &AttemptObservation,
classified: ClassifiedQuotaFailure,
) -> AttemptDecision {
if observation.routing.can_retry_next_candidate {
return AttemptDecision::RetryNextCandidate {
exhaustion: Some(classified.exhaustion),
on_exhaustion: None,
};
}
AttemptDecision::Return(AttemptReturnKind::Observed)
}
pub(super) async fn apply_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
classified: &ClassifiedQuotaFailure,
) {
let Some(account_id) = classified.exhaustion.account_id.as_deref() else {
return;
};
if let Some(effect) = &classified.effect {
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
}
}
pub(super) fn client_failure(failure: &ResponsesSseFailure) -> ClientFailure {
ClientFailure::new(failure.clone(), 429, false)
}
}
pub(super) struct ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord,
effect: Option<AccountStateEffect>,
}
fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
let account_id = observation.account.as_ref()?.id.as_str();
ControllerFailureFact::from_attempt(observation)
.and_then(|failure| classify_failure(account_id, failure))
}
fn classify_failure(
account_id: &str,
failure: ControllerFailureFact<'_>,
) -> Option<ClassifiedQuotaFailure> {
match failure {
ControllerFailureFact::Upstream(facts) => classify_upstream(account_id, facts),
ControllerFailureFact::Response(failure) => classify_stream(account_id, failure),
}
}
fn classify_upstream(
account_id: &str,
facts: &UpstreamFailureFacts,
) -> Option<ClassifiedQuotaFailure> {
classified_quota_failure(account_id, &facts.body, classify_upstream_failure(facts)?)
}
fn classify_stream(
account_id: &str,
failure: &ResponsesSseFailure,
) -> Option<ClassifiedQuotaFailure> {
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
classified_quota_failure(account_id, &body, classify_response_failure(failure)?)
}
fn classified_quota_failure(
account_id: &str,
message: &str,
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
) -> Option<ClassifiedQuotaFailure> {
let kind = match classified.kind {
AccountFailureKind::RateLimited => ExhaustedAccountKind::RateLimited,
AccountFailureKind::QuotaExhausted => ExhaustedAccountKind::QuotaExhausted,
AccountFailureKind::ModelUnsupported
| AccountFailureKind::Expired
| AccountFailureKind::Disabled
| AccountFailureKind::Banned => return None,
};
Some(ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(account_id, kind, message),
effect: classified.effect,
})
}
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
let limit_reached = quota_snapshot_limit_reached(quota);
acquired.account.quota_verify_required = false;
acquired.account.quota_limit_reached = limit_reached;
acquired.account.quota_cooldown_until = limit_reached
.then_some(quota_snapshot_reset_at(quota))
.flatten();
if let Some(reset_at) = quota_snapshot_reset_at(quota) {
acquired.account.window_reset_at = Some(reset_at);
if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) {
acquired.account.limit_window_seconds = Some(limit_window_seconds);
}
}
acquired
}
@@ -27,17 +27,16 @@ use crate::{
},
};
use super::{
use crate::dispatch::{
errors::{
ResponseDispatchError, dispatch_error_metadata, enrich_response_dispatch_error_metadata,
},
stream::{
live::{LiveResponseStreamContext, latest_response_id},
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
failure::sse::stream_failure_metadata,
lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
};
use super::StreamContext;
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
pub(super) recorder: &'a Recorder,
pub(super) request_id: &'a str,
@@ -65,6 +64,7 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> {
pub(super) transport: CodexBackendTransport,
pub(super) request: &'a CodexResponsesRequest,
pub(super) failure: &'a ResponsesSseFailure,
pub(super) status_code: i64,
pub(super) diagnostics: &'a CodexUpstreamDiagnostics,
pub(super) rate_limit_headers: &'a [(String, String)],
pub(super) prefetched: &'a [u8],
@@ -86,14 +86,6 @@ pub(super) struct ResponseDispatchErrorEventRecord<'a> {
pub(super) error: &'a ResponseDispatchError,
}
pub(super) struct ResponseDispatchErrorDetails<'a> {
pub(super) client_api_key_id: Option<&'a str>,
pub(super) account_id: Option<&'a str>,
pub(super) stream: bool,
pub(super) compact: bool,
pub(super) transport: Option<&'a str>,
}
pub(super) async fn record_response_dispatch_error_event(
record: ResponseDispatchErrorEventRecord<'_>,
) {
@@ -181,14 +173,14 @@ pub(super) async fn record_response_upstream_error_event(
enrich_response_request_semantics(&mut metadata, record.request);
event.metadata = metadata;
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "failed to record upstream error event");
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record upstream error event");
}
}
pub(super) async fn record_prefetched_response_stream_failure_event(
record: ResponseStreamFailureEventRecord<'_>,
) {
let event_status_code = status_code_for_stream_failure(record.failure);
let event_status_code = record.status_code;
let mut metadata = stream_failure_metadata(record.failure, None);
if let Some(object) = metadata.as_object_mut() {
if record.transport == CodexBackendTransport::WebSocket {
@@ -251,7 +243,7 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
}
event.metadata = metadata;
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "failed to record prefetched stream error event");
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record prefetched stream error event");
}
}
@@ -265,8 +257,16 @@ fn ensure_stream_metadata_flag(metadata: &mut Value) {
.or_insert(Value::Bool(true));
}
fn insert_optional_i64(object: &mut serde_json::Map<String, Value>, key: &str, value: Option<i64>) {
if let Some(value) = value {
object
.entry(key.to_string())
.or_insert_with(|| Value::Number(value.into()));
}
}
fn enrich_live_response_stream_metadata(
context: &LiveResponseStreamContext,
context: &StreamContext<'_>,
rate_limit_headers: &[(String, String)],
metadata: &mut Value,
status_code: i64,
@@ -288,15 +288,45 @@ fn enrich_live_response_stream_metadata(
.entry("websocketPool".to_string())
.or_insert_with(|| decision.metadata_value());
}
if let Some(decision) = context.transport_metrics.decision {
object
.entry("transportDecision".to_string())
.or_insert_with(|| Value::String(decision.as_str().to_string()));
}
insert_optional_i64(
object,
"wsConnectMs",
context.transport_metrics.ws_connect_ms,
);
insert_optional_i64(
object,
"transportDecisionWaitMs",
context.transport_metrics.transport_decision_wait_ms,
);
insert_optional_i64(
object,
"upstreamHeadersMs",
context.transport_metrics.upstream_headers_ms,
);
insert_optional_i64(
object,
"firstEventMs",
context.transport_metrics.first_event_ms,
);
if let Some(http_version) = context.transport_metrics.http_version.as_deref() {
object
.entry("httpVersion".to_string())
.or_insert_with(|| Value::String(http_version.to_string()));
}
insert_response_status_metadata_object(
object,
status_code,
200,
diagnostics_status_code(&context.diagnostics),
diagnostics_status_code(context.diagnostics),
);
insert_response_request_summary_object(object, &context.request, context.transport);
insert_upstream_diagnostics_metadata(object, &context.diagnostics);
insert_response_trace_metadata_object(object, &context.attempts, Some(&context.attempt));
insert_response_request_summary_object(object, context.request, context.transport);
insert_upstream_diagnostics_metadata(object, context.diagnostics);
insert_response_trace_metadata_object(object, context.attempts, Some(context.attempt));
if context.recorder.captures_body() {
object
.entry("requestBody".to_string())
@@ -476,7 +506,7 @@ fn upstream_failure_client_status_code(upstream_status: i64) -> i64 {
}
pub(super) async fn record_live_response_stream_event(
context: &LiveResponseStreamContext,
context: &StreamContext<'_>,
status_code: i64,
level: UsageRecordLevel,
message: &str,
@@ -488,7 +518,7 @@ pub(super) async fn record_live_response_stream_event(
.response_metadata
.effective_model
.as_deref()
.unwrap_or(&context.display_model);
.unwrap_or(context.display_model);
ensure_stream_metadata_flag(&mut metadata);
if let Some(object) = metadata.as_object_mut() {
object.insert(
@@ -508,7 +538,7 @@ pub(super) async fn record_live_response_stream_event(
Value::Bool(context.response_metadata.reasoning_included),
);
}
enrich_event_route_metadata(&mut metadata, &context.route);
enrich_event_route_metadata(&mut metadata, context.route);
enrich_live_response_stream_metadata(
context,
rate_limit_headers,
@@ -516,31 +546,34 @@ pub(super) async fn record_live_response_stream_event(
status_code,
body,
);
log_live_response_stream_finalized(context, status_code, level, message, &metadata, body);
if !is_success_usage_event(level, status_code) {
ensure_stream_failure_class(&mut metadata);
}
log_live_response_stream_finalized(context, status_code, level, message, &metadata);
enrich_usage_record_identity(
&mut metadata,
Some(&context.requested_model),
Some(context.requested_model),
effective_model,
context.client_ip.as_deref(),
context.request.client_ip.as_deref(),
context.request.client_user_agent.as_deref(),
reasoning_effort_from_request(&context.request),
reasoning_effort_from_request(context.request),
context.request.service_tier(),
);
enrich_response_request_semantics(&mut metadata, &context.request);
enrich_response_request_semantics(&mut metadata, context.request);
if is_success_usage_event(level, status_code) {
crate::telemetry::recorder::record_response_event(
crate::telemetry::usage::types::ResponseUsageRecord {
recorder: &context.recorder,
request_id: &context.request_id,
recorder: context.recorder,
request_id: context.request_id,
client_api_key_id: context.request.client_api_key_id.as_deref(),
account_id: &context.account_id,
route: &context.route,
account_id: context.account_id,
route: context.route,
model: effective_model,
requested_model: Some(&context.requested_model),
client_ip: context.client_ip.as_deref(),
requested_model: Some(context.requested_model),
client_ip: context.request.client_ip.as_deref(),
client_user_agent: context.request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(&context.request),
reasoning_effort: reasoning_effort_from_request(context.request),
service_tier: context.request.service_tier(),
started_at: context.started_at,
status_code,
@@ -551,18 +584,18 @@ pub(super) async fn record_live_response_stream_event(
)
.await;
} else {
let mut event = OpsErrorLog::new(response_event_kind(&context.route), message);
event.request_id = Some(context.request_id.clone());
let mut event = OpsErrorLog::new(response_event_kind(context.route), message);
event.request_id = Some(context.request_id.to_string());
event.client_api_key_id = context.request.client_api_key_id.clone();
event.provider = Some("openai".to_string());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.account_id = Some(context.account_id.to_string());
event.route = Some(context.route.to_string());
event.model = Some(effective_model.to_string());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
event.metadata = metadata;
if let Err(error) = context.recorder.record_error(event).await {
tracing::error!(account_id = %context.account_id, error = %error, "failed to record live response stream error event");
tracing::error!(account_id = %context.account_id, error = %error, "Failed to record live response stream error event");
}
}
}
@@ -571,24 +604,31 @@ fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
level != UsageRecordLevel::Error && status_code < 400
}
fn ensure_stream_failure_class(metadata: &mut Value) {
let failure_class = ["failureClass", "upstreamCode", "terminal"]
.into_iter()
.find_map(|field| metadata_string_field(metadata, field))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string);
let (Some(failure_class), Some(object)) = (failure_class, metadata.as_object_mut()) else {
return;
};
object.insert("failureClass".to_string(), Value::String(failure_class));
}
fn log_live_response_stream_finalized(
context: &LiveResponseStreamContext,
context: &StreamContext<'_>,
status_code: i64,
level: UsageRecordLevel,
message: &str,
metadata: &Value,
body: &str,
) {
let response_id = metadata_string_field(metadata, "responseId")
.map(ToString::to_string)
.or_else(|| latest_response_id(body));
let response_id = metadata_string_field(metadata, "responseId").map(ToString::to_string);
let first_token_ms = metadata.get("firstTokenMs").and_then(Value::as_i64);
let websocket_pool_kind = context
.websocket_pool_decision
.map(WebSocketPoolDecision::kind);
let websocket_pool_reason = context
.websocket_pool_decision
.and_then(WebSocketPoolDecision::reason);
let completed = metadata.get("completed").and_then(Value::as_bool);
let failed = metadata.get("failed").and_then(Value::as_bool);
let upstream_code = metadata_string_field(metadata, "upstreamCode");
@@ -609,7 +649,6 @@ fn log_live_response_stream_finalized(
event_message = %message,
transport = %backend_transport_name(context.transport),
websocket_pool_kind = ?websocket_pool_kind,
websocket_pool_reason = ?websocket_pool_reason,
response_id = response_id.as_deref().unwrap_or(""),
first_token_ms = ?first_token_ms,
latency_ms = elapsed_millis_i64(context.started_at),
@@ -619,7 +658,7 @@ fn log_live_response_stream_finalized(
failure_class = ?failure_class,
failure_source = ?failure_source,
failure_detail = ?failure_detail,
"live response stream finalized"
"Live response stream finalized"
);
};
}
@@ -644,25 +683,6 @@ fn metadata_string_field<'a>(metadata: &'a Value, field: &str) -> Option<&'a str
metadata.get(field).and_then(Value::as_str)
}
pub(super) async fn live_response_rate_limit_headers(
context: &LiveResponseStreamContext,
) -> Vec<(String, String)> {
let mut headers = context.rate_limit_headers.clone();
if let Some(updates) = &context.rate_limit_header_updates {
headers.extend(updates.lock().await.iter().cloned());
}
headers
}
pub(super) async fn live_response_turn_state(
context: &LiveResponseStreamContext,
) -> Option<String> {
if let Some(update) = &context.turn_state_update {
return update.lock().await.clone();
}
context.turn_state.clone()
}
pub(super) fn insert_first_token_ms(metadata: &mut Value, first_token_ms: Option<i64>) {
let Some(first_token_ms) = first_token_ms else {
return;
@@ -0,0 +1,359 @@
//! 请求终态遥测的唯一 feature owner。
mod events;
use serde_json::Value;
use crate::{
dispatch::{
controllers::telemetry::events::{
ResponseUpstreamErrorEventRecord, insert_first_token_ms,
insert_response_status_metadata, insert_response_trace_metadata,
insert_response_upstream_diagnostics, insert_websocket_pool_decision,
record_live_response_stream_event, record_response_upstream_error_event,
},
controllers::{
AttemptUpstreamErrorContext, CompleteExit, DispatchErrorObservation,
PrefetchedStreamFailureObservation,
},
errors::{ResponseDispatchError, backend_transport_name},
failure::sse::{
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, stream_failure_metadata,
stream_failure_source, synthetic_stream_disconnected_detail,
},
lifecycle::{
stream::{StreamSummary, StreamTerminal},
trace::ResponseDispatchAttempt,
},
},
infra::time::elapsed_millis_i64,
telemetry::{
recorder::{
Recorder, enrich_response_request_semantics, reasoning_effort_from_request,
record_response_event,
},
usage::types::{ResponseUsageRecord, UsageRecordLevel},
},
upstream::openai::{
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
transport::{
CodexBackendTransport, CodexResponseMetadata, CodexTransportMetrics,
CodexUpstreamDiagnostics, WebSocketPoolDecision,
},
},
};
pub(super) struct TelemetryController;
pub(super) struct StreamContext<'a> {
pub recorder: &'a Recorder,
pub account_id: &'a str,
pub request_id: &'a str,
pub route: &'a str,
pub display_model: &'a str,
pub requested_model: &'a str,
pub request: &'a CodexResponsesRequest,
pub transport: CodexBackendTransport,
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
pub diagnostics: &'a CodexUpstreamDiagnostics,
pub response_metadata: &'a CodexResponseMetadata,
pub transport_metrics: &'a CodexTransportMetrics,
pub attempt: &'a ResponseDispatchAttempt,
pub attempts: &'a [ResponseDispatchAttempt],
pub started_at: std::time::Instant,
}
pub(super) struct StreamExit<'a> {
pub context: StreamContext<'a>,
pub summary: &'a StreamSummary,
pub rate_limit_headers: &'a [(String, String)],
pub failure_status: Option<i64>,
pub body: &'a str,
}
impl TelemetryController {
pub(super) async fn observe_dispatch_error(
recorder: &Recorder,
observation: DispatchErrorObservation<'_>,
error: &ResponseDispatchError,
) {
events::record_response_dispatch_error_event(events::ResponseDispatchErrorEventRecord {
recorder,
request_id: observation.request_id,
client_api_key_id: observation.client_api_key_id,
account_id: observation.account_id,
route: observation.route,
model: observation.model,
started_at: observation.started_at,
stream: observation.stream,
compact: observation.compact,
transport: observation.transport,
error,
})
.await;
}
pub(super) async fn observe_prefetched_stream_failure(
recorder: &Recorder,
observation: PrefetchedStreamFailureObservation<'_>,
) {
events::record_prefetched_response_stream_failure_event(
events::ResponseStreamFailureEventRecord {
recorder,
request_id: observation.request_id,
account_id: observation.account_id,
route: observation.route,
model: observation.model,
requested_model: observation.requested_model,
started_at: observation.started_at,
transport: observation.transport,
request: observation.request,
failure: observation.failure,
status_code: i64::from(observation.error.client_http_status_code()),
diagnostics: observation.diagnostics,
rate_limit_headers: observation.rate_limit_headers,
prefetched: observation.prefetched,
trace: observation.trace,
attempt: observation.attempt,
},
)
.await;
}
pub(super) async fn observe_upstream_error(
recorder: &Recorder,
context: &AttemptUpstreamErrorContext<'_>,
) {
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
recorder,
request_id: context.request_id,
account_id: &context.account.id,
account_email: context.account.email.as_deref(),
route: context.route,
model: context.model,
started_at: context.started_at,
stream: context.stream,
transport: context.transport,
request: context.request,
error: context.error,
trace: context.trace,
attempt: Some(context.attempt),
})
.await;
}
pub(super) async fn leave_complete(exit: &CompleteExit<'_>) {
let effective_model = exit
.response
.response_metadata
.effective_model
.as_deref()
.unwrap_or(exit.display_model);
let mut metadata = serde_json::json!({
"responseId": exit.response_id,
"stream": false,
"completed": exit.completed,
"incomplete": !exit.completed,
"transport": backend_transport_name(exit.response.transport),
"firstTokenMs": exit.response.first_token_ms,
"usage": exit.response.usage,
"effectiveModel": effective_model,
"modelsEtag": exit.response.response_metadata.models_etag.as_deref(),
"reasoningIncluded": exit.response.response_metadata.reasoning_included,
"transportDecision": exit.response.transport_metrics.decision.map(|decision| decision.as_str()),
"wsConnectMs": exit.response.transport_metrics.ws_connect_ms,
"transportDecisionWaitMs": exit.response.transport_metrics.transport_decision_wait_ms,
"upstreamHeadersMs": exit.response.transport_metrics.upstream_headers_ms,
"firstEventMs": exit.response.transport_metrics.first_event_ms,
"httpVersion": exit.response.transport_metrics.http_version.as_deref(),
});
insert_response_status_metadata(
&mut metadata,
200,
200,
exit.response.diagnostics.status_code.map(i64::from),
);
insert_response_upstream_diagnostics(&mut metadata, &exit.response.diagnostics);
insert_response_trace_metadata(&mut metadata, exit.trace, Some(exit.attempt));
insert_websocket_pool_decision(&mut metadata, exit.response.websocket_pool_decision);
enrich_response_request_semantics(&mut metadata, exit.request);
record_response_event(ResponseUsageRecord {
recorder: exit.recorder,
request_id: exit.request_id,
client_api_key_id: exit.request.client_api_key_id.as_deref(),
account_id: exit.account_id,
route: exit.route,
model: effective_model,
requested_model: Some(exit.requested_model),
client_ip: exit.request.client_ip.as_deref(),
client_user_agent: exit.request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(exit.request),
service_tier: exit.request.service_tier(),
started_at: exit.started_at,
status_code: 200,
message: if exit.completed {
"v1 responses completed"
} else {
"v1 responses incomplete"
},
metadata,
rate_limit_headers: &exit.response.rate_limit_headers,
})
.await;
}
pub(super) async fn leave_stream(exit: StreamExit<'_>) {
let context = &exit.context;
let summary = exit.summary;
let rate_limit_headers = exit.rate_limit_headers;
let body = exit.body;
match &summary.terminal {
StreamTerminal::Completed { response } => {
record_success(context, summary, response, true, rate_limit_headers, body).await;
}
StreamTerminal::Incomplete { response } => {
record_success(context, summary, response, false, rate_limit_headers, body).await;
}
StreamTerminal::Failed { failure } => {
record_failure(
context,
summary,
failure,
exit.failure_status.unwrap_or(502),
rate_limit_headers,
body,
)
.await;
}
terminal @ (StreamTerminal::UpstreamClosed
| StreamTerminal::UpstreamError { .. }
| StreamTerminal::ProtocolError { .. }
| StreamTerminal::CaptureLimitExceeded) => {
let failure = synthetic_failure(terminal);
record_failure(context, summary, &failure, 502, rate_limit_headers, body).await;
}
terminal @ (StreamTerminal::Cancelled
| StreamTerminal::DownstreamClosed
| StreamTerminal::Shutdown) => {
let status_code = if matches!(terminal, StreamTerminal::Shutdown) {
503
} else {
499
};
let mut metadata = serde_json::json!({
"stream": true,
"cancelled": true,
"terminal": terminal.name(),
"responseId": summary.last_response_id,
"firstEventMs": summary.first_event_ms,
"usage": summary.usage,
});
insert_first_token_ms(&mut metadata, summary.first_token_ms);
record_live_response_stream_event(
context,
status_code,
UsageRecordLevel::Warn,
"v1 responses stream cancelled",
metadata,
rate_limit_headers,
body,
)
.await;
}
}
}
}
async fn record_success(
context: &StreamContext<'_>,
summary: &StreamSummary,
response: &Value,
completed: bool,
rate_limit_headers: &[(String, String)],
body: &str,
) {
record_live_response_stream_event(
context,
200,
UsageRecordLevel::Info,
if completed {
"v1 responses stream completed"
} else {
"v1 responses stream incomplete"
},
serde_json::json!({
"stream": true,
"completed": completed,
"incomplete": !completed,
"responseId": response.get("id").and_then(Value::as_str),
"firstEventMs": summary.first_event_ms,
"firstTokenMs": summary.first_token_ms,
"usage": summary.usage,
}),
rate_limit_headers,
body,
)
.await;
}
fn synthetic_failure(terminal: &StreamTerminal) -> ResponsesSseFailure {
let detail = terminal.synthetic_failure_detail().unwrap_or_default();
ResponsesSseFailure {
event: "response.failed".to_string(),
message: if detail.trim().is_empty() {
STREAM_DISCONNECTED_MESSAGE.to_string()
} else {
format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}")
},
upstream_code: Some(STREAM_DISCONNECTED_CODE.to_string()),
upstream_type: Some("server_error".to_string()),
explicit_status_code: None,
retry_after_seconds: None,
}
}
async fn record_failure(
context: &StreamContext<'_>,
summary: &StreamSummary,
failure: &ResponsesSseFailure,
status_code: i64,
rate_limit_headers: &[(String, String)],
body: &str,
) {
let failure_source = stream_failure_source(failure);
let failure_detail = synthetic_stream_disconnected_detail(failure);
tracing::warn!(
account_id = %context.account_id,
request_id = %context.request_id,
response_id = summary.last_response_id.as_deref().unwrap_or(""),
transport = %backend_transport_name(context.transport),
websocket_pool_kind = ?context.websocket_pool_decision.map(WebSocketPoolDecision::kind),
first_token_ms = ?summary.first_token_ms,
first_event_ms = summary.first_event_ms,
latency_ms = elapsed_millis_i64(context.started_at),
event = %failure.event,
code = ?failure.upstream_code.as_deref(),
failure_source,
failure_detail = ?failure_detail.as_deref(),
"Live upstream stream ended with response.failed"
);
let mut metadata = stream_failure_metadata(failure, summary.usage);
metadata["responseId"] = summary
.last_response_id
.as_ref()
.map_or(Value::Null, |response_id| {
Value::String(response_id.clone())
});
insert_first_token_ms(&mut metadata, summary.first_token_ms);
metadata["firstEventMs"] = Value::Number(summary.first_event_ms.into());
record_live_response_stream_event(
context,
status_code,
UsageRecordLevel::Error,
"v1 responses stream failed",
metadata,
rate_limit_headers,
body,
)
.await;
}
+106
View File
@@ -0,0 +1,106 @@
//! 配额响应头与 token 用量写入的唯一 feature owner。
use crate::{
dispatch::lifecycle::contract::{
AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
},
fleet::pool::AccountPoolService,
upstream::openai::{
protocol::{events::TokenUsage, responses::CodexResponsesRequest},
transport::{CodexRateLimitHeaderUpdates, CodexTurnStateUpdate},
},
};
pub(super) struct UsageScope {
image_generation_requested: bool,
}
pub(super) struct StreamUsageExit {
pub(super) rate_limit_headers: Vec<(String, String)>,
pub(super) turn_state: Option<String>,
}
pub(super) struct StreamExit<'a> {
pub account_pool: &'a AccountPoolService,
pub account_id: &'a str,
pub account_plan_type: Option<&'a str>,
pub rate_limit_headers: &'a [(String, String)],
pub rate_limit_header_updates: Option<&'a CodexRateLimitHeaderUpdates>,
pub turn_state_update: Option<&'a CodexTurnStateUpdate>,
pub turn_state: Option<&'a str>,
pub usage: Option<TokenUsage>,
}
pub(super) struct UsageController;
impl UsageController {
pub(super) fn enter(request: &CodexResponsesRequest) -> UsageScope {
UsageScope {
image_generation_requested: request.expects_image_generation(),
}
}
pub(super) async fn observe_attempt(
account_pool: &AccountPoolService,
scope: &UsageScope,
observation: &AttemptObservation,
) {
if !matches!(
observation.kind,
AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Empty)
) {
return;
}
let Some(account_id) = observation
.account
.as_ref()
.map(|account| account.id.as_str())
else {
return;
};
account_pool
.record_empty_response_attempt(account_id, scope.image_generation_requested)
.await;
}
pub(super) async fn leave_complete(
account_pool: &AccountPoolService,
account_id: &str,
usage: Option<TokenUsage>,
image_generation_requested: bool,
) {
if let Some(usage) = usage {
account_pool
.record_response_usage(account_id, usage, image_generation_requested)
.await;
}
}
pub(super) async fn leave_stream(exit: StreamExit<'_>) -> StreamUsageExit {
let mut rate_limit_headers = exit.rate_limit_headers.to_vec();
if let Some(updates) = exit.rate_limit_header_updates {
rate_limit_headers.extend(updates.lock().await.iter().cloned());
}
exit.account_pool
.sync_passive_rate_limit_headers_for_account(
exit.account_id,
exit.account_plan_type,
&rate_limit_headers,
)
.await;
let turn_state = if let Some(update) = exit.turn_state_update {
update.lock().await.clone()
} else {
exit.turn_state.map(ToString::to_string)
};
if let Some(usage) = exit.usage {
exit.account_pool
.record_token_usage(exit.account_id, &usage)
.await;
}
StreamUsageExit {
rate_limit_headers,
turn_state,
}
}
}
+42 -140
View File
@@ -1,13 +1,12 @@
//! Shared upstream error classification for dispatch routes.
use chrono::{DateTime, Duration, Utc};
use serde_json::{Map, Value, json};
use thiserror::Error;
use crate::{
dispatch::{
recovery::exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
stream::sse_failure::{sse_failure_error_body, stream_failure_http_status},
dispatch::failure::{
exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
sse::sse_failure_error_body,
},
upstream::openai::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
@@ -15,6 +14,41 @@ use crate::{
},
};
/// controller 已经完成分类、可直接交给 API 编码的客户端失败事实。
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ClientFailure {
failure: ResponsesSseFailure,
status_code: u16,
expose_upstream: bool,
}
impl ClientFailure {
pub(in crate::dispatch) fn new(
failure: ResponsesSseFailure,
status_code: u16,
expose_upstream: bool,
) -> Self {
Self {
failure,
status_code,
expose_upstream,
}
}
pub(crate) fn failure(&self) -> &ResponsesSseFailure {
&self.failure
}
pub(crate) fn status_code(&self) -> u16 {
self.status_code
}
/// 只有 controller 明确允许时,API 才透传上游 code/message。
pub(crate) fn exposed_upstream(&self) -> Option<&ResponsesSseFailure> {
self.expose_upstream.then_some(&self.failure)
}
}
#[derive(Clone, Debug)]
pub(crate) struct DispatchErrorMetadata {
pub failure_class: DispatchFailureClass,
@@ -188,80 +222,7 @@ pub(crate) fn upstream_error_diagnostics(
}
}
pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, .. } if status_code_is_rate_limited(status.as_u16())
)
}
pub(crate) fn is_retryable_upstream_5xx_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, .. }
if status_code_is_transient_upstream(status.as_u16())
)
}
pub(crate) fn is_retryable_account_transport_error(error: &CodexClientError) -> bool {
match error {
CodexClientError::Http(error) => error.is_connect() || error.is_timeout(),
CodexClientError::StreamIdleTimeout { .. } => true,
CodexClientError::WebSocket(
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::Transport(_)
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ConnectTimeout { .. }
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::SendTimeout { .. }
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ClosedBeforeTerminal
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. }
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::InitialEventTimeout { .. },
) => true,
_ => is_retryable_upstream_5xx_error(error),
}
}
pub(crate) fn is_quota_exhausted_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, .. } if status_code_is_quota_exhausted(status.as_u16())
)
}
pub(crate) fn is_model_unsupported_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, body, .. }
if status.is_client_error()
&& !matches!(status.as_u16(), 401 | 402 | 403 | 404 | 429)
&& is_model_unsupported_signal(body)
)
}
pub(crate) fn is_history_recovery_upstream_error(error: &CodexClientError) -> bool {
upstream_error_code(error).is_some_and(|code| is_history_recovery_code(&code))
|| matches!(
error,
CodexClientError::WebSocket(
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ContinuationUnavailable { reason }
) if !matches!(
reason,
crate::upstream::openai::transport::websocket::PreviousResponseUnavailableReason::ConnectionBusy
)
)
}
pub(crate) fn is_continuation_busy_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::WebSocket(
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ContinuationUnavailable {
reason: crate::upstream::openai::transport::websocket::PreviousResponseUnavailableReason::ConnectionBusy,
}
)
)
}
pub(crate) fn upstream_error_body(error: &CodexClientError) -> String {
fn upstream_error_body(error: &CodexClientError) -> String {
match error {
CodexClientError::Upstream { body, .. } => body.clone(),
error => error.to_string(),
@@ -277,54 +238,6 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S
}
}
pub(crate) fn is_model_unsupported_signal(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("model_not_supported")
|| value.contains("model_not_available")
|| (value.contains("model")
&& (value.contains("not supported")
|| value.contains("not available")
|| value.contains("not_supported")
|| value.contains("not_available")))
}
pub(crate) fn is_history_recovery_code(code: &str) -> bool {
matches!(
code,
"previous_response_not_found"
| "invalid_encrypted_content"
| "missing_tool_output"
| "no_tool_output"
)
}
pub(crate) fn upstream_error_code(error: &CodexClientError) -> Option<String> {
let CodexClientError::Upstream { body, .. } = error else {
return None;
};
let value = serde_json::from_str::<Value>(body).ok()?;
value
.pointer("/response/error/code")
.or_else(|| value.pointer("/error/code"))
.or_else(|| value.get("code"))
.and_then(Value::as_str)
.map(ToString::to_string)
}
pub(crate) fn rate_limit_cooldown_until(
error: &CodexClientError,
now: DateTime<Utc>,
) -> DateTime<Utc> {
let retry_after_seconds = match error {
CodexClientError::Upstream {
retry_after_seconds,
..
} => retry_after_seconds.unwrap_or(60),
_ => 60,
};
now + Duration::seconds(retry_after_seconds.min(i64::MAX as u64) as i64)
}
pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 {
match error {
CodexClientError::Upstream { status, .. } => status.as_u16(),
@@ -339,17 +252,6 @@ pub(crate) fn backend_transport_name(transport: CodexBackendTransport) -> &'stat
}
}
fn status_code_is_rate_limited(status_code: u16) -> bool {
status_code == 429
}
fn status_code_is_quota_exhausted(status_code: u16) -> bool {
status_code == 402
}
fn status_code_is_transient_upstream(status_code: u16) -> bool {
matches!(status_code, 500..=599)
}
/// Responses 调度错误。
#[derive(Debug, Error)]
pub enum ResponseDispatchError {
@@ -414,7 +316,7 @@ pub enum ResponseDispatchError {
#[error("upstream response did not include visible output")]
EmptyUpstreamResponse,
#[error("upstream response failed: {0:?}")]
Failed(ResponsesSseFailure),
Failed(ClientFailure),
}
impl ResponseDispatchError {
@@ -474,7 +376,7 @@ impl ResponseDispatchError {
| Self::EmptyUpstreamResponse => 502,
Self::UpstreamUnavailable { .. } => 503,
Self::ContinuationBusy => 503,
Self::Failed(failure) => stream_failure_http_status(failure),
Self::Failed(failure) => failure.status_code(),
Self::ModelUnsupported { .. } | Self::HistoryUnavailable { .. } => 400,
Self::Upstream(error) => upstream_error_http_status(error),
}
@@ -570,7 +472,7 @@ impl ResponseDispatchError {
Self::Failed(failure) => DispatchErrorMetadata {
failure_class: DispatchFailureClass::ResponseFailed,
exhausted_count: None,
upstream_error: Some(sse_failure_error_body(failure)),
upstream_error: Some(sse_failure_error_body(failure.failure())),
upstream_status: None,
diagnostics: None,
},
@@ -1,8 +1,6 @@
//! 账号耗尽状态聚合。
use crate::fleet::account::AccountStatus;
#[derive(Clone, Copy, Debug)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ExhaustedAccountKind {
QuotaExhausted,
RateLimited,
@@ -15,6 +13,35 @@ pub(crate) enum ExhaustedAccountKind {
UpstreamUnavailable,
}
/// Controller 交给 attempt runner 聚合的单账号耗尽事实。
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) struct AccountExhaustionRecord {
pub account_id: Option<String>,
pub kind: ExhaustedAccountKind,
pub upstream_error: String,
pub status_code: Option<u16>,
}
impl AccountExhaustionRecord {
pub(in crate::dispatch) fn new(
account_id: impl Into<String>,
kind: ExhaustedAccountKind,
upstream_error: impl Into<String>,
) -> Self {
Self {
account_id: Some(account_id.into()),
kind,
upstream_error: upstream_error.into(),
status_code: None,
}
}
pub(in crate::dispatch) fn with_status_code(mut self, status_code: u16) -> Self {
self.status_code = Some(status_code);
self
}
}
impl ExhaustedAccountKind {
pub(crate) fn message_reason(self) -> &'static str {
match self {
@@ -31,7 +58,7 @@ impl ExhaustedAccountKind {
}
}
#[derive(Clone, Debug)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct ExhaustedAccount {
pub kind: ExhaustedAccountKind,
pub count: usize,
@@ -69,100 +96,12 @@ struct ExhaustedCounter {
}
impl AccountExhaustionTracker {
pub(crate) fn last_account_id(&self) -> Option<&str> {
self.last_account_id.as_deref()
}
pub(crate) fn record_rate_limited(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
pub(in crate::dispatch) fn record_exhaustion(&mut self, record: AccountExhaustionRecord) {
self.record(
account_id,
ExhaustedAccountKind::RateLimited,
upstream_error,
None,
);
}
pub(crate) fn record_quota_exhausted(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::QuotaExhausted,
upstream_error,
None,
);
}
pub(crate) fn record_auth_failure(
&mut self,
account_id: Option<&str>,
account_status: AccountStatus,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
let kind = match account_status {
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
AccountStatus::Banned => ExhaustedAccountKind::Banned,
_ => ExhaustedAccountKind::Expired,
};
self.record(account_id, kind, upstream_error, status_code);
}
pub(crate) fn record_cloudflare_challenge(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflareChallenge,
upstream_error,
None,
);
}
pub(crate) fn record_cloudflare_path_blocked(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflarePathBlocked,
upstream_error,
None,
);
}
pub(crate) fn record_model_unsupported(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::ModelUnsupported,
upstream_error,
None,
);
}
pub(crate) fn record_upstream_unavailable(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::UpstreamUnavailable,
upstream_error,
None,
record.account_id.as_deref(),
record.kind,
record.upstream_error,
record.status_code,
);
}
@@ -177,6 +116,10 @@ impl AccountExhaustionTracker {
})
}
pub(crate) fn last_account_id(&self) -> Option<&str> {
self.last_account_id.as_deref()
}
fn record(
&mut self,
account_id: Option<&str>,
+4
View File
@@ -0,0 +1,4 @@
//! 已分类账号失败的聚合、状态副作用和 SSE 终态事实。
pub(crate) mod exhaustion;
pub(in crate::dispatch) mod sse;
+78
View File
@@ -0,0 +1,78 @@
//! Responses SSE 失败事实和流终态元数据。
use serde_json::Value;
use crate::upstream::openai::protocol::{events::TokenUsage, responses::ResponsesSseFailure};
pub(in crate::dispatch) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
pub(in crate::dispatch) const STREAM_DISCONNECTED_MESSAGE: &str =
"Upstream stream closed before response.completed";
pub(in crate::dispatch) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
match failure.upstream_code.as_deref() {
Some(code) => serde_json::json!({
"error": {
"code": code,
"message": failure.message.as_str(),
}
})
.to_string(),
None => failure.message.clone(),
}
}
pub(in crate::dispatch) fn stream_failure_metadata(
failure: &ResponsesSseFailure,
usage: Option<TokenUsage>,
) -> Value {
let mut metadata = serde_json::json!({
"stream": true,
"failed": true,
"failureEvent": failure.event,
"failureMessage": failure.message,
"upstreamCode": failure.upstream_code,
"usage": usage,
});
enrich_stream_failure_source_metadata(&mut metadata, failure);
metadata
}
fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object.insert(
"failureSource".to_string(),
Value::String(stream_failure_source(failure).to_string()),
);
if let Some(detail) = synthetic_stream_disconnected_detail(failure) {
object.insert("synthetic".to_string(), Value::Bool(true));
if !detail.is_empty() {
object.insert("failureDetail".to_string(), Value::String(detail));
}
}
}
pub(in crate::dispatch) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
if synthetic_stream_disconnected_detail(failure).is_some() {
"proxy"
} else {
"upstream"
}
}
pub(in crate::dispatch) fn synthetic_stream_disconnected_detail(
failure: &ResponsesSseFailure,
) -> Option<String> {
if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) {
return None;
}
let detail = failure
.message
.strip_prefix(STREAM_DISCONNECTED_MESSAGE)?
.strip_prefix(": ")
.unwrap_or_default()
.trim()
.to_string();
Some(detail)
}
+990
View File
@@ -0,0 +1,990 @@
//! 单账号 attempt 的 stepwise 执行器。
use std::time::Instant;
use serde_json::Value;
use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::{
AttemptAccountPreparation, AttemptAccountPreparationContext, AttemptRoutePreparation,
AttemptRoutePreparationContext, AttemptUpstreamErrorContext, ControllerRequestScope,
ControllerSet,
},
errors::{ClientFailure, ResponseDispatchError},
failure::exhaustion::AccountExhaustionTracker,
routing::candidates::AccountAttemptLedger,
transport::{
account::{
AccountUpstreamContext, create_response_stream_with_account,
create_response_with_account, prepare_response_transport_with_account,
},
canonical::{CanonicalStreamDecoder, normalize_complete_response},
prefetch::{PrefetchedStream, StreamPrefetchError, prefetch_until_commit},
},
},
fleet::{
account::Account,
pool::{AccountCandidateLease, AccountPoolService},
},
upstream::openai::{
failure::upstream_failure_facts,
protocol::responses::{
CodexResponsesRequest, CollectedResponse, TransportRequirement, transport_requirement,
},
transport::{CodexBackendClient, CodexBackendStreamingResponse, CodexBackendTransport},
},
};
use super::{
contract::{
AttemptAccountFacts, AttemptApplyOutcome, AttemptContractError, AttemptDecision,
AttemptObservation, AttemptObservationKind, AttemptRejection, AttemptReturnKind,
CandidateLedgerFacts, CompleteResponseFacts, EstablishedAttemptContext,
EstablishedComplete, EstablishedCompleteBody, EstablishedResponse, EstablishedStream,
PendingAttempt, PendingCompleteResponse, PendingProtocolFailure, PendingStreamResponse,
PinnedCandidateAcquireFailureKind, ProtocolFailureKind,
},
trace::ResponseDispatchTrace,
};
pub(in crate::dispatch) enum AttemptMode {
Complete {
started_at: Instant,
tuple_schema: Option<Value>,
},
Stream {
tuple_schema: Option<Value>,
},
}
pub(in crate::dispatch) struct AttemptRunnerDependencies<'a> {
pub account_pool: &'a AccountPoolService,
pub codex: &'a CodexBackendClient,
pub controllers: &'a ControllerSet,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
pub route: &'a str,
pub requested_model: &'a str,
pub started_at: Instant,
}
/// 尚未越过提交边界的 attempt capability;只有该类型接受 retry decision。
pub(in crate::dispatch) struct OpenAttempt<'runner, 'dependencies> {
runner: &'runner mut AttemptRunner<'dependencies>,
observation: Box<AttemptObservation>,
}
impl OpenAttempt<'_, '_> {
pub(in crate::dispatch) fn observation(&self) -> &AttemptObservation {
&self.observation
}
pub(in crate::dispatch) fn controller_scope(
&self,
) -> Result<&ControllerRequestScope, AttemptContractError> {
self.runner.controller_scope()
}
pub(in crate::dispatch) async fn apply(
self,
decision: AttemptDecision,
) -> Result<AttemptApplyOutcome, AttemptContractError> {
self.runner.apply_open(decision).await
}
}
/// 已越过提交边界的 attempt capability;该类型只允许建立 stream。
pub(in crate::dispatch) struct CommittedAttempt {
established: EstablishedResponse,
}
impl CommittedAttempt {
pub(in crate::dispatch) fn accept(self) -> EstablishedResponse {
self.established
}
}
/// `next` 产出的提交 typestate;Committed 分支没有 decision/retry API。
pub(in crate::dispatch) enum AttemptStep<'runner, 'dependencies> {
Open(OpenAttempt<'runner, 'dependencies>),
Committed(CommittedAttempt),
}
pub(in crate::dispatch) struct AttemptRunner<'a> {
dependencies: AttemptRunnerDependencies<'a>,
mode: AttemptMode,
request: Option<CodexResponsesRequest>,
controller_scope: Option<ControllerRequestScope>,
candidates: AccountAttemptLedger,
exhaustion: AccountExhaustionTracker,
on_exhaustion: Option<ClientFailure>,
trace: Option<ResponseDispatchTrace>,
required_account_id: Option<String>,
pending: Option<PendingAttempt>,
}
impl<'a> AttemptRunner<'a> {
pub(in crate::dispatch) fn new(
dependencies: AttemptRunnerDependencies<'a>,
mode: AttemptMode,
request: CodexResponsesRequest,
controller_scope: ControllerRequestScope,
candidates: AccountAttemptLedger,
) -> Self {
Self {
dependencies,
mode,
request: Some(request),
controller_scope: Some(controller_scope),
candidates,
exhaustion: AccountExhaustionTracker::default(),
on_exhaustion: None,
trace: Some(ResponseDispatchTrace::default()),
required_account_id: None,
pending: None,
}
}
pub(in crate::dispatch) async fn next(
&mut self,
) -> Result<AttemptStep<'_, 'a>, AttemptContractError> {
if self.pending.is_some() {
return Err(AttemptContractError::DecisionRequired);
}
let request = self
.request
.as_ref()
.ok_or(AttemptContractError::Terminal)?;
let acquired = if let Some(account_id) = self.required_account_id.take() {
match self
.dependencies
.account_pool
.acquire_candidate(request.model(), &account_id, chrono::Utc::now())
.await
{
AccountCandidateLease::Acquired(acquired) => Some(*acquired),
AccountCandidateLease::Busy => {
self.pending = Some(PendingAttempt::PinnedCandidateUnavailable {
account_id,
kind: PinnedCandidateAcquireFailureKind::Busy,
});
return self.step();
}
AccountCandidateLease::Unavailable => {
self.pending = Some(PendingAttempt::PinnedCandidateUnavailable {
account_id,
kind: PinnedCandidateAcquireFailureKind::Unavailable,
});
return self.step();
}
}
} else {
self.candidates
.acquire_next(self.dependencies.account_pool)
.await
};
let Some(acquired) = acquired else {
self.pending = Some(PendingAttempt::NoCandidate);
return self.step();
};
let account = acquired.account.clone();
let controller_scope = self
.controller_scope
.as_mut()
.ok_or(AttemptContractError::Terminal)?;
let prepared_account = match self
.dependencies
.controllers
.prepare_attempt_account(
AttemptAccountPreparationContext {
codex: self.dependencies.codex,
account_identity: self.dependencies.account_identity,
request_id: self.dependencies.request_id,
},
acquired,
)
.await
{
AttemptAccountPreparation::Ready(prepared) => *prepared,
AttemptAccountPreparation::Rejected => {
self.pending = Some(PendingAttempt::CandidatePreparationRejected { account });
return self.step();
}
};
let acquired = prepared_account.lease;
let cookie_header = prepared_account.cookie_header;
let account_scoped_request = match self.dependencies.controllers.prepare_attempt_route(
controller_scope,
request,
AttemptRoutePreparationContext {
account_identity: self.dependencies.account_identity,
account_id: &account.id,
},
) {
AttemptRoutePreparation::Ready(request) => *request,
AttemptRoutePreparation::Unavailable { message } => {
acquired.release_without_usage().await;
self.pending = Some(PendingAttempt::RoutePreparationRejected { account, message });
return self.step();
}
};
let attempt = self
.trace
.as_mut()
.ok_or(AttemptContractError::Terminal)?
.start_attempt(&account.id);
let upstream_context = AccountUpstreamContext {
codex: self.dependencies.codex,
request_id: self.dependencies.request_id,
account: &account,
cookie_header: cookie_header.as_deref(),
};
let ((), prepared_transport) = tokio::join!(
self.dependencies
.account_pool
.wait_for_request_interval(&acquired),
prepare_response_transport_with_account(upstream_context, &account_scoped_request),
);
let prepared_transport = match prepared_transport {
Ok(prepared) => prepared,
Err(error) => {
let attempt_request = account_scoped_request.into_request();
acquired.complete().await;
let transport = error
.transport()
.unwrap_or_else(|| intended_transport(&attempt_request));
observe_upstream_error(
self.dependencies.controllers,
AttemptUpstreamErrorContext {
request_id: self.dependencies.request_id,
account: &account,
route: self.dependencies.route,
model: self.dependencies.requested_model,
started_at: self.dependencies.started_at,
stream: matches!(self.mode, AttemptMode::Stream { .. }),
transport,
request: &attempt_request,
error: &error,
trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?,
attempt: &attempt,
},
)
.await;
self.pending = Some(PendingAttempt::UpstreamFailure {
account,
attempt_request,
attempt,
transport,
error,
});
return self.step();
}
};
match &self.mode {
AttemptMode::Complete {
started_at,
tuple_schema,
} => {
let result = create_response_with_account(
upstream_context,
&account_scoped_request,
prepared_transport,
*started_at,
)
.await;
let attempt_request = account_scoped_request.into_request();
acquired.complete().await;
self.pending = Some(match result {
Ok(response) => {
self.dependencies
.controllers
.observe_complete_upstream(
&account,
&response.set_cookie_headers,
&response.rate_limit_headers,
)
.await;
match normalize_complete_response(&response.body, tuple_schema.as_ref()) {
Ok(collected) => PendingAttempt::CompleteResponse(Box::new(
PendingCompleteResponse {
account,
attempt_request,
attempt,
response,
collected,
},
)),
Err(error) => PendingAttempt::ProtocolFailure {
account,
attempt_request,
attempt,
transport: response.transport,
error: PendingProtocolFailure::InvalidSse(error),
},
}
}
Err(error) => {
let transport = error
.transport()
.unwrap_or_else(|| intended_transport(request));
observe_upstream_error(
self.dependencies.controllers,
AttemptUpstreamErrorContext {
request_id: self.dependencies.request_id,
account: &account,
route: self.dependencies.route,
model: self.dependencies.requested_model,
started_at: self.dependencies.started_at,
stream: matches!(self.mode, AttemptMode::Stream { .. }),
transport,
request: &attempt_request,
error: &error,
trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?,
attempt: &attempt,
},
)
.await;
PendingAttempt::UpstreamFailure {
account,
attempt_request,
attempt,
transport,
error,
}
}
});
}
AttemptMode::Stream { tuple_schema } => {
let result = create_response_stream_with_account(
upstream_context,
&account_scoped_request,
prepared_transport,
)
.await;
let attempt_request = account_scoped_request.into_request();
self.pending = Some(match result {
Err(error) => {
acquired.complete().await;
let transport = error
.transport()
.unwrap_or_else(|| intended_transport(request));
observe_upstream_error(
self.dependencies.controllers,
AttemptUpstreamErrorContext {
request_id: self.dependencies.request_id,
account: &account,
route: self.dependencies.route,
model: self.dependencies.requested_model,
started_at: self.dependencies.started_at,
stream: matches!(self.mode, AttemptMode::Stream { .. }),
transport,
request: &attempt_request,
error: &error,
trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?,
attempt: &attempt,
},
)
.await;
PendingAttempt::UpstreamFailure {
account,
attempt_request,
attempt,
transport,
error,
}
}
Ok(response) => {
let CodexBackendStreamingResponse {
body,
transport,
turn_state,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates,
turn_state_update,
websocket_pool_decision,
diagnostics,
response_metadata,
transport_metrics,
connection_local_continuation,
} = response;
let mut decoder = CanonicalStreamDecoder::new(tuple_schema.clone());
match prefetch_until_commit(
body,
&mut decoder,
attempt_request.stream_commit_policy,
self.dependencies.started_at,
)
.await
{
Err(StreamPrefetchError::Upstream(error)) => {
acquired.complete().await;
observe_upstream_error(
self.dependencies.controllers,
AttemptUpstreamErrorContext {
request_id: self.dependencies.request_id,
account: &account,
route: self.dependencies.route,
model: self.dependencies.requested_model,
started_at: self.dependencies.started_at,
stream: matches!(self.mode, AttemptMode::Stream { .. }),
transport,
request: &attempt_request,
error: &error,
trace: self
.trace
.as_ref()
.ok_or(AttemptContractError::Terminal)?,
attempt: &attempt,
},
)
.await;
PendingAttempt::UpstreamFailure {
account,
attempt_request,
attempt,
transport,
error,
}
}
Err(StreamPrefetchError::Empty) => {
acquired.complete().await;
PendingAttempt::ProtocolFailure {
account,
attempt_request,
attempt,
transport,
error: PendingProtocolFailure::EmptyStream,
}
}
Err(StreamPrefetchError::NoCommitBoundary) => {
acquired.complete().await;
PendingAttempt::ProtocolFailure {
account,
attempt_request,
attempt,
transport,
error: PendingProtocolFailure::NoCommitBoundary,
}
}
Err(StreamPrefetchError::InvalidSse(error)) => {
acquired.complete().await;
PendingAttempt::ProtocolFailure {
account,
attempt_request,
attempt,
transport,
error: PendingProtocolFailure::InvalidSse(error),
}
}
Ok(PrefetchedStream {
bytes: prefetched,
initial_batch,
body,
first_event_ms,
}) => {
let first_failure = initial_batch.first_failure();
PendingAttempt::StreamResponse(Box::new(PendingStreamResponse {
account,
attempt_request,
attempt,
lease: acquired,
response: CodexBackendStreamingResponse {
body,
transport,
turn_state,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates,
turn_state_update,
websocket_pool_decision,
diagnostics,
response_metadata,
transport_metrics,
connection_local_continuation,
},
prefetched,
decoder,
initial_batch,
first_failure,
first_event_ms,
}))
}
}
}
});
}
}
self.step()
}
pub(in crate::dispatch) fn controller_scope(
&self,
) -> Result<&ControllerRequestScope, AttemptContractError> {
self.controller_scope
.as_ref()
.ok_or(AttemptContractError::Terminal)
}
async fn apply_open(
&mut self,
decision: AttemptDecision,
) -> Result<AttemptApplyOutcome, AttemptContractError> {
let pending = self
.pending
.take()
.ok_or(AttemptContractError::DecisionRequired)?;
match decision {
AttemptDecision::Accept => self
.establish(pending)
.map(AttemptApplyOutcome::Established),
AttemptDecision::RetrySameAccount => {
let account_id = pending_account(&pending)
.ok_or_else(|| invalid("retry_same", &pending))?
.to_string();
if !self.dependencies.controllers.prepare_same_account_retry(
self.controller_scope
.as_mut()
.ok_or(AttemptContractError::Terminal)?,
&account_id,
) {
return Err(invalid("retry_same", &pending));
}
release_pending_stream(pending).await;
self.required_account_id = Some(account_id);
Ok(AttemptApplyOutcome::Continue)
}
AttemptDecision::RetryNextCandidate {
exhaustion,
on_exhaustion,
} => {
release_pending_stream(pending).await;
self.on_exhaustion = on_exhaustion;
if let Some(exhaustion) = exhaustion {
self.exhaustion.record_exhaustion(exhaustion);
}
Ok(AttemptApplyOutcome::Continue)
}
AttemptDecision::Return(kind) => self.reject(pending, kind).await,
}
}
fn step(&mut self) -> Result<AttemptStep<'_, 'a>, AttemptContractError> {
let kind = {
let pending = self
.pending
.as_ref()
.ok_or(AttemptContractError::DecisionRequired)?;
observation_kind(pending, &self.candidates, &self.exhaustion)
};
let Some(kind) = kind else {
let pending = self
.pending
.take()
.ok_or(AttemptContractError::DecisionRequired)?;
return Ok(AttemptStep::Committed(CommittedAttempt {
established: self.establish(pending)?,
}));
};
let observation = self.observation(kind)?;
Ok(AttemptStep::Open(OpenAttempt {
runner: self,
observation: Box::new(observation),
}))
}
fn observation(
&self,
kind: AttemptObservationKind,
) -> Result<AttemptObservation, AttemptContractError> {
let pending = self
.pending
.as_ref()
.ok_or(AttemptContractError::DecisionRequired)?;
let request = self
.request
.as_ref()
.ok_or(AttemptContractError::Terminal)?;
let controller_scope = self
.controller_scope
.as_ref()
.ok_or(AttemptContractError::Terminal)?;
let account = pending_account_value(pending);
let account_id = account.map(|account| account.id.as_str());
Ok(AttemptObservation {
account: account.map(AttemptAccountFacts::from),
attempt: pending_attempt(pending).cloned(),
transport: pending_transport(pending).unwrap_or_else(|| intended_transport(request)),
routing: self
.dependencies
.controllers
.attempt_routing_facts(controller_scope, account_id),
kind,
})
}
fn establish(
&mut self,
pending: PendingAttempt,
) -> Result<EstablishedResponse, AttemptContractError> {
let established = match pending {
PendingAttempt::CompleteResponse(complete) => {
let PendingCompleteResponse {
account,
attempt,
response,
collected,
..
} = *complete;
let body = EstablishedCompleteBody::try_from(collected).map_err(|()| {
AttemptContractError::InvalidDecision {
decision: "establish",
observation: "non_successful_complete_response",
}
})?;
EstablishedResponse::Complete(Box::new(EstablishedComplete {
context: self.take_established_context(account, attempt)?,
response,
body,
}))
}
PendingAttempt::StreamResponse(stream) => {
let PendingStreamResponse {
account,
attempt,
lease,
response,
decoder,
initial_batch,
first_event_ms,
..
} = *stream;
EstablishedResponse::Stream(Box::new(EstablishedStream {
context: self.take_established_context(account, attempt)?,
lease,
response,
decoder,
initial_batch,
first_event_ms,
}))
}
other => return Err(invalid("establish", &other)),
};
Ok(established)
}
async fn reject(
&mut self,
pending: PendingAttempt,
kind: AttemptReturnKind,
) -> Result<AttemptApplyOutcome, AttemptContractError> {
let account = pending_account_value(&pending).cloned();
let account_id = account
.as_ref()
.map(|account| account.id.clone())
.or_else(|| self.exhaustion.last_account_id().map(ToString::to_string));
let attempt_request = pending_attempt_request(&pending).cloned();
let attempt = pending_attempt(&pending).cloned();
let request = self
.request
.as_ref()
.ok_or(AttemptContractError::Terminal)?;
let transport = pending_transport(&pending).unwrap_or_else(|| intended_transport(request));
let stream_failure = match &pending {
PendingAttempt::StreamResponse(stream) => {
stream.first_failure.as_ref().map(|failure| {
super::contract::RejectedStreamFailure {
failure: failure.clone(),
prefetched: stream.prefetched.clone(),
diagnostics: stream.response.diagnostics.clone(),
rate_limit_headers: stream.response.rate_limit_headers.clone(),
}
})
}
_ => None,
};
let error =
rejection_error(pending, kind, &self.exhaustion, self.on_exhaustion.take()).await?;
let request = self.request.take().ok_or(AttemptContractError::Terminal)?;
self.controller_scope
.take()
.ok_or(AttemptContractError::Terminal)?;
let trace = self.trace.take().ok_or(AttemptContractError::Terminal)?;
Ok(AttemptApplyOutcome::Rejected(Box::new(AttemptRejection {
request,
trace,
account_id,
account,
attempt_request,
attempt,
transport,
stream_failure,
error,
})))
}
fn take_established_context(
&mut self,
account: Account,
attempt: super::trace::ResponseDispatchAttempt,
) -> Result<EstablishedAttemptContext, AttemptContractError> {
Ok(EstablishedAttemptContext {
request: self.request.take().ok_or(AttemptContractError::Terminal)?,
controller_scope: self
.controller_scope
.take()
.ok_or(AttemptContractError::Terminal)?,
trace: self.trace.take().ok_or(AttemptContractError::Terminal)?,
account,
attempt,
})
}
}
async fn observe_upstream_error(
controllers: &ControllerSet,
context: AttemptUpstreamErrorContext<'_>,
) {
controllers.observe_upstream_error(context).await;
}
fn observation_kind(
pending: &PendingAttempt,
candidates: &AccountAttemptLedger,
exhaustion: &AccountExhaustionTracker,
) -> Option<AttemptObservationKind> {
Some(match pending {
PendingAttempt::NoCandidate => AttemptObservationKind::NoCandidate {
ledger: CandidateLedgerFacts {
candidates: candidates.candidate_count(),
attempted: candidates.attempted_count(),
state_excluded: candidates.state_excluded_count(),
},
last_exhausted: exhaustion.last_exhausted(),
},
PendingAttempt::PinnedCandidateUnavailable { account_id, kind } => {
AttemptObservationKind::PinnedCandidateUnavailable {
account_id: account_id.clone(),
kind: *kind,
}
}
PendingAttempt::CandidatePreparationRejected { .. } => {
AttemptObservationKind::CandidatePreparationRejected
}
PendingAttempt::RoutePreparationRejected { message, .. } => {
AttemptObservationKind::RoutePreparationRejected {
message: message.clone(),
}
}
PendingAttempt::UpstreamFailure { error, .. } => {
AttemptObservationKind::UpstreamFailure(upstream_failure_facts(error))
}
PendingAttempt::ProtocolFailure { error, .. } => AttemptObservationKind::ProtocolFailure {
kind: match error {
PendingProtocolFailure::InvalidSse(_) => ProtocolFailureKind::InvalidSse,
PendingProtocolFailure::EmptyStream => ProtocolFailureKind::EmptyStream,
PendingProtocolFailure::NoCommitBoundary => ProtocolFailureKind::NoCommitBoundary,
},
message: match error {
PendingProtocolFailure::InvalidSse(error) => error.to_string(),
PendingProtocolFailure::EmptyStream => {
"upstream stream ended without data".to_string()
}
PendingProtocolFailure::NoCommitBoundary => {
"upstream stream ended before output or a terminal event".to_string()
}
},
},
PendingAttempt::CompleteResponse(complete) => {
AttemptObservationKind::CompleteResponse(match &complete.collected {
CollectedResponse::Completed(_) => CompleteResponseFacts::Completed,
CollectedResponse::Incomplete(_) => CompleteResponseFacts::Incomplete,
CollectedResponse::Failed(failure) => {
CompleteResponseFacts::Failed(failure.clone())
}
CollectedResponse::MissingCompleted => CompleteResponseFacts::MissingCompleted,
CollectedResponse::Empty => CompleteResponseFacts::Empty,
})
}
PendingAttempt::StreamResponse(stream) => {
return stream
.first_failure
.clone()
.map(AttemptObservationKind::StreamFailure);
}
})
}
async fn release_pending_stream(pending: PendingAttempt) {
if let PendingAttempt::StreamResponse(stream) = pending {
stream.lease.complete().await;
}
}
async fn rejection_error(
pending: PendingAttempt,
kind: AttemptReturnKind,
exhaustion: &AccountExhaustionTracker,
on_exhaustion: Option<ClientFailure>,
) -> Result<ResponseDispatchError, AttemptContractError> {
match kind {
AttemptReturnKind::Failed(failure) => {
release_pending_stream(pending).await;
Ok(ResponseDispatchError::Failed(failure))
}
AttemptReturnKind::ContinuationBusy => Ok(ResponseDispatchError::ContinuationBusy),
AttemptReturnKind::RouteUnavailable { message } => {
Ok(ResponseDispatchError::HistoryUnavailable {
upstream_error: message,
})
}
AttemptReturnKind::Observed => match pending {
PendingAttempt::NoCandidate | PendingAttempt::CandidatePreparationRejected { .. } => {
Ok(on_exhaustion.map_or_else(
|| {
exhaustion
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount)
},
ResponseDispatchError::Failed,
))
}
PendingAttempt::RoutePreparationRejected { message, .. } => {
Ok(ResponseDispatchError::HistoryUnavailable {
upstream_error: message,
})
}
PendingAttempt::PinnedCandidateUnavailable { .. } => {
Err(invalid("return", &PendingAttempt::NoCandidate))
}
PendingAttempt::UpstreamFailure { error, .. } => {
Ok(ResponseDispatchError::Upstream(error))
}
PendingAttempt::ProtocolFailure { error, .. } => Ok(match error {
PendingProtocolFailure::InvalidSse(error) => {
ResponseDispatchError::InvalidSse(error)
}
PendingProtocolFailure::EmptyStream => ResponseDispatchError::EmptyUpstreamResponse,
PendingProtocolFailure::NoCommitBoundary => ResponseDispatchError::MissingCompleted,
}),
PendingAttempt::CompleteResponse(complete) => match complete.collected {
CollectedResponse::Failed(_) => Err(invalid(
"return_untyped_failure",
&PendingAttempt::NoCandidate,
)),
CollectedResponse::MissingCompleted => Ok(ResponseDispatchError::MissingCompleted),
CollectedResponse::Empty => Ok(ResponseDispatchError::EmptyUpstreamResponse),
CollectedResponse::Completed(_) | CollectedResponse::Incomplete(_) => {
Err(invalid("return", &PendingAttempt::NoCandidate))
}
},
PendingAttempt::StreamResponse(stream) => match stream.first_failure {
Some(failure) => {
stream.lease.complete().await;
let _ = failure;
Err(invalid(
"return_untyped_failure",
&PendingAttempt::NoCandidate,
))
}
None => Err(invalid("return", &PendingAttempt::NoCandidate)),
},
},
}
}
fn pending_account(pending: &PendingAttempt) -> Option<&str> {
pending_account_value(pending).map(|account| account.id.as_str())
}
fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> {
match pending {
PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None,
PendingAttempt::CandidatePreparationRejected { account }
| PendingAttempt::RoutePreparationRejected { account, .. }
| PendingAttempt::UpstreamFailure { account, .. }
| PendingAttempt::ProtocolFailure { account, .. } => Some(account),
PendingAttempt::CompleteResponse(complete) => Some(&complete.account),
PendingAttempt::StreamResponse(stream) => Some(&stream.account),
}
}
fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDispatchAttempt> {
match pending {
PendingAttempt::UpstreamFailure { attempt, .. }
| PendingAttempt::ProtocolFailure { attempt, .. } => Some(attempt),
PendingAttempt::CompleteResponse(complete) => Some(&complete.attempt),
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRequest> {
match pending {
PendingAttempt::UpstreamFailure {
attempt_request, ..
}
| PendingAttempt::ProtocolFailure {
attempt_request, ..
} => Some(attempt_request),
PendingAttempt::CompleteResponse(complete) => Some(&complete.attempt_request),
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
fn pending_transport(
pending: &PendingAttempt,
) -> Option<crate::upstream::openai::transport::CodexBackendTransport> {
match pending {
PendingAttempt::UpstreamFailure { transport, .. }
| PendingAttempt::ProtocolFailure { transport, .. } => Some(*transport),
PendingAttempt::CompleteResponse(complete) => Some(complete.response.transport),
PendingAttempt::StreamResponse(stream) => Some(stream.response.transport),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
fn intended_transport(request: &CodexResponsesRequest) -> CodexBackendTransport {
match transport_requirement(request) {
TransportRequirement::HttpRequired => CodexBackendTransport::HttpSse,
TransportRequirement::ExplicitWebSocketWarmup
| TransportRequirement::ExactWebSocketContinuation
| TransportRequirement::PersistedContinuation
| TransportRequirement::ExternalUnknown
| TransportRequirement::NewChain => CodexBackendTransport::WebSocket,
}
}
fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractError {
AttemptContractError::InvalidDecision {
decision,
observation: match pending {
PendingAttempt::NoCandidate => "no_candidate",
PendingAttempt::PinnedCandidateUnavailable { .. } => "pinned_candidate_unavailable",
PendingAttempt::CandidatePreparationRejected { .. } => "candidate_preparation_rejected",
PendingAttempt::RoutePreparationRejected { .. } => "route_preparation_rejected",
PendingAttempt::UpstreamFailure { .. } => "upstream_failure",
PendingAttempt::ProtocolFailure { .. } => "protocol_failure",
PendingAttempt::CompleteResponse(_) => "complete_response",
PendingAttempt::StreamResponse(_) => "stream_response",
},
}
}
+282
View File
@@ -0,0 +1,282 @@
//! Attempt lifecycle 的类型化 observation、decision 与结果。
use bytes::Bytes;
use crate::{
dispatch::{
errors::{ClientFailure, ResponseDispatchError},
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccount},
transport::canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
},
fleet::{account::Account, pool::AccountLease},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::{
responses::{CodexResponsesRequest, CollectedResponse, ResponsesSseFailure},
sse::SseError,
},
transport::{
CodexBackendResponse, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexUpstreamDiagnostics,
},
},
};
use super::trace::{ResponseDispatchAttempt, ResponseDispatchTrace};
/// Request/Attempt/Stream 共同消费的最终结果类别。
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum FinalOutcome {
Completed,
Incomplete,
Failed,
Cancelled,
Shutdown,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) struct AttemptAccountFacts {
pub id: String,
pub email: Option<String>,
pub plan_type: Option<String>,
}
impl From<&Account> for AttemptAccountFacts {
fn from(account: &Account) -> Self {
Self {
id: account.id.clone(),
email: account.email.clone(),
plan_type: account.plan_type.clone(),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::dispatch) struct AttemptRoutingFacts {
pub external_origin: bool,
pub can_retry_same_account: bool,
pub can_retry_next_candidate: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::dispatch) struct CandidateLedgerFacts {
pub candidates: usize,
pub attempted: usize,
pub state_excluded: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum CompleteResponseFacts {
Completed,
Incomplete,
Failed(ResponsesSseFailure),
MissingCompleted,
Empty,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum ProtocolFailureKind {
InvalidSse,
EmptyStream,
NoCommitBoundary,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum PinnedCandidateAcquireFailureKind {
Busy,
Unavailable,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum AttemptObservationKind {
NoCandidate {
ledger: CandidateLedgerFacts,
last_exhausted: Option<ExhaustedAccount>,
},
CandidatePreparationRejected,
RoutePreparationRejected {
message: String,
},
PinnedCandidateUnavailable {
account_id: String,
kind: PinnedCandidateAcquireFailureKind,
},
UpstreamFailure(UpstreamFailureFacts),
ProtocolFailure {
kind: ProtocolFailureKind,
message: String,
},
CompleteResponse(CompleteResponseFacts),
StreamFailure(ResponsesSseFailure),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) struct AttemptObservation {
pub account: Option<AttemptAccountFacts>,
pub attempt: Option<ResponseDispatchAttempt>,
pub transport: CodexBackendTransport,
pub routing: AttemptRoutingFacts,
pub kind: AttemptObservationKind,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum AttemptDecision {
Accept,
RetrySameAccount,
RetryNextCandidate {
exhaustion: Option<AccountExhaustionRecord>,
on_exhaustion: Option<ClientFailure>,
},
Return(AttemptReturnKind),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum AttemptReturnKind {
Observed,
Failed(ClientFailure),
ContinuationBusy,
RouteUnavailable { message: String },
}
pub(in crate::dispatch) struct EstablishedAttemptContext {
pub request: CodexResponsesRequest,
pub controller_scope: crate::dispatch::controllers::ControllerRequestScope,
pub trace: ResponseDispatchTrace,
pub account: Account,
pub attempt: ResponseDispatchAttempt,
}
pub(in crate::dispatch) struct EstablishedComplete {
pub context: EstablishedAttemptContext,
pub response: CodexBackendResponse,
pub body: EstablishedCompleteBody,
}
pub(in crate::dispatch) enum EstablishedCompleteBody {
Completed(serde_json::Value),
Incomplete(serde_json::Value),
}
impl TryFrom<CollectedResponse> for EstablishedCompleteBody {
type Error = ();
fn try_from(value: CollectedResponse) -> Result<Self, Self::Error> {
match value {
CollectedResponse::Completed(body) => Ok(Self::Completed(body)),
CollectedResponse::Incomplete(body) => Ok(Self::Incomplete(body)),
CollectedResponse::Failed(_)
| CollectedResponse::MissingCompleted
| CollectedResponse::Empty => Err(()),
}
}
}
pub(in crate::dispatch) struct EstablishedStream {
pub context: EstablishedAttemptContext,
pub lease: AccountLease,
pub response: CodexBackendStreamingResponse,
pub decoder: CanonicalStreamDecoder,
pub initial_batch: CanonicalStreamBatch,
pub first_event_ms: i64,
}
pub(in crate::dispatch) enum EstablishedResponse {
Complete(Box<EstablishedComplete>),
Stream(Box<EstablishedStream>),
}
pub(in crate::dispatch) struct AttemptRejection {
pub request: CodexResponsesRequest,
pub trace: ResponseDispatchTrace,
pub account_id: Option<String>,
pub account: Option<Account>,
pub attempt_request: Option<CodexResponsesRequest>,
pub attempt: Option<ResponseDispatchAttempt>,
pub transport: CodexBackendTransport,
pub stream_failure: Option<RejectedStreamFailure>,
pub error: ResponseDispatchError,
}
pub(in crate::dispatch) struct RejectedStreamFailure {
pub failure: ResponsesSseFailure,
pub prefetched: Bytes,
pub diagnostics: CodexUpstreamDiagnostics,
pub rate_limit_headers: Vec<(String, String)>,
}
pub(in crate::dispatch) enum AttemptApplyOutcome {
Continue,
Established(EstablishedResponse),
Rejected(Box<AttemptRejection>),
}
#[derive(Debug, thiserror::Error)]
pub(in crate::dispatch) enum AttemptContractError {
#[error("the previous attempt observation has not been decided")]
DecisionRequired,
#[error("attempt runner has already reached a terminal outcome")]
Terminal,
#[error("decision `{decision}` is invalid for observation `{observation}`")]
InvalidDecision {
decision: &'static str,
observation: &'static str,
},
}
pub(super) enum PendingProtocolFailure {
InvalidSse(SseError),
EmptyStream,
NoCommitBoundary,
}
pub(super) struct PendingStreamResponse {
pub account: Account,
pub attempt_request: CodexResponsesRequest,
pub attempt: ResponseDispatchAttempt,
pub lease: AccountLease,
pub response: CodexBackendStreamingResponse,
pub prefetched: Bytes,
pub decoder: CanonicalStreamDecoder,
pub initial_batch: CanonicalStreamBatch,
pub first_failure: Option<ResponsesSseFailure>,
pub first_event_ms: i64,
}
pub(super) struct PendingCompleteResponse {
pub account: Account,
pub attempt_request: CodexResponsesRequest,
pub attempt: ResponseDispatchAttempt,
pub response: CodexBackendResponse,
pub collected: CollectedResponse,
}
pub(super) enum PendingAttempt {
NoCandidate,
PinnedCandidateUnavailable {
account_id: String,
kind: PinnedCandidateAcquireFailureKind,
},
CandidatePreparationRejected {
account: Account,
},
RoutePreparationRejected {
account: Account,
message: String,
},
UpstreamFailure {
account: Account,
attempt_request: CodexResponsesRequest,
attempt: ResponseDispatchAttempt,
transport: CodexBackendTransport,
error: CodexClientError,
},
ProtocolFailure {
account: Account,
attempt_request: CodexResponsesRequest,
attempt: ResponseDispatchAttempt,
transport: CodexBackendTransport,
error: PendingProtocolFailure,
},
CompleteResponse(Box<PendingCompleteResponse>),
StreamResponse(Box<PendingStreamResponse>),
}
+202
View File
@@ -0,0 +1,202 @@
//! Responses live stream 的唯一结算出口。
use std::time::{Duration, Instant};
use bytes::Bytes;
use tokio::sync::mpsc;
use tokio::time::timeout;
use crate::{
dispatch::{
controllers::{ControllerRequestScope, ControllerSet, StreamControllerContext},
errors::ResponseDispatchStreamError,
failure::sse::{STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE},
lifecycle::stream::StreamSummary,
transport::canonical::{CanonicalResponseChunk, CanonicalResponseEvent},
},
fleet::pool::AccountLease,
upstream::openai::protocol::sse::{
DONE_SSE_FRAME, encode_sse_event, response_failed_sse_data_with_id,
},
};
const LEASE_FINALIZE_TIMEOUT: Duration = Duration::from_millis(100);
/// 消费 stream context,类型层面保证每个 live loop 只能结算一次。
pub(in crate::dispatch) struct StreamFinalizer {
controllers: ControllerSet,
controller_scope: ControllerRequestScope,
controller_context: StreamControllerContext,
account_lease: AccountLease,
}
impl StreamFinalizer {
pub(in crate::dispatch) fn new(
controllers: ControllerSet,
controller_scope: ControllerRequestScope,
controller_context: StreamControllerContext,
account_lease: AccountLease,
) -> Self {
Self {
controllers,
controller_scope,
controller_context,
account_lease,
}
}
pub(in crate::dispatch) fn response_headers(&self) -> &[(String, String)] {
&self.controller_context.response_metadata.client_headers
}
pub(in crate::dispatch) fn started_at(&self) -> Instant {
self.controller_context.started_at
}
pub(in crate::dispatch) async fn finalize(
self,
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
mut summary: StreamSummary,
) {
let Self {
controllers,
controller_scope,
controller_context,
account_lease,
} = self;
let body = String::from_utf8_lossy(&summary.body).into_owned();
controllers
.leave_stream(controller_scope, controller_context, &summary, &body)
.await;
if timeout(LEASE_FINALIZE_TIMEOUT, account_lease.complete())
.await
.is_err()
{
tracing::warn!(
timeout_ms = LEASE_FINALIZE_TIMEOUT.as_millis(),
"Timed out finalizing response account lease"
);
}
finish_client_stream(sender, event_sender, &mut summary).await;
}
}
async fn finish_client_stream(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
summary: &mut StreamSummary,
) {
if !summary.terminal.should_finish_client_stream() {
return;
}
if flush_terminal_chunks(sender, event_sender, &mut summary.terminal_chunks)
.await
.is_err()
{
return;
}
if let Some(detail) = summary.terminal.synthetic_failure_detail() {
if append_separator_if_needed(sender, event_sender, &mut summary.body)
.await
.is_err()
{
return;
}
let failure = response_failed_sse_data_with_id(
summary.last_response_id.as_deref(),
"server_error",
STREAM_DISCONNECTED_CODE,
&synthetic_failure_message(detail),
);
let bytes = Bytes::from(encode_sse_event("response.failed", &failure.to_string()));
if append_client_chunk(
sender,
event_sender,
&mut summary.body,
CanonicalResponseChunk::new(
bytes,
vec![CanonicalResponseEvent::proxy_failure(failure)],
),
)
.await
.is_err()
{
return;
}
}
if !body_has_done(&summary.body) {
let _ = append_client_chunk(
sender,
event_sender,
&mut summary.body,
CanonicalResponseChunk::new(Bytes::from_static(DONE_SSE_FRAME.as_bytes()), Vec::new()),
)
.await;
}
}
async fn flush_terminal_chunks(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
terminal_chunks: &mut Vec<CanonicalResponseChunk>,
) -> Result<(), ()> {
for chunk in terminal_chunks.drain(..) {
let (bytes, events) = chunk.into_parts();
let _ = event_sender.send(events);
sender.send(Ok(bytes)).await.map_err(|_| ())?;
}
Ok(())
}
async fn append_separator_if_needed(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
body: &mut Vec<u8>,
) -> Result<(), ()> {
let separator = if body.is_empty()
|| body.ends_with(b"\n\n")
|| body.ends_with(b"\r\n\r\n")
|| body.ends_with(b"\r\r")
{
return Ok(());
} else if body.ends_with(b"\n") || body.ends_with(b"\r") {
Bytes::from_static(b"\n")
} else {
Bytes::from_static(b"\n\n")
};
append_client_chunk(
sender,
event_sender,
body,
CanonicalResponseChunk::new(separator, Vec::new()),
)
.await
}
async fn append_client_chunk(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
body: &mut Vec<u8>,
chunk: CanonicalResponseChunk,
) -> Result<(), ()> {
body.extend_from_slice(chunk.bytes());
let (bytes, events) = chunk.into_parts();
let _ = event_sender.send(events);
sender.send(Ok(bytes)).await.map_err(|_| ())
}
fn body_has_done(body: &[u8]) -> bool {
String::from_utf8_lossy(body)
.trim_end_matches(['\r', '\n'])
.ends_with(DONE_SSE_FRAME.trim_end_matches(['\r', '\n']))
}
fn synthetic_failure_message(detail: &str) -> String {
if detail.trim().is_empty() {
STREAM_DISCONNECTED_MESSAGE.to_string()
} else {
format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}")
}
}
+13
View File
@@ -0,0 +1,13 @@
//! `v1/*` 请求的生命周期基础类型与入口准备。
pub(in crate::dispatch) mod attempt;
pub(in crate::dispatch) mod contract;
pub(in crate::dispatch) mod finalizer;
pub(in crate::dispatch) mod pipeline;
pub(in crate::dispatch) mod request;
pub(in crate::dispatch) mod stream;
pub(in crate::dispatch) mod trace;
pub(in crate::dispatch) use request::{
RequestContext, RequestEnterDependencies, RequestMode, enter_request,
};
@@ -0,0 +1,53 @@
//! Attempt 生命周期的唯一静态 pipeline。
use crate::{
dispatch::{controllers::ControllerRequestScope, routing::candidates::AccountAttemptLedger},
upstream::openai::protocol::responses::CodexResponsesRequest,
};
use super::{
attempt::{AttemptMode, AttemptRunner, AttemptRunnerDependencies, AttemptStep},
contract::{AttemptApplyOutcome, AttemptContractError, AttemptRejection, EstablishedResponse},
};
pub(in crate::dispatch) enum AttemptPipelineOutcome {
Established(EstablishedResponse),
Rejected(Box<AttemptRejection>),
}
/// 执行完整的 attempt/retry/commit 状态机。
///
/// mode 适配器只能消费最终结果,不能自行复制 retry loop 或越过 commit 边界。
pub(in crate::dispatch) async fn run_attempt_pipeline(
dependencies: AttemptRunnerDependencies<'_>,
mode: AttemptMode,
request: CodexResponsesRequest,
controller_scope: ControllerRequestScope,
candidates: AccountAttemptLedger,
) -> Result<AttemptPipelineOutcome, AttemptContractError> {
let controllers = dependencies.controllers;
let mut runner = AttemptRunner::new(dependencies, mode, request, controller_scope, candidates);
loop {
let outcome = match runner.next().await? {
AttemptStep::Open(attempt) => {
let decision = controllers
.handle_attempt(attempt.controller_scope()?, attempt.observation())
.await;
attempt.apply(decision).await?
}
AttemptStep::Committed(attempt) => {
return Ok(AttemptPipelineOutcome::Established(attempt.accept()));
}
};
match outcome {
AttemptApplyOutcome::Continue => {}
AttemptApplyOutcome::Established(established) => {
return Ok(AttemptPipelineOutcome::Established(established));
}
AttemptApplyOutcome::Rejected(rejection) => {
return Ok(AttemptPipelineOutcome::Rejected(rejection));
}
}
}
}
+101
View File
@@ -0,0 +1,101 @@
//! Responses 请求进入 dispatch 的共享生命周期阶段。
use chrono::Utc;
use serde_json::Value;
use crate::{
dispatch::{
affinity::{AccountIdentityService, prepare_variant_identity},
controllers::{ControllerEnter, ControllerRequestScope, ControllerSet},
routing::candidates::AccountAttemptLedger,
},
fleet::pool::{AccountAcquireRequest, AccountPoolService},
models::service::ModelService,
upstream::openai::protocol::responses::CodexResponsesRequest,
};
/// Responses 请求的执行模式。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(in crate::dispatch) enum RequestMode {
Complete,
Stream,
}
impl RequestMode {
fn stream(self) -> bool {
matches!(self, Self::Stream)
}
}
/// 请求 enter 阶段所需的窄依赖集合。
pub(in crate::dispatch) struct RequestEnterDependencies<'a> {
pub account_pool: &'a AccountPoolService,
pub models: &'a ModelService,
pub account_identity: &'a AccountIdentityService,
pub controllers: ControllerSet,
}
/// 已完成请求级 enter 的 Responses 上下文。
///
/// 该类型只承载请求、会话恢复、会话策略和候选快照;账号 attempt
/// 与流生命周期由后续 pipeline 接管。
pub(in crate::dispatch) struct RequestContext {
pub request: CodexResponsesRequest,
pub display_model: String,
pub compact: bool,
pub tuple_schema: Option<Value>,
pub image_generation_requested: bool,
pub controllers: ControllerSet,
pub controller_scope: ControllerRequestScope,
pub candidates: AccountAttemptLedger,
}
/// 执行所有 Responses 模式共享的请求 enter 阶段。
pub(in crate::dispatch) async fn enter_request(
dependencies: RequestEnterDependencies<'_>,
mut request: CodexResponsesRequest,
requested_model: &str,
mode: RequestMode,
) -> RequestContext {
let catalog = dependencies.models.catalog().await;
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
request.set_stream(mode.stream());
let compact = request.semantics().compact;
let tuple_schema = request.tuple_schema.clone();
let image_generation_requested = request.expects_image_generation();
let now = Utc::now();
prepare_variant_identity(&mut request);
dependencies
.account_identity
.prepare_local_identity(&mut request);
let controllers = dependencies.controllers;
let ControllerEnter {
scope,
preferred_account_id,
excluded_account_ids,
} = controllers.enter(&request, now).await;
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
if let Some(preferred_account_id) = preferred_account_id {
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
}
let candidates = AccountAttemptLedger::freeze(
dependencies.account_pool,
&acquire_request,
&excluded_account_ids,
)
.await;
RequestContext {
request,
display_model,
compact,
tuple_schema,
image_generation_requested,
controllers,
controller_scope: scope,
candidates,
}
}
+85
View File
@@ -0,0 +1,85 @@
//! Responses live stream 的规范事件与终态。
use serde_json::Value;
use crate::{
dispatch::transport::canonical::{CanonicalResponseChunk, CanonicalStreamTerminal},
upstream::openai::protocol::{events::TokenUsage, responses::ResponsesSseFailure},
};
/// live stream 离开事件循环的唯一终态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub(in crate::dispatch) enum StreamTerminal {
Completed { response: Value },
Incomplete { response: Value },
Failed { failure: ResponsesSseFailure },
UpstreamClosed,
UpstreamError { detail: String },
ProtocolError { detail: String },
CaptureLimitExceeded,
Cancelled,
DownstreamClosed,
Shutdown,
}
/// live loop 交给唯一 finalizer 的完整事实快照。
pub(in crate::dispatch) struct StreamSummary {
pub terminal: StreamTerminal,
pub body: Vec<u8>,
pub terminal_chunks: Vec<CanonicalResponseChunk>,
pub first_token_ms: Option<i64>,
pub first_event_ms: i64,
pub usage: Option<TokenUsage>,
pub last_response_id: Option<String>,
}
impl From<CanonicalStreamTerminal> for StreamTerminal {
fn from(terminal: CanonicalStreamTerminal) -> Self {
match terminal {
CanonicalStreamTerminal::Completed(response) => Self::Completed { response },
CanonicalStreamTerminal::Incomplete(response) => Self::Incomplete { response },
CanonicalStreamTerminal::Failed(failure) => Self::Failed { failure },
}
}
}
impl StreamTerminal {
/// 只有上游没有给出业务终态时,代理才需要合成 `response.failed`。
pub(in crate::dispatch) fn synthetic_failure_detail(&self) -> Option<&str> {
match self {
Self::UpstreamClosed => Some(""),
Self::UpstreamError { detail } | Self::ProtocolError { detail } => Some(detail),
Self::CaptureLimitExceeded => {
Some("upstream response exceeded the 16 MiB proxy capture limit")
}
Self::Completed { .. }
| Self::Incomplete { .. }
| Self::Failed { .. }
| Self::Cancelled
| Self::DownstreamClosed
| Self::Shutdown => None,
}
}
pub(in crate::dispatch) fn should_finish_client_stream(&self) -> bool {
!matches!(
self,
Self::Cancelled | Self::DownstreamClosed | Self::Shutdown
)
}
pub(in crate::dispatch) fn name(&self) -> &'static str {
match self {
Self::Completed { .. } => "completed",
Self::Incomplete { .. } => "incomplete",
Self::Failed { .. } => "failed",
Self::UpstreamClosed => "upstream_closed",
Self::UpstreamError { .. } => "upstream_error",
Self::ProtocolError { .. } => "protocol_error",
Self::CaptureLimitExceeded => "capture_limit_exceeded",
Self::Cancelled => "cancelled",
Self::DownstreamClosed => "downstream_closed",
Self::Shutdown => "shutdown",
}
}
}
@@ -1,4 +1,4 @@
//! Request-level trace state for Responses dispatch.
//! Responses 请求级账号尝试轨迹。
#[derive(Debug, Default)]
pub(in crate::dispatch) struct ResponseDispatchTrace {
+5 -4
View File
@@ -1,10 +1,11 @@
//! 请求调度模块(账号选择、回退、恢复、会话亲和性)。
pub mod affinity;
mod attempts;
pub(crate) mod controllers;
pub mod errors;
mod recording;
pub mod recovery;
pub(crate) mod failure;
pub(crate) mod lifecycle;
pub(crate) mod routing;
pub(crate) mod service;
pub(crate) mod stream;
pub mod upstream_call;
pub(crate) mod transport;
@@ -1,143 +0,0 @@
//! 可换号上游失败的账号风险剥离。
use chrono::Utc;
use crate::{
dispatch::{
errors::{
is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error,
is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body,
upstream_error_http_status,
},
recovery::{
auth::{auth_failure_account_status, is_auth_upstream_error},
cloudflare::{
CloudflareRecovery, cloudflare_challenge_error_message,
cloudflare_path_block_error_message, is_cloudflare_challenge_upstream_error,
is_cloudflare_path_block_upstream_error,
},
exhaustion::AccountExhaustionTracker,
},
stream::sse_failure::{
auth_sse_failure_account_status, is_auth_sse_failure, is_model_unsupported_sse_failure,
is_quota_exhausted_sse_failure, sse_failure_error_body, stream_failure_http_status,
},
},
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{
protocol::responses::ResponsesSseFailure,
transport::{CodexClientError, is_banned_upstream_error},
},
};
/// 处理需要隔离当前账号并交回调度器选择下一候选的上游失败。
///
/// 本函数不选择下一个账号,也不实现固定轮换策略;调用方继续使用请求级候选账本,
/// 因而始终服从当前运行时调度策略。
pub(crate) async fn isolate_rotatable_account_failure(
account_pool: &AccountPoolService,
cloudflare: &CloudflareRecovery,
exhausted_accounts: &mut AccountExhaustionTracker,
account_id: &str,
error: &CodexClientError,
) -> bool {
if is_rate_limit_upstream_error(error) {
exhausted_accounts.record_rate_limited(Some(account_id), upstream_error_body(error));
account_pool
.mark_quota_limited_until(account_id, rate_limit_cooldown_until(error, Utc::now()))
.await;
return true;
}
if is_quota_exhausted_upstream_error(error) {
exhausted_accounts.record_quota_exhausted(Some(account_id), upstream_error_body(error));
account_pool
.set_status(account_id, AccountStatus::QuotaExhausted)
.await;
return true;
}
if is_auth_upstream_error(error) {
let account_status = auth_failure_account_status(error);
exhausted_accounts.record_auth_failure(
Some(account_id),
account_status,
upstream_error_body(error),
Some(upstream_error_http_status(error)),
);
account_pool.set_status(account_id, account_status).await;
return true;
}
if is_cloudflare_challenge_upstream_error(error) {
exhausted_accounts
.record_cloudflare_challenge(Some(account_id), cloudflare_challenge_error_message());
cloudflare.apply_challenge(account_pool, account_id).await;
return true;
}
if is_cloudflare_path_block_upstream_error(error) {
exhausted_accounts.record_cloudflare_path_blocked(
Some(account_id),
cloudflare_path_block_error_message(),
);
cloudflare.apply_path_block(account_pool, account_id).await;
return true;
}
if is_model_unsupported_upstream_error(error) {
exhausted_accounts.record_model_unsupported(Some(account_id), upstream_error_body(error));
return true;
}
if is_banned_upstream_error(error) {
exhausted_accounts.record_auth_failure(
Some(account_id),
AccountStatus::Banned,
upstream_error_body(error),
Some(upstream_error_http_status(error)),
);
account_pool
.set_status(account_id, AccountStatus::Banned)
.await;
return true;
}
false
}
/// 处理 SSE `response.failed` 中需要隔离当前账号的失败。
pub(crate) async fn isolate_sse_account_failure(
account_pool: &AccountPoolService,
exhausted_accounts: &mut AccountExhaustionTracker,
account_id: &str,
failure: &ResponsesSseFailure,
) -> bool {
if is_model_unsupported_sse_failure(failure) {
exhausted_accounts
.record_model_unsupported(Some(account_id), sse_failure_error_body(failure));
return true;
}
if is_quota_exhausted_sse_failure(failure) {
exhausted_accounts.record_quota_exhausted(Some(account_id), failure.message.clone());
account_pool
.set_status(account_id, AccountStatus::QuotaExhausted)
.await;
return true;
}
if is_auth_sse_failure(failure) {
let account_status = auth_sse_failure_account_status(failure);
exhausted_accounts.record_auth_failure(
Some(account_id),
account_status,
sse_failure_error_body(failure),
Some(stream_failure_http_status(failure)),
);
account_pool.set_status(account_id, account_status).await;
return true;
}
false
}
-22
View File
@@ -1,22 +0,0 @@
//! 上游认证失败分类。
use crate::{
fleet::account::AccountStatus,
upstream::openai::transport::{CodexClientError, is_banned_auth_signal},
};
pub(crate) fn is_auth_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, .. } if status.as_u16() == 401
)
}
pub(crate) fn auth_failure_account_status(error: &CodexClientError) -> AccountStatus {
match error {
CodexClientError::Upstream { body, .. } if is_banned_auth_signal(body) => {
AccountStatus::Banned
}
_ => AccountStatus::Expired,
}
}
-146
View File
@@ -1,146 +0,0 @@
//! Cloudflare challenge and path-block recovery shared by dispatch routes.
use chrono::Utc;
use crate::{
fleet::{
account::AccountStatus,
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
pool::AccountPoolService,
},
upstream::openai::transport::CodexClientError,
};
#[derive(Clone)]
pub struct CloudflareRecovery {
path_block_tracker: CloudflarePathBlockTracker,
challenge_tracker: CloudflareChallengeCooldownTracker,
cookie_store: PgCookieStore,
}
impl CloudflareRecovery {
pub fn new(cookie_store: PgCookieStore) -> Self {
Self {
path_block_tracker: CloudflarePathBlockTracker::default(),
challenge_tracker: CloudflareChallengeCooldownTracker::default(),
cookie_store,
}
}
pub async fn cookie_header_for_request(&self, account_id: &str, path: &str) -> Option<String> {
self.cookie_store
.cookie_header_for_request(account_id, "chatgpt.com", path)
.await
.ok()?
}
pub async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) {
for header in headers {
if let Err(error) = self
.cookie_store
.capture_set_cookie(account_id, header)
.await
{
tracing::warn!(
account_id,
error = %error,
"failed to persist upstream set-cookie header"
);
}
}
}
pub async fn apply_challenge(&self, account_pool: &AccountPoolService, account_id: &str) {
let now = Utc::now();
let cooldown = self
.challenge_tracker
.record_challenge(account_id, now)
.await;
account_pool
.set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until)
.await;
if let Err(error) = self
.cookie_store
.expire_account_cookies_at(account_id, cooldown.cooldown_until)
.await
{
tracing::warn!(
account_id,
error = %error,
"failed to persist Cloudflare challenge cookie cleanup deadline"
);
}
}
pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) {
self.delete_account_cookies(account_id, "Cloudflare path-block")
.await;
let now = Utc::now();
self.path_block_tracker
.record_path_block(account_id, now)
.await;
if self
.path_block_tracker
.should_disable(account_id, now)
.await
{
account_pool
.set_status(account_id, AccountStatus::Disabled)
.await;
}
}
pub async fn reset_account_recovery(&self, account_id: &str) {
self.path_block_tracker.reset(account_id).await;
self.challenge_tracker.reset(account_id).await;
}
async fn delete_account_cookies(&self, account_id: &str, reason: &str) {
delete_account_cookies(&self.cookie_store, account_id, reason).await;
}
}
async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) {
if let Err(error) = cookie_store.delete_account_cookies(account_id).await {
tracing::warn!(
account_id,
reason,
error = %error,
"failed to delete account cookies after Cloudflare recovery signal"
);
}
}
pub fn is_cloudflare_challenge_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, body, .. }
if status.as_u16() == 403 && is_cloudflare_challenge_signal(body)
)
}
pub fn is_cloudflare_path_block_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, body, .. }
if status.as_u16() == 404 && body.trim().is_empty()
)
}
pub fn cloudflare_challenge_error_message() -> &'static str {
"Upstream blocked the request (Cloudflare challenge)"
}
pub fn cloudflare_path_block_error_message() -> &'static str {
"Upstream blocked the request (Cloudflare path-block)"
}
fn is_cloudflare_challenge_signal(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("cf-mitigated")
|| value.contains("cf-chl-bypass")
|| value.contains("_cf_chl")
|| value.contains("cf_chl")
|| value.contains("attention required")
|| value.contains("just a moment")
}
-192
View File
@@ -1,192 +0,0 @@
//! previous response 所有权与连接内上下文重放。
use std::sync::Arc;
use chrono::Utc;
use serde_json::Value;
use crate::{
dispatch::affinity::{SessionAffinityEntry, SessionAffinityService},
upstream::openai::protocol::responses::{CodexResponsesRequest, StreamCommitPolicy},
};
/// 当前请求的 previous response 来源。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(in crate::dispatch) enum HistorySource {
None,
Managed,
ConnectionLocal,
ExternalUnknown,
}
/// 请求级 previous response 恢复计划。
pub(in crate::dispatch) struct HistoryRecoveryPlan {
source: HistorySource,
managed_entry: Option<SessionAffinityEntry>,
local_replay_input: Option<Arc<Vec<Value>>>,
external_attempted_account_id: Option<String>,
replay_owner_account: bool,
}
impl HistoryRecoveryPlan {
pub(in crate::dispatch) async fn load(
session_affinity: &SessionAffinityService,
request: &CodexResponsesRequest,
) -> Self {
let Some(previous_response_id) = request.previous_response_id() else {
return Self {
source: HistorySource::None,
managed_entry: None,
local_replay_input: None,
external_attempted_account_id: None,
replay_owner_account: false,
};
};
let managed_entry = session_affinity
.lookup(previous_response_id, Utc::now())
.await;
let local_replay_input = request.local_replay_input.clone();
let source = match (&managed_entry, &local_replay_input) {
(Some(_), _) => HistorySource::Managed,
(None, Some(_)) => HistorySource::ConnectionLocal,
(None, None) => HistorySource::ExternalUnknown,
};
Self {
source,
managed_entry,
local_replay_input,
external_attempted_account_id: None,
replay_owner_account: false,
}
}
pub(in crate::dispatch) fn preferred_account_id(&self) -> Option<&str> {
self.managed_entry
.as_ref()
.map(|entry| entry.account_id.as_str())
}
pub(in crate::dispatch) fn prepare_attempt(
&mut self,
original: &CodexResponsesRequest,
account_id: &str,
) -> Option<CodexResponsesRequest> {
match self.source {
HistorySource::None => Some(original.clone()),
HistorySource::ExternalUnknown => match self.external_attempted_account_id.as_deref() {
None => {
self.external_attempted_account_id = Some(account_id.to_string());
Some(external_history_request(original))
}
Some(first_account_id) if first_account_id == account_id => {
Some(external_history_request(original))
}
Some(_) => None,
},
HistorySource::ConnectionLocal => self
.local_replay_input
.as_deref()
.map(|replay_input| full_replay_request(original, replay_input)),
HistorySource::Managed => {
let entry = self.managed_entry.as_ref()?;
if entry.account_id == account_id && !self.replay_owner_account {
let mut request = original.clone();
request.local_conversation_id = Some(entry.conversation_id.clone());
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
request.previous_response_scope = Some(entry.continuation_scope);
if request.turn_state.as_deref().is_none_or(str::is_empty) {
request.turn_state.clone_from(&entry.turn_state);
}
Some(request)
} else {
self.local_replay_input.as_deref().map(|replay_input| {
let mut request = full_replay_request(original, replay_input);
request.local_conversation_id = Some(entry.conversation_id.clone());
request
})
}
}
}
}
/// previous ID 在原连接不可用时,切换为同账号完整重放。
pub(in crate::dispatch) fn recover_managed_history(&mut self, account_id: &str) -> bool {
let Some(entry) = self.managed_entry.as_ref() else {
return false;
};
if entry.account_id != account_id
|| self.local_replay_input.is_none()
|| self.replay_owner_account
{
return false;
}
self.replay_owner_account = true;
true
}
pub(in crate::dispatch) fn can_failover(&self) -> bool {
match self.source {
HistorySource::None | HistorySource::ConnectionLocal => true,
HistorySource::Managed => self.local_replay_input.is_some(),
HistorySource::ExternalUnknown => false,
}
}
pub(in crate::dispatch) fn is_external_unknown(&self) -> bool {
self.source == HistorySource::ExternalUnknown
}
pub(in crate::dispatch) fn conversation_id<'a>(
&'a self,
original: &'a CodexResponsesRequest,
) -> Option<&'a str> {
self.managed_entry
.as_ref()
.map(|entry| entry.conversation_id.as_str())
.or_else(|| non_empty(original.local_conversation_id.as_deref()))
}
}
fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesRequest {
let mut request = original.clone();
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
request.previous_response_scope =
Some(crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown);
request
}
fn full_replay_request(
original: &CodexResponsesRequest,
replay_input: &[Value],
) -> CodexResponsesRequest {
let mut request = original.clone();
let mut input = Vec::with_capacity(replay_input.len() + original.input().len());
input.extend_from_slice(replay_input);
input.extend_from_slice(original.input());
sanitize_replay_items(&mut input);
request.set_input(input);
request.set_previous_response_id(None);
request.turn_state = None;
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
request
}
pub(crate) fn sanitize_replay_items(values: &mut [Value]) {
for value in values {
match value {
Value::Array(values) => sanitize_replay_items(values),
Value::Object(object) => {
object.remove("id");
object.remove("encrypted_content");
for value in object.values_mut() {
sanitize_replay_items(std::slice::from_mut(value));
}
}
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
}
}
}
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
-7
View File
@@ -1,7 +0,0 @@
//! 调度失败恢复策略。
pub(crate) mod account_failure;
pub(crate) mod auth;
pub mod cloudflare;
pub(crate) mod exhaustion;
pub(crate) mod history;
@@ -26,8 +26,10 @@ impl AccountAttemptLedger {
pub(in crate::dispatch) async fn freeze(
account_pool: &AccountPoolService,
request: &AccountAcquireRequest,
excluded_account_ids: &BTreeSet<String>,
) -> Self {
let candidate_ids = account_pool.candidate_snapshot(request).await;
let mut candidate_ids = account_pool.candidate_snapshot(request).await;
candidate_ids.retain(|account_id| !excluded_account_ids.contains(account_id));
Self {
model: request.model.clone(),
pending: candidate_ids.iter().cloned().collect(),
@@ -80,7 +82,7 @@ impl AccountAttemptLedger {
attempted_count = self.attempted.len(),
busy_count = self.busy.len(),
timeout_ms = ACCOUNT_CANDIDATE_WAIT_TIMEOUT.as_millis(),
"timed out waiting for a busy account candidate"
"Timed out waiting for a busy account candidate"
);
return None;
}
+3
View File
@@ -0,0 +1,3 @@
//! 请求级通用账号候选路由。
pub(in crate::dispatch) mod candidates;
+209 -489
View File
@@ -10,69 +10,44 @@
use std::{pin::Pin, sync::Arc, time::Instant};
use bytes::Bytes;
use chrono::Utc;
use futures::stream::Stream;
use serde_json::{Value, json};
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use crate::{
dispatch::{
affinity::resolve::{prepare_variant_identity, record_response_affinity},
affinity::{AccountIdentityService, SessionAffinityService},
attempts::AccountAttemptLedger,
errors::{
backend_transport_name, is_continuation_busy_error, is_history_recovery_upstream_error,
is_retryable_account_transport_error, upstream_error_body,
upstream_error_set_cookie_headers,
},
recovery::{
account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure},
controllers::{
CompleteExit, ControllerSet, DispatchErrorObservation,
cloudflare::CloudflareRecovery,
exhaustion::AccountExhaustionTracker,
history::HistoryRecoveryPlan,
history::{
ConnectionReplayPlan as HistoryConnectionReplayPlan,
ConnectionReplaySnapshot as HistoryConnectionReplaySnapshot, HistoryController,
},
},
upstream_call::{
AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext,
QuotaVerificationDecision, create_response_with_account,
verify_acquired_quota_if_required,
errors::backend_transport_name,
lifecycle::{
RequestContext, RequestEnterDependencies, RequestMode,
attempt::{AttemptMode, AttemptRunnerDependencies},
contract::{
AttemptContractError, EstablishedAttemptContext, EstablishedComplete,
EstablishedCompleteBody, EstablishedResponse, FinalOutcome,
},
enter_request,
pipeline::{AttemptPipelineOutcome, run_attempt_pipeline},
},
transport::canonical::CanonicalResponseEvent,
},
fleet::{
account::Account,
pool::{AccountAcquireRequest, AccountPoolService},
},
fleet::pool::AccountPoolService,
models::service::ModelService,
telemetry::{
recorder::{
Recorder, enrich_response_request_semantics, reasoning_effort_from_request,
record_response_event,
},
usage::types::ResponseUsageRecord,
},
telemetry::recorder::Recorder,
upstream::openai::{
protocol::{
events::TokenUsage,
responses::{CodexResponsesRequest, CollectedResponse, response_from_codex_sse},
},
transport::{
CodexBackendClient, CodexBackendResponse, backend_transport_for_response_request,
},
protocol::{responses::CodexResponsesRequest, sse::SseError},
transport::CodexBackendClient,
},
};
use super::{
errors::{ResponseDispatchError, ResponseDispatchStreamError},
recording::{
ResponseDispatchErrorDetails, ResponseDispatchErrorEventRecord,
ResponseUpstreamErrorEventRecord, insert_response_status_metadata,
insert_response_trace_metadata, insert_response_upstream_diagnostics,
insert_websocket_pool_decision, record_response_dispatch_error_event,
record_response_upstream_error_event,
},
stream::{
sse_failure::is_history_recovery_sse_failure,
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
};
use super::errors::{ResponseDispatchError, ResponseDispatchStreamError};
/// OpenAI Responses 调度服务。
#[derive(Clone)]
@@ -84,6 +59,7 @@ pub struct ResponseDispatchService {
pub(in crate::dispatch) account_identity: Arc<AccountIdentityService>,
pub(in crate::dispatch) recorder: Arc<Recorder>,
pub(in crate::dispatch) cloudflare: CloudflareRecovery,
pub(in crate::dispatch) shutdown: CancellationToken,
}
pub(crate) struct ResponseDispatchServiceParts {
@@ -94,6 +70,7 @@ pub(crate) struct ResponseDispatchServiceParts {
pub account_identity: Arc<AccountIdentityService>,
pub recorder: Arc<Recorder>,
pub cloudflare: CloudflareRecovery,
pub shutdown: CancellationToken,
}
/// Responses live SSE 响应体流。
@@ -103,6 +80,7 @@ pub type ResponseBodyStream =
/// Responses live SSE 调度结果。
pub struct ResponseDispatchStream {
pub body: ResponseBodyStream,
pub(crate) canonical_events: tokio::sync::mpsc::UnboundedReceiver<Vec<CanonicalResponseEvent>>,
pub response_headers: Vec<(String, String)>,
}
@@ -112,6 +90,43 @@ pub struct ResponseDispatchResponse {
pub response_headers: Vec<(String, String)>,
}
/// 入站 Responses WebSocket 连接持有的窄重放快照。
pub(crate) struct ConnectionReplaySnapshot {
inner: HistoryConnectionReplaySnapshot,
}
/// 单次 `response.create` 对快照的待提交重放计划。
pub(crate) struct ConnectionReplayPlan {
inner: HistoryConnectionReplayPlan,
}
/// API 从 canonical 响应事件采集的原始 transcript 事实。
pub(crate) struct ConnectionTranscriptFacts {
response_id: String,
output: Vec<Value>,
}
impl ConnectionTranscriptFacts {
pub(crate) fn new(response_id: String, output: Vec<Value>) -> Self {
Self {
response_id,
output,
}
}
}
fn attempt_contract_error(error: AttemptContractError) -> ResponseDispatchError {
ResponseDispatchError::InvalidSse(SseError::ParseError(format!(
"attempt lifecycle contract violation: {error}"
)))
}
fn attempt_mode_mismatch(expected: &str) -> ResponseDispatchError {
ResponseDispatchError::InvalidSse(SseError::ParseError(format!(
"attempt lifecycle returned a response for the wrong mode; expected {expected}"
)))
}
impl ResponseDispatchService {
pub(crate) fn new(parts: ResponseDispatchServiceParts) -> Self {
Self {
@@ -122,471 +137,176 @@ impl ResponseDispatchService {
account_identity: parts.account_identity,
recorder: parts.recorder,
cloudflare: parts.cloudflare,
shutdown: parts.shutdown,
}
}
async fn record_response_affinity(
pub(crate) fn connection_replay_snapshot(&self) -> ConnectionReplaySnapshot {
ConnectionReplaySnapshot {
inner: HistoryController::new_connection_replay_snapshot(),
}
}
pub(crate) fn prepare_connection_replay(
&self,
history: &HistoryRecoveryPlan,
original_request: &CodexResponsesRequest,
account_id: &str,
body: &str,
turn_state: Option<String>,
usage: Option<TokenUsage>,
snapshot: &ConnectionReplaySnapshot,
request: &mut CodexResponsesRequest,
) -> ConnectionReplayPlan {
ConnectionReplayPlan {
inner: HistoryController::prepare_connection_replay(&snapshot.inner, request),
}
}
pub(crate) fn commit_connection_replay(
&self,
snapshot: &mut ConnectionReplaySnapshot,
plan: ConnectionReplayPlan,
transcript: ConnectionTranscriptFacts,
) {
record_response_affinity(
&self.session_affinity,
history,
original_request,
account_id,
body,
turn_state,
usage,
)
.await;
HistoryController::commit_connection_replay(
&mut snapshot.inner,
plan.inner,
transcript.response_id,
transcript.output,
);
}
}
impl ResponseDispatchService {
/// 调度非流式 Responses 请求到 Codex Responses 上游。
pub async fn complete(
&self,
request_id: &str,
route: &str,
mut request: CodexResponsesRequest,
request: CodexResponsesRequest,
requested_model: &str,
) -> Result<ResponseDispatchResponse, ResponseDispatchError> {
let started_at = Instant::now();
let catalog = self.models.catalog().await;
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
let compact = request.semantics().compact;
let tuple_schema = request.tuple_schema.clone();
let image_generation_requested = request.expects_image_generation();
let now = Utc::now();
prepare_variant_identity(&mut request);
self.account_identity.prepare_local_identity(&mut request);
let mut history = HistoryRecoveryPlan::load(&self.session_affinity, &request).await;
let preferred_account_id = self
.preferred_account_id_for_request(&request, &history, now)
.await;
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
if let Some(preferred_account_id) = preferred_account_id {
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
}
let mut candidates =
AccountAttemptLedger::freeze(&self.account_pool, &acquire_request).await;
let mut exhausted_accounts = AccountExhaustionTracker::default();
let mut next_required_account_id: Option<String> = None;
let mut trace = ResponseDispatchTrace::default();
let (account, response, collected_response, attempt): (
Account,
CodexBackendResponse,
CollectedResponse,
ResponseDispatchAttempt,
) = loop {
let acquired = if let Some(account_id) = next_required_account_id.take() {
match self
.account_pool
.acquire_with(
&AccountAcquireRequest::new(request.model(), Utc::now())
.with_required_account_id(account_id),
)
.await
{
Some(acquired) => Some(acquired),
None => candidates.acquire_next(&self.account_pool).await,
}
} else {
candidates.acquire_next(&self.account_pool).await
};
let Some(acquired) = acquired else {
tracing::info!(
candidate_count = candidates.candidate_count(),
attempted = candidates.attempted_count(),
state_excluded = candidates.state_excluded_count(),
"Responses account candidate ledger exhausted"
);
let error = exhausted_accounts
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount);
self.record_response_dispatch_error(
request_id,
route,
requested_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: exhausted_accounts.last_account_id(),
stream: false,
compact,
transport: Some(backend_transport_name(
backend_transport_for_response_request(&request),
)),
},
&error,
)
.await;
return Err(error);
};
let acquired_account_id = acquired.account.id.clone();
// 配额验证
let acquired = match verify_acquired_quota_if_required(
QuotaVerificationContext {
account_pool: self.account_pool.as_ref(),
codex: self.codex.as_ref(),
cloudflare: &self.cloudflare,
account_identity: &self.account_identity,
request_id,
},
acquired,
)
.await
{
QuotaVerificationDecision::Ready(acquired) => *acquired,
QuotaVerificationDecision::RetryWithAnotherAccount => {
exhausted_accounts.record_rate_limited(
Some(&acquired_account_id),
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
);
continue;
}
};
let Some(attempt_request) = history.prepare_attempt(&request, &acquired.account.id)
else {
acquired.release_without_usage().await;
return Err(ResponseDispatchError::HistoryUnavailable {
upstream_error: "previous response history cannot be sent to another account"
.to_string(),
});
};
self.account_pool.wait_for_request_interval(&acquired).await;
let account = acquired.account.clone();
let release_account_id = account.id.clone();
let attempt = trace.start_attempt(&release_account_id);
let response_result = create_response_with_account(
AccountUpstreamContext {
codex: &self.codex,
account_identity: &self.account_identity,
cloudflare: &self.cloudflare,
request_id,
account: &account,
},
&attempt_request,
let controllers = ControllerSet::new(
Arc::clone(&self.session_affinity),
Arc::clone(&self.account_pool),
Arc::clone(&self.codex),
self.cloudflare.clone(),
Arc::clone(&self.recorder),
);
let context = enter_request(
RequestEnterDependencies {
account_pool: &self.account_pool,
models: &self.models,
account_identity: &self.account_identity,
controllers,
},
request,
requested_model,
RequestMode::Complete,
)
.await;
let RequestContext {
request,
display_model,
compact,
tuple_schema,
image_generation_requested,
controllers,
controller_scope,
candidates,
} = context;
let established = match run_attempt_pipeline(
AttemptRunnerDependencies {
account_pool: self.account_pool.as_ref(),
codex: self.codex.as_ref(),
controllers: &controllers,
account_identity: self.account_identity.as_ref(),
request_id,
route,
requested_model,
started_at,
)
.await;
acquired.complete().await;
if let Err(error) = &response_result {
self.cloudflare
.capture_set_cookie_headers(
&release_account_id,
upstream_error_set_cookie_headers(error),
},
AttemptMode::Complete {
started_at,
tuple_schema,
},
request,
controller_scope,
candidates,
)
.await
.map_err(attempt_contract_error)?
{
AttemptPipelineOutcome::Established(established) => match established {
EstablishedResponse::Complete(established) => *established,
EstablishedResponse::Stream(_) => {
return Err(attempt_mode_mismatch("complete"));
}
},
AttemptPipelineOutcome::Rejected(rejection) => {
let rejection = *rejection;
controllers
.observe_dispatch_error(
DispatchErrorObservation {
request_id,
client_api_key_id: rejection.request.client_api_key_id.as_deref(),
account_id: rejection.account_id.as_deref(),
route,
model: requested_model,
started_at,
stream: false,
compact,
transport: Some(backend_transport_name(rejection.transport)),
},
&rejection.error,
)
.await;
}
match response_result {
Ok(response) => {
self.cloudflare
.capture_set_cookie_headers(
&release_account_id,
&response.set_cookie_headers,
)
.await;
self.account_pool
.sync_passive_rate_limit_headers(&account, &response.rate_limit_headers)
.await;
let collected_response =
match response_from_codex_sse(&response.body, tuple_schema.as_ref()) {
Ok(collected_response) => collected_response,
Err(error) => {
let error = ResponseDispatchError::InvalidSse(error);
self.record_response_dispatch_error(
request_id,
route,
requested_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&release_account_id),
stream: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
)
.await;
return Err(error);
}
};
if matches!(collected_response, CollectedResponse::Empty) {
self.account_pool
.record_empty_response_attempt(
&release_account_id,
image_generation_requested,
)
.await;
exhausted_accounts.record_upstream_unavailable(
Some(&release_account_id),
"upstream response did not include visible output",
);
if !history.can_failover() {
break (account, response, collected_response, attempt);
}
continue;
}
if let CollectedResponse::Failed(failure) = &collected_response {
if is_history_recovery_sse_failure(failure)
&& history.recover_managed_history(&release_account_id)
{
next_required_account_id = Some(release_account_id);
continue;
}
if isolate_sse_account_failure(
self.account_pool.as_ref(),
&mut exhausted_accounts,
&release_account_id,
failure,
)
.await
{
if !history.can_failover() {
break (account, response, collected_response, attempt);
}
continue;
}
}
break (account, response, collected_response, attempt);
}
Err(error) => {
if isolate_rotatable_account_failure(
self.account_pool.as_ref(),
&self.cloudflare,
&mut exhausted_accounts,
&release_account_id,
&error,
)
.await
{
if !history.can_failover() {
return Err(ResponseDispatchError::Upstream(error));
}
continue;
}
if is_continuation_busy_error(&error) {
if history.recover_managed_history(&release_account_id) {
next_required_account_id = Some(release_account_id);
continue;
}
return Err(ResponseDispatchError::ContinuationBusy);
}
let history_unavailable = is_history_recovery_upstream_error(&error);
if history_unavailable && history.recover_managed_history(&release_account_id) {
next_required_account_id = Some(release_account_id);
continue;
}
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
recorder: &self.recorder,
request_id,
account_id: &release_account_id,
account_email: account.email.as_deref(),
route,
model: requested_model,
started_at,
stream: false,
transport: backend_transport_for_response_request(&attempt_request),
request: &attempt_request,
error: &error,
trace: &trace,
attempt: Some(&attempt),
})
.await;
if is_retryable_account_transport_error(&error) {
if !history.can_failover() {
return Err(ResponseDispatchError::Upstream(error));
}
exhausted_accounts.record_upstream_unavailable(
Some(&release_account_id),
upstream_error_body(&error),
);
continue;
}
if history_unavailable {
if history.is_external_unknown() {
return Err(ResponseDispatchError::Upstream(error));
}
return Err(ResponseDispatchError::HistoryUnavailable {
upstream_error: upstream_error_body(&error),
});
}
return Err(ResponseDispatchError::Upstream(error));
}
let error = rejection.error;
return Err(error);
}
};
let completed = matches!(&collected_response, CollectedResponse::Completed(_));
match collected_response {
CollectedResponse::Completed(body) | CollectedResponse::Incomplete(body) => {
let response_id = body.get("id").and_then(Value::as_str);
self.cloudflare.reset_account_recovery(&account.id).await;
if let Some(usage) = response.usage {
self.account_pool
.record_response_usage(&account.id, usage, image_generation_requested)
.await;
}
if completed {
self.record_response_affinity(
&history,
&request,
&account.id,
&response.body,
response.turn_state.clone(),
response.usage,
)
.await;
}
let effective_model = response
.response_metadata
.effective_model
.as_deref()
.unwrap_or(&display_model);
let mut metadata = json!({
"responseId": response_id,
"stream": false,
"completed": completed,
"incomplete": !completed,
"transport": backend_transport_name(response.transport),
"firstTokenMs": response.first_token_ms,
"usage": response.usage,
"effectiveModel": effective_model,
"modelsEtag": response.response_metadata.models_etag.as_deref(),
"reasoningIncluded": response.response_metadata.reasoning_included,
});
insert_response_status_metadata(
&mut metadata,
200,
200,
response.diagnostics.status_code.map(i64::from),
);
insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics);
insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt));
insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision);
enrich_response_request_semantics(&mut metadata, &request);
record_response_event(ResponseUsageRecord {
recorder: &self.recorder,
request_id,
client_api_key_id: request.client_api_key_id.as_deref(),
let EstablishedComplete {
context,
response,
body,
} = established;
let EstablishedAttemptContext {
request,
controller_scope,
trace,
account,
attempt,
..
} = context;
let (outcome, body) = match body {
EstablishedCompleteBody::Completed(body) => (FinalOutcome::Completed, body),
EstablishedCompleteBody::Incomplete(body) => (FinalOutcome::Incomplete, body),
};
let response_id = body.get("id").and_then(Value::as_str);
let completed = outcome == FinalOutcome::Completed;
controllers
.finalize_complete(
controller_scope,
CompleteExit {
request: &request,
account_id: &account.id,
route,
model: effective_model,
requested_model: Some(requested_model),
client_ip: request.client_ip.as_deref(),
client_user_agent: request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(&request),
service_tier: request.service_tier(),
started_at,
status_code: 200,
message: if completed {
"v1 responses completed"
} else {
"v1 responses incomplete"
},
metadata,
rate_limit_headers: &response.rate_limit_headers,
})
.await;
self.models
.observe_models_etag(response.response_metadata.models_etag.as_deref());
Ok(ResponseDispatchResponse {
body,
response_headers: response.response_metadata.client_headers,
})
}
CollectedResponse::Failed(failure) => {
let error = ResponseDispatchError::Failed(failure);
self.record_response_dispatch_error(
body: &response.body,
turn_state: response.turn_state.clone(),
usage: response.usage,
image_generation_requested,
completed,
recorder: &self.recorder,
models: &self.models,
request_id,
route,
requested_model,
display_model: &display_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&account.id),
stream: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
)
.await;
Err(error)
}
CollectedResponse::MissingCompleted => {
let error = ResponseDispatchError::MissingCompleted;
self.record_response_dispatch_error(
request_id,
route,
requested_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&account.id),
stream: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
)
.await;
Err(error)
}
CollectedResponse::Empty => {
let error = ResponseDispatchError::EmptyUpstreamResponse;
self.record_response_dispatch_error(
request_id,
route,
requested_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&account.id),
stream: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
)
.await;
Err(error)
}
}
}
pub(super) async fn record_response_dispatch_error(
&self,
request_id: &str,
route: &str,
requested_model: &str,
started_at: Instant,
details: ResponseDispatchErrorDetails<'_>,
error: &ResponseDispatchError,
) {
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
recorder: &self.recorder,
request_id,
client_api_key_id: details.client_api_key_id,
account_id: details.account_id,
route,
model: requested_model,
started_at,
stream: details.stream,
compact: details.compact,
transport: details.transport,
error,
response_id,
response: &response,
trace: &trace,
attempt: &attempt,
},
outcome,
)
.await;
Ok(ResponseDispatchResponse {
body,
response_headers: response.response_metadata.client_headers,
})
.await;
}
}
+185 -459
View File
@@ -1,46 +1,29 @@
use std::{sync::Arc, time::Instant};
//! 流式 Responses 在提交前共享的 request/attempt 生命周期。
use chrono::Utc;
use std::{sync::Arc, time::Instant};
use crate::{
dispatch::{
affinity::prepare_variant_identity,
attempts::AccountAttemptLedger,
errors::{
ResponseDispatchError, backend_transport_name, is_continuation_busy_error,
is_history_recovery_upstream_error, is_retryable_account_transport_error,
upstream_error_body, upstream_error_set_cookie_headers,
controllers::{
ControllerSet, DispatchErrorObservation, PrefetchedStreamFailureObservation,
StreamControllerContext,
},
recording::{
ResponseDispatchErrorDetails, ResponseStreamFailureEventRecord,
ResponseUpstreamErrorEventRecord, record_prefetched_response_stream_failure_event,
record_response_upstream_error_event,
},
recovery::{
account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure},
exhaustion::AccountExhaustionTracker,
history::HistoryRecoveryPlan,
errors::{ResponseDispatchError, backend_transport_name},
lifecycle::{
RequestContext, RequestEnterDependencies, RequestMode,
attempt::{AttemptMode, AttemptRunnerDependencies},
contract::{
AttemptContractError, EstablishedAttemptContext, EstablishedResponse,
EstablishedStream,
},
enter_request,
finalizer::StreamFinalizer,
pipeline::{AttemptPipelineOutcome, run_attempt_pipeline},
},
service::{ResponseDispatchService, ResponseDispatchStream},
stream::{
live::{LiveResponseStreamContext, spawn_live_response_stream},
prefetch::prefetch_first_sse_chunk,
sse_failure::{
first_sse_failure, is_history_recovery_sse_failure, sse_failure_error_body,
},
trace::ResponseDispatchTrace,
},
upstream_call::{
AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext,
QuotaVerificationDecision, create_response_stream_with_account,
verify_acquired_quota_if_required,
},
},
fleet::pool::AccountAcquireRequest,
upstream::openai::{
protocol::responses::CodexResponsesRequest,
transport::backend_transport_for_response_request,
stream::live::spawn_live_response_stream,
},
upstream::openai::protocol::{responses::CodexResponsesRequest, sse::SseError},
};
impl ResponseDispatchService {
@@ -49,437 +32,180 @@ impl ResponseDispatchService {
&self,
request_id: &str,
route: &str,
mut request: CodexResponsesRequest,
request: CodexResponsesRequest,
requested_model: &str,
) -> Result<ResponseDispatchStream, ResponseDispatchError> {
let started_at = Instant::now();
let catalog = self.models.catalog().await;
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
request.set_stream(true);
let compact = request.semantics().compact;
let tuple_schema = request.tuple_schema.clone();
let now = Utc::now();
prepare_variant_identity(&mut request);
self.account_identity.prepare_local_identity(&mut request);
let mut history = HistoryRecoveryPlan::load(&self.session_affinity, &request).await;
let preferred_account_id = self
.preferred_account_id_for_request(&request, &history, now)
.await;
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
if let Some(preferred_account_id) = preferred_account_id {
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
}
let mut candidates =
AccountAttemptLedger::freeze(&self.account_pool, &acquire_request).await;
let mut exhausted_accounts = AccountExhaustionTracker::default();
let mut next_required_account_id: Option<String> = None;
let mut trace = ResponseDispatchTrace::default();
macro_rules! return_stream_dispatch_error {
($error:expr) => {{
let error = $error;
self.record_response_dispatch_error(
request_id,
route,
requested_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: exhausted_accounts.last_account_id(),
stream: true,
compact,
transport: Some(backend_transport_name(
backend_transport_for_response_request(&request),
)),
},
&error,
)
.await;
return Err(error);
}};
($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{
let error = $error;
self.record_response_dispatch_error(
request_id,
route,
requested_model,
started_at,
ResponseDispatchErrorDetails {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: $account_id,
stream: true,
compact,
transport: $transport,
},
&error,
)
.await;
return Err(error);
}};
}
loop {
let acquired = if let Some(account_id) = next_required_account_id.take() {
match self
.account_pool
.acquire_with(
&AccountAcquireRequest::new(request.model(), Utc::now())
.with_required_account_id(account_id),
)
.await
{
Some(acquired) => Some(acquired),
None => candidates.acquire_next(&self.account_pool).await,
}
} else {
candidates.acquire_next(&self.account_pool).await
};
let Some(acquired) = acquired else {
tracing::info!(
candidate_count = candidates.candidate_count(),
attempted = candidates.attempted_count(),
state_excluded = candidates.state_excluded_count(),
"Responses stream account candidate ledger exhausted"
);
let error = exhausted_accounts
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount);
return_stream_dispatch_error!(error);
};
let acquired_account_id = acquired.account.id.clone();
let acquired = match verify_acquired_quota_if_required(
QuotaVerificationContext {
account_pool: self.account_pool.as_ref(),
codex: self.codex.as_ref(),
cloudflare: &self.cloudflare,
account_identity: &self.account_identity,
request_id,
},
acquired,
)
.await
{
QuotaVerificationDecision::Ready(acquired) => *acquired,
QuotaVerificationDecision::RetryWithAnotherAccount => {
exhausted_accounts.record_rate_limited(
Some(&acquired_account_id),
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
);
continue;
}
};
let Some(attempt_request) = history.prepare_attempt(&request, &acquired.account.id)
else {
acquired.release_without_usage().await;
return_stream_dispatch_error!(ResponseDispatchError::HistoryUnavailable {
upstream_error: "previous response history cannot be sent to another account"
.to_string(),
});
};
self.account_pool.wait_for_request_interval(&acquired).await;
let account = acquired.account.clone();
let release_account_id = account.id.clone();
let attempt = trace.start_attempt(&release_account_id);
let response_result = create_response_stream_with_account(
AccountUpstreamContext {
codex: &self.codex,
account_identity: &self.account_identity,
cloudflare: &self.cloudflare,
request_id,
account: &account,
},
&attempt_request,
)
.await;
if let Err(error) = &response_result {
self.cloudflare
.capture_set_cookie_headers(
&release_account_id,
upstream_error_set_cookie_headers(error),
)
.await;
}
match response_result {
Ok(response) => {
let transport = response.transport;
let set_cookie_headers = response.set_cookie_headers;
let rate_limit_headers = response.rate_limit_headers;
let rate_limit_header_updates = response.rate_limit_header_updates;
let turn_state_update = response.turn_state_update;
let websocket_pool_decision = response.websocket_pool_decision;
let turn_state = response.turn_state;
let diagnostics = response.diagnostics;
let response_metadata = response.response_metadata;
self.models
.observe_models_etag(response_metadata.models_etag.as_deref());
let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await {
Ok(prefetched) => prefetched,
Err(ResponseDispatchError::Upstream(error)) => {
acquired.complete().await;
if isolate_rotatable_account_failure(
self.account_pool.as_ref(),
&self.cloudflare,
&mut exhausted_accounts,
&release_account_id,
&error,
)
.await
{
if history.is_external_unknown() {
return_stream_dispatch_error!(
ResponseDispatchError::Upstream(error),
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(transport))
);
}
continue;
}
if is_continuation_busy_error(&error) {
if history.recover_managed_history(&release_account_id) {
next_required_account_id = Some(release_account_id);
continue;
}
return_stream_dispatch_error!(
ResponseDispatchError::ContinuationBusy,
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(transport))
);
}
let history_unavailable = is_history_recovery_upstream_error(&error);
if history_unavailable
&& history.recover_managed_history(&release_account_id)
{
next_required_account_id = Some(release_account_id);
continue;
}
record_response_upstream_error_event(
ResponseUpstreamErrorEventRecord {
recorder: &self.recorder,
request_id,
account_id: &release_account_id,
account_email: account.email.as_deref(),
route,
model: requested_model,
started_at,
stream: true,
transport,
request: &attempt_request,
error: &error,
trace: &trace,
attempt: Some(&attempt),
},
)
.await;
if is_retryable_account_transport_error(&error) {
exhausted_accounts.record_upstream_unavailable(
Some(&release_account_id),
upstream_error_body(&error),
);
if history.is_external_unknown() {
return Err(ResponseDispatchError::Upstream(error));
}
continue;
}
if history_unavailable {
if history.is_external_unknown() {
return Err(ResponseDispatchError::Upstream(error));
}
return Err(ResponseDispatchError::HistoryUnavailable {
upstream_error: upstream_error_body(&error),
});
}
return Err(ResponseDispatchError::Upstream(error));
}
Err(error) => {
acquired.complete().await;
return_stream_dispatch_error!(
error,
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(transport))
);
}
};
let first_failure = match first_sse_failure(&prefetched) {
Ok(failure) => failure,
Err(error) => {
acquired.complete().await;
return_stream_dispatch_error!(
ResponseDispatchError::InvalidSse(error),
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(transport))
);
}
};
if let Some(failure) = first_failure {
if is_history_recovery_sse_failure(&failure)
&& history.recover_managed_history(&release_account_id)
{
acquired.complete().await;
next_required_account_id = Some(release_account_id);
continue;
}
if isolate_sse_account_failure(
self.account_pool.as_ref(),
&mut exhausted_accounts,
&release_account_id,
&failure,
)
.await
{
acquired.complete().await;
if history.is_external_unknown() {
return_stream_dispatch_error!(
ResponseDispatchError::Failed(failure),
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(transport))
);
}
continue;
}
acquired.complete().await;
record_prefetched_response_stream_failure_event(
ResponseStreamFailureEventRecord {
recorder: &self.recorder,
request_id,
account_id: &release_account_id,
route,
model: &display_model,
requested_model,
started_at,
transport,
request: &attempt_request,
failure: &failure,
diagnostics: &diagnostics,
rate_limit_headers: &rate_limit_headers,
prefetched: &prefetched,
trace: &trace,
attempt: &attempt,
},
)
let controllers = ControllerSet::new(
Arc::clone(&self.session_affinity),
Arc::clone(&self.account_pool),
Arc::clone(&self.codex),
self.cloudflare.clone(),
Arc::clone(&self.recorder),
);
let context = enter_request(
RequestEnterDependencies {
account_pool: &self.account_pool,
models: &self.models,
account_identity: &self.account_identity,
controllers,
},
request,
requested_model,
RequestMode::Stream,
)
.await;
let RequestContext {
request,
display_model,
compact,
tuple_schema,
controllers,
controller_scope,
candidates,
..
} = context;
let established = match run_attempt_pipeline(
AttemptRunnerDependencies {
account_pool: self.account_pool.as_ref(),
codex: self.codex.as_ref(),
controllers: &controllers,
account_identity: self.account_identity.as_ref(),
request_id,
route,
requested_model,
started_at,
},
AttemptMode::Stream {
tuple_schema: tuple_schema.clone(),
},
request,
controller_scope,
candidates,
)
.await
.map_err(attempt_contract_error)?
{
AttemptPipelineOutcome::Established(established) => match established {
EstablishedResponse::Stream(established) => *established,
EstablishedResponse::Complete(_) => return Err(attempt_mode_mismatch()),
},
AttemptPipelineOutcome::Rejected(rejection) => {
let rejection = *rejection;
if let (Some(account), Some(attempt_request), Some(attempt), Some(stream_failure)) = (
rejection.account.as_ref(),
rejection.attempt_request.as_ref(),
rejection.attempt.as_ref(),
rejection.stream_failure.as_ref(),
) {
controllers
.observe_prefetched_stream_failure(PrefetchedStreamFailureObservation {
request_id,
account_id: &account.id,
route,
model: &display_model,
requested_model,
started_at,
transport: rejection.transport,
request: attempt_request,
failure: &stream_failure.failure,
error: &rejection.error,
diagnostics: &stream_failure.diagnostics,
rate_limit_headers: &stream_failure.rate_limit_headers,
prefetched: &stream_failure.prefetched,
trace: &rejection.trace,
attempt,
})
.await;
let error = if is_history_recovery_sse_failure(&failure)
&& !history.is_external_unknown()
{
ResponseDispatchError::HistoryUnavailable {
upstream_error: sse_failure_error_body(&failure),
}
} else {
ResponseDispatchError::Failed(failure.clone())
};
return Err(error);
}
let context = LiveResponseStreamContext {
account_pool: Arc::clone(&self.account_pool),
account_lease: acquired,
session_affinity: Arc::clone(&self.session_affinity),
history,
recorder: Arc::clone(&self.recorder),
cloudflare: self.cloudflare.clone(),
account_id: account.id,
account_plan_type: account.plan_type,
request_id: request_id.to_string(),
route: route.to_string(),
display_model: display_model.clone(),
requested_model: requested_model.to_string(),
client_ip: request.client_ip.clone(),
request,
tuple_schema,
transport,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates,
turn_state_update,
websocket_pool_decision,
turn_state,
diagnostics,
response_metadata,
attempt: attempt.clone(),
attempts: trace.attempts().to_vec(),
started_at,
};
return Ok(spawn_live_response_stream(context, prefetched, body));
}
Err(error) => {
acquired.complete().await;
if isolate_rotatable_account_failure(
self.account_pool.as_ref(),
&self.cloudflare,
&mut exhausted_accounts,
&release_account_id,
&error,
controllers
.observe_dispatch_error(
DispatchErrorObservation {
request_id,
client_api_key_id: rejection.request.client_api_key_id.as_deref(),
account_id: rejection.account_id.as_deref(),
route,
model: requested_model,
started_at,
stream: true,
compact,
transport: Some(backend_transport_name(rejection.transport)),
},
&rejection.error,
)
.await
{
if history.is_external_unknown() {
return_stream_dispatch_error!(
ResponseDispatchError::Upstream(error),
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(
backend_transport_for_response_request(&attempt_request)
))
);
}
continue;
}
if is_continuation_busy_error(&error) {
if history.recover_managed_history(&release_account_id) {
next_required_account_id = Some(release_account_id);
continue;
}
return_stream_dispatch_error!(
ResponseDispatchError::ContinuationBusy,
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(
backend_transport_for_response_request(&attempt_request)
))
);
}
let history_unavailable = is_history_recovery_upstream_error(&error);
if history_unavailable && history.recover_managed_history(&release_account_id) {
next_required_account_id = Some(release_account_id);
continue;
}
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
recorder: &self.recorder,
request_id,
account_id: &release_account_id,
account_email: account.email.as_deref(),
route,
model: requested_model,
started_at,
stream: true,
transport: backend_transport_for_response_request(&attempt_request),
request: &attempt_request,
error: &error,
trace: &trace,
attempt: Some(&attempt),
})
.await;
if is_retryable_account_transport_error(&error) {
exhausted_accounts.record_upstream_unavailable(
Some(&release_account_id),
upstream_error_body(&error),
);
if history.is_external_unknown() {
return Err(ResponseDispatchError::Upstream(error));
}
continue;
}
if history_unavailable {
if history.is_external_unknown() {
return Err(ResponseDispatchError::Upstream(error));
}
return Err(ResponseDispatchError::HistoryUnavailable {
upstream_error: upstream_error_body(&error),
});
}
return Err(ResponseDispatchError::Upstream(error));
}
let error = rejection.error;
return Err(error);
}
}
};
let EstablishedStream {
context,
lease,
response,
decoder,
initial_batch,
first_event_ms,
} = established;
let EstablishedAttemptContext {
request,
controller_scope,
trace,
account,
attempt,
..
} = context;
self.models
.observe_models_etag(response.response_metadata.models_etag.as_deref());
let attempts = trace.attempts().to_vec();
let transport = response.transport;
let finalizer = StreamFinalizer::new(
controllers,
controller_scope,
StreamControllerContext {
account_id: account.id,
account_plan_type: account.plan_type,
request_id: request_id.to_string(),
route: route.to_string(),
display_model,
requested_model: requested_model.to_string(),
request,
transport,
set_cookie_headers: response.set_cookie_headers,
rate_limit_headers: response.rate_limit_headers,
rate_limit_header_updates: response.rate_limit_header_updates,
turn_state_update: response.turn_state_update,
websocket_pool_decision: response.websocket_pool_decision,
turn_state: response.turn_state,
diagnostics: response.diagnostics,
response_metadata: response.response_metadata,
transport_metrics: response.transport_metrics,
connection_local_continuation: response.connection_local_continuation,
attempt,
attempts,
started_at,
},
lease,
);
Ok(spawn_live_response_stream(
finalizer,
decoder,
initial_batch,
first_event_ms,
response.body,
self.shutdown.clone(),
))
}
}
fn attempt_contract_error(error: AttemptContractError) -> ResponseDispatchError {
ResponseDispatchError::InvalidSse(SseError::ParseError(format!(
"attempt lifecycle contract violation: {error}"
)))
}
fn attempt_mode_mismatch() -> ResponseDispatchError {
ResponseDispatchError::InvalidSse(SseError::ParseError(
"attempt lifecycle returned a non-stream response for stream mode".to_string(),
))
}
+248 -587
View File
@@ -1,61 +1,32 @@
use std::{
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Instant,
};
use bytes::Bytes;
use futures::{StreamExt, stream::Stream};
use serde_json::Value;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use crate::{
dispatch::{
affinity::{SessionAffinityService, resolve::record_response_affinity},
errors::{ResponseDispatchStreamError, backend_transport_name},
recording::{
insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state,
record_live_response_stream_event,
errors::ResponseDispatchStreamError,
lifecycle::{
finalizer::StreamFinalizer,
stream::{StreamSummary, StreamTerminal},
},
recovery::{cloudflare::CloudflareRecovery, history::HistoryRecoveryPlan},
service::ResponseDispatchStream,
transport::canonical::{
CanonicalResponseChunk, CanonicalResponseEvent, CanonicalStreamBatch,
CanonicalStreamDecoder,
},
},
fleet::pool::{AccountLease, AccountPoolService},
infra::time::elapsed_millis_i64,
telemetry::{recorder::Recorder, usage::types::UsageRecordLevel},
upstream::openai::{
protocol::{
events::extract_sse_usage,
responses::{
CodexResponsesRequest, CollectedResponse,
reconvert_responses_sse_event_tuple_values, response_body_has_first_output,
response_from_codex_sse, response_sse_event_is_terminal,
},
sse::{
DONE_SSE_FRAME, encode_sse_event, parse_sse_events,
response_failed_sse_event_with_id, sse_body_has_done, sse_frame_end,
},
},
transport::{
CodexBackendSseStream, CodexBackendTransport, CodexRateLimitHeaderUpdates,
CodexResponseMetadata, CodexTurnStateUpdate, CodexUpstreamDiagnostics,
WebSocketPoolDecision,
},
},
upstream::openai::transport::CodexBackendSseStream,
};
const MAX_LIVE_RESPONSE_CAPTURE_BYTES: usize = 16 * 1024 * 1024;
const LIVE_RESPONSE_CAPTURE_LIMIT_DETAIL: &str =
"upstream response exceeded the 16 MiB proxy capture limit";
use super::{
sse_failure::{
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, status_code_for_stream_failure,
stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail,
},
trace::ResponseDispatchAttempt,
};
struct MpscResponseBodyStream {
receiver: mpsc::Receiver<Result<Bytes, ResponseDispatchStreamError>>,
@@ -79,144 +50,34 @@ impl Stream for MpscResponseBodyStream {
}
pub(in crate::dispatch) fn spawn_live_response_stream(
context: LiveResponseStreamContext,
prefetched: Bytes,
mut body: CodexBackendSseStream,
finalizer: StreamFinalizer,
decoder: CanonicalStreamDecoder,
initial_batch: CanonicalStreamBatch,
first_event_ms: i64,
body: CodexBackendSseStream,
shutdown: CancellationToken,
) -> ResponseDispatchStream {
let response_headers = context.response_metadata.client_headers.clone();
let response_headers = finalizer.response_headers().to_vec();
let started_at = finalizer.started_at();
let (sender, receiver) = mpsc::channel(8);
let (cancel_sender, mut cancel_receiver) = oneshot::channel();
let (event_sender, canonical_events) = mpsc::unbounded_channel();
let (cancel_sender, cancel_receiver) = oneshot::channel();
tokio::spawn(async move {
let mut tuple_transformer = context
.tuple_schema
.clone()
.map(TupleSseEventTransformer::new);
let mut body_bytes = Vec::new();
let mut first_token_ms = None;
let mut first_output_detector = FirstOutputDetector::default();
first_output_detector.observe(context.started_at, &prefetched, &mut first_token_ms);
match send_live_response_stream_chunk(
let summary = run_live_response_stream(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
prefetched,
&event_sender,
LiveRunInputs {
initial_batch,
first_event_ms,
body,
cancel: cancel_receiver,
decoder,
started_at,
shutdown,
},
)
.await
{
Ok(()) => {}
Err(LiveStreamWriteError::DownstreamClosed) => {
context.account_lease.complete().await;
return;
}
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
terminate_oversized_live_stream(context, &sender, &mut body_bytes, first_token_ms)
.await;
return;
}
}
loop {
let next = tokio::select! {
_ = &mut cancel_receiver => {
context.account_lease.complete().await;
return;
}
next = body.next() => next,
};
let Some(next) = next else {
break;
};
match next {
Ok(chunk) => {
first_output_detector.observe(context.started_at, &chunk, &mut first_token_ms);
match send_live_response_stream_chunk(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
chunk,
)
.await
{
Ok(()) => {}
Err(LiveStreamWriteError::DownstreamClosed) => {
context.account_lease.complete().await;
return;
}
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
terminate_oversized_live_stream(
context,
&sender,
&mut body_bytes,
first_token_ms,
)
.await;
return;
}
}
}
Err(error) => {
match flush_live_response_stream_transformer(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
)
.await
{
Ok(()) => {}
Err(LiveStreamWriteError::DownstreamClosed) => {
context.account_lease.complete().await;
return;
}
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
terminate_oversized_live_stream(
context,
&sender,
&mut body_bytes,
first_token_ms,
)
.await;
return;
}
}
let detail = error.to_string();
let Some(body_text) =
send_live_response_stream_tail(&sender, &mut body_bytes, Some(&detail))
.await
else {
context.account_lease.complete().await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
return;
}
}
}
match flush_live_response_stream_transformer(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
)
.await
{
Ok(()) => {}
Err(LiveStreamWriteError::DownstreamClosed) => {
context.account_lease.complete().await;
return;
}
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
terminate_oversized_live_stream(context, &sender, &mut body_bytes, first_token_ms)
.await;
return;
}
}
let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await
else {
context.account_lease.complete().await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
.await;
finalizer.finalize(&sender, &event_sender, summary).await;
});
ResponseDispatchStream {
@@ -224,445 +85,245 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
receiver,
cancel: Some(cancel_sender),
}),
canonical_events,
response_headers,
}
}
async fn send_live_response_stream_chunk(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
transformer: Option<&mut TupleSseEventTransformer>,
chunk: Bytes,
) -> Result<(), LiveStreamWriteError> {
let chunks = match transformer {
Some(transformer) => transformer.push(&chunk)?,
None => vec![chunk],
};
send_live_response_stream_chunks(sender, body_bytes, chunks).await
struct LiveRunInputs {
initial_batch: CanonicalStreamBatch,
first_event_ms: i64,
body: CodexBackendSseStream,
cancel: oneshot::Receiver<()>,
decoder: CanonicalStreamDecoder,
started_at: Instant,
shutdown: CancellationToken,
}
async fn flush_live_response_stream_transformer(
async fn run_live_response_stream(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
transformer: Option<&mut TupleSseEventTransformer>,
) -> Result<(), LiveStreamWriteError> {
let Some(transformer) = transformer else {
return Ok(());
};
send_live_response_stream_chunks(sender, body_bytes, transformer.finish()?).await
}
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
inputs: LiveRunInputs,
) -> StreamSummary {
let LiveRunInputs {
initial_batch,
first_event_ms,
mut body,
mut cancel,
mut decoder,
started_at,
shutdown,
} = inputs;
let mut body_bytes = Vec::new();
let mut terminal_chunks = Vec::new();
let mut first_token_ms = None;
async fn send_live_response_stream_chunks(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
chunks: Vec<Bytes>,
) -> Result<(), LiveStreamWriteError> {
for chunk in chunks {
if body_bytes.len().saturating_add(chunk.len()) > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
return Err(LiveStreamWriteError::CaptureLimitExceeded);
let decoded_terminal = take_stream_terminal(&mut decoder);
let terminal = if decoded_terminal.is_some() {
stage_terminal_batch(&mut body_bytes, &mut terminal_chunks, initial_batch)
.err()
.or(decoded_terminal)
} else {
forward_batch(sender, event_sender, &mut body_bytes, initial_batch)
.await
.err()
};
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
if let Some(terminal) = terminal {
return stream_summary(
terminal,
body_bytes,
terminal_chunks,
first_token_ms,
first_event_ms,
&decoder,
);
}
let mut shutdown = Box::pin(shutdown.cancelled_owned());
let terminal = loop {
let next = tokio::select! {
_ = &mut cancel => break StreamTerminal::Cancelled,
() = &mut shutdown => break StreamTerminal::Shutdown,
next = body.next() => next,
};
match next {
Some(Ok(chunk)) => match decoder.push(chunk) {
Ok(batch) => {
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
if let Some(terminal) = take_stream_terminal(&mut decoder) {
match stage_terminal_batch(&mut body_bytes, &mut terminal_chunks, batch) {
Ok(()) => break terminal,
Err(terminal) => break terminal,
}
}
if let Err(terminal) =
forward_batch(sender, event_sender, &mut body_bytes, batch).await
{
break terminal;
}
}
Err(error) => {
break StreamTerminal::ProtocolError {
detail: error.to_string(),
};
}
},
Some(Err(error)) => {
break finish_decoder(
sender,
event_sender,
&mut body_bytes,
&mut terminal_chunks,
&mut decoder,
StreamTerminal::UpstreamError {
detail: error.to_string(),
},
)
.await;
}
None => {
break finish_decoder(
sender,
event_sender,
&mut body_bytes,
&mut terminal_chunks,
&mut decoder,
StreamTerminal::UpstreamClosed,
)
.await;
}
}
body_bytes.extend_from_slice(&chunk);
if sender.send(Ok(chunk)).await.is_err() {
return Err(LiveStreamWriteError::DownstreamClosed);
};
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
stream_summary(
terminal,
body_bytes,
terminal_chunks,
first_token_ms,
first_event_ms,
&decoder,
)
}
async fn finish_decoder(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
body_bytes: &mut Vec<u8>,
terminal_chunks: &mut Vec<CanonicalResponseChunk>,
decoder: &mut CanonicalStreamDecoder,
fallback: StreamTerminal,
) -> StreamTerminal {
let batch = match decoder.finish() {
Ok(batch) => batch,
Err(error) => {
return StreamTerminal::ProtocolError {
detail: error.to_string(),
};
}
};
if let Some(terminal) = take_stream_terminal(decoder) {
return match stage_terminal_batch(body_bytes, terminal_chunks, batch) {
Ok(()) => terminal,
Err(terminal) => terminal,
};
}
if let Err(terminal) = forward_batch(sender, event_sender, body_bytes, batch).await {
return terminal;
}
fallback
}
fn take_stream_terminal(decoder: &mut CanonicalStreamDecoder) -> Option<StreamTerminal> {
let terminal = decoder.take_terminal().map(StreamTerminal::from);
let transport_done = decoder.take_transport_done();
terminal.or_else(|| transport_done.then_some(StreamTerminal::UpstreamClosed))
}
fn stage_terminal_batch(
body_bytes: &mut Vec<u8>,
terminal_chunks: &mut Vec<CanonicalResponseChunk>,
batch: CanonicalStreamBatch,
) -> Result<(), StreamTerminal> {
let batch_bytes = batch.chunks.iter().fold(0usize, |total, chunk| {
total.saturating_add(chunk.bytes().len())
});
ensure_capture_capacity(body_bytes.len(), batch_bytes)?;
for chunk in batch.chunks {
body_bytes.extend_from_slice(chunk.bytes());
terminal_chunks.push(chunk);
}
Ok(())
}
async fn forward_batch(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
event_sender: &mpsc::UnboundedSender<Vec<CanonicalResponseEvent>>,
body_bytes: &mut Vec<u8>,
batch: CanonicalStreamBatch,
) -> Result<(), StreamTerminal> {
for chunk in batch.chunks {
capture_chunk(body_bytes, &chunk)?;
let (bytes, events) = chunk.into_parts();
let _ = event_sender.send(events);
if sender.send(Ok(bytes)).await.is_err() {
return Err(StreamTerminal::DownstreamClosed);
}
}
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum LiveStreamWriteError {
DownstreamClosed,
CaptureLimitExceeded,
fn capture_chunk(
body_bytes: &mut Vec<u8>,
chunk: &CanonicalResponseChunk,
) -> Result<(), StreamTerminal> {
ensure_capture_capacity(body_bytes.len(), chunk.bytes().len())?;
body_bytes.extend_from_slice(chunk.bytes());
Ok(())
}
#[derive(Default)]
struct FirstOutputDetector {
pending: Vec<u8>,
}
impl FirstOutputDetector {
fn observe(&mut self, started_at: Instant, chunk: &[u8], first_token_ms: &mut Option<i64>) {
if first_token_ms.is_some() {
return;
}
self.pending.extend_from_slice(chunk);
while let Some(frame_end) = sse_frame_end(&self.pending) {
let frame = self.pending.drain(..frame_end).collect::<Vec<_>>();
if response_body_has_first_output(&frame) {
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
self.pending.clear();
return;
}
}
if self.pending.len() > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
self.pending.clear();
}
}
}
struct TupleSseEventTransformer {
tuple_schema: Value,
pending: Vec<u8>,
}
impl TupleSseEventTransformer {
fn new(tuple_schema: Value) -> Self {
Self {
tuple_schema,
pending: Vec::new(),
}
}
fn push(&mut self, chunk: &[u8]) -> Result<Vec<Bytes>, LiveStreamWriteError> {
self.pending.extend_from_slice(chunk);
if self.pending.len() > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
return Err(LiveStreamWriteError::CaptureLimitExceeded);
}
let mut chunks = Vec::new();
while let Some(frame_end) = sse_frame_end(&self.pending) {
let frame = self.pending.drain(..frame_end).collect::<Vec<_>>();
chunks.push(self.transform_frame(&frame));
}
Ok(chunks)
}
fn finish(&mut self) -> Result<Vec<Bytes>, LiveStreamWriteError> {
if self.pending.is_empty() {
return Ok(Vec::new());
}
let frame = std::mem::take(&mut self.pending);
Ok(vec![self.transform_frame(&frame)])
}
fn transform_frame(&self, frame: &[u8]) -> Bytes {
let frame_text = String::from_utf8_lossy(frame);
let Ok(events) = parse_sse_events(&frame_text) else {
return Bytes::copy_from_slice(frame);
};
let [event] = events.as_slice() else {
return Bytes::copy_from_slice(frame);
};
let Ok(data) = serde_json::from_str::<Value>(&event.data) else {
return Bytes::copy_from_slice(frame);
};
let transformed = reconvert_responses_sse_event_tuple_values(
event.event.as_deref(),
data,
&self.tuple_schema,
fn ensure_capture_capacity(
captured_bytes: usize,
incoming_bytes: usize,
) -> Result<(), StreamTerminal> {
if captured_bytes.saturating_add(incoming_bytes) > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
tracing::warn!(
captured_bytes,
incoming_bytes,
capture_limit_bytes = MAX_LIVE_RESPONSE_CAPTURE_BYTES,
"Live response capture limit exceeded"
);
Bytes::from(encode_sse_event(
event.event.as_deref().unwrap_or_default(),
&transformed.to_string(),
))
return Err(StreamTerminal::CaptureLimitExceeded);
}
Ok(())
}
async fn terminate_oversized_live_stream(
context: LiveResponseStreamContext,
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
first_token_ms: Option<i64>,
fn update_first_token_ms(
decoder: &CanonicalStreamDecoder,
started_at: Instant,
first_token_ms: &mut Option<i64>,
) {
tracing::warn!(
request_id = %context.request_id,
account_id = %context.account_id,
captured_bytes = body_bytes.len(),
capture_limit_bytes = MAX_LIVE_RESPONSE_CAPTURE_BYTES,
"live response capture limit exceeded"
);
let Some(body_text) = send_live_response_stream_tail(
sender,
body_bytes,
Some(LIVE_RESPONSE_CAPTURE_LIMIT_DETAIL),
)
.await
else {
context.account_lease.complete().await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
}
async fn send_live_response_stream_tail(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
failure_detail: Option<&str>,
) -> Option<String> {
let mut body_text = String::from_utf8_lossy(body_bytes).to_string();
if !sse_body_has_terminal_event(&body_text) {
if let Some(separator) = missing_sse_event_separator(&body_text) {
body_text.push_str(separator);
body_bytes.extend_from_slice(separator.as_bytes());
if sender
.send(Ok(Bytes::copy_from_slice(separator.as_bytes())))
.await
.is_err()
{
return None;
}
}
let failure =
premature_close_failed_event(latest_response_id(&body_text).as_deref(), failure_detail);
body_text.push_str(&failure);
body_bytes.extend_from_slice(failure.as_bytes());
if sender.send(Ok(Bytes::from(failure))).await.is_err() {
return None;
}
}
if !sse_body_has_done(&body_text) {
body_text.push_str(DONE_SSE_FRAME);
body_bytes.extend_from_slice(DONE_SSE_FRAME.as_bytes());
if sender
.send(Ok(Bytes::from_static(DONE_SSE_FRAME.as_bytes())))
.await
.is_err()
{
return None;
}
}
Some(body_text)
}
fn sse_body_has_terminal_event(body: &str) -> bool {
parse_sse_events(body).is_ok_and(|events| events.iter().any(response_sse_event_is_terminal))
}
fn missing_sse_event_separator(body: &str) -> Option<&'static str> {
if body.is_empty()
|| body.ends_with("\n\n")
|| body.ends_with("\r\n\r\n")
|| body.ends_with("\r\r")
{
None
} else if body.ends_with('\n') || body.ends_with('\r') {
Some("\n")
} else {
Some("\n\n")
if first_token_ms.is_none() && decoder.first_output_seen() {
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
}
}
pub(in crate::dispatch) struct LiveResponseStreamContext {
pub(in crate::dispatch) account_pool: Arc<AccountPoolService>,
pub(in crate::dispatch) account_lease: AccountLease,
pub(in crate::dispatch) session_affinity: Arc<SessionAffinityService>,
pub(in crate::dispatch) history: HistoryRecoveryPlan,
pub(in crate::dispatch) recorder: Arc<Recorder>,
pub(in crate::dispatch) cloudflare: CloudflareRecovery,
pub(in crate::dispatch) account_id: String,
pub(in crate::dispatch) account_plan_type: Option<String>,
pub(in crate::dispatch) request_id: String,
pub(in crate::dispatch) route: String,
pub(in crate::dispatch) display_model: String,
pub(in crate::dispatch) requested_model: String,
pub(in crate::dispatch) client_ip: Option<String>,
pub(in crate::dispatch) request: CodexResponsesRequest,
pub(in crate::dispatch) tuple_schema: Option<Value>,
pub(in crate::dispatch) transport: CodexBackendTransport,
pub(in crate::dispatch) set_cookie_headers: Vec<String>,
pub(in crate::dispatch) rate_limit_headers: Vec<(String, String)>,
pub(in crate::dispatch) rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
pub(in crate::dispatch) turn_state_update: Option<CodexTurnStateUpdate>,
pub(in crate::dispatch) websocket_pool_decision: Option<WebSocketPoolDecision>,
pub(in crate::dispatch) turn_state: Option<String>,
pub(in crate::dispatch) diagnostics: CodexUpstreamDiagnostics,
pub(in crate::dispatch) response_metadata: CodexResponseMetadata,
pub(in crate::dispatch) attempt: ResponseDispatchAttempt,
pub(in crate::dispatch) attempts: Vec<ResponseDispatchAttempt>,
pub(in crate::dispatch) started_at: Instant,
}
pub(in crate::dispatch) fn latest_response_id(body: &str) -> Option<String> {
parse_sse_events(body).ok().and_then(|events| {
events.iter().rev().find_map(|event| {
serde_json::from_str::<Value>(&event.data)
.ok()
.and_then(|data| {
data.pointer("/response/id")
.and_then(Value::as_str)
.filter(|id| !id.trim().is_empty())
.map(ToString::to_string)
})
})
})
}
pub(in crate::dispatch) fn premature_close_failed_event(
response_id: Option<&str>,
detail: Option<&str>,
) -> String {
let message = match detail.filter(|value| !value.trim().is_empty()) {
Some(detail) => format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}"),
None => STREAM_DISCONNECTED_MESSAGE.to_string(),
};
response_failed_sse_event_with_id(
response_id,
"server_error",
STREAM_DISCONNECTED_CODE,
&message,
)
}
pub(in crate::dispatch) async fn finalize_live_response_stream(
context: LiveResponseStreamContext,
body: String,
fn stream_summary(
terminal: StreamTerminal,
body: Vec<u8>,
terminal_chunks: Vec<CanonicalResponseChunk>,
first_token_ms: Option<i64>,
) {
context
.cloudflare
.capture_set_cookie_headers(&context.account_id, &context.set_cookie_headers)
.await;
let rate_limit_headers = live_response_rate_limit_headers(&context).await;
context
.account_pool
.sync_passive_rate_limit_headers_for_account(
&context.account_id,
context.account_plan_type.as_deref(),
&rate_limit_headers,
)
.await;
let turn_state = live_response_turn_state(&context).await;
let usage = match extract_sse_usage(&body) {
Ok(Some(usage)) => {
context
.account_pool
.record_token_usage(&context.account_id, &usage)
.await;
Some(usage)
}
Ok(None) => None,
Err(error) => {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to extract streaming token usage");
None
}
};
let collected_response = response_from_codex_sse(&body, context.tuple_schema.as_ref());
let completed_terminal = matches!(&collected_response, Ok(CollectedResponse::Completed(_)));
match collected_response {
Ok(CollectedResponse::Completed(terminal) | CollectedResponse::Incomplete(terminal)) => {
context
.cloudflare
.reset_account_recovery(&context.account_id)
.await;
let response_id = terminal.get("id").and_then(Value::as_str);
if completed_terminal {
record_response_affinity(
&context.session_affinity,
&context.history,
&context.request,
&context.account_id,
&body,
turn_state,
usage,
)
.await;
}
record_live_response_stream_event(
&context,
200,
UsageRecordLevel::Info,
if completed_terminal {
"v1 responses stream completed"
} else {
"v1 responses stream incomplete"
},
serde_json::json!({
"stream": true,
"completed": completed_terminal,
"incomplete": !completed_terminal,
"responseId": response_id,
"firstTokenMs": first_token_ms,
"usage": usage,
}),
&rate_limit_headers,
&body,
)
.await;
}
Ok(CollectedResponse::Failed(failure)) => {
let response_id = latest_response_id(&body);
let latency_ms = elapsed_millis_i64(context.started_at);
let failure_source = stream_failure_source(&failure);
let failure_detail = synthetic_stream_disconnected_detail(&failure);
let websocket_pool_kind = context
.websocket_pool_decision
.map(|decision| decision.kind());
let websocket_pool_reason = context
.websocket_pool_decision
.and_then(|decision| decision.reason());
tracing::warn!(
account_id = %context.account_id,
request_id = %context.request_id,
response_id = response_id.as_deref().unwrap_or(""),
transport = %backend_transport_name(context.transport),
websocket_pool_kind = ?websocket_pool_kind,
websocket_pool_reason = ?websocket_pool_reason,
first_token_ms = ?first_token_ms,
latency_ms,
event = %failure.event,
code = ?failure.upstream_code.as_deref(),
failure_source = %failure_source,
failure_detail = ?failure_detail.as_deref(),
"live upstream stream ended with response.failed"
);
let mut metadata = stream_failure_metadata(&failure, usage);
insert_first_token_ms(&mut metadata, first_token_ms);
record_live_response_stream_event(
&context,
status_code_for_stream_failure(&failure),
UsageRecordLevel::Error,
"v1 responses stream failed",
metadata,
&rate_limit_headers,
&body,
)
.await;
}
Ok(CollectedResponse::MissingCompleted | CollectedResponse::Empty) => {
tracing::warn!(
account_id = %context.account_id,
"live upstream stream ended without response.completed"
);
let mut metadata = serde_json::json!({
"stream": true,
"failed": true,
"upstreamCode": "missing_completed",
"usage": usage,
});
insert_first_token_ms(&mut metadata, first_token_ms);
record_live_response_stream_event(
&context,
502,
UsageRecordLevel::Error,
"v1 responses stream ended without response.completed",
metadata,
&rate_limit_headers,
&body,
)
.await;
}
Err(error) => {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to parse completed live stream");
let mut metadata = serde_json::json!({
"stream": true,
"sseParseError": error.to_string(),
"usage": usage,
});
insert_first_token_ms(&mut metadata, first_token_ms);
record_live_response_stream_event(
&context,
502,
UsageRecordLevel::Warn,
"v1 responses stream SSE response invalid",
metadata,
&rate_limit_headers,
&body,
)
.await;
}
first_event_ms: i64,
decoder: &CanonicalStreamDecoder,
) -> StreamSummary {
StreamSummary {
terminal,
body,
terminal_chunks,
first_token_ms,
first_event_ms,
usage: decoder.usage(),
last_response_id: decoder.last_response_id().map(ToString::to_string),
}
context.account_lease.complete().await;
}
-3
View File
@@ -2,6 +2,3 @@
pub(crate) mod lifecycle;
pub(crate) mod live;
pub(crate) mod prefetch;
pub(crate) mod sse_failure;
pub(crate) mod trace;
-36
View File
@@ -1,36 +0,0 @@
use bytes::Bytes;
use futures::StreamExt;
use crate::upstream::openai::{
protocol::sse::{SseError, sse_frame_end},
transport::CodexBackendSseStream,
};
use crate::dispatch::errors::ResponseDispatchError;
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
pub(in crate::dispatch) async fn prefetch_first_sse_chunk(
mut body: CodexBackendSseStream,
) -> Result<(Bytes, CodexBackendSseStream), ResponseDispatchError> {
let mut prefetched = Vec::new();
while sse_frame_end(&prefetched).is_none() {
let Some(next) = body.next().await else {
if prefetched.is_empty() {
return Err(ResponseDispatchError::EmptyUpstreamResponse);
}
return Err(ResponseDispatchError::MissingCompleted);
};
let chunk = next.map_err(ResponseDispatchError::Upstream)?;
prefetched.extend_from_slice(&chunk);
if prefetched.len() > MAX_STREAM_PREFETCH_BYTES {
return Err(ResponseDispatchError::InvalidSse(
SseError::BufferExceeded {
max_bytes: MAX_STREAM_PREFETCH_BYTES,
},
));
}
}
Ok((Bytes::from(prefetched), body))
}
-221
View File
@@ -1,221 +0,0 @@
use serde_json::Value;
use crate::{
dispatch::errors::{is_history_recovery_code, is_model_unsupported_signal},
fleet::account::AccountStatus,
upstream::openai::{
protocol::{
events::TokenUsage,
responses::{CollectedResponse, ResponsesSseFailure, response_from_codex_sse},
sse::SseError,
},
transport::is_banned_auth_signal,
},
};
pub(in crate::dispatch) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
pub(in crate::dispatch) const STREAM_DISCONNECTED_MESSAGE: &str =
"Upstream stream closed before response.completed";
pub(in crate::dispatch) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
match failure.upstream_code.as_deref() {
Some(code) => serde_json::json!({
"error": {
"code": code,
"message": failure.message.as_str(),
}
})
.to_string(),
None => failure.message.clone(),
}
}
fn sse_failure_matches<F>(failure: &ResponsesSseFailure, signal: F) -> bool
where
F: Fn(&str) -> bool,
{
failure.upstream_code.as_deref().is_some_and(&signal) || signal(&failure.message)
}
fn sse_failure_matches_parts<C, M>(
failure: &ResponsesSseFailure,
code_signal: C,
message_signal: M,
) -> bool
where
C: Fn(&str) -> bool,
M: Fn(&str) -> bool,
{
failure.upstream_code.as_deref().is_some_and(code_signal) || message_signal(&failure.message)
}
pub(crate) fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches_parts(
failure,
|code| matches!(code, "quota_exceeded" | "insufficient_quota"),
|message| message.to_ascii_lowercase().contains("quota"),
)
}
fn is_auth_sse_failure_code(code: &str) -> bool {
let code = code.to_ascii_lowercase();
matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
)
}
fn is_auth_sse_failure_message(message: &str) -> bool {
let message = message.to_ascii_lowercase();
message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired")
}
pub(crate) fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches_parts(
failure,
is_auth_sse_failure_code,
is_auth_sse_failure_message,
)
}
pub(in crate::dispatch) fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches(failure, is_model_unsupported_signal)
}
pub(in crate::dispatch) fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool {
failure
.upstream_code
.as_deref()
.is_some_and(is_history_recovery_code)
}
pub(crate) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus {
if sse_failure_matches(failure, is_banned_auth_signal) {
AccountStatus::Banned
} else {
AccountStatus::Expired
}
}
pub(in crate::dispatch) fn first_sse_failure(
prefetched: &[u8],
) -> Result<Option<ResponsesSseFailure>, SseError> {
let body = String::from_utf8_lossy(prefetched);
match response_from_codex_sse(&body, None)? {
CollectedResponse::Failed(failure) => Ok(Some(failure)),
CollectedResponse::Completed(_)
| CollectedResponse::Incomplete(_)
| CollectedResponse::MissingCompleted
| CollectedResponse::Empty => Ok(None),
}
}
pub(in crate::dispatch) fn stream_failure_metadata(
failure: &ResponsesSseFailure,
usage: Option<TokenUsage>,
) -> Value {
let mut metadata = serde_json::json!({
"stream": true,
"failed": true,
"failureEvent": failure.event,
"failureMessage": failure.message,
"upstreamCode": failure.upstream_code,
"usage": usage,
});
enrich_stream_failure_source_metadata(&mut metadata, failure);
metadata
}
fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object.insert(
"failureSource".to_string(),
Value::String(stream_failure_source(failure).to_string()),
);
if let Some(detail) = synthetic_stream_disconnected_detail(failure) {
object.insert("synthetic".to_string(), Value::Bool(true));
if !detail.is_empty() {
object.insert("failureDetail".to_string(), Value::String(detail));
}
}
}
pub(in crate::dispatch) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
if synthetic_stream_disconnected_detail(failure).is_some() {
"proxy"
} else {
"upstream"
}
}
pub(in crate::dispatch) fn synthetic_stream_disconnected_detail(
failure: &ResponsesSseFailure,
) -> Option<String> {
if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) {
return None;
}
let detail = failure
.message
.strip_prefix(STREAM_DISCONNECTED_MESSAGE)?
.strip_prefix(": ")
.unwrap_or_default()
.trim()
.to_string();
Some(detail)
}
pub(in crate::dispatch) fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 {
let code = failure
.upstream_code
.as_deref()
.unwrap_or("error")
.to_ascii_lowercase();
if code.contains("model") && (code.contains("not_supported") || code.contains("not_available"))
{
return 400;
}
if code.contains("invalid_request") || code.contains("not_found") {
return 400;
}
if code.contains("context_window")
|| code.contains("invalid_prompt")
|| code.contains("cyber_policy")
|| code.contains("bad_request")
{
return 400;
}
if code.contains("rate_limit") || code.contains("usage_limit") {
return 429;
}
if code.contains("unauthorized")
|| code.contains("invalid_api_key")
|| code == "token_invalid"
|| code == "token_expired"
|| code == "account_deactivated"
{
return 401;
}
if code.contains("forbidden") || code.contains("banned") {
return 403;
}
if code.contains("payment") || code.contains("quota") {
return 429;
}
if code.contains("server_overloaded") {
return 503;
}
502
}
pub(in crate::dispatch) fn stream_failure_http_status(failure: &ResponsesSseFailure) -> u16 {
u16::try_from(status_code_for_stream_failure(failure)).unwrap_or(502)
}
+92
View File
@@ -0,0 +1,92 @@
//! 调度层共享的账号级 Codex 上游调用。
use std::time::Instant;
use crate::{
dispatch::affinity::AccountScopedRequest,
fleet::account::Account,
upstream::openai::transport::{
CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse, CodexClientError,
CodexRequestContext, PreparedResponseTransport,
},
};
#[derive(Clone, Copy)]
pub(in crate::dispatch) struct AccountUpstreamContext<'a> {
pub codex: &'a CodexBackendClient,
pub request_id: &'a str,
pub account: &'a Account,
pub cookie_header: Option<&'a str>,
}
pub(in crate::dispatch) async fn create_response_with_account(
context: AccountUpstreamContext<'_>,
request: &AccountScopedRequest,
prepared: PreparedResponseTransport,
started_at: Instant,
) -> Result<CodexBackendResponse, CodexClientError> {
context
.codex
.create_response_with_prepared(
request.request(),
codex_request_context(context, request),
prepared,
started_at,
)
.await
}
pub(in crate::dispatch) async fn prepare_response_transport_with_account(
context: AccountUpstreamContext<'_>,
request: &AccountScopedRequest,
) -> Result<PreparedResponseTransport, CodexClientError> {
context
.codex
.prepare_response_transport_with_pool_account(
request.request(),
codex_request_context(context, request),
Some(&context.account.id),
)
.await
}
pub(in crate::dispatch) async fn create_response_stream_with_account(
context: AccountUpstreamContext<'_>,
request: &AccountScopedRequest,
prepared: PreparedResponseTransport,
) -> Result<CodexBackendStreamingResponse, CodexClientError> {
context
.codex
.create_response_stream_with_prepared(
request.request(),
codex_request_context(context, request),
prepared,
)
.await
}
fn codex_request_context<'a>(
context: AccountUpstreamContext<'a>,
request: &'a AccountScopedRequest,
) -> CodexRequestContext<'a> {
let request_body = request.request();
let identity = request.identity();
CodexRequestContext {
access_token: &context.account.access_token,
account_id: context.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: request_body.turn_state.as_deref(),
turn_metadata: request_body.turn_metadata.as_deref(),
beta_features: request_body.beta_features.as_deref(),
include_timing_metrics: request_body.include_timing_metrics.as_deref(),
version: request_body.version.as_deref(),
codex_window_id: request_body.codex_window_id.as_deref(),
parent_thread_id: request_body.parent_thread_id.as_deref(),
cookie_header: context.cookie_header,
installation_id: Some(&identity.installation_id),
session_id: request_body.client_session_id.as_deref(),
thread_id: request_body.client_thread_id.as_deref(),
client_request_id: request_body.client_request_id.as_deref(),
turn_id: request_body.client_turn_id.as_deref(),
}
}
+424
View File
@@ -0,0 +1,424 @@
//! Responses SSE 原始帧到规范事件的唯一解码入口。
use bytes::Bytes;
use serde_json::Value;
use crate::upstream::openai::protocol::{
events::{TokenUsage, extract_usage},
responses::{
CollectedResponse, ResponsesSseFailure, StreamCommitPolicy,
reconvert_responses_sse_event_tuple_values,
},
sse::{
MAX_SSE_EVENT_BUFFER_BYTES, SseError, encode_sse_event, parse_sse_events, sse_frame_end,
sse_frame_is_done,
},
};
/// 将 complete transport 的原始 SSE 归一化为类型化终态事实。
pub(in crate::dispatch) fn normalize_complete_response(
body: &str,
tuple_schema: Option<&Value>,
) -> Result<CollectedResponse, SseError> {
let events = parse_sse_events(body)?;
let mut output_text = String::new();
let mut output_items = Vec::new();
let mut terminal_response = None;
let mut terminal_incomplete = false;
let mut failed_response = None;
for event in events {
let Ok(mut value) = serde_json::from_str::<Value>(&event.data) else {
continue;
};
if let Some(tuple_schema) = tuple_schema {
value = reconvert_responses_sse_event_tuple_values(
event.event.as_deref(),
value,
tuple_schema,
);
}
match event.event.as_deref() {
Some("response.output_text.delta") => {
if let Some(delta) = value.get("delta").and_then(Value::as_str) {
output_text.push_str(delta);
}
}
Some("response.output_item.done") => {
if let Some(item) = value.get("item") {
output_items.push(item.clone());
}
}
Some("response.completed") => {
terminal_response = value.get("response").cloned();
terminal_incomplete = false;
}
Some("response.incomplete") => {
terminal_response = value.get("response").cloned();
terminal_incomplete = true;
}
Some(event_name @ ("error" | "response.failed")) if failed_response.is_none() => {
failed_response = Some(ResponsesSseFailure::from_event(event_name, &value));
}
_ => {}
}
}
if let Some(failure) = failed_response {
return Ok(CollectedResponse::Failed(failure));
}
let Some(response) = terminal_response else {
return Ok(CollectedResponse::MissingCompleted);
};
if terminal_incomplete {
return Ok(CollectedResponse::Incomplete(response));
}
if complete_response_is_empty(&response, &output_text, &output_items) {
return Ok(CollectedResponse::Empty);
}
Ok(CollectedResponse::Completed(response))
}
fn complete_response_is_empty(response: &Value, output_text: &str, output_items: &[Value]) -> bool {
if !output_text.trim().is_empty() || !output_items.is_empty() {
return false;
}
if response.get("status").and_then(Value::as_str) == Some("incomplete") {
return false;
}
response
.pointer("/usage/output_tokens")
.and_then(Value::as_u64)
.unwrap_or_default()
== 0
}
/// 一次解码后可同时供 HTTP SSE、Responses WebSocket 与生命周期控制器消费的事件。
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct CanonicalResponseEvent {
event_type: Option<String>,
data: Value,
}
impl CanonicalResponseEvent {
pub(crate) fn event_type(&self) -> Option<&str> {
self.event_type.as_deref()
}
pub(crate) fn data(&self) -> &Value {
&self.data
}
pub(in crate::dispatch) fn response_id(&self) -> Option<&str> {
self.data
.pointer("/response/id")
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
}
pub(in crate::dispatch) fn usage(&self) -> Option<TokenUsage> {
self.data
.get("response")
.and_then(extract_usage)
.or_else(|| extract_usage(&self.data))
}
pub(in crate::dispatch) fn has_first_output(&self) -> bool {
match self.event_type() {
Some(
"response.output_text.delta"
| "response.reasoning_summary_text.delta"
| "response.reasoning_text.delta"
| "response.function_call_arguments.delta"
| "response.custom_tool_call_input.delta",
) => self
.data
.get("delta")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty()),
Some("response.output_text.done") => self
.data
.get("text")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty()),
Some("response.function_call_arguments.done") => self
.data
.get("arguments")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty()),
Some("response.output_item.done") => {
self.data.get("item").is_some_and(Value::is_object)
}
_ => false,
}
}
pub(in crate::dispatch) fn terminal(&self) -> Option<CanonicalStreamTerminal> {
match self.event_type() {
Some("response.completed") => self
.data
.get("response")
.filter(|response| response.is_object())
.cloned()
.map(CanonicalStreamTerminal::Completed),
Some("response.incomplete") => self
.data
.get("response")
.filter(|response| response.is_object())
.cloned()
.map(CanonicalStreamTerminal::Incomplete),
Some(event @ ("response.failed" | "error")) => Some(CanonicalStreamTerminal::Failed(
ResponsesSseFailure::from_event(event, &self.data),
)),
_ => None,
}
}
pub(in crate::dispatch) fn proxy_failure(data: Value) -> Self {
Self {
event_type: Some("response.failed".to_string()),
data,
}
}
}
/// 上游明确发出的业务终态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub(in crate::dispatch) enum CanonicalStreamTerminal {
Completed(Value),
Incomplete(Value),
Failed(ResponsesSseFailure),
}
/// 已解码事件与对应下游字节的单个 fanout 单元。
#[derive(Debug, Clone)]
pub(crate) struct CanonicalResponseChunk {
bytes: Bytes,
events: Vec<CanonicalResponseEvent>,
}
impl CanonicalResponseChunk {
pub(in crate::dispatch) fn new(bytes: Bytes, events: Vec<CanonicalResponseEvent>) -> Self {
Self { bytes, events }
}
pub(in crate::dispatch) fn bytes(&self) -> &Bytes {
&self.bytes
}
pub(in crate::dispatch) fn into_parts(self) -> (Bytes, Vec<CanonicalResponseEvent>) {
(self.bytes, self.events)
}
pub(crate) fn events(&self) -> &[CanonicalResponseEvent] {
&self.events
}
}
/// 单次输入产生的所有下游 fanout 单元。
#[derive(Debug)]
pub(in crate::dispatch) struct CanonicalStreamBatch {
pub chunks: Vec<CanonicalResponseChunk>,
}
impl CanonicalStreamBatch {
pub(in crate::dispatch) fn append(&mut self, mut other: Self) {
self.chunks.append(&mut other.chunks);
}
/// 只把真实输出前出现的失败视为可在 attempt 阶段处理的失败。
/// 同一上游 chunk 可能同时包含输出与终态,事件顺序不能被网络分块抹掉。
pub(in crate::dispatch) fn first_failure(&self) -> Option<ResponsesSseFailure> {
for event in self.chunks.iter().flat_map(|chunk| chunk.events()) {
if event.has_first_output() {
return None;
}
if let Some(terminal) = event.terminal() {
return match terminal {
CanonicalStreamTerminal::Failed(failure) => Some(failure),
CanonicalStreamTerminal::Completed(_)
| CanonicalStreamTerminal::Incomplete(_) => None,
};
}
}
None
}
}
/// 原始 SSE 在整个 attempt/live 生命周期中唯一持有的增量 decoder。
pub(in crate::dispatch) struct CanonicalStreamDecoder {
tuple_schema: Option<Value>,
pending: Vec<u8>,
forwardable_frame_seen: bool,
first_output_seen: bool,
terminal: Option<CanonicalStreamTerminal>,
transport_done_seen: bool,
completed_usage: Option<TokenUsage>,
fallback_usage: Option<TokenUsage>,
last_response_id: Option<String>,
}
impl CanonicalStreamDecoder {
pub(in crate::dispatch) fn new(tuple_schema: Option<Value>) -> Self {
Self {
tuple_schema,
pending: Vec::new(),
forwardable_frame_seen: false,
first_output_seen: false,
terminal: None,
transport_done_seen: false,
completed_usage: None,
fallback_usage: None,
last_response_id: None,
}
}
pub(in crate::dispatch) fn push(
&mut self,
chunk: Bytes,
) -> Result<CanonicalStreamBatch, SseError> {
self.pending.extend_from_slice(&chunk);
let mut transformed_chunks = Vec::new();
let mut consumed = 0usize;
while let Some(frame_len) = sse_frame_end(&self.pending[consumed..]) {
let end = consumed.saturating_add(frame_len);
let frame = self.pending[consumed..end].to_vec();
if let Some((transformed, events)) = self.normalize_frame(&frame)? {
transformed_chunks.push(CanonicalResponseChunk::new(transformed, events));
}
consumed = end;
if self.terminal.is_some() || self.transport_done_seen {
consumed = self.pending.len();
break;
}
}
if consumed != 0 {
self.pending.drain(..consumed);
}
if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES {
return Err(SseError::BufferExceeded {
max_bytes: MAX_SSE_EVENT_BUFFER_BYTES,
});
}
Ok(CanonicalStreamBatch {
chunks: transformed_chunks,
})
}
pub(in crate::dispatch) fn finish(&mut self) -> Result<CanonicalStreamBatch, SseError> {
if self.pending.is_empty() {
return Ok(CanonicalStreamBatch { chunks: Vec::new() });
}
let frame = std::mem::take(&mut self.pending);
let Some((transformed, events)) = self.normalize_frame(&frame)? else {
return Ok(CanonicalStreamBatch { chunks: Vec::new() });
};
Ok(CanonicalStreamBatch {
chunks: vec![CanonicalResponseChunk::new(transformed, events)],
})
}
pub(in crate::dispatch) fn take_terminal(&mut self) -> Option<CanonicalStreamTerminal> {
self.terminal.take()
}
pub(in crate::dispatch) fn take_transport_done(&mut self) -> bool {
std::mem::take(&mut self.transport_done_seen)
}
pub(in crate::dispatch) fn transport_done_seen(&self) -> bool {
self.transport_done_seen
}
pub(in crate::dispatch) fn first_output_seen(&self) -> bool {
self.first_output_seen
}
pub(in crate::dispatch) fn commit_boundary_reached(&self, policy: StreamCommitPolicy) -> bool {
match policy {
StreamCommitPolicy::FirstForwardableEvent => self.forwardable_frame_seen,
StreamCommitPolicy::UntilOutputOrTerminal => {
self.first_output_seen || self.terminal.is_some()
}
}
}
pub(in crate::dispatch) fn usage(&self) -> Option<TokenUsage> {
self.completed_usage.or(self.fallback_usage)
}
pub(in crate::dispatch) fn last_response_id(&self) -> Option<&str> {
self.last_response_id.as_deref()
}
fn normalize_frame(
&mut self,
frame: &[u8],
) -> Result<Option<(Bytes, Vec<CanonicalResponseEvent>)>, SseError> {
let frame_text =
std::str::from_utf8(frame).map_err(|error| SseError::ParseError(error.to_string()))?;
if sse_frame_is_done(frame_text) {
self.transport_done_seen = true;
return Ok(None);
}
let events = parse_sse_events(frame_text)?;
let mut canonical_events = Vec::new();
let mut transformed = String::new();
for event in events {
let Ok(mut data) = serde_json::from_str::<Value>(&event.data) else {
continue;
};
let event_type = event.event.or_else(|| {
data.get("type")
.and_then(Value::as_str)
.map(ToString::to_string)
});
if let Some(tuple_schema) = self.tuple_schema.as_ref() {
data = reconvert_responses_sse_event_tuple_values(
event_type.as_deref(),
data,
tuple_schema,
);
transformed.push_str(&encode_sse_event(
event_type.as_deref().unwrap_or_default(),
&data.to_string(),
));
}
let canonical = CanonicalResponseEvent { event_type, data };
self.observe(&canonical);
canonical_events.push(canonical);
}
self.forwardable_frame_seen = true;
Ok(Some((
if self.tuple_schema.is_some() && !transformed.is_empty() {
Bytes::from(transformed)
} else {
Bytes::copy_from_slice(frame)
},
canonical_events,
)))
}
fn observe(&mut self, event: &CanonicalResponseEvent) {
self.first_output_seen |= event.has_first_output();
if let Some(response_id) = event.response_id() {
self.last_response_id = Some(response_id.to_string());
}
if let Some(usage) = event.usage() {
if event.event_type() == Some("response.completed") {
self.completed_usage = Some(usage);
} else {
self.fallback_usage = Some(usage);
}
}
if self.terminal.is_none() {
self.terminal = event.terminal();
}
}
}
+5
View File
@@ -0,0 +1,5 @@
//! Attempt 阶段的账号作用域上游调用与流规范化。
pub(in crate::dispatch) mod account;
pub(crate) mod canonical;
pub(in crate::dispatch) mod prefetch;
@@ -0,0 +1,80 @@
//! 上游 SSE 提交边界前的规范事件预取。
use bytes::Bytes;
use std::time::Instant;
use futures::StreamExt;
use crate::{
dispatch::transport::canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
upstream::openai::{
protocol::{responses::StreamCommitPolicy, sse::SseError},
transport::{CodexBackendSseStream, CodexClientError},
},
};
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
#[derive(Debug, thiserror::Error)]
pub(in crate::dispatch) enum StreamPrefetchError {
#[error("upstream stream failed before the commit boundary: {0}")]
Upstream(#[from] CodexClientError),
#[error("upstream stream ended without any data")]
Empty,
#[error("upstream stream ended before output or a terminal event")]
NoCommitBoundary,
#[error("invalid upstream SSE response: {0}")]
InvalidSse(#[from] SseError),
}
pub(in crate::dispatch) struct PrefetchedStream {
pub bytes: Bytes,
pub initial_batch: CanonicalStreamBatch,
pub body: CodexBackendSseStream,
pub first_event_ms: i64,
}
pub(in crate::dispatch) async fn prefetch_until_commit(
mut body: CodexBackendSseStream,
decoder: &mut CanonicalStreamDecoder,
policy: StreamCommitPolicy,
started_at: Instant,
) -> Result<PrefetchedStream, StreamPrefetchError> {
let mut prefetched = Vec::new();
let mut initial_batch = CanonicalStreamBatch { chunks: Vec::new() };
let mut first_event_ms = None;
while !decoder.commit_boundary_reached(policy) {
let Some(next) = body.next().await else {
if prefetched.is_empty() {
return Err(StreamPrefetchError::Empty);
}
initial_batch.append(decoder.finish()?);
if decoder.commit_boundary_reached(policy) {
break;
}
return Err(StreamPrefetchError::NoCommitBoundary);
};
let chunk = next?;
first_event_ms
.get_or_insert_with(|| crate::infra::time::elapsed_millis_i64(started_at).max(1));
prefetched.extend_from_slice(&chunk);
let batch = decoder.push(chunk)?;
initial_batch.append(batch);
if decoder.transport_done_seen() && !decoder.commit_boundary_reached(policy) {
return Err(StreamPrefetchError::NoCommitBoundary);
}
if !decoder.commit_boundary_reached(policy) && prefetched.len() > MAX_STREAM_PREFETCH_BYTES
{
return Err(StreamPrefetchError::InvalidSse(SseError::BufferExceeded {
max_bytes: MAX_STREAM_PREFETCH_BYTES,
}));
}
}
Ok(PrefetchedStream {
bytes: Bytes::from(prefetched),
initial_batch,
body,
first_event_ms: first_event_ms.ok_or(StreamPrefetchError::Empty)?,
})
}
-221
View File
@@ -1,221 +0,0 @@
//! 调度层共享的账号级 Codex 上游调用。
use std::time::Instant;
use serde_json::Value;
use crate::{
dispatch::{affinity::AccountIdentityService, recovery::cloudflare::CloudflareRecovery},
fleet::{
account::Account,
pool::{AccountLease, AccountPoolService},
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
},
},
upstream::openai::{
protocol::responses::CodexResponsesRequest,
transport::{
CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse,
CodexClientError, CodexRequestContext,
},
},
};
pub(crate) const QUOTA_VERIFY_LIMIT_REACHED_MESSAGE: &str =
"Upstream usage quota still reports limit_reached";
pub(crate) enum QuotaVerificationDecision {
Ready(Box<AccountLease>),
RetryWithAnotherAccount,
}
pub(crate) struct QuotaVerificationContext<'a> {
pub account_pool: &'a AccountPoolService,
pub codex: &'a CodexBackendClient,
pub cloudflare: &'a CloudflareRecovery,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
}
#[derive(Clone, Copy)]
pub(crate) struct AccountUpstreamContext<'a> {
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub cloudflare: &'a CloudflareRecovery,
pub request_id: &'a str,
pub account: &'a Account,
}
pub(crate) async fn verify_acquired_quota_if_required(
context: QuotaVerificationContext<'_>,
acquired: AccountLease,
) -> QuotaVerificationDecision {
if !acquired.account.quota_verify_required {
return QuotaVerificationDecision::Ready(Box::new(acquired));
}
let account_id = acquired.account.id.clone();
let identity = context
.account_identity
.scope_auxiliary(&account_id, context.request_id);
let cookie_header = context
.cloudflare
.cookie_header_for_request(&account_id, "/codex/usage")
.await;
let usage = context
.codex
.fetch_usage(CodexRequestContext {
access_token: &acquired.account.access_token,
account_id: acquired.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: None,
parent_thread_id: None,
cookie_header: cookie_header.as_deref(),
installation_id: Some(&identity.installation_id),
session_id: None,
thread_id: None,
prompt_cache_key: None,
client_request_id: Some(&identity.client_request_id),
turn_id: None,
})
.await;
let raw = match usage {
Ok(raw) => raw,
Err(error) => {
tracing::warn!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "upstream_error",
retry_with_another_account = false,
error = %error,
"failed to verify stale quota state before upstream request"
);
return QuotaVerificationDecision::Ready(Box::new(acquired));
}
};
let quota = quota_from_usage(&raw);
context
.account_pool
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
acquired.release_without_usage().await;
tracing::info!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "limit_reached",
retry_with_another_account = true,
"quota verification reported exhausted account before upstream request"
);
return QuotaVerificationDecision::RetryWithAnotherAccount;
}
QuotaVerificationDecision::Ready(Box::new(acquired_with_verified_quota(acquired, &quota)))
}
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
let limit_reached = quota_snapshot_limit_reached(quota);
acquired.account.quota_verify_required = false;
acquired.account.quota_limit_reached = limit_reached;
acquired.account.quota_cooldown_until = limit_reached
.then_some(quota_snapshot_reset_at(quota))
.flatten();
if let Some(reset_at) = quota_snapshot_reset_at(quota) {
acquired.account.window_reset_at = Some(reset_at);
if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) {
acquired.account.limit_window_seconds = Some(limit_window_seconds);
}
}
acquired
}
pub(crate) async fn create_response_with_account(
context: AccountUpstreamContext<'_>,
request: &CodexResponsesRequest,
started_at: Instant,
) -> Result<CodexBackendResponse, CodexClientError> {
let cookie_header = context
.cloudflare
.cookie_header_for_request(&context.account.id, "/codex/responses")
.await;
let identity = context
.account_identity
.scope(request, &context.account.id, context.request_id);
context
.codex
.create_response_with_pool_account_started_at(
request,
CodexRequestContext {
access_token: &context.account.access_token,
account_id: context.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: request.turn_state.as_deref(),
turn_metadata: request.turn_metadata.as_deref(),
beta_features: request.beta_features.as_deref(),
include_timing_metrics: request.include_timing_metrics.as_deref(),
version: request.version.as_deref(),
codex_window_id: identity.window_id.as_deref(),
parent_thread_id: identity.parent_thread_id.as_deref(),
cookie_header: cookie_header.as_deref(),
installation_id: Some(&identity.installation_id),
session_id: identity.session_id.as_deref(),
thread_id: identity.thread_id.as_deref(),
prompt_cache_key: identity.prompt_cache_key.as_deref(),
client_request_id: Some(&identity.client_request_id),
turn_id: identity.turn_id.as_deref(),
},
Some(&context.account.id),
started_at,
)
.await
}
pub(crate) async fn create_response_stream_with_account(
context: AccountUpstreamContext<'_>,
request: &CodexResponsesRequest,
) -> Result<CodexBackendStreamingResponse, CodexClientError> {
let cookie_header = context
.cloudflare
.cookie_header_for_request(&context.account.id, "/codex/responses")
.await;
let identity = context
.account_identity
.scope(request, &context.account.id, context.request_id);
context
.codex
.create_response_stream_with_pool_account(
request,
CodexRequestContext {
access_token: &context.account.access_token,
account_id: context.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: request.turn_state.as_deref(),
turn_metadata: request.turn_metadata.as_deref(),
beta_features: request.beta_features.as_deref(),
include_timing_metrics: request.include_timing_metrics.as_deref(),
version: request.version.as_deref(),
codex_window_id: identity.window_id.as_deref(),
parent_thread_id: identity.parent_thread_id.as_deref(),
cookie_header: cookie_header.as_deref(),
installation_id: Some(&identity.installation_id),
session_id: identity.session_id.as_deref(),
thread_id: identity.thread_id.as_deref(),
prompt_cache_key: identity.prompt_cache_key.as_deref(),
client_request_id: Some(&identity.client_request_id),
turn_id: identity.turn_id.as_deref(),
},
Some(&context.account.id),
)
.await
}
+317
View File
@@ -0,0 +1,317 @@
//! 上游失败到账号状态 effect 的唯一领域分类。
use chrono::{Duration, Utc};
use crate::{
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{
failure::{UpstreamFailureFacts, upstream_failure_facts},
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexClientError},
},
};
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountFailureKind {
ModelUnsupported,
Expired,
Disabled,
Banned,
QuotaExhausted,
RateLimited,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AccountStateEffect {
SetStatus(AccountStatus),
MarkQuotaLimitedUntil(chrono::DateTime<Utc>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClassifiedAccountFailure {
pub kind: AccountFailureKind,
pub effect: Option<AccountStateEffect>,
}
impl ClassifiedAccountFailure {
pub fn account_status(&self) -> Option<AccountStatus> {
match self.effect.as_ref()? {
AccountStateEffect::SetStatus(status) => Some(*status),
AccountStateEffect::MarkQuotaLimitedUntil(_) => None,
}
}
}
pub fn classify_client_failure(error: &CodexClientError) -> Option<ClassifiedAccountFailure> {
classify_upstream_failure(&upstream_failure_facts(error))
}
pub fn classify_upstream_failure(facts: &UpstreamFailureFacts) -> Option<ClassifiedAccountFailure> {
classify_failure(FailureFields {
status_code: facts.status_code,
code: facts.code.as_deref(),
error_type: facts.error_type.as_deref(),
message: &facts.message,
body: &facts.body,
retry_after_seconds: facts.retry_after_seconds,
})
}
pub fn classify_response_failure(
failure: &ResponsesSseFailure,
) -> Option<ClassifiedAccountFailure> {
classify_failure(FailureFields {
status_code: failure.explicit_status_code,
code: failure.upstream_code.as_deref(),
error_type: failure.upstream_type.as_deref(),
message: &failure.message,
body: &failure.message,
retry_after_seconds: failure.retry_after_seconds,
})
}
pub async fn apply_account_state_effect_immediately(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
effect: &AccountStateEffect,
) {
match effect {
AccountStateEffect::SetStatus(status) => {
account_pool
.set_status_immediately(account_id, *status)
.await;
}
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until_immediately(account_id, *until)
.await;
}
}
codex.evict_websocket_account(account_id).await;
}
pub async fn apply_account_state_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
effect: &AccountStateEffect,
) -> bool {
let updated = match effect {
AccountStateEffect::SetStatus(status) => account_pool.set_status(account_id, *status).await,
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until(account_id, *until)
.await
}
};
codex.evict_websocket_account(account_id).await;
updated
}
pub async fn apply_client_failure_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
error: &CodexClientError,
) -> Option<ClassifiedAccountFailure> {
let classified = classify_client_failure(error)?;
if let Some(effect) = &classified.effect {
apply_account_state_effect(account_pool, codex, account_id, effect).await;
}
Some(classified)
}
struct FailureFields<'a> {
status_code: Option<u16>,
code: Option<&'a str>,
error_type: Option<&'a str>,
message: &'a str,
body: &'a str,
retry_after_seconds: Option<u64>,
}
fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailure> {
let code = fields.code.unwrap_or_default();
let error_type = fields.error_type.unwrap_or_default();
if [code, fields.message, fields.body]
.into_iter()
.any(is_model_unsupported)
{
return Some(classified(AccountFailureKind::ModelUnsupported, None));
}
if let Some(status) = explicit_account_status(code, error_type, fields.message)
.or_else(|| explicit_account_status(code, error_type, fields.body))
{
return Some(status_failure(status));
}
if fields.status_code == Some(403) && !is_html(fields.body) {
return Some(status_failure(AccountStatus::Banned));
}
if fields.status_code == Some(401) {
return Some(status_failure(AccountStatus::Expired));
}
let quota_kind = [code, error_type]
.into_iter()
.find_map(classify_quota_signal)
.or_else(|| classify_quota_message(fields.message))
.or_else(|| classify_quota_message(fields.body))
.or(match fields.status_code {
Some(429) => Some(AccountFailureKind::RateLimited),
Some(402) => Some(AccountFailureKind::QuotaExhausted),
_ => None,
})?;
Some(match quota_kind {
AccountFailureKind::RateLimited => {
let seconds = fields
.retry_after_seconds
.unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS)
.min(i64::MAX as u64) as i64;
classified(
AccountFailureKind::RateLimited,
Some(AccountStateEffect::MarkQuotaLimitedUntil(
Utc::now() + Duration::seconds(seconds),
)),
)
}
AccountFailureKind::QuotaExhausted => classified(
AccountFailureKind::QuotaExhausted,
Some(AccountStateEffect::SetStatus(AccountStatus::QuotaExhausted)),
),
AccountFailureKind::ModelUnsupported
| AccountFailureKind::Expired
| AccountFailureKind::Disabled
| AccountFailureKind::Banned => unreachable!("quota classifier returned account status"),
})
}
fn status_failure(status: AccountStatus) -> ClassifiedAccountFailure {
let kind = match status {
AccountStatus::Expired => AccountFailureKind::Expired,
AccountStatus::Disabled => AccountFailureKind::Disabled,
AccountStatus::Banned => AccountFailureKind::Banned,
AccountStatus::Active | AccountStatus::QuotaExhausted => {
unreachable!("explicit account failure must make the account unavailable")
}
};
classified(kind, Some(AccountStateEffect::SetStatus(status)))
}
fn classified(
kind: AccountFailureKind,
effect: Option<AccountStateEffect>,
) -> ClassifiedAccountFailure {
ClassifiedAccountFailure { kind, effect }
}
fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Option<AccountStatus> {
let code = code.trim().to_ascii_lowercase();
let error_type = error_type.trim().to_ascii_lowercase();
let message = message.to_ascii_lowercase();
if matches!(
code.as_str(),
"identity_verification_required" | "verification_required"
) || message.contains("identity verification is required")
{
return Some(AccountStatus::Disabled);
}
if matches!(
code.as_str(),
"account_banned"
| "account_deactivated"
| "account_disabled"
| "account_suspended"
| "deactivated_workspace"
| "organization_disabled"
| "workspace_deactivated"
) || message.contains("account is banned")
|| message.contains("account has been banned")
|| message.contains("account deactivated")
|| message.contains("account has been deactivated")
|| message.contains("account disabled")
|| message.contains("account has been disabled")
|| message.contains("account suspended")
|| message.contains("organization has been disabled")
|| message.contains("workspace has been deactivated")
|| message.contains("deactivated_workspace")
|| code == "forbidden"
|| error_type == "permission_error"
{
return Some(AccountStatus::Banned);
}
is_auth_failure(&code, &error_type, &message).then_some(AccountStatus::Expired)
}
fn is_auth_failure(code: &str, error_type: &str, message: &str) -> bool {
matches!(
code,
"token_invalid"
| "token_invalidated"
| "token_expired"
| "token_revoked"
| "refresh_token_invalidated"
| "unauthorized"
| "invalid_api_key"
| "authentication_error"
) || error_type == "authentication_error"
|| message.contains("token revoked")
|| message.contains("token invalidated")
|| message.contains("token invalid")
|| message.contains("token expired")
|| message.contains("unauthorized")
|| message.contains("invalid api key")
}
fn classify_quota_signal(signal: &str) -> Option<AccountFailureKind> {
let signal = signal.trim().to_ascii_lowercase();
match signal.as_str() {
"usage_limit_reached"
| "rate_limit_exceeded"
| "rate_limit_reached"
| "rate_limit_error"
| "workspace_owner_usage_limit_reached"
| "workspace_member_usage_limit_reached" => Some(AccountFailureKind::RateLimited),
"quota_exhausted"
| "quota_exceeded"
| "payment_required"
| "insufficient_quota"
| "workspace_owner_credits_depleted"
| "workspace_member_credits_depleted" => Some(AccountFailureKind::QuotaExhausted),
signal if signal.starts_with("billing_limit") => Some(AccountFailureKind::QuotaExhausted),
_ => None,
}
}
fn classify_quota_message(message: &str) -> Option<AccountFailureKind> {
let message = message.to_ascii_lowercase();
if message.contains("rate limit") || message.contains("usage limit") {
return Some(AccountFailureKind::RateLimited);
}
if message.contains("quota")
|| message.contains("payment required")
|| message.contains("billing limit")
{
return Some(AccountFailureKind::QuotaExhausted);
}
None
}
fn is_model_unsupported(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("model_not_supported")
|| value.contains("model_not_available")
|| (value.contains("model")
&& (value.contains("not supported")
|| value.contains("not available")
|| value.contains("not_supported")
|| value.contains("not_available")))
}
fn is_html(value: &str) -> bool {
let value = value.trim_start().to_ascii_lowercase();
value.starts_with("<!doctype") || value.starts_with("<html") || value.contains("<html")
}
+1 -2
View File
@@ -208,7 +208,6 @@ impl AccountManageService {
installation_id: installation_id.as_deref(),
session_id: None,
thread_id: None,
prompt_cache_key: None,
client_request_id: None,
turn_id: None,
};
@@ -229,7 +228,7 @@ impl AccountManageService {
Err(error) => {
tracing::warn!(
error = %error,
"failed to fetch supplemental account information during import"
"Failed to fetch supplemental account information during import"
);
ImportSupplementalAccountInfo {
status: import_status_from_usage_error(&error),
+12 -54
View File
@@ -7,16 +7,18 @@ use serde_json::{Value, json};
use tokio::sync::mpsc;
use crate::{
fleet::{account::AccountStatus, store::StoredAccount},
fleet::{
account::AccountStatus,
account_failure::{classify_client_failure, classify_response_failure},
store::StoredAccount,
},
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{SseEvent, encode_sse_event, parse_sse_events},
},
transport::{
CodexClientError, CodexRequestContext, is_banned_auth_signal, is_banned_upstream_error,
},
transport::{CodexClientError, CodexRequestContext},
},
};
@@ -103,7 +105,7 @@ impl AccountManageService {
account_id = %account.id,
plan_type,
error = %error,
"failed to refresh account plan models"
"Failed to refresh account plan models"
);
}
}
@@ -171,7 +173,6 @@ impl AccountManageService {
installation_id: Some(&installation_id),
session_id: None,
thread_id: None,
prompt_cache_key: None,
client_request_id: None,
turn_id: None,
};
@@ -214,7 +215,7 @@ impl AccountManageService {
tracing::warn!(
account_id,
error = %error,
"failed to inspect account after connection test"
"Failed to inspect account after connection test"
);
return None;
}
@@ -231,7 +232,7 @@ impl AccountManageService {
account_id,
status = %status,
error = %error,
"failed to persist account status after connection test"
"Failed to persist account status after connection test"
);
return None;
}
@@ -243,7 +244,7 @@ impl AccountManageService {
tracing::warn!(
account_id,
error = %error,
"failed to clear token refresh schedule after connection test"
"Failed to clear token refresh schedule after connection test"
);
}
self.sync_account_pool_best_effort(account_id, "connection test")
@@ -367,54 +368,11 @@ async fn process_sse_event(
}
fn account_status_from_client_error(error: &CodexClientError) -> Option<AccountStatus> {
if is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else {
match error {
CodexClientError::Upstream { status, .. } if status.as_u16() == 402 => {
Some(AccountStatus::QuotaExhausted)
}
CodexClientError::Upstream { status, body, .. } if status.as_u16() == 401 => {
Some(if is_banned_auth_signal(body) {
AccountStatus::Banned
} else {
AccountStatus::Expired
})
}
_ => None,
}
}
classify_client_failure(error)?.account_status()
}
fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
let code = failure
.upstream_code
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let message = failure.message.to_ascii_lowercase();
if matches!(code.as_str(), "quota_exceeded" | "insufficient_quota") || message.contains("quota")
{
return Some(AccountStatus::QuotaExhausted);
}
let authentication_failed = matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
) || message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired");
authentication_failed.then(|| {
if is_banned_auth_signal(&failure.message) {
AccountStatus::Banned
} else {
AccountStatus::Expired
}
})
classify_response_failure(failure)?.account_status()
}
fn take_sse_frame(buffer: &mut String) -> Option<String> {
+9 -6
View File
@@ -1,6 +1,7 @@
use secrecy::ExposeSecret;
use super::{AccountManageService, types::AccountManageError};
use crate::fleet::account_failure::apply_client_failure_effect;
impl AccountManageService {
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
@@ -50,15 +51,17 @@ impl AccountManageService {
installation_id: Some(&installation_id),
session_id: None,
thread_id: None,
prompt_cache_key: None,
client_request_id: None,
turn_id: None,
};
let raw = self
.codex
.fetch_usage(context)
.await
.map_err(|e| AccountManageError::FetchQuota(e.to_string()))?;
let raw = match self.codex.fetch_usage(context).await {
Ok(raw) => raw,
Err(error) => {
apply_client_failure_effect(&self.account_pool, &self.codex, account_id, &error)
.await;
return Err(AccountManageError::FetchQuota(error.to_string()));
}
};
let normalized = crate::fleet::quota::quota_from_usage(&raw);
self.account_pool
.apply_quota_snapshot(account_id, &normalized)
+2 -2
View File
@@ -121,7 +121,7 @@ impl AccountManageService {
account_id,
operation,
error = %error,
"failed to sync runtime account pool after admin account update"
"Failed to sync runtime account pool after admin account update"
);
}
}
@@ -145,7 +145,7 @@ impl AccountManageService {
tracing::warn!(
account_id,
error = %error,
"failed to inspect account while evicting websocket pool"
"Failed to inspect account while evicting WebSocket pool"
);
}
}
+4 -5
View File
@@ -7,6 +7,7 @@ use thiserror::Error;
use crate::{
fleet::{
account::{AccountStatus, AccountStatus as AcctStatus},
account_failure::classify_client_failure,
store::StoredAccountMetadata,
},
upstream::openai::{token_client::RefreshFailure, transport::CodexClientError},
@@ -323,11 +324,9 @@ fn normalized_plan_type(value: &str) -> Option<String> {
}
pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option<AccountStatus> {
if crate::upstream::openai::transport::is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else {
None
}
classify_client_failure(error)?
.account_status()
.filter(|status| *status == AccountStatus::Banned)
}
pub(super) fn parse_account_status(status: &str) -> Result<AcctStatus, AccountManageError> {
+1
View File
@@ -1,6 +1,7 @@
//! 账号、调度、配额与刷新领域。
pub mod account;
pub mod account_failure;
pub mod cookies;
pub mod import;
pub mod manage;
+56 -11
View File
@@ -33,8 +33,11 @@ use crate::upstream::openai::protocol::events::{
pub use crate::fleet::scheduler::RotationStrategy;
mod filters;
mod persistence;
mod state;
use persistence::AccountStatePersistence;
pub use filters::*;
pub use state::*;
@@ -71,6 +74,7 @@ pub struct AccountPoolService {
usage_store: Arc<dyn AccountUsageStore>,
request_interval: Arc<std::sync::RwLock<StdDuration>>,
candidate_changed: Arc<tokio::sync::Notify>,
state_persistence: AccountStatePersistence,
}
/// 一次请求对账号池 slot 的唯一所有权凭据。
@@ -191,6 +195,7 @@ impl AccountPoolService {
options: AccountPoolOptions,
request_interval_ms: u64,
) -> Self {
let state_persistence = AccountStatePersistence::new(Arc::clone(&store));
Self {
pool: Arc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))),
store,
@@ -199,6 +204,7 @@ impl AccountPoolService {
request_interval_ms,
))),
candidate_changed: Arc::new(tokio::sync::Notify::new()),
state_persistence,
}
}
@@ -321,7 +327,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist account request usage"
"Failed to persist account request usage"
);
}
}
@@ -334,15 +340,32 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist account status"
"Failed to persist account status"
);
false
}
};
let in_memory = self.pool.lock().await.set_status(account_id, status);
if in_memory {
self.candidate_changed.notify_waiters();
}
persisted || in_memory
}
/// 先在线性化的内存池中更新账号状态,再将持久化写入后台队列。
pub async fn set_status_immediately(&self, account_id: &str, status: AccountStatus) -> bool {
let in_memory = {
let mut pool = self.pool.lock().await;
let in_memory = pool.set_status(account_id, status);
self.state_persistence.set_status(account_id, status);
in_memory
};
if in_memory {
self.candidate_changed.notify_waiters();
}
in_memory
}
/// 标记账号 quota 冷却状态。
pub async fn mark_quota_limited_until(&self, account_id: &str, until: DateTime<Utc>) -> bool {
let persisted = match self.store.mark_quota_limited_until(account_id, until).await {
@@ -351,7 +374,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist quota cooldown state"
"Failed to persist quota cooldown state"
);
false
}
@@ -361,9 +384,31 @@ impl AccountPoolService {
.lock()
.await
.mark_quota_limited_until(account_id, until);
if in_memory {
self.candidate_changed.notify_waiters();
}
persisted || in_memory
}
/// 先在线性化的内存池中标记配额冷却,再将持久化写入后台队列。
pub async fn mark_quota_limited_until_immediately(
&self,
account_id: &str,
until: DateTime<Utc>,
) -> bool {
let in_memory = {
let mut pool = self.pool.lock().await;
let in_memory = pool.mark_quota_limited_until(account_id, until);
self.state_persistence
.mark_quota_limited_until(account_id, until);
in_memory
};
if in_memory {
self.candidate_changed.notify_waiters();
}
in_memory
}
/// 按请求上下文获取可用账号。
pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option<AccountLease> {
let refresh = self.pool.lock().await.acquire_with_status_refresh(request);
@@ -454,7 +499,7 @@ impl AccountPoolService {
tracing::error!(
account_id = %refreshed.account.id,
error = %error,
"failed to persist refreshed runtime account state"
"Failed to persist refreshed runtime account state"
);
}
if refreshed.sync_usage_window
@@ -469,7 +514,7 @@ impl AccountPoolService {
tracing::error!(
account_id = %refreshed.account.id,
error = %error,
"failed to persist refreshed runtime account usage window"
"Failed to persist refreshed runtime account usage window"
);
}
}
@@ -530,7 +575,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist account token usage"
"Failed to persist account token usage"
);
}
self.pool.lock().await.record_window_token_usage(
@@ -562,7 +607,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist empty response usage"
"Failed to persist empty response usage"
);
}
if image_generation_requested {
@@ -607,7 +652,7 @@ impl AccountPoolService {
tracing::error!(
account_id = %account_id,
error = %error,
"failed to read existing quota json before passive rate-limit sync"
"Failed to read existing quota JSON before passive rate-limit sync"
);
None
}
@@ -634,7 +679,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist verified quota snapshot"
"Failed to persist verified quota snapshot"
);
false
}
@@ -649,7 +694,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist verified quota window"
"Failed to persist verified quota window"
);
}
@@ -712,7 +757,7 @@ impl AccountPoolService {
tracing::error!(
account_id,
error = %error,
"failed to persist Cloudflare cooldown state"
"Failed to persist Cloudflare cooldown state"
);
false
}
+94
View File
@@ -0,0 +1,94 @@
//! 账号运行状态的有序后台持久化。
use std::sync::Arc;
use chrono::{DateTime, Utc};
use tokio::sync::mpsc;
use crate::fleet::{account::AccountStatus, store::AccountStore};
#[derive(Clone)]
pub(super) struct AccountStatePersistence {
sender: mpsc::UnboundedSender<AccountStateWrite>,
}
enum AccountStateWrite {
SetStatus {
account_id: String,
status: AccountStatus,
},
MarkQuotaLimitedUntil {
account_id: String,
until: DateTime<Utc>,
},
}
impl AccountStatePersistence {
pub(super) fn new(store: Arc<dyn AccountStore>) -> Self {
let (sender, receiver) = mpsc::unbounded_channel();
tokio::spawn(run(receiver, store));
Self { sender }
}
pub(super) fn set_status(&self, account_id: &str, status: AccountStatus) {
self.enqueue(AccountStateWrite::SetStatus {
account_id: account_id.to_owned(),
status,
});
}
pub(super) fn mark_quota_limited_until(&self, account_id: &str, until: DateTime<Utc>) {
self.enqueue(AccountStateWrite::MarkQuotaLimitedUntil {
account_id: account_id.to_owned(),
until,
});
}
fn enqueue(&self, write: AccountStateWrite) {
if self.sender.send(write).is_err() {
tracing::error!("Account state persistence worker stopped unexpectedly");
}
}
}
async fn run(
mut receiver: mpsc::UnboundedReceiver<AccountStateWrite>,
store: Arc<dyn AccountStore>,
) {
while let Some(write) = receiver.recv().await {
match write {
AccountStateWrite::SetStatus { account_id, status } => {
match store.set_status(&account_id, status).await {
Ok(true) => {}
Ok(false) => tracing::warn!(
account_id,
?status,
"Account disappeared before deferred status persistence"
),
Err(error) => tracing::error!(
account_id,
?status,
error = %error,
"Failed to persist deferred account status"
),
}
}
AccountStateWrite::MarkQuotaLimitedUntil { account_id, until } => {
match store.mark_quota_limited_until(&account_id, until).await {
Ok(true) => {}
Ok(false) => tracing::warn!(
account_id,
%until,
"Account disappeared before deferred quota persistence"
),
Err(error) => tracing::error!(
account_id,
%until,
error = %error,
"Failed to persist deferred quota state"
),
}
}
}
}
}
+2 -2
View File
@@ -418,7 +418,7 @@ impl AccountPool {
quota_limit_reached = selected.quota_limit_reached,
quota_cooldown_until = ?selected.quota_cooldown_until,
previous_slot_at = ?previous_slot_at,
"account selected for upstream request"
"Account selected for upstream request"
);
let selected_id = selected.id.clone();
@@ -645,7 +645,7 @@ impl AccountPool {
last_used_at = account.last_used_at.as_deref().unwrap_or_default(),
quota_limit_reached = account.quota_limit_reached,
quota_cooldown_until = ?account.quota_cooldown_until,
"account selected for model refresh"
"Account selected for model refresh"
);
selected_accounts.push((plan_type, account));
}
+16 -16
View File
@@ -13,6 +13,7 @@ use tracing::warn;
use crate::fleet::{
account::{Account, AccountStatus},
account_failure::apply_client_failure_effect,
pool::AccountPoolService,
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
@@ -41,7 +42,7 @@ pub struct QuotaRefreshService {
request_spacing: Duration,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
cookie_store: Option<PgCookieStore>,
account_pool: Option<Arc<AccountPoolService>>,
account_pool: Arc<AccountPoolService>,
}
impl QuotaRefreshService {
@@ -50,6 +51,7 @@ impl QuotaRefreshService {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
) -> Self {
Self {
@@ -60,7 +62,7 @@ impl QuotaRefreshService {
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
cookie_store: None,
account_pool: None,
account_pool,
}
}
@@ -73,6 +75,7 @@ impl QuotaRefreshService {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
min_refresh_interval_secs: u64,
) -> Self {
@@ -84,7 +87,7 @@ impl QuotaRefreshService {
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
cookie_store: None,
account_pool: None,
account_pool,
}
}
@@ -100,12 +103,6 @@ impl QuotaRefreshService {
self
}
/// 设置运行时账号池,用于刷新后同步内存调度状态。
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.account_pool = Some(account_pool);
self
}
/// 执行一次配额锁定或待验证账号刷新,复用调用方传入的最近刷新记录。
pub async fn refresh_locked_accounts(
&self,
@@ -176,7 +173,6 @@ impl QuotaRefreshService {
installation_id: Some(&installation_id),
session_id: None,
thread_id: None,
prompt_cache_key: None,
client_request_id: None,
turn_id: None,
})
@@ -212,10 +208,17 @@ impl QuotaRefreshService {
self.sync_runtime_pool_account(&account.id).await;
}
Err(error) => {
apply_client_failure_effect(
&self.account_pool,
&self.codex,
&account.id,
&error,
)
.await;
warn!(
account_id = %account.id,
error = %error,
"quota usage 拉取失败"
"Quota usage 拉取失败"
);
summary.failed += 1;
}
@@ -239,14 +242,11 @@ impl QuotaRefreshService {
}
async fn sync_runtime_pool_account(&self, account_id: &str) {
let Some(account_pool) = &self.account_pool else {
return;
};
if let Err(error) = account_pool.sync_account_from_store(account_id).await {
if let Err(error) = self.account_pool.sync_account_from_store(account_id).await {
warn!(
account_id = %account_id,
error = %error,
"failed to sync runtime account pool after quota refresh"
"Failed to sync runtime account pool after quota refresh"
);
}
}
+5 -5
View File
@@ -148,7 +148,7 @@ impl Drop for RefreshLeaseGuard {
drop(tasks.spawn_on(
async move {
if let Err(error) = store.release(&account_id, &owner).await {
warn!(account_id, error = %error, "failed to release cancelled token refresh lease");
warn!(account_id, error = %error, "Failed to release cancelled token refresh lease");
}
},
&runtime,
@@ -452,7 +452,7 @@ where
warn!(
account_id = %timer_account_id,
error = %error,
"scheduled token refresh failed"
"Scheduled token refresh failed"
);
}
});
@@ -503,7 +503,7 @@ where
warn!(
remaining_tasks = self.tasks.len(),
timeout_secs = REFRESH_TASK_SHUTDOWN_TIMEOUT.as_secs(),
"waiting for in-flight token refresh tasks timed out"
"Waiting for in-flight token refresh tasks timed out"
);
}
}
@@ -561,7 +561,7 @@ where
if delay.is_zero() {
debug!(
account_id = %account.id,
"refreshed token is already within refresh margin; skipping next timer"
"Refreshed token is already within refresh margin; skipping next timer"
);
return;
}
@@ -617,7 +617,7 @@ where
next_refresh_at = ?updated.next_refresh_at,
refresh_token_present = updated.refresh_token.is_some(),
refresh_token_rotated,
"token refresh succeeded"
"Token refresh succeeded"
);
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
+35 -5
View File
@@ -4,6 +4,9 @@ use std::time::Duration;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Postgres, Row, Transaction, postgres::PgPoolOptions};
use tracing::info;
const MIGRATION_LOCK_KEY: i64 = i64::from_be_bytes(*b"CPR-MIGR");
const MIGRATIONS: &[Migration] = &[Migration {
version: 1,
@@ -50,7 +53,7 @@ pub async fn migrate(pool: &PgPool) -> Result<(), sqlx::Error> {
validate_migration_catalog()?;
let mut tx = pool.begin().await?;
ensure_migrations_table(&mut tx).await?;
let migration_table_created = ensure_migrations_table(&mut tx).await?;
let applied = applied_migrations(&mut tx).await?;
if applied.is_empty() {
@@ -63,7 +66,14 @@ pub async fn migrate(pool: &PgPool) -> Result<(), sqlx::Error> {
record_migration(&mut tx, migration).await?;
}
tx.commit().await
tx.commit().await?;
if migration_table_created {
info!("数据库迁移表已创建");
} else {
info!("数据库迁移表已存在,跳过创建");
}
Ok(())
}
fn validate_migration_catalog() -> Result<(), sqlx::Error> {
@@ -86,9 +96,29 @@ fn validate_migration_catalog() -> Result<(), sqlx::Error> {
Ok(())
}
async fn ensure_migrations_table(tx: &mut Transaction<'_, Postgres>) -> Result<(), sqlx::Error> {
async fn ensure_migrations_table(tx: &mut Transaction<'_, Postgres>) -> Result<bool, sqlx::Error> {
sqlx::query("select pg_advisory_xact_lock($1)")
.bind(MIGRATION_LOCK_KEY)
.execute(&mut **tx)
.await?;
let exists = sqlx::query_scalar(
"select exists (
select 1
from information_schema.tables
where table_schema = current_schema()
and table_name = 'schema_migrations'
and table_type = 'BASE TABLE'
)",
)
.fetch_one(&mut **tx)
.await?;
if exists {
return Ok(false);
}
sqlx::query(
"create table if not exists schema_migrations (
"create table schema_migrations (
version bigint not null check (version > 0),
name text not null,
checksum text not null,
@@ -98,7 +128,7 @@ async fn ensure_migrations_table(tx: &mut Transaction<'_, Postgres>) -> Result<(
)
.execute(&mut **tx)
.await?;
Ok(())
Ok(true)
}
async fn applied_migrations(
+45 -5
View File
@@ -4,7 +4,7 @@ use std::{
backtrace::Backtrace,
env, fmt, fs,
fs::{File, OpenOptions},
io::{self, Write},
io::{self, IsTerminal, Write},
panic,
path::{Path, PathBuf},
sync::mpsc,
@@ -22,7 +22,7 @@ use tracing_subscriber::{
util::SubscriberInitExt,
};
use crate::infra::time::{china_date, china_rfc3339_millis};
use crate::infra::time::{china_date, china_datetime, china_rfc3339_millis};
const LOG_PREFIX: &str = "codex-proxy-rs";
const LOG_SUFFIX: &str = "log";
@@ -169,6 +169,7 @@ pub fn init_tracing(config: &TracingConfig) -> Result<LogGuard, LogError> {
let mut writer_guards = Vec::new();
let mut counters = Vec::new();
let stdout_is_terminal = io::stdout().is_terminal();
let stdout_writer = config.stdout.then(|| {
let (writer, guard) = NonBlockingBuilder::default()
.lossy(true)
@@ -194,13 +195,20 @@ pub fn init_tracing(config: &TracingConfig) -> Result<LogGuard, LogError> {
writer
});
let stdout_layer = stdout_writer.map(json_layer);
let (compact_stdout_writer, json_stdout_writer) = match stdout_writer {
Some(writer) if stdout_is_terminal => (Some(writer), None),
Some(writer) => (None, Some(writer)),
None => (None, None),
};
let compact_stdout_layer = compact_stdout_writer.map(compact_layer);
let json_stdout_layer = json_stdout_writer.map(json_layer);
let file_layer = file_writer.map(json_layer);
let drop_monitor = DroppedLogMonitor::start(counters)?;
let init_result = tracing_subscriber::registry()
.with(filter)
.with(stdout_layer)
.with(compact_stdout_layer)
.with(json_stdout_layer)
.with(file_layer)
.try_init()
.map_err(|_| LogError::SubscriberAlreadyInitialized);
@@ -217,13 +225,36 @@ pub fn init_tracing(config: &TracingConfig) -> Result<LogGuard, LogError> {
})
}
fn compact_layer<S>(
writer: tracing_appender::non_blocking::NonBlocking,
) -> tracing_subscriber::fmt::Layer<
S,
tracing_subscriber::fmt::format::DefaultFields,
tracing_subscriber::fmt::format::Format<
tracing_subscriber::fmt::format::Compact,
ChinaConsoleTimer,
>,
tracing_appender::non_blocking::NonBlocking,
>
where
S: tracing::Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
{
tracing_fmt::layer()
.compact()
.with_timer(ChinaConsoleTimer)
.with_writer(writer)
.with_ansi(true)
.with_target(true)
.with_level(true)
}
fn install_panic_hook() {
let previous_hook = panic::take_hook();
panic::set_hook(Box::new(move |panic_info| {
tracing::error!(
panic = %panic_info,
backtrace = %Backtrace::force_capture(),
"process panicked"
"Process panicked"
);
previous_hook(panic_info);
}));
@@ -316,6 +347,15 @@ impl Write for ChinaDailyFileAppender {
}
}
#[derive(Debug, Clone, Copy)]
struct ChinaConsoleTimer;
impl FormatTime for ChinaConsoleTimer {
fn format_time(&self, writer: &mut Writer<'_>) -> fmt::Result {
write!(writer, "{}", china_datetime(&Utc::now()))
}
}
#[derive(Debug, Clone, Copy)]
struct ChinaLogTimer;
+4 -8
View File
@@ -10,15 +10,11 @@ use rand::Rng;
const IDENTITY_SECRET_FILE_NAME: &str = "identity_hmac_secret";
fn data_dir() -> PathBuf {
dirs::data_local_dir().unwrap_or_else(|| PathBuf::from(".runtime/data"))
}
/// 确保本地数据目录存在。
pub fn ensure_data_dir() -> std::io::Result<PathBuf> {
let dir = data_dir();
std::fs::create_dir_all(&dir)?;
Ok(dir)
pub fn ensure_data_dir(directory: impl AsRef<Path>) -> std::io::Result<PathBuf> {
let directory = directory.as_ref();
std::fs::create_dir_all(directory)?;
Ok(directory.to_path_buf())
}
/// 读取或创建账号身份隔离使用的 256-bit HMAC 密钥。
+1 -1
View File
@@ -86,7 +86,7 @@ impl KeyVerifier {
pub async fn flush_pending_last_used(&self) {
let updates = self.take_pending_last_used();
if let Err(error) = self.store.touch_last_used_batch(&updates).await {
tracing::error!(error = %error, "failed to flush client key last_used_at batch");
tracing::error!(error = %error, "Failed to flush client key last_used_at batch");
merge_last_used_updates(
&mut self
.pending_last_used
+23 -13
View File
@@ -1,25 +1,29 @@
use std::{env, error::Error, io};
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
use codex_proxy_rs::bootstrap::config::BootstrapConfig;
fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
let mut arguments = env::args_os().skip(1);
match arguments.next().as_deref().and_then(|value| value.to_str()) {
None | Some("serve") => {
reject_extra_arguments(arguments)?;
codex_proxy_rs::bootstrap::services::run().await
let config = BootstrapConfig::load()?;
runtime()?.block_on(codex_proxy_rs::bootstrap::services::run(config))
}
Some("rebuild-buckets") => {
reject_extra_arguments(arguments)?;
let config = codex_proxy_rs::bootstrap::config::AppConfig::load()?;
let pool = codex_proxy_rs::infra::database::connect(&config.database.url).await?;
let report = codex_proxy_rs::telemetry::rebuild::rebuild_buckets(&pool).await?;
println!(
"request_time_buckets rebuilt: cutoff={}, deleted={}, rebuilt={}",
report.cutoff.to_rfc3339(),
report.deleted_rows,
report.rebuilt_rows
);
Ok(())
let config = BootstrapConfig::load()?;
runtime()?.block_on(async move {
let pool = codex_proxy_rs::infra::database::connect(config.database_url()).await?;
let report = codex_proxy_rs::telemetry::rebuild::rebuild_buckets(&pool).await?;
println!(
"request_time_buckets rebuilt: cutoff={}, deleted={}, rebuilt={}",
report.cutoff.to_rfc3339(),
report.deleted_rows,
report.rebuilt_rows
);
Ok(())
})
}
Some(command) => Err(invalid_cli(&format!(
"unknown command {command:?}; expected serve or rebuild-buckets"
@@ -28,6 +32,12 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
}
}
fn runtime() -> Result<tokio::runtime::Runtime, io::Error> {
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
}
fn reject_extra_arguments(mut arguments: impl Iterator) -> Result<(), io::Error> {
if arguments.next().is_some() {
return Err(invalid_cli("unexpected extra command arguments"));
+28 -70
View File
@@ -6,7 +6,7 @@ use thiserror::Error;
use crate::infra::format::nonnegative_i64_to_u64;
use crate::telemetry::{
account_usage::store::{PgAccountUsageStore, UsageListRecord, UsageSummary},
account_usage::store::{PgAccountUsageStore, UsageListRecord},
billing,
buckets::query::{
ModelBucketUsage, ModelUsageWindow, PgRequestBucketQuery, UsageBucketTotals,
@@ -42,13 +42,25 @@ impl AccountUsageQueryService {
.map_err(|_| AccountUsageQueryError::List)
}
/// 汇总账号用量。
pub async fn summary(&self) -> Result<AccountUsageSummary, AccountUsageQueryError> {
self.store
.usage_summary()
/// 汇总请求时间桶保留期内的全局用量与计费。
pub async fn retained_summary(&self) -> Result<RetainedUsageSummary, AccountUsageQueryError> {
let buckets = self
.buckets
.retained_usage_buckets()
.await
.map(AccountUsageSummary::from)
.map_err(|_| AccountUsageQueryError::Summary)
.map_err(|_| AccountUsageQueryError::RetainedSummary)?;
let mut summary = RetainedUsageSummary::default();
for bucket in buckets {
summary.billing_amount_usd += billing::calculate_billing_amount(
nonnegative_i64_to_u64(bucket.totals.input_tokens),
nonnegative_i64_to_u64(bucket.totals.output_tokens),
nonnegative_i64_to_u64(bucket.totals.cached_tokens),
&bucket.model,
bucket.service_tier.as_deref(),
);
summary.totals.add(bucket.totals);
}
Ok(summary)
}
/// 列出指定时间范围内的聚合时间桶。
@@ -95,12 +107,19 @@ impl AccountUsageQueryService {
pub enum AccountUsageQueryError {
#[error("failed to list account usage")]
List,
#[error("failed to summarize account usage")]
Summary,
#[error("failed to summarize retained usage")]
RetainedSummary,
#[error("failed to list usage time buckets")]
TimeBuckets,
}
/// 请求时间桶保留期内的全局用量与计费汇总。
#[derive(Debug, Clone, Default, PartialEq)]
pub struct RetainedUsageSummary {
pub totals: UsageBucketTotals,
pub billing_amount_usd: f64,
}
/// 管理端账号用量记录。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUsageRecord {
@@ -129,23 +148,6 @@ pub struct AccountUsageRecord {
pub last_used_at: Option<DateTime<Utc>>,
}
/// 管理端账号用量汇总。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageSummary {
pub account_count: i64,
pub request_count: i64,
pub empty_response_count: i64,
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub reasoning_tokens: i64,
pub total_tokens: i64,
pub image_input_tokens: i64,
pub image_output_tokens: i64,
pub image_request_count: i64,
pub image_request_failed_count: i64,
}
/// 管理端时间桶用量记录。
#[derive(Debug, Clone, PartialEq)]
pub struct AccountUsageTimeBucket {
@@ -163,31 +165,10 @@ pub struct AccountUsageTimeBucket {
pub latency_count: i64,
pub max_latency_ms: i64,
pub min_latency_ms: i64,
pub billing_amount_usd: Option<f64>,
}
impl From<UsageSummary> for AccountUsageSummary {
fn from(s: UsageSummary) -> Self {
Self {
account_count: s.account_count,
request_count: s.request_count,
empty_response_count: s.empty_response_count,
input_tokens: s.input_tokens,
output_tokens: s.output_tokens,
cached_tokens: s.cached_tokens,
reasoning_tokens: s.reasoning_tokens,
total_tokens: s.total_tokens,
image_input_tokens: s.image_input_tokens,
image_output_tokens: s.image_output_tokens,
image_request_count: s.image_request_count,
image_request_failed_count: s.image_request_failed_count,
}
}
}
impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
fn from(record: UsageTimeBucketRecord) -> Self {
let billing_amount_usd = time_bucket_billing_amount_usd(&record);
Self {
bucket_start: record.bucket_start,
model: record.model,
@@ -203,33 +184,10 @@ impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
latency_count: record.latency_count,
max_latency_ms: record.max_latency_ms,
min_latency_ms: record.min_latency_ms,
billing_amount_usd,
}
}
}
fn time_bucket_billing_amount_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(record.input_tokens);
let output_tokens = nonnegative_i64_to_u64(record.output_tokens);
let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens);
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
return None;
}
let model = record.model.trim();
if model.is_empty() {
return None;
}
Some(billing::calculate_billing_amount(
input_tokens,
output_tokens,
cached_tokens,
model,
record.service_tier.as_deref(),
))
}
impl From<UsageListRecord> for AccountUsageRecord {
fn from(usage: UsageListRecord) -> Self {
Self {
@@ -129,22 +129,6 @@ select
from account_usage au
left join accounts a on a.id = au.account_id";
const USAGE_SUMMARY_SQL: &str = r"
select
count(*) as account_count,
coalesce(sum(request_count), 0)::bigint as request_count,
coalesce(sum(empty_response_count), 0)::bigint as empty_response_count,
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(reasoning_tokens), 0)::bigint as reasoning_tokens,
coalesce(sum(total_tokens), 0)::bigint as total_tokens,
coalesce(sum(image_input_tokens), 0)::bigint as image_input_tokens,
coalesce(sum(image_output_tokens), 0)::bigint as image_output_tokens,
coalesce(sum(image_request_count), 0)::bigint as image_request_count,
coalesce(sum(image_request_failed_count), 0)::bigint as image_request_failed_count
from account_usage";
/// PostgreSQL 用量存储错误。
#[derive(Debug, Error)]
pub enum PgAccountUsageStoreError {
@@ -312,35 +296,6 @@ pub struct UsageListRecord {
pub last_used_at: Option<DateTime<Utc>>,
}
/// 账号用量汇总。
#[derive(Debug, Clone)]
pub struct UsageSummary {
/// 账号数。
pub account_count: i64,
/// 总请求数。
pub request_count: i64,
/// 总空响应数。
pub empty_response_count: i64,
/// 总输入 token。
pub input_tokens: i64,
/// 总输出 token。
pub output_tokens: i64,
/// 总缓存 token。
pub cached_tokens: i64,
/// 总 reasoning token。
pub reasoning_tokens: i64,
/// 总 token。
pub total_tokens: i64,
/// 总图片输入 token。
pub image_input_tokens: i64,
/// 总图片输出 token。
pub image_output_tokens: i64,
/// 总图片请求成功数。
pub image_request_count: i64,
/// 总图片请求失败数。
pub image_request_failed_count: i64,
}
/// PostgreSQL 用量存储。
#[derive(Clone)]
pub struct PgAccountUsageStore {
@@ -492,25 +447,6 @@ where account_id in (",
let rows = builder.build().fetch_all(&self.pool).await?;
rows.iter().map(usage_list_from_row).collect()
}
/// 汇总账号用量。
pub async fn usage_summary(&self) -> PgAccountUsageStoreResult<UsageSummary> {
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
Ok(UsageSummary {
account_count: row.get("account_count"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
reasoning_tokens: row.get("reasoning_tokens"),
total_tokens: row.get("total_tokens"),
image_input_tokens: row.get("image_input_tokens"),
image_output_tokens: row.get("image_output_tokens"),
image_request_count: row.get("image_request_count"),
image_request_failed_count: row.get("image_request_failed_count"),
})
}
}
#[async_trait]
+85 -1
View File
@@ -1,6 +1,6 @@
//! 管理端用量计费规则。
const LONG_CONTEXT_THRESHOLD: u64 = 272_000;
pub(crate) const LONG_CONTEXT_THRESHOLD: u64 = 272_000;
#[derive(Debug, Clone, Copy)]
struct ModelPricing {
@@ -226,6 +226,24 @@ fn calculate_billing_with_pricing(
) -> BillingBreakdown {
let is_long = input_tokens > LONG_CONTEXT_THRESHOLD;
calculate_billing_with_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
pricing,
service_tier,
is_long,
)
}
fn calculate_billing_with_pricing_for_context(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
pricing: ModelPricing,
service_tier: Option<&str>,
is_long: bool,
) -> BillingBreakdown {
let mut input_price = pricing.input_price_per_mtoken;
let mut output_price = pricing.output_price_per_mtoken;
let mut cache_read_price = pricing.cache_read_price_per_mtoken;
@@ -284,6 +302,36 @@ fn calculate_billing_with_pricing(
}
}
/// 按模型、档位和已确定的上下文长度档位计算聚合用量费用。
///
/// 聚合查询必须显式传入 `is_long`,避免许多短请求相加后被误判为单次长上下文请求。
pub(crate) fn calculate_aggregate_billing(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
model: &str,
service_tier: Option<&str>,
is_long: bool,
) -> BillingBreakdown {
let pricing = model_pricing(model).unwrap_or_else(|| {
highest_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
service_tier,
is_long,
)
});
calculate_billing_with_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
pricing,
service_tier,
is_long,
)
}
/// 按已登记的官方价格计算美元计费金额,未知模型按最高费率计算。
pub(crate) fn calculate_billing_amount(
input_tokens: u64,
@@ -346,6 +394,42 @@ fn highest_pricing_for_request(
})
}
fn highest_pricing_for_context(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
service_tier: Option<&str>,
is_long: bool,
) -> ModelPricing {
MODEL_PRICING_RULES
.iter()
.fold(ModelPricing::new(0.0, 0.0), |highest, rule| {
let highest_amount = calculate_billing_with_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
highest,
service_tier,
is_long,
)
.total_amount;
let candidate_amount = calculate_billing_with_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
rule.pricing,
service_tier,
is_long,
)
.total_amount;
if candidate_amount >= highest_amount {
rule.pricing
} else {
highest
}
})
}
fn normalize_billing_model_name(model: &str) -> String {
let mut model = model.trim().to_ascii_lowercase();
model = model.trim_start_matches('/').to_string();
+42 -1
View File
@@ -29,6 +29,17 @@ where bucket_start >= $1 and bucket_start <= $2
group by bucket_start, model, service_tier
order by bucket_start asc, model asc, service_tier asc";
const RETAINED_USAGE_BUCKETS_SQL: &str = r"
select
model,
nullif(service_tier, '__unknown__') as service_tier,
coalesce(sum(success_count + error_count), 0)::bigint as request_count,
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens
from request_time_buckets
group by bucket_start, model, service_tier";
/// 请求时间桶查询错误。
#[derive(Debug, Error)]
pub enum PgRequestBucketQueryError {
@@ -72,8 +83,16 @@ pub struct UsageBucketTotals {
pub cached_tokens: i64,
}
/// 保留期内单个计费桶的用量事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedUsageBucket {
pub model: String,
pub service_tier: Option<String>,
pub totals: UsageBucketTotals,
}
impl UsageBucketTotals {
fn add(&mut self, other: Self) {
pub(crate) fn add(&mut self, other: Self) {
self.request_count = self.request_count.saturating_add(other.request_count);
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
@@ -127,6 +146,28 @@ impl PgRequestBucketQuery {
Ok(rows.iter().map(usage_time_bucket_from_row).collect())
}
/// 列出请求时间桶保留期内的全局计费桶,不受当前账号集合影响。
pub async fn retained_usage_buckets(
&self,
) -> Result<Vec<RetainedUsageBucket>, PgRequestBucketQueryError> {
let rows = sqlx::query(RETAINED_USAGE_BUCKETS_SQL)
.fetch_all(&self.pool)
.await?;
Ok(rows
.iter()
.map(|row| RetainedUsageBucket {
model: row.get("model"),
service_tier: row.get("service_tier"),
totals: UsageBucketTotals {
request_count: row.get("request_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
},
})
.collect())
}
pub async fn usage_by_windows(
&self,
windows: &[UsageBucketWindow],
+85 -86
View File
@@ -8,14 +8,15 @@ use crate::infra::{
format_billing_amount, format_compact_number, format_duration_ms, format_percent,
format_rate, format_tokens, nonnegative_i64_to_u64,
},
time::{china_day_start, china_hour, china_hour_start, china_quarter_hour_start},
time::{china_day_start, china_quarter_hour_start},
};
use super::account_usage::query::{AccountUsageSummary, AccountUsageTimeBucket};
use super::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
const TREND_SLOT_MINUTES: i64 = 15;
#[derive(Debug, Clone, Copy, Default)]
pub struct DashboardAccountCounts {
@@ -78,7 +79,7 @@ struct DashboardTokenCardData {
today_tokens_value: u64,
yesterday_tokens_value: u64,
total_tokens: String,
today_billing_amount_usd: String,
total_billing_amount_usd: String,
}
#[derive(Debug, Serialize)]
@@ -117,13 +118,13 @@ struct DashboardTrendPointData {
errors: String,
errors_value: u64,
latency: String,
latency_value: u64,
latency_value: Option<u64>,
max_latency: String,
max_latency_value: u64,
max_latency_value: Option<u64>,
min_latency: String,
min_latency_value: u64,
min_latency_value: Option<u64>,
success_rate: String,
success_rate_value: f64,
success_rate_value: Option<f64>,
}
#[derive(Debug, Serialize)]
@@ -173,25 +174,25 @@ impl UsageWindow {
pub fn dashboard_cards(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
retained_usage: &RetainedUsageSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, buckets, lifetime_usage, Utc::now())
dashboard_cards_at(accounts, buckets, retained_usage, Utc::now())
}
pub fn dashboard_cards_at(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
retained_usage: &RetainedUsageSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
let yesterday_start = today_start - Duration::days(1);
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_billing_amount = billing_window(buckets, today_start, now).unwrap_or(0.0);
let total_requests = nonnegative_i64_to_u64(lifetime_usage.request_count);
let total_tokens = nonnegative_i64_to_u64(lifetime_usage.total_tokens);
let total_cached_tokens = nonnegative_i64_to_u64(lifetime_usage.cached_tokens);
let total_requests = nonnegative_i64_to_u64(retained_usage.totals.request_count);
let total_tokens = nonnegative_i64_to_u64(retained_usage.totals.input_tokens)
.saturating_add(nonnegative_i64_to_u64(retained_usage.totals.output_tokens));
let total_cached_tokens = nonnegative_i64_to_u64(retained_usage.totals.cached_tokens);
DashboardCardsData {
accounts: DashboardAccountsCardData {
@@ -213,13 +214,13 @@ pub fn dashboard_cards_at(
today_tokens_value: today.tokens(),
yesterday_tokens_value: yesterday.tokens(),
total_tokens: format_tokens(total_tokens),
today_billing_amount_usd: format_billing_amount(today_billing_amount),
total_billing_amount_usd: format_billing_amount(retained_usage.billing_amount_usd),
},
cache: DashboardCacheCardData {
today_hit_rate: format_rate(today.cache_hit_rate()),
today_hit_rate_value: today.cache_hit_rate(),
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
total_hit_rate: format_rate(summary_cache_hit_rate(lifetime_usage)),
total_hit_rate: format_rate(summary_cache_hit_rate(retained_usage)),
total_cached_tokens: format_tokens(total_cached_tokens),
average_first_token_latency_ms: format_optional_duration_ms(
today.avg_first_token_latency(),
@@ -240,41 +241,51 @@ pub fn dashboard_trend_data_at(
kind: DashboardTrendKind,
now: DateTime<Utc>,
) -> DashboardTrendData {
let current_hour = china_hour_start(now);
let current_slot = china_quarter_hour_start(now);
let start = china_day_start(now);
let elapsed_hours = current_hour.signed_duration_since(start).num_hours();
let mut buckets = (0..=elapsed_hours)
.map(|index| (start + Duration::hours(index), UsageWindow::default()))
let elapsed_slots =
current_slot.signed_duration_since(start).num_minutes() / TREND_SLOT_MINUTES;
let mut buckets = (0..=elapsed_slots)
.map(|index| {
(
start + Duration::minutes(TREND_SLOT_MINUTES * index),
UsageWindow::default(),
)
})
.collect::<Vec<_>>();
for record in records {
if record.bucket_start < start || record.bucket_start > now {
continue;
}
let record_hour = china_hour_start(record.bucket_start);
if let Some((_, bucket)) = buckets
.iter_mut()
.find(|(bucket_start, _)| *bucket_start == record_hour)
{
apply_bucket(bucket, record);
}
let record_slot = china_quarter_hour_start(record.bucket_start);
let slot_index =
record_slot.signed_duration_since(start).num_minutes() / TREND_SLOT_MINUTES;
let Ok(slot_index) = usize::try_from(slot_index) else {
continue;
};
let Some((_, bucket)) = buckets.get_mut(slot_index) else {
continue;
};
apply_bucket(bucket, record);
}
let points = buckets
.iter()
.map(|(bucket_start, bucket)| {
let latency = bucket.avg_first_token_latency().unwrap_or(0);
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
let latency = bucket.avg_first_token_latency();
let max_latency = bucket.max_first_token_bucket_latency();
let min_latency = bucket.min_first_token_bucket_latency();
let cache_hit_rate = bucket.cache_hit_rate().unwrap_or(0.0);
let success_rate = if bucket.requests > 0 {
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
let success_rate = (bucket.requests > 0).then(|| {
(bucket.requests.saturating_sub(bucket.errors) as f64 / bucket.requests as f64
* 1000.0)
.round()
/ 10.0
} else {
0.0
};
});
let elapsed_minutes = bucket_start.signed_duration_since(start).num_minutes();
DashboardTrendPointData {
time: format!("{:02}", china_hour(bucket_start)),
time: format!("{:02}:{:02}", elapsed_minutes / 60, elapsed_minutes % 60),
requests: format_compact_number(bucket.requests),
requests_value: bucket.requests,
input_tokens: format_tokens(bucket.input_tokens),
@@ -287,13 +298,15 @@ pub fn dashboard_trend_data_at(
tokens_value: bucket.tokens(),
errors: format_compact_number(bucket.errors),
errors_value: bucket.errors,
latency: format_optional_duration_ms(Some(latency)),
latency: format_optional_duration_ms(latency),
latency_value: latency,
max_latency: format_optional_duration_ms(Some(max_latency)),
max_latency: format_optional_duration_ms(max_latency),
max_latency_value: max_latency,
min_latency: format_optional_duration_ms(Some(min_latency)),
min_latency: format_optional_duration_ms(min_latency),
min_latency_value: min_latency,
success_rate: format_percent(success_rate),
success_rate: success_rate
.map(format_percent)
.unwrap_or_else(|| "—".to_string()),
success_rate_value: success_rate,
}
})
@@ -379,52 +392,48 @@ fn trend_summary(
trend_summary_item(
kind,
"输入",
points.iter().map(|p| p.input_tokens_value).sum(),
Some(points.iter().map(|p| p.input_tokens_value).sum()),
None,
),
trend_summary_item(
kind,
"输出",
points.iter().map(|p| p.output_tokens_value).sum(),
Some(points.iter().map(|p| p.output_tokens_value).sum()),
None,
),
trend_summary_item(
kind,
"缓存",
points.iter().map(|p| p.cached_tokens_value).sum(),
Some(points.iter().map(|p| p.cached_tokens_value).sum()),
None,
),
],
DashboardTrendKind::Latency => {
let samples = points
let (latency_sum, latency_count) = points
.iter()
.filter(|point| point.latency_value > 0)
.collect::<Vec<_>>();
let average = if samples.is_empty() {
0
} else {
samples.iter().map(|point| point.latency_value).sum::<u64>() / samples.len() as u64
};
.filter_map(|point| point.latency_value)
.fold((0_u64, 0_u64), |(sum, count), latency| {
(sum + latency, count + 1)
});
let average = (latency_count > 0).then(|| latency_sum / latency_count);
vec![
trend_summary_item(kind, "平均", average, None),
trend_summary_item(
kind,
"最高",
samples
points
.iter()
.map(|p| p.max_latency_value)
.max()
.unwrap_or(0),
.filter_map(|point| point.max_latency_value)
.max(),
None,
),
trend_summary_item(
kind,
"最低",
samples
points
.iter()
.filter_map(|p| (p.min_latency_value > 0).then_some(p.min_latency_value))
.min()
.unwrap_or(0),
.filter_map(|point| point.min_latency_value)
.min(),
None,
),
]
@@ -432,12 +441,13 @@ fn trend_summary(
DashboardTrendKind::Errors => {
let errors = points.iter().map(|point| point.errors_value).sum::<u64>();
let requests = points.iter().map(|point| point.requests_value).sum::<u64>();
let success_rate = (requests > 0)
.then(|| ((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0);
let success_rate = (requests > 0).then(|| {
(requests.saturating_sub(errors) as f64 / requests as f64 * 1000.0).round() / 10.0
});
vec![
trend_summary_item(kind, "错误数", errors, None),
trend_summary_item(kind, "成功率", 0, success_rate),
trend_summary_item(kind, "总请求", requests, None),
trend_summary_item(kind, "错误数", Some(errors), None),
trend_summary_item(kind, "成功率", None, success_rate),
trend_summary_item(kind, "总请求", Some(requests), None),
]
}
}
@@ -446,16 +456,18 @@ fn trend_summary(
fn trend_summary_item(
kind: DashboardTrendKind,
label: &str,
value: u64,
value: Option<u64>,
ratio: Option<f64>,
) -> DashboardTrendSummaryData {
DashboardTrendSummaryData {
label: label.to_string(),
value: match kind {
DashboardTrendKind::Usage => format_tokens(value),
DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)),
DashboardTrendKind::Errors => format_compact_number(value),
},
value: value
.map(|value| match kind {
DashboardTrendKind::Usage => format_tokens(value),
DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)),
DashboardTrendKind::Errors => format_compact_number(value),
})
.unwrap_or_else(|| "—".to_string()),
ratio: ratio.map(format_percent),
}
}
@@ -474,19 +486,6 @@ fn usage_window(
})
}
fn billing_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Option<f64> {
let billing_amounts = records
.iter()
.filter(|record| record.bucket_start >= start && record.bucket_start < end)
.filter_map(|record| record.billing_amount_usd)
.collect::<Vec<_>>();
(!billing_amounts.is_empty()).then(|| billing_amounts.into_iter().sum())
}
fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
@@ -528,10 +527,10 @@ fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char
}
}
fn summary_cache_hit_rate(summary: &AccountUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
fn summary_cache_hit_rate(summary: &RetainedUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.totals.input_tokens);
(input_tokens > 0)
.then(|| nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
.then(|| nonnegative_i64_to_u64(summary.totals.cached_tokens) as f64 / input_tokens as f64)
}
fn format_optional_duration_ms(value: Option<u64>) -> String {
+3 -3
View File
@@ -69,7 +69,7 @@ impl Recorder {
account_id = event.account_id,
model = event.model,
status_code = event.status_code,
"rejected invalid success usage fact"
"Rejected invalid success usage fact"
);
return Err(RecorderError::InvalidUsageFact);
}
@@ -124,7 +124,7 @@ pub(crate) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'
tracing::error!(
account_id = record.account_id.unwrap_or(""),
error = %error,
"failed to record dispatch error event"
"Failed to record dispatch error event"
);
}
}
@@ -169,7 +169,7 @@ pub(crate) async fn record_response_event(record: ResponseUsageRecord<'_>) {
event.metadata = metadata;
if let Err(error) = record.recorder.record_usage(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "failed to record response event");
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record response event");
}
}
+801
View File
@@ -0,0 +1,801 @@
//! 管理端请求观测聚合。
use std::collections::{BTreeMap, HashMap};
use chrono::{DateTime, Duration, Utc};
use serde::Serialize;
use sqlx::{AssertSqlSafe, PgPool, Row};
use crate::telemetry::billing::{LONG_CONTEXT_THRESHOLD, calculate_aggregate_billing};
use super::store::PgUsageRecordStoreResult;
const DEFAULT_RANGE_DAYS: i64 = 7;
// 一个请求可能写入多个错误事实。查询时优先把存在成功事实的 request_id 视为成功,
// 否则只保留该 request_id 的最新错误;缺少 request_id 时退化到事实 id。
const TERMINAL_REQUESTS_CTE: &str = r#"
success_terminals as (
select distinct on (request_key)
request_key, id, request_id, terminal_at, true as is_success,
client_api_key_id, provider, account_id, transport, model_name,
null::text as failure_class, latency_ms, first_token_ms, billing_model,
service_tier, input_tokens, output_tokens, cached_tokens
from (
select
case when request_id is null then 'event:' || id else 'request:' || request_id end as request_key,
id, request_id, created_at as terminal_at, client_api_key_id, provider,
account_id, transport,
coalesce(nullif(requested_model, ''), model) || ' → ' ||
coalesce(nullif(upstream_model, ''), model) as model_name,
latency_ms, first_token_ms,
coalesce(nullif(upstream_model, ''), model) as billing_model,
service_tier, input_tokens, output_tokens, cached_tokens
from usage_records
where created_at >= $1 and created_at < $2
) candidates
order by request_key, terminal_at desc, id desc
),
error_terminals as (
select distinct on (request_key)
request_key, id, request_id, terminal_at, false as is_success,
client_api_key_id, provider, account_id, transport, model_name,
failure_class, latency_ms, null::bigint as first_token_ms,
null::text as billing_model, null::text as service_tier,
null::bigint as input_tokens, null::bigint as output_tokens,
null::bigint as cached_tokens
from (
select
case when request_id is null then 'event:' || id else 'request:' || request_id end as request_key,
id, request_id, created_at as terminal_at, client_api_key_id, provider,
account_id, transport, coalesce(nullif(model, ''), '未知模型') as model_name,
coalesce(
nullif(failure_class, ''),
nullif(metadata_json->>'failureClass', ''),
nullif(metadata_json->>'upstreamCode', ''),
nullif(metadata_json->>'terminal', ''),
'未分类错误'
) as failure_class,
latency_ms
from ops_error_logs
where created_at >= $1 and created_at < $2
) candidates
order by request_key, terminal_at desc, id desc
),
final_errors as (
select errors.*
from error_terminals errors
where not exists (
select 1
from usage_records successes
where
(errors.request_id is not null and successes.request_id = errors.request_id)
or (errors.request_id is null and successes.id = errors.id)
)
),
terminal_requests as (
select * from success_terminals
union all
select * from final_errors
)
"#;
/// 观测时间粒度。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InsightsGranularity {
QuarterHour,
Hour,
Day,
}
impl InsightsGranularity {
pub fn for_range(start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
let range = end.signed_duration_since(start);
if range <= Duration::days(2) {
Self::QuarterHour
} else if range <= Duration::days(14) {
Self::Hour
} else {
Self::Day
}
}
fn as_str(self) -> &'static str {
match self {
Self::QuarterHour => "15m",
Self::Hour => "1h",
Self::Day => "1d",
}
}
fn step_seconds(self) -> i64 {
match self {
Self::QuarterHour => 15 * 60,
Self::Hour => 60 * 60,
Self::Day => 24 * 60 * 60,
}
}
fn bucket_expression(self, column: &str) -> String {
let interval = match self {
Self::QuarterHour => "15 minutes",
Self::Hour => "1 hour",
Self::Day => "1 day",
};
let origin = if self == Self::Day {
"1970-01-01 16:00:00+00"
} else {
"1970-01-01 00:00:00+00"
};
format!("date_bin(interval '{interval}', {column}, timestamptz '{origin}')")
}
}
/// 热点诊断维度。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UsageDiagnosticsDimension {
Model,
Account,
ApiKey,
Provider,
Transport,
FailureClass,
}
impl UsageDiagnosticsDimension {
pub fn parse(value: &str) -> Option<Self> {
match value {
"model" => Some(Self::Model),
"account" => Some(Self::Account),
"apiKey" | "api_key" => Some(Self::ApiKey),
"provider" => Some(Self::Provider),
"transport" => Some(Self::Transport),
"failureClass" | "failure_class" => Some(Self::FailureClass),
_ => None,
}
}
fn as_str(self) -> &'static str {
match self {
Self::Model => "model",
Self::Account => "account",
Self::ApiKey => "apiKey",
Self::Provider => "provider",
Self::Transport => "transport",
Self::FailureClass => "failureClass",
}
}
fn expression(self) -> &'static str {
match self {
Self::Model => "terminal_requests.model_name",
Self::Account => {
"coalesce(nullif(concat_ws(' → ', nullif(diagnostic_accounts.email, ''), nullif(terminal_requests.account_id, '')), ''), '未知账号')"
}
Self::ApiKey => "coalesce(nullif(diagnostic_api_keys.name, ''), '未知密钥')",
Self::Provider => "coalesce(nullif(terminal_requests.provider, ''), '未知 Provider')",
Self::Transport => "coalesce(nullif(terminal_requests.transport, ''), '未知传输')",
Self::FailureClass => "terminal_requests.failure_class",
}
}
fn source(self) -> &'static str {
match self {
Self::Account => {
"terminal_requests left join accounts diagnostic_accounts on diagnostic_accounts.id = terminal_requests.account_id"
}
Self::ApiKey => {
"terminal_requests left join client_api_keys diagnostic_api_keys on diagnostic_api_keys.id = terminal_requests.client_api_key_id"
}
Self::Model | Self::Provider | Self::Transport | Self::FailureClass => {
"terminal_requests"
}
}
}
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageInsightsOverview {
granularity: String,
health: UsageHealthInsights,
performance: UsagePerformanceInsights,
cost: UsageCostInsights,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsageHealthInsights {
total_requests: u64,
success_requests: u64,
failed_requests: u64,
success_rate: f64,
request_change_rate: Option<f64>,
success_rate_change: Option<f64>,
points: Vec<UsageHealthPoint>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsageHealthPoint {
bucket: DateTime<Utc>,
label: String,
success_requests: u64,
failed_requests: u64,
error_rate: f64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsagePerformanceInsights {
latency_p50_ms: Option<f64>,
latency_p95_ms: Option<f64>,
latency_p99_ms: Option<f64>,
ttft_p50_ms: Option<f64>,
ttft_p95_ms: Option<f64>,
ttft_p99_ms: Option<f64>,
latency_coverage: f64,
ttft_coverage: f64,
points: Vec<UsagePerformancePoint>,
}
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsagePerformancePoint {
bucket: DateTime<Utc>,
label: String,
latency_p50_ms: Option<f64>,
latency_p95_ms: Option<f64>,
latency_p99_ms: Option<f64>,
ttft_p50_ms: Option<f64>,
ttft_p95_ms: Option<f64>,
ttft_p99_ms: Option<f64>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsageCostInsights {
estimated_cost: f64,
standard_cost: f64,
cost_per_request: f64,
tokens_per_request: f64,
cached_token_rate: f64,
cache_hit_request_rate: f64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
total_tokens: u64,
points: Vec<UsageCostPoint>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsageCostPoint {
bucket: DateTime<Utc>,
label: String,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
total_tokens: u64,
estimated_cost: f64,
standard_cost: f64,
cached_token_rate: f64,
cache_hit_request_rate: f64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageDiagnosticsInsights {
dimension: String,
items: Vec<UsageDiagnosticsItem>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UsageDiagnosticsItem {
name: String,
request_count: u64,
success_count: u64,
error_count: u64,
error_rate: f64,
request_share: f64,
latency_p95_ms: Option<f64>,
estimated_cost: f64,
}
pub async fn overview(
pool: &PgPool,
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> PgUsageRecordStoreResult<UsageInsightsOverview> {
let granularity = InsightsGranularity::for_range(start, end);
let (health, performance, cost) = tokio::try_join!(
load_health(pool, start, end, granularity),
load_performance(pool, start, end, granularity),
load_cost(pool, start, end, granularity),
)?;
Ok(UsageInsightsOverview {
granularity: granularity.as_str().to_string(),
health,
performance,
cost,
})
}
pub async fn diagnostics(
pool: &PgPool,
start: DateTime<Utc>,
end: DateTime<Utc>,
dimension: UsageDiagnosticsDimension,
) -> PgUsageRecordStoreResult<UsageDiagnosticsInsights> {
let expression = dimension.expression();
let source = dimension.source();
let failure_filter = if dimension == UsageDiagnosticsDimension::FailureClass {
"where not terminal_requests.is_success"
} else {
""
};
let sql = format!(
"with {TERMINAL_REQUESTS_CTE}\n\
select {expression} as name,\n\
count(*)::bigint as request_count,\n\
count(*) filter (where is_success)::bigint as success_count,\n\
count(*) filter (where not is_success)::bigint as error_count,\n\
percentile_cont(0.95) within group (order by latency_ms)\n\
filter (where is_success and latency_ms is not null) as latency_p95_ms,\n\
(select count(*) from terminal_requests)::bigint as total_requests\n\
from {source} {failure_filter}\n\
group by name\n\
order by request_count desc, name asc\n\
limit 8"
);
// SAFETY: SQL fragments come only from closed enums and module constants; values remain bound.
let rows = sqlx::query(AssertSqlSafe(sql))
.bind(start)
.bind(end)
.fetch_all(pool)
.await?;
let costs = if dimension == UsageDiagnosticsDimension::FailureClass {
HashMap::new()
} else {
load_diagnostic_costs(pool, start, end, dimension).await?
};
let items = rows
.into_iter()
.map(|row| {
let name: String = row.get("name");
let request_count = nonnegative(row.get("request_count"));
let success_count = nonnegative(row.get("success_count"));
let error_count = nonnegative(row.get("error_count"));
let total_requests = nonnegative(row.get("total_requests"));
UsageDiagnosticsItem {
estimated_cost: costs.get(&name).copied().unwrap_or_default(),
name,
request_count,
success_count,
error_count,
error_rate: rate(error_count, request_count),
request_share: rate(request_count, total_requests),
latency_p95_ms: row.get("latency_p95_ms"),
}
})
.collect();
Ok(UsageDiagnosticsInsights {
dimension: dimension.as_str().to_string(),
items,
})
}
async fn load_health(
pool: &PgPool,
start: DateTime<Utc>,
end: DateTime<Utc>,
granularity: InsightsGranularity,
) -> PgUsageRecordStoreResult<UsageHealthInsights> {
let range = end.signed_duration_since(start);
let previous_start = start.checked_sub_signed(range).unwrap_or(start);
let bucket = granularity.bucket_expression("terminal_at");
let sql = format!(
"with {TERMINAL_REQUESTS_CTE}\n\
select case when terminal_at >= $3 then 'current' else 'previous' end as period,\n\
{bucket} as bucket_start,\n\
count(*) filter (where is_success)::bigint as success_requests,\n\
count(*) filter (where not is_success)::bigint as failed_requests\n\
from terminal_requests\n\
group by period, bucket_start\n\
order by bucket_start"
);
// SAFETY: granularity selects one of three module-owned date_bin expressions.
let rows = sqlx::query(AssertSqlSafe(sql))
.bind(previous_start)
.bind(end)
.bind(start)
.fetch_all(pool)
.await?;
let mut current = HashMap::new();
let mut current_success = 0_u64;
let mut current_failed = 0_u64;
let mut previous_success = 0_u64;
let mut previous_failed = 0_u64;
for row in rows {
let success = nonnegative(row.get("success_requests"));
let failed = nonnegative(row.get("failed_requests"));
if row.get::<String, _>("period") == "current" {
current_success = current_success.saturating_add(success);
current_failed = current_failed.saturating_add(failed);
current.insert(
row.get::<DateTime<Utc>, _>("bucket_start"),
(success, failed),
);
} else {
previous_success = previous_success.saturating_add(success);
previous_failed = previous_failed.saturating_add(failed);
}
}
let total = current_success.saturating_add(current_failed);
let previous_total = previous_success.saturating_add(previous_failed);
let success_rate = rate(current_success, total);
let previous_success_rate = rate(previous_success, previous_total);
let points = bucket_starts(start, end, granularity)
.into_iter()
.map(|bucket| {
let (success_requests, failed_requests) =
current.get(&bucket).copied().unwrap_or_default();
let point_total = success_requests.saturating_add(failed_requests);
UsageHealthPoint {
bucket,
label: bucket_label(bucket, granularity),
success_requests,
failed_requests,
error_rate: rate(failed_requests, point_total),
}
})
.collect();
Ok(UsageHealthInsights {
total_requests: total,
success_requests: current_success,
failed_requests: current_failed,
success_rate,
request_change_rate: relative_change(total, previous_total),
success_rate_change: (previous_total > 0).then_some(success_rate - previous_success_rate),
points,
})
}
async fn load_performance(
pool: &PgPool,
start: DateTime<Utc>,
end: DateTime<Utc>,
granularity: InsightsGranularity,
) -> PgUsageRecordStoreResult<UsagePerformanceInsights> {
let summary = sqlx::query(
r#"
select
count(*)::bigint as request_count,
count(latency_ms)::bigint as latency_count,
count(first_token_ms)::bigint as ttft_count,
percentile_cont(0.50) within group (order by latency_ms) as latency_p50_ms,
percentile_cont(0.95) within group (order by latency_ms) as latency_p95_ms,
percentile_cont(0.99) within group (order by latency_ms) as latency_p99_ms,
percentile_cont(0.50) within group (order by first_token_ms) as ttft_p50_ms,
percentile_cont(0.95) within group (order by first_token_ms) as ttft_p95_ms,
percentile_cont(0.99) within group (order by first_token_ms) as ttft_p99_ms
from usage_records
where created_at >= $1 and created_at < $2
"#,
)
.bind(start)
.bind(end)
.fetch_one(pool)
.await?;
let bucket = granularity.bucket_expression("created_at");
let point_sql = format!(
"select {bucket} as bucket_start,\n\
percentile_cont(0.50) within group (order by latency_ms) as latency_p50_ms,\n\
percentile_cont(0.95) within group (order by latency_ms) as latency_p95_ms,\n\
percentile_cont(0.99) within group (order by latency_ms) as latency_p99_ms,\n\
percentile_cont(0.50) within group (order by first_token_ms) as ttft_p50_ms,\n\
percentile_cont(0.95) within group (order by first_token_ms) as ttft_p95_ms,\n\
percentile_cont(0.99) within group (order by first_token_ms) as ttft_p99_ms\n\
from usage_records\n\
where created_at >= $1 and created_at < $2\n\
group by bucket_start order by bucket_start"
);
// SAFETY: granularity selects one of three module-owned date_bin expressions.
let rows = sqlx::query(AssertSqlSafe(point_sql))
.bind(start)
.bind(end)
.fetch_all(pool)
.await?;
let mut by_bucket = HashMap::new();
for row in rows {
let bucket: DateTime<Utc> = row.get("bucket_start");
by_bucket.insert(
bucket,
UsagePerformancePoint {
bucket,
label: bucket_label(bucket, granularity),
latency_p50_ms: row.get("latency_p50_ms"),
latency_p95_ms: row.get("latency_p95_ms"),
latency_p99_ms: row.get("latency_p99_ms"),
ttft_p50_ms: row.get("ttft_p50_ms"),
ttft_p95_ms: row.get("ttft_p95_ms"),
ttft_p99_ms: row.get("ttft_p99_ms"),
},
);
}
let points = bucket_starts(start, end, granularity)
.into_iter()
.map(|bucket| {
by_bucket
.remove(&bucket)
.unwrap_or_else(|| UsagePerformancePoint {
bucket,
label: bucket_label(bucket, granularity),
..Default::default()
})
})
.collect();
let request_count = nonnegative(summary.get("request_count"));
Ok(UsagePerformanceInsights {
latency_p50_ms: summary.get("latency_p50_ms"),
latency_p95_ms: summary.get("latency_p95_ms"),
latency_p99_ms: summary.get("latency_p99_ms"),
ttft_p50_ms: summary.get("ttft_p50_ms"),
ttft_p95_ms: summary.get("ttft_p95_ms"),
ttft_p99_ms: summary.get("ttft_p99_ms"),
latency_coverage: rate(nonnegative(summary.get("latency_count")), request_count),
ttft_coverage: rate(nonnegative(summary.get("ttft_count")), request_count),
points,
})
}
async fn load_cost(
pool: &PgPool,
start: DateTime<Utc>,
end: DateTime<Utc>,
granularity: InsightsGranularity,
) -> PgUsageRecordStoreResult<UsageCostInsights> {
let bucket = granularity.bucket_expression("created_at");
let sql = format!(
"select {bucket} as bucket_start,\n\
coalesce(nullif(upstream_model, ''), model) as billing_model,\n\
service_tier, coalesce(input_tokens, 0) > $3 as is_long,\n\
count(*)::bigint as request_count,\n\
count(*) filter (where coalesce(input_tokens, 0) > 0)::bigint as cache_eligible_count,\n\
count(*) filter (\n\
where coalesce(input_tokens, 0) > 0 and coalesce(cached_tokens, 0) > 0\n\
)::bigint as cache_hit_count,\n\
coalesce(sum(input_tokens), 0)::bigint as input_tokens,\n\
coalesce(sum(output_tokens), 0)::bigint as output_tokens,\n\
coalesce(sum(least(coalesce(cached_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cached_tokens\n\
from usage_records\n\
where created_at >= $1 and created_at < $2\n\
group by bucket_start, billing_model, service_tier, is_long\n\
order by bucket_start"
);
// SAFETY: granularity selects one of three module-owned date_bin expressions.
let rows = sqlx::query(AssertSqlSafe(sql))
.bind(start)
.bind(end)
.bind(i64::try_from(LONG_CONTEXT_THRESHOLD).unwrap_or(i64::MAX))
.fetch_all(pool)
.await?;
let mut total = CostAccumulator::default();
let mut by_bucket = BTreeMap::<DateTime<Utc>, CostAccumulator>::new();
for row in rows {
let sample = cost_sample(&row);
total.push(sample);
by_bucket
.entry(row.get("bucket_start"))
.or_default()
.push(sample);
}
let points = bucket_starts(start, end, granularity)
.into_iter()
.map(|bucket| {
let item = by_bucket.remove(&bucket).unwrap_or_default();
UsageCostPoint {
bucket,
label: bucket_label(bucket, granularity),
input_tokens: item.input_tokens,
output_tokens: item.output_tokens,
cached_tokens: item.cached_tokens,
total_tokens: item.total_tokens(),
estimated_cost: item.estimated_cost,
standard_cost: item.standard_cost,
cached_token_rate: rate(item.cached_tokens, item.input_tokens),
cache_hit_request_rate: rate(item.cache_hit_count, item.cache_eligible_count),
}
})
.collect();
Ok(UsageCostInsights {
estimated_cost: total.estimated_cost,
standard_cost: total.standard_cost,
cost_per_request: amount_per(total.estimated_cost, total.request_count),
tokens_per_request: amount_per(total.total_tokens() as f64, total.request_count),
cached_token_rate: rate(total.cached_tokens, total.input_tokens),
cache_hit_request_rate: rate(total.cache_hit_count, total.cache_eligible_count),
input_tokens: total.input_tokens,
output_tokens: total.output_tokens,
cached_tokens: total.cached_tokens,
total_tokens: total.total_tokens(),
points,
})
}
async fn load_diagnostic_costs(
pool: &PgPool,
start: DateTime<Utc>,
end: DateTime<Utc>,
dimension: UsageDiagnosticsDimension,
) -> PgUsageRecordStoreResult<HashMap<String, f64>> {
let expression = dimension.expression();
let source = dimension.source();
let sql = format!(
"with {TERMINAL_REQUESTS_CTE}\n\
select {expression} as name, billing_model, service_tier,\n\
coalesce(input_tokens, 0) > $3 as is_long,\n\
count(*)::bigint as request_count,\n\
count(*) filter (where coalesce(input_tokens, 0) > 0)::bigint as cache_eligible_count,\n\
count(*) filter (\n\
where coalesce(input_tokens, 0) > 0 and coalesce(cached_tokens, 0) > 0\n\
)::bigint as cache_hit_count,\n\
coalesce(sum(input_tokens), 0)::bigint as input_tokens,\n\
coalesce(sum(output_tokens), 0)::bigint as output_tokens,\n\
coalesce(sum(least(coalesce(cached_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cached_tokens\n\
from {source}\n\
where terminal_requests.is_success\n\
group by name, billing_model, service_tier, is_long"
);
// SAFETY: the grouping expression comes from the closed diagnostics dimension enum.
let rows = sqlx::query(AssertSqlSafe(sql))
.bind(start)
.bind(end)
.bind(i64::try_from(LONG_CONTEXT_THRESHOLD).unwrap_or(i64::MAX))
.fetch_all(pool)
.await?;
let mut costs = HashMap::new();
for row in rows {
let name: String = row.get("name");
let sample = cost_sample(&row);
*costs.entry(name).or_insert(0.0) += sample.estimated_cost;
}
Ok(costs)
}
#[derive(Debug, Clone, Copy, Default)]
struct CostSample {
request_count: u64,
cache_eligible_count: u64,
cache_hit_count: u64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
estimated_cost: f64,
standard_cost: f64,
}
#[derive(Debug, Default)]
struct CostAccumulator {
request_count: u64,
cache_eligible_count: u64,
cache_hit_count: u64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
estimated_cost: f64,
standard_cost: f64,
}
impl CostAccumulator {
fn push(&mut self, sample: CostSample) {
self.request_count = self.request_count.saturating_add(sample.request_count);
self.cache_eligible_count = self
.cache_eligible_count
.saturating_add(sample.cache_eligible_count);
self.cache_hit_count = self.cache_hit_count.saturating_add(sample.cache_hit_count);
self.input_tokens = self.input_tokens.saturating_add(sample.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(sample.output_tokens);
self.cached_tokens = self.cached_tokens.saturating_add(sample.cached_tokens);
self.estimated_cost += sample.estimated_cost;
self.standard_cost += sample.standard_cost;
}
fn total_tokens(&self) -> u64 {
self.input_tokens.saturating_add(self.output_tokens)
}
}
fn cost_sample(row: &sqlx::postgres::PgRow) -> CostSample {
let input_tokens = nonnegative(row.get("input_tokens"));
let output_tokens = nonnegative(row.get("output_tokens"));
let cached_tokens = nonnegative(row.get("cached_tokens")).min(input_tokens);
let billing_model: String = row.get("billing_model");
let service_tier: Option<String> = row.get("service_tier");
let billing = calculate_aggregate_billing(
input_tokens,
output_tokens,
cached_tokens,
&billing_model,
service_tier.as_deref(),
row.get("is_long"),
);
CostSample {
request_count: nonnegative(row.get("request_count")),
cache_eligible_count: nonnegative(row.get("cache_eligible_count")),
cache_hit_count: nonnegative(row.get("cache_hit_count")),
input_tokens,
output_tokens,
cached_tokens,
estimated_cost: billing.total_amount,
standard_cost: billing.standard_amount,
}
}
fn bucket_starts(
start: DateTime<Utc>,
end: DateTime<Utc>,
granularity: InsightsGranularity,
) -> Vec<DateTime<Utc>> {
let step = granularity.step_seconds();
let offset = if granularity == InsightsGranularity::Day {
8 * 60 * 60
} else {
0
};
let aligned = (start.timestamp() + offset).div_euclid(step) * step - offset;
let Some(mut cursor) = DateTime::<Utc>::from_timestamp(aligned, 0) else {
return Vec::new();
};
let mut points = Vec::new();
while cursor < end {
points.push(cursor);
let Some(next) = cursor.checked_add_signed(Duration::seconds(step)) else {
break;
};
cursor = next;
}
points
}
fn bucket_label(bucket: DateTime<Utc>, granularity: InsightsGranularity) -> String {
let china = bucket + Duration::hours(8);
match granularity {
InsightsGranularity::Day => china.format("%Y-%m-%d").to_string(),
InsightsGranularity::QuarterHour | InsightsGranularity::Hour => {
china.format("%m-%d %H:%M").to_string()
}
}
}
fn nonnegative(value: Option<i64>) -> u64 {
value.unwrap_or_default().max(0) as u64
}
fn rate(numerator: u64, denominator: u64) -> f64 {
if denominator == 0 {
0.0
} else {
numerator as f64 / denominator as f64
}
}
fn amount_per(amount: f64, count: u64) -> f64 {
if count == 0 {
0.0
} else {
amount / count as f64
}
}
fn relative_change(current: u64, previous: u64) -> Option<f64> {
(previous > 0).then(|| (current as f64 - previous as f64) / previous as f64)
}
/// 返回缺省观测时间范围。
pub fn default_time_range(now: DateTime<Utc>) -> (DateTime<Utc>, DateTime<Utc>) {
(now - Duration::days(DEFAULT_RANGE_DAYS), now)
}
+1
View File
@@ -1,5 +1,6 @@
//! 成功请求事实。
pub mod insights;
pub mod query;
pub mod store;
pub mod types;
+30 -3
View File
@@ -1,4 +1,4 @@
//! 管理端使用记录服务。
//! 管理端使用统计与请求记录查询服务。
use std::collections::{BTreeMap, HashMap};
@@ -10,6 +10,10 @@ use crate::{
infra::{format::optional_nonnegative_i64_to_u64, json::NumberedPage},
telemetry::{
billing,
usage::insights::{
UsageDiagnosticsDimension, UsageDiagnosticsInsights, UsageInsightsOverview,
diagnostics, overview,
},
usage::store::{
PgUsageRecordStore, PgUsageRecordStoreResult, UsageRecordFilter, push_filter,
},
@@ -114,14 +118,14 @@ pub enum UsageQueryError {
Accounts,
}
/// 管理端使用记录服务。
/// 管理端使用统计与请求记录查询服务。
#[derive(Clone)]
pub struct UsageQueryService {
store: PgUsageRecordStore,
}
impl UsageQueryService {
/// 构造管理端使用记录服务。
/// 构造管理端使用统计与请求记录查询服务。
pub fn new(store: PgUsageRecordStore) -> Self {
Self { store }
}
@@ -224,6 +228,29 @@ impl UsageQueryService {
.map_err(|_| UsageQueryError::List)
}
/// 读取请求健康、性能与成本观测概览。
pub async fn insights_overview(
&self,
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Result<UsageInsightsOverview, UsageQueryError> {
overview(self.store.pool(), start, end)
.await
.map_err(|_| UsageQueryError::List)
}
/// 按维度读取请求热点诊断。
pub async fn insights_diagnostics(
&self,
start: DateTime<Utc>,
end: DateTime<Utc>,
dimension: UsageDiagnosticsDimension,
) -> Result<UsageDiagnosticsInsights, UsageQueryError> {
diagnostics(self.store.pool(), start, end, dimension)
.await
.map_err(|_| UsageQueryError::List)
}
/// 清空日志。
pub async fn clear(&self) -> Result<ClearUsageRecords, UsageQueryError> {
self.store

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