mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
05ec1a4d38 | ||
|
|
9808e5939c | ||
|
|
a58911c127 | ||
|
|
0b9e080f1b | ||
|
|
f218e31c72 | ||
|
|
83f2f55e4a | ||
|
|
0384e39289 | ||
|
|
9036288d69 | ||
|
|
77b18eec68 | ||
|
|
3f7d473bea | ||
|
|
99cf39c76d | ||
|
|
f4d0576e6b | ||
|
|
2d9510d9e3 | ||
|
|
36d6e26291 | ||
|
|
eff522f2b6 | ||
|
|
65a175d48c | ||
|
|
f1872df293 |
@@ -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"
|
||||
|
||||
+1
-7
@@ -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/
|
||||
|
||||
@@ -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&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&logo=vuedotjs&logoColor=white" /></a>
|
||||
<a href="https://vite.dev/"><img alt="Vite 8" src="https://img.shields.io/badge/Vite-8-646cff?style=flat-square&logo=vite&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&logo=postgresql&logoColor=white" /></a>
|
||||
<a href="https://redis.io/"><img alt="Redis 8" src="https://img.shields.io/badge/Redis-8-ff4438?style=flat-square&logo=redis&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 多账号网关代理**
|
||||
|
||||
[](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://www.rust-lang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://vite.dev/)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](#开源许可)
|
||||
|
||||
[快速开始](#快速开始) · [首次使用](#首次使用) · [客户端接入](#客户端接入) · [常见问题](#常见问题) · [部署文档](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) 开源。
|
||||
|
||||
Generated
+1
-49
@@ -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
@@ -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"] }
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -102,6 +102,7 @@ struct DashboardUsageRecordData {
|
||||
client_ip: Option<String>,
|
||||
user_agent: Option<String>,
|
||||
reasoning_effort: Option<String>,
|
||||
reasoning_preset: Option<String>,
|
||||
compact: bool,
|
||||
request_kind: Option<String>,
|
||||
subagent_kind: Option<String>,
|
||||
@@ -418,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")
|
||||
}
|
||||
|
||||
@@ -517,6 +518,7 @@ fn dashboard_usage_record_data(
|
||||
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
|
||||
let reasoning_effort =
|
||||
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
|
||||
let reasoning_preset = metadata_string(&record.metadata, &["reasoningPreset"]);
|
||||
let compact = record
|
||||
.metadata
|
||||
.get("compact")
|
||||
@@ -541,6 +543,7 @@ fn dashboard_usage_record_data(
|
||||
client_ip,
|
||||
user_agent,
|
||||
reasoning_effort,
|
||||
reasoning_preset,
|
||||
compact,
|
||||
request_kind,
|
||||
subagent_kind,
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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 {
|
||||
@@ -88,6 +104,7 @@ pub(crate) struct UsageRecordData {
|
||||
client_ip: Option<String>,
|
||||
user_agent: Option<String>,
|
||||
reasoning_effort: Option<String>,
|
||||
reasoning_preset: Option<String>,
|
||||
compact: bool,
|
||||
request_kind: Option<String>,
|
||||
subagent_kind: Option<String>,
|
||||
@@ -303,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,
|
||||
@@ -404,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 => {
|
||||
@@ -442,6 +511,7 @@ fn usage_record_data(
|
||||
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
|
||||
let reasoning_effort =
|
||||
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
|
||||
let reasoning_preset = metadata_string(&record.metadata, &["reasoningPreset"]);
|
||||
let compact = record
|
||||
.metadata
|
||||
.get("compact")
|
||||
@@ -463,6 +533,7 @@ fn usage_record_data(
|
||||
client_ip,
|
||||
user_agent,
|
||||
reasoning_effort,
|
||||
reasoning_preset,
|
||||
compact,
|
||||
request_kind,
|
||||
subagent_kind,
|
||||
@@ -613,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
|
||||
};
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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(¤t_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),
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
"默认管理员账号已完成初始化检查"
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 定期版本检查失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,8 +68,8 @@ impl QuotaRefreshTask {
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let config = PeriodicTaskConfig::new(
|
||||
self.interval_secs,
|
||||
"quota 刷新任务已启动",
|
||||
"quota 刷新任务已关闭",
|
||||
"Quota 刷新任务已启动",
|
||||
"Quota 刷新任务已关闭",
|
||||
);
|
||||
spawn_periodic_task(self, config)
|
||||
}
|
||||
@@ -96,7 +96,7 @@ impl PeriodicTaskRunner for QuotaRefreshTask {
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = %error, "quota 刷新任务失败");
|
||||
warn!(error = %error, "Quota 刷新任务失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
@@ -107,6 +107,6 @@ fn log_refreshed_summary(summary: QuotaRefreshSummary) {
|
||||
info!(
|
||||
refreshed = summary.refreshed,
|
||||
failed = summary.failed,
|
||||
"quota 刷新任务完成"
|
||||
"Quota 刷新任务完成"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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 刷新定时器调度失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,334 @@
|
||||
//! 账号级失败分类、隔离副作用与换号决策 owner。
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
controllers::{
|
||||
ControllerFailureFact,
|
||||
account_state::{AccountStateEffect, AccountStateEffects},
|
||||
},
|
||||
errors::ClientFailure,
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
|
||||
lifecycle::contract::{
|
||||
AttemptDecision, AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
|
||||
},
|
||||
transport::observation::{UpstreamFailureFacts, UpstreamFailureKind},
|
||||
},
|
||||
fleet::{account::AccountStatus, pool::AccountPoolService},
|
||||
upstream::openai::{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 {
|
||||
AccountStateEffects::apply(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();
|
||||
let error_type = failure
|
||||
.upstream_type
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
let message = failure.message.to_ascii_lowercase();
|
||||
if is_model_unsupported(&code)
|
||||
|| is_model_unsupported(&message)
|
||||
|| 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 is_auth_failure(&code, &error_type, &message) {
|
||||
return Some(ClientFailure::new(failure.clone(), 401, false));
|
||||
}
|
||||
if code.contains("forbidden") || code.contains("banned") || error_type == "permission_error"
|
||||
{
|
||||
return Some(ClientFailure::new(failure.clone(), 403, false));
|
||||
}
|
||||
if code.contains("server_overloaded") {
|
||||
return Some(ClientFailure::new(failure.clone(), 503, true));
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
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> {
|
||||
let status = facts.status_code;
|
||||
let body = facts.body.as_str();
|
||||
let code = facts.code.as_deref().unwrap_or_default();
|
||||
let error_type = facts.error_type.as_deref().unwrap_or_default();
|
||||
let message = facts.message.as_str();
|
||||
|
||||
if is_model_unsupported(code) || is_model_unsupported(message) || is_model_unsupported(body) {
|
||||
return Some(model_unsupported_failure(account_id, body));
|
||||
}
|
||||
if let Some(account_status) = explicit_account_status(code, error_type, body) {
|
||||
return Some(account_status_failure(
|
||||
account_id,
|
||||
account_status,
|
||||
body,
|
||||
status,
|
||||
));
|
||||
}
|
||||
if status == Some(403) && !is_html(body) {
|
||||
return Some(account_status_failure(
|
||||
account_id,
|
||||
AccountStatus::Banned,
|
||||
body,
|
||||
status,
|
||||
));
|
||||
}
|
||||
if status == Some(401) {
|
||||
return Some(account_status_failure(
|
||||
account_id,
|
||||
AccountStatus::Expired,
|
||||
body,
|
||||
status,
|
||||
));
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<ClassifiedFailure> {
|
||||
let code = failure.upstream_code.as_deref().unwrap_or_default();
|
||||
let error_type = failure.upstream_type.as_deref().unwrap_or_default();
|
||||
if is_model_unsupported(code) || is_model_unsupported(&failure.message) {
|
||||
return Some(model_unsupported_failure(
|
||||
account_id,
|
||||
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
|
||||
));
|
||||
}
|
||||
let account_status =
|
||||
explicit_account_status(code, error_type, &failure.message).or_else(|| {
|
||||
(code.eq_ignore_ascii_case("forbidden") || error_type == "permission_error")
|
||||
.then_some(AccountStatus::Banned)
|
||||
})?;
|
||||
Some(account_status_failure(
|
||||
account_id,
|
||||
account_status,
|
||||
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
|
||||
failure.explicit_status_code,
|
||||
))
|
||||
}
|
||||
|
||||
fn model_unsupported_failure(account_id: &str, message: &str) -> ClassifiedFailure {
|
||||
ClassifiedFailure {
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account_id,
|
||||
ExhaustedAccountKind::ModelUnsupported,
|
||||
message,
|
||||
),
|
||||
effect: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn account_status_failure(
|
||||
account_id: &str,
|
||||
status: AccountStatus,
|
||||
message: &str,
|
||||
status_code: Option<u16>,
|
||||
) -> ClassifiedFailure {
|
||||
let kind = match status {
|
||||
AccountStatus::Expired => ExhaustedAccountKind::Expired,
|
||||
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
|
||||
AccountStatus::Banned => ExhaustedAccountKind::Banned,
|
||||
AccountStatus::Active | AccountStatus::QuotaExhausted => {
|
||||
unreachable!("account failure controller only invalidates unavailable statuses")
|
||||
}
|
||||
};
|
||||
let mut exhaustion = AccountExhaustionRecord::new(account_id, kind, message);
|
||||
if let Some(status_code) = status_code {
|
||||
exhaustion = exhaustion.with_status_code(status_code);
|
||||
}
|
||||
ClassifiedFailure {
|
||||
exhaustion,
|
||||
effect: Some(AccountStateEffect::SetStatus(status)),
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
{
|
||||
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 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")
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
//! 明确账号事实触发的运行时失效与连接驱逐。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::{
|
||||
fleet::{account::AccountStatus, pool::AccountPoolService},
|
||||
upstream::openai::transport::CodexBackendClient,
|
||||
};
|
||||
|
||||
pub(super) enum AccountStateEffect {
|
||||
SetStatus(AccountStatus),
|
||||
MarkQuotaLimitedUntil(DateTime<Utc>),
|
||||
}
|
||||
|
||||
pub(super) struct AccountStateEffects;
|
||||
|
||||
impl AccountStateEffects {
|
||||
pub(super) async fn apply(
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
transport::observation::UpstreamFailureFacts,
|
||||
},
|
||||
fleet::{
|
||||
account::AccountStatus,
|
||||
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
};
|
||||
|
||||
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,225 @@
|
||||
//! `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},
|
||||
transport::observation::UpstreamFailureFacts,
|
||||
},
|
||||
upstream::openai::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,
|
||||
}
|
||||
@@ -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,
|
||||
},
|
||||
transport::observation::{UpstreamFailureFacts, UpstreamFailureKind},
|
||||
},
|
||||
upstream::openai::{
|
||||
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())
|
||||
}
|
||||
@@ -0,0 +1,997 @@
|
||||
//! 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,
|
||||
},
|
||||
transport::observation::UpstreamFailureFacts,
|
||||
},
|
||||
fleet::pool::AccountPoolService,
|
||||
upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
};
|
||||
|
||||
mod account_failure;
|
||||
mod account_state;
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,345 @@
|
||||
//! 请求前配额验证规则 owner。
|
||||
|
||||
use chrono::{Duration, Utc};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::AccountIdentityService,
|
||||
controllers::{
|
||||
ControllerFailureFact,
|
||||
account_state::{AccountStateEffect, AccountStateEffects},
|
||||
},
|
||||
errors::ClientFailure,
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
|
||||
lifecycle::contract::{
|
||||
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
|
||||
},
|
||||
transport::observation::UpstreamFailureFacts,
|
||||
},
|
||||
fleet::{
|
||||
pool::{AccountLease, AccountPoolService},
|
||||
quota::{
|
||||
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
|
||||
quota_snapshot_reset_at,
|
||||
},
|
||||
},
|
||||
upstream::openai::{
|
||||
protocol::responses::ResponsesSseFailure,
|
||||
transport::{CodexBackendClient, CodexRequestContext},
|
||||
},
|
||||
};
|
||||
|
||||
const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached";
|
||||
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
|
||||
|
||||
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, "a)
|
||||
.await;
|
||||
if quota_snapshot_limit_reached("a) {
|
||||
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, "a)))
|
||||
}
|
||||
|
||||
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 {
|
||||
AccountStateEffects::apply(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> {
|
||||
let signaled_kind = facts
|
||||
.code
|
||||
.as_deref()
|
||||
.and_then(classify_quota_signal)
|
||||
.or_else(|| facts.error_type.as_deref().and_then(classify_quota_signal))
|
||||
.or_else(|| classify_quota_message(&facts.message));
|
||||
match signaled_kind.or(match facts.status_code {
|
||||
Some(429) => Some(StreamQuotaKind::RateLimited),
|
||||
Some(402) => Some(StreamQuotaKind::QuotaExhausted),
|
||||
_ => None,
|
||||
})? {
|
||||
StreamQuotaKind::RateLimited => Some(rate_limited_failure(
|
||||
account_id,
|
||||
facts.body.clone(),
|
||||
facts.retry_after_seconds,
|
||||
)),
|
||||
StreamQuotaKind::QuotaExhausted => {
|
||||
Some(quota_exhausted_failure(account_id, facts.body.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_stream(
|
||||
account_id: &str,
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
let kind = classify_stream_signal(failure).or(match failure.explicit_status_code {
|
||||
Some(429) => Some(StreamQuotaKind::RateLimited),
|
||||
Some(402) => Some(StreamQuotaKind::QuotaExhausted),
|
||||
_ => None,
|
||||
})?;
|
||||
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
|
||||
Some(match kind {
|
||||
StreamQuotaKind::RateLimited => {
|
||||
rate_limited_failure(account_id, body, failure.retry_after_seconds)
|
||||
}
|
||||
StreamQuotaKind::QuotaExhausted => quota_exhausted_failure(account_id, body),
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum StreamQuotaKind {
|
||||
RateLimited,
|
||||
QuotaExhausted,
|
||||
}
|
||||
|
||||
fn classify_stream_signal(failure: &ResponsesSseFailure) -> Option<StreamQuotaKind> {
|
||||
failure
|
||||
.upstream_code
|
||||
.as_deref()
|
||||
.and_then(classify_quota_signal)
|
||||
.or_else(|| {
|
||||
failure
|
||||
.upstream_type
|
||||
.as_deref()
|
||||
.and_then(classify_quota_signal)
|
||||
})
|
||||
.or_else(|| classify_quota_message(&failure.message))
|
||||
}
|
||||
|
||||
fn classify_quota_signal(signal: &str) -> Option<StreamQuotaKind> {
|
||||
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(StreamQuotaKind::RateLimited),
|
||||
"quota_exhausted"
|
||||
| "quota_exceeded"
|
||||
| "payment_required"
|
||||
| "insufficient_quota"
|
||||
| "workspace_owner_credits_depleted"
|
||||
| "workspace_member_credits_depleted" => Some(StreamQuotaKind::QuotaExhausted),
|
||||
signal if signal.starts_with("billing_limit") => Some(StreamQuotaKind::QuotaExhausted),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_quota_message(message: &str) -> Option<StreamQuotaKind> {
|
||||
let message = message.to_ascii_lowercase();
|
||||
if message.contains("rate limit") || message.contains("usage limit") {
|
||||
return Some(StreamQuotaKind::RateLimited);
|
||||
}
|
||||
if message.contains("quota")
|
||||
|| message.contains("payment required")
|
||||
|| message.contains("billing limit")
|
||||
{
|
||||
return Some(StreamQuotaKind::QuotaExhausted);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn rate_limited_failure(
|
||||
account_id: &str,
|
||||
message: String,
|
||||
retry_after_seconds: Option<u64>,
|
||||
) -> ClassifiedQuotaFailure {
|
||||
let seconds = retry_after_seconds
|
||||
.unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS)
|
||||
.min(i64::MAX as u64) as i64;
|
||||
ClassifiedQuotaFailure {
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account_id,
|
||||
ExhaustedAccountKind::RateLimited,
|
||||
message,
|
||||
),
|
||||
effect: Some(AccountStateEffect::MarkQuotaLimitedUntil(
|
||||
Utc::now() + Duration::seconds(seconds),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_exhausted_failure(account_id: &str, message: String) -> ClassifiedQuotaFailure {
|
||||
ClassifiedQuotaFailure {
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account_id,
|
||||
ExhaustedAccountKind::QuotaExhausted,
|
||||
message,
|
||||
),
|
||||
effect: Some(AccountStateEffect::SetStatus(
|
||||
crate::fleet::account::AccountStatus::QuotaExhausted,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
+76
-71
@@ -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())
|
||||
@@ -424,6 +454,7 @@ fn response_request_summary(
|
||||
"compact": semantics.compact,
|
||||
"requestKind": semantics.request_kind,
|
||||
"subagentKind": semantics.subagent_kind,
|
||||
"reasoningPreset": semantics.reasoning_preset,
|
||||
"transport": backend_transport_name(transport),
|
||||
"inputType": json_value_kind(input),
|
||||
"inputItemsCount": input.and_then(Value::as_array).map(Vec::len),
|
||||
@@ -475,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,
|
||||
@@ -487,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(
|
||||
@@ -507,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,
|
||||
@@ -515,31 +546,31 @@ pub(super) async fn record_live_response_stream_event(
|
||||
status_code,
|
||||
body,
|
||||
);
|
||||
log_live_response_stream_finalized(context, status_code, level, message, &metadata, body);
|
||||
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,
|
||||
@@ -550,18 +581,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,23 +602,17 @@ fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
|
||||
}
|
||||
|
||||
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");
|
||||
@@ -608,7 +633,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),
|
||||
@@ -618,7 +642,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"
|
||||
);
|
||||
};
|
||||
}
|
||||
@@ -643,25 +667,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;
|
||||
}
|
||||
@@ -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
@@ -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,
|
||||
},
|
||||
|
||||
+40
-97
@@ -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>,
|
||||
@@ -0,0 +1,4 @@
|
||||
//! 已分类账号失败的聚合、状态副作用和 SSE 终态事实。
|
||||
|
||||
pub(crate) mod exhaustion;
|
||||
pub(in crate::dispatch) mod sse;
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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},
|
||||
observation::upstream_failure_facts,
|
||||
prefetch::{PrefetchedStream, StreamPrefetchError, prefetch_until_commit},
|
||||
},
|
||||
},
|
||||
fleet::{
|
||||
account::Account,
|
||||
pool::{AccountCandidateLease, AccountPoolService},
|
||||
},
|
||||
upstream::openai::{
|
||||
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",
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
//! Attempt lifecycle 的类型化 observation、decision 与结果。
|
||||
|
||||
use bytes::Bytes;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
errors::{ClientFailure, ResponseDispatchError},
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccount},
|
||||
transport::{
|
||||
canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
|
||||
observation::UpstreamFailureFacts,
|
||||
},
|
||||
},
|
||||
fleet::{account::Account, pool::AccountLease},
|
||||
upstream::openai::{
|
||||
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>),
|
||||
}
|
||||
@@ -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}")
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
//! 请求级通用账号候选路由。
|
||||
|
||||
pub(in crate::dispatch) mod candidates;
|
||||
+209
-489
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
//! Attempt 阶段的账号作用域上游调用与事实规范化。
|
||||
|
||||
pub(in crate::dispatch) mod account;
|
||||
pub(crate) mod canonical;
|
||||
pub(in crate::dispatch) mod observation;
|
||||
pub(in crate::dispatch) mod prefetch;
|
||||
@@ -0,0 +1,177 @@
|
||||
//! 上游原始错误到 attempt 通用事实的单次规范化。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::upstream::openai::transport::{
|
||||
CodexClientError,
|
||||
websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason},
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) enum UpstreamFailureKind {
|
||||
HttpStatus,
|
||||
HttpConnect,
|
||||
HttpTimeout,
|
||||
HttpTransport,
|
||||
StreamIdle,
|
||||
WebSocketUpstream,
|
||||
WebSocketTransport,
|
||||
WebSocketTimeout,
|
||||
WebSocketProtocol,
|
||||
PostSendAmbiguous,
|
||||
ContinuationUnavailable(PreviousResponseUnavailableReason),
|
||||
RequestEncoding,
|
||||
Protocol,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(in crate::dispatch) struct UpstreamFailureFacts {
|
||||
pub kind: UpstreamFailureKind,
|
||||
pub status_code: Option<u16>,
|
||||
pub code: Option<String>,
|
||||
pub error_type: Option<String>,
|
||||
pub message: String,
|
||||
pub body: String,
|
||||
pub retry_after_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn upstream_failure_facts(
|
||||
error: &CodexClientError,
|
||||
) -> UpstreamFailureFacts {
|
||||
match error {
|
||||
CodexClientError::Http(error) => UpstreamFailureFacts {
|
||||
kind: if error.is_connect() {
|
||||
UpstreamFailureKind::HttpConnect
|
||||
} else if error.is_timeout() {
|
||||
UpstreamFailureKind::HttpTimeout
|
||||
} else {
|
||||
UpstreamFailureKind::HttpTransport
|
||||
},
|
||||
status_code: error.status().map(|status| status.as_u16()),
|
||||
code: None,
|
||||
error_type: None,
|
||||
message: error.to_string(),
|
||||
body: error.to_string(),
|
||||
retry_after_seconds: None,
|
||||
},
|
||||
CodexClientError::Upstream {
|
||||
status,
|
||||
body,
|
||||
retry_after_seconds,
|
||||
transport,
|
||||
..
|
||||
} => {
|
||||
let (code, error_type, message) = error_fields(body);
|
||||
UpstreamFailureFacts {
|
||||
kind: match transport {
|
||||
crate::upstream::openai::transport::CodexBackendTransport::HttpSse => {
|
||||
UpstreamFailureKind::HttpStatus
|
||||
}
|
||||
crate::upstream::openai::transport::CodexBackendTransport::WebSocket => {
|
||||
UpstreamFailureKind::WebSocketUpstream
|
||||
}
|
||||
},
|
||||
status_code: Some(status.as_u16()),
|
||||
code,
|
||||
error_type,
|
||||
message: message.unwrap_or_else(|| body.clone()),
|
||||
body: body.clone(),
|
||||
retry_after_seconds: *retry_after_seconds,
|
||||
}
|
||||
}
|
||||
CodexClientError::StreamIdleTimeout { .. } => {
|
||||
simple_facts(UpstreamFailureKind::StreamIdle, error)
|
||||
}
|
||||
CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => {
|
||||
let (body_code, body_type, body_message) = error_fields(&upstream.body);
|
||||
UpstreamFailureFacts {
|
||||
kind: UpstreamFailureKind::WebSocketUpstream,
|
||||
status_code: Some(upstream.status_code),
|
||||
code: upstream.code.clone().or(body_code),
|
||||
error_type: upstream.error_type.clone().or(body_type),
|
||||
message: upstream
|
||||
.message
|
||||
.clone()
|
||||
.or(body_message)
|
||||
.unwrap_or_else(|| upstream.body.clone()),
|
||||
body: upstream.body.clone(),
|
||||
retry_after_seconds: upstream.retry_after_seconds,
|
||||
}
|
||||
}
|
||||
CodexClientError::WebSocket(CodexWebSocketExchangeError::ContinuationUnavailable {
|
||||
reason,
|
||||
}) => simple_facts(UpstreamFailureKind::ContinuationUnavailable(*reason), error),
|
||||
CodexClientError::WebSocket(
|
||||
CodexWebSocketExchangeError::ConnectTimeout { .. }
|
||||
| CodexWebSocketExchangeError::FastPathTimeout { .. }
|
||||
| CodexWebSocketExchangeError::SendTimeout { .. }
|
||||
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
|
||||
| CodexWebSocketExchangeError::InitialEventTimeout { .. },
|
||||
) => simple_facts(UpstreamFailureKind::WebSocketTimeout, error),
|
||||
CodexClientError::WebSocket(
|
||||
CodexWebSocketExchangeError::Transport(_)
|
||||
| CodexWebSocketExchangeError::Connect(_)
|
||||
| CodexWebSocketExchangeError::OriginCircuitOpen
|
||||
| CodexWebSocketExchangeError::OriginHalfOpenBusy
|
||||
| CodexWebSocketExchangeError::SharedConnectFailed
|
||||
| CodexWebSocketExchangeError::ClosedBeforeTerminal
|
||||
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. },
|
||||
) => simple_facts(UpstreamFailureKind::WebSocketTransport, error),
|
||||
CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous { .. }) => {
|
||||
simple_facts(UpstreamFailureKind::PostSendAmbiguous, error)
|
||||
}
|
||||
CodexClientError::WebSocket(
|
||||
CodexWebSocketExchangeError::InvalidSse(_)
|
||||
| CodexWebSocketExchangeError::InvalidCompletedResponse { .. }
|
||||
| CodexWebSocketExchangeError::UnexpectedBinaryEvent,
|
||||
) => simple_facts(UpstreamFailureKind::WebSocketProtocol, error),
|
||||
CodexClientError::WebSocket(CodexWebSocketExchangeError::InvalidRequest(_))
|
||||
| CodexClientError::WebSocketEncode(_)
|
||||
| CodexClientError::InvalidHeaderName(_)
|
||||
| CodexClientError::InvalidHeaderValue(_)
|
||||
| CodexClientError::CustomCa(_) => {
|
||||
simple_facts(UpstreamFailureKind::RequestEncoding, error)
|
||||
}
|
||||
CodexClientError::InvalidSse(_) => simple_facts(UpstreamFailureKind::Protocol, error),
|
||||
}
|
||||
}
|
||||
|
||||
fn simple_facts(kind: UpstreamFailureKind, error: &CodexClientError) -> UpstreamFailureFacts {
|
||||
let message = error.to_string();
|
||||
UpstreamFailureFacts {
|
||||
kind,
|
||||
status_code: None,
|
||||
code: None,
|
||||
error_type: None,
|
||||
body: message.clone(),
|
||||
message,
|
||||
retry_after_seconds: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn error_fields(body: &str) -> (Option<String>, Option<String>, Option<String>) {
|
||||
let Ok(value) = serde_json::from_str::<Value>(body) else {
|
||||
return (None, None, None);
|
||||
};
|
||||
let error = value
|
||||
.pointer("/response/error")
|
||||
.or_else(|| value.get("error"))
|
||||
.unwrap_or(&value);
|
||||
(
|
||||
error
|
||||
.get("code")
|
||||
.or_else(|| value.get("code"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string),
|
||||
error
|
||||
.get("type")
|
||||
.or_else(|| value.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string),
|
||||
error
|
||||
.get("message")
|
||||
.or_else(|| value.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string),
|
||||
)
|
||||
}
|
||||
@@ -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)?,
|
||||
})
|
||||
}
|
||||
@@ -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, "a)
|
||||
.await;
|
||||
if quota_snapshot_limit_reached("a) {
|
||||
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, "a)))
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -103,7 +103,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 +171,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 +213,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 +230,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 +242,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")
|
||||
|
||||
@@ -50,7 +50,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,
|
||||
};
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
@@ -176,7 +176,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,
|
||||
})
|
||||
@@ -215,7 +214,7 @@ impl QuotaRefreshService {
|
||||
warn!(
|
||||
account_id = %account.id,
|
||||
error = %error,
|
||||
"quota usage 拉取失败"
|
||||
"Quota usage 拉取失败"
|
||||
);
|
||||
summary.failed += 1;
|
||||
}
|
||||
@@ -246,7 +245,7 @@ impl QuotaRefreshService {
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)));
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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 密钥。
|
||||
|
||||
@@ -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
@@ -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"));
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -8,7 +8,7 @@ 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};
|
||||
@@ -16,6 +16,7 @@ use super::account_usage::query::{AccountUsageSummary, AccountUsageTimeBucket};
|
||||
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 {
|
||||
@@ -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)]
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -439,6 +439,7 @@ pub(crate) fn enrich_response_request_semantics(
|
||||
object.insert("compact".to_string(), Value::Bool(semantics.compact));
|
||||
insert_trimmed_string(object, "requestKind", semantics.request_kind.as_deref());
|
||||
insert_trimmed_string(object, "subagentKind", semantics.subagent_kind.as_deref());
|
||||
insert_trimmed_string(object, "reasoningPreset", semantics.reasoning_preset);
|
||||
}
|
||||
|
||||
pub(crate) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
|
||||
|
||||
@@ -0,0 +1,777 @@
|
||||
//! 管理端请求观测聚合。
|
||||
|
||||
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, ''), '未分类错误') 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 => "model_name",
|
||||
Self::Account => "coalesce(nullif(account_id, ''), '未知账号')",
|
||||
Self::ApiKey => "coalesce(nullif(client_api_key_id, ''), '未知 API Key')",
|
||||
Self::Provider => "coalesce(nullif(provider, ''), '未知 Provider')",
|
||||
Self::Transport => "coalesce(nullif(transport, ''), '未知传输')",
|
||||
Self::FailureClass => "failure_class",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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 failure_filter = if dimension == UsageDiagnosticsDimension::FailureClass {
|
||||
"where not 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 terminal_requests {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 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 terminal_requests\n\
|
||||
where 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,5 +1,6 @@
|
||||
//! 成功请求事实。
|
||||
|
||||
pub mod insights;
|
||||
pub mod query;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -366,7 +366,7 @@ impl SystemUpdateConfig {
|
||||
});
|
||||
let state_file = env_string("CPR_UPDATE_STATE_FILE")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| PathBuf::from("/app/data/update-state.json"));
|
||||
.unwrap_or_else(|| PathBuf::from("/app/.runtime/data/update-state.json"));
|
||||
let update_lock_file = env_string("CPR_UPDATE_LOCK_FILE")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| state_file.with_extension("lock"));
|
||||
|
||||
@@ -127,20 +127,26 @@ pub fn extract_sse_usage(body: &str) -> Result<Option<TokenUsage>, SseError> {
|
||||
/// `try again in 11.054s` / `try again in 28ms` 文本。
|
||||
pub fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
|
||||
let value = serde_json::from_str::<Value>(body).ok()?;
|
||||
retry_after_seconds_from_value(&value)
|
||||
}
|
||||
|
||||
/// 从已解析的上游错误事件中提取重试秒数,不做业务错误分类。
|
||||
pub(crate) fn retry_after_seconds_from_value(value: &Value) -> Option<u64> {
|
||||
let error = value
|
||||
.pointer("/response/error")
|
||||
.or_else(|| value.get("error"))
|
||||
.unwrap_or(&value);
|
||||
if let Some(seconds) = error
|
||||
.get("resets_in_seconds")
|
||||
.and_then(Value::as_u64)
|
||||
.filter(|seconds| *seconds > 0)
|
||||
{
|
||||
return Some(seconds);
|
||||
}
|
||||
retry_after_seconds_from_resets_at(error)
|
||||
.unwrap_or(value);
|
||||
retry_after_seconds_field(error)
|
||||
.or_else(|| {
|
||||
error
|
||||
.get("resets_in_seconds")
|
||||
.and_then(Value::as_u64)
|
||||
.filter(|seconds| *seconds > 0)
|
||||
})
|
||||
.or_else(|| retry_after_seconds_from_resets_at(error))
|
||||
.or_else(|| retry_after_seconds_field(value))
|
||||
.or_else(|| retry_after_seconds_header(value))
|
||||
.or_else(|| retry_after_seconds_from_rate_limit_message(error))
|
||||
.or_else(|| retry_after_seconds_field(&value))
|
||||
}
|
||||
|
||||
fn retry_after_seconds_field(value: &Value) -> Option<u64> {
|
||||
@@ -155,6 +161,31 @@ fn retry_after_seconds_field(value: &Value) -> Option<u64> {
|
||||
.filter(|seconds| *seconds > 0)
|
||||
}
|
||||
|
||||
fn retry_after_seconds_header(value: &Value) -> Option<u64> {
|
||||
value
|
||||
.get("headers")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|headers| {
|
||||
headers.iter().find_map(|(name, value)| {
|
||||
if name.eq_ignore_ascii_case("retry-after") {
|
||||
json_value_as_positive_u64(value)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn json_value_as_positive_u64(value: &Value) -> Option<u64> {
|
||||
let seconds = match value {
|
||||
Value::Number(value) => value.as_u64()?,
|
||||
Value::String(value) => value.trim().parse::<u64>().ok()?,
|
||||
Value::Array(values) => values.first().and_then(json_value_as_positive_u64)?,
|
||||
Value::Null | Value::Bool(_) | Value::Object(_) => return None,
|
||||
};
|
||||
(seconds > 0).then_some(seconds)
|
||||
}
|
||||
|
||||
/// 标准化的单个限流窗口。
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct RateLimitWindow {
|
||||
@@ -489,7 +520,14 @@ fn retry_after_seconds_from_rate_limit_message(error: &Value) -> Option<u64> {
|
||||
.get("code")
|
||||
.or_else(|| error.get("type"))
|
||||
.and_then(Value::as_str)?;
|
||||
if code != "rate_limit_exceeded" {
|
||||
if !matches!(
|
||||
code,
|
||||
"rate_limit_exceeded"
|
||||
| "rate_limit_reached"
|
||||
| "usage_limit_reached"
|
||||
| "workspace_owner_usage_limit_reached"
|
||||
| "workspace_member_usage_limit_reached"
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
let message = error.get("message").and_then(Value::as_str)?;
|
||||
|
||||
@@ -6,11 +6,16 @@ use serde_json::{Map, Value};
|
||||
use crate::{
|
||||
infra::time::elapsed_millis_i64,
|
||||
upstream::openai::protocol::{
|
||||
events,
|
||||
schema::reconvert_tuple_values,
|
||||
sse::{SseError, SseEvent, parse_sse_events},
|
||||
},
|
||||
};
|
||||
|
||||
const MULTI_AGENT_MODE_OPEN_TAG: &str = "<multi_agent_mode>";
|
||||
const MULTI_AGENT_MODE_CLOSE_TAG: &str = "</multi_agent_mode>";
|
||||
const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegation is active.";
|
||||
|
||||
/// Codex Responses 上游请求体。
|
||||
///
|
||||
/// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object,逐字段(含顺序、
|
||||
@@ -70,6 +75,8 @@ pub struct CodexResponsesRequest {
|
||||
pub previous_response_scope: Option<PreviousResponseScope>,
|
||||
/// 当前下游 WebSocket 持有的完整历史前缀,仅用于连接内换号恢复。
|
||||
pub local_replay_input: Option<Arc<Vec<Value>>>,
|
||||
/// 当前下游连接能否在本轮完成后形成完整 canonical replay snapshot。
|
||||
pub local_replay_available: bool,
|
||||
/// 流式响应在交给下游前的本地提交策略。
|
||||
pub stream_commit_policy: StreamCommitPolicy,
|
||||
}
|
||||
@@ -80,6 +87,8 @@ pub struct CodexResponsesRequest {
|
||||
pub enum PreviousResponseScope {
|
||||
Persisted,
|
||||
ConnectionLocal,
|
||||
ReplayRequired,
|
||||
Unavailable,
|
||||
ExternalUnknown,
|
||||
}
|
||||
|
||||
@@ -97,6 +106,8 @@ pub struct CodexRequestSemantics {
|
||||
pub request_kind: Option<String>,
|
||||
/// Codex turn metadata 中的子代理类型。
|
||||
pub subagent_kind: Option<String>,
|
||||
/// Codex 客户端选择的推理预设。
|
||||
pub reasoning_preset: Option<&'static str>,
|
||||
/// 请求是否为远端压缩。
|
||||
pub compact: bool,
|
||||
}
|
||||
@@ -111,40 +122,77 @@ impl Serialize for CodexResponsesRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex Responses 请求的上游传输决策。
|
||||
/// Codex Responses 请求对上游传输的显式要求。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CodexTransport {
|
||||
/// HTTP SSE。
|
||||
HttpSse,
|
||||
/// 优先 WebSocket,失败后可按条件回退 HTTP SSE。
|
||||
WebSocketPreferred,
|
||||
/// 必须使用 WebSocket。
|
||||
WebSocketRequired,
|
||||
pub enum TransportRequirement {
|
||||
/// 客户端显式要求 HTTP。
|
||||
HttpRequired,
|
||||
/// `generate=false + store=false` 预热必须保留在同一条 WebSocket。
|
||||
ExplicitWebSocketWarmup,
|
||||
/// 只能使用持有指定 connection-local response 的精确 WebSocket。
|
||||
ExactWebSocketContinuation,
|
||||
/// previous response 已持久化,允许 WebSocket 或 HTTP/2。
|
||||
PersistedContinuation,
|
||||
/// previous response 的所有权未知,只允许当前选定账号原样尝试。
|
||||
ExternalUnknown,
|
||||
/// 没有 previous response 的普通新链。
|
||||
NewChain,
|
||||
}
|
||||
|
||||
/// 判断 Responses 请求应使用哪种上游传输。
|
||||
pub fn transport_for_request(request: &CodexResponsesRequest) -> CodexTransport {
|
||||
impl TransportRequirement {
|
||||
/// 是否必须使用 WebSocket,且禁止 HTTP fallback。
|
||||
pub fn requires_websocket(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::ExplicitWebSocketWarmup | Self::ExactWebSocketContinuation
|
||||
)
|
||||
}
|
||||
|
||||
/// WebSocket 尚未发送 payload 时失败,是否允许切到同账号 HTTP/2。
|
||||
pub fn allows_pre_send_http_fallback(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::PersistedContinuation | Self::ExternalUnknown | Self::NewChain
|
||||
)
|
||||
}
|
||||
|
||||
/// 用于审计与遥测的稳定名称。
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::HttpRequired => "http_required",
|
||||
Self::ExplicitWebSocketWarmup => "explicit_websocket_warmup",
|
||||
Self::ExactWebSocketContinuation => "exact_websocket_continuation",
|
||||
Self::PersistedContinuation => "persisted_continuation",
|
||||
Self::ExternalUnknown => "external_unknown",
|
||||
Self::NewChain => "new_chain",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 将已完成 history preparation 的请求规范化为唯一 transport requirement。
|
||||
pub fn transport_requirement(request: &CodexResponsesRequest) -> TransportRequirement {
|
||||
if !request.generate() && !request.store() {
|
||||
return TransportRequirement::ExplicitWebSocketWarmup;
|
||||
}
|
||||
if request.force_http_sse {
|
||||
return CodexTransport::HttpSse;
|
||||
return TransportRequirement::HttpRequired;
|
||||
}
|
||||
|
||||
if request.previous_response_id().is_some() {
|
||||
return CodexTransport::WebSocketRequired;
|
||||
match request.previous_response_id() {
|
||||
Some(_) => match request.previous_response_scope {
|
||||
Some(PreviousResponseScope::Persisted) => TransportRequirement::PersistedContinuation,
|
||||
Some(PreviousResponseScope::ConnectionLocal) => {
|
||||
TransportRequirement::ExactWebSocketContinuation
|
||||
}
|
||||
Some(PreviousResponseScope::ExternalUnknown) | None => {
|
||||
TransportRequirement::ExternalUnknown
|
||||
}
|
||||
Some(PreviousResponseScope::ReplayRequired | PreviousResponseScope::Unavailable) => {
|
||||
TransportRequirement::ExactWebSocketContinuation
|
||||
}
|
||||
},
|
||||
None if request.use_websocket => TransportRequirement::NewChain,
|
||||
None => TransportRequirement::HttpRequired,
|
||||
}
|
||||
|
||||
if request.use_websocket {
|
||||
return CodexTransport::WebSocketPreferred;
|
||||
}
|
||||
|
||||
CodexTransport::HttpSse
|
||||
}
|
||||
|
||||
/// 判断请求在 WebSocket 失败后是否允许 HTTP SSE 回退。
|
||||
pub fn http_sse_fallback_allowed(request: &CodexResponsesRequest) -> bool {
|
||||
!matches!(
|
||||
transport_for_request(request),
|
||||
CodexTransport::WebSocketRequired
|
||||
)
|
||||
}
|
||||
|
||||
/// 判断已收到的 Responses SSE 内容是否包含首个完整的真实输出事件。
|
||||
@@ -244,6 +292,12 @@ pub struct ResponsesSseFailure {
|
||||
pub message: String,
|
||||
/// 上游错误码。
|
||||
pub upstream_code: Option<String>,
|
||||
/// 上游显式错误类型;不从业务码推导。
|
||||
pub upstream_type: Option<String>,
|
||||
/// 上游显式状态码;不从业务码或错误类型推导。
|
||||
pub explicit_status_code: Option<u16>,
|
||||
/// 上游显式重试间隔,或从官方限流消息中解析出的重试间隔。
|
||||
pub retry_after_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
impl ResponsesSseFailure {
|
||||
@@ -252,6 +306,9 @@ impl ResponsesSseFailure {
|
||||
event: event.to_string(),
|
||||
message: failure_message(value).unwrap_or_else(|| "Codex upstream SSE failed".into()),
|
||||
upstream_code: failure_code(value),
|
||||
upstream_type: failure_type(value),
|
||||
explicit_status_code: failure_explicit_status_code(value),
|
||||
retry_after_seconds: events::retry_after_seconds_from_value(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -267,66 +324,6 @@ pub struct CompletedResponseMetadata {
|
||||
pub output: Vec<Value>,
|
||||
}
|
||||
|
||||
/// 将 Codex Responses SSE 完成响应收集为非流式 Responses JSON。
|
||||
pub fn response_from_codex_sse(
|
||||
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(value) = serde_json::from_str::<Value>(&event.data) else {
|
||||
continue;
|
||||
};
|
||||
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(mut response) = terminal_response else {
|
||||
return Ok(CollectedResponse::MissingCompleted);
|
||||
};
|
||||
if terminal_incomplete {
|
||||
reconvert_completed_response_tuple_values(&mut response, tuple_schema);
|
||||
return Ok(CollectedResponse::Incomplete(response));
|
||||
}
|
||||
if is_empty_response(&response, &output_text, &output_items) {
|
||||
return Ok(CollectedResponse::Empty);
|
||||
}
|
||||
|
||||
reconvert_completed_response_tuple_values(&mut response, tuple_schema);
|
||||
Ok(CollectedResponse::Completed(response))
|
||||
}
|
||||
|
||||
/// 从 Codex Responses SSE 中提取完成响应元数据。
|
||||
pub fn completed_response_metadata(
|
||||
body: &str,
|
||||
@@ -431,21 +428,6 @@ fn reconvert_output_text_delta_tuple_values(data: &mut Value, tuple_schema: &Val
|
||||
data["delta"] = Value::String(reconverted.to_string());
|
||||
}
|
||||
|
||||
fn is_empty_response(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
|
||||
}
|
||||
|
||||
fn reconvert_completed_response_tuple_values(response: &mut Value, tuple_schema: Option<&Value>) {
|
||||
let Some(tuple_schema) = tuple_schema else {
|
||||
return;
|
||||
@@ -509,6 +491,28 @@ fn failure_code(value: &Value) -> Option<String> {
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn failure_type(value: &Value) -> Option<String> {
|
||||
failure_error(value)
|
||||
.and_then(|error| error.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|error_type| !error_type.trim().is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn failure_explicit_status_code(value: &Value) -> Option<u16> {
|
||||
value
|
||||
.get("status")
|
||||
.or_else(|| value.get("status_code"))
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|status| u16::try_from(status).ok())
|
||||
}
|
||||
|
||||
fn failure_error(value: &Value) -> Option<&Value> {
|
||||
value
|
||||
.pointer("/response/error")
|
||||
.or_else(|| value.get("error"))
|
||||
}
|
||||
|
||||
impl CodexResponsesRequest {
|
||||
/// 从客户端原始 Responses JSON object 构造上游请求。
|
||||
///
|
||||
@@ -540,6 +544,7 @@ impl CodexResponsesRequest {
|
||||
parent_thread_id: None,
|
||||
previous_response_scope: None,
|
||||
local_replay_input: None,
|
||||
local_replay_available: false,
|
||||
stream_commit_policy: StreamCommitPolicy::FirstForwardableEvent,
|
||||
}
|
||||
}
|
||||
@@ -623,6 +628,14 @@ impl CodexResponsesRequest {
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// 是否实际生成模型响应;官方预热请求会显式传入 `false`。
|
||||
pub fn generate(&self) -> bool {
|
||||
self.body
|
||||
.get("generate")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(true)
|
||||
}
|
||||
|
||||
/// 设置上游存储标志。
|
||||
pub fn set_store(&mut self, store: bool) {
|
||||
self.body.insert("store".to_string(), Value::Bool(store));
|
||||
@@ -696,7 +709,7 @@ impl CodexResponsesRequest {
|
||||
self.body.get("client_metadata")
|
||||
}
|
||||
|
||||
/// 提取 Codex 请求类型、子代理类型与压缩语义。
|
||||
/// 提取 Codex 请求类型、子代理类型、推理预设与压缩语义。
|
||||
pub fn semantics(&self) -> CodexRequestSemantics {
|
||||
let turn_metadata = self
|
||||
.turn_metadata
|
||||
@@ -717,6 +730,16 @@ impl CodexResponsesRequest {
|
||||
.and_then(|metadata| metadata.get("subagent_kind"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(non_empty_owned_string);
|
||||
let reasoning_preset = (subagent_kind.is_none()
|
||||
&& self
|
||||
.reasoning()
|
||||
.and_then(|value| value.get("effort"))
|
||||
.and_then(Value::as_str)
|
||||
== Some("max")
|
||||
&& self
|
||||
.latest_multi_agent_mode()
|
||||
.is_some_and(|mode| mode.starts_with(PROACTIVE_MULTI_AGENT_MODE_PREFIX)))
|
||||
.then_some("ultra");
|
||||
let compact = request_kind.as_deref() == Some("compaction")
|
||||
|| self
|
||||
.input()
|
||||
@@ -726,10 +749,19 @@ impl CodexResponsesRequest {
|
||||
CodexRequestSemantics {
|
||||
request_kind,
|
||||
subagent_kind,
|
||||
reasoning_preset,
|
||||
compact,
|
||||
}
|
||||
}
|
||||
|
||||
fn latest_multi_agent_mode(&self) -> Option<&str> {
|
||||
latest_multi_agent_mode(self.input()).or_else(|| {
|
||||
self.local_replay_input
|
||||
.as_deref()
|
||||
.and_then(|input| latest_multi_agent_mode(input.as_slice()))
|
||||
})
|
||||
}
|
||||
|
||||
/// 设置 / 合并 client metadata。
|
||||
pub fn set_client_metadata(&mut self, client_metadata: Option<Value>) {
|
||||
match client_metadata {
|
||||
@@ -742,9 +774,9 @@ impl CodexResponsesRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// 仅在客户端原本提供字符串字段时替换其账号作用域 wire 值。
|
||||
/// 替换客户端原本提供的账号身份字段;无法安全重建时删除该字段。
|
||||
pub(crate) fn replace_existing_identity_field(&mut self, key: &str, value: Option<&str>) {
|
||||
if !self.body.get(key).is_some_and(Value::is_string) {
|
||||
if !self.body.contains_key(key) {
|
||||
return;
|
||||
}
|
||||
match value.filter(|value| !value.trim().is_empty()) {
|
||||
@@ -772,3 +804,25 @@ fn non_empty_owned_string(value: &str) -> Option<String> {
|
||||
let value = value.trim();
|
||||
(!value.is_empty()).then(|| value.to_string())
|
||||
}
|
||||
|
||||
fn latest_multi_agent_mode(input: &[Value]) -> Option<&str> {
|
||||
input.iter().rev().find_map(|item| {
|
||||
if item.get("role").and_then(Value::as_str) != Some("developer") {
|
||||
return None;
|
||||
}
|
||||
item.get("content")?
|
||||
.as_array()?
|
||||
.iter()
|
||||
.rev()
|
||||
.filter_map(|content| content.get("text").and_then(Value::as_str))
|
||||
.find_map(multi_agent_mode_from_text)
|
||||
})
|
||||
}
|
||||
|
||||
fn multi_agent_mode_from_text(text: &str) -> Option<&str> {
|
||||
let close = text.rfind(MULTI_AGENT_MODE_CLOSE_TAG)?;
|
||||
let before_close = &text[..close];
|
||||
let open = before_close.rfind(MULTI_AGENT_MODE_OPEN_TAG)?;
|
||||
let body = &before_close[open + MULTI_AGENT_MODE_OPEN_TAG.len()..];
|
||||
Some(body.trim())
|
||||
}
|
||||
|
||||
@@ -22,6 +22,16 @@ pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024;
|
||||
/// SSE 流结束标记帧。
|
||||
pub const DONE_SSE_FRAME: &str = "data: [DONE]\n\n";
|
||||
|
||||
/// 判断一个完整 SSE frame 是否为传输层 `[DONE]` 控制帧。
|
||||
pub(crate) fn sse_frame_is_done(frame: &str) -> bool {
|
||||
let mut data = frame.lines().filter_map(|raw_line| {
|
||||
let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
|
||||
let (field, value) = split_sse_field(line);
|
||||
(field == "data").then_some(value)
|
||||
});
|
||||
matches!((data.next(), data.next()), (Some("[DONE]"), None))
|
||||
}
|
||||
|
||||
/// 用于处理任意分块边界的增量 SSE 解码器。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct SseEventDecoder {
|
||||
@@ -78,6 +88,17 @@ pub fn response_failed_sse_event_with_id(
|
||||
code: &str,
|
||||
message: &str,
|
||||
) -> String {
|
||||
let data = response_failed_sse_data_with_id(response_id, error_type, code, message);
|
||||
encode_sse_event("response.failed", &data.to_string())
|
||||
}
|
||||
|
||||
/// 构造 OpenAI Responses `response.failed` 的 JSON 数据。
|
||||
pub fn response_failed_sse_data_with_id(
|
||||
response_id: Option<&str>,
|
||||
error_type: &str,
|
||||
code: &str,
|
||||
message: &str,
|
||||
) -> serde_json::Value {
|
||||
let error = json!({
|
||||
"type": error_type,
|
||||
"code": code,
|
||||
@@ -89,7 +110,7 @@ pub fn response_failed_sse_event_with_id(
|
||||
|| format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()),
|
||||
ToString::to_string,
|
||||
);
|
||||
let data = json!({
|
||||
json!({
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
@@ -97,8 +118,7 @@ pub fn response_failed_sse_event_with_id(
|
||||
"error": error,
|
||||
},
|
||||
"error": error,
|
||||
});
|
||||
encode_sse_event("response.failed", &data.to_string())
|
||||
})
|
||||
}
|
||||
|
||||
/// 判断 SSE 文本是否已经包含结束标记。
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::upstream::openai::protocol::responses::{
|
||||
CodexResponsesRequest, CodexTransport, http_sse_fallback_allowed, transport_for_request,
|
||||
};
|
||||
use crate::upstream::openai::protocol::responses::{CodexResponsesRequest, transport_requirement};
|
||||
use crate::upstream::openai::protocol::sse::encode_sse_event;
|
||||
|
||||
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
|
||||
@@ -181,10 +179,6 @@ pub fn websocket_response_completed_id(raw: &str) -> Result<Option<String>, Stri
|
||||
Ok(Some(completed.id))
|
||||
}
|
||||
|
||||
#[path = "websocket_errors.rs"]
|
||||
mod errors;
|
||||
pub use errors::*;
|
||||
|
||||
/// 判断 WebSocket 事件是否会结束当前响应流。
|
||||
pub fn is_terminal_websocket_event(event: &str) -> bool {
|
||||
matches!(
|
||||
@@ -209,21 +203,13 @@ pub fn websocket_audit_artifact_from_attempt(
|
||||
payload: PayloadAuditSnapshot,
|
||||
) -> WebSocketAuditArtifact {
|
||||
WebSocketAuditArtifact {
|
||||
transport_mode: websocket_transport_mode_name(request).to_string(),
|
||||
fallback_allowed: http_sse_fallback_allowed(request),
|
||||
transport_mode: transport_requirement(request).as_str().to_string(),
|
||||
fallback_allowed: transport_requirement(request).allows_pre_send_http_fallback(),
|
||||
opening: Some(opening),
|
||||
payload: Some(payload),
|
||||
}
|
||||
}
|
||||
|
||||
fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static str {
|
||||
match transport_for_request(request) {
|
||||
CodexTransport::HttpSse => "http_sse",
|
||||
CodexTransport::WebSocketPreferred => "websocket_preferred",
|
||||
CodexTransport::WebSocketRequired => "websocket_required",
|
||||
}
|
||||
}
|
||||
|
||||
/// 生成 Responses WebSocket `response.create` payload。
|
||||
///
|
||||
/// payload = `{"type": "response.create"}` 加上原始上游 body 的全部字段
|
||||
|
||||
@@ -1,154 +0,0 @@
|
||||
//! Responses WebSocket 错误帧解析。
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::json_value_as_string;
|
||||
|
||||
/// 从 WebSocket 错误帧推导出的上游错误分类。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ClassifiedWebSocketError {
|
||||
/// 应映射给 HTTP 层的状态码。
|
||||
pub status_code: u16,
|
||||
/// 上游错误类型。
|
||||
pub error_type: Option<String>,
|
||||
/// 上游错误码。
|
||||
pub code: Option<String>,
|
||||
/// 上游错误消息。
|
||||
pub message: Option<String>,
|
||||
/// 上游错误参数。
|
||||
pub param: Option<String>,
|
||||
/// 错误帧携带的安全响应头。
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
/// 将 WebSocket 错误帧映射为结构化上游错误。
|
||||
pub fn classify_websocket_error_frame(raw: &str) -> Option<ClassifiedWebSocketError> {
|
||||
let value = serde_json::from_str::<Value>(raw).ok()?;
|
||||
let event_type = value.get("type").and_then(Value::as_str)?;
|
||||
if event_type != "error" && event_type != "response.failed" {
|
||||
return None;
|
||||
}
|
||||
|
||||
let error = websocket_error_value(&value);
|
||||
let error_type = error
|
||||
.and_then(|error| error.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string);
|
||||
let code = error
|
||||
.and_then(|error| error.get("code"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string);
|
||||
let message = error
|
||||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string);
|
||||
let param = error
|
||||
.and_then(|error| error.get("param"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string);
|
||||
let explicit_status =
|
||||
explicit_error_status_code(&value).filter(|status_code| !(200..=299).contains(status_code));
|
||||
if explicit_error_status_code(&value).is_some() && explicit_status.is_none() {
|
||||
return None;
|
||||
}
|
||||
let normalized_code = code.as_deref().map(str::to_ascii_lowercase);
|
||||
let normalized_type = error_type.as_deref().map(str::to_ascii_lowercase);
|
||||
let status_code = explicit_status
|
||||
.or_else(|| normalized_code.as_deref().and_then(error_status_code))
|
||||
.or_else(|| normalized_type.as_deref().and_then(error_type_status_code))?;
|
||||
|
||||
Some(ClassifiedWebSocketError {
|
||||
status_code,
|
||||
error_type,
|
||||
code,
|
||||
message,
|
||||
param,
|
||||
headers: websocket_error_headers(&value),
|
||||
})
|
||||
}
|
||||
|
||||
fn websocket_error_value(value: &Value) -> Option<&Map<String, Value>> {
|
||||
value
|
||||
.pointer("/response/error")
|
||||
.or_else(|| value.get("error"))
|
||||
.and_then(Value::as_object)
|
||||
}
|
||||
|
||||
fn explicit_error_status_code(value: &Value) -> Option<u16> {
|
||||
value
|
||||
.get("status")
|
||||
.or_else(|| value.get("status_code"))
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|status| u16::try_from(status).ok())
|
||||
}
|
||||
|
||||
fn error_status_code(code: &str) -> Option<u16> {
|
||||
match code {
|
||||
"usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => Some(429),
|
||||
"quota_exhausted" | "payment_required" => Some(402),
|
||||
"unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => Some(401),
|
||||
"forbidden" | "account_banned" | "banned" => Some(403),
|
||||
"previous_response_not_found"
|
||||
| "invalid_encrypted_content"
|
||||
| "missing_tool_output"
|
||||
| "no_tool_output" => Some(400),
|
||||
"websocket_connection_limit_reached" => Some(503),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn error_type_status_code(error_type: &str) -> Option<u16> {
|
||||
match error_type {
|
||||
"invalid_request_error" => Some(400),
|
||||
"authentication_error" => Some(401),
|
||||
"permission_error" => Some(403),
|
||||
"rate_limit_error" => Some(429),
|
||||
"server_error" | "api_error" => Some(500),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn websocket_error_headers(value: &Value) -> Vec<(String, String)> {
|
||||
value
|
||||
.get("headers")
|
||||
.and_then(Value::as_object)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|(name, value)| json_value_as_string(value).map(|value| (name.clone(), value)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 从包裹在 WebSocket `error.headers` 中的 retry-after 值提取秒数。
|
||||
pub fn retry_after_seconds_from_wrapped_error_headers(raw: &str) -> Option<u64> {
|
||||
let value = serde_json::from_str::<Value>(raw).ok()?;
|
||||
let header_retry_after = value
|
||||
.get("headers")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|headers| {
|
||||
headers.iter().find_map(|(name, value)| {
|
||||
if name.eq_ignore_ascii_case("retry-after") {
|
||||
retry_after_seconds_from_json_header_value(value)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
});
|
||||
header_retry_after.or_else(|| {
|
||||
value
|
||||
.pointer("/error/retry_after_seconds")
|
||||
.or_else(|| value.get("retry_after_seconds"))
|
||||
.and_then(Value::as_u64)
|
||||
})
|
||||
}
|
||||
|
||||
fn retry_after_seconds_from_json_header_value(value: &Value) -> Option<u64> {
|
||||
let seconds = match value {
|
||||
Value::String(value) => value.trim().parse::<u64>().ok()?,
|
||||
Value::Number(value) => value.as_u64()?,
|
||||
Value::Array(values) => values
|
||||
.first()
|
||||
.and_then(retry_after_seconds_from_json_header_value)?,
|
||||
_ => return None,
|
||||
};
|
||||
(seconds > 0).then_some(seconds)
|
||||
}
|
||||
@@ -24,7 +24,7 @@ use tokio_tungstenite::tungstenite::handshake::client::generate_key;
|
||||
use crate::upstream::openai::fingerprint::{Fingerprint, RuntimeFingerprint};
|
||||
use crate::upstream::openai::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
|
||||
use crate::upstream::openai::protocol::responses::{
|
||||
CodexResponsesRequest, CodexTransport, http_sse_fallback_allowed, transport_for_request,
|
||||
CodexResponsesRequest, TransportRequirement, transport_requirement,
|
||||
};
|
||||
use crate::upstream::openai::protocol::sse::SseError;
|
||||
use crate::upstream::openai::protocol::websocket::{
|
||||
@@ -41,9 +41,12 @@ use super::response_meta;
|
||||
use super::tls::{CustomCaError, build_reqwest_client_with_custom_ca, custom_ca_env_cache_key};
|
||||
use super::websocket::{
|
||||
CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketRateLimitHeaderUpdates,
|
||||
CodexWebSocketTurnStateUpdate, execute_response_create_request_stream_with_pool,
|
||||
execute_response_create_request_with_pool, write_websocket_audit_artifact_from_env,
|
||||
CodexWebSocketRequest, CodexWebSocketTurnStateUpdate, PreparedWebSocket,
|
||||
WEBSOCKET_FAST_PATH_BUDGET, execute_prepared_response_create_request,
|
||||
execute_prepared_response_create_request_stream, post_send_ambiguous,
|
||||
prepare_response_create_request_with_pool, write_websocket_audit_artifact_from_env,
|
||||
};
|
||||
use super::websocket_breaker::WebSocketOriginBreaker;
|
||||
use super::websocket_pool::{
|
||||
CodexWebSocketPool, CodexWebSocketPoolKey, DEFAULT_INITIAL_EVENT_TIMEOUT, WebSocketPoolDecision,
|
||||
};
|
||||
@@ -57,12 +60,12 @@ const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
|
||||
"x-codex-ws-stream-request-start-ms";
|
||||
type ReqwestClientCacheKey = (bool, Option<String>);
|
||||
type ReqwestClientCacheKey = Option<String>;
|
||||
type ReqwestClientCache = Mutex<HashMap<ReqwestClientCacheKey, Client>>;
|
||||
|
||||
/// 构建带缓存的 reqwest Client。若 `force_http11` 为 true 则强制 HTTP/1.1。
|
||||
pub fn build_reqwest_client(force_http11: bool) -> Result<Client, CustomCaError> {
|
||||
let cache_key = (force_http11, custom_ca_env_cache_key());
|
||||
/// 构建带缓存、自动协商 HTTP/2 的 reqwest Client。
|
||||
pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
|
||||
let cache_key = custom_ca_env_cache_key();
|
||||
static CLIENTS: OnceLock<ReqwestClientCache> = OnceLock::new();
|
||||
let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new()));
|
||||
if let Some(client) = cache
|
||||
@@ -77,17 +80,16 @@ pub fn build_reqwest_client(force_http11: bool) -> Result<Client, CustomCaError>
|
||||
.use_rustls_tls()
|
||||
.no_proxy()
|
||||
.pool_max_idle_per_host(4)
|
||||
.pool_idle_timeout(None::<Duration>)
|
||||
.connect_timeout(UPSTREAM_CONNECT_TIMEOUT)
|
||||
.tcp_keepalive(Duration::from_secs(30))
|
||||
.http2_keep_alive_interval(Duration::from_secs(30))
|
||||
.http2_keep_alive_timeout(Duration::from_secs(5))
|
||||
.http2_keep_alive_while_idle(true)
|
||||
.gzip(true)
|
||||
.brotli(true)
|
||||
.zstd(true)
|
||||
.deflate(true);
|
||||
let builder = if force_http11 {
|
||||
builder.http1_only()
|
||||
} else {
|
||||
builder
|
||||
};
|
||||
let client = build_reqwest_client_with_custom_ca(builder)?;
|
||||
let mut clients = cache
|
||||
.lock()
|
||||
@@ -142,9 +144,28 @@ pub enum CodexClientError {
|
||||
diagnostics: CodexUpstreamDiagnostics,
|
||||
/// 上游透传的 `set-cookie` 列表。
|
||||
set_cookie_headers: Vec<String>,
|
||||
/// 实际收到该上游响应的 transport。
|
||||
transport: CodexBackendTransport,
|
||||
},
|
||||
}
|
||||
|
||||
impl CodexClientError {
|
||||
/// 返回错误实际发生的 transport;请求编码等本地错误没有 transport。
|
||||
pub fn transport(&self) -> Option<CodexBackendTransport> {
|
||||
match self {
|
||||
Self::Http(_) | Self::StreamIdleTimeout { .. } | Self::InvalidSse(_) => {
|
||||
Some(CodexBackendTransport::HttpSse)
|
||||
}
|
||||
Self::WebSocket(_) => Some(CodexBackendTransport::WebSocket),
|
||||
Self::Upstream { transport, .. } => Some(*transport),
|
||||
Self::CustomCa(_)
|
||||
| Self::InvalidHeaderName(_)
|
||||
| Self::InvalidHeaderValue(_)
|
||||
| Self::WebSocketEncode(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断上游错误正文是否表示账号已封禁或停用。
|
||||
pub fn is_banned_auth_signal(value: &str) -> bool {
|
||||
let value = value.to_ascii_lowercase();
|
||||
@@ -155,12 +176,21 @@ pub fn is_banned_auth_signal(value: &str) -> bool {
|
||||
|| value.contains("banned")
|
||||
}
|
||||
|
||||
/// 判断 402 错误正文是否带有 OpenAI 工作区停用标记。
|
||||
pub fn is_deactivated_workspace_error_body(value: &str) -> bool {
|
||||
serde_json::from_str::<Value>(value).is_ok_and(|value| {
|
||||
value.pointer("/detail/code").and_then(Value::as_str) == Some("deactivated_workspace")
|
||||
})
|
||||
}
|
||||
|
||||
/// 判断 Codex 上游错误是否表示账号已封禁或停用。
|
||||
pub fn is_banned_upstream_error(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, body, .. }
|
||||
if status.as_u16() == 403 && !is_html_error_body(body)
|
||||
if (status.as_u16() == 403
|
||||
&& !is_html_error_body(body))
|
||||
|| (status.as_u16() == 402 && is_deactivated_workspace_error_body(body))
|
||||
)
|
||||
}
|
||||
|
||||
@@ -223,11 +253,11 @@ pub struct CodexRequestContext<'a> {
|
||||
pub access_token: &'a str,
|
||||
/// ChatGPT account id。
|
||||
pub account_id: Option<&'a str>,
|
||||
/// 请求 ID。
|
||||
/// 代理请求 ID。
|
||||
pub request_id: &'a str,
|
||||
/// x-codex-turn-state。
|
||||
/// 当前账号同一 turn 内的 opaque sticky-routing 状态。
|
||||
pub turn_state: Option<&'a str>,
|
||||
/// x-codex-turn-metadata。
|
||||
/// 客户端 turn metadata;其中 installation ID 已按当前账号处理。
|
||||
pub turn_metadata: Option<&'a str>,
|
||||
/// x-codex-beta-features。
|
||||
pub beta_features: Option<&'a str>,
|
||||
@@ -235,23 +265,21 @@ pub struct CodexRequestContext<'a> {
|
||||
pub include_timing_metrics: Option<&'a str>,
|
||||
/// version。
|
||||
pub version: Option<&'a str>,
|
||||
/// x-codex-window-id。
|
||||
/// 客户端 window ID。
|
||||
pub codex_window_id: Option<&'a str>,
|
||||
/// x-codex-parent-thread-id。
|
||||
/// 客户端 parent thread ID。
|
||||
pub parent_thread_id: Option<&'a str>,
|
||||
/// cookie 头。
|
||||
pub cookie_header: Option<&'a str>,
|
||||
/// x-codex-installation-id。
|
||||
/// 当前账号稳定派生的 installation ID。
|
||||
pub installation_id: Option<&'a str>,
|
||||
/// session_id。
|
||||
/// 客户端 session ID。
|
||||
pub session_id: Option<&'a str>,
|
||||
/// thread_id。
|
||||
/// 客户端 thread ID。
|
||||
pub thread_id: Option<&'a str>,
|
||||
/// 账号作用域 prompt cache key。
|
||||
pub prompt_cache_key: Option<&'a str>,
|
||||
/// 账号作用域 client request ID。
|
||||
/// 客户端逻辑 request ID;缺失时使用代理请求 ID。
|
||||
pub client_request_id: Option<&'a str>,
|
||||
/// 账号作用域 turn ID。
|
||||
/// 客户端 turn ID。
|
||||
pub turn_id: Option<&'a str>,
|
||||
}
|
||||
|
||||
@@ -291,6 +319,10 @@ pub struct CodexBackendResponse {
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
/// 安全响应元数据。
|
||||
pub response_metadata: CodexResponseMetadata,
|
||||
/// 传输选择与低延迟阶段耗时。
|
||||
pub transport_metrics: CodexTransportMetrics,
|
||||
/// 当前响应是否由池中 WebSocket 保留 connection-local continuation。
|
||||
pub connection_local_continuation: bool,
|
||||
}
|
||||
|
||||
/// Codex Responses 实际使用的上游传输。
|
||||
@@ -302,18 +334,47 @@ pub enum CodexBackendTransport {
|
||||
WebSocket,
|
||||
}
|
||||
|
||||
/// Map a Codex Responses request to the concrete backend transport.
|
||||
pub fn backend_transport_for_response_request(
|
||||
request: &CodexResponsesRequest,
|
||||
) -> CodexBackendTransport {
|
||||
match transport_for_request(request) {
|
||||
CodexTransport::HttpSse => CodexBackendTransport::HttpSse,
|
||||
CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => {
|
||||
CodexBackendTransport::WebSocket
|
||||
/// transport owner 最终做出的稳定决策。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CodexTransportDecision {
|
||||
HttpRequired,
|
||||
ReusedWebSocket,
|
||||
ConnectedWebSocket,
|
||||
ExactWebSocket,
|
||||
RequiredWebSocket,
|
||||
Http2WebSocketSlow,
|
||||
Http2BreakerOpen,
|
||||
Http2PoolUnavailable,
|
||||
Http2PreSendFailure,
|
||||
}
|
||||
|
||||
impl CodexTransportDecision {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::HttpRequired => "http_required",
|
||||
Self::ReusedWebSocket => "ws_reused",
|
||||
Self::ConnectedWebSocket => "ws_connected_fast",
|
||||
Self::ExactWebSocket => "ws_exact_required",
|
||||
Self::RequiredWebSocket => "ws_required",
|
||||
Self::Http2WebSocketSlow => "http2_ws_slow",
|
||||
Self::Http2BreakerOpen => "http2_breaker_open",
|
||||
Self::Http2PoolUnavailable => "http2_pool_unavailable",
|
||||
Self::Http2PreSendFailure => "http2_ws_pre_send_failure",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// transport 决策与握手阶段观测值。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct CodexTransportMetrics {
|
||||
pub decision: Option<CodexTransportDecision>,
|
||||
pub ws_connect_ms: Option<i64>,
|
||||
pub transport_decision_wait_ms: Option<i64>,
|
||||
pub upstream_headers_ms: Option<i64>,
|
||||
pub first_event_ms: Option<i64>,
|
||||
pub http_version: Option<String>,
|
||||
}
|
||||
|
||||
/// Live stream rate-limit updates captured after the response headers.
|
||||
pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates;
|
||||
|
||||
@@ -342,6 +403,10 @@ pub struct CodexBackendStreamingResponse {
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
/// 安全响应元数据。
|
||||
pub response_metadata: CodexResponseMetadata,
|
||||
/// 传输选择与低延迟阶段耗时。
|
||||
pub transport_metrics: CodexTransportMetrics,
|
||||
/// terminal completed 后是否由池中 WebSocket 保留 connection-local continuation。
|
||||
pub connection_local_continuation: bool,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -356,6 +421,26 @@ pub struct CodexBackendClient {
|
||||
fingerprint: RuntimeFingerprint,
|
||||
websocket_pool: Option<Arc<CodexWebSocketPool>>,
|
||||
websocket_initial_event_timeout: Option<Duration>,
|
||||
websocket_fast_path_budget: Duration,
|
||||
websocket_origin_breaker: WebSocketOriginBreaker,
|
||||
websocket_origin_key: String,
|
||||
}
|
||||
|
||||
/// 已完成账号级 opening 准备、但尚未发送 payload 的 transport。
|
||||
pub(crate) struct PreparedResponseTransport {
|
||||
requirement: TransportRequirement,
|
||||
route: PreparedResponseRoute,
|
||||
metrics: CodexTransportMetrics,
|
||||
}
|
||||
|
||||
enum PreparedResponseRoute {
|
||||
Http,
|
||||
WebSocket(Box<PreparedWebSocketRoute>),
|
||||
}
|
||||
|
||||
struct PreparedWebSocketRoute {
|
||||
request: CodexWebSocketRequest,
|
||||
prepared: PreparedWebSocket,
|
||||
}
|
||||
|
||||
#[path = "client_sse.rs"]
|
||||
@@ -382,7 +467,6 @@ impl CodexModelCatalogClient for CodexBackendClient {
|
||||
installation_id: request.installation_id,
|
||||
session_id: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
client_request_id: None,
|
||||
turn_id: None,
|
||||
})
|
||||
@@ -407,28 +491,15 @@ fn log_websocket_pool_decision(
|
||||
return;
|
||||
};
|
||||
let rid_short = request_id.chars().take(8).collect::<String>();
|
||||
if let Some(reason) = decision.reason() {
|
||||
tracing::info!(
|
||||
request_id = %request_id,
|
||||
rid = %rid_short,
|
||||
account_id = account_id.unwrap_or_default(),
|
||||
ws_pool = decision.kind(),
|
||||
ws_pool_reason = reason,
|
||||
conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()),
|
||||
ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()),
|
||||
"websocket pool decision"
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
request_id = %request_id,
|
||||
rid = %rid_short,
|
||||
account_id = account_id.unwrap_or_default(),
|
||||
ws_pool = decision.kind(),
|
||||
conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()),
|
||||
ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()),
|
||||
"websocket pool decision"
|
||||
);
|
||||
}
|
||||
tracing::info!(
|
||||
request_id = %request_id,
|
||||
rid = %rid_short,
|
||||
account_id = account_id.unwrap_or_default(),
|
||||
ws_pool = decision.kind(),
|
||||
conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()),
|
||||
ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()),
|
||||
"WebSocket pool decision"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -508,29 +579,6 @@ fn is_model_entry(value: &Value) -> bool {
|
||||
// Request helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn response_upstream_request(
|
||||
request: &CodexResponsesRequest,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexResponsesRequest {
|
||||
let mut upstream = request.clone();
|
||||
if request.prompt_cache_key().is_some() {
|
||||
upstream.set_prompt_cache_key(context.prompt_cache_key.map(ToString::to_string));
|
||||
}
|
||||
for (key, value) in [
|
||||
("session_id", context.session_id),
|
||||
("thread_id", context.thread_id),
|
||||
("turn_id", context.turn_id),
|
||||
("x-client-request-id", context.client_request_id),
|
||||
("x-codex-window-id", context.codex_window_id),
|
||||
("x-codex-parent-thread-id", context.parent_thread_id),
|
||||
("x-codex-installation-id", context.installation_id),
|
||||
] {
|
||||
upstream.replace_existing_identity_field(key, value);
|
||||
}
|
||||
upstream.set_client_metadata(response_client_metadata(request.client_metadata(), context));
|
||||
upstream
|
||||
}
|
||||
|
||||
fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponsesRequest {
|
||||
let mut request = request.clone();
|
||||
stamp_ws_stream_request_start_ms(&mut request);
|
||||
@@ -557,66 +605,6 @@ fn now_unix_timestamp_millis() -> u128 {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn response_client_metadata(
|
||||
client_metadata: Option<&Value>,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> Option<Value> {
|
||||
// 以客户端原始 client_metadata 为基础(保留 number/bool/object 等所有值类型),
|
||||
// 在其上追加代理自身的上下文字段。
|
||||
let mut metadata = match client_metadata? {
|
||||
Value::Object(input) => input.clone(),
|
||||
value => return Some(value.clone()),
|
||||
};
|
||||
|
||||
replace_metadata_string(
|
||||
&mut metadata,
|
||||
"x-codex-installation-id",
|
||||
context.installation_id,
|
||||
);
|
||||
replace_metadata_string(&mut metadata, "session_id", context.session_id);
|
||||
replace_metadata_string(&mut metadata, "thread_id", context.thread_id);
|
||||
replace_metadata_string(
|
||||
&mut metadata,
|
||||
"x-client-request-id",
|
||||
context.client_request_id,
|
||||
);
|
||||
replace_metadata_string(&mut metadata, "turn_id", context.turn_id);
|
||||
replace_metadata_string(&mut metadata, "x-codex-window-id", context.codex_window_id);
|
||||
insert_metadata_string(
|
||||
&mut metadata,
|
||||
"x-codex-turn-metadata",
|
||||
context.turn_metadata,
|
||||
);
|
||||
replace_metadata_string(
|
||||
&mut metadata,
|
||||
"x-codex-parent-thread-id",
|
||||
context.parent_thread_id,
|
||||
);
|
||||
|
||||
if metadata.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(Value::Object(metadata))
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_metadata_string(metadata: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
if let Some(value) = value.filter(|value| !value.trim().is_empty()) {
|
||||
metadata.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
fn replace_metadata_string(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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option<String> {
|
||||
let value = client_metadata?
|
||||
.as_object()?
|
||||
@@ -650,12 +638,83 @@ fn websocket_exchange_error_to_client_error(
|
||||
retry_after_seconds: upstream.retry_after_seconds,
|
||||
diagnostics: upstream.diagnostics,
|
||||
set_cookie_headers: upstream.set_cookie_headers,
|
||||
transport: CodexBackendTransport::WebSocket,
|
||||
}
|
||||
}
|
||||
error => CodexClientError::WebSocket(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn websocket_error_allows_http_fallback(error: &CodexClientError) -> bool {
|
||||
!matches!(error, CodexClientError::Upstream { .. })
|
||||
fn websocket_success_decision(
|
||||
requirement: TransportRequirement,
|
||||
prepared: &PreparedWebSocket,
|
||||
) -> CodexTransportDecision {
|
||||
match requirement {
|
||||
TransportRequirement::ExactWebSocketContinuation => CodexTransportDecision::ExactWebSocket,
|
||||
TransportRequirement::ExplicitWebSocketWarmup => CodexTransportDecision::RequiredWebSocket,
|
||||
_ if prepared.reused() => CodexTransportDecision::ReusedWebSocket,
|
||||
_ => CodexTransportDecision::ConnectedWebSocket,
|
||||
}
|
||||
}
|
||||
|
||||
fn http_fallback_decision(error: &CodexWebSocketExchangeError) -> CodexTransportDecision {
|
||||
match error {
|
||||
CodexWebSocketExchangeError::FastPathTimeout { .. } => {
|
||||
CodexTransportDecision::Http2WebSocketSlow
|
||||
}
|
||||
CodexWebSocketExchangeError::OriginCircuitOpen
|
||||
| CodexWebSocketExchangeError::OriginHalfOpenBusy => {
|
||||
CodexTransportDecision::Http2BreakerOpen
|
||||
}
|
||||
CodexWebSocketExchangeError::ContinuationUnavailable { .. }
|
||||
| CodexWebSocketExchangeError::SharedConnectFailed => {
|
||||
CodexTransportDecision::Http2PoolUnavailable
|
||||
}
|
||||
_ => CodexTransportDecision::Http2PreSendFailure,
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_preparation_metrics(
|
||||
response: &mut CodexTransportMetrics,
|
||||
preparation: CodexTransportMetrics,
|
||||
) {
|
||||
response.decision = preparation.decision;
|
||||
response.ws_connect_ms = preparation.ws_connect_ms;
|
||||
response.transport_decision_wait_ms = preparation.transport_decision_wait_ms;
|
||||
}
|
||||
|
||||
fn elapsed_duration_millis(duration: Duration) -> i64 {
|
||||
i64::try_from(duration.as_millis())
|
||||
.unwrap_or(i64::MAX)
|
||||
.max(1)
|
||||
}
|
||||
|
||||
fn http_version_name(version: reqwest::Version) -> &'static str {
|
||||
match version {
|
||||
reqwest::Version::HTTP_09 => "HTTP/0.9",
|
||||
reqwest::Version::HTTP_10 => "HTTP/1.0",
|
||||
reqwest::Version::HTTP_11 => "HTTP/1.1",
|
||||
reqwest::Version::HTTP_2 => "HTTP/2",
|
||||
reqwest::Version::HTTP_3 => "HTTP/3",
|
||||
_ => "unknown",
|
||||
}
|
||||
}
|
||||
|
||||
fn websocket_origin_key(base_url: &str) -> String {
|
||||
let origin = reqwest::Url::parse(base_url)
|
||||
.ok()
|
||||
.and_then(|url| {
|
||||
let host = url.host_str()?;
|
||||
Some(format!(
|
||||
"{}://{}:{}",
|
||||
url.scheme(),
|
||||
host,
|
||||
url.port_or_known_default().unwrap_or_default()
|
||||
))
|
||||
})
|
||||
.unwrap_or_else(|| base_url.trim_end_matches('/').to_string());
|
||||
match custom_ca_env_cache_key() {
|
||||
Some(tls_profile) => format!("{origin}\0{tls_profile}"),
|
||||
None => origin,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,12 +7,16 @@ impl CodexBackendClient {
|
||||
base_url: impl Into<String>,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
) -> Self {
|
||||
let base_url = base_url.into().trim_end_matches('/').to_string();
|
||||
Self {
|
||||
client,
|
||||
base_url: base_url.into().trim_end_matches('/').to_string(),
|
||||
websocket_origin_key: websocket_origin_key(&base_url),
|
||||
base_url,
|
||||
fingerprint,
|
||||
websocket_pool: None,
|
||||
websocket_initial_event_timeout: Some(DEFAULT_INITIAL_EVENT_TIMEOUT),
|
||||
websocket_fast_path_budget: WEBSOCKET_FAST_PATH_BUDGET,
|
||||
websocket_origin_breaker: WebSocketOriginBreaker::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -21,6 +25,16 @@ impl CodexBackendClient {
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_websocket_fast_path_budget(mut self, budget: Duration) -> Self {
|
||||
self.websocket_fast_path_budget = budget.max(Duration::from_millis(1));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_websocket_origin_breaker(mut self, breaker: WebSocketOriginBreaker) -> Self {
|
||||
self.websocket_origin_breaker = breaker;
|
||||
self
|
||||
}
|
||||
|
||||
/// 为 Responses WebSocket 请求启用连接池。
|
||||
pub fn with_websocket_pool(mut self, pool: Arc<CodexWebSocketPool>) -> Self {
|
||||
self.websocket_pool = Some(pool);
|
||||
@@ -43,6 +57,7 @@ impl CodexBackendClient {
|
||||
started_at: Instant,
|
||||
) -> CodexClientResult<CodexBackendResponse> {
|
||||
let headers = self.request_headers_for_http_response(upstream_request, context)?;
|
||||
let headers_started_at = Instant::now();
|
||||
let response = self
|
||||
.client
|
||||
.post(endpoint_url(&self.base_url, CODEX_RESPONSES_PATH))
|
||||
@@ -50,6 +65,8 @@ impl CodexBackendClient {
|
||||
.json(&upstream_request)
|
||||
.send()
|
||||
.await?;
|
||||
let upstream_headers_ms = elapsed_duration_millis(headers_started_at.elapsed());
|
||||
let http_version = http_version_name(response.version()).to_string();
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let turn_state = response_meta::turn_state(response.headers());
|
||||
@@ -67,13 +84,16 @@ impl CodexBackendClient {
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers,
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
});
|
||||
}
|
||||
|
||||
let mut body_bytes = Vec::new();
|
||||
let mut first_token_ms = None;
|
||||
let mut first_event_ms = None;
|
||||
let mut stream = http_sse_stream(response);
|
||||
while let Some(chunk) = stream.try_next().await? {
|
||||
first_event_ms.get_or_insert_with(|| elapsed_duration_millis(started_at.elapsed()));
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
response_meta::update_first_token_ms(started_at, &body_bytes, &mut first_token_ms);
|
||||
}
|
||||
@@ -90,6 +110,13 @@ impl CodexBackendClient {
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
response_metadata,
|
||||
transport_metrics: CodexTransportMetrics {
|
||||
upstream_headers_ms: Some(upstream_headers_ms),
|
||||
first_event_ms,
|
||||
http_version: Some(http_version),
|
||||
..CodexTransportMetrics::default()
|
||||
},
|
||||
connection_local_continuation: false,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -100,6 +127,7 @@ impl CodexBackendClient {
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexClientResult<CodexBackendStreamingResponse> {
|
||||
let headers = self.request_headers_for_http_response(upstream_request, context)?;
|
||||
let headers_started_at = Instant::now();
|
||||
let response = self
|
||||
.client
|
||||
.post(endpoint_url(&self.base_url, CODEX_RESPONSES_PATH))
|
||||
@@ -107,6 +135,8 @@ impl CodexBackendClient {
|
||||
.json(&upstream_request)
|
||||
.send()
|
||||
.await?;
|
||||
let upstream_headers_ms = elapsed_duration_millis(headers_started_at.elapsed());
|
||||
let http_version = http_version_name(response.version()).to_string();
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let turn_state = response_meta::turn_state(response.headers());
|
||||
@@ -124,6 +154,7 @@ impl CodexBackendClient {
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers,
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -138,6 +169,12 @@ impl CodexBackendClient {
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
response_metadata,
|
||||
transport_metrics: CodexTransportMetrics {
|
||||
upstream_headers_ms: Some(upstream_headers_ms),
|
||||
http_version: Some(http_version),
|
||||
..CodexTransportMetrics::default()
|
||||
},
|
||||
connection_local_continuation: false,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -167,43 +204,11 @@ impl CodexBackendClient {
|
||||
pool_account_id: Option<&str>,
|
||||
started_at: Instant,
|
||||
) -> CodexClientResult<CodexBackendResponse> {
|
||||
let upstream_request = response_upstream_request(request, context);
|
||||
match transport_for_request(&upstream_request) {
|
||||
CodexTransport::HttpSse => {
|
||||
self.create_response_http_sse(&upstream_request, context, started_at)
|
||||
.await
|
||||
}
|
||||
CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => {
|
||||
match self
|
||||
.create_response_websocket(
|
||||
&upstream_request,
|
||||
context,
|
||||
pool_account_id,
|
||||
started_at,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => Ok(response),
|
||||
Err(error)
|
||||
if http_sse_fallback_allowed(&upstream_request)
|
||||
&& websocket_error_allows_http_fallback(&error) =>
|
||||
{
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = pool_account_id.or(context.account_id).unwrap_or_default(),
|
||||
transport = "websocket",
|
||||
fallback_transport = "http_sse",
|
||||
fallback_reason = "websocket_error",
|
||||
error = %error,
|
||||
"websocket response failed; falling back to HTTP SSE"
|
||||
);
|
||||
self.create_response_http_sse(&upstream_request, context, started_at)
|
||||
.await
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
let prepared = self
|
||||
.prepare_response_transport_with_pool_account(request, context, pool_account_id)
|
||||
.await?;
|
||||
self.create_response_with_prepared(request, context, prepared, started_at)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 发送 Responses SSE 请求并返回 live SSE 流。
|
||||
@@ -222,50 +227,35 @@ impl CodexBackendClient {
|
||||
context: CodexRequestContext<'_>,
|
||||
pool_account_id: Option<&str>,
|
||||
) -> CodexClientResult<CodexBackendStreamingResponse> {
|
||||
let upstream_request = response_upstream_request(request, context);
|
||||
match transport_for_request(&upstream_request) {
|
||||
CodexTransport::HttpSse => {
|
||||
self.create_response_stream_http_sse(&upstream_request, context)
|
||||
.await
|
||||
}
|
||||
CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => {
|
||||
match self
|
||||
.create_response_websocket_stream(&upstream_request, context, pool_account_id)
|
||||
.await
|
||||
{
|
||||
Ok(response) => Ok(response),
|
||||
Err(error)
|
||||
if http_sse_fallback_allowed(&upstream_request)
|
||||
&& websocket_error_allows_http_fallback(&error) =>
|
||||
{
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = pool_account_id.or(context.account_id).unwrap_or_default(),
|
||||
transport = "websocket",
|
||||
fallback_transport = "http_sse",
|
||||
fallback_reason = "websocket_error",
|
||||
error = %error,
|
||||
"websocket response stream failed; falling back to HTTP SSE"
|
||||
);
|
||||
self.create_response_stream_http_sse(&upstream_request, context)
|
||||
.await
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
let prepared = self
|
||||
.prepare_response_transport_with_pool_account(request, context, pool_account_id)
|
||||
.await?;
|
||||
self.create_response_stream_with_prepared(request, context, prepared)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_response_websocket(
|
||||
/// 在发送 payload 前完成 transport 选择和可取消的 WebSocket opening。
|
||||
pub(crate) async fn prepare_response_transport_with_pool_account(
|
||||
&self,
|
||||
upstream_request: &CodexResponsesRequest,
|
||||
request: &CodexResponsesRequest,
|
||||
context: CodexRequestContext<'_>,
|
||||
pool_account_id: Option<&str>,
|
||||
started_at: Instant,
|
||||
) -> CodexClientResult<CodexBackendResponse> {
|
||||
let websocket_request = websocket_upstream_request(upstream_request);
|
||||
) -> CodexClientResult<PreparedResponseTransport> {
|
||||
let requirement = transport_requirement(request);
|
||||
if requirement == TransportRequirement::HttpRequired {
|
||||
return Ok(PreparedResponseTransport {
|
||||
requirement,
|
||||
route: PreparedResponseRoute::Http,
|
||||
metrics: CodexTransportMetrics {
|
||||
decision: Some(CodexTransportDecision::HttpRequired),
|
||||
..CodexTransportMetrics::default()
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
let websocket_request = websocket_upstream_request(request);
|
||||
let headers = self.request_headers_for_http_response(&websocket_request, context)?;
|
||||
let prepared = CodexWebSocketConnection::responses_create_request(
|
||||
let websocket_create = CodexWebSocketConnection::responses_create_request(
|
||||
&self.base_url,
|
||||
&generate_key(),
|
||||
websocket_header_pairs(&headers),
|
||||
@@ -274,129 +264,220 @@ impl CodexBackendClient {
|
||||
.map_err(CodexClientError::WebSocketEncode)?;
|
||||
let artifact = websocket_audit_artifact_from_attempt(
|
||||
&websocket_request,
|
||||
prepared.connection().opening_audit_snapshot(),
|
||||
websocket_create.connection().opening_audit_snapshot(),
|
||||
websocket_payload_audit_snapshot(&websocket_request),
|
||||
);
|
||||
if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await {
|
||||
tracing::warn!(error = %error, "failed to write Codex WebSocket audit artifact");
|
||||
tracing::warn!(error = %error, "Failed to write Codex WebSocket audit artifact");
|
||||
}
|
||||
let pool_key = self.websocket_pool_key(upstream_request, context, pool_account_id);
|
||||
let pool_key = self.websocket_pool_key(request, context, pool_account_id);
|
||||
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
|
||||
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
|
||||
(Some(pool), Some(key)) => {
|
||||
execute_response_create_request_with_pool(
|
||||
&prepared,
|
||||
Some((pool, key)),
|
||||
started_at,
|
||||
self.websocket_initial_event_timeout,
|
||||
)
|
||||
.await
|
||||
let pool = self.websocket_pool.as_deref().zip(pool_key);
|
||||
let fast_path_budget = match requirement {
|
||||
TransportRequirement::PersistedContinuation
|
||||
| TransportRequirement::ExternalUnknown
|
||||
| TransportRequirement::NewChain => Some(self.websocket_fast_path_budget),
|
||||
TransportRequirement::ExplicitWebSocketWarmup
|
||||
| TransportRequirement::ExactWebSocketContinuation => None,
|
||||
TransportRequirement::HttpRequired => None,
|
||||
};
|
||||
let prepare_started_at = Instant::now();
|
||||
let prepared = prepare_response_create_request_with_pool(
|
||||
&websocket_create,
|
||||
pool,
|
||||
&self.websocket_origin_breaker,
|
||||
&self.websocket_origin_key,
|
||||
fast_path_budget,
|
||||
requirement.requires_websocket(),
|
||||
self.websocket_initial_event_timeout,
|
||||
)
|
||||
.await;
|
||||
let prepared = match prepared {
|
||||
Ok(prepared) => prepared,
|
||||
Err(error)
|
||||
if requirement.allows_pre_send_http_fallback()
|
||||
&& error.allows_pre_send_http_fallback() =>
|
||||
{
|
||||
let decision = http_fallback_decision(&error);
|
||||
let wait_ms = elapsed_duration_millis(prepare_started_at.elapsed());
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = pool_account_id.or(context.account_id).unwrap_or_default(),
|
||||
transport_requirement = requirement.as_str(),
|
||||
transport_decision = decision.as_str(),
|
||||
transport_decision_wait_ms = wait_ms,
|
||||
error = %error,
|
||||
"WebSocket preparation failed before payload send; using same-account HTTP"
|
||||
);
|
||||
return Ok(PreparedResponseTransport {
|
||||
requirement,
|
||||
route: PreparedResponseRoute::Http,
|
||||
metrics: CodexTransportMetrics {
|
||||
decision: Some(decision),
|
||||
ws_connect_ms: matches!(
|
||||
error,
|
||||
CodexWebSocketExchangeError::FastPathTimeout { .. }
|
||||
)
|
||||
.then_some(wait_ms),
|
||||
transport_decision_wait_ms: Some(wait_ms),
|
||||
..CodexTransportMetrics::default()
|
||||
},
|
||||
});
|
||||
}
|
||||
_ => {
|
||||
execute_response_create_request_with_pool(
|
||||
&prepared,
|
||||
None,
|
||||
started_at,
|
||||
self.websocket_initial_event_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(websocket_exchange_error_to_client_error)?;
|
||||
Err(error) => return Err(websocket_exchange_error_to_client_error(error)),
|
||||
};
|
||||
let decision = websocket_success_decision(requirement, &prepared);
|
||||
let metrics = CodexTransportMetrics {
|
||||
decision: Some(decision),
|
||||
ws_connect_ms: prepared.connect_elapsed().map(elapsed_duration_millis),
|
||||
transport_decision_wait_ms: Some(elapsed_duration_millis(
|
||||
prepared.decision_wait_elapsed(),
|
||||
)),
|
||||
upstream_headers_ms: prepared.connect_elapsed().map(elapsed_duration_millis),
|
||||
first_event_ms: None,
|
||||
http_version: Some("HTTP/1.1".to_string()),
|
||||
};
|
||||
log_websocket_pool_decision(
|
||||
context.request_id,
|
||||
pool_account_id.or(context.account_id),
|
||||
pool_log_context.as_ref(),
|
||||
exchange.pool_decision,
|
||||
prepared.pool_decision(),
|
||||
);
|
||||
|
||||
Ok(CodexBackendResponse {
|
||||
body: exchange.body,
|
||||
transport: CodexBackendTransport::WebSocket,
|
||||
usage: exchange.usage,
|
||||
turn_state: exchange.turn_state,
|
||||
set_cookie_headers: exchange.set_cookie_headers,
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
first_token_ms: exchange.first_token_ms,
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
response_metadata: exchange.response_metadata,
|
||||
})
|
||||
}
|
||||
|
||||
async fn create_response_websocket_stream(
|
||||
&self,
|
||||
upstream_request: &CodexResponsesRequest,
|
||||
context: CodexRequestContext<'_>,
|
||||
pool_account_id: Option<&str>,
|
||||
) -> CodexClientResult<CodexBackendStreamingResponse> {
|
||||
let websocket_request = websocket_upstream_request(upstream_request);
|
||||
let headers = self.request_headers_for_http_response(&websocket_request, context)?;
|
||||
let prepared = CodexWebSocketConnection::responses_create_request(
|
||||
&self.base_url,
|
||||
&generate_key(),
|
||||
websocket_header_pairs(&headers),
|
||||
&websocket_request,
|
||||
)
|
||||
.map_err(CodexClientError::WebSocketEncode)?;
|
||||
let artifact = websocket_audit_artifact_from_attempt(
|
||||
&websocket_request,
|
||||
prepared.connection().opening_audit_snapshot(),
|
||||
websocket_payload_audit_snapshot(&websocket_request),
|
||||
);
|
||||
if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await {
|
||||
tracing::warn!(error = %error, "failed to write Codex WebSocket audit artifact");
|
||||
}
|
||||
let pool_key = self.websocket_pool_key(upstream_request, context, pool_account_id);
|
||||
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
|
||||
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
|
||||
(Some(pool), Some(key)) => {
|
||||
execute_response_create_request_stream_with_pool(
|
||||
&prepared,
|
||||
Some((pool, key)),
|
||||
self.websocket_initial_event_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => {
|
||||
execute_response_create_request_stream_with_pool(
|
||||
&prepared,
|
||||
None,
|
||||
self.websocket_initial_event_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(websocket_exchange_error_to_client_error)?;
|
||||
tracing::info!(
|
||||
request_id = %context.request_id,
|
||||
account_id = pool_account_id.or(context.account_id).unwrap_or_default(),
|
||||
websocket_connection_id = %exchange.websocket_connection_id,
|
||||
ws_pool = exchange.pool_decision.map_or("unpooled", WebSocketPoolDecision::kind),
|
||||
"websocket response stream established"
|
||||
);
|
||||
log_websocket_pool_decision(
|
||||
context.request_id,
|
||||
pool_account_id.or(context.account_id),
|
||||
pool_log_context.as_ref(),
|
||||
exchange.pool_decision,
|
||||
transport_requirement = requirement.as_str(),
|
||||
transport_decision = decision.as_str(),
|
||||
ws_connect_ms = ?metrics.ws_connect_ms,
|
||||
transport_decision_wait_ms = ?metrics.transport_decision_wait_ms,
|
||||
"Responses transport prepared"
|
||||
);
|
||||
Ok(PreparedResponseTransport {
|
||||
requirement,
|
||||
route: PreparedResponseRoute::WebSocket(Box::new(PreparedWebSocketRoute {
|
||||
request: websocket_create,
|
||||
prepared,
|
||||
})),
|
||||
metrics,
|
||||
})
|
||||
}
|
||||
|
||||
Ok(CodexBackendStreamingResponse {
|
||||
body: Box::pin(
|
||||
exchange
|
||||
.body
|
||||
.map_err(websocket_exchange_error_to_client_error),
|
||||
),
|
||||
transport: CodexBackendTransport::WebSocket,
|
||||
turn_state: exchange.turn_state,
|
||||
set_cookie_headers: exchange.set_cookie_headers,
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
|
||||
turn_state_update: Some(exchange.turn_state_update),
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
response_metadata: exchange.response_metadata,
|
||||
pub(crate) async fn create_response_with_prepared(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
context: CodexRequestContext<'_>,
|
||||
prepared: PreparedResponseTransport,
|
||||
started_at: Instant,
|
||||
) -> CodexClientResult<CodexBackendResponse> {
|
||||
let PreparedResponseTransport {
|
||||
requirement,
|
||||
route,
|
||||
metrics,
|
||||
} = prepared;
|
||||
let result = match route {
|
||||
PreparedResponseRoute::Http => self
|
||||
.create_response_http_sse(request, context, started_at)
|
||||
.await
|
||||
.map(|mut response| {
|
||||
merge_preparation_metrics(&mut response.transport_metrics, metrics);
|
||||
response
|
||||
}),
|
||||
PreparedResponseRoute::WebSocket(route) => {
|
||||
let PreparedWebSocketRoute { request, prepared } = *route;
|
||||
execute_prepared_response_create_request(&request, prepared, started_at)
|
||||
.await
|
||||
.map_err(websocket_exchange_error_to_client_error)
|
||||
.map(|exchange| CodexBackendResponse {
|
||||
body: exchange.body,
|
||||
transport: CodexBackendTransport::WebSocket,
|
||||
usage: exchange.usage,
|
||||
turn_state: exchange.turn_state,
|
||||
set_cookie_headers: exchange.set_cookie_headers,
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
first_token_ms: exchange.first_token_ms,
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
response_metadata: exchange.response_metadata,
|
||||
transport_metrics: CodexTransportMetrics {
|
||||
first_event_ms: exchange.first_event_ms,
|
||||
..metrics
|
||||
},
|
||||
connection_local_continuation: exchange.connection_local_continuation,
|
||||
})
|
||||
}
|
||||
};
|
||||
result.inspect_err(|error| {
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
transport_requirement = requirement.as_str(),
|
||||
failure_phase = "post_send_or_explicit_response",
|
||||
error = %error,
|
||||
"Responses transport failed after preparation; automatic fallback is disabled"
|
||||
);
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn create_response_stream_with_prepared(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
context: CodexRequestContext<'_>,
|
||||
prepared: PreparedResponseTransport,
|
||||
) -> CodexClientResult<CodexBackendStreamingResponse> {
|
||||
let PreparedResponseTransport {
|
||||
requirement,
|
||||
route,
|
||||
metrics,
|
||||
} = prepared;
|
||||
let result = match route {
|
||||
PreparedResponseRoute::Http => self
|
||||
.create_response_stream_http_sse(request, context)
|
||||
.await
|
||||
.map(|mut response| {
|
||||
merge_preparation_metrics(&mut response.transport_metrics, metrics);
|
||||
response
|
||||
}),
|
||||
PreparedResponseRoute::WebSocket(route) => {
|
||||
let PreparedWebSocketRoute { request, prepared } = *route;
|
||||
execute_prepared_response_create_request_stream(&request, prepared)
|
||||
.await
|
||||
.map_err(websocket_exchange_error_to_client_error)
|
||||
.map(|exchange| {
|
||||
tracing::info!(
|
||||
request_id = %context.request_id,
|
||||
websocket_connection_id = %exchange.websocket_connection_id,
|
||||
ws_pool = exchange.pool_decision.map_or("unpooled", WebSocketPoolDecision::kind),
|
||||
"WebSocket response stream established"
|
||||
);
|
||||
CodexBackendStreamingResponse {
|
||||
body: Box::pin(
|
||||
exchange
|
||||
.body
|
||||
.map_err(post_send_ambiguous)
|
||||
.map_err(websocket_exchange_error_to_client_error),
|
||||
),
|
||||
transport: CodexBackendTransport::WebSocket,
|
||||
turn_state: exchange.turn_state,
|
||||
set_cookie_headers: exchange.set_cookie_headers,
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
|
||||
turn_state_update: Some(exchange.turn_state_update),
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
response_metadata: exchange.response_metadata,
|
||||
transport_metrics: metrics,
|
||||
connection_local_continuation: exchange.connection_local_continuation,
|
||||
}
|
||||
})
|
||||
}
|
||||
};
|
||||
result.inspect_err(|error| {
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
transport_requirement = requirement.as_str(),
|
||||
failure_phase = "post_send_or_explicit_response",
|
||||
error = %error,
|
||||
"Responses stream transport failed after preparation; automatic fallback is disabled"
|
||||
);
|
||||
})
|
||||
}
|
||||
|
||||
@@ -446,6 +527,7 @@ impl CodexBackendClient {
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers,
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
});
|
||||
}
|
||||
let parsed =
|
||||
@@ -460,6 +542,7 @@ impl CodexBackendClient {
|
||||
StatusCode::BAD_GATEWAY.as_u16(),
|
||||
),
|
||||
set_cookie_headers,
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
})?;
|
||||
let models = extract_model_entries(&parsed);
|
||||
if !models.is_empty() {
|
||||
@@ -472,6 +555,7 @@ impl CodexBackendClient {
|
||||
body: "model catalog response contains no models".to_string(),
|
||||
diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()),
|
||||
set_cookie_headers: Vec::new(),
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user