mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
28f2879a05 | ||
|
|
2474259c22 | ||
|
|
79c2c5948f | ||
|
|
3c4bd8757f | ||
|
|
06f3443e56 | ||
|
|
88fd21f169 | ||
|
|
521201c926 | ||
|
|
be2123ea30 | ||
|
|
4b6d369225 | ||
|
|
76b81a6ceb | ||
|
|
29c8b679ba | ||
|
|
4091a0c2d8 | ||
|
|
f0ccc981fb | ||
|
|
d01ea7f775 | ||
|
|
fd749ea8f8 | ||
|
|
f0543b637e | ||
|
|
549bb9f62b | ||
|
|
9e57f78bab | ||
|
|
73c0ecdc4c | ||
|
|
bfd05722d5 | ||
|
|
9dc2ea4a2f | ||
|
|
6fb27e869e | ||
|
|
81c58b1cdc | ||
|
|
c8dc2e293a | ||
|
|
25081de913 | ||
|
|
6eb0917ff1 | ||
|
|
72d892d5e3 | ||
|
|
0b2da3087b | ||
|
|
5391fdb31c | ||
|
|
24036f1f84 | ||
|
|
b043fced0c |
@@ -12,12 +12,23 @@ on:
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
run_backend:
|
||||
description: Run backend quality checks
|
||||
required: false
|
||||
default: true
|
||||
type: boolean
|
||||
run_frontend:
|
||||
description: Run frontend quality checks
|
||||
required: false
|
||||
default: true
|
||||
type: boolean
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
backend:
|
||||
if: inputs.run_backend
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
@@ -77,6 +88,7 @@ jobs:
|
||||
run: cargo test --test main --locked
|
||||
|
||||
frontend:
|
||||
if: inputs.run_frontend
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
@@ -93,7 +105,7 @@ jobs:
|
||||
version: 11.7.0
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
|
||||
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
|
||||
@@ -20,6 +20,8 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
outputs:
|
||||
backend: ${{ steps.filter.outputs.backend }}
|
||||
frontend: ${{ steps.filter.outputs.frontend }}
|
||||
container: ${{ steps.filter.outputs.container }}
|
||||
steps:
|
||||
- name: Checkout
|
||||
@@ -28,11 +30,18 @@ jobs:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Detect container changes
|
||||
- name: Detect affected checks
|
||||
id: filter
|
||||
uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2
|
||||
with:
|
||||
filters: |
|
||||
backend:
|
||||
- 'backend/**'
|
||||
- 'release/version.yaml'
|
||||
- '.github/workflows/**'
|
||||
frontend:
|
||||
- 'frontend/**'
|
||||
- '.github/workflows/**'
|
||||
container:
|
||||
- 'backend/**'
|
||||
- 'frontend/**'
|
||||
@@ -43,12 +52,15 @@ jobs:
|
||||
- '.github/workflows/ci.yml'
|
||||
|
||||
quality:
|
||||
needs: changes
|
||||
uses: ./.github/workflows/_quality.yml
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
run_backend: ${{ needs.changes.outputs.backend == 'true' }}
|
||||
run_frontend: ${{ needs.changes.outputs.frontend == 'true' }}
|
||||
|
||||
container:
|
||||
needs: [changes, quality]
|
||||
needs: changes
|
||||
if: needs.changes.outputs.container == 'true'
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
@@ -165,7 +165,7 @@ jobs:
|
||||
|
||||
build-docker-binary:
|
||||
name: build-docker-binary (${{ matrix.os }}-${{ matrix.arch }})
|
||||
needs: [release-info, quality]
|
||||
needs: release-info
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix: ${{ fromJson(needs.release-info.outputs.docker_matrix) }}
|
||||
@@ -225,7 +225,7 @@ jobs:
|
||||
|
||||
build-binary:
|
||||
name: build-binary (${{ matrix.os }}-${{ matrix.arch }})
|
||||
needs: [release-info, quality]
|
||||
needs: release-info
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix: ${{ fromJson(needs.release-info.outputs.non_docker_matrix) }}
|
||||
@@ -621,7 +621,7 @@ jobs:
|
||||
cosign sign --yes "${images[@]}"
|
||||
|
||||
- name: Publish GitHub Release
|
||||
uses: softprops/action-gh-release@718ea10b132b3b2eba29c1007bb80653f286566b # v3.0.1
|
||||
uses: softprops/action-gh-release@3d0d9888cb7fd7b750713d6e236d1fcb99157228 # v3.0.2
|
||||
with:
|
||||
tag_name: ${{ needs.release-info.outputs.tag }}
|
||||
generate_release_notes: true
|
||||
|
||||
@@ -57,7 +57,7 @@ jobs:
|
||||
version: 11.7.0
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
|
||||
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
|
||||
@@ -1,41 +1,41 @@
|
||||
<!-- prettier-ignore -->
|
||||
<div align="center">
|
||||
|
||||
<img src="frontend/public/favicon.svg" alt="Codex Proxy RS" width="88" height="88" />
|
||||
<img src="frontend/public/favicon.svg" alt="Codex Proxy RS" width="80" height="80" />
|
||||
|
||||
# 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)
|
||||
[](#开源许可)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](https://opensource.org/license/mit)
|
||||
|
||||
[快速开始](#快速开始) · [首次使用](#首次使用) · [客户端接入](#客户端接入) · [常见问题](#常见问题) · [部署文档](deploy/README.md)
|
||||
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运维](#运维) · [部署文档](deploy/README.md) · [架构](docs/architecture.md)
|
||||
|
||||
</div>
|
||||
|
||||
> [!NOTE]
|
||||
> 当前兼容 OpenAI Responses API,不提供 `/v1/chat/completions`。
|
||||
> 只支持 OpenAI Responses API,不提供 `/v1/chat/completions`。
|
||||
|
||||
## 功能
|
||||
## 能力
|
||||
|
||||
- 支持 Responses HTTP、SSE 和官方 WebSocket 请求。
|
||||
- 提供智能调度、额度重置优先、轮询和粘滞四种账号策略。
|
||||
- 某个账号遇到额度、认证、模型或传输错误时,继续尝试其他账号。
|
||||
- 在网页中管理账号、客户端 Key、模型别名和运行参数。
|
||||
- 按请求、模型、账号或客户端 Key 查看 Token、费用、缓存、TTFT、延迟和错误。
|
||||
- 支持账号 OAuth、Token 导入以及 CPR、Sub2API、CLIProxyAPI JSON 导入。
|
||||
- 可直接生成 Codex CLI 配置,或通过 deeplink 导入 CCSwitch。
|
||||
| 领域 | 实现 |
|
||||
| -------- | ------------------------------------------------ |
|
||||
| 协议 | Responses HTTP、SSE、官方 WebSocket |
|
||||
| 传输 | 会话级 WebSocket 池、共享 HTTP/2 后备 |
|
||||
| 调度 | 智能调度、额度优先、轮询、粘滞 |
|
||||
| 故障处理 | 额度、认证、封禁、模型与传输故障切换 |
|
||||
| 身份 | 账号身份隔离、稳定 installation ID、会话拓扑透传 |
|
||||
| 管理 | 账号、客户端 Key、模型别名、运行设置与在线更新 |
|
||||
| 遥测 | Token、费用、缓存、TTFT、延迟与错误 |
|
||||
|
||||
## 快速开始
|
||||
|
||||
Docker Compose 是最简单的启动方式。开始前请安装 Docker Engine 和 Docker Compose Plugin。
|
||||
需要 Docker Engine 与 Docker Compose Plugin。
|
||||
|
||||
### 1. 准备配置
|
||||
### 1. 准备
|
||||
|
||||
```bash
|
||||
git clone https://github.com/zyycn/codex-proxy-rs.git
|
||||
@@ -48,74 +48,63 @@ sudo chown "$(id -u):10001" deploy/config.yaml
|
||||
chmod 0640 deploy/config.yaml
|
||||
```
|
||||
|
||||
分别执行三次以下命令生成密码:
|
||||
分别执行三次:
|
||||
|
||||
```bash
|
||||
openssl rand -hex 24
|
||||
```
|
||||
|
||||
把结果填入 `deploy/config.yaml` 的 `database.password`、`redis.password` 和
|
||||
`admin.default_password`。PostgreSQL 与 Redis 密码必须是 48 位十六进制字符;管理员密码
|
||||
至少 12 位且不能包含 `$`。
|
||||
将结果写入 `deploy/config.yaml`:
|
||||
|
||||
> [!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 通常也不需要调整数据目录权限。
|
||||
| 配置 | 约束 |
|
||||
| ------------------------------ | ------------------------ |
|
||||
| `x-cpr.database.password` | 48 位十六进制 |
|
||||
| `x-cpr.redis.password` | 48 位十六进制 |
|
||||
| `x-cpr.admin.default_password` | 至少 12 位,不能包含 `$` |
|
||||
|
||||
### 2. 启动服务
|
||||
Linux 需要允许容器组写入运行目录:
|
||||
|
||||
```bash
|
||||
sudo chown -R "$(id -u):10001" .runtime/data .runtime/logs
|
||||
chmod 0770 .runtime/data .runtime/logs
|
||||
```
|
||||
|
||||
### 2. 启动
|
||||
|
||||
```bash
|
||||
docker compose -f deploy/compose.yaml config --quiet
|
||||
docker compose -f deploy/compose.yaml pull
|
||||
docker compose -f deploy/compose.yaml up -d --no-build
|
||||
```
|
||||
|
||||
启动后检查:
|
||||
|
||||
```bash
|
||||
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>。
|
||||
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址:
|
||||
`http://127.0.0.1:8080`。
|
||||
|
||||
## 首次使用
|
||||
### 3. 初始化
|
||||
|
||||
1. 使用 `admin@cpr.local` 和 `x-cpr.admin.default_password` 中设置的密码登录。
|
||||
2. 打开“账号”,用 OAuth、Access/Refresh Token 或 JSON 文件添加账号。
|
||||
3. 点一次连接测试,能看到模型和额度即可。
|
||||
4. 打开“API 密钥”,创建一个 `sk_...` 客户端 Key。
|
||||
5. 点击“使用密钥”复制 Codex CLI 配置,或一键导入 CCSwitch。
|
||||
1. 使用 `admin@cpr.local` 与初始密码登录。
|
||||
2. 添加 Codex 账号并执行连接测试。
|
||||
3. 创建 `sk_...` 客户端 Key。
|
||||
4. 复制 Codex CLI 配置,或导入 CCSwitch。
|
||||
|
||||
> [!TIP]
|
||||
> `admin.default_password` 只用于首次创建管理员。已有管理员后修改该字段不会重置登录密码。
|
||||
> [!NOTE]
|
||||
> `admin.default_password` 只在首次创建管理员时生效。
|
||||
|
||||
完整部署、权限和密码轮换规则见 [deploy/README.md](deploy/README.md)。
|
||||
|
||||
## 客户端接入
|
||||
|
||||
| 配置 | 值 |
|
||||
| -------- | ---------------------------------------- |
|
||||
| Base URL | `http://127.0.0.1:8080/v1` |
|
||||
| API Key | 管理端创建的 `sk_...` 客户端 Key |
|
||||
| API Key | 管理端创建的 `sk_...` |
|
||||
| 鉴权 | `Authorization: Bearer <client-api-key>` |
|
||||
|
||||
先查询当前可用模型:
|
||||
|
||||
```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 \
|
||||
-H 'Authorization: Bearer <client-api-key>' \
|
||||
-H 'Content-Type: application/json' \
|
||||
@@ -127,89 +116,45 @@ curl http://127.0.0.1:8080/v1/responses \
|
||||
```
|
||||
|
||||
<details>
|
||||
<summary>查看全部客户端 API 路由</summary>
|
||||
<summary>API 路由</summary>
|
||||
|
||||
| 路由 | 说明 |
|
||||
| -------------------------------- | ------------------------------- |
|
||||
| `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` | 模型运行信息 |
|
||||
| 路由 | 用途 |
|
||||
| --------------------------- | ------------------- |
|
||||
| `POST /v1/responses` | JSON 或 SSE |
|
||||
| `GET /v1/responses` | Responses WebSocket |
|
||||
| `POST /v1/responses/review` | Review |
|
||||
| `GET /v1/models` | 模型列表 |
|
||||
| `GET /v1/models/catalog` | 模型目录 |
|
||||
| `GET /v1/models/{model_id}` | 模型详情 |
|
||||
|
||||
所有 `/v1/*` 路由都需要客户端 API Key。
|
||||
|
||||
</details>
|
||||
|
||||
## 管理端
|
||||
|
||||
| 页面 | 内容 |
|
||||
| -------- | ---------------------------------------------------------------------- |
|
||||
| 概览 | 查看服务状态、请求趋势、Token、账号池容量、请求健康和最近调用 |
|
||||
| 账号 | 添加、导入、导出、测试、启停账号,以及刷新 Token 和额度 |
|
||||
| API 密钥 | 创建、复制、启停或删除 Key,生成 Codex CLI 配置并导入 CCSwitch |
|
||||
| 使用统计 | 按时间、模型、账号和客户端 Key 查看用量、费用、缓存、延迟和错误 |
|
||||
| 系统设置 | 配置模型别名、调度策略、单账号并发、请求间隔、刷新参数和管理员 API Key |
|
||||
| 系统更新 | 查看版本和发布说明;Release 版本可以在线更新并重启 |
|
||||
|
||||
## 常用操作
|
||||
|
||||
### 升级
|
||||
## 运维
|
||||
|
||||
```bash
|
||||
# 升级
|
||||
docker compose -f deploy/compose.yaml pull codex-proxy-rs
|
||||
docker compose -f deploy/compose.yaml up -d --no-build
|
||||
```
|
||||
|
||||
### 查看日志
|
||||
|
||||
```bash
|
||||
# 日志
|
||||
docker compose -f deploy/compose.yaml logs -f codex-proxy-rs
|
||||
```
|
||||
|
||||
### 从当前源码构建
|
||||
|
||||
```bash
|
||||
# 从源码构建
|
||||
docker compose -f deploy/compose.yaml build codex-proxy-rs
|
||||
docker compose -f deploy/compose.yaml up -d
|
||||
```
|
||||
|
||||
> [!WARNING]
|
||||
> PostgreSQL、Redis、应用数据和文件日志都保存在仓库根目录 `.runtime/`。普通
|
||||
> `docker compose down` 不会删除这些绑定目录;删除 `.runtime/` 才会清除数据。升级前请备份
|
||||
> 整个目录,数据库备份中包含客户端凭据,应按敏感数据保管。
|
||||
> [!IMPORTANT]
|
||||
> `.runtime/` 保存数据库、Redis、身份密钥和日志。删除该目录会永久清除运行状态。
|
||||
|
||||
## 常见问题
|
||||
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket
|
||||
upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
|
||||
|
||||
### `/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)
|
||||
- [部署](deploy/README.md)
|
||||
- [架构](docs/architecture.md)
|
||||
- [Release](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
+17
-6
@@ -371,6 +371,7 @@ dependencies = [
|
||||
"rand_core 0.6.4",
|
||||
"redis",
|
||||
"reqwest",
|
||||
"roxmltree",
|
||||
"rustls",
|
||||
"rustls-native-certs",
|
||||
"rustls-pki-types",
|
||||
@@ -381,6 +382,7 @@ dependencies = [
|
||||
"sha2 0.11.0",
|
||||
"sqlx",
|
||||
"subtle",
|
||||
"syn",
|
||||
"tar",
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
@@ -1784,9 +1786,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.3.0"
|
||||
version = "1.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2fa6f8e4b491d7a8ef3a9550a4d71969bd0064f46e32b8dbbcc7fc60dad94fed"
|
||||
checksum = "3cb5358643f48330db5c78856982faf39c93ba50c60d682b43d0344a3b649769"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
@@ -1905,6 +1907,15 @@ dependencies = [
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "roxmltree"
|
||||
version = "0.21.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1964b10c76125c36f8afe190065a4bf9a87bf324842c05701330bba9f1cacbb"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "2.1.3"
|
||||
@@ -1926,9 +1937,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.23.36"
|
||||
version = "0.23.42"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b"
|
||||
checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138"
|
||||
dependencies = [
|
||||
"once_cell",
|
||||
"ring",
|
||||
@@ -2882,9 +2893,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.23.4"
|
||||
version = "1.23.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53"
|
||||
checksum = "ea5fab0d6c3c01ae70085a09cb03d4c7a1d6314e2b3e075392783396d724ca0a"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"js-sys",
|
||||
|
||||
+3
-1
@@ -25,9 +25,10 @@ rand = "0.10.1"
|
||||
# Enables password_hash::rand_core::OsRng for argon2 0.5.
|
||||
rand_core = { version = "0.6.4", features = ["getrandom"] }
|
||||
redis = { version = "1.3.0", default-features = false, features = ["tokio-comp", "connection-manager", "script"] }
|
||||
roxmltree = "0.21.1"
|
||||
# Pinned to the Codex Desktop 26.707.41301 HTTP stack; re-audit a newer Desktop before upgrading.
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] }
|
||||
rustls = { version = "=0.23.36", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
rustls = { version = "=0.23.42", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
rustls-native-certs = "0.8.4"
|
||||
rustls-pki-types = "1.14.1"
|
||||
secrecy = { version = "0.10.3", features = ["serde"] }
|
||||
@@ -53,6 +54,7 @@ uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
|
||||
|
||||
[dev-dependencies]
|
||||
sqlx = { version = "0.9.0", default-features = false, features = ["derive"] }
|
||||
syn = { version = "2.0.118", features = ["full", "visit"] }
|
||||
tempfile = "3.27.0"
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["test-util"] }
|
||||
wiremock = "0.6.5"
|
||||
|
||||
@@ -10,9 +10,7 @@ use axum::{
|
||||
},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::Utc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
@@ -30,11 +28,7 @@ use crate::{
|
||||
manage::{
|
||||
AccountHealthCheck, AccountManageError, AccountRefreshOutcome, AccountRefreshResult,
|
||||
AccountUpdate, ManagedAccount, OAuthExchangeInput,
|
||||
quota_view::{
|
||||
AccountQuotaData, AccountQuotaUsageWindow, AccountQuotaWindowLocalUsage, quota_data,
|
||||
},
|
||||
},
|
||||
refresh::token_refresh_status_eligible,
|
||||
},
|
||||
infra::{
|
||||
format::{
|
||||
@@ -44,14 +38,9 @@ use crate::{
|
||||
json::{SortDirection, clamp_limit, clamp_page, total_pages},
|
||||
time::{china_datetime, china_relative_time, china_rfc3339},
|
||||
},
|
||||
telemetry::{
|
||||
account_usage::query::AccountUsageRecord,
|
||||
billing,
|
||||
buckets::query::{ModelUsageWindow, UsageBucketWindow},
|
||||
},
|
||||
telemetry::account_usage::query::AccountUsageRecord,
|
||||
};
|
||||
|
||||
const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200;
|
||||
const ACCOUNT_EXPORT_CONFIRMATION: &str = "export_sensitive_accounts";
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -195,21 +184,14 @@ impl AdminAccountData {
|
||||
struct AdminAccountUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
empty_response_count: i64,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
reasoning_tokens: i64,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
image_input_tokens: i64,
|
||||
image_output_tokens: i64,
|
||||
image_tokens_display: String,
|
||||
image_request_count: i64,
|
||||
image_request_failed_count: i64,
|
||||
created_tokens: i64,
|
||||
created_tokens_display: String,
|
||||
read_tokens: i64,
|
||||
@@ -225,16 +207,10 @@ impl AdminAccountUsageData {
|
||||
models: Vec<AdminAccountModelUsageData>,
|
||||
) -> Self {
|
||||
let request_count = usage.map_or(0, |usage| usage.window_request_count);
|
||||
let empty_response_count = 0;
|
||||
let input_tokens = usage.map_or(0, |usage| usage.window_input_tokens);
|
||||
let output_tokens = usage.map_or(0, |usage| usage.window_output_tokens);
|
||||
let cached_tokens = usage.map_or(0, |usage| usage.window_cached_tokens);
|
||||
let reasoning_tokens = 0;
|
||||
let total_tokens = input_tokens.saturating_add(output_tokens);
|
||||
let image_input_tokens = 0;
|
||||
let image_output_tokens = 0;
|
||||
let image_request_count = 0;
|
||||
let image_request_failed_count = 0;
|
||||
let created_tokens = input_tokens.saturating_sub(cached_tokens);
|
||||
let read_tokens = cached_tokens;
|
||||
let last_used_at = usage.and_then(|usage| usage.last_used_at);
|
||||
@@ -242,23 +218,14 @@ impl AdminAccountUsageData {
|
||||
Self {
|
||||
request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(request_count)),
|
||||
empty_response_count,
|
||||
input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(input_tokens)),
|
||||
output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(output_tokens)),
|
||||
cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(cached_tokens)),
|
||||
reasoning_tokens,
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
image_input_tokens,
|
||||
image_output_tokens,
|
||||
image_tokens_display: format_tokens(nonnegative_i64_to_u64(
|
||||
image_input_tokens + image_output_tokens,
|
||||
)),
|
||||
image_request_count,
|
||||
image_request_failed_count,
|
||||
created_tokens,
|
||||
created_tokens_display: format_tokens(nonnegative_i64_to_u64(created_tokens)),
|
||||
read_tokens,
|
||||
@@ -292,50 +259,6 @@ struct AdminAccountModelUsageData {
|
||||
last_used_at_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct AccountListStats {
|
||||
usage_by_account: HashMap<String, AccountUsageRecord>,
|
||||
quota_by_account: HashMap<String, AccountQuotaData>,
|
||||
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
|
||||
refreshing_account_ids: HashSet<String>,
|
||||
}
|
||||
|
||||
struct AccountModelUsageRecord {
|
||||
account_id: String,
|
||||
model: String,
|
||||
request_count: i64,
|
||||
error_count: i64,
|
||||
input_tokens: i64,
|
||||
output_tokens: i64,
|
||||
cached_tokens: i64,
|
||||
billing_amount_usd: f64,
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct AccountQuotaWindowSelection {
|
||||
account_id: String,
|
||||
window: AccountQuotaUsageWindow,
|
||||
}
|
||||
|
||||
impl AccountListStats {
|
||||
fn data_for(&self, account: ManagedAccount) -> AdminAccountData {
|
||||
let account_id = account.id.clone();
|
||||
let token_refreshing = token_refresh_status_eligible(account.status)
|
||||
&& self.refreshing_account_ids.contains(&account_id);
|
||||
AdminAccountData::from_parts(
|
||||
account,
|
||||
self.usage_by_account.get(&account_id),
|
||||
self.quota_by_account.get(&account_id).cloned(),
|
||||
self.models_by_account
|
||||
.get(&account_id)
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
token_refreshing,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminAccountSummaryData {
|
||||
@@ -345,6 +268,17 @@ struct AdminAccountSummaryData {
|
||||
attention: u64,
|
||||
}
|
||||
|
||||
impl From<AccountSummary> for AdminAccountSummaryData {
|
||||
fn from(summary: AccountSummary) -> Self {
|
||||
Self {
|
||||
total: summary.total,
|
||||
active: summary.active,
|
||||
quota_exhausted: summary.quota_exhausted,
|
||||
attention: summary.attention,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminAccountPageData {
|
||||
@@ -511,34 +445,33 @@ pub(crate) async fn accounts(
|
||||
let page_size = clamp_limit(params.page_size.unwrap_or(50));
|
||||
let status = account_status_filter(params.status)?;
|
||||
let sort = account_list_sort(params.sort_by, params.sort_direction)?;
|
||||
let quota_by_account = quota_snapshots_by_account(&state).await;
|
||||
let summary = account_summary_data(&state).await;
|
||||
|
||||
match state
|
||||
let result = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.list_page(page, page_size, params.search, status, sort)
|
||||
.account_list
|
||||
.query(AccountListQuery {
|
||||
page,
|
||||
page_size,
|
||||
search: params.search,
|
||||
status,
|
||||
sort,
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await?;
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|item| stats.data_for(item))
|
||||
.collect(),
|
||||
total: page.total,
|
||||
page: page.page,
|
||||
page_size: page.page_size,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminAccountPageEnvelope::ok(page, summary),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(account_error(&error)),
|
||||
}
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: result
|
||||
.page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(account_list_item_data)
|
||||
.collect(),
|
||||
total: result.page.total,
|
||||
page: result.page.page,
|
||||
page_size: result.page.page_size,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminAccountPageEnvelope::ok(page, result.summary.into()),
|
||||
))
|
||||
}
|
||||
|
||||
/// `POST /api/admin/accounts`
|
||||
@@ -566,7 +499,8 @@ mod import_routes;
|
||||
mod lifecycle_routes;
|
||||
mod oauth_routes;
|
||||
mod probe_routes;
|
||||
mod query;
|
||||
pub(crate) mod query;
|
||||
mod quota_presenter;
|
||||
mod quota_routes;
|
||||
|
||||
pub(crate) use export_routes::*;
|
||||
@@ -575,4 +509,5 @@ pub(crate) use lifecycle_routes::*;
|
||||
pub(crate) use oauth_routes::*;
|
||||
pub(crate) use probe_routes::*;
|
||||
use query::*;
|
||||
use quota_presenter::*;
|
||||
pub(crate) use quota_routes::*;
|
||||
|
||||
@@ -1,4 +1,11 @@
|
||||
use super::*;
|
||||
use std::convert::Infallible;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::StreamExt;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::fleet::manage::AccountTestEvent;
|
||||
|
||||
/// `GET /api/admin/accounts/test?id=...&modelId=...`
|
||||
pub async fn test_account_connection(
|
||||
@@ -21,6 +28,12 @@ pub async fn test_account_connection(
|
||||
.test_connection_stream(&account_id, model)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?;
|
||||
let stream = stream.map(|event| {
|
||||
Ok::<Bytes, Infallible>(Bytes::from(format!(
|
||||
"data: {}\n\n",
|
||||
account_test_event_json(event)
|
||||
)))
|
||||
});
|
||||
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
@@ -32,6 +45,43 @@ pub async fn test_account_connection(
|
||||
.map_err(|_| AdminError::internal("Failed to build account test stream"))
|
||||
}
|
||||
|
||||
fn account_test_event_json(event: AccountTestEvent) -> Value {
|
||||
match event {
|
||||
AccountTestEvent::Started { model } => json!({
|
||||
"type": "test_start",
|
||||
"model": model,
|
||||
"text": "正在连接 Codex Responses"
|
||||
}),
|
||||
AccountTestEvent::Request { payload } => json!({
|
||||
"type": "request",
|
||||
"payload": payload
|
||||
}),
|
||||
AccountTestEvent::Content { text } => json!({ "type": "content", "text": text }),
|
||||
AccountTestEvent::Complete { account_status } => account_test_terminal_event(
|
||||
json!({ "type": "test_complete", "success": true }),
|
||||
account_status,
|
||||
),
|
||||
AccountTestEvent::Error {
|
||||
error,
|
||||
account_status,
|
||||
} => {
|
||||
account_test_terminal_event(json!({ "type": "error", "error": error }), account_status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn account_test_terminal_event(mut event: Value, account_status: Option<AccountStatus>) -> Value {
|
||||
if let Some(account_status) = account_status
|
||||
&& let Some(event) = event.as_object_mut()
|
||||
{
|
||||
event.insert(
|
||||
"accountStatus".to_string(),
|
||||
Value::String(account_status.as_str().to_string()),
|
||||
);
|
||||
}
|
||||
event
|
||||
}
|
||||
|
||||
/// `GET /api/admin/accounts/models?id=...`
|
||||
pub async fn account_models(
|
||||
State(state): State<AppState>,
|
||||
@@ -60,3 +110,29 @@ pub async fn account_models(
|
||||
Err(error) => Err(account_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/accounts/models`
|
||||
pub async fn refresh_account_models(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
Json(payload): Json<AccountActionRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let models = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.refresh_account_models(&payload.id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountModelsData {
|
||||
models: models
|
||||
.into_iter()
|
||||
.map(|model| AccountModelData {
|
||||
id: model.id,
|
||||
label: model.label,
|
||||
})
|
||||
.collect(),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,211 @@
|
||||
//! typed quota read model 的管理端 presenter。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::{
|
||||
api::admin::accounts_routes::query::{AccountQuotaReadModel, AccountQuotaWindowLocalUsage},
|
||||
fleet::quota::{QuotaWindowKind, QuotaWindowRole},
|
||||
infra::{
|
||||
format::{
|
||||
format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
|
||||
},
|
||||
time::{china_datetime, china_relative_time},
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(super) struct AccountQuotaData {
|
||||
refreshed_at_display: String,
|
||||
windows: Vec<AccountQuotaWindowData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaWindowData {
|
||||
key: String,
|
||||
group: &'static str,
|
||||
window_seconds: Option<u64>,
|
||||
label_display: String,
|
||||
used_percent: Option<f64>,
|
||||
used_percent_display: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
local_usage: Option<AccountQuotaWindowLocalUsageData>,
|
||||
reset_at_display: String,
|
||||
window_used_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaWindowLocalUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
|
||||
impl From<AccountQuotaWindowLocalUsage> for AccountQuotaWindowLocalUsageData {
|
||||
fn from(usage: AccountQuotaWindowLocalUsage) -> Self {
|
||||
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
|
||||
Self {
|
||||
request_count: usage.request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
|
||||
input_tokens: usage.input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
|
||||
output_tokens: usage.output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
|
||||
cached_tokens: usage.cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AccountQuotaData {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
refreshed_at_display: "-".to_string(),
|
||||
windows: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn quota_data(read: AccountQuotaReadModel) -> AccountQuotaData {
|
||||
let now = Utc::now();
|
||||
AccountQuotaData {
|
||||
refreshed_at_display: china_relative_time(read.fetched_at, now),
|
||||
windows: read
|
||||
.windows
|
||||
.into_iter()
|
||||
.map(|read| {
|
||||
let window = read.quota.window;
|
||||
let window_seconds = window.window_seconds();
|
||||
let reset_at = window.reset_datetime();
|
||||
let used_percent = window.used_percent;
|
||||
AccountQuotaWindowData {
|
||||
key: read.quota.key,
|
||||
group: quota_window_group(window.kind),
|
||||
window_seconds,
|
||||
label_display: quota_window_label(
|
||||
read.quota.source.as_str(),
|
||||
read.quota.limit_name.as_deref(),
|
||||
read.quota.metered_feature.as_deref(),
|
||||
read.quota.role,
|
||||
window.kind,
|
||||
window_seconds,
|
||||
),
|
||||
used_percent,
|
||||
used_percent_display: used_percent
|
||||
.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
local_usage: read.local_usage.map(Into::into),
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
window_used_display: quota_window_used_display(reset_at, window_seconds, now),
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_group(kind: QuotaWindowKind) -> &'static str {
|
||||
match kind {
|
||||
QuotaWindowKind::Monthly => "monthly",
|
||||
QuotaWindowKind::ShortTerm | QuotaWindowKind::Weekly => "shortTerm",
|
||||
QuotaWindowKind::Other => "other",
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_label(
|
||||
source: &str,
|
||||
limit_name: Option<&str>,
|
||||
metered_feature: Option<&str>,
|
||||
role: QuotaWindowRole,
|
||||
kind: QuotaWindowKind,
|
||||
window_seconds: Option<u64>,
|
||||
) -> String {
|
||||
if role == QuotaWindowRole::Monthly {
|
||||
return "月限额".to_string();
|
||||
}
|
||||
let base = match kind {
|
||||
QuotaWindowKind::ShortTerm => "5小时限额".to_string(),
|
||||
QuotaWindowKind::Weekly => "周限额".to_string(),
|
||||
QuotaWindowKind::Monthly => "月限额".to_string(),
|
||||
QuotaWindowKind::Other => custom_window_label(window_seconds),
|
||||
};
|
||||
let label = limit_name
|
||||
.or(metered_feature)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty() && !value.eq_ignore_ascii_case("codex"));
|
||||
if source == "core" || label.is_none() {
|
||||
return base;
|
||||
}
|
||||
let label = label.unwrap_or_default();
|
||||
let label = if is_review_limit_label(label) {
|
||||
"代码审查"
|
||||
} else {
|
||||
label
|
||||
};
|
||||
format!("{label} · {base}")
|
||||
}
|
||||
|
||||
fn custom_window_label(window_seconds: Option<u64>) -> String {
|
||||
let Some(seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
|
||||
return "额度".to_string();
|
||||
};
|
||||
if seconds % 86_400 == 0 {
|
||||
format!("{}天限额", seconds / 86_400)
|
||||
} else if seconds % 3_600 == 0 {
|
||||
format!("{}小时限额", seconds / 3_600)
|
||||
} else {
|
||||
format!("{}分钟限额", seconds.div_ceil(60))
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_used_display(
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
window_seconds: Option<u64>,
|
||||
now: DateTime<Utc>,
|
||||
) -> String {
|
||||
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
|
||||
return "-".to_string();
|
||||
};
|
||||
let remaining = reset_at
|
||||
.signed_duration_since(now)
|
||||
.num_seconds()
|
||||
.max(0)
|
||||
.cast_unsigned();
|
||||
let used = window_seconds.saturating_sub(remaining);
|
||||
format!(
|
||||
"{} / {}",
|
||||
format_duration_days(used),
|
||||
format_duration_days(window_seconds)
|
||||
)
|
||||
}
|
||||
|
||||
fn format_duration_days(seconds: u64) -> String {
|
||||
let days = seconds as f64 / 86_400.0;
|
||||
if days >= 1.0 {
|
||||
format!("{days:.1}d")
|
||||
} else {
|
||||
format!("{:.1}h", seconds as f64 / 3_600.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn is_review_limit_label(value: &str) -> bool {
|
||||
let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_");
|
||||
normalized == "review"
|
||||
|| normalized == "code_review"
|
||||
|| normalized == "codex_review"
|
||||
|| normalized == "codex_code_review"
|
||||
|| normalized.contains("code_review")
|
||||
|| normalized.contains("codex_review")
|
||||
}
|
||||
@@ -14,13 +14,25 @@ pub async fn account_quota(
|
||||
.await
|
||||
{
|
||||
Ok(data) => {
|
||||
let quota = data.get("quota").cloned().unwrap_or(Value::Null);
|
||||
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
|
||||
let quota_json = quota.to_string();
|
||||
let mut quota_data = quota_data("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
let quota = serde_json::to_value(&data.quota)
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let raw = data.raw;
|
||||
let quota_read = AccountQuotaReadModel::from_snapshot(data.quota, Some(Utc::now()));
|
||||
let account = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.get(&account_id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?
|
||||
.ok_or_else(account_not_found)?;
|
||||
let item = state
|
||||
.services
|
||||
.account_list
|
||||
.enrich_account(account, quota_read)
|
||||
.await
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let quota_data = item.quota.clone().map(quota_data).unwrap_or_default();
|
||||
let account = account_list_item_data(item);
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account(
|
||||
@@ -39,23 +51,3 @@ pub async fn account_quota(
|
||||
Err(e) => Err(account_error(&e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn account_data_for_quota_refresh(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota_data: AccountQuotaData,
|
||||
) -> Result<AdminAccountData, AdminError> {
|
||||
let Some(account) = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.get(account_id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?
|
||||
else {
|
||||
return Err(account_not_found());
|
||||
};
|
||||
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
|
||||
let stats =
|
||||
account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?;
|
||||
Ok(stats.data_for(account))
|
||||
}
|
||||
|
||||
+124
-39
@@ -1,4 +1,4 @@
|
||||
//! Dashboard 遥测聚合与展示模型。
|
||||
//! Dashboard 遥测事实的 HTTP presenter。
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -11,11 +11,14 @@ use crate::infra::{
|
||||
time::{china_day_start, china_quarter_hour_start},
|
||||
};
|
||||
|
||||
use super::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
|
||||
use crate::telemetry::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
|
||||
use crate::telemetry::usage::insights::RequestHealthTimeBucket;
|
||||
|
||||
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
|
||||
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
|
||||
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
|
||||
const HEALTH_TIMELINE_MIN_SAMPLE_SIZE: u64 = 10;
|
||||
const HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD: u64 = 3;
|
||||
const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 99.0;
|
||||
const TREND_SLOT_MINUTES: i64 = 15;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
@@ -40,7 +43,35 @@ pub struct DashboardHealthTimelineData {
|
||||
title: String,
|
||||
description: String,
|
||||
reliability_display: String,
|
||||
points: String,
|
||||
status: DashboardHealthTimelineStatus,
|
||||
success_requests: u64,
|
||||
failed_requests: u64,
|
||||
cancelled_requests: u64,
|
||||
caller_error_requests: u64,
|
||||
points: Vec<DashboardHealthTimelinePointData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DashboardHealthTimelinePointData {
|
||||
time: String,
|
||||
status: DashboardHealthTimelineStatus,
|
||||
reliability_display: String,
|
||||
success_requests: u64,
|
||||
failed_requests: u64,
|
||||
cancelled_requests: u64,
|
||||
caller_error_requests: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum DashboardHealthTimelineStatus {
|
||||
Future,
|
||||
NoData,
|
||||
Unavailable,
|
||||
Unstable,
|
||||
LowSample,
|
||||
Stable,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -148,6 +179,14 @@ struct UsageWindow {
|
||||
min_first_token_bucket_latency: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
struct HealthWindow {
|
||||
success_requests: u64,
|
||||
failed_requests: u64,
|
||||
cancelled_requests: u64,
|
||||
caller_error_requests: u64,
|
||||
}
|
||||
|
||||
impl UsageWindow {
|
||||
fn tokens(self) -> u64 {
|
||||
self.input_tokens + self.output_tokens
|
||||
@@ -320,13 +359,13 @@ pub fn dashboard_trend_data_at(
|
||||
}
|
||||
|
||||
pub fn dashboard_health_timeline_data(
|
||||
buckets: &[AccountUsageTimeBucket],
|
||||
buckets: &[RequestHealthTimeBucket],
|
||||
) -> DashboardHealthTimelineData {
|
||||
dashboard_health_timeline_data_at(buckets, Utc::now())
|
||||
}
|
||||
|
||||
pub fn dashboard_health_timeline_data_at(
|
||||
records: &[AccountUsageTimeBucket],
|
||||
records: &[RequestHealthTimeBucket],
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardHealthTimelineData {
|
||||
let current_slot = china_quarter_hour_start(now);
|
||||
@@ -335,7 +374,7 @@ pub fn dashboard_health_timeline_data_at(
|
||||
.map(|index| {
|
||||
(
|
||||
start + Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * index),
|
||||
UsageWindow::default(),
|
||||
HealthWindow::default(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -348,36 +387,62 @@ pub fn dashboard_health_timeline_data_at(
|
||||
.iter_mut()
|
||||
.find(|(bucket_start, _)| *bucket_start == record_slot)
|
||||
{
|
||||
apply_bucket(bucket, record);
|
||||
bucket.success_requests = bucket
|
||||
.success_requests
|
||||
.saturating_add(record.success_requests);
|
||||
bucket.failed_requests = bucket
|
||||
.failed_requests
|
||||
.saturating_add(record.failed_requests);
|
||||
bucket.cancelled_requests = bucket
|
||||
.cancelled_requests
|
||||
.saturating_add(record.cancelled_requests);
|
||||
bucket.caller_error_requests = bucket
|
||||
.caller_error_requests
|
||||
.saturating_add(record.caller_error_requests);
|
||||
}
|
||||
}
|
||||
|
||||
let (requests, errors) = buckets
|
||||
let totals = buckets
|
||||
.iter()
|
||||
.filter(|(bucket_start, _)| *bucket_start <= current_slot)
|
||||
.fold((0, 0), |(requests, errors), (_, bucket)| {
|
||||
(requests + bucket.requests, errors + bucket.errors)
|
||||
.fold(HealthWindow::default(), |mut totals, (_, bucket)| {
|
||||
totals.success_requests = totals
|
||||
.success_requests
|
||||
.saturating_add(bucket.success_requests);
|
||||
totals.failed_requests = totals
|
||||
.failed_requests
|
||||
.saturating_add(bucket.failed_requests);
|
||||
totals.cancelled_requests = totals
|
||||
.cancelled_requests
|
||||
.saturating_add(bucket.cancelled_requests);
|
||||
totals.caller_error_requests = totals
|
||||
.caller_error_requests
|
||||
.saturating_add(bucket.caller_error_requests);
|
||||
totals
|
||||
});
|
||||
let reliability = (requests > 0)
|
||||
.then(|| ((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0);
|
||||
DashboardHealthTimelineData {
|
||||
title: "请求健康时间线".to_string(),
|
||||
description: "请求可靠性".to_string(),
|
||||
reliability_display: reliability
|
||||
.map(|value| format!("{value:.1}%"))
|
||||
.unwrap_or_else(|| "-".to_string()),
|
||||
description: "有效请求可用性".to_string(),
|
||||
reliability_display: format_health_reliability(totals),
|
||||
status: health_status(totals, false),
|
||||
success_requests: totals.success_requests,
|
||||
failed_requests: totals.failed_requests,
|
||||
cancelled_requests: totals.cancelled_requests,
|
||||
caller_error_requests: totals.caller_error_requests,
|
||||
points: buckets
|
||||
.into_iter()
|
||||
.map(|(bucket_start, bucket)| {
|
||||
let success_rate = if bucket.requests > 0 {
|
||||
(bucket.requests.saturating_sub(bucket.errors) as f64 / bucket.requests as f64
|
||||
* 1000.0)
|
||||
.round()
|
||||
/ 10.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
health_tone(bucket, success_rate, bucket_start > current_slot)
|
||||
.enumerate()
|
||||
.map(|(index, (bucket_start, bucket))| {
|
||||
let elapsed_minutes = index as i64 * HEALTH_TIMELINE_SLOT_MINUTES;
|
||||
DashboardHealthTimelinePointData {
|
||||
time: format!("{:02}:{:02}", elapsed_minutes / 60, elapsed_minutes % 60),
|
||||
status: health_status(bucket, bucket_start > current_slot),
|
||||
reliability_display: format_health_reliability(bucket),
|
||||
success_requests: bucket.success_requests,
|
||||
failed_requests: bucket.failed_requests,
|
||||
cancelled_requests: bucket.cancelled_requests,
|
||||
caller_error_requests: bucket.caller_error_requests,
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
@@ -510,23 +575,43 @@ fn first_token_bucket_latency(record: &AccountUsageTimeBucket) -> Option<u64> {
|
||||
.filter(|latency| *latency > 0)
|
||||
}
|
||||
|
||||
fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char {
|
||||
let successes = bucket.requests.saturating_sub(bucket.errors);
|
||||
fn health_status(bucket: HealthWindow, is_future: bool) -> DashboardHealthTimelineStatus {
|
||||
let eligible_requests = bucket
|
||||
.success_requests
|
||||
.saturating_add(bucket.failed_requests);
|
||||
if is_future {
|
||||
'0'
|
||||
} else if bucket.requests == 0 {
|
||||
'1'
|
||||
} else if successes == 0 {
|
||||
'2'
|
||||
} else if success_rate < 90.0 {
|
||||
'3'
|
||||
} else if successes < HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD {
|
||||
'4'
|
||||
DashboardHealthTimelineStatus::Future
|
||||
} else if eligible_requests == 0 {
|
||||
DashboardHealthTimelineStatus::NoData
|
||||
} else if bucket.success_requests == 0
|
||||
&& bucket.failed_requests >= HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD
|
||||
{
|
||||
DashboardHealthTimelineStatus::Unavailable
|
||||
} else if eligible_requests < HEALTH_TIMELINE_MIN_SAMPLE_SIZE {
|
||||
DashboardHealthTimelineStatus::LowSample
|
||||
} else if health_reliability(bucket)
|
||||
.is_some_and(|reliability| reliability < HEALTH_TIMELINE_STABLE_RELIABILITY)
|
||||
{
|
||||
DashboardHealthTimelineStatus::Unstable
|
||||
} else {
|
||||
'5'
|
||||
DashboardHealthTimelineStatus::Stable
|
||||
}
|
||||
}
|
||||
|
||||
fn health_reliability(bucket: HealthWindow) -> Option<f64> {
|
||||
let eligible_requests = bucket
|
||||
.success_requests
|
||||
.saturating_add(bucket.failed_requests);
|
||||
(eligible_requests > 0)
|
||||
.then(|| bucket.success_requests as f64 / eligible_requests as f64 * 100.0)
|
||||
}
|
||||
|
||||
fn format_health_reliability(bucket: HealthWindow) -> String {
|
||||
health_reliability(bucket)
|
||||
.map(|value| format!("{value:.1}%"))
|
||||
.unwrap_or_else(|| "-".to_string())
|
||||
}
|
||||
|
||||
fn summary_cache_hit_rate(summary: &RetainedUsageSummary) -> Option<f64> {
|
||||
let input_tokens = nonnegative_i64_to_u64(summary.totals.input_tokens);
|
||||
(input_tokens > 0)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,6 +2,7 @@
|
||||
|
||||
pub mod accounts_routes;
|
||||
pub mod auth_routes;
|
||||
mod dashboard_presenter;
|
||||
pub mod dashboard_routes;
|
||||
pub mod keys_routes;
|
||||
pub mod ops_routes;
|
||||
|
||||
@@ -17,7 +17,10 @@ use crate::{
|
||||
json::{NumberedPage, clamp_limit, clamp_page},
|
||||
time::china_datetime,
|
||||
},
|
||||
telemetry::{ops::query::OpsQueryError, ops::store::OpsErrorFilter, ops::types::OpsErrorLog},
|
||||
telemetry::{
|
||||
ops::query::OpsQueryError,
|
||||
ops::types::{OpsErrorFilter, OpsErrorLog},
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
|
||||
@@ -14,7 +14,8 @@ use super::{
|
||||
accounts_routes::{
|
||||
account_models, account_quota, accounts, batch_delete_accounts, create_account,
|
||||
export_accounts, health_check_accounts, import_accounts, oauth_authorize_account,
|
||||
oauth_exchange_account, refresh_account, test_account_connection, update_account,
|
||||
oauth_exchange_account, refresh_account, refresh_account_models, test_account_connection,
|
||||
update_account,
|
||||
},
|
||||
auth_routes::{login, logout, session_status},
|
||||
dashboard_routes::{dashboard_summary, dashboard_trend},
|
||||
@@ -45,7 +46,7 @@ pub fn router() -> Router<AppState> {
|
||||
.route("/api/admin/settings", get(settings).post(update_settings))
|
||||
.route(
|
||||
"/api/admin/settings/admin-api-key",
|
||||
get(admin_api_key_status).delete(delete_admin_api_key),
|
||||
get(admin_api_key_status).post(delete_admin_api_key),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/settings/admin-api-key/regenerate",
|
||||
@@ -72,7 +73,10 @@ pub fn router() -> Router<AppState> {
|
||||
post(oauth_exchange_account),
|
||||
)
|
||||
.route("/api/admin/accounts/test", get(test_account_connection))
|
||||
.route("/api/admin/accounts/models", get(account_models))
|
||||
.route(
|
||||
"/api/admin/accounts/models",
|
||||
get(account_models).post(refresh_account_models),
|
||||
)
|
||||
.route("/api/admin/accounts/delete", post(batch_delete_accounts))
|
||||
.route("/api/admin/accounts/update", post(update_account))
|
||||
.route("/api/admin/accounts/refresh", post(refresh_account))
|
||||
|
||||
@@ -115,7 +115,7 @@ pub(crate) async fn regenerate_admin_api_key(
|
||||
))
|
||||
}
|
||||
|
||||
/// `DELETE /api/admin/settings/admin-api-key`
|
||||
/// `POST /api/admin/settings/admin-api-key`
|
||||
pub(crate) async fn delete_admin_api_key(
|
||||
State(state): State<AppState>,
|
||||
_auth: AdminAuth,
|
||||
|
||||
@@ -121,11 +121,13 @@ pub(crate) struct UsageRecordTokenDetailsData {
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
reasoning_tokens: u64,
|
||||
total_tokens: u64,
|
||||
input_tokens_display: String,
|
||||
output_tokens_display: String,
|
||||
cached_tokens_display: String,
|
||||
cache_write_tokens_display: String,
|
||||
reasoning_tokens_display: String,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
@@ -136,22 +138,26 @@ pub(crate) struct UsageRecordBillingData {
|
||||
input_amount: f64,
|
||||
output_amount: f64,
|
||||
cache_read_amount: f64,
|
||||
cache_write_amount: f64,
|
||||
standard_amount: f64,
|
||||
total_amount: f64,
|
||||
input_price_per_mtoken: f64,
|
||||
output_price_per_mtoken: f64,
|
||||
cache_read_price_per_mtoken: f64,
|
||||
cache_write_price_per_mtoken: f64,
|
||||
service_tier: Option<String>,
|
||||
service_tier_display: String,
|
||||
multiplier: f64,
|
||||
input_amount_display: String,
|
||||
output_amount_display: String,
|
||||
cache_read_amount_display: String,
|
||||
cache_write_amount_display: String,
|
||||
standard_amount_display: String,
|
||||
total_amount_display: String,
|
||||
input_price_display: String,
|
||||
output_price_display: String,
|
||||
cache_read_price_display: String,
|
||||
cache_write_price_display: String,
|
||||
multiplier_display: String,
|
||||
}
|
||||
|
||||
@@ -162,6 +168,7 @@ struct UsageRecordSummaryData {
|
||||
input_tokens: String,
|
||||
output_tokens: String,
|
||||
cached_tokens: String,
|
||||
cache_write_tokens: String,
|
||||
total_tokens: String,
|
||||
average_latency_ms: String,
|
||||
}
|
||||
@@ -175,6 +182,7 @@ struct UsageRecordBreakdownData {
|
||||
input_tokens: String,
|
||||
output_tokens: String,
|
||||
cached_tokens: String,
|
||||
cache_write_tokens: String,
|
||||
total_tokens: String,
|
||||
total_tokens_value: u64,
|
||||
total_tokens_total: String,
|
||||
@@ -195,6 +203,8 @@ struct UsageRecordTrendPointData {
|
||||
output_tokens_value: u64,
|
||||
cached_tokens: String,
|
||||
cached_tokens_value: u64,
|
||||
cache_write_tokens: String,
|
||||
cache_write_tokens_value: u64,
|
||||
total_tokens: String,
|
||||
total_tokens_value: u64,
|
||||
cache_hit_rate: String,
|
||||
@@ -574,6 +584,10 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
|
||||
.cached_tokens
|
||||
.and_then(|value| value.try_into().ok())
|
||||
.unwrap_or(0);
|
||||
let cache_write_tokens: u64 = record
|
||||
.cache_write_tokens
|
||||
.and_then(|value| value.try_into().ok())
|
||||
.unwrap_or(0);
|
||||
let reasoning_tokens: u64 = record
|
||||
.reasoning_tokens
|
||||
.and_then(|value| value.try_into().ok())
|
||||
@@ -584,11 +598,13 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
reasoning_tokens,
|
||||
total_tokens,
|
||||
input_tokens_display: format_number(input_tokens),
|
||||
output_tokens_display: format_number(output_tokens),
|
||||
cached_tokens_display: format_compact_number(cached_tokens),
|
||||
cache_write_tokens_display: format_compact_number(cache_write_tokens),
|
||||
reasoning_tokens_display: format_number(reasoning_tokens),
|
||||
total_tokens_display: format_number(total_tokens),
|
||||
}
|
||||
@@ -605,28 +621,33 @@ pub(crate) fn usage_billing(
|
||||
tokens.input_tokens,
|
||||
tokens.output_tokens,
|
||||
tokens.cached_tokens,
|
||||
tokens.cache_write_tokens,
|
||||
)?;
|
||||
|
||||
Some(UsageRecordBillingData {
|
||||
input_amount: billing.input_amount,
|
||||
output_amount: billing.output_amount,
|
||||
cache_read_amount: billing.cache_read_amount,
|
||||
cache_write_amount: billing.cache_write_amount,
|
||||
standard_amount: billing.standard_amount,
|
||||
total_amount: billing.total_amount,
|
||||
input_price_per_mtoken: billing.input_price_per_mtoken,
|
||||
output_price_per_mtoken: billing.output_price_per_mtoken,
|
||||
cache_read_price_per_mtoken: billing.cache_read_price_per_mtoken,
|
||||
cache_write_price_per_mtoken: billing.cache_write_price_per_mtoken,
|
||||
service_tier: billing.service_tier,
|
||||
service_tier_display: billing.service_tier_display,
|
||||
multiplier: billing.multiplier,
|
||||
input_amount_display: format_billing_amount(billing.input_amount),
|
||||
output_amount_display: format_billing_amount(billing.output_amount),
|
||||
cache_read_amount_display: format_billing_amount(billing.cache_read_amount),
|
||||
cache_write_amount_display: format_billing_amount(billing.cache_write_amount),
|
||||
standard_amount_display: format_billing_amount(billing.standard_amount),
|
||||
total_amount_display: format_billing_amount(billing.total_amount),
|
||||
input_price_display: format_token_price(billing.input_price_per_mtoken),
|
||||
output_price_display: format_token_price(billing.output_price_per_mtoken),
|
||||
cache_read_price_display: format_token_price(billing.cache_read_price_per_mtoken),
|
||||
cache_write_price_display: format_token_price(billing.cache_write_price_per_mtoken),
|
||||
multiplier_display: format_multiplier(billing.multiplier),
|
||||
})
|
||||
}
|
||||
@@ -638,6 +659,7 @@ impl From<UsageRecordSummary> for UsageRecordSummaryData {
|
||||
input_tokens: format_tokens(summary.input_tokens),
|
||||
output_tokens: format_tokens(summary.output_tokens),
|
||||
cached_tokens: format_tokens(summary.cached_tokens),
|
||||
cache_write_tokens: format_tokens(summary.cache_write_tokens),
|
||||
total_tokens: format_tokens(summary.total_tokens),
|
||||
average_latency_ms: format_duration_ms_f64(summary.average_latency_ms),
|
||||
}
|
||||
@@ -659,6 +681,7 @@ impl UsageRecordBreakdownData {
|
||||
input_tokens: format_tokens(item.input_tokens),
|
||||
output_tokens: format_tokens(item.output_tokens),
|
||||
cached_tokens: format_tokens(item.cached_tokens),
|
||||
cache_write_tokens: format_tokens(item.cache_write_tokens),
|
||||
total_tokens: format_tokens(item.total_tokens),
|
||||
total_tokens_value: item.total_tokens,
|
||||
total_tokens_total: format_tokens(total_tokens_total),
|
||||
@@ -698,6 +721,8 @@ impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
|
||||
output_tokens_value: point.output_tokens,
|
||||
cached_tokens: format_tokens(point.cached_tokens),
|
||||
cached_tokens_value: point.cached_tokens,
|
||||
cache_write_tokens: format_tokens(point.cache_write_tokens),
|
||||
cache_write_tokens_value: point.cache_write_tokens,
|
||||
total_tokens: format_tokens(point.total_tokens),
|
||||
total_tokens_value: point.total_tokens,
|
||||
cache_hit_rate: format_percent(cache_hit_rate_value),
|
||||
|
||||
@@ -83,7 +83,7 @@ const MODEL_NOT_FOUND_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
};
|
||||
const HISTORY_UNAVAILABLE_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
error_type: "invalid_request_error",
|
||||
code: "previous_response_unavailable",
|
||||
code: "previous_response_not_found",
|
||||
};
|
||||
const CONTINUATION_BUSY_ERROR: OpenAiErrorKind = OpenAiErrorKind {
|
||||
error_type: "server_error",
|
||||
|
||||
@@ -19,7 +19,13 @@ use crate::{
|
||||
api::AppState,
|
||||
api::middleware::request_id::{ClientIp, RequestId},
|
||||
dispatch::{errors::ResponseDispatchError, service::ResponseDispatchStream},
|
||||
upstream::openai::protocol::{responses::CodexResponsesRequest, sse::sse_body_has_done},
|
||||
upstream::openai::protocol::{
|
||||
responses::{
|
||||
CodexResponsesRequest, WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY,
|
||||
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
|
||||
},
|
||||
sse::sse_body_has_done,
|
||||
},
|
||||
};
|
||||
|
||||
use self::sse::{done_sse_frame, event_stream_response as sse_event_stream_response};
|
||||
@@ -103,11 +109,29 @@ pub fn build_codex_request(
|
||||
None => {}
|
||||
}
|
||||
apply_context_header_fields(&mut request, headers);
|
||||
apply_request_semantic_headers(&mut request, headers);
|
||||
apply_session_fields(&mut request, headers);
|
||||
request.client_user_agent = user_agent_from_headers(headers);
|
||||
request
|
||||
}
|
||||
|
||||
fn apply_request_semantic_headers(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
|
||||
request.responses_lite = header_string(headers, X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER)
|
||||
.or_else(|| {
|
||||
request
|
||||
.client_metadata()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| {
|
||||
metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY)
|
||||
})
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
});
|
||||
request.memgen_request = header_string(headers, X_OPENAI_MEMGEN_REQUEST_HEADER);
|
||||
}
|
||||
|
||||
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"))
|
||||
|
||||
@@ -131,8 +131,10 @@ struct ResponsesWebSocketSession {
|
||||
connection_id: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ConnectionTranscriptCollector {
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
completed_items: Vec<Value>,
|
||||
completed: Option<ConnectionTranscriptFacts>,
|
||||
}
|
||||
@@ -143,6 +145,20 @@ enum ForwardDispatchOutcome {
|
||||
}
|
||||
|
||||
impl ConnectionTranscriptCollector {
|
||||
fn new(
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
completed_items: Vec::new(),
|
||||
completed: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe(&mut self, event: &CanonicalResponseEvent) {
|
||||
match event.event_type() {
|
||||
Some("response.output_item.done") => {
|
||||
@@ -166,7 +182,13 @@ impl ConnectionTranscriptCollector {
|
||||
.filter(|output| !output.is_empty())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| self.completed_items.clone());
|
||||
self.completed = Some(ConnectionTranscriptFacts::new(response_id, output));
|
||||
self.completed = Some(ConnectionTranscriptFacts::new(
|
||||
self.account_id.clone(),
|
||||
self.request_input.clone(),
|
||||
self.continued_from_previous_response,
|
||||
response_id,
|
||||
output,
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
@@ -346,6 +368,9 @@ async fn forward_dispatch_stream(
|
||||
request_id: &str,
|
||||
) -> ForwardDispatchOutcome {
|
||||
let ResponseDispatchStream {
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
body,
|
||||
mut canonical_events,
|
||||
response_headers,
|
||||
@@ -361,7 +386,11 @@ async fn forward_dispatch_stream(
|
||||
|
||||
let mut body = body;
|
||||
let mut terminal_sent = false;
|
||||
let mut completed = ConnectionTranscriptCollector::default();
|
||||
let mut completed = ConnectionTranscriptCollector::new(
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
);
|
||||
loop {
|
||||
let next = if terminal_sent {
|
||||
ActiveResponseInput::Upstream(body.next().await)
|
||||
@@ -516,7 +545,7 @@ fn parse_response_create(payload: &str) -> Result<Map<String, Value>, ResponseCr
|
||||
Some(Value::String(message_type)) if message_type == "response.create" => {}
|
||||
_ => return Err(ResponseCreateFrameError::UnsupportedType),
|
||||
}
|
||||
if body.get("stream").and_then(Value::as_bool) != Some(true) {
|
||||
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
|
||||
return Err(ResponseCreateFrameError::StreamingRequired);
|
||||
}
|
||||
Ok(body)
|
||||
|
||||
@@ -10,32 +10,30 @@ use axum::{Router, extract::State, http::StatusCode, middleware, routing::get};
|
||||
|
||||
use crate::{
|
||||
api::{
|
||||
admin, assets, client,
|
||||
admin::{
|
||||
self, accounts_routes::query::AccountListQueryService,
|
||||
dashboard_routes::DashboardQueryService,
|
||||
},
|
||||
assets, client,
|
||||
middleware::{request_id::attach_request_id, trace::http_trace_layer},
|
||||
},
|
||||
auth::service::SessionService,
|
||||
dispatch::{affinity::SessionAffinityService, service::ResponseDispatchService},
|
||||
fleet::{
|
||||
manage::AccountManageService, pool::AccountPoolService, refresh::TokenRefreshService,
|
||||
store::AccountStore,
|
||||
},
|
||||
fleet::{manage::AccountManageService, pool::AccountPoolService},
|
||||
keys::{manage::KeyManageService, service::KeyVerifier},
|
||||
models::service::ModelService,
|
||||
settings::service::SettingsService,
|
||||
telemetry::{
|
||||
account_usage::query::AccountUsageQueryService, ops::query::OpsQueryService,
|
||||
usage::query::UsageQueryService,
|
||||
},
|
||||
telemetry::{ops::query::OpsQueryService, usage::query::UsageQueryService},
|
||||
update::service::SystemUpdateService,
|
||||
upstream::openai::{fingerprint::RuntimeFingerprint, token_client::OpenAiTokenClient},
|
||||
};
|
||||
|
||||
/// HTTP API 所需的领域服务集合。
|
||||
#[derive(Clone)]
|
||||
pub struct ApiServices {
|
||||
pub(crate) account_list: Arc<AccountListQueryService>,
|
||||
pub(crate) dashboard: Arc<DashboardQueryService>,
|
||||
pub health_probe: Arc<dyn HealthProbe>,
|
||||
pub models: Arc<ModelService>,
|
||||
pub accounts: Arc<dyn AccountStore>,
|
||||
pub client_keys: Arc<KeyVerifier>,
|
||||
pub admin_client_keys: Arc<KeyManageService>,
|
||||
pub admin_sessions: Arc<SessionService>,
|
||||
@@ -43,12 +41,9 @@ pub struct ApiServices {
|
||||
pub admin_accounts: Arc<AccountManageService>,
|
||||
pub usage_records: Arc<UsageQueryService>,
|
||||
pub ops_errors: Arc<OpsQueryService>,
|
||||
pub usage: Arc<AccountUsageQueryService>,
|
||||
pub account_pool: Arc<AccountPoolService>,
|
||||
pub token_refresh: Arc<TokenRefreshService<OpenAiTokenClient>>,
|
||||
pub responses: Arc<ResponseDispatchService>,
|
||||
pub session_affinity: Arc<SessionAffinityService>,
|
||||
pub fingerprint: RuntimeFingerprint,
|
||||
pub process_control: Arc<dyn ProcessControl>,
|
||||
pub system_update: Arc<SystemUpdateService>,
|
||||
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
use std::{
|
||||
collections::{BTreeMap, HashSet},
|
||||
collections::BTreeMap,
|
||||
env,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use axum::http::HeaderName;
|
||||
use chrono::{DateTime, Utc};
|
||||
use secrecy::{ExposeSecret, SecretBox};
|
||||
use serde::{Deserialize, de::IgnoredAny};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
@@ -106,7 +106,7 @@ impl BootstrapConfig {
|
||||
runtime,
|
||||
quota,
|
||||
ws_pool,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
admin,
|
||||
logging,
|
||||
telemetry,
|
||||
@@ -122,7 +122,7 @@ impl BootstrapConfig {
|
||||
redis: RedisConfig { url: redis.url },
|
||||
runtime,
|
||||
ws_pool,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: admin.session_ttl_minutes,
|
||||
default_username: admin.default_username,
|
||||
@@ -241,8 +241,8 @@ pub struct AppConfig {
|
||||
pub runtime: RuntimePathsConfig,
|
||||
/// WebSocket 连接池启动设置。
|
||||
pub ws_pool: WebSocketPoolSettings,
|
||||
/// 上游请求指纹默认配置。
|
||||
pub fingerprint: FingerprintConfig,
|
||||
/// 经审计固定的 Codex Desktop 上游请求画像。
|
||||
pub wire_profile: WireProfileConfig,
|
||||
/// 管理员初始化配置,不包含密码。
|
||||
pub admin: AdminConfig,
|
||||
/// 日志配置。
|
||||
@@ -350,42 +350,36 @@ pub struct WebSocketPoolSettings {
|
||||
pub max_age_ms: u64,
|
||||
/// 单账号最大连接数。
|
||||
pub max_per_account: usize,
|
||||
/// 所有账号合计最大连接数。
|
||||
pub max_total: usize,
|
||||
/// 所有账号合计最大并发建连数。
|
||||
pub max_connecting: usize,
|
||||
/// 首个真实输出(首事件)到达前的绝对超时;`0` 表示禁用。
|
||||
pub initial_event_timeout_ms: u64,
|
||||
}
|
||||
|
||||
/// 上游请求指纹默认配置。
|
||||
/// 经审计固定的 Codex Desktop 上游请求画像配置。
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct FingerprintConfig {
|
||||
/// 客户端来源名。
|
||||
pub struct WireProfileConfig {
|
||||
/// `originator` 请求头及 User-Agent 产品名。
|
||||
pub originator: String,
|
||||
/// 应用版本。
|
||||
pub app_version: String,
|
||||
/// 构建号。
|
||||
pub build_number: String,
|
||||
/// 平台名。
|
||||
pub platform: String,
|
||||
/// 架构名。
|
||||
/// bundled Codex Core 版本。
|
||||
pub codex_version: String,
|
||||
/// Desktop 应用版本。
|
||||
pub desktop_version: String,
|
||||
/// Desktop 制品构建号。
|
||||
pub desktop_build: String,
|
||||
/// Codex Core UA 中的目标操作系统类型。
|
||||
pub os_type: String,
|
||||
/// Codex Core UA 中的目标操作系统版本。
|
||||
pub os_version: String,
|
||||
/// Codex Core UA 中的目标架构。
|
||||
pub arch: String,
|
||||
/// Chromium 主版本。
|
||||
pub chromium_version: String,
|
||||
/// User-Agent 模板。
|
||||
pub user_agent_template: String,
|
||||
/// 默认请求头。
|
||||
pub default_headers: Vec<FingerprintHeaderConfig>,
|
||||
/// 请求头排序优先级。
|
||||
pub header_order: Vec<String>,
|
||||
}
|
||||
|
||||
/// 指纹默认请求头配置项。
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct FingerprintHeaderConfig {
|
||||
/// 请求头名称。
|
||||
pub name: String,
|
||||
/// 请求头值。
|
||||
pub value: String,
|
||||
/// Codex Core UA 中的终端标记。
|
||||
pub terminal: String,
|
||||
/// 此画像最后一次经制品与源码核验的时间。
|
||||
pub verified_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 管理员初始化与会话配置,不包含初始化密码。
|
||||
@@ -453,7 +447,7 @@ struct FileAppConfig {
|
||||
runtime: RuntimePathsConfig,
|
||||
quota: QuotaConfig,
|
||||
ws_pool: WebSocketPoolSettings,
|
||||
fingerprint: FingerprintConfig,
|
||||
wire_profile: WireProfileConfig,
|
||||
admin: FileAdminConfig,
|
||||
logging: LoggingConfig,
|
||||
telemetry: TelemetryConfig,
|
||||
@@ -549,7 +543,15 @@ fn validate_file_config(config: &FileAppConfig) -> Result<(), ConfigError> {
|
||||
if config.ws_pool.max_per_account == 0 {
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_per_account"));
|
||||
}
|
||||
validate_fingerprint(&config.fingerprint)?;
|
||||
if config.ws_pool.max_total == 0 {
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_total"));
|
||||
}
|
||||
if config.ws_pool.max_connecting == 0
|
||||
|| config.ws_pool.max_connecting > config.ws_pool.max_total
|
||||
{
|
||||
return Err(ConfigError::InvalidField("ws_pool.max_connecting"));
|
||||
}
|
||||
validate_wire_profile(&config.wire_profile)?;
|
||||
validate_logging(&config.logging)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -594,54 +596,48 @@ fn validate_admin_password(password: &str) -> Result<(), ConfigError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_fingerprint(config: &FingerprintConfig) -> Result<(), ConfigError> {
|
||||
fn validate_wire_profile(config: &WireProfileConfig) -> 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()),
|
||||
("wire_profile.originator", config.originator.as_str()),
|
||||
("wire_profile.codex_version", config.codex_version.as_str()),
|
||||
(
|
||||
"fingerprint.chromium_version",
|
||||
config.chromium_version.as_str(),
|
||||
),
|
||||
(
|
||||
"fingerprint.user_agent_template",
|
||||
config.user_agent_template.as_str(),
|
||||
"wire_profile.desktop_version",
|
||||
config.desktop_version.as_str(),
|
||||
),
|
||||
("wire_profile.desktop_build", config.desktop_build.as_str()),
|
||||
("wire_profile.os_type", config.os_type.as_str()),
|
||||
("wire_profile.os_version", config.os_version.as_str()),
|
||||
("wire_profile.arch", config.arch.as_str()),
|
||||
("wire_profile.terminal", config.terminal.as_str()),
|
||||
] {
|
||||
validate_nonempty(field, value)?;
|
||||
}
|
||||
|
||||
if config.default_headers.is_empty() || config.header_order.is_empty() {
|
||||
return Err(ConfigError::InvalidField("fingerprint"));
|
||||
if semver::Version::parse(&config.codex_version).is_err() {
|
||||
return Err(ConfigError::InvalidField("wire_profile.codex_version"));
|
||||
}
|
||||
|
||||
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"));
|
||||
}
|
||||
if !is_numeric_dotted_version(&config.desktop_version) {
|
||||
return Err(ConfigError::InvalidField("wire_profile.desktop_version"));
|
||||
}
|
||||
|
||||
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"));
|
||||
}
|
||||
if !config
|
||||
.desktop_build
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit())
|
||||
{
|
||||
return Err(ConfigError::InvalidField("wire_profile.desktop_build"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_numeric_dotted_version(value: &str) -> bool {
|
||||
let mut parts = value.split('.');
|
||||
let valid_parts = parts
|
||||
.by_ref()
|
||||
.filter(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit()))
|
||||
.count();
|
||||
valid_parts >= 2 && valid_parts == value.split('.').count()
|
||||
}
|
||||
|
||||
fn validate_logging(config: &LoggingConfig) -> Result<(), ConfigError> {
|
||||
EnvFilter::try_new(&config.level).map_err(|_| ConfigError::InvalidField("logging.level"))?;
|
||||
if !config.stdout && !config.file.enabled {
|
||||
|
||||
@@ -6,6 +6,14 @@ use axum::Router;
|
||||
use secrecy::ExposeSecret;
|
||||
|
||||
use crate::{
|
||||
api::admin::{
|
||||
accounts_routes::query::{AccountListQueryService, RefreshActivityQuery},
|
||||
dashboard_routes::{
|
||||
DashboardDesktopRelease, DashboardDesktopReleaseQuery, DashboardDesktopReleaseSnapshot,
|
||||
DashboardQueryService, DashboardQueryServiceParts, DashboardWireProfile,
|
||||
DashboardWireProfileQuery,
|
||||
},
|
||||
},
|
||||
auth::{
|
||||
service::SessionService,
|
||||
store::{PgAdminUserStore, RedisAdminSessionStore},
|
||||
@@ -18,11 +26,15 @@ use crate::{
|
||||
fleet::{
|
||||
cookies::PgCookieStore,
|
||||
manage::{AccountManageService, AccountManageServiceParts},
|
||||
pool::{AccountPoolOptions, AccountPoolService, AccountPoolStaticSettings},
|
||||
pool::{
|
||||
AccountPoolOptions, AccountPoolRuntimeOptions, AccountPoolService,
|
||||
AccountPoolStaticSettings, RotationStrategy,
|
||||
},
|
||||
refresh::{
|
||||
RedisRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, TokenRefreshService,
|
||||
},
|
||||
store::{AccountStore, PgAccountStore},
|
||||
usage::AccountUsageStore,
|
||||
},
|
||||
infra::{
|
||||
database::connect,
|
||||
@@ -33,6 +45,7 @@ use crate::{
|
||||
},
|
||||
keys::{manage::KeyManageService, service::KeyVerifier, store::PgClientKeyStore},
|
||||
models::{
|
||||
gateway::ModelCatalogSource,
|
||||
service::{ModelService, ModelServiceError},
|
||||
store::{ModelSnapshotStore, RedisModelSnapshotStore},
|
||||
types::ModelConfig,
|
||||
@@ -40,7 +53,7 @@ use crate::{
|
||||
settings::{SettingsSnapshot, service::SettingsService},
|
||||
telemetry::{
|
||||
account_usage::query::AccountUsageQueryService,
|
||||
account_usage::store::{AccountUsageStore, PgAccountUsageStore},
|
||||
account_usage::store::PgAccountUsageStore,
|
||||
buckets::store::PgRequestBucketStore,
|
||||
ops::query::OpsQueryService,
|
||||
ops::store::PgOpsErrorLogStore,
|
||||
@@ -50,11 +63,12 @@ use crate::{
|
||||
},
|
||||
update::service::SystemUpdateService,
|
||||
upstream::openai::{
|
||||
fingerprint::{Fingerprint, PgFingerprintStore, RuntimeFingerprint},
|
||||
desktop_release::DesktopReleaseStatus,
|
||||
profile::{CodexWireProfile, CodexWireProfileState},
|
||||
token_client::{OpenAiTokenClient, TokenClientConfig, default_openai_token_client},
|
||||
transport::{
|
||||
CodexBackendClient, CodexModelCatalogClient, CodexWebSocketPool,
|
||||
CodexWebSocketPoolConfig, build_reqwest_client, tls::CustomCaError,
|
||||
CodexBackendClient, CodexWebSocketPool, CodexWebSocketPoolConfig, build_reqwest_client,
|
||||
tls::CustomCaError,
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -80,7 +94,6 @@ pub struct BackgroundTaskStores {
|
||||
pub admin_users: PgAdminUserStore,
|
||||
pub admin_sessions: RedisAdminSessionStore,
|
||||
pub cookies: PgCookieStore,
|
||||
pub fingerprints: PgFingerprintStore,
|
||||
pub session_affinity: RedisSessionAffinityStore,
|
||||
pub refresh_leases: RedisRefreshLeaseStore,
|
||||
pub model_snapshots: RedisModelSnapshotStore,
|
||||
@@ -117,7 +130,8 @@ pub struct Services {
|
||||
pub session_affinity: Arc<SessionAffinityService>,
|
||||
pub codex: Arc<CodexBackendClient>,
|
||||
pub websocket_pool: Option<Arc<CodexWebSocketPool>>,
|
||||
pub fingerprint: RuntimeFingerprint,
|
||||
pub wire_profile: CodexWireProfileState,
|
||||
pub desktop_release: DesktopReleaseStatus,
|
||||
pub account_pseudonymizer: Arc<AccountPseudonymizer>,
|
||||
pub system_update: Arc<SystemUpdateService>,
|
||||
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
|
||||
@@ -141,25 +155,64 @@ impl UsageRecordOptions {
|
||||
}
|
||||
}
|
||||
|
||||
impl DashboardWireProfileQuery for CodexWireProfileState {
|
||||
fn snapshot(&self) -> DashboardWireProfile {
|
||||
let profile = CodexWireProfileState::snapshot(self);
|
||||
let user_agent = profile.user_agent();
|
||||
DashboardWireProfile {
|
||||
originator: profile.originator,
|
||||
codex_version: profile.codex_version,
|
||||
desktop_version: profile.desktop_version,
|
||||
desktop_build: profile.desktop_build,
|
||||
os_type: profile.os_type,
|
||||
os_version: profile.os_version,
|
||||
arch: profile.arch,
|
||||
terminal: profile.terminal,
|
||||
user_agent,
|
||||
verified_at: profile.verified_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DashboardDesktopReleaseQuery for DesktopReleaseStatus {
|
||||
fn snapshot(&self) -> DashboardDesktopReleaseSnapshot {
|
||||
let snapshot = DesktopReleaseStatus::snapshot(self);
|
||||
DashboardDesktopReleaseSnapshot {
|
||||
checked_at: snapshot.checked_at,
|
||||
latest: snapshot.latest.map(|release| DashboardDesktopRelease {
|
||||
version: release.version,
|
||||
build: release.build,
|
||||
published_at: release.published_at,
|
||||
minimum_system_version: release.minimum_system_version,
|
||||
hardware_requirements: release.hardware_requirements,
|
||||
download_url: release.download_url,
|
||||
download_size: release.download_size,
|
||||
signature_present: release.signature_present,
|
||||
}),
|
||||
last_error: snapshot.last_error,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Services {
|
||||
pub fn new(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
) -> Self {
|
||||
Self::try_new(config, stores, fingerprint)
|
||||
Self::try_new(config, stores, wire_profile)
|
||||
.expect("failed to build runtime services with configured TLS transport")
|
||||
}
|
||||
|
||||
pub fn try_new(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
Self::try_with_usage_record_options(
|
||||
config,
|
||||
stores,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
UsageRecordOptions::from_config(config),
|
||||
)
|
||||
}
|
||||
@@ -167,7 +220,7 @@ impl Services {
|
||||
pub fn try_with_usage_record_options(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
usage_record_options: UsageRecordOptions,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
let mut identity_secret = [0u8; 32];
|
||||
@@ -175,7 +228,7 @@ impl Services {
|
||||
Self::try_with_identity_secret_and_usage_record_options(
|
||||
config,
|
||||
stores,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
identity_secret,
|
||||
usage_record_options,
|
||||
)
|
||||
@@ -184,7 +237,7 @@ impl Services {
|
||||
fn try_with_identity_secret_and_usage_record_options(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
wire_profile: CodexWireProfileState,
|
||||
identity_secret: [u8; 32],
|
||||
usage_record_options: UsageRecordOptions,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
@@ -202,6 +255,8 @@ impl Services {
|
||||
enabled: true,
|
||||
max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms),
|
||||
max_per_account: config.ws_pool.max_per_account,
|
||||
max_total: config.ws_pool.max_total,
|
||||
max_connecting: config.ws_pool.max_connecting,
|
||||
initial_event_timeout: (config.ws_pool.initial_event_timeout_ms > 0).then(|| {
|
||||
std::time::Duration::from_millis(config.ws_pool.initial_event_timeout_ms)
|
||||
}),
|
||||
@@ -214,7 +269,7 @@ impl Services {
|
||||
let client = CodexBackendClient::new(
|
||||
build_reqwest_client()?,
|
||||
config.api.base_url.clone(),
|
||||
fingerprint.clone(),
|
||||
wire_profile.clone(),
|
||||
)
|
||||
.with_websocket_initial_event_timeout(websocket_initial_event_timeout);
|
||||
if let Some(pool) = &websocket_pool {
|
||||
@@ -266,9 +321,10 @@ impl Services {
|
||||
refresh_policy.clone(),
|
||||
token_client.clone(),
|
||||
)
|
||||
.with_refresh_lease_store(stores.refresh_leases.clone()),
|
||||
.with_refresh_lease_store(stores.refresh_leases.clone())
|
||||
.with_account_pool_sync(account_pool.clone()),
|
||||
);
|
||||
let upstream_client: Arc<dyn CodexModelCatalogClient> = codex.clone();
|
||||
let upstream_client: Arc<dyn ModelCatalogSource> = codex.clone();
|
||||
let store: Arc<dyn ModelSnapshotStore> = Arc::new(stores.model_snapshots.clone());
|
||||
let models = Arc::new(ModelService::new(
|
||||
ModelConfig {
|
||||
@@ -280,7 +336,7 @@ impl Services {
|
||||
let admin_accounts = Arc::new(AccountManageService::new(AccountManageServiceParts {
|
||||
store: stores.accounts.clone(),
|
||||
cookies: stores.cookies.clone(),
|
||||
codex: codex.clone(),
|
||||
upstream: codex.clone(),
|
||||
models: models.clone(),
|
||||
account_pool: account_pool.clone(),
|
||||
token_refresher: Arc::new(token_client),
|
||||
@@ -313,6 +369,8 @@ impl Services {
|
||||
}));
|
||||
let system_update = Arc::new(SystemUpdateService::from_env());
|
||||
|
||||
let desktop_release = DesktopReleaseStatus::default();
|
||||
|
||||
Ok(Self {
|
||||
database,
|
||||
redis,
|
||||
@@ -333,7 +391,8 @@ impl Services {
|
||||
session_affinity,
|
||||
codex,
|
||||
websocket_pool,
|
||||
fingerprint,
|
||||
wire_profile,
|
||||
desktop_release,
|
||||
account_pseudonymizer,
|
||||
system_update,
|
||||
connection_drain,
|
||||
@@ -351,13 +410,30 @@ impl Services {
|
||||
|
||||
impl From<&Services> for crate::api::router::ApiServices {
|
||||
fn from(services: &Services) -> Self {
|
||||
let refresh_activity = services.token_refresh.clone() as Arc<dyn RefreshActivityQuery>;
|
||||
let account_list = Arc::new(AccountListQueryService::new(
|
||||
services.admin_accounts.clone(),
|
||||
services.usage.clone(),
|
||||
refresh_activity.clone(),
|
||||
));
|
||||
let dashboard = Arc::new(DashboardQueryService::new(DashboardQueryServiceParts {
|
||||
accounts: services.accounts.clone(),
|
||||
usage: services.usage.clone(),
|
||||
usage_records: services.usage_records.clone(),
|
||||
account_pool: services.account_pool.clone(),
|
||||
refresh_activity,
|
||||
settings: services.settings.clone(),
|
||||
wire_profile: Arc::new(services.wire_profile.clone()),
|
||||
desktop_release: Arc::new(services.desktop_release.clone()),
|
||||
}));
|
||||
Self {
|
||||
account_list,
|
||||
dashboard,
|
||||
health_probe: Arc::new(crate::bootstrap::state::RuntimeHealthProbe::new(
|
||||
services.database.clone(),
|
||||
services.redis.clone(),
|
||||
)),
|
||||
models: services.models.clone(),
|
||||
accounts: services.accounts.clone(),
|
||||
client_keys: services.client_keys.clone(),
|
||||
admin_client_keys: services.admin_client_keys.clone(),
|
||||
admin_sessions: services.admin_sessions.clone(),
|
||||
@@ -365,12 +441,9 @@ impl From<&Services> for crate::api::router::ApiServices {
|
||||
admin_accounts: services.admin_accounts.clone(),
|
||||
usage_records: services.usage_records.clone(),
|
||||
ops_errors: services.ops_errors.clone(),
|
||||
usage: services.usage.clone(),
|
||||
account_pool: services.account_pool.clone(),
|
||||
token_refresh: services.token_refresh.clone(),
|
||||
responses: services.responses.clone(),
|
||||
session_affinity: services.session_affinity.clone(),
|
||||
fingerprint: services.fingerprint.clone(),
|
||||
process_control: Arc::new(crate::bootstrap::shutdown::RuntimeProcessControl),
|
||||
system_update: services.system_update.clone(),
|
||||
connection_drain: services.connection_drain.clone(),
|
||||
@@ -419,33 +492,50 @@ pub fn apply_settings_to_config(config: &mut AppConfig, settings: &SettingsSnaps
|
||||
config.auth.rotation_strategy = settings.rotation_strategy.clone();
|
||||
}
|
||||
|
||||
pub fn fingerprint_from_config(
|
||||
config: &crate::bootstrap::config::FingerprintConfig,
|
||||
) -> Fingerprint {
|
||||
Fingerprint {
|
||||
pub fn wire_profile_from_config(
|
||||
config: &crate::bootstrap::config::WireProfileConfig,
|
||||
) -> CodexWireProfile {
|
||||
CodexWireProfile {
|
||||
originator: config.originator.clone(),
|
||||
app_version: config.app_version.clone(),
|
||||
build_number: config.build_number.clone(),
|
||||
platform: config.platform.clone(),
|
||||
codex_version: config.codex_version.clone(),
|
||||
desktop_version: config.desktop_version.clone(),
|
||||
desktop_build: config.desktop_build.clone(),
|
||||
os_type: config.os_type.clone(),
|
||||
os_version: config.os_version.clone(),
|
||||
arch: config.arch.clone(),
|
||||
chromium_version: config.chromium_version.clone(),
|
||||
user_agent_template: config.user_agent_template.clone(),
|
||||
default_headers: config
|
||||
.default_headers
|
||||
.iter()
|
||||
.map(|header| (header.name.clone(), header.value.clone()))
|
||||
.collect(),
|
||||
header_order: config.header_order.clone(),
|
||||
updated_at: None,
|
||||
terminal: config.terminal.clone(),
|
||||
verified_at: config.verified_at,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
|
||||
AccountPoolStaticSettings {
|
||||
let static_settings = AccountPoolStaticSettings {
|
||||
skip_quota_limited: config.quota.skip_exhausted,
|
||||
tier_priority: config.auth.tier_priority.clone(),
|
||||
};
|
||||
account_pool_runtime_options(&settings_snapshot_from_config(config), &static_settings).options
|
||||
}
|
||||
|
||||
pub fn account_pool_runtime_options(
|
||||
settings: &SettingsSnapshot,
|
||||
static_settings: &AccountPoolStaticSettings,
|
||||
) -> AccountPoolRuntimeOptions {
|
||||
AccountPoolRuntimeOptions {
|
||||
options: AccountPoolOptions {
|
||||
max_concurrent_per_account: settings.max_concurrent_per_account,
|
||||
rotation_strategy: match settings.rotation_strategy.as_str() {
|
||||
"smart" => RotationStrategy::Smart,
|
||||
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
|
||||
"round_robin" => RotationStrategy::RoundRobin,
|
||||
"sticky" => RotationStrategy::Sticky,
|
||||
_ => RotationStrategy::Smart,
|
||||
},
|
||||
skip_quota_limited: static_settings.skip_quota_limited,
|
||||
tier_priority: static_settings.tier_priority.clone(),
|
||||
..AccountPoolOptions::default()
|
||||
},
|
||||
request_interval_ms: settings.request_interval_ms,
|
||||
}
|
||||
.pool_options(&settings_snapshot_from_config(config))
|
||||
}
|
||||
|
||||
pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
@@ -457,6 +547,7 @@ pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + S
|
||||
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);
|
||||
let mut restart_path = None;
|
||||
|
||||
let serve_result = {
|
||||
let (graceful_shutdown_tx, graceful_shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
@@ -471,46 +562,63 @@ pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + S
|
||||
tokio::pin!(serve);
|
||||
tokio::select! {
|
||||
result = &mut serve => result,
|
||||
() = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
let _ = graceful_shutdown_tx.send(());
|
||||
let active_websocket_connections = connection_drain.begin_shutdown();
|
||||
tracing::info!(
|
||||
active_websocket_connections,
|
||||
"正在排空入站连接并关闭运行时任务"
|
||||
);
|
||||
let coordinator = task_coordinator.take();
|
||||
let connection_drain = connection_drain.clone();
|
||||
let runtime_shutdown = async move {
|
||||
let shutdown_tasks = async move {
|
||||
if let Some(coordinator) = coordinator {
|
||||
coordinator.shutdown().await;
|
||||
}
|
||||
};
|
||||
tokio::join!(shutdown_tasks, connection_drain.wait());
|
||||
};
|
||||
let drain_runtime = async {
|
||||
let (serve_result, ()) = tokio::join!(&mut serve, runtime_shutdown);
|
||||
serve_result
|
||||
};
|
||||
tokio::pin!(drain_runtime);
|
||||
tokio::select! {
|
||||
result = &mut drain_runtime => result,
|
||||
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
|
||||
tracing::warn!(
|
||||
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
|
||||
"HTTP graceful shutdown timed out; cancelling remaining connections"
|
||||
);
|
||||
action = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
match action {
|
||||
crate::bootstrap::shutdown::ShutdownAction::Restart(executable_path) => {
|
||||
restart_path = Some(executable_path);
|
||||
Ok(())
|
||||
}
|
||||
() = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
tracing::warn!("Received a second shutdown signal; cancelling remaining connections");
|
||||
Ok(())
|
||||
crate::bootstrap::shutdown::ShutdownAction::Graceful => {
|
||||
let _ = graceful_shutdown_tx.send(());
|
||||
let active_websocket_connections = connection_drain.begin_shutdown();
|
||||
tracing::info!(
|
||||
active_websocket_connections,
|
||||
"正在排空入站连接并关闭运行时任务"
|
||||
);
|
||||
let coordinator = task_coordinator.take();
|
||||
let connection_drain = connection_drain.clone();
|
||||
let runtime_shutdown = async move {
|
||||
let shutdown_tasks = async move {
|
||||
if let Some(coordinator) = coordinator {
|
||||
coordinator.shutdown().await;
|
||||
}
|
||||
};
|
||||
tokio::join!(shutdown_tasks, connection_drain.wait());
|
||||
};
|
||||
let drain_runtime = async {
|
||||
let (serve_result, ()) = tokio::join!(&mut serve, runtime_shutdown);
|
||||
serve_result
|
||||
};
|
||||
tokio::pin!(drain_runtime);
|
||||
tokio::select! {
|
||||
result = &mut drain_runtime => result,
|
||||
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
|
||||
tracing::warn!(
|
||||
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
|
||||
"HTTP graceful shutdown timed out; cancelling remaining connections"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
_ = crate::bootstrap::shutdown::shutdown_signal() => {
|
||||
tracing::warn!("Received a second shutdown signal; cancelling remaining connections");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(executable_path) = restart_path {
|
||||
tracing::info!(
|
||||
executable = %executable_path.display(),
|
||||
"正在以更新后的二进制替换当前进程"
|
||||
);
|
||||
drop(log_guard);
|
||||
return exec_replacement_process(executable_path);
|
||||
}
|
||||
|
||||
if let Some(coordinator) = task_coordinator.take() {
|
||||
let active_websocket_connections = connection_drain.begin_shutdown();
|
||||
tracing::info!(
|
||||
@@ -521,15 +629,6 @@ pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + S
|
||||
}
|
||||
serve_result?;
|
||||
|
||||
if let Some(executable_path) = crate::bootstrap::shutdown::restart_executable_path() {
|
||||
tracing::info!(
|
||||
executable = %executable_path.display(),
|
||||
"正在以更新后的二进制替换当前进程"
|
||||
);
|
||||
drop(log_guard);
|
||||
exec_replacement_process(executable_path)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -560,14 +659,12 @@ async fn build_router(
|
||||
SettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool).await?;
|
||||
apply_settings_to_config(&mut config, &settings);
|
||||
|
||||
let fingerprint_store = PgFingerprintStore::new(pool.clone());
|
||||
let stores = BackgroundTaskStores {
|
||||
redis: redis.clone(),
|
||||
accounts: PgAccountStore::new(pool.clone()),
|
||||
admin_users: PgAdminUserStore::new(pool.clone()),
|
||||
admin_sessions: RedisAdminSessionStore::new(redis.clone()),
|
||||
cookies: PgCookieStore::new(pool.clone()),
|
||||
fingerprints: fingerprint_store.clone(),
|
||||
session_affinity: RedisSessionAffinityStore::new(redis.clone()),
|
||||
refresh_leases: RedisRefreshLeaseStore::new(redis.clone()),
|
||||
model_snapshots: RedisModelSnapshotStore::new(redis.clone()),
|
||||
@@ -580,16 +677,12 @@ async fn build_router(
|
||||
|
||||
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 wire_profile = CodexWireProfileState::new(wire_profile_from_config(&config.wire_profile));
|
||||
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(&config);
|
||||
let services = Services::try_with_identity_secret_and_usage_record_options(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint,
|
||||
wire_profile,
|
||||
identity_secret,
|
||||
UsageRecordOptions::from_config(&config),
|
||||
)?;
|
||||
|
||||
@@ -1,19 +1,17 @@
|
||||
//! 优雅关闭信号处理。
|
||||
//! 进程关闭与重启信号处理。
|
||||
|
||||
use std::{
|
||||
path::PathBuf,
|
||||
sync::{
|
||||
OnceLock,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
};
|
||||
use std::{path::PathBuf, sync::OnceLock};
|
||||
|
||||
use tokio::signal;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<()>> = OnceLock::new();
|
||||
static RESTART_REQUESTED: AtomicBool = AtomicBool::new(false);
|
||||
static RESTART_EXECUTABLE_PATH: OnceLock<PathBuf> = OnceLock::new();
|
||||
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<ShutdownAction>> = OnceLock::new();
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum ShutdownAction {
|
||||
Graceful,
|
||||
Restart(PathBuf),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct RuntimeProcessControl;
|
||||
@@ -24,12 +22,12 @@ impl crate::api::router::ProcessControl for RuntimeProcessControl {
|
||||
}
|
||||
|
||||
fn request_restart(&self, executable_path: PathBuf) {
|
||||
request_process_restart(executable_path);
|
||||
request_immediate_restart(executable_path);
|
||||
}
|
||||
}
|
||||
|
||||
/// 等待进程关闭信号(Ctrl+C 或 SIGTERM)。
|
||||
pub async fn shutdown_signal() {
|
||||
/// 等待进程关闭或重启信号。
|
||||
pub(crate) async fn shutdown_signal() -> ShutdownAction {
|
||||
let mut requested = shutdown_sender().subscribe();
|
||||
let ctrl_c = async {
|
||||
signal::ctrl_c()
|
||||
@@ -49,33 +47,25 @@ pub async fn shutdown_signal() {
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = requested.recv() => {},
|
||||
() = ctrl_c => {},
|
||||
() = terminate => {},
|
||||
action = requested.recv() => match action {
|
||||
Ok(action) => action,
|
||||
Err(_) => ShutdownAction::Graceful,
|
||||
},
|
||||
() = ctrl_c => ShutdownAction::Graceful,
|
||||
() = terminate => ShutdownAction::Graceful,
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求进程按优雅关闭路径退出。
|
||||
pub fn request_shutdown() {
|
||||
let _ = shutdown_sender().send(());
|
||||
let _ = shutdown_sender().send(ShutdownAction::Graceful);
|
||||
}
|
||||
|
||||
/// 请求进程完成优雅关闭后以新二进制替换当前进程。
|
||||
pub fn request_process_restart(executable_path: PathBuf) {
|
||||
let _ = RESTART_EXECUTABLE_PATH.set(executable_path);
|
||||
RESTART_REQUESTED.store(true, Ordering::SeqCst);
|
||||
request_shutdown();
|
||||
fn request_immediate_restart(executable_path: PathBuf) {
|
||||
let _ = shutdown_sender().send(ShutdownAction::Restart(executable_path));
|
||||
}
|
||||
|
||||
/// 获取待 exec 的重启目标。
|
||||
pub fn restart_executable_path() -> Option<PathBuf> {
|
||||
RESTART_REQUESTED
|
||||
.load(Ordering::SeqCst)
|
||||
.then(|| RESTART_EXECUTABLE_PATH.get().cloned())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
fn shutdown_sender() -> &'static broadcast::Sender<()> {
|
||||
fn shutdown_sender() -> &'static broadcast::Sender<ShutdownAction> {
|
||||
SHUTDOWN_REQUESTS.get_or_init(|| {
|
||||
let (sender, _receiver) = broadcast::channel(16);
|
||||
sender
|
||||
|
||||
@@ -1,19 +1,13 @@
|
||||
//! 后台任务协调器。
|
||||
|
||||
use std::{future::Future, path::PathBuf, sync::Arc, time::Duration};
|
||||
use std::{future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
use super::{
|
||||
cookie_cleanup::CookieCleanupTask,
|
||||
fingerprint_update::{
|
||||
CODEX_DESKTOP_APPCAST_URL, DEFAULT_EXTRACTED_FINGERPRINT_PATH,
|
||||
start_fingerprint_update_task,
|
||||
},
|
||||
model_refresh::ModelRefreshTask,
|
||||
quota_refresh::QuotaRefreshTask,
|
||||
retention_trim::RetentionTrimTask,
|
||||
token_refresh::TokenRefreshTask,
|
||||
cookie_cleanup::CookieCleanupTask, desktop_release_update::start_desktop_release_update_task,
|
||||
model_refresh::ModelRefreshTask, quota_refresh::QuotaRefreshTask,
|
||||
retention_trim::RetentionTrimTask, token_refresh::TokenRefreshTask,
|
||||
};
|
||||
|
||||
/// 后台任务协调器。
|
||||
@@ -46,7 +40,6 @@ impl TaskCoordinator {
|
||||
stores.usage_records.clone(),
|
||||
stores.ops_errors.clone(),
|
||||
stores.request_buckets.clone(),
|
||||
stores.fingerprints.clone(),
|
||||
)
|
||||
.start(),
|
||||
);
|
||||
@@ -94,17 +87,14 @@ impl TaskCoordinator {
|
||||
.with_cookie_store(stores.cookies.clone())
|
||||
.start(),
|
||||
);
|
||||
coordinator.push("fingerprint_update", {
|
||||
let fingerprint = services.fingerprint.snapshot();
|
||||
start_fingerprint_update_task(
|
||||
stores.fingerprints.clone(),
|
||||
services.fingerprint.clone(),
|
||||
CODEX_DESKTOP_APPCAST_URL.to_string(),
|
||||
PathBuf::from(DEFAULT_EXTRACTED_FINGERPRINT_PATH),
|
||||
fingerprint.app_version,
|
||||
fingerprint.build_number,
|
||||
)
|
||||
});
|
||||
coordinator.push(
|
||||
"desktop_release_update",
|
||||
start_desktop_release_update_task(
|
||||
services.desktop_release.clone(),
|
||||
services.wire_profile.clone(),
|
||||
crate::upstream::openai::desktop_release::CODEX_DESKTOP_APPCAST_URL.to_string(),
|
||||
),
|
||||
);
|
||||
if let Some(pool) = &services.websocket_pool {
|
||||
coordinator.push(
|
||||
"websocket_pool",
|
||||
@@ -118,30 +108,52 @@ impl TaskCoordinator {
|
||||
/// 启动各领域的设置订阅任务。
|
||||
pub fn start_settings_subscriptions(services: &crate::bootstrap::services::Services) -> Self {
|
||||
let mut coordinator = TaskCoordinator::default();
|
||||
let mut settings = services.settings.subscribe();
|
||||
let initial = settings.borrow().clone();
|
||||
let (model_tx, model_rx) = tokio::sync::watch::channel(crate::models::types::ModelConfig {
|
||||
model_aliases: initial.model_aliases.clone(),
|
||||
});
|
||||
let (pool_tx, pool_rx) =
|
||||
tokio::sync::watch::channel(crate::bootstrap::services::account_pool_runtime_options(
|
||||
&initial,
|
||||
&services.account_pool_static,
|
||||
));
|
||||
let (refresh_tx, refresh_rx) =
|
||||
tokio::sync::watch::channel(crate::fleet::refresh::RefreshPolicy {
|
||||
refresh_margin_seconds: initial.refresh_margin_seconds,
|
||||
refresh_concurrency: initial.refresh_concurrency,
|
||||
});
|
||||
coordinator.push(
|
||||
"model_settings",
|
||||
SchedulerHandle::spawn(
|
||||
services
|
||||
.models
|
||||
.clone()
|
||||
.subscribe_settings(services.settings.subscribe()),
|
||||
),
|
||||
SchedulerHandle::spawn(services.models.clone().subscribe_config(model_rx)),
|
||||
);
|
||||
coordinator.push(
|
||||
"account_pool_settings",
|
||||
SchedulerHandle::spawn(services.account_pool.clone().subscribe_settings(
|
||||
services.settings.subscribe(),
|
||||
services.account_pool_static.clone(),
|
||||
)),
|
||||
SchedulerHandle::spawn(services.account_pool.clone().subscribe_options(pool_rx)),
|
||||
);
|
||||
coordinator.push(
|
||||
"refresh_policy_settings",
|
||||
SchedulerHandle::spawn(
|
||||
services
|
||||
.refresh_policy
|
||||
.clone()
|
||||
.subscribe_settings(services.settings.subscribe()),
|
||||
),
|
||||
SchedulerHandle::spawn(services.refresh_policy.clone().subscribe_policy(refresh_rx)),
|
||||
);
|
||||
let static_settings = services.account_pool_static.clone();
|
||||
coordinator.push(
|
||||
"settings_fanout",
|
||||
SchedulerHandle::spawn(async move {
|
||||
while settings.changed().await.is_ok() {
|
||||
let snapshot = settings.borrow_and_update().clone();
|
||||
model_tx.send_replace(crate::models::types::ModelConfig {
|
||||
model_aliases: snapshot.model_aliases.clone(),
|
||||
});
|
||||
pool_tx.send_replace(crate::bootstrap::services::account_pool_runtime_options(
|
||||
&snapshot,
|
||||
&static_settings,
|
||||
));
|
||||
refresh_tx.send_replace(crate::fleet::refresh::RefreshPolicy {
|
||||
refresh_margin_seconds: snapshot.refresh_margin_seconds,
|
||||
refresh_concurrency: snapshot.refresh_concurrency,
|
||||
});
|
||||
}
|
||||
}),
|
||||
);
|
||||
coordinator
|
||||
}
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
//! Codex Desktop 官方发布观测任务接线。
|
||||
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::upstream::openai::desktop_release::{
|
||||
APPCAST_POLL_INTERVAL, DesktopReleaseChecker, DesktopReleaseStatus,
|
||||
};
|
||||
use crate::upstream::openai::profile::CodexWireProfileState;
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
|
||||
};
|
||||
|
||||
/// 启动 Desktop 官方发布检查后台任务。
|
||||
pub fn start_desktop_release_update_task(
|
||||
status: DesktopReleaseStatus,
|
||||
wire_profile: CodexWireProfileState,
|
||||
appcast_url: String,
|
||||
) -> SchedulerHandle {
|
||||
let checker = DesktopReleaseChecker::with_client(
|
||||
reqwest::Client::new(),
|
||||
appcast_url,
|
||||
status,
|
||||
wire_profile,
|
||||
);
|
||||
spawn_periodic_task(
|
||||
DesktopReleaseUpdateTask {
|
||||
checker,
|
||||
first_tick: true,
|
||||
},
|
||||
PeriodicTaskConfig::new(
|
||||
APPCAST_POLL_INTERVAL.as_secs(),
|
||||
"Codex Desktop 发布检查器已启动",
|
||||
"Codex Desktop 发布检查器已关闭",
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
struct DesktopReleaseUpdateTask {
|
||||
checker: DesktopReleaseChecker,
|
||||
first_tick: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for DesktopReleaseUpdateTask {
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let first_tick = std::mem::take(&mut self.first_tick);
|
||||
match self.checker.check_and_record().await {
|
||||
Ok(release) => info!(
|
||||
version = %release.version,
|
||||
build = %release.build,
|
||||
"Codex Desktop 最新发布信息已刷新并已同步请求身份"
|
||||
),
|
||||
Err(error) if first_tick => {
|
||||
warn!(error = %error, "Codex Desktop 首次发布检查失败");
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = %error, "Codex Desktop 定期发布检查失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,85 +0,0 @@
|
||||
//! 指纹更新任务接线。
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::upstream::openai::fingerprint::{
|
||||
APPCAST_POLL_INTERVAL, PgFingerprintStore, RuntimeFingerprint, UpdateChecker,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
|
||||
};
|
||||
|
||||
/// Codex Desktop 官方 appcast 地址。
|
||||
pub const CODEX_DESKTOP_APPCAST_URL: &str =
|
||||
"https://persistent.oaistatic.com/codex-app-prod/appcast.xml";
|
||||
/// 本地可选完整指纹提取文件路径。
|
||||
pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint.json";
|
||||
|
||||
/// 启动指纹自动更新后台任务。
|
||||
pub fn start_fingerprint_update_task(
|
||||
store: PgFingerprintStore,
|
||||
runtime_fingerprint: RuntimeFingerprint,
|
||||
appcast_url: String,
|
||||
extracted_fingerprint_path: PathBuf,
|
||||
current_version: String,
|
||||
current_build: String,
|
||||
) -> SchedulerHandle {
|
||||
let update_checker = UpdateChecker::with_client(
|
||||
store,
|
||||
reqwest::Client::new(),
|
||||
appcast_url,
|
||||
extracted_fingerprint_path,
|
||||
current_version,
|
||||
current_build,
|
||||
);
|
||||
spawn_periodic_task(
|
||||
FingerprintUpdateTask {
|
||||
checker: update_checker,
|
||||
runtime_fingerprint,
|
||||
first_tick: true,
|
||||
},
|
||||
PeriodicTaskConfig::new(
|
||||
APPCAST_POLL_INTERVAL.as_secs(),
|
||||
"Fingerprint 后台版本检查器已启动",
|
||||
"Fingerprint 后台版本检查器已关闭",
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
struct FingerprintUpdateTask {
|
||||
checker: UpdateChecker,
|
||||
runtime_fingerprint: RuntimeFingerprint,
|
||||
first_tick: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for FingerprintUpdateTask {
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let first_tick = std::mem::take(&mut self.first_tick);
|
||||
match self.checker.check_and_apply_update().await {
|
||||
Ok(Some(fingerprint)) => {
|
||||
let app_version = fingerprint.app_version.clone();
|
||||
let build_number = fingerprint.build_number.clone();
|
||||
self.runtime_fingerprint.replace(fingerprint);
|
||||
info!(
|
||||
app_version = %app_version,
|
||||
build_number = %build_number,
|
||||
"Fingerprint 运行时版本已更新"
|
||||
);
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(error) => {
|
||||
if first_tick {
|
||||
warn!(error = %error, "Fingerprint 首次版本检查失败");
|
||||
} else {
|
||||
warn!(error = %error, "Fingerprint 定期版本检查失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@
|
||||
mod cleanup;
|
||||
pub mod cookie_cleanup;
|
||||
pub mod coordinator;
|
||||
pub mod fingerprint_update;
|
||||
pub mod desktop_release_update;
|
||||
pub mod model_refresh;
|
||||
mod periodic;
|
||||
pub mod quota_refresh;
|
||||
|
||||
@@ -13,23 +13,18 @@ use crate::models::service::{
|
||||
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup, spawn_periodic_task},
|
||||
};
|
||||
use super::coordinator::SchedulerHandle;
|
||||
|
||||
/// 模型刷新任务接线器。
|
||||
pub struct ModelRefreshTask {
|
||||
model_service: Arc<ModelService>,
|
||||
account_pool: Arc<AccountPoolService>,
|
||||
refresh_interval_secs: u64,
|
||||
initial_delay_ms: u64,
|
||||
retry_delay_ms: u64,
|
||||
max_retries: u32,
|
||||
account_pseudonymizer: Arc<AccountPseudonymizer>,
|
||||
}
|
||||
|
||||
const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600;
|
||||
const INITIAL_DELAY_MS: u64 = 1000;
|
||||
const RETRY_DELAY_MS: u64 = 10_000;
|
||||
const MAX_RETRIES: u32 = 12;
|
||||
@@ -44,7 +39,6 @@ impl ModelRefreshTask {
|
||||
Self {
|
||||
model_service,
|
||||
account_pool,
|
||||
refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS,
|
||||
initial_delay_ms: INITIAL_DELAY_MS,
|
||||
retry_delay_ms: RETRY_DELAY_MS,
|
||||
max_retries: MAX_RETRIES,
|
||||
@@ -52,15 +46,13 @@ impl ModelRefreshTask {
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动后台刷新任务。
|
||||
/// 启动一次性首次刷新任务。
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let config = PeriodicTaskConfig::new(
|
||||
self.refresh_interval_secs,
|
||||
"模型刷新任务已启动",
|
||||
"模型刷新任务正在关闭",
|
||||
)
|
||||
.wait_first_interval();
|
||||
spawn_periodic_task(self, config)
|
||||
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
|
||||
let handle = tokio::spawn(async move {
|
||||
self.run_initial_refresh(&mut shutdown_rx).await;
|
||||
});
|
||||
SchedulerHandle::new(shutdown_tx, handle)
|
||||
}
|
||||
|
||||
/// 在 Responses 响应声明模型目录 ETag 变化时异步刷新。
|
||||
@@ -115,84 +107,52 @@ impl ModelRefreshTask {
|
||||
.await;
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for ModelRefreshTask {
|
||||
fn before_loop<'a>(
|
||||
&'a mut self,
|
||||
shutdown_rx: &'a mut mpsc::Receiver<()>,
|
||||
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
|
||||
Box::pin(async move {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return PeriodicTaskStartup::Stop;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
|
||||
async fn run_initial_refresh(&self, shutdown_rx: &mut mpsc::Receiver<()>) {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
|
||||
}
|
||||
|
||||
let mut attempt = 0;
|
||||
let mut has_fetched_once = false;
|
||||
let mut deferred_until_accounts_available = false;
|
||||
|
||||
while attempt < self.max_retries {
|
||||
match self.refresh_once().await {
|
||||
Ok(result) => {
|
||||
info!(
|
||||
refreshed_plans = result.refreshed_plans,
|
||||
model_count = result.model_count,
|
||||
"首次模型刷新成功"
|
||||
);
|
||||
has_fetched_once = true;
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
attempt += 1;
|
||||
if matches!(error, ModelServiceError::NoAccounts) {
|
||||
info!("首次模型刷新跳过:没有可用账户,将进入周期重试");
|
||||
deferred_until_accounts_available = true;
|
||||
break;
|
||||
}
|
||||
|
||||
warn!(
|
||||
attempt,
|
||||
max_retries = self.max_retries,
|
||||
error = ?error,
|
||||
"首次模型刷新失败"
|
||||
);
|
||||
|
||||
if attempt < self.max_retries {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return PeriodicTaskStartup::Stop;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !has_fetched_once && !deferred_until_accounts_available {
|
||||
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
|
||||
}
|
||||
|
||||
PeriodicTaskStartup::Continue
|
||||
})
|
||||
}
|
||||
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let mut attempt = 0;
|
||||
while attempt < self.max_retries {
|
||||
match self.refresh_once().await {
|
||||
Ok(_) => {}
|
||||
Ok(result) => {
|
||||
info!(
|
||||
refreshed_plans = result.refreshed_plans,
|
||||
model_count = result.model_count,
|
||||
"首次模型刷新成功"
|
||||
);
|
||||
return;
|
||||
}
|
||||
Err(ModelServiceError::NoAccounts) => {
|
||||
info!("模型刷新跳过:没有可用账户");
|
||||
info!("首次模型刷新跳过:没有可用账户,后续由 ETag 或手动操作触发");
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = ?error, "刷新模型列表失败");
|
||||
attempt += 1;
|
||||
warn!(
|
||||
attempt,
|
||||
max_retries = self.max_retries,
|
||||
error = ?error,
|
||||
"首次模型刷新失败"
|
||||
);
|
||||
|
||||
if attempt < self.max_retries {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
warn!("模型刷新任务首次尝试全部失败,等待 ETag 或手动操作再次触发");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,20 +10,11 @@ use super::coordinator::SchedulerHandle;
|
||||
|
||||
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
|
||||
|
||||
/// 周期任务进入循环前的决策。
|
||||
pub(crate) enum PeriodicTaskStartup {
|
||||
/// 继续进入周期循环。
|
||||
Continue,
|
||||
/// 任务已自行结束,不再进入周期循环。
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// 周期任务调度配置。
|
||||
pub(crate) struct PeriodicTaskConfig {
|
||||
interval_secs: u64,
|
||||
started_message: &'static str,
|
||||
stopped_message: &'static str,
|
||||
run_immediately: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskConfig {
|
||||
@@ -36,25 +27,12 @@ impl PeriodicTaskConfig {
|
||||
interval_secs,
|
||||
started_message,
|
||||
stopped_message,
|
||||
run_immediately: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn wait_first_interval(mut self) -> Self {
|
||||
self.run_immediately = false;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 周期任务行为。
|
||||
pub(crate) trait PeriodicTaskRunner: Send + 'static {
|
||||
fn before_loop<'a>(
|
||||
&'a mut self,
|
||||
_shutdown_rx: &'a mut mpsc::Receiver<()>,
|
||||
) -> TaskFuture<'a, PeriodicTaskStartup> {
|
||||
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
|
||||
}
|
||||
|
||||
fn tick(&mut self) -> TaskFuture<'_, ()>;
|
||||
|
||||
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
|
||||
@@ -75,17 +53,7 @@ where
|
||||
"{}", config.started_message
|
||||
);
|
||||
|
||||
if matches!(
|
||||
task.before_loop(&mut shutdown_rx).await,
|
||||
PeriodicTaskStartup::Stop
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut ticker = interval(Duration::from_secs(config.interval_secs));
|
||||
if !config.run_immediately {
|
||||
ticker.tick().await;
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
|
||||
@@ -10,7 +10,8 @@ use crate::{
|
||||
cookies::PgCookieStore,
|
||||
pool::AccountPoolService,
|
||||
quota::{QuotaRefreshService, QuotaRefreshSummary},
|
||||
store::PgAccountStore,
|
||||
store::{AccountStore, PgAccountStore},
|
||||
usage::AccountUsageStore,
|
||||
},
|
||||
infra::identity::AccountPseudonymizer,
|
||||
telemetry::account_usage::store::PgAccountUsageStore,
|
||||
@@ -42,8 +43,8 @@ impl QuotaRefreshTask {
|
||||
) -> Self {
|
||||
Self {
|
||||
service: QuotaRefreshService::with_min_refresh_interval_secs(
|
||||
store,
|
||||
usage_store,
|
||||
Arc::new(store) as Arc<dyn AccountStore>,
|
||||
Arc::new(usage_store) as Arc<dyn AccountUsageStore>,
|
||||
codex,
|
||||
account_pool,
|
||||
account_pseudonymizer,
|
||||
|
||||
@@ -3,12 +3,9 @@
|
||||
use chrono::Utc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::{
|
||||
telemetry::{
|
||||
buckets::store::PgRequestBucketStore, ops::store::PgOpsErrorLogStore,
|
||||
usage::store::PgUsageRecordStore,
|
||||
},
|
||||
upstream::openai::fingerprint::PgFingerprintStore,
|
||||
use crate::telemetry::{
|
||||
buckets::store::PgRequestBucketStore, ops::store::PgOpsErrorLogStore,
|
||||
usage::store::PgUsageRecordStore,
|
||||
};
|
||||
|
||||
use super::{
|
||||
@@ -17,13 +14,11 @@ use super::{
|
||||
};
|
||||
|
||||
const RETENTION_TRIM_INTERVAL_SECS: u64 = 60 * 60;
|
||||
const FINGERPRINT_HISTORY_LIMIT: std::num::NonZeroU32 = std::num::NonZeroU32::new(100).unwrap();
|
||||
|
||||
pub struct RetentionTrimTask {
|
||||
usage_records: PgUsageRecordStore,
|
||||
ops_errors: PgOpsErrorLogStore,
|
||||
buckets: PgRequestBucketStore,
|
||||
fingerprints: PgFingerprintStore,
|
||||
}
|
||||
|
||||
impl RetentionTrimTask {
|
||||
@@ -31,13 +26,11 @@ impl RetentionTrimTask {
|
||||
usage_records: PgUsageRecordStore,
|
||||
ops_errors: PgOpsErrorLogStore,
|
||||
buckets: PgRequestBucketStore,
|
||||
fingerprints: PgFingerprintStore,
|
||||
) -> Self {
|
||||
Self {
|
||||
usage_records,
|
||||
ops_errors,
|
||||
buckets,
|
||||
fingerprints,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -70,17 +63,6 @@ impl RetentionTrimTask {
|
||||
Ok(_) => {}
|
||||
Err(error) => warn!(error = %error, "清理请求时间桶失败"),
|
||||
}
|
||||
match self
|
||||
.fingerprints
|
||||
.trim_update_history_to_limit(FINGERPRINT_HISTORY_LIMIT)
|
||||
.await
|
||||
{
|
||||
Ok(deleted) if deleted > 0 => {
|
||||
info!(deleted, "已清理超限指纹更新历史")
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(error) => warn!(error = %error, "清理超限指纹更新历史失败"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,11 +2,11 @@
|
||||
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::fleet::refresh::TokenRefresher;
|
||||
use crate::fleet::{
|
||||
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy, TokenRefreshService},
|
||||
store::PgAccountStore,
|
||||
};
|
||||
use crate::upstream::openai::token_client::TokenRefresher;
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
|
||||
@@ -3,12 +3,56 @@
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
dispatch::controllers::history::{CrossAccountReplay, sanitize_cross_account_input},
|
||||
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";
|
||||
const INSTALLATION_IDENTITY_KEYS: [&str; 3] = [
|
||||
"installation_id",
|
||||
"installationId",
|
||||
"x-codex-installation-id",
|
||||
];
|
||||
const CROSS_ACCOUNT_IDENTITY_KEYS: [&str; 26] = [
|
||||
"authorization",
|
||||
"Authorization",
|
||||
"cookie",
|
||||
"Cookie",
|
||||
"chatgpt-account-id",
|
||||
"chatgpt_account_id",
|
||||
"chatgptAccountId",
|
||||
"account_id",
|
||||
"accountId",
|
||||
"user_id",
|
||||
"userId",
|
||||
"chatgpt_user_id",
|
||||
"chatgptUserId",
|
||||
"access_token",
|
||||
"accessToken",
|
||||
"session_token",
|
||||
"sessionToken",
|
||||
"refresh_token",
|
||||
"refreshToken",
|
||||
"id_token",
|
||||
"idToken",
|
||||
"token",
|
||||
"cookies",
|
||||
"cookie_header",
|
||||
"cookieHeader",
|
||||
"cf_clearance",
|
||||
];
|
||||
const CROSS_ACCOUNT_OPAQUE_STATE_KEYS: [&str; 8] = [
|
||||
"turnState",
|
||||
"turn_state",
|
||||
"x-codex-turn-state",
|
||||
"previous_response_id",
|
||||
"previousResponseId",
|
||||
"response_id",
|
||||
"responseId",
|
||||
"conversation",
|
||||
];
|
||||
|
||||
/// 持久化密钥驱动的账号身份隔离服务。
|
||||
#[derive(Clone)]
|
||||
@@ -32,24 +76,21 @@ pub(in crate::dispatch) struct AccountScopedIdentity {
|
||||
|
||||
/// 单个请求内账号绑定 opaque 状态的 owner。
|
||||
pub(in crate::dispatch) struct AccountIdentityScope {
|
||||
turn_state_account_id: Option<String>,
|
||||
account_bound_state_owner_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,
|
||||
account_bound_state_owner_id: preferred_account_id,
|
||||
}
|
||||
}
|
||||
|
||||
fn preserves_turn_state(&mut self, account_id: &str, has_turn_state: bool) -> bool {
|
||||
if !has_turn_state {
|
||||
return false;
|
||||
}
|
||||
fn bind_and_is_cross_account(&mut self, account_id: &str) -> bool {
|
||||
let owner = self
|
||||
.turn_state_account_id
|
||||
.account_bound_state_owner_id
|
||||
.get_or_insert_with(|| account_id.to_string());
|
||||
owner == account_id
|
||||
owner != account_id
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,15 +141,19 @@ impl AccountIdentityService {
|
||||
scope: &mut AccountIdentityScope,
|
||||
mut request: CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
cross_account_replay: Option<&CrossAccountReplay>,
|
||||
) -> AccountScopedRequest {
|
||||
if let Some(replay) = cross_account_replay {
|
||||
replay.validate_target(account_id);
|
||||
}
|
||||
let cross_account =
|
||||
cross_account_replay.is_some() || scope.bind_and_is_cross_account(account_id);
|
||||
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,
|
||||
request.turn_state.is_some() || client_metadata_turn_state.is_some(),
|
||||
);
|
||||
let preserve_turn_state = !cross_account
|
||||
&& (request.turn_state.is_some() || client_metadata_turn_state.is_some());
|
||||
let turn_state = if preserve_turn_state {
|
||||
request.turn_state.clone()
|
||||
} else {
|
||||
@@ -122,9 +167,14 @@ impl AccountIdentityService {
|
||||
let turn_metadata = request
|
||||
.turn_metadata
|
||||
.as_deref()
|
||||
.map(|raw| Self::scope_turn_metadata(raw, &identity));
|
||||
.and_then(|raw| Self::scope_turn_metadata(raw, &identity, cross_account));
|
||||
let client_metadata_turn_metadata = client_metadata_string(&request, TURN_METADATA_KEY)
|
||||
.map(|raw| Self::scope_turn_metadata(&raw, &identity));
|
||||
.and_then(|raw| Self::scope_turn_metadata(&raw, &identity, cross_account));
|
||||
|
||||
if cross_account {
|
||||
sanitize_cross_account_input(&mut request);
|
||||
remove_cross_account_request_state(&mut request);
|
||||
}
|
||||
|
||||
apply_identity_to_request(
|
||||
&mut request,
|
||||
@@ -133,6 +183,7 @@ impl AccountIdentityService {
|
||||
turn_metadata.as_deref(),
|
||||
client_metadata_turn_state.as_deref(),
|
||||
client_metadata_turn_metadata.as_deref(),
|
||||
cross_account,
|
||||
);
|
||||
request.turn_state = turn_state;
|
||||
request.turn_metadata = turn_metadata;
|
||||
@@ -140,37 +191,33 @@ impl AccountIdentityService {
|
||||
AccountScopedRequest { request, identity }
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn scope_auxiliary(&self, account_id: &str) -> AccountScopedIdentity {
|
||||
AccountScopedIdentity {
|
||||
installation_id: self.pseudonymizer.installation_id(account_id),
|
||||
}
|
||||
}
|
||||
|
||||
fn scope_turn_metadata(raw: &str, identity: &AccountScopedIdentity) -> String {
|
||||
fn scope_turn_metadata(
|
||||
raw: &str,
|
||||
identity: &AccountScopedIdentity,
|
||||
cross_account: bool,
|
||||
) -> Option<String> {
|
||||
let Ok(Value::Object(mut metadata)) = serde_json::from_str::<Value>(raw) else {
|
||||
return raw.to_string();
|
||||
return (!cross_account).then(|| 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();
|
||||
if cross_account {
|
||||
remove_cross_account_metadata_state(&mut metadata);
|
||||
} else if !INSTALLATION_IDENTITY_KEYS
|
||||
.into_iter()
|
||||
.any(|key| metadata.contains_key(key))
|
||||
{
|
||||
return Some(raw.to_string());
|
||||
}
|
||||
|
||||
for key in [
|
||||
"installation_id",
|
||||
"installationId",
|
||||
"x-codex-installation-id",
|
||||
] {
|
||||
for key in INSTALLATION_IDENTITY_KEYS {
|
||||
replace_existing_metadata_field(&mut metadata, key, Some(&identity.installation_id));
|
||||
}
|
||||
|
||||
serde_json::to_string(&metadata).unwrap_or_else(|_| raw.to_string())
|
||||
match serde_json::to_string(&metadata) {
|
||||
Ok(metadata) => Some(metadata),
|
||||
Err(_) if cross_account => None,
|
||||
Err(_) => Some(raw.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn pseudonym(&self, domain: &str, account_id: Option<&str>, value: &str) -> String {
|
||||
@@ -194,6 +241,7 @@ fn apply_identity_to_request(
|
||||
turn_metadata: Option<&str>,
|
||||
client_metadata_turn_state: Option<&str>,
|
||||
client_metadata_turn_metadata: Option<&str>,
|
||||
cross_account: bool,
|
||||
) {
|
||||
for (key, value) in [
|
||||
(
|
||||
@@ -215,6 +263,7 @@ fn apply_identity_to_request(
|
||||
identity,
|
||||
client_metadata_turn_state,
|
||||
client_metadata_turn_metadata,
|
||||
cross_account,
|
||||
);
|
||||
request.set_client_metadata(client_metadata);
|
||||
}
|
||||
@@ -224,12 +273,17 @@ fn scoped_client_metadata(
|
||||
identity: &AccountScopedIdentity,
|
||||
turn_state: Option<&str>,
|
||||
turn_metadata: Option<&str>,
|
||||
cross_account: bool,
|
||||
) -> Option<Value> {
|
||||
let client_metadata = client_metadata?;
|
||||
let Value::Object(mut metadata) = client_metadata.clone() else {
|
||||
return Some(client_metadata.clone());
|
||||
return (!cross_account).then(|| client_metadata.clone());
|
||||
};
|
||||
|
||||
if cross_account {
|
||||
remove_cross_account_metadata_state(&mut metadata);
|
||||
}
|
||||
|
||||
for (key, value) in [
|
||||
(
|
||||
"x-codex-installation-id",
|
||||
@@ -254,6 +308,28 @@ fn scoped_client_metadata(
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
fn remove_cross_account_request_state(request: &mut CodexResponsesRequest) {
|
||||
for key in CROSS_ACCOUNT_IDENTITY_KEYS
|
||||
.into_iter()
|
||||
.chain(CROSS_ACCOUNT_OPAQUE_STATE_KEYS)
|
||||
{
|
||||
request.replace_existing_identity_field(key, None);
|
||||
}
|
||||
for key in ["turn_metadata", TURN_METADATA_KEY] {
|
||||
request.replace_existing_identity_field(key, None);
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_cross_account_metadata_state(metadata: &mut Map<String, Value>) {
|
||||
for key in CROSS_ACCOUNT_IDENTITY_KEYS
|
||||
.into_iter()
|
||||
.chain(CROSS_ACCOUNT_OPAQUE_STATE_KEYS)
|
||||
.chain(["turnMetadata", "turn_metadata", TURN_METADATA_KEY])
|
||||
{
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn replace_existing_metadata_field(
|
||||
metadata: &mut Map<String, Value>,
|
||||
key: &str,
|
||||
|
||||
@@ -12,7 +12,7 @@ use crate::{
|
||||
fleet::{
|
||||
account_failure::{
|
||||
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
|
||||
classify_response_failure, classify_upstream_failure,
|
||||
classify_account_failure,
|
||||
},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
@@ -119,23 +119,31 @@ impl AccountFailureController {
|
||||
if code.contains("invalid_request")
|
||||
|| code.contains("not_found")
|
||||
|| code.contains("context_window")
|
||||
|| code == "context_length_exceeded"
|
||||
|| code.contains("invalid_prompt")
|
||||
|| code.contains("bad_request")
|
||||
{
|
||||
return Some(ClientFailure::new(failure.clone(), 400, false));
|
||||
}
|
||||
if code.contains("server_overloaded") {
|
||||
if matches!(code.as_str(), "server_is_overloaded" | "slow_down")
|
||||
|| code.contains("server_overloaded")
|
||||
{
|
||||
return Some(ClientFailure::new(failure.clone(), 503, true));
|
||||
}
|
||||
classify_response_failure(failure).map(|classified| {
|
||||
let status = match classified.kind {
|
||||
AccountFailureKind::ModelUnsupported => 400,
|
||||
AccountFailureKind::Expired => 401,
|
||||
AccountFailureKind::Disabled | AccountFailureKind::Banned => 403,
|
||||
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => 429,
|
||||
};
|
||||
ClientFailure::new(failure.clone(), status, false)
|
||||
})
|
||||
if code == "usage_not_included" {
|
||||
return Some(ClientFailure::new(failure.clone(), 403, false));
|
||||
}
|
||||
classify_account_failure(&super::failure_observation::from_response(failure)).map(
|
||||
|classified| {
|
||||
let status = match classified.kind {
|
||||
AccountFailureKind::ModelUnsupported => 400,
|
||||
AccountFailureKind::Expired => 401,
|
||||
AccountFailureKind::Disabled | AccountFailureKind::Banned => 403,
|
||||
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => 429,
|
||||
};
|
||||
ClientFailure::new(failure.clone(), status, false)
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -171,7 +179,7 @@ fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<C
|
||||
account_id,
|
||||
&facts.body,
|
||||
facts.status_code,
|
||||
classify_upstream_failure(facts)?,
|
||||
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -180,7 +188,7 @@ fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<Class
|
||||
account_id,
|
||||
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
|
||||
failure.explicit_status_code,
|
||||
classify_response_failure(failure)?,
|
||||
classify_account_failure(&super::failure_observation::from_response(failure))?,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
//! 上游 typed failure 到 fleet 账号事实的边界映射。
|
||||
|
||||
use crate::{
|
||||
fleet::account_gateway::AccountFailureObservation,
|
||||
upstream::openai::{failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure},
|
||||
};
|
||||
|
||||
pub(super) fn from_upstream(facts: &UpstreamFailureFacts) -> AccountFailureObservation {
|
||||
AccountFailureObservation {
|
||||
status_code: facts.status_code,
|
||||
code: facts.code.clone(),
|
||||
error_type: facts.error_type.clone(),
|
||||
identity_authorization_error: facts.identity_authorization_error.clone(),
|
||||
identity_error_code: facts.identity_error_code.clone(),
|
||||
message: facts.message.clone(),
|
||||
body: facts.body.clone(),
|
||||
retry_after_seconds: facts.retry_after_seconds,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn from_response(failure: &ResponsesSseFailure) -> AccountFailureObservation {
|
||||
AccountFailureObservation {
|
||||
status_code: failure.explicit_status_code,
|
||||
code: failure.upstream_code.clone(),
|
||||
error_type: failure.upstream_type.clone(),
|
||||
message: failure.message.clone(),
|
||||
body: failure.message.clone(),
|
||||
retry_after_seconds: failure.retry_after_seconds,
|
||||
..AccountFailureObservation::default()
|
||||
}
|
||||
}
|
||||
@@ -16,7 +16,10 @@ use crate::{
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::{UpstreamFailureFacts, UpstreamFailureKind},
|
||||
protocol::responses::{CodexResponsesRequest, PreviousResponseScope, StreamCommitPolicy},
|
||||
protocol::responses::{
|
||||
CodexResponsesRequest, LocalReplayItem, LocalReplayTranscript, PreviousResponseScope,
|
||||
StreamCommitPolicy,
|
||||
},
|
||||
transport::CodexBackendTransport,
|
||||
transport::websocket::PreviousResponseUnavailableReason,
|
||||
},
|
||||
@@ -38,7 +41,8 @@ enum HistorySource {
|
||||
pub(super) struct HistoryState {
|
||||
source: HistorySource,
|
||||
managed_entry: Option<SessionAffinityEntry>,
|
||||
local_replay_input: Option<Arc<Vec<Value>>>,
|
||||
local_replay_transcript: Option<Arc<LocalReplayTranscript>>,
|
||||
local_replay_account_id: Option<String>,
|
||||
external_attempted_account_id: Option<String>,
|
||||
replay_owner_account: bool,
|
||||
}
|
||||
@@ -46,7 +50,8 @@ pub(super) struct HistoryState {
|
||||
#[derive(Default)]
|
||||
pub(in crate::dispatch) struct ConnectionReplaySnapshot {
|
||||
last_response_id: Option<String>,
|
||||
input: Option<Arc<Vec<Value>>>,
|
||||
account_id: Option<String>,
|
||||
transcript: Option<Arc<LocalReplayTranscript>>,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct ConnectionReplayPlan {
|
||||
@@ -54,16 +59,51 @@ pub(in crate::dispatch) struct ConnectionReplayPlan {
|
||||
}
|
||||
|
||||
enum ConnectionReplayUpdate {
|
||||
Replace(Vec<Value>),
|
||||
Append {
|
||||
history: Arc<Vec<Value>>,
|
||||
input: Vec<Value>,
|
||||
},
|
||||
Replace,
|
||||
Append { history: Arc<LocalReplayTranscript> },
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
/// 已确认来源和目标账号不同的完整 transcript 重放边界。
|
||||
pub(in crate::dispatch) struct CrossAccountReplay {
|
||||
source_account_id: String,
|
||||
target_account_id: String,
|
||||
}
|
||||
|
||||
impl CrossAccountReplay {
|
||||
fn between(source_account_id: &str, target_account_id: &str) -> Option<Self> {
|
||||
(source_account_id != target_account_id).then(|| Self {
|
||||
source_account_id: source_account_id.to_string(),
|
||||
target_account_id: target_account_id.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn validate_target(&self, account_id: &str) {
|
||||
debug_assert_ne!(self.source_account_id, self.target_account_id);
|
||||
debug_assert_eq!(self.target_account_id, account_id);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct PreparedHistoryRequest {
|
||||
pub(super) request: CodexResponsesRequest,
|
||||
pub(super) cross_account_replay: Option<CrossAccountReplay>,
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct HistoryController;
|
||||
|
||||
pub(in crate::dispatch) fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) {
|
||||
if request.input().is_empty() {
|
||||
return;
|
||||
}
|
||||
let input = request
|
||||
.input()
|
||||
.iter()
|
||||
.cloned()
|
||||
.filter_map(sanitize_cross_account_item)
|
||||
.collect();
|
||||
request.set_input(input);
|
||||
}
|
||||
|
||||
impl HistoryController {
|
||||
pub(in crate::dispatch) fn new_connection_replay_snapshot() -> ConnectionReplaySnapshot {
|
||||
ConnectionReplaySnapshot::default()
|
||||
@@ -73,28 +113,29 @@ impl HistoryController {
|
||||
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)
|
||||
clear_local_replay(request);
|
||||
ConnectionReplayUpdate::Replace
|
||||
}
|
||||
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));
|
||||
if let (Some(account_id), Some(history)) =
|
||||
(snapshot.account_id.as_ref(), snapshot.transcript.as_ref())
|
||||
{
|
||||
request.local_replay_transcript = Some(Arc::clone(history));
|
||||
request.local_replay_account_id = Some(account_id.clone());
|
||||
ConnectionReplayUpdate::Append {
|
||||
history: Arc::clone(history),
|
||||
input: turn_input,
|
||||
}
|
||||
} else {
|
||||
request.local_replay_input = None;
|
||||
clear_local_replay(request);
|
||||
ConnectionReplayUpdate::Unavailable
|
||||
}
|
||||
}
|
||||
Some(_) => {
|
||||
request.local_replay_input = None;
|
||||
clear_local_replay(request);
|
||||
ConnectionReplayUpdate::Unavailable
|
||||
}
|
||||
};
|
||||
@@ -107,18 +148,38 @@ impl HistoryController {
|
||||
plan: ConnectionReplayPlan,
|
||||
response_id: String,
|
||||
output: Vec<Value>,
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
) {
|
||||
snapshot.last_response_id = Some(response_id);
|
||||
snapshot.account_id = Some(account_id.clone());
|
||||
match plan.update {
|
||||
ConnectionReplayUpdate::Replace(input) => {
|
||||
snapshot.input = Some(Arc::new(sanitized_transcript(input, output)));
|
||||
ConnectionReplayUpdate::Replace => {
|
||||
snapshot.transcript = Some(Arc::new(transcript_for_turn(
|
||||
request_input,
|
||||
output,
|
||||
&account_id,
|
||||
)));
|
||||
}
|
||||
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::Append { mut history } => {
|
||||
if continued_from_previous_response {
|
||||
append_turn(
|
||||
Arc::make_mut(&mut history),
|
||||
request_input,
|
||||
output,
|
||||
&account_id,
|
||||
);
|
||||
snapshot.transcript = Some(history);
|
||||
} else {
|
||||
snapshot.transcript = Some(Arc::new(transcript_for_turn(
|
||||
request_input,
|
||||
output,
|
||||
&account_id,
|
||||
)));
|
||||
}
|
||||
}
|
||||
ConnectionReplayUpdate::Unavailable => snapshot.input = None,
|
||||
ConnectionReplayUpdate::Unavailable => snapshot.transcript = None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -145,51 +206,67 @@ impl HistoryController {
|
||||
None
|
||||
}
|
||||
};
|
||||
let local_replay_input = request.local_replay_input.clone();
|
||||
let source = match (&managed_entry, &local_replay_input) {
|
||||
let (local_replay_transcript, local_replay_account_id) = match (
|
||||
request.local_replay_transcript.clone(),
|
||||
request.local_replay_account_id.clone(),
|
||||
) {
|
||||
(Some(transcript), Some(account_id)) => (Some(transcript), Some(account_id)),
|
||||
(None, None) => (None, None),
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
previous_response_id,
|
||||
"Ignoring incomplete connection-local replay provenance"
|
||||
);
|
||||
(None, None)
|
||||
}
|
||||
};
|
||||
let has_complete_local_replay = local_replay_transcript.is_some();
|
||||
let source = match (&managed_entry, has_complete_local_replay) {
|
||||
(Some(_), _) => HistorySource::Managed,
|
||||
(None, Some(_)) => HistorySource::ConnectionLocal,
|
||||
(None, None) => HistorySource::ExternalUnknown,
|
||||
(None, true) => HistorySource::ConnectionLocal,
|
||||
(None, false) => HistorySource::ExternalUnknown,
|
||||
};
|
||||
HistoryState {
|
||||
source,
|
||||
managed_entry,
|
||||
local_replay_input,
|
||||
local_replay_transcript,
|
||||
local_replay_account_id,
|
||||
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())
|
||||
state.local_replay_account_id.as_deref().or_else(|| {
|
||||
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> {
|
||||
) -> Result<PreparedHistoryRequest, String> {
|
||||
let prepared = match state.source {
|
||||
HistorySource::None => Some(request.clone()),
|
||||
HistorySource::None => Some(PreparedHistoryRequest::plain(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(PreparedHistoryRequest::plain(external_history_request(
|
||||
request,
|
||||
)))
|
||||
}
|
||||
Some(first_account_id) if first_account_id == account_id => Some(
|
||||
PreparedHistoryRequest::plain(external_history_request(request)),
|
||||
),
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
HistorySource::ConnectionLocal => state
|
||||
.local_replay_input
|
||||
.as_deref()
|
||||
.map(|replay_input| full_replay_request(request, replay_input)),
|
||||
HistorySource::ConnectionLocal => full_replay_from_state(state, request, account_id),
|
||||
HistorySource::Managed => prepare_managed_attempt(state, request, account_id),
|
||||
};
|
||||
prepared.ok_or_else(|| CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE.to_string())
|
||||
@@ -213,7 +290,7 @@ impl HistoryController {
|
||||
.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::Managed => state.local_replay_transcript.is_some(),
|
||||
HistorySource::ExternalUnknown => false,
|
||||
},
|
||||
}
|
||||
@@ -294,7 +371,7 @@ impl HistoryController {
|
||||
return false;
|
||||
};
|
||||
entry.account_id == account_id
|
||||
&& state.local_replay_input.is_some()
|
||||
&& state.local_replay_transcript.is_some()
|
||||
&& !state.replay_owner_account
|
||||
}
|
||||
}
|
||||
@@ -303,10 +380,10 @@ fn prepare_managed_attempt(
|
||||
state: &HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
) -> Option<CodexResponsesRequest> {
|
||||
) -> Option<PreparedHistoryRequest> {
|
||||
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);
|
||||
return replay_for_conversation(state, request, account_id, &entry.conversation_id);
|
||||
}
|
||||
|
||||
match entry.continuation_scope {
|
||||
@@ -320,10 +397,10 @@ fn prepare_managed_attempt(
|
||||
if prepared.turn_state.as_deref().is_none_or(str::is_empty) {
|
||||
prepared.turn_state.clone_from(&entry.turn_state);
|
||||
}
|
||||
Some(prepared)
|
||||
Some(PreparedHistoryRequest::plain(prepared))
|
||||
}
|
||||
PreviousResponseScope::ReplayRequired => {
|
||||
replay_for_conversation(state, request, &entry.conversation_id)
|
||||
replay_for_conversation(state, request, account_id, &entry.conversation_id)
|
||||
}
|
||||
PreviousResponseScope::Unavailable => None,
|
||||
}
|
||||
@@ -332,21 +409,37 @@ fn prepare_managed_attempt(
|
||||
fn replay_for_conversation(
|
||||
state: &HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
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());
|
||||
) -> Option<PreparedHistoryRequest> {
|
||||
full_replay_from_state(state, request, account_id).map(|mut prepared| {
|
||||
prepared.request.local_conversation_id = Some(conversation_id.to_string());
|
||||
prepared
|
||||
})
|
||||
}
|
||||
|
||||
fn full_replay_from_state(
|
||||
state: &HistoryState,
|
||||
request: &CodexResponsesRequest,
|
||||
account_id: &str,
|
||||
) -> Option<PreparedHistoryRequest> {
|
||||
let transcript = state.local_replay_transcript.as_deref()?;
|
||||
let source_account_id = state.local_replay_account_id.as_deref()?;
|
||||
Some(full_replay_request(
|
||||
request,
|
||||
transcript,
|
||||
source_account_id,
|
||||
account_id,
|
||||
))
|
||||
}
|
||||
|
||||
impl HistoryState {
|
||||
fn empty() -> Self {
|
||||
Self {
|
||||
source: HistorySource::None,
|
||||
managed_entry: None,
|
||||
local_replay_input: None,
|
||||
local_replay_transcript: None,
|
||||
local_replay_account_id: None,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
}
|
||||
@@ -363,43 +456,138 @@ fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesR
|
||||
|
||||
fn full_replay_request(
|
||||
original: &CodexResponsesRequest,
|
||||
replay_input: &[Value],
|
||||
) -> CodexResponsesRequest {
|
||||
transcript: &LocalReplayTranscript,
|
||||
source_account_id: &str,
|
||||
target_account_id: &str,
|
||||
) -> PreparedHistoryRequest {
|
||||
let cross_account_replay = CrossAccountReplay::between(source_account_id, target_account_id);
|
||||
let mut request = original.clone();
|
||||
let mut input = Vec::with_capacity(replay_input.len() + original.input().len());
|
||||
input.extend_from_slice(replay_input);
|
||||
let mut input = replay_input_for_account(transcript, target_account_id);
|
||||
input.reserve(original.input().len());
|
||||
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
|
||||
PreparedHistoryRequest {
|
||||
request,
|
||||
cross_account_replay,
|
||||
}
|
||||
}
|
||||
|
||||
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(_) => {}
|
||||
impl PreparedHistoryRequest {
|
||||
fn plain(request: CodexResponsesRequest) -> Self {
|
||||
Self {
|
||||
request,
|
||||
cross_account_replay: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_local_replay(request: &mut CodexResponsesRequest) {
|
||||
request.local_replay_transcript = None;
|
||||
request.local_replay_account_id = None;
|
||||
}
|
||||
|
||||
fn transcript_for_turn(
|
||||
input: Vec<Value>,
|
||||
output: Vec<Value>,
|
||||
account_id: &str,
|
||||
) -> LocalReplayTranscript {
|
||||
let mut transcript = LocalReplayTranscript::default();
|
||||
append_client_input(&mut transcript, input);
|
||||
append_account_output(&mut transcript, output, account_id);
|
||||
transcript
|
||||
}
|
||||
|
||||
fn append_turn(
|
||||
transcript: &mut LocalReplayTranscript,
|
||||
input: Vec<Value>,
|
||||
output: Vec<Value>,
|
||||
account_id: &str,
|
||||
) {
|
||||
project_transcript_to_account(transcript, account_id);
|
||||
append_client_input(transcript, input);
|
||||
append_account_output(transcript, output, account_id);
|
||||
}
|
||||
|
||||
fn append_client_input(transcript: &mut LocalReplayTranscript, input: Vec<Value>) {
|
||||
transcript
|
||||
.items
|
||||
.extend(input.into_iter().map(LocalReplayItem::ClientInput));
|
||||
}
|
||||
|
||||
fn append_account_output(
|
||||
transcript: &mut LocalReplayTranscript,
|
||||
output: Vec<Value>,
|
||||
account_id: &str,
|
||||
) {
|
||||
transcript.items.extend(
|
||||
output
|
||||
.into_iter()
|
||||
.map(|item| LocalReplayItem::AccountOutput {
|
||||
account_id: account_id.to_string(),
|
||||
item,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
fn project_transcript_to_account(transcript: &mut LocalReplayTranscript, account_id: &str) {
|
||||
transcript.items = transcript
|
||||
.items
|
||||
.drain(..)
|
||||
.filter_map(|item| match item {
|
||||
LocalReplayItem::AccountOutput {
|
||||
account_id: owner,
|
||||
item,
|
||||
} if owner != account_id => {
|
||||
sanitize_cross_account_item(item).map(LocalReplayItem::SanitizedOutput)
|
||||
}
|
||||
item => Some(item),
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
fn replay_input_for_account(transcript: &LocalReplayTranscript, account_id: &str) -> Vec<Value> {
|
||||
transcript
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
LocalReplayItem::ClientInput(value) | LocalReplayItem::SanitizedOutput(value) => {
|
||||
Some(value.clone())
|
||||
}
|
||||
LocalReplayItem::AccountOutput {
|
||||
account_id: owner,
|
||||
item,
|
||||
} if owner == account_id => Some(without_output_id(item.clone())),
|
||||
LocalReplayItem::AccountOutput { item, .. } => {
|
||||
sanitize_cross_account_item(item.clone())
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn without_output_id(mut item: Value) -> Value {
|
||||
if let Value::Object(object) = &mut item {
|
||||
object.remove("id");
|
||||
}
|
||||
item
|
||||
}
|
||||
|
||||
fn sanitize_cross_account_item(mut item: Value) -> Option<Value> {
|
||||
if let Value::Object(object) = &mut item {
|
||||
if matches!(
|
||||
object.get("type").and_then(Value::as_str),
|
||||
Some("compaction" | "compaction_summary" | "context_compaction")
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
object.remove("id");
|
||||
object.remove("encrypted_content");
|
||||
}
|
||||
Some(item)
|
||||
}
|
||||
|
||||
fn is_continuation_busy(facts: &UpstreamFailureFacts) -> bool {
|
||||
matches!(
|
||||
facts.kind,
|
||||
|
||||
@@ -19,7 +19,10 @@ use crate::{
|
||||
AccountIdentityScope, AccountIdentityService, AccountScopedRequest,
|
||||
SessionAffinityService,
|
||||
},
|
||||
errors::{ClientFailure, ResponseDispatchError, upstream_error_set_cookie_headers},
|
||||
errors::{
|
||||
ClientFailure, ResponseDispatchError, upstream_error_rate_limit_headers,
|
||||
upstream_error_set_cookie_headers,
|
||||
},
|
||||
lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
|
||||
},
|
||||
fleet::{account::Account, pool::AccountLease},
|
||||
@@ -51,6 +54,7 @@ use crate::{
|
||||
},
|
||||
},
|
||||
fleet::pool::AccountPoolService,
|
||||
fleet::scheduler::AttemptFeedback,
|
||||
upstream::openai::{
|
||||
failure::UpstreamFailureFacts,
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
@@ -61,6 +65,7 @@ mod account_failure;
|
||||
mod affinity;
|
||||
pub(crate) mod cloudflare;
|
||||
pub(crate) mod cyber_policy;
|
||||
mod failure_observation;
|
||||
pub(in crate::dispatch) mod history;
|
||||
mod quota;
|
||||
pub(in crate::dispatch) mod telemetry;
|
||||
@@ -270,22 +275,11 @@ pub(in crate::dispatch) struct ControllerEnter {
|
||||
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,
|
||||
@@ -299,38 +293,16 @@ pub(in crate::dispatch) enum AttemptRoutePreparation {
|
||||
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;
|
||||
}
|
||||
};
|
||||
) -> PreparedAttemptAccount {
|
||||
let cookie_header = self
|
||||
.cloudflare
|
||||
.cookie_header_for_request(&acquired.account.id, "/codex/responses")
|
||||
.await;
|
||||
AttemptAccountPreparation::Ready(Box::new(PreparedAttemptAccount {
|
||||
PreparedAttemptAccount {
|
||||
lease: acquired,
|
||||
cookie_header,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn prepare_attempt_route(
|
||||
@@ -339,18 +311,19 @@ impl ControllerSet {
|
||||
request: &CodexResponsesRequest,
|
||||
context: AttemptRoutePreparationContext<'_>,
|
||||
) -> AttemptRoutePreparation {
|
||||
let request = match HistoryController::prepare_attempt(
|
||||
let prepared = match HistoryController::prepare_attempt(
|
||||
&mut scope.history,
|
||||
request,
|
||||
context.account_id,
|
||||
) {
|
||||
Ok(request) => request,
|
||||
Ok(prepared) => prepared,
|
||||
Err(message) => return AttemptRoutePreparation::Unavailable { message },
|
||||
};
|
||||
AttemptRoutePreparation::Ready(Box::new(context.account_identity.scope_request(
|
||||
&mut scope.identity,
|
||||
request,
|
||||
prepared.request,
|
||||
context.account_id,
|
||||
prepared.cross_account_replay.as_ref(),
|
||||
)))
|
||||
}
|
||||
|
||||
@@ -384,8 +357,12 @@ impl ControllerSet {
|
||||
.await;
|
||||
let _ = best_effort_controller_io(
|
||||
"quota.complete_headers",
|
||||
self.account_pool
|
||||
.sync_passive_rate_limit_headers(account, rate_limit_headers),
|
||||
usage::UsageController::sync_passive_quota(
|
||||
&self.account_pool,
|
||||
&account.id,
|
||||
account.plan_type.as_deref(),
|
||||
rate_limit_headers,
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -394,19 +371,28 @@ impl ControllerSet {
|
||||
&self,
|
||||
context: AttemptUpstreamErrorContext<'_>,
|
||||
) {
|
||||
let _ = best_effort_controller_io(
|
||||
let cookie = 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(
|
||||
);
|
||||
let rate_limit_headers = upstream_error_rate_limit_headers(context.error);
|
||||
let quota = best_effort_controller_io(
|
||||
"quota.attempt_error_headers",
|
||||
usage::UsageController::sync_passive_quota(
|
||||
&self.account_pool,
|
||||
&context.account.id,
|
||||
context.account.plan_type.as_deref(),
|
||||
&rate_limit_headers,
|
||||
),
|
||||
);
|
||||
let telemetry = best_effort_controller_io(
|
||||
"telemetry.attempt_error",
|
||||
telemetry::TelemetryController::observe_upstream_error(&self.recorder, &context),
|
||||
)
|
||||
.await;
|
||||
);
|
||||
let _ = tokio::join!(cookie, quota, telemetry);
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn observe_dispatch_error(
|
||||
@@ -453,6 +439,18 @@ impl ControllerSet {
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let feedback = if exit.completed {
|
||||
AttemptFeedback::Completed {
|
||||
first_token_ms: nonnegative_u64(exit.response.first_token_ms),
|
||||
}
|
||||
} else {
|
||||
AttemptFeedback::Incomplete {
|
||||
first_token_ms: nonnegative_u64(exit.response.first_token_ms),
|
||||
}
|
||||
};
|
||||
self.account_pool
|
||||
.report_attempt(exit.account_id, feedback)
|
||||
.await;
|
||||
|
||||
if successful {
|
||||
let _ = best_effort_controller_io(
|
||||
@@ -657,6 +655,12 @@ impl ControllerSet {
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
self.account_pool
|
||||
.report_attempt(
|
||||
&context.account_id,
|
||||
stream_feedback(summary, context.started_at, context.attempt.started_at()),
|
||||
)
|
||||
.await;
|
||||
|
||||
let fallback = usage::StreamUsageExit {
|
||||
rate_limit_headers: context.rate_limit_headers.clone(),
|
||||
@@ -726,7 +730,6 @@ impl ControllerSet {
|
||||
}),
|
||||
)
|
||||
.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)
|
||||
@@ -750,6 +753,13 @@ impl ControllerSet {
|
||||
usage::UsageController::observe_attempt(&self.account_pool, &scope.usage, observation),
|
||||
)
|
||||
.await;
|
||||
if let (Some(account), Some(feedback)) =
|
||||
(observation.account.as_ref(), attempt_feedback(observation))
|
||||
{
|
||||
self.account_pool
|
||||
.report_attempt(&account.id, feedback)
|
||||
.await;
|
||||
}
|
||||
|
||||
if matches!(
|
||||
observation.kind,
|
||||
@@ -918,9 +928,6 @@ impl ControllerSet {
|
||||
AttemptObservationKind::NoCandidate { .. } => {
|
||||
AttemptDecision::Return(AttemptReturnKind::Observed)
|
||||
}
|
||||
AttemptObservationKind::CandidatePreparationRejected => {
|
||||
unreachable!("candidate preparation controller owns it")
|
||||
}
|
||||
AttemptObservationKind::RoutePreparationRejected { message } => {
|
||||
AttemptDecision::Return(AttemptReturnKind::RouteUnavailable {
|
||||
message: message.clone(),
|
||||
@@ -953,6 +960,71 @@ impl ControllerSet {
|
||||
}
|
||||
}
|
||||
|
||||
fn attempt_feedback(observation: &AttemptObservation) -> Option<AttemptFeedback> {
|
||||
observation.attempt.as_ref()?;
|
||||
if matches!(
|
||||
observation.kind,
|
||||
AttemptObservationKind::CompleteResponse(
|
||||
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Completed
|
||||
| crate::dispatch::lifecycle::contract::CompleteResponseFacts::Incomplete,
|
||||
)
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
Some(AttemptFeedback::FailedBeforeFirstToken)
|
||||
}
|
||||
|
||||
fn stream_feedback(
|
||||
summary: &crate::dispatch::lifecycle::stream::StreamSummary,
|
||||
request_started_at: Instant,
|
||||
attempt_started_at: Instant,
|
||||
) -> AttemptFeedback {
|
||||
let first_token_ms = attempt_first_token_ms(
|
||||
summary.first_token_ms,
|
||||
request_started_at,
|
||||
attempt_started_at,
|
||||
);
|
||||
match summary.terminal {
|
||||
StreamTerminal::Completed { .. } => AttemptFeedback::Completed { first_token_ms },
|
||||
StreamTerminal::Incomplete { .. } => AttemptFeedback::Incomplete { first_token_ms },
|
||||
StreamTerminal::Cancelled | StreamTerminal::DownstreamClosed => AttemptFeedback::Cancelled,
|
||||
StreamTerminal::Shutdown => AttemptFeedback::Shutdown,
|
||||
StreamTerminal::Failed { .. }
|
||||
| StreamTerminal::UpstreamClosed
|
||||
| StreamTerminal::UpstreamError { .. }
|
||||
| StreamTerminal::ProtocolError { .. }
|
||||
| StreamTerminal::CaptureLimitExceeded => {
|
||||
summary
|
||||
.first_token_ms
|
||||
.map_or(AttemptFeedback::FailedBeforeFirstToken, |_| {
|
||||
AttemptFeedback::FailedAfterFirstToken {
|
||||
first_token_ms: first_token_ms.unwrap_or(0),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn attempt_first_token_ms(
|
||||
request_first_token_ms: Option<i64>,
|
||||
request_started_at: Instant,
|
||||
attempt_started_at: Instant,
|
||||
) -> Option<u64> {
|
||||
let offset_ms = i64::try_from(
|
||||
attempt_started_at
|
||||
.saturating_duration_since(request_started_at)
|
||||
.as_millis(),
|
||||
)
|
||||
.ok()?;
|
||||
request_first_token_ms
|
||||
.and_then(|value| value.checked_sub(offset_ms))
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
}
|
||||
|
||||
fn nonnegative_u64(value: Option<i64>) -> Option<u64> {
|
||||
value.and_then(|value| u64::try_from(value).ok())
|
||||
}
|
||||
|
||||
fn response_failure(observation: &AttemptObservation) -> Option<&ResponsesSseFailure> {
|
||||
match &observation.kind {
|
||||
AttemptObservationKind::CompleteResponse(
|
||||
|
||||
@@ -1,139 +1,29 @@
|
||||
//! 请求前配额验证规则 owner。
|
||||
|
||||
use serde_json::Value;
|
||||
//! 配额失败规则 owner。
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::AccountIdentityService,
|
||||
controllers::ControllerFailureFact,
|
||||
errors::ClientFailure,
|
||||
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
|
||||
lifecycle::contract::{
|
||||
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
|
||||
},
|
||||
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
|
||||
},
|
||||
fleet::{
|
||||
account_failure::{
|
||||
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
|
||||
classify_response_failure, classify_upstream_failure,
|
||||
},
|
||||
pool::{AccountLease, AccountPoolService},
|
||||
quota::{
|
||||
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
|
||||
quota_snapshot_reset_at,
|
||||
classify_account_failure,
|
||||
},
|
||||
pool::AccountPoolService,
|
||||
},
|
||||
upstream::openai::{
|
||||
failure::UpstreamFailureFacts,
|
||||
protocol::responses::ResponsesSseFailure,
|
||||
transport::{CodexBackendClient, CodexRequestContext},
|
||||
failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure,
|
||||
transport::CodexBackendClient,
|
||||
},
|
||||
};
|
||||
|
||||
const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached";
|
||||
|
||||
pub(super) enum QuotaEnterOutcome {
|
||||
Ready(Box<AccountLease>),
|
||||
LimitReached,
|
||||
}
|
||||
|
||||
pub(super) struct QuotaEnterContext<'a> {
|
||||
pub account_pool: &'a AccountPoolService,
|
||||
pub codex: &'a CodexBackendClient,
|
||||
pub account_identity: &'a AccountIdentityService,
|
||||
pub request_id: &'a str,
|
||||
pub cookie_header: Option<&'a str>,
|
||||
}
|
||||
|
||||
pub(super) struct QuotaController;
|
||||
|
||||
impl QuotaController {
|
||||
pub(super) async fn enter(
|
||||
context: QuotaEnterContext<'_>,
|
||||
acquired: AccountLease,
|
||||
) -> QuotaEnterOutcome {
|
||||
if !acquired.account.quota_verify_required {
|
||||
return QuotaEnterOutcome::Ready(Box::new(acquired));
|
||||
}
|
||||
|
||||
let account_id = acquired.account.id.clone();
|
||||
let identity = context.account_identity.scope_auxiliary(&account_id);
|
||||
let usage = context
|
||||
.codex
|
||||
.fetch_usage(CodexRequestContext {
|
||||
access_token: &acquired.account.access_token,
|
||||
account_id: acquired.account.account_id.as_deref(),
|
||||
request_id: context.request_id,
|
||||
turn_state: None,
|
||||
turn_metadata: None,
|
||||
beta_features: None,
|
||||
include_timing_metrics: None,
|
||||
version: None,
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: context.cookie_header,
|
||||
installation_id: Some(&identity.installation_id),
|
||||
session_id: None,
|
||||
thread_id: None,
|
||||
client_request_id: None,
|
||||
turn_id: None,
|
||||
})
|
||||
.await;
|
||||
|
||||
let raw = match usage {
|
||||
Ok(raw) => raw,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = %account_id,
|
||||
quota_verify_required = true,
|
||||
quota_verify_result = "upstream_error",
|
||||
retry_with_another_account = false,
|
||||
error = %error,
|
||||
"Failed to verify stale quota state before upstream request"
|
||||
);
|
||||
return QuotaEnterOutcome::Ready(Box::new(acquired));
|
||||
}
|
||||
};
|
||||
|
||||
let quota = quota_from_usage(&raw);
|
||||
context
|
||||
.account_pool
|
||||
.apply_quota_snapshot(&account_id, "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)
|
||||
}
|
||||
|
||||
@@ -200,7 +90,11 @@ fn classify_upstream(
|
||||
account_id: &str,
|
||||
facts: &UpstreamFailureFacts,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
classified_quota_failure(account_id, &facts.body, classify_upstream_failure(facts)?)
|
||||
classified_quota_failure(
|
||||
account_id,
|
||||
&facts.body,
|
||||
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
|
||||
)
|
||||
}
|
||||
|
||||
fn classify_stream(
|
||||
@@ -208,7 +102,11 @@ fn classify_stream(
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> Option<ClassifiedQuotaFailure> {
|
||||
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
|
||||
classified_quota_failure(account_id, &body, classify_response_failure(failure)?)
|
||||
classified_quota_failure(
|
||||
account_id,
|
||||
&body,
|
||||
classify_account_failure(&super::failure_observation::from_response(failure))?,
|
||||
)
|
||||
}
|
||||
|
||||
fn classified_quota_failure(
|
||||
@@ -229,19 +127,3 @@ fn classified_quota_failure(
|
||||
effect: classified.effect,
|
||||
})
|
||||
}
|
||||
|
||||
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
|
||||
let limit_reached = quota_snapshot_limit_reached(quota);
|
||||
acquired.account.quota_verify_required = false;
|
||||
acquired.account.quota_limit_reached = limit_reached;
|
||||
acquired.account.quota_cooldown_until = limit_reached
|
||||
.then_some(quota_snapshot_reset_at(quota))
|
||||
.flatten();
|
||||
if let Some(reset_at) = quota_snapshot_reset_at(quota) {
|
||||
acquired.account.window_reset_at = Some(reset_at);
|
||||
if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) {
|
||||
acquired.account.limit_window_seconds = Some(limit_window_seconds);
|
||||
}
|
||||
}
|
||||
acquired
|
||||
}
|
||||
|
||||
@@ -10,23 +10,25 @@ use crate::{
|
||||
infra::time::elapsed_millis_i64,
|
||||
telemetry::{
|
||||
ops::types::OpsErrorLog,
|
||||
recorder::{
|
||||
DispatchErrorLogRecord, Recorder, enrich_event_route_metadata,
|
||||
enrich_response_request_semantics, enrich_usage_record_identity,
|
||||
event_kind as response_event_kind, reasoning_effort_from_request,
|
||||
record_dispatch_error_event,
|
||||
recorder::Recorder,
|
||||
usage::types::{
|
||||
UsageRecord, UsageRecordLevel, metadata_i64, metadata_service_tier,
|
||||
metadata_string as metadata_string_any,
|
||||
},
|
||||
usage::types::UsageRecordLevel,
|
||||
},
|
||||
upstream::openai::{
|
||||
protocol::events::TokenUsage,
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
transport::{
|
||||
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
|
||||
WebSocketPoolDecision,
|
||||
CodexBackendTransport, CodexClientError, CodexResponseMetadata,
|
||||
CodexUpstreamDiagnostics, WebSocketPoolDecision,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const RESPONSE_EVENT_KIND: &str = "v1.response";
|
||||
const RESPONSES_API_KIND: &str = "responses";
|
||||
|
||||
use crate::dispatch::{
|
||||
errors::{
|
||||
ResponseDispatchError, dispatch_error_metadata, enrich_response_dispatch_error_metadata,
|
||||
@@ -86,6 +88,119 @@ pub(super) struct ResponseDispatchErrorEventRecord<'a> {
|
||||
pub(super) error: &'a ResponseDispatchError,
|
||||
}
|
||||
|
||||
pub(super) struct ResponseUsageEventRecord<'a> {
|
||||
pub(super) recorder: &'a Recorder,
|
||||
pub(super) request_id: &'a str,
|
||||
pub(super) client_api_key_id: Option<&'a str>,
|
||||
pub(super) account_id: &'a str,
|
||||
pub(super) route: &'a str,
|
||||
pub(super) model: &'a str,
|
||||
pub(super) requested_model: Option<&'a str>,
|
||||
pub(super) client_ip: Option<&'a str>,
|
||||
pub(super) client_user_agent: Option<&'a str>,
|
||||
pub(super) reasoning_effort: Option<&'a str>,
|
||||
pub(super) service_tier: Option<&'a str>,
|
||||
pub(super) started_at: Instant,
|
||||
pub(super) status_code: i64,
|
||||
pub(super) message: &'a str,
|
||||
pub(super) usage: Option<TokenUsage>,
|
||||
pub(super) metadata: Value,
|
||||
pub(super) rate_limit_headers: &'a [(String, String)],
|
||||
}
|
||||
|
||||
pub(super) struct LiveResponseStreamEventRecord<'a> {
|
||||
pub(super) context: &'a StreamContext<'a>,
|
||||
pub(super) status_code: i64,
|
||||
pub(super) level: UsageRecordLevel,
|
||||
pub(super) message: &'a str,
|
||||
pub(super) usage: Option<TokenUsage>,
|
||||
pub(super) metadata: Value,
|
||||
pub(super) rate_limit_headers: &'a [(String, String)],
|
||||
pub(super) body: &'a str,
|
||||
}
|
||||
|
||||
struct DispatchErrorLogRecord<'a> {
|
||||
recorder: &'a Recorder,
|
||||
request_id: &'a str,
|
||||
client_api_key_id: Option<&'a str>,
|
||||
provider: Option<&'a str>,
|
||||
account_id: Option<&'a str>,
|
||||
route: &'a str,
|
||||
model: &'a str,
|
||||
started_at: Instant,
|
||||
status_code: i64,
|
||||
message: &'a str,
|
||||
metadata: Value,
|
||||
}
|
||||
|
||||
pub(super) async fn record_response_event(record: ResponseUsageEventRecord<'_>) {
|
||||
let mut metadata = record.metadata;
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
|
||||
let mut event = UsageRecord::new(
|
||||
RESPONSE_EVENT_KIND,
|
||||
record.message,
|
||||
record.account_id,
|
||||
record.model,
|
||||
record.status_code,
|
||||
);
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
|
||||
event.route = Some(record.route.to_string());
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
event.requested_model = normalized_string(record.requested_model.or(Some(record.model)));
|
||||
event.upstream_model = normalized_string(Some(record.model));
|
||||
event.service_tier = normalized_string(record.service_tier);
|
||||
enrich_usage_record_identity(
|
||||
&mut metadata,
|
||||
record.requested_model,
|
||||
record.model,
|
||||
record.client_ip,
|
||||
record.client_user_agent,
|
||||
record.reasoning_effort,
|
||||
record.service_tier,
|
||||
);
|
||||
if !record.rate_limit_headers.is_empty()
|
||||
&& let Some(object) = metadata.as_object_mut()
|
||||
{
|
||||
object.insert(
|
||||
"rateLimitHeaders".to_string(),
|
||||
serde_json::json!(record.rate_limit_headers),
|
||||
);
|
||||
}
|
||||
lift_success_fact_fields(&mut event, &mut metadata, record.usage);
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
|
||||
let mut metadata = record.metadata;
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
|
||||
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, record.message);
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
|
||||
event.provider = record.provider.map(ToString::to_string);
|
||||
event.account_id = record.account_id.map(ToString::to_string);
|
||||
event.route = Some(record.route.to_string());
|
||||
event.model = Some(record.model.to_string());
|
||||
event.status_code = Some(record.status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
|
||||
if let Err(error) = record.recorder.record_error(event).await {
|
||||
tracing::error!(
|
||||
account_id = record.account_id.unwrap_or(""),
|
||||
error = %error,
|
||||
"Failed to record dispatch error event"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn record_response_dispatch_error_event(
|
||||
record: ResponseDispatchErrorEventRecord<'_>,
|
||||
) {
|
||||
@@ -172,6 +287,7 @@ pub(super) async fn record_response_upstream_error_event(
|
||||
);
|
||||
enrich_response_request_semantics(&mut metadata, record.request);
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
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");
|
||||
}
|
||||
@@ -221,10 +337,7 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
record.request.service_tier(),
|
||||
);
|
||||
enrich_response_request_semantics(&mut metadata, record.request);
|
||||
let mut event = OpsErrorLog::new(
|
||||
response_event_kind(record.route),
|
||||
"v1 responses stream failed",
|
||||
);
|
||||
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, "v1 responses stream failed");
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.client_api_key_id = record.request.client_api_key_id.clone();
|
||||
event.provider = Some("openai".to_string());
|
||||
@@ -242,11 +355,209 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
);
|
||||
}
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
||||
fn lift_error_fact_fields(event: &mut OpsErrorLog) {
|
||||
event.client_status_code = event
|
||||
.client_status_code
|
||||
.or_else(|| metadata_i64(&event.metadata, &["clientStatusCode", "client_status_code"]));
|
||||
event.upstream_status_code = event
|
||||
.upstream_status_code
|
||||
.or_else(|| metadata_i64(&event.metadata, &["upstreamStatusCode", "upstreamStatus"]));
|
||||
event.transport = event
|
||||
.transport
|
||||
.take()
|
||||
.or_else(|| metadata_string_any(&event.metadata, &["transport"]));
|
||||
event.attempt_index = event
|
||||
.attempt_index
|
||||
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"]));
|
||||
event.failure_class = event
|
||||
.failure_class
|
||||
.take()
|
||||
.or_else(|| metadata_string_any(&event.metadata, &["failureClass", "failure_class"]));
|
||||
event.response_id = event
|
||||
.response_id
|
||||
.take()
|
||||
.or_else(|| metadata_string_any(&event.metadata, &["responseId", "response_id"]));
|
||||
event.upstream_request_id = event.upstream_request_id.take().or_else(|| {
|
||||
metadata_string_any(
|
||||
&event.metadata,
|
||||
&[
|
||||
"upstreamRequestId",
|
||||
"upstream_request_id",
|
||||
"openaiRequestId",
|
||||
],
|
||||
)
|
||||
});
|
||||
event.service_tier = event
|
||||
.service_tier
|
||||
.take()
|
||||
.or_else(|| metadata_service_tier(&event.metadata).map(ToString::to_string));
|
||||
|
||||
let Some(metadata) = event.metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in [
|
||||
"clientStatusCode",
|
||||
"client_status_code",
|
||||
"upstreamStatusCode",
|
||||
"upstreamStatus",
|
||||
"transport",
|
||||
"attemptIndex",
|
||||
"attempt_index",
|
||||
"failureClass",
|
||||
"failure_class",
|
||||
"responseId",
|
||||
"response_id",
|
||||
"upstreamRequestId",
|
||||
"upstream_request_id",
|
||||
"openaiRequestId",
|
||||
"serviceTier",
|
||||
] {
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn lift_success_fact_fields(
|
||||
event: &mut UsageRecord,
|
||||
metadata: &mut Value,
|
||||
usage: Option<TokenUsage>,
|
||||
) {
|
||||
event.transport = metadata_string(metadata, "transport");
|
||||
event.attempt_index = metadata_nonnegative_i64(metadata, "attemptIndex");
|
||||
event.response_id = metadata_string(metadata, "responseId");
|
||||
event.upstream_request_id = metadata_string(metadata, "upstreamRequestId")
|
||||
.or_else(|| metadata_string(metadata, "openaiRequestId"));
|
||||
event.first_token_ms = metadata_nonnegative_i64(metadata, "firstTokenMs");
|
||||
if let Some(usage) = usage {
|
||||
event.input_tokens = i64::try_from(usage.input_tokens).ok();
|
||||
event.output_tokens = i64::try_from(usage.output_tokens).ok();
|
||||
event.cached_tokens = i64::try_from(usage.cached_tokens).ok();
|
||||
event.cache_write_tokens = i64::try_from(usage.cache_write_tokens).ok();
|
||||
event.reasoning_tokens = i64::try_from(usage.reasoning_tokens).ok();
|
||||
}
|
||||
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in [
|
||||
"usage",
|
||||
"route",
|
||||
"apiKind",
|
||||
"requestedModel",
|
||||
"upstreamModel",
|
||||
"serviceTier",
|
||||
"statusCode",
|
||||
"transport",
|
||||
"attemptIndex",
|
||||
"responseId",
|
||||
"upstreamRequestId",
|
||||
"openaiRequestId",
|
||||
"firstTokenMs",
|
||||
] {
|
||||
object.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn metadata_string(metadata: &Value, field: &str) -> Option<String> {
|
||||
metadata
|
||||
.get(field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn metadata_nonnegative_i64(metadata: &Value, field: &str) -> Option<i64> {
|
||||
metadata
|
||||
.get(field)
|
||||
.and_then(Value::as_i64)
|
||||
.filter(|value| *value >= 0)
|
||||
}
|
||||
|
||||
fn normalized_string(value: Option<&str>) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
object
|
||||
.entry("route".to_string())
|
||||
.or_insert_with(|| Value::String(route.to_string()));
|
||||
object
|
||||
.entry("apiKind".to_string())
|
||||
.or_insert_with(|| Value::String(RESPONSES_API_KIND.to_string()));
|
||||
}
|
||||
|
||||
fn enrich_usage_record_identity(
|
||||
metadata: &mut Value,
|
||||
requested_model: Option<&str>,
|
||||
upstream_model: &str,
|
||||
client_ip: Option<&str>,
|
||||
client_user_agent: Option<&str>,
|
||||
reasoning_effort: Option<&str>,
|
||||
service_tier: Option<&str>,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
let upstream_model = upstream_model.trim();
|
||||
let requested_model = requested_model
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(upstream_model);
|
||||
object.insert(
|
||||
"requestedModel".to_string(),
|
||||
Value::String(requested_model.to_string()),
|
||||
);
|
||||
object.insert(
|
||||
"upstreamModel".to_string(),
|
||||
Value::String(upstream_model.to_string()),
|
||||
);
|
||||
insert_trimmed_string(object, "clientIp", client_ip);
|
||||
insert_trimmed_string(object, "userAgent", client_user_agent);
|
||||
insert_trimmed_string(object, "reasoningEffort", reasoning_effort);
|
||||
insert_trimmed_string(object, "serviceTier", service_tier);
|
||||
}
|
||||
|
||||
pub(super) fn enrich_response_request_semantics(
|
||||
metadata: &mut Value,
|
||||
request: &CodexResponsesRequest,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
let semantics = 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(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
|
||||
request
|
||||
.reasoning()?
|
||||
.get("effort")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn insert_trimmed_string(object: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
object.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_stream_metadata_flag(metadata: &mut Value) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
*metadata = serde_json::json!({ "stream": true });
|
||||
@@ -313,6 +624,11 @@ fn enrich_live_response_stream_metadata(
|
||||
"firstEventMs",
|
||||
context.transport_metrics.first_event_ms,
|
||||
);
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"openaiProcessingMs",
|
||||
openai_processing_ms(context.response_metadata),
|
||||
);
|
||||
if let Some(http_version) = context.transport_metrics.http_version.as_deref() {
|
||||
object
|
||||
.entry("httpVersion".to_string())
|
||||
@@ -505,15 +821,17 @@ fn upstream_failure_client_status_code(upstream_status: i64) -> i64 {
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) async fn record_live_response_stream_event(
|
||||
context: &StreamContext<'_>,
|
||||
status_code: i64,
|
||||
level: UsageRecordLevel,
|
||||
message: &str,
|
||||
mut metadata: Value,
|
||||
rate_limit_headers: &[(String, String)],
|
||||
body: &str,
|
||||
) {
|
||||
pub(super) async fn record_live_response_stream_event(record: LiveResponseStreamEventRecord<'_>) {
|
||||
let LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code,
|
||||
level,
|
||||
message,
|
||||
usage,
|
||||
mut metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
} = record;
|
||||
let effective_model = context
|
||||
.response_metadata
|
||||
.effective_model
|
||||
@@ -562,29 +880,28 @@ pub(super) async fn record_live_response_stream_event(
|
||||
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,
|
||||
client_api_key_id: context.request.client_api_key_id.as_deref(),
|
||||
account_id: context.account_id,
|
||||
route: context.route,
|
||||
model: effective_model,
|
||||
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),
|
||||
service_tier: context.request.service_tier(),
|
||||
started_at: context.started_at,
|
||||
status_code,
|
||||
message,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
},
|
||||
)
|
||||
record_response_event(ResponseUsageEventRecord {
|
||||
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,
|
||||
model: effective_model,
|
||||
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),
|
||||
service_tier: context.request.service_tier(),
|
||||
started_at: context.started_at,
|
||||
status_code,
|
||||
message,
|
||||
usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
})
|
||||
.await;
|
||||
} else {
|
||||
let mut event = OpsErrorLog::new(response_event_kind(context.route), message);
|
||||
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, 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());
|
||||
@@ -594,6 +911,7 @@ pub(super) async fn record_live_response_stream_event(
|
||||
event.status_code = Some(status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
|
||||
event.metadata = metadata;
|
||||
lift_error_fact_fields(&mut event);
|
||||
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");
|
||||
}
|
||||
@@ -683,16 +1001,40 @@ fn metadata_string_field<'a>(metadata: &'a Value, field: &str) -> Option<&'a str
|
||||
metadata.get(field).and_then(Value::as_str)
|
||||
}
|
||||
|
||||
pub(super) fn insert_first_token_ms(metadata: &mut Value, first_token_ms: Option<i64>) {
|
||||
let Some(first_token_ms) = first_token_ms else {
|
||||
pub(super) fn insert_output_timing_ms(
|
||||
metadata: &mut Value,
|
||||
first_token_ms: Option<i64>,
|
||||
first_reasoning_ms: Option<i64>,
|
||||
first_text_ms: Option<i64>,
|
||||
) {
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
insert_optional_i64(object, "firstTokenMs", first_token_ms);
|
||||
insert_optional_i64(object, "firstReasoningMs", first_reasoning_ms);
|
||||
insert_optional_i64(object, "firstTextMs", first_text_ms);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn insert_openai_processing_ms(
|
||||
metadata: &mut Value,
|
||||
response_metadata: &CodexResponseMetadata,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
object.insert(
|
||||
"firstTokenMs".to_string(),
|
||||
Value::Number(first_token_ms.into()),
|
||||
);
|
||||
}
|
||||
insert_optional_i64(
|
||||
object,
|
||||
"openaiProcessingMs",
|
||||
openai_processing_ms(response_metadata),
|
||||
);
|
||||
}
|
||||
|
||||
fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option<i64> {
|
||||
response_metadata
|
||||
.client_headers
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms"))
|
||||
.and_then(|(_, value)| value.parse::<i64>().ok())
|
||||
.filter(|value| *value >= 0)
|
||||
}
|
||||
|
||||
pub(super) fn insert_websocket_pool_decision(
|
||||
|
||||
@@ -2,15 +2,20 @@
|
||||
|
||||
mod events;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
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,
|
||||
LiveResponseStreamEventRecord, ResponseUpstreamErrorEventRecord,
|
||||
ResponseUsageEventRecord, enrich_response_request_semantics,
|
||||
insert_openai_processing_ms, insert_output_timing_ms, insert_response_status_metadata,
|
||||
insert_response_trace_metadata, insert_response_upstream_diagnostics,
|
||||
insert_websocket_pool_decision, reasoning_effort_from_request,
|
||||
record_live_response_stream_event, record_response_event,
|
||||
record_response_upstream_error_event,
|
||||
},
|
||||
controllers::{
|
||||
AttemptUpstreamErrorContext, CompleteExit, DispatchErrorObservation,
|
||||
@@ -27,13 +32,7 @@ use crate::{
|
||||
},
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
telemetry::{
|
||||
recorder::{
|
||||
Recorder, enrich_response_request_semantics, reasoning_effort_from_request,
|
||||
record_response_event,
|
||||
},
|
||||
usage::types::{ResponseUsageRecord, UsageRecordLevel},
|
||||
},
|
||||
telemetry::{recorder::Recorder, usage::types::UsageRecordLevel},
|
||||
upstream::openai::{
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
transport::{
|
||||
@@ -149,13 +148,27 @@ impl TelemetryController {
|
||||
.effective_model
|
||||
.as_deref()
|
||||
.unwrap_or(exit.display_model);
|
||||
let first_token_ms = request_timing_ms(
|
||||
exit.started_at,
|
||||
exit.attempt.started_at(),
|
||||
exit.response.first_token_ms,
|
||||
);
|
||||
let first_reasoning_ms = request_timing_ms(
|
||||
exit.started_at,
|
||||
exit.attempt.started_at(),
|
||||
exit.response.first_reasoning_ms,
|
||||
);
|
||||
let first_text_ms = request_timing_ms(
|
||||
exit.started_at,
|
||||
exit.attempt.started_at(),
|
||||
exit.response.first_text_ms,
|
||||
);
|
||||
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(),
|
||||
@@ -167,6 +180,13 @@ impl TelemetryController {
|
||||
"firstEventMs": exit.response.transport_metrics.first_event_ms,
|
||||
"httpVersion": exit.response.transport_metrics.http_version.as_deref(),
|
||||
});
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
first_token_ms,
|
||||
first_reasoning_ms,
|
||||
first_text_ms,
|
||||
);
|
||||
insert_openai_processing_ms(&mut metadata, &exit.response.response_metadata);
|
||||
insert_response_status_metadata(
|
||||
&mut metadata,
|
||||
200,
|
||||
@@ -177,7 +197,7 @@ impl TelemetryController {
|
||||
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 {
|
||||
record_response_event(ResponseUsageEventRecord {
|
||||
recorder: exit.recorder,
|
||||
request_id: exit.request_id,
|
||||
client_api_key_id: exit.request.client_api_key_id.as_deref(),
|
||||
@@ -196,6 +216,7 @@ impl TelemetryController {
|
||||
} else {
|
||||
"v1 responses incomplete"
|
||||
},
|
||||
usage: exit.response.usage,
|
||||
metadata,
|
||||
rate_limit_headers: &exit.response.rate_limit_headers,
|
||||
})
|
||||
@@ -248,22 +269,44 @@ impl TelemetryController {
|
||||
"firstEventMs": summary.first_event_ms,
|
||||
"usage": summary.usage,
|
||||
});
|
||||
insert_first_token_ms(&mut metadata, summary.first_token_ms);
|
||||
record_live_response_stream_event(
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
summary.first_token_ms,
|
||||
summary.first_reasoning_ms,
|
||||
summary.first_text_ms,
|
||||
);
|
||||
record_live_response_stream_event(LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code,
|
||||
UsageRecordLevel::Warn,
|
||||
"v1 responses stream cancelled",
|
||||
level: UsageRecordLevel::Warn,
|
||||
message: "v1 responses stream cancelled",
|
||||
usage: summary.usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
)
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn request_timing_ms(
|
||||
request_started_at: Instant,
|
||||
attempt_started_at: Instant,
|
||||
attempt_first_token_ms: Option<i64>,
|
||||
) -> Option<i64> {
|
||||
let attempt_first_token_ms = u64::try_from(attempt_first_token_ms?).ok()?;
|
||||
let output_at =
|
||||
attempt_started_at.checked_add(Duration::from_millis(attempt_first_token_ms))?;
|
||||
i64::try_from(
|
||||
output_at
|
||||
.saturating_duration_since(request_started_at)
|
||||
.as_millis(),
|
||||
)
|
||||
.ok()
|
||||
}
|
||||
|
||||
async fn record_success(
|
||||
context: &StreamContext<'_>,
|
||||
summary: &StreamSummary,
|
||||
@@ -272,27 +315,34 @@ async fn record_success(
|
||||
rate_limit_headers: &[(String, String)],
|
||||
body: &str,
|
||||
) {
|
||||
record_live_response_stream_event(
|
||||
let mut metadata = serde_json::json!({
|
||||
"stream": true,
|
||||
"completed": completed,
|
||||
"incomplete": !completed,
|
||||
"responseId": response.get("id").and_then(Value::as_str),
|
||||
"firstEventMs": summary.first_event_ms,
|
||||
"usage": summary.usage,
|
||||
});
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
summary.first_token_ms,
|
||||
summary.first_reasoning_ms,
|
||||
summary.first_text_ms,
|
||||
);
|
||||
record_live_response_stream_event(LiveResponseStreamEventRecord {
|
||||
context,
|
||||
200,
|
||||
UsageRecordLevel::Info,
|
||||
if completed {
|
||||
status_code: 200,
|
||||
level: UsageRecordLevel::Info,
|
||||
message: 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,
|
||||
}),
|
||||
usage: summary.usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
)
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -344,16 +394,22 @@ async fn record_failure(
|
||||
.map_or(Value::Null, |response_id| {
|
||||
Value::String(response_id.clone())
|
||||
});
|
||||
insert_first_token_ms(&mut metadata, summary.first_token_ms);
|
||||
insert_output_timing_ms(
|
||||
&mut metadata,
|
||||
summary.first_token_ms,
|
||||
summary.first_reasoning_ms,
|
||||
summary.first_text_ms,
|
||||
);
|
||||
metadata["firstEventMs"] = Value::Number(summary.first_event_ms.into());
|
||||
record_live_response_stream_event(
|
||||
record_live_response_stream_event(LiveResponseStreamEventRecord {
|
||||
context,
|
||||
status_code,
|
||||
UsageRecordLevel::Error,
|
||||
"v1 responses stream failed",
|
||||
level: UsageRecordLevel::Error,
|
||||
message: "v1 responses stream failed",
|
||||
usage: summary.usage,
|
||||
metadata,
|
||||
rate_limit_headers,
|
||||
body,
|
||||
)
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -1,12 +1,24 @@
|
||||
//! 配额响应头与 token 用量写入的唯一 feature owner。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::lifecycle::contract::{
|
||||
AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
|
||||
},
|
||||
fleet::pool::AccountPoolService,
|
||||
fleet::{
|
||||
pool::AccountPoolService,
|
||||
quota::{
|
||||
QuotaCreditsObservation, QuotaLimitObservation, QuotaObservation,
|
||||
QuotaWindowObservation,
|
||||
},
|
||||
usage::ResponseUsage,
|
||||
},
|
||||
upstream::openai::{
|
||||
protocol::{events::TokenUsage, responses::CodexResponsesRequest},
|
||||
protocol::{
|
||||
events::{RateLimitWindow, TokenUsage, parse_rate_limit_headers},
|
||||
responses::CodexResponsesRequest,
|
||||
},
|
||||
transport::{CodexRateLimitHeaderUpdates, CodexTurnStateUpdate},
|
||||
},
|
||||
};
|
||||
@@ -71,23 +83,41 @@ impl UsageController {
|
||||
) {
|
||||
if let Some(usage) = usage {
|
||||
account_pool
|
||||
.record_response_usage(account_id, usage, image_generation_requested)
|
||||
.record_response_usage(
|
||||
account_id,
|
||||
response_usage(usage),
|
||||
image_generation_requested,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn sync_passive_quota(
|
||||
account_pool: &AccountPoolService,
|
||||
account_id: &str,
|
||||
plan_type: Option<&str>,
|
||||
headers: &[(String, String)],
|
||||
) {
|
||||
let Some(observation) = quota_observation(headers) else {
|
||||
return;
|
||||
};
|
||||
account_pool
|
||||
.sync_passive_quota_observation_for_account(account_id, plan_type, &observation)
|
||||
.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;
|
||||
Self::sync_passive_quota(
|
||||
exit.account_pool,
|
||||
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 {
|
||||
@@ -95,7 +125,7 @@ impl UsageController {
|
||||
};
|
||||
if let Some(usage) = exit.usage {
|
||||
exit.account_pool
|
||||
.record_token_usage(exit.account_id, &usage)
|
||||
.record_response_usage(exit.account_id, response_usage(usage), false)
|
||||
.await;
|
||||
}
|
||||
StreamUsageExit {
|
||||
@@ -104,3 +134,60 @@ impl UsageController {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn response_usage(usage: TokenUsage) -> ResponseUsage {
|
||||
ResponseUsage {
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_tokens: usage.cached_tokens,
|
||||
reasoning_tokens: usage.reasoning_tokens,
|
||||
total_tokens: usage.total_tokens,
|
||||
image_input_tokens: usage.image_input_tokens,
|
||||
image_output_tokens: usage.image_output_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_observation(headers: &[(String, String)]) -> Option<QuotaObservation> {
|
||||
let parsed = parse_rate_limit_headers(headers)?;
|
||||
Some(QuotaObservation {
|
||||
limits: parsed
|
||||
.limits
|
||||
.into_iter()
|
||||
.map(|(limit_id, details)| {
|
||||
(
|
||||
limit_id,
|
||||
QuotaLimitObservation {
|
||||
limit_id: details.limit_id,
|
||||
limit_name: details.limit_name,
|
||||
metered_feature: None,
|
||||
allowed: details.allowed,
|
||||
limit_reached: details.limit_reached,
|
||||
primary: details.primary.map(quota_window_observation),
|
||||
secondary: details.secondary.map(quota_window_observation),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
active_limit: parsed.active_limit,
|
||||
credits: parsed.credits.map(|credits| QuotaCreditsObservation {
|
||||
has_credits: credits.has_credits,
|
||||
unlimited: credits.unlimited,
|
||||
overage_limit_reached: false,
|
||||
balance: credits.balance.map(Value::String),
|
||||
}),
|
||||
spend_control: None,
|
||||
plan_type: parsed.plan_type,
|
||||
promo_message: parsed.promo_message,
|
||||
rate_limit_reached_type: parsed.rate_limit_reached_type,
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_window_observation(window: RateLimitWindow) -> QuotaWindowObservation {
|
||||
QuotaWindowObservation {
|
||||
used_percent: window.used_percent,
|
||||
window_minutes: window.window_minutes,
|
||||
reset_at: window.reset_at,
|
||||
used: None,
|
||||
limit: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -238,6 +238,15 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_error_rate_limit_headers(error: &CodexClientError) -> Vec<(String, String)> {
|
||||
match error {
|
||||
CodexClientError::Upstream {
|
||||
rate_limit_headers, ..
|
||||
} => rate_limit_headers.clone(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 {
|
||||
match error {
|
||||
CodexClientError::Upstream { status, .. } => status.as_u16(),
|
||||
|
||||
@@ -8,9 +8,8 @@ use crate::{
|
||||
dispatch::{
|
||||
affinity::AccountIdentityService,
|
||||
controllers::{
|
||||
AttemptAccountPreparation, AttemptAccountPreparationContext, AttemptRoutePreparation,
|
||||
AttemptRoutePreparationContext, AttemptUpstreamErrorContext, ControllerRequestScope,
|
||||
ControllerSet,
|
||||
AttemptRoutePreparation, AttemptRoutePreparationContext, AttemptUpstreamErrorContext,
|
||||
ControllerRequestScope, ControllerSet,
|
||||
},
|
||||
errors::{ClientFailure, ResponseDispatchError},
|
||||
failure::exhaustion::AccountExhaustionTracker,
|
||||
@@ -50,13 +49,8 @@ use super::{
|
||||
};
|
||||
|
||||
pub(in crate::dispatch) enum AttemptMode {
|
||||
Complete {
|
||||
started_at: Instant,
|
||||
tuple_schema: Option<Value>,
|
||||
},
|
||||
Stream {
|
||||
tuple_schema: Option<Value>,
|
||||
},
|
||||
Complete { tuple_schema: Option<Value> },
|
||||
Stream { tuple_schema: Option<Value> },
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) struct AttemptRunnerDependencies<'a> {
|
||||
@@ -196,25 +190,11 @@ impl<'a> AttemptRunner<'a> {
|
||||
.controller_scope
|
||||
.as_mut()
|
||||
.ok_or(AttemptContractError::Terminal)?;
|
||||
let prepared_account = match self
|
||||
let prepared_account = 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();
|
||||
}
|
||||
};
|
||||
.prepare_attempt_account(acquired)
|
||||
.await;
|
||||
let acquired = prepared_account.lease;
|
||||
let cookie_header = prepared_account.cookie_header;
|
||||
|
||||
@@ -288,15 +268,12 @@ impl<'a> AttemptRunner<'a> {
|
||||
};
|
||||
|
||||
match &self.mode {
|
||||
AttemptMode::Complete {
|
||||
started_at,
|
||||
tuple_schema,
|
||||
} => {
|
||||
AttemptMode::Complete { tuple_schema } => {
|
||||
let result = create_response_with_account(
|
||||
upstream_context,
|
||||
&account_scoped_request,
|
||||
prepared_transport,
|
||||
*started_at,
|
||||
attempt.started_at(),
|
||||
)
|
||||
.await;
|
||||
let attempt_request = account_scoped_request.into_request();
|
||||
@@ -420,7 +397,7 @@ impl<'a> AttemptRunner<'a> {
|
||||
body,
|
||||
&mut decoder,
|
||||
attempt_request.stream_commit_policy,
|
||||
self.dependencies.started_at,
|
||||
attempt.started_at(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -658,6 +635,7 @@ impl<'a> AttemptRunner<'a> {
|
||||
PendingAttempt::StreamResponse(stream) => {
|
||||
let PendingStreamResponse {
|
||||
account,
|
||||
attempt_request,
|
||||
attempt,
|
||||
lease,
|
||||
response,
|
||||
@@ -668,6 +646,7 @@ impl<'a> AttemptRunner<'a> {
|
||||
} = *stream;
|
||||
EstablishedResponse::Stream(Box::new(EstablishedStream {
|
||||
context: self.take_established_context(account, attempt)?,
|
||||
attempt_request,
|
||||
lease,
|
||||
response,
|
||||
decoder,
|
||||
@@ -775,9 +754,6 @@ fn observation_kind(
|
||||
kind: *kind,
|
||||
}
|
||||
}
|
||||
PendingAttempt::CandidatePreparationRejected { .. } => {
|
||||
AttemptObservationKind::CandidatePreparationRejected
|
||||
}
|
||||
PendingAttempt::RoutePreparationRejected { message, .. } => {
|
||||
AttemptObservationKind::RoutePreparationRejected {
|
||||
message: message.clone(),
|
||||
@@ -846,17 +822,15 @@ async fn rejection_error(
|
||||
})
|
||||
}
|
||||
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::NoCandidate => 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,
|
||||
@@ -908,8 +882,7 @@ fn pending_account(pending: &PendingAttempt) -> Option<&str> {
|
||||
fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> {
|
||||
match pending {
|
||||
PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None,
|
||||
PendingAttempt::CandidatePreparationRejected { account }
|
||||
| PendingAttempt::RoutePreparationRejected { account, .. }
|
||||
PendingAttempt::RoutePreparationRejected { account, .. }
|
||||
| PendingAttempt::UpstreamFailure { account, .. }
|
||||
| PendingAttempt::ProtocolFailure { account, .. } => Some(account),
|
||||
PendingAttempt::CompleteResponse(complete) => Some(&complete.account),
|
||||
@@ -925,7 +898,6 @@ fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDi
|
||||
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt),
|
||||
PendingAttempt::NoCandidate
|
||||
| PendingAttempt::PinnedCandidateUnavailable { .. }
|
||||
| PendingAttempt::CandidatePreparationRejected { .. }
|
||||
| PendingAttempt::RoutePreparationRejected { .. } => None,
|
||||
}
|
||||
}
|
||||
@@ -942,7 +914,6 @@ fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRe
|
||||
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request),
|
||||
PendingAttempt::NoCandidate
|
||||
| PendingAttempt::PinnedCandidateUnavailable { .. }
|
||||
| PendingAttempt::CandidatePreparationRejected { .. }
|
||||
| PendingAttempt::RoutePreparationRejected { .. } => None,
|
||||
}
|
||||
}
|
||||
@@ -957,7 +928,6 @@ fn pending_transport(
|
||||
PendingAttempt::StreamResponse(stream) => Some(stream.response.transport),
|
||||
PendingAttempt::NoCandidate
|
||||
| PendingAttempt::PinnedCandidateUnavailable { .. }
|
||||
| PendingAttempt::CandidatePreparationRejected { .. }
|
||||
| PendingAttempt::RoutePreparationRejected { .. } => None,
|
||||
}
|
||||
}
|
||||
@@ -979,7 +949,6 @@ fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractE
|
||||
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",
|
||||
|
||||
@@ -93,7 +93,6 @@ pub(in crate::dispatch) enum AttemptObservationKind {
|
||||
ledger: CandidateLedgerFacts,
|
||||
last_exhausted: Option<ExhaustedAccount>,
|
||||
},
|
||||
CandidatePreparationRejected,
|
||||
RoutePreparationRejected {
|
||||
message: String,
|
||||
},
|
||||
@@ -173,6 +172,7 @@ impl TryFrom<CollectedResponse> for EstablishedCompleteBody {
|
||||
|
||||
pub(in crate::dispatch) struct EstablishedStream {
|
||||
pub context: EstablishedAttemptContext,
|
||||
pub attempt_request: CodexResponsesRequest,
|
||||
pub lease: AccountLease,
|
||||
pub response: CodexBackendStreamingResponse,
|
||||
pub decoder: CanonicalStreamDecoder,
|
||||
@@ -256,9 +256,6 @@ pub(super) enum PendingAttempt {
|
||||
account_id: String,
|
||||
kind: PinnedCandidateAcquireFailureKind,
|
||||
},
|
||||
CandidatePreparationRejected {
|
||||
account: Account,
|
||||
},
|
||||
RoutePreparationRejected {
|
||||
account: Account,
|
||||
message: String,
|
||||
|
||||
@@ -27,6 +27,8 @@ pub(in crate::dispatch) struct StreamFinalizer {
|
||||
controllers: ControllerSet,
|
||||
controller_scope: ControllerRequestScope,
|
||||
controller_context: StreamControllerContext,
|
||||
request_input: Vec<serde_json::Value>,
|
||||
continued_from_previous_response: bool,
|
||||
account_lease: AccountLease,
|
||||
}
|
||||
|
||||
@@ -35,12 +37,16 @@ impl StreamFinalizer {
|
||||
controllers: ControllerSet,
|
||||
controller_scope: ControllerRequestScope,
|
||||
controller_context: StreamControllerContext,
|
||||
request_input: Vec<serde_json::Value>,
|
||||
continued_from_previous_response: bool,
|
||||
account_lease: AccountLease,
|
||||
) -> Self {
|
||||
Self {
|
||||
controllers,
|
||||
controller_scope,
|
||||
controller_context,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
account_lease,
|
||||
}
|
||||
}
|
||||
@@ -53,6 +59,18 @@ impl StreamFinalizer {
|
||||
self.controller_context.started_at
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn account_id(&self) -> &str {
|
||||
&self.controller_context.account_id
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn request_input(&self) -> &[serde_json::Value] {
|
||||
&self.request_input
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn continued_from_previous_response(&self) -> bool {
|
||||
self.continued_from_previous_response
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) async fn finalize(
|
||||
self,
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
@@ -64,19 +82,23 @@ impl StreamFinalizer {
|
||||
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())
|
||||
let mut lease_completion = Box::pin(account_lease.complete());
|
||||
if timeout(LEASE_FINALIZE_TIMEOUT, lease_completion.as_mut())
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
tracing::warn!(
|
||||
timeout_ms = LEASE_FINALIZE_TIMEOUT.as_millis(),
|
||||
"Timed out finalizing response account lease"
|
||||
"Response account lease is continuing finalization in the background"
|
||||
);
|
||||
// 超时只约束客户端收尾,不取消已经开始的 slot 释放与请求用量持久化。
|
||||
drop(tokio::spawn(lease_completion));
|
||||
}
|
||||
finish_client_stream(sender, event_sender, &mut summary).await;
|
||||
}
|
||||
|
||||
@@ -60,7 +60,7 @@ pub(in crate::dispatch) async fn enter_request(
|
||||
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());
|
||||
debug_assert_eq!(request.stream(), mode.stream());
|
||||
|
||||
let compact = request.semantics().compact;
|
||||
let tuple_schema = request.tuple_schema.clone();
|
||||
|
||||
@@ -28,6 +28,8 @@ pub(in crate::dispatch) struct StreamSummary {
|
||||
pub body: Vec<u8>,
|
||||
pub terminal_chunks: Vec<CanonicalResponseChunk>,
|
||||
pub first_token_ms: Option<i64>,
|
||||
pub first_reasoning_ms: Option<i64>,
|
||||
pub first_text_ms: Option<i64>,
|
||||
pub first_event_ms: i64,
|
||||
pub usage: Option<TokenUsage>,
|
||||
pub last_response_id: Option<String>,
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! Responses 请求级账号尝试轨迹。
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(in crate::dispatch) struct ResponseDispatchTrace {
|
||||
next_attempt_index: i64,
|
||||
@@ -10,6 +12,7 @@ pub(in crate::dispatch) struct ResponseDispatchTrace {
|
||||
pub(in crate::dispatch) struct ResponseDispatchAttempt {
|
||||
index: i64,
|
||||
account_id: String,
|
||||
started_at: Instant,
|
||||
}
|
||||
|
||||
impl ResponseDispatchTrace {
|
||||
@@ -20,6 +23,7 @@ impl ResponseDispatchTrace {
|
||||
let attempt = ResponseDispatchAttempt {
|
||||
index: self.next_attempt_index,
|
||||
account_id: account_id.to_string(),
|
||||
started_at: Instant::now(),
|
||||
};
|
||||
self.next_attempt_index += 1;
|
||||
self.attempts.push(attempt.clone());
|
||||
@@ -39,4 +43,8 @@ impl ResponseDispatchAttempt {
|
||||
pub(in crate::dispatch) fn account_id(&self) -> &str {
|
||||
&self.account_id
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn started_at(&self) -> Instant {
|
||||
self.started_at
|
||||
}
|
||||
}
|
||||
|
||||
@@ -79,6 +79,9 @@ pub type ResponseBodyStream =
|
||||
|
||||
/// Responses live SSE 调度结果。
|
||||
pub struct ResponseDispatchStream {
|
||||
pub(crate) account_id: String,
|
||||
pub(crate) request_input: Vec<Value>,
|
||||
pub(crate) continued_from_previous_response: bool,
|
||||
pub body: ResponseBodyStream,
|
||||
pub(crate) canonical_events: tokio::sync::mpsc::UnboundedReceiver<Vec<CanonicalResponseEvent>>,
|
||||
pub response_headers: Vec<(String, String)>,
|
||||
@@ -102,13 +105,25 @@ pub(crate) struct ConnectionReplayPlan {
|
||||
|
||||
/// API 从 canonical 响应事件采集的原始 transcript 事实。
|
||||
pub(crate) struct ConnectionTranscriptFacts {
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
response_id: String,
|
||||
output: Vec<Value>,
|
||||
}
|
||||
|
||||
impl ConnectionTranscriptFacts {
|
||||
pub(crate) fn new(response_id: String, output: Vec<Value>) -> Self {
|
||||
pub(crate) fn new(
|
||||
account_id: String,
|
||||
request_input: Vec<Value>,
|
||||
continued_from_previous_response: bool,
|
||||
response_id: String,
|
||||
output: Vec<Value>,
|
||||
) -> Self {
|
||||
Self {
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
response_id,
|
||||
output,
|
||||
}
|
||||
@@ -168,6 +183,9 @@ impl ResponseDispatchService {
|
||||
plan.inner,
|
||||
transcript.response_id,
|
||||
transcript.output,
|
||||
transcript.account_id,
|
||||
transcript.request_input,
|
||||
transcript.continued_from_previous_response,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -220,10 +238,7 @@ impl ResponseDispatchService {
|
||||
requested_model,
|
||||
started_at,
|
||||
},
|
||||
AttemptMode::Complete {
|
||||
started_at,
|
||||
tuple_schema,
|
||||
},
|
||||
AttemptMode::Complete { tuple_schema },
|
||||
request,
|
||||
controller_scope,
|
||||
candidates,
|
||||
|
||||
@@ -141,6 +141,7 @@ impl ResponseDispatchService {
|
||||
|
||||
let EstablishedStream {
|
||||
context,
|
||||
attempt_request,
|
||||
lease,
|
||||
response,
|
||||
decoder,
|
||||
@@ -159,6 +160,8 @@ impl ResponseDispatchService {
|
||||
.observe_models_etag(response.response_metadata.models_etag.as_deref());
|
||||
let attempts = trace.attempts().to_vec();
|
||||
let transport = response.transport;
|
||||
let request_input = attempt_request.input().to_vec();
|
||||
let continued_from_previous_response = attempt_request.previous_response_id().is_some();
|
||||
let finalizer = StreamFinalizer::new(
|
||||
controllers,
|
||||
controller_scope,
|
||||
@@ -185,6 +188,8 @@ impl ResponseDispatchService {
|
||||
attempts,
|
||||
started_at,
|
||||
},
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
lease,
|
||||
);
|
||||
Ok(spawn_live_response_stream(
|
||||
|
||||
@@ -57,6 +57,9 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
body: CodexBackendSseStream,
|
||||
shutdown: CancellationToken,
|
||||
) -> ResponseDispatchStream {
|
||||
let account_id = finalizer.account_id().to_string();
|
||||
let request_input = finalizer.request_input().to_vec();
|
||||
let continued_from_previous_response = finalizer.continued_from_previous_response();
|
||||
let response_headers = finalizer.response_headers().to_vec();
|
||||
let started_at = finalizer.started_at();
|
||||
let (sender, receiver) = mpsc::channel(8);
|
||||
@@ -81,6 +84,9 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
});
|
||||
|
||||
ResponseDispatchStream {
|
||||
account_id,
|
||||
request_input,
|
||||
continued_from_previous_response,
|
||||
body: Box::pin(MpscResponseBodyStream {
|
||||
receiver,
|
||||
cancel: Some(cancel_sender),
|
||||
@@ -117,6 +123,8 @@ async fn run_live_response_stream(
|
||||
let mut body_bytes = Vec::new();
|
||||
let mut terminal_chunks = Vec::new();
|
||||
let mut first_token_ms = None;
|
||||
let mut first_reasoning_ms = None;
|
||||
let mut first_text_ms = None;
|
||||
|
||||
let decoded_terminal = take_stream_terminal(&mut decoder);
|
||||
let terminal = if decoded_terminal.is_some() {
|
||||
@@ -128,13 +136,19 @@ async fn run_live_response_stream(
|
||||
.await
|
||||
.err()
|
||||
};
|
||||
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
|
||||
update_output_timing_ms(
|
||||
&decoder,
|
||||
started_at,
|
||||
&mut first_token_ms,
|
||||
&mut first_reasoning_ms,
|
||||
&mut first_text_ms,
|
||||
);
|
||||
if let Some(terminal) = terminal {
|
||||
return stream_summary(
|
||||
terminal,
|
||||
body_bytes,
|
||||
terminal_chunks,
|
||||
first_token_ms,
|
||||
(first_token_ms, first_reasoning_ms, first_text_ms),
|
||||
first_event_ms,
|
||||
&decoder,
|
||||
);
|
||||
@@ -150,7 +164,13 @@ async fn run_live_response_stream(
|
||||
match next {
|
||||
Some(Ok(chunk)) => match decoder.push(chunk) {
|
||||
Ok(batch) => {
|
||||
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
|
||||
update_output_timing_ms(
|
||||
&decoder,
|
||||
started_at,
|
||||
&mut first_token_ms,
|
||||
&mut first_reasoning_ms,
|
||||
&mut first_text_ms,
|
||||
);
|
||||
if let Some(terminal) = take_stream_terminal(&mut decoder) {
|
||||
match stage_terminal_batch(&mut body_bytes, &mut terminal_chunks, batch) {
|
||||
Ok(()) => break terminal,
|
||||
@@ -196,12 +216,18 @@ async fn run_live_response_stream(
|
||||
}
|
||||
};
|
||||
|
||||
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
|
||||
update_output_timing_ms(
|
||||
&decoder,
|
||||
started_at,
|
||||
&mut first_token_ms,
|
||||
&mut first_reasoning_ms,
|
||||
&mut first_text_ms,
|
||||
);
|
||||
stream_summary(
|
||||
terminal,
|
||||
body_bytes,
|
||||
terminal_chunks,
|
||||
first_token_ms,
|
||||
(first_token_ms, first_reasoning_ms, first_text_ms),
|
||||
first_event_ms,
|
||||
&decoder,
|
||||
)
|
||||
@@ -299,13 +325,29 @@ fn ensure_capture_capacity(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn update_first_token_ms(
|
||||
fn update_output_timing_ms(
|
||||
decoder: &CanonicalStreamDecoder,
|
||||
started_at: Instant,
|
||||
first_token_ms: &mut Option<i64>,
|
||||
first_reasoning_ms: &mut Option<i64>,
|
||||
first_text_ms: &mut Option<i64>,
|
||||
) {
|
||||
if first_token_ms.is_none() && decoder.first_output_seen() {
|
||||
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
|
||||
let first_token_missing = first_token_ms.is_none() && decoder.first_semantic_output_seen();
|
||||
let first_reasoning_missing =
|
||||
first_reasoning_ms.is_none() && decoder.first_reasoning_output_seen();
|
||||
let first_text_missing = first_text_ms.is_none() && decoder.first_text_output_seen();
|
||||
if !(first_token_missing || first_reasoning_missing || first_text_missing) {
|
||||
return;
|
||||
}
|
||||
let elapsed_ms = elapsed_millis_i64(started_at).max(1);
|
||||
if first_token_missing {
|
||||
*first_token_ms = Some(elapsed_ms);
|
||||
}
|
||||
if first_reasoning_missing {
|
||||
*first_reasoning_ms = Some(elapsed_ms);
|
||||
}
|
||||
if first_text_missing {
|
||||
*first_text_ms = Some(elapsed_ms);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -313,15 +355,18 @@ fn stream_summary(
|
||||
terminal: StreamTerminal,
|
||||
body: Vec<u8>,
|
||||
terminal_chunks: Vec<CanonicalResponseChunk>,
|
||||
first_token_ms: Option<i64>,
|
||||
output_timing: (Option<i64>, Option<i64>, Option<i64>),
|
||||
first_event_ms: i64,
|
||||
decoder: &CanonicalStreamDecoder,
|
||||
) -> StreamSummary {
|
||||
let (first_token_ms, first_reasoning_ms, first_text_ms) = output_timing;
|
||||
StreamSummary {
|
||||
terminal,
|
||||
body,
|
||||
terminal_chunks,
|
||||
first_token_ms,
|
||||
first_reasoning_ms,
|
||||
first_text_ms,
|
||||
first_event_ms,
|
||||
usage: decoder.usage(),
|
||||
last_response_id: decoder.last_response_id().map(ToString::to_string),
|
||||
|
||||
@@ -6,8 +6,8 @@ use serde_json::Value;
|
||||
use crate::upstream::openai::protocol::{
|
||||
events::{TokenUsage, extract_usage},
|
||||
responses::{
|
||||
CollectedResponse, ResponsesSseFailure, StreamCommitPolicy,
|
||||
reconvert_responses_sse_event_tuple_values,
|
||||
CollectedResponse, ResponseEventSignals, ResponsesSseFailure, StreamCommitPolicy,
|
||||
reconvert_responses_sse_event_tuple_values, response_event_signals,
|
||||
},
|
||||
sse::{
|
||||
MAX_SSE_EVENT_BUFFER_BYTES, SseError, encode_sse_event, parse_sse_events, sse_frame_end,
|
||||
@@ -123,34 +123,8 @@ impl CanonicalResponseEvent {
|
||||
.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 signals(&self) -> ResponseEventSignals {
|
||||
response_event_signals(self.event_type(), &self.data)
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn terminal(&self) -> Option<CanonicalStreamTerminal> {
|
||||
@@ -230,7 +204,7 @@ impl CanonicalStreamBatch {
|
||||
/// 同一上游 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() {
|
||||
if event.signals().semantic_output {
|
||||
return None;
|
||||
}
|
||||
if let Some(terminal) = event.terminal() {
|
||||
@@ -250,7 +224,9 @@ pub(in crate::dispatch) struct CanonicalStreamDecoder {
|
||||
tuple_schema: Option<Value>,
|
||||
pending: Vec<u8>,
|
||||
forwardable_frame_seen: bool,
|
||||
first_output_seen: bool,
|
||||
first_semantic_output_seen: bool,
|
||||
first_reasoning_output_seen: bool,
|
||||
first_text_output_seen: bool,
|
||||
terminal: Option<CanonicalStreamTerminal>,
|
||||
transport_done_seen: bool,
|
||||
completed_usage: Option<TokenUsage>,
|
||||
@@ -264,7 +240,9 @@ impl CanonicalStreamDecoder {
|
||||
tuple_schema,
|
||||
pending: Vec::new(),
|
||||
forwardable_frame_seen: false,
|
||||
first_output_seen: false,
|
||||
first_semantic_output_seen: false,
|
||||
first_reasoning_output_seen: false,
|
||||
first_text_output_seen: false,
|
||||
terminal: None,
|
||||
transport_done_seen: false,
|
||||
completed_usage: None,
|
||||
@@ -333,15 +311,23 @@ impl CanonicalStreamDecoder {
|
||||
self.transport_done_seen
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn first_output_seen(&self) -> bool {
|
||||
self.first_output_seen
|
||||
pub(in crate::dispatch) fn first_semantic_output_seen(&self) -> bool {
|
||||
self.first_semantic_output_seen
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn first_reasoning_output_seen(&self) -> bool {
|
||||
self.first_reasoning_output_seen
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn first_text_output_seen(&self) -> bool {
|
||||
self.first_text_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()
|
||||
self.first_semantic_output_seen || self.terminal.is_some()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -406,7 +392,10 @@ impl CanonicalStreamDecoder {
|
||||
}
|
||||
|
||||
fn observe(&mut self, event: &CanonicalResponseEvent) {
|
||||
self.first_output_seen |= event.has_first_output();
|
||||
let signals = event.signals();
|
||||
self.first_semantic_output_seen |= signals.semantic_output;
|
||||
self.first_reasoning_output_seen |= signals.reasoning_output;
|
||||
self.first_text_output_seen |= signals.text_output;
|
||||
if let Some(response_id) = event.response_id() {
|
||||
self.last_response_id = Some(response_id.to_string());
|
||||
}
|
||||
|
||||
@@ -2,13 +2,10 @@
|
||||
|
||||
use chrono::{Duration, Utc};
|
||||
|
||||
use crate::{
|
||||
fleet::{account::AccountStatus, pool::AccountPoolService},
|
||||
upstream::openai::{
|
||||
failure::{UpstreamFailureFacts, upstream_failure_facts},
|
||||
protocol::responses::ResponsesSseFailure,
|
||||
transport::{CodexBackendClient, CodexClientError},
|
||||
},
|
||||
use crate::fleet::{
|
||||
account::AccountStatus,
|
||||
account_gateway::{AccountFailureObservation, AccountUpstreamGateway},
|
||||
pool::AccountPoolService,
|
||||
};
|
||||
|
||||
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
|
||||
@@ -44,37 +41,24 @@ impl ClassifiedAccountFailure {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn classify_client_failure(error: &CodexClientError) -> Option<ClassifiedAccountFailure> {
|
||||
classify_upstream_failure(&upstream_failure_facts(error))
|
||||
}
|
||||
|
||||
pub fn classify_upstream_failure(facts: &UpstreamFailureFacts) -> Option<ClassifiedAccountFailure> {
|
||||
pub fn classify_account_failure(
|
||||
facts: &AccountFailureObservation,
|
||||
) -> Option<ClassifiedAccountFailure> {
|
||||
classify_failure(FailureFields {
|
||||
status_code: facts.status_code,
|
||||
code: facts.code.as_deref(),
|
||||
error_type: facts.error_type.as_deref(),
|
||||
identity_authorization_error: facts.identity_authorization_error.as_deref(),
|
||||
identity_error_code: facts.identity_error_code.as_deref(),
|
||||
message: &facts.message,
|
||||
body: &facts.body,
|
||||
retry_after_seconds: facts.retry_after_seconds,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn classify_response_failure(
|
||||
failure: &ResponsesSseFailure,
|
||||
) -> Option<ClassifiedAccountFailure> {
|
||||
classify_failure(FailureFields {
|
||||
status_code: failure.explicit_status_code,
|
||||
code: failure.upstream_code.as_deref(),
|
||||
error_type: failure.upstream_type.as_deref(),
|
||||
message: &failure.message,
|
||||
body: &failure.message,
|
||||
retry_after_seconds: failure.retry_after_seconds,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn apply_account_state_effect_immediately(
|
||||
account_pool: &AccountPoolService,
|
||||
codex: &CodexBackendClient,
|
||||
upstream: &dyn AccountUpstreamGateway,
|
||||
account_id: &str,
|
||||
effect: &AccountStateEffect,
|
||||
) {
|
||||
@@ -90,12 +74,12 @@ pub async fn apply_account_state_effect_immediately(
|
||||
.await;
|
||||
}
|
||||
}
|
||||
codex.evict_websocket_account(account_id).await;
|
||||
upstream.evict_account_connections(account_id).await;
|
||||
}
|
||||
|
||||
pub async fn apply_account_state_effect(
|
||||
account_pool: &AccountPoolService,
|
||||
codex: &CodexBackendClient,
|
||||
upstream: &dyn AccountUpstreamGateway,
|
||||
account_id: &str,
|
||||
effect: &AccountStateEffect,
|
||||
) -> bool {
|
||||
@@ -107,27 +91,16 @@ pub async fn apply_account_state_effect(
|
||||
.await
|
||||
}
|
||||
};
|
||||
codex.evict_websocket_account(account_id).await;
|
||||
upstream.evict_account_connections(account_id).await;
|
||||
updated
|
||||
}
|
||||
|
||||
pub async fn apply_client_failure_effect(
|
||||
account_pool: &AccountPoolService,
|
||||
codex: &CodexBackendClient,
|
||||
account_id: &str,
|
||||
error: &CodexClientError,
|
||||
) -> Option<ClassifiedAccountFailure> {
|
||||
let classified = classify_client_failure(error)?;
|
||||
if let Some(effect) = &classified.effect {
|
||||
apply_account_state_effect(account_pool, codex, account_id, effect).await;
|
||||
}
|
||||
Some(classified)
|
||||
}
|
||||
|
||||
struct FailureFields<'a> {
|
||||
status_code: Option<u16>,
|
||||
code: Option<&'a str>,
|
||||
error_type: Option<&'a str>,
|
||||
identity_authorization_error: Option<&'a str>,
|
||||
identity_error_code: Option<&'a str>,
|
||||
message: &'a str,
|
||||
body: &'a str,
|
||||
retry_after_seconds: Option<u64>,
|
||||
@@ -136,6 +109,8 @@ struct FailureFields<'a> {
|
||||
fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailure> {
|
||||
let code = fields.code.unwrap_or_default();
|
||||
let error_type = fields.error_type.unwrap_or_default();
|
||||
let identity_error_code = fields.identity_error_code.unwrap_or_default();
|
||||
let identity_authorization_error = fields.identity_authorization_error.unwrap_or_default();
|
||||
|
||||
if [code, fields.message, fields.body]
|
||||
.into_iter()
|
||||
@@ -143,14 +118,13 @@ fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailur
|
||||
{
|
||||
return Some(classified(AccountFailureKind::ModelUnsupported, None));
|
||||
}
|
||||
if let Some(status) = explicit_account_status(code, error_type, fields.message)
|
||||
.or_else(|| explicit_account_status(code, error_type, fields.body))
|
||||
if let Some(status) =
|
||||
explicit_account_status(identity_error_code, "", identity_authorization_error)
|
||||
.or_else(|| explicit_account_status(code, error_type, fields.message))
|
||||
.or_else(|| explicit_account_status(code, error_type, fields.body))
|
||||
{
|
||||
return Some(status_failure(status));
|
||||
}
|
||||
if fields.status_code == Some(403) && !is_html(fields.body) {
|
||||
return Some(status_failure(AccountStatus::Banned));
|
||||
}
|
||||
if fields.status_code == Some(401) {
|
||||
return Some(status_failure(AccountStatus::Expired));
|
||||
}
|
||||
@@ -238,8 +212,6 @@ fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Optio
|
||||
|| message.contains("organization has been disabled")
|
||||
|| message.contains("workspace has been deactivated")
|
||||
|| message.contains("deactivated_workspace")
|
||||
|| code == "forbidden"
|
||||
|| error_type == "permission_error"
|
||||
{
|
||||
return Some(AccountStatus::Banned);
|
||||
}
|
||||
@@ -310,8 +282,3 @@ fn is_model_unsupported(value: &str) -> bool {
|
||||
|| 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,98 @@
|
||||
//! 账号管理与配额刷新用例拥有的 OpenAI IO 端口。
|
||||
|
||||
use std::pin::Pin;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::Stream;
|
||||
use secrecy::SecretString;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::fleet::quota::QuotaSnapshot;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountUpstreamContext {
|
||||
pub access_token: SecretString,
|
||||
pub account_id: Option<String>,
|
||||
pub request_id: String,
|
||||
pub cookie_header: Option<String>,
|
||||
pub installation_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AccountFailureObservation {
|
||||
pub status_code: Option<u16>,
|
||||
pub code: Option<String>,
|
||||
pub error_type: Option<String>,
|
||||
pub identity_authorization_error: Option<String>,
|
||||
pub identity_error_code: Option<String>,
|
||||
pub message: String,
|
||||
pub body: String,
|
||||
pub retry_after_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
#[error("{message}")]
|
||||
pub struct AccountGatewayError {
|
||||
message: String,
|
||||
failure: Option<AccountFailureObservation>,
|
||||
}
|
||||
|
||||
impl AccountGatewayError {
|
||||
pub fn new(message: impl Into<String>, failure: Option<AccountFailureObservation>) -> Self {
|
||||
Self {
|
||||
message: message.into(),
|
||||
failure,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn failure(&self) -> Option<&AccountFailureObservation> {
|
||||
self.failure.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountUsageResult {
|
||||
pub quota: QuotaSnapshot,
|
||||
pub raw: Value,
|
||||
pub account_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub email: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountProbeRequest {
|
||||
pub model: String,
|
||||
pub instructions: String,
|
||||
pub input_text: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AccountProbeEvent {
|
||||
Content(String),
|
||||
Complete,
|
||||
Failed(AccountGatewayError),
|
||||
}
|
||||
|
||||
pub type AccountProbeEventStream = Pin<Box<dyn Stream<Item = AccountProbeEvent> + Send>>;
|
||||
|
||||
pub struct AccountProbeSession {
|
||||
pub request_payload: Value,
|
||||
pub events: AccountProbeEventStream,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AccountUpstreamGateway: Send + Sync + 'static {
|
||||
async fn fetch_usage(
|
||||
&self,
|
||||
context: AccountUpstreamContext,
|
||||
) -> Result<AccountUsageResult, AccountGatewayError>;
|
||||
|
||||
async fn probe_response(
|
||||
&self,
|
||||
context: AccountUpstreamContext,
|
||||
request: AccountProbeRequest,
|
||||
) -> Result<AccountProbeSession, AccountGatewayError>;
|
||||
|
||||
async fn evict_account_connections(&self, account_id: &str);
|
||||
}
|
||||
@@ -1,12 +1,14 @@
|
||||
use chrono::Utc;
|
||||
use secrecy::SecretString;
|
||||
|
||||
use crate::fleet::account_gateway::AccountUpstreamContext;
|
||||
|
||||
use super::{
|
||||
AccountManageService,
|
||||
types::{
|
||||
AccountManageError, ImportSupplementalAccountInfo, ImportSupplementalNeeds,
|
||||
ImportedAccountState, ImportedAccounts, ResolvedImportTokens, import_quota_plan_type,
|
||||
import_status_from_usage_error, import_usage_plan_type, import_usage_string,
|
||||
import_status_from_usage_error,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -190,41 +192,26 @@ impl AccountManageService {
|
||||
return ImportSupplementalAccountInfo::default();
|
||||
}
|
||||
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let installation_id =
|
||||
account_id.map(|account_id| self.account_pseudonymizer.installation_id(account_id));
|
||||
let context = crate::upstream::openai::transport::CodexRequestContext {
|
||||
access_token,
|
||||
account_id,
|
||||
request_id: &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,
|
||||
let context = AccountUpstreamContext {
|
||||
access_token: SecretString::new(access_token.to_string().into()),
|
||||
account_id: account_id.map(ToString::to_string),
|
||||
request_id: uuid::Uuid::new_v4().to_string(),
|
||||
cookie_header: None,
|
||||
installation_id: installation_id.as_deref(),
|
||||
session_id: None,
|
||||
thread_id: None,
|
||||
client_request_id: None,
|
||||
turn_id: None,
|
||||
installation_id,
|
||||
};
|
||||
|
||||
match self.codex.fetch_usage(context).await {
|
||||
Ok(raw) => {
|
||||
let normalized = crate::fleet::quota::quota_from_usage(&raw);
|
||||
ImportSupplementalAccountInfo {
|
||||
account_id: import_usage_string(&raw, "account_id"),
|
||||
user_id: import_usage_string(&raw, "user_id"),
|
||||
email: import_usage_string(&raw, "email"),
|
||||
plan_type: import_usage_plan_type(&raw),
|
||||
quota_json: serde_json::to_string(&normalized).ok(),
|
||||
quota_fetched_at: Some(Utc::now()),
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
match self.upstream.fetch_usage(context).await {
|
||||
Ok(result) => ImportSupplementalAccountInfo {
|
||||
account_id: result.account_id,
|
||||
user_id: result.user_id,
|
||||
email: result.email,
|
||||
plan_type: result.plan_type,
|
||||
quota_json: result.quota.to_json().ok(),
|
||||
quota_fetched_at: Some(Utc::now()),
|
||||
status: None,
|
||||
},
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
|
||||
@@ -10,12 +10,12 @@ use tokio::time::sleep;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{
|
||||
fleet::refresh::RefreshFailure,
|
||||
fleet::{
|
||||
account::{AccountListSort, AccountStatus},
|
||||
store::AccountStore,
|
||||
},
|
||||
infra::{json::NumberedPage, time::elapsed_millis_i64},
|
||||
upstream::openai::token_client::RefreshFailure,
|
||||
};
|
||||
|
||||
use super::{
|
||||
|
||||
@@ -6,9 +6,9 @@ mod lifecycle;
|
||||
pub(crate) mod oauth;
|
||||
mod probe;
|
||||
mod quota;
|
||||
pub mod quota_view;
|
||||
mod service;
|
||||
mod types;
|
||||
|
||||
pub use probe::{AccountModelOption, AccountTestEvent, AccountTestStream};
|
||||
pub use service::*;
|
||||
pub use types::*;
|
||||
|
||||
+148
-223
@@ -1,66 +1,86 @@
|
||||
use std::{convert::Infallible, pin::Pin};
|
||||
use std::pin::Pin;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::{Stream, StreamExt};
|
||||
use secrecy::ExposeSecret;
|
||||
use serde_json::{Value, json};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::{
|
||||
fleet::{
|
||||
account::AccountStatus,
|
||||
account_failure::{classify_client_failure, classify_response_failure},
|
||||
account_failure::classify_account_failure,
|
||||
account_gateway::{
|
||||
AccountGatewayError, AccountProbeEvent, AccountProbeRequest, AccountProbeSession,
|
||||
AccountUpstreamContext,
|
||||
},
|
||||
store::StoredAccount,
|
||||
},
|
||||
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
|
||||
upstream::openai::{
|
||||
protocol::{
|
||||
responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
sse::{SseEvent, encode_sse_event, parse_sse_events},
|
||||
},
|
||||
transport::{CodexClientError, CodexRequestContext},
|
||||
models::{
|
||||
service::{ModelRefreshPlanAccount, ModelRefreshResult},
|
||||
types::CodexModelInfo,
|
||||
},
|
||||
};
|
||||
|
||||
use super::{AccountManageService, types::AccountManageError};
|
||||
|
||||
pub(super) type AccountTestStream = Pin<Box<dyn Stream<Item = Result<Bytes, Infallible>> + Send>>;
|
||||
pub type AccountTestStream = Pin<Box<dyn Stream<Item = AccountTestEvent> + Send>>;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AccountTestEvent {
|
||||
Started {
|
||||
model: String,
|
||||
},
|
||||
Request {
|
||||
payload: Value,
|
||||
},
|
||||
Content {
|
||||
text: String,
|
||||
},
|
||||
Complete {
|
||||
account_status: Option<AccountStatus>,
|
||||
},
|
||||
Error {
|
||||
error: String,
|
||||
account_status: Option<AccountStatus>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountModelOption {
|
||||
pub id: String,
|
||||
pub label: String,
|
||||
}
|
||||
|
||||
struct AccountTestOutcome {
|
||||
event: Value,
|
||||
error: Option<String>,
|
||||
status: Option<AccountStatus>,
|
||||
}
|
||||
|
||||
impl AccountTestOutcome {
|
||||
fn success() -> Self {
|
||||
Self {
|
||||
event: json!({ "type": "test_complete", "success": true }),
|
||||
error: None,
|
||||
status: Some(AccountStatus::Active),
|
||||
}
|
||||
}
|
||||
|
||||
fn error(error: impl Into<String>, status: Option<AccountStatus>) -> Self {
|
||||
Self {
|
||||
event: json!({ "type": "error", "error": error.into() }),
|
||||
error: Some(error.into()),
|
||||
status,
|
||||
}
|
||||
}
|
||||
|
||||
fn attach_account_status(&mut self, status: AccountStatus) {
|
||||
let Some(event) = self.event.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
event.insert(
|
||||
"accountStatus".to_string(),
|
||||
Value::String(status.as_str().to_string()),
|
||||
);
|
||||
fn into_event(self) -> AccountTestEvent {
|
||||
match self.error {
|
||||
Some(error) => AccountTestEvent::Error {
|
||||
error,
|
||||
account_status: self.status,
|
||||
},
|
||||
None => AccountTestEvent::Complete {
|
||||
account_status: self.status,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,11 +96,7 @@ impl AccountManageService {
|
||||
.map_err(|_| AccountManageError::Inspect)?
|
||||
.ok_or(AccountManageError::NotFound)?;
|
||||
let plan_type = account_plan_type(&account);
|
||||
let mut models = self.models.catalog().await.models_for_plan(&plan_type);
|
||||
if models.is_empty() {
|
||||
self.refresh_account_plan_models(&account, &plan_type).await;
|
||||
models = self.models.catalog().await.models_for_plan(&plan_type);
|
||||
}
|
||||
let models = self.models.catalog().await.models_for_plan(&plan_type);
|
||||
let models = models.iter().map(account_model_option).collect::<Vec<_>>();
|
||||
if models.is_empty() {
|
||||
return Err(AccountManageError::NoModels);
|
||||
@@ -88,7 +104,32 @@ impl AccountManageService {
|
||||
Ok(models)
|
||||
}
|
||||
|
||||
async fn refresh_account_plan_models(&self, account: &StoredAccount, plan_type: &str) {
|
||||
pub async fn refresh_account_models(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<Vec<AccountModelOption>, AccountManageError> {
|
||||
let account = self
|
||||
.store
|
||||
.get(account_id)
|
||||
.await
|
||||
.map_err(|_| AccountManageError::Inspect)?
|
||||
.ok_or(AccountManageError::NotFound)?;
|
||||
let plan_type = account_plan_type(&account);
|
||||
self.refresh_account_plan_models(&account, &plan_type)
|
||||
.await?;
|
||||
let models = self.models.catalog().await.models_for_plan(&plan_type);
|
||||
let models = models.iter().map(account_model_option).collect::<Vec<_>>();
|
||||
if models.is_empty() {
|
||||
return Err(AccountManageError::NoModels);
|
||||
}
|
||||
Ok(models)
|
||||
}
|
||||
|
||||
async fn refresh_account_plan_models(
|
||||
&self,
|
||||
account: &StoredAccount,
|
||||
plan_type: &str,
|
||||
) -> Result<ModelRefreshResult, AccountManageError> {
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let plan_account = ModelRefreshPlanAccount {
|
||||
plan_type: plan_type.to_string(),
|
||||
@@ -96,18 +137,20 @@ impl AccountManageService {
|
||||
account_id: account.account_id.clone(),
|
||||
installation_id: self.account_pseudonymizer.installation_id(&account.id),
|
||||
};
|
||||
if let Err(error) = self
|
||||
let result = self
|
||||
.models
|
||||
.refresh_backend_models(&[plan_account], &request_id)
|
||||
.refresh_selected_plan_models(&[plan_account], &request_id)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id = %account.id,
|
||||
plan_type,
|
||||
error = %error,
|
||||
"Failed to refresh account plan models"
|
||||
);
|
||||
}
|
||||
.map_err(|error| AccountManageError::RefreshModels(error.to_string()))?;
|
||||
let routing = self
|
||||
.models
|
||||
.model_plan_routing()
|
||||
.await
|
||||
.map_err(|error| AccountManageError::RefreshModels(error.to_string()))?;
|
||||
self.account_pool
|
||||
.apply_model_plan_routing(routing.allowlist, routing.fetched_plan_types)
|
||||
.await;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub async fn test_connection_stream(
|
||||
@@ -121,84 +164,55 @@ impl AccountManageService {
|
||||
.await
|
||||
.map_err(|_| AccountManageError::Inspect)?
|
||||
.ok_or(AccountManageError::NotFound)?;
|
||||
|
||||
let token = account.access_token.expose_secret().to_string();
|
||||
let upstream_account_id = account.account_id.clone();
|
||||
let cookie_header = self
|
||||
.cookies
|
||||
.cookie_header_for_request(&account.id, "chatgpt.com", "/codex/responses")
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
let installation_id = self.account_pseudonymizer.installation_id(&account.id);
|
||||
let codex = self.codex.clone();
|
||||
let context = AccountUpstreamContext {
|
||||
access_token: account.access_token.clone(),
|
||||
account_id: account.account_id.clone(),
|
||||
request_id: uuid::Uuid::new_v4().to_string(),
|
||||
cookie_header: self
|
||||
.cookies
|
||||
.cookie_header_for_request(&account.id, "chatgpt.com", "/codex/responses")
|
||||
.await
|
||||
.ok()
|
||||
.flatten(),
|
||||
installation_id: Some(self.account_pseudonymizer.installation_id(&account.id)),
|
||||
};
|
||||
let upstream = self.upstream.clone();
|
||||
let service = self.clone();
|
||||
let stored_account_id = account.id.clone();
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let (tx, rx) = mpsc::channel::<Bytes>(16);
|
||||
let stored_account_id = account.id;
|
||||
let (tx, rx) = mpsc::channel(16);
|
||||
|
||||
tokio::spawn(async move {
|
||||
send_test_event(
|
||||
&tx,
|
||||
json!({
|
||||
"type": "test_start",
|
||||
"model": model.clone(),
|
||||
"text": "正在连接 Codex Responses"
|
||||
}),
|
||||
AccountTestEvent::Started {
|
||||
model: model.clone(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let request = test_responses_request(model);
|
||||
send_test_event(
|
||||
&tx,
|
||||
json!({
|
||||
"type": "request",
|
||||
"payload": serde_json::to_value(&request).unwrap_or_else(|_| json!({}))
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let context = CodexRequestContext {
|
||||
access_token: &token,
|
||||
account_id: upstream_account_id.as_deref(),
|
||||
request_id: &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(&installation_id),
|
||||
session_id: None,
|
||||
thread_id: None,
|
||||
client_request_id: None,
|
||||
turn_id: None,
|
||||
let request = AccountProbeRequest {
|
||||
model,
|
||||
instructions:
|
||||
"You are checking whether this Codex account can answer. Reply with ok."
|
||||
.to_string(),
|
||||
input_text: "hi".to_string(),
|
||||
};
|
||||
|
||||
let mut outcome = match codex.create_response_stream(&request, context).await {
|
||||
Ok(response) => process_upstream_test_stream(response.body, &tx).await,
|
||||
let mut outcome = match upstream.probe_response(context, request).await {
|
||||
Ok(session) => process_probe_session(session, &tx).await,
|
||||
Err(error) => AccountTestOutcome::error(
|
||||
error.to_string(),
|
||||
account_status_from_client_error(&error),
|
||||
account_status_from_gateway_error(&error),
|
||||
),
|
||||
};
|
||||
|
||||
if let Some(status) = outcome.status
|
||||
&& let Some(status) = service
|
||||
if let Some(status) = outcome.status {
|
||||
outcome.status = service
|
||||
.apply_connection_test_status(&stored_account_id, status)
|
||||
.await
|
||||
{
|
||||
outcome.attach_account_status(status);
|
||||
.await;
|
||||
}
|
||||
send_test_event(&tx, outcome.event).await;
|
||||
send_test_event(&tx, outcome.into_event()).await;
|
||||
});
|
||||
|
||||
let stream = futures::stream::unfold(rx, |mut rx| async {
|
||||
rx.recv()
|
||||
.await
|
||||
.map(|bytes| (Ok::<Bytes, Infallible>(bytes), rx))
|
||||
rx.recv().await.map(|event| (event, rx))
|
||||
});
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
@@ -223,7 +237,6 @@ impl AccountManageService {
|
||||
if current.status == AccountStatus::Disabled {
|
||||
return Some(AccountStatus::Disabled);
|
||||
}
|
||||
|
||||
match self.store.set_status(account_id, status).await {
|
||||
Ok(true) => {}
|
||||
Ok(false) => return None,
|
||||
@@ -237,7 +250,6 @@ impl AccountManageService {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
if matches!(status, AccountStatus::Expired | AccountStatus::Banned)
|
||||
&& let Err(error) = self.store.set_next_refresh_at(account_id, None).await
|
||||
{
|
||||
@@ -253,6 +265,34 @@ impl AccountManageService {
|
||||
}
|
||||
}
|
||||
|
||||
async fn process_probe_session(
|
||||
mut session: AccountProbeSession,
|
||||
tx: &mpsc::Sender<AccountTestEvent>,
|
||||
) -> AccountTestOutcome {
|
||||
send_test_event(
|
||||
tx,
|
||||
AccountTestEvent::Request {
|
||||
payload: session.request_payload,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
while let Some(event) = session.events.next().await {
|
||||
match event {
|
||||
AccountProbeEvent::Content(text) => {
|
||||
send_test_event(tx, AccountTestEvent::Content { text }).await;
|
||||
}
|
||||
AccountProbeEvent::Complete => return AccountTestOutcome::success(),
|
||||
AccountProbeEvent::Failed(error) => {
|
||||
return AccountTestOutcome::error(
|
||||
error.to_string(),
|
||||
account_status_from_gateway_error(&error),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
AccountTestOutcome::error("Stream ended before response.completed", None)
|
||||
}
|
||||
|
||||
fn account_model_option(model: &CodexModelInfo) -> AccountModelOption {
|
||||
AccountModelOption {
|
||||
id: model.id.clone(),
|
||||
@@ -270,125 +310,10 @@ fn account_plan_type(account: &StoredAccount) -> String {
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn test_responses_request(model: String) -> CodexResponsesRequest {
|
||||
let mut request = CodexResponsesRequest::new_http_sse(
|
||||
model,
|
||||
"You are checking whether this Codex account can answer. Reply with ok.",
|
||||
vec![json!({
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_text",
|
||||
"text": "hi"
|
||||
}
|
||||
]
|
||||
})],
|
||||
);
|
||||
request.set_stream(true);
|
||||
request.set_store(false);
|
||||
request.force_http_sse = true;
|
||||
request
|
||||
fn account_status_from_gateway_error(error: &AccountGatewayError) -> Option<AccountStatus> {
|
||||
classify_account_failure(error.failure()?)?.account_status()
|
||||
}
|
||||
|
||||
async fn process_upstream_test_stream(
|
||||
mut body: crate::upstream::openai::transport::CodexBackendSseStream,
|
||||
tx: &mpsc::Sender<Bytes>,
|
||||
) -> AccountTestOutcome {
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(chunk) = body.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(bytes) => bytes,
|
||||
Err(error) => {
|
||||
return AccountTestOutcome::error(
|
||||
error.to_string(),
|
||||
account_status_from_client_error(&error),
|
||||
);
|
||||
}
|
||||
};
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
while let Some(frame) = take_sse_frame(&mut buffer) {
|
||||
if let Some(outcome) = process_sse_frame(tx, &frame).await {
|
||||
return outcome;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !buffer.trim().is_empty()
|
||||
&& let Some(outcome) = process_sse_frame(tx, &buffer).await
|
||||
{
|
||||
return outcome;
|
||||
}
|
||||
|
||||
AccountTestOutcome::error("Stream ended before response.completed", None)
|
||||
}
|
||||
|
||||
async fn process_sse_frame(tx: &mpsc::Sender<Bytes>, frame: &str) -> Option<AccountTestOutcome> {
|
||||
let events = match parse_sse_events(frame) {
|
||||
Ok(events) => events,
|
||||
Err(error) => {
|
||||
return Some(AccountTestOutcome::error(error.to_string(), None));
|
||||
}
|
||||
};
|
||||
|
||||
for event in events {
|
||||
if let Some(outcome) = process_sse_event(tx, &event).await {
|
||||
return Some(outcome);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
async fn process_sse_event(
|
||||
tx: &mpsc::Sender<Bytes>,
|
||||
event: &SseEvent,
|
||||
) -> Option<AccountTestOutcome> {
|
||||
let value: Value = match serde_json::from_str(&event.data) {
|
||||
Ok(value) => value,
|
||||
Err(_) => return None,
|
||||
};
|
||||
match value.get("type").and_then(Value::as_str) {
|
||||
Some("response.output_text.delta") => {
|
||||
if let Some(delta) = value.get("delta").and_then(Value::as_str)
|
||||
&& !delta.is_empty()
|
||||
{
|
||||
send_test_event(tx, json!({ "type": "content", "text": delta })).await;
|
||||
}
|
||||
None
|
||||
}
|
||||
Some("response.completed" | "response.done") => Some(AccountTestOutcome::success()),
|
||||
Some(event_name @ ("response.failed" | "error")) => {
|
||||
let failure = ResponsesSseFailure::from_event(event_name, &value);
|
||||
let status = sse_failure_account_status(&failure);
|
||||
Some(AccountTestOutcome::error(failure.message, status))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn account_status_from_client_error(error: &CodexClientError) -> Option<AccountStatus> {
|
||||
classify_client_failure(error)?.account_status()
|
||||
}
|
||||
|
||||
fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
|
||||
classify_response_failure(failure)?.account_status()
|
||||
}
|
||||
|
||||
fn take_sse_frame(buffer: &mut String) -> Option<String> {
|
||||
let index = buffer.find("\n\n").or_else(|| buffer.find("\r\n\r\n"))?;
|
||||
let delimiter_len = if buffer[index..].starts_with("\r\n\r\n") {
|
||||
4
|
||||
} else {
|
||||
2
|
||||
};
|
||||
let frame = buffer[..index + delimiter_len].to_string();
|
||||
buffer.drain(..index + delimiter_len);
|
||||
Some(frame)
|
||||
}
|
||||
|
||||
async fn send_test_event(tx: &mpsc::Sender<Bytes>, event: Value) {
|
||||
let _ = tx
|
||||
.send(Bytes::from(encode_sse_event("", &event.to_string())))
|
||||
.await;
|
||||
async fn send_test_event(tx: &mpsc::Sender<AccountTestEvent>, event: AccountTestEvent) {
|
||||
let _ = tx.send(event).await;
|
||||
}
|
||||
|
||||
@@ -1,7 +1,16 @@
|
||||
use secrecy::ExposeSecret;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{AccountManageService, types::AccountManageError};
|
||||
use crate::fleet::account_failure::apply_client_failure_effect;
|
||||
use crate::fleet::{
|
||||
account_failure::{apply_account_state_effect, classify_account_failure},
|
||||
account_gateway::AccountUpstreamContext,
|
||||
quota::QuotaSnapshot,
|
||||
};
|
||||
|
||||
pub struct AccountQuotaRefresh {
|
||||
pub quota: QuotaSnapshot,
|
||||
pub raw: Value,
|
||||
}
|
||||
|
||||
impl AccountManageService {
|
||||
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
|
||||
@@ -24,7 +33,7 @@ impl AccountManageService {
|
||||
pub async fn account_quota(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<serde_json::Value, AccountManageError> {
|
||||
) -> Result<AccountQuotaRefresh, AccountManageError> {
|
||||
let stored = self
|
||||
.store
|
||||
.get(account_id)
|
||||
@@ -32,42 +41,38 @@ impl AccountManageService {
|
||||
.map_err(|_| AccountManageError::NotFound)?
|
||||
.ok_or(AccountManageError::NotFound)?;
|
||||
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let token = stored.access_token.expose_secret().to_string();
|
||||
let cookie_header = self.usage_cookie_header(account_id).await;
|
||||
let installation_id = self.account_pseudonymizer.installation_id(account_id);
|
||||
let context = crate::upstream::openai::transport::CodexRequestContext {
|
||||
access_token: &token,
|
||||
account_id: stored.account_id.as_deref(),
|
||||
request_id: &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(&installation_id),
|
||||
session_id: None,
|
||||
thread_id: None,
|
||||
client_request_id: None,
|
||||
turn_id: None,
|
||||
let context = AccountUpstreamContext {
|
||||
access_token: stored.access_token.clone(),
|
||||
account_id: stored.account_id.clone(),
|
||||
request_id: uuid::Uuid::new_v4().to_string(),
|
||||
cookie_header: self.usage_cookie_header(account_id).await,
|
||||
installation_id: Some(self.account_pseudonymizer.installation_id(account_id)),
|
||||
};
|
||||
let raw = match self.codex.fetch_usage(context).await {
|
||||
Ok(raw) => raw,
|
||||
let result = match self.upstream.fetch_usage(context).await {
|
||||
Ok(result) => result,
|
||||
Err(error) => {
|
||||
apply_client_failure_effect(&self.account_pool, &self.codex, account_id, &error)
|
||||
if let Some(classified) = error.failure().and_then(classify_account_failure)
|
||||
&& let Some(effect) = &classified.effect
|
||||
{
|
||||
apply_account_state_effect(
|
||||
&self.account_pool,
|
||||
self.upstream.as_ref(),
|
||||
account_id,
|
||||
effect,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
return Err(AccountManageError::FetchQuota(error.to_string()));
|
||||
}
|
||||
};
|
||||
let normalized = crate::fleet::quota::quota_from_usage(&raw);
|
||||
self.account_pool
|
||||
.apply_quota_snapshot(account_id, &normalized)
|
||||
.apply_quota_snapshot(account_id, &result.quota)
|
||||
.await;
|
||||
self.sync_account_pool_best_effort(account_id, "account quota refresh")
|
||||
.await;
|
||||
Ok(serde_json::json!({ "quota": normalized, "raw": raw }))
|
||||
Ok(AccountQuotaRefresh {
|
||||
quota: result.quota,
|
||||
raw: result.raw,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,471 +0,0 @@
|
||||
//! 管理端账号 quota 展示视图。
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::infra::{
|
||||
format::{format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64},
|
||||
time::{china_datetime, china_relative_time},
|
||||
};
|
||||
|
||||
const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000;
|
||||
const WEEK_WINDOW_SECONDS: u64 = 604_800;
|
||||
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct AccountQuotaData {
|
||||
refreshed_at_display: String,
|
||||
windows: Vec<AccountQuotaWindowData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaWindowData {
|
||||
key: String,
|
||||
group: String,
|
||||
window_seconds: Option<u64>,
|
||||
label_display: String,
|
||||
used_percent: Option<f64>,
|
||||
used_percent_display: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
local_usage: Option<AccountQuotaWindowLocalUsageData>,
|
||||
reset_at_display: String,
|
||||
window_used_display: String,
|
||||
#[serde(skip)]
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaWindowLocalUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct AccountQuotaUsageWindow {
|
||||
pub(crate) key: String,
|
||||
pub(crate) start: DateTime<Utc>,
|
||||
pub(crate) end: DateTime<Utc>,
|
||||
pub(crate) window_seconds: u64,
|
||||
}
|
||||
|
||||
impl AccountQuotaUsageWindow {
|
||||
pub(crate) fn duration_seconds(&self) -> u64 {
|
||||
self.end
|
||||
.signed_duration_since(self.start)
|
||||
.num_seconds()
|
||||
.max(0) as u64
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub(crate) struct AccountQuotaWindowLocalUsage {
|
||||
pub(crate) request_count: i64,
|
||||
pub(crate) input_tokens: i64,
|
||||
pub(crate) output_tokens: i64,
|
||||
pub(crate) cached_tokens: i64,
|
||||
}
|
||||
|
||||
impl From<AccountQuotaWindowLocalUsage> for AccountQuotaWindowLocalUsageData {
|
||||
fn from(usage: AccountQuotaWindowLocalUsage) -> Self {
|
||||
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
|
||||
Self {
|
||||
request_count: usage.request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
|
||||
input_tokens: usage.input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
|
||||
output_tokens: usage.output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
|
||||
cached_tokens: usage.cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AccountQuotaData {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
refreshed_at_display: "-".to_string(),
|
||||
windows: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AccountQuotaData {
|
||||
pub(crate) fn usage_windows(&self) -> Vec<AccountQuotaUsageWindow> {
|
||||
self.windows
|
||||
.iter()
|
||||
.filter_map(AccountQuotaWindowData::usage_window)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn apply_local_usage(
|
||||
&mut self,
|
||||
usage_by_window: &HashMap<String, AccountQuotaWindowLocalUsage>,
|
||||
) {
|
||||
for window in &mut self.windows {
|
||||
window.local_usage = window.usage_window().map(|_| {
|
||||
usage_by_window
|
||||
.get(&window.key)
|
||||
.copied()
|
||||
.unwrap_or_default()
|
||||
.into()
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AccountQuotaWindowData {
|
||||
fn usage_window(&self) -> Option<AccountQuotaUsageWindow> {
|
||||
let end = self.reset_at?;
|
||||
let window_seconds = self.window_seconds?;
|
||||
let seconds = i64::try_from(window_seconds).ok()?;
|
||||
let start = end.checked_sub_signed(Duration::seconds(seconds))?;
|
||||
(start <= end).then(|| AccountQuotaUsageWindow {
|
||||
key: self.key.clone(),
|
||||
start,
|
||||
end,
|
||||
window_seconds,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn quota_data(quota_json: &str, fetched_at: Option<DateTime<Utc>>) -> AccountQuotaData {
|
||||
let quota = serde_json::from_str::<Value>(quota_json).unwrap_or(Value::Null);
|
||||
let windows = quota_windows("a);
|
||||
|
||||
AccountQuotaData {
|
||||
refreshed_at_display: china_relative_time(fetched_at, Utc::now()),
|
||||
windows,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_windows(quota: &Value) -> Vec<AccountQuotaWindowData> {
|
||||
let mut windows = Vec::new();
|
||||
|
||||
let has_monthly_limit = push_monthly_quota_window(&mut windows, quota.get("monthly_limit"));
|
||||
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
|
||||
for snapshot in snapshots {
|
||||
push_snapshot_quota_windows(&mut windows, snapshot, has_monthly_limit);
|
||||
}
|
||||
}
|
||||
|
||||
windows.sort_by_key(quota_window_sort_key);
|
||||
windows
|
||||
}
|
||||
|
||||
fn push_monthly_quota_window(
|
||||
windows: &mut Vec<AccountQuotaWindowData>,
|
||||
monthly_limit: Option<&Value>,
|
||||
) -> bool {
|
||||
let Some(monthly_limit) = monthly_limit.filter(|value| !value.is_null()) else {
|
||||
return false;
|
||||
};
|
||||
let used_percent = monthly_limit
|
||||
.get("used_percent")
|
||||
.and_then(number_value)
|
||||
.map(|value| value.clamp(0.0, 100.0));
|
||||
let reset_at = monthly_limit
|
||||
.get("reset_at")
|
||||
.and_then(Value::as_i64)
|
||||
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
|
||||
let window_seconds = monthly_limit
|
||||
.get("window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
.or(Some(MONTH_WINDOW_SECONDS));
|
||||
if used_percent.is_none() && reset_at.is_none() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let key = monthly_limit
|
||||
.get("key")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("monthly");
|
||||
windows.push(AccountQuotaWindowData {
|
||||
key: quota_key_segment(key),
|
||||
group: "monthly".to_string(),
|
||||
window_seconds,
|
||||
label_display: "月限额".to_string(),
|
||||
used_percent,
|
||||
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
local_usage: None,
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
window_used_display: quota_window_used_display(reset_at, window_seconds),
|
||||
reset_at,
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
fn push_snapshot_quota_windows(
|
||||
windows: &mut Vec<AccountQuotaWindowData>,
|
||||
snapshot: &Value,
|
||||
skip_core_monthly: bool,
|
||||
) {
|
||||
let source = snapshot
|
||||
.get("source")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("quota");
|
||||
let source_key = snapshot_source_key(source, snapshot);
|
||||
let label_prefix = snapshot_label(snapshot);
|
||||
for role in ["primary", "secondary"] {
|
||||
let Some(window) = snapshot.get(role).filter(|value| !value.is_null()) else {
|
||||
continue;
|
||||
};
|
||||
let window_seconds = window
|
||||
.get("window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|minutes| minutes.checked_mul(60));
|
||||
if skip_core_monthly
|
||||
&& source == "core"
|
||||
&& window_seconds
|
||||
.is_some_and(|seconds| quota_window_matches(seconds, MONTH_WINDOW_SECONDS))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
push_quota_window(
|
||||
windows,
|
||||
&source_key,
|
||||
role,
|
||||
label_prefix.as_deref(),
|
||||
Some(window),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn push_quota_window(
|
||||
windows: &mut Vec<AccountQuotaWindowData>,
|
||||
source_key: &str,
|
||||
role: &str,
|
||||
label_prefix: Option<&str>,
|
||||
window: Option<&Value>,
|
||||
) {
|
||||
let Some(window) = window.filter(|value| !value.is_null()) else {
|
||||
return;
|
||||
};
|
||||
let used_percent = window
|
||||
.get("used_percent")
|
||||
.and_then(number_value)
|
||||
.map(|value| value.clamp(0.0, 100.0));
|
||||
let reset_at = window
|
||||
.get("reset_at")
|
||||
.and_then(Value::as_i64)
|
||||
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
|
||||
let window_seconds = window
|
||||
.get("window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|minutes| minutes.checked_mul(60));
|
||||
let label_prefix = label_prefix
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
if reset_at.is_none()
|
||||
&& window_seconds.is_none()
|
||||
&& (used_percent.is_none() || label_prefix.is_none())
|
||||
{
|
||||
return;
|
||||
}
|
||||
let base_label = quota_window_label_display(window_seconds);
|
||||
let label_display = label_prefix
|
||||
.map(|value| format!("{value} · {base_label}"))
|
||||
.unwrap_or(base_label);
|
||||
|
||||
windows.push(AccountQuotaWindowData {
|
||||
key: unique_quota_window_key(windows, source_key, role, window_seconds),
|
||||
group: quota_window_group(window_seconds).to_string(),
|
||||
window_seconds,
|
||||
label_display,
|
||||
used_percent,
|
||||
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
local_usage: None,
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
window_used_display: quota_window_used_display(reset_at, window_seconds),
|
||||
reset_at,
|
||||
});
|
||||
}
|
||||
|
||||
fn snapshot_source_key(source: &str, snapshot: &Value) -> String {
|
||||
let label = snapshot
|
||||
.get("limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
|
||||
.unwrap_or(source);
|
||||
format!("{}-{}", quota_key_segment(source), quota_key_segment(label))
|
||||
}
|
||||
|
||||
fn snapshot_label(snapshot: &Value) -> Option<String> {
|
||||
let source = snapshot.get("source").and_then(Value::as_str);
|
||||
if source == Some("core") {
|
||||
return None;
|
||||
}
|
||||
let label = snapshot
|
||||
.get("limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
if label.eq_ignore_ascii_case("codex") {
|
||||
return None;
|
||||
}
|
||||
if is_review_limit_label(Some(label)) {
|
||||
return Some("代码审查".to_string());
|
||||
}
|
||||
Some(label.to_string())
|
||||
}
|
||||
|
||||
fn unique_quota_window_key(
|
||||
windows: &[AccountQuotaWindowData],
|
||||
source_key: &str,
|
||||
role: &str,
|
||||
window_seconds: Option<u64>,
|
||||
) -> String {
|
||||
let bucket = quota_window_key_part(window_seconds).unwrap_or(role);
|
||||
let key = format!("{source_key}-{bucket}");
|
||||
if windows.iter().any(|window| window.key == key) {
|
||||
format!("{key}-{role}")
|
||||
} else {
|
||||
key
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_key_part(window_seconds: Option<u64>) -> Option<&'static str> {
|
||||
match window_seconds {
|
||||
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
|
||||
Some("five-hour")
|
||||
}
|
||||
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => Some("weekly"),
|
||||
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => Some("monthly"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_key_segment(value: &str) -> String {
|
||||
let mut segment = String::with_capacity(value.len());
|
||||
for ch in value.chars() {
|
||||
if ch.is_ascii_alphanumeric() {
|
||||
segment.push(ch.to_ascii_lowercase());
|
||||
} else if !segment.ends_with('-') {
|
||||
segment.push('-');
|
||||
}
|
||||
}
|
||||
let segment = segment.trim_matches('-');
|
||||
if segment.is_empty() {
|
||||
"quota".to_string()
|
||||
} else {
|
||||
segment.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_sort_key(window: &AccountQuotaWindowData) -> (u8, u64, String) {
|
||||
let group_order = match window.group.as_str() {
|
||||
"monthly" => 0,
|
||||
"shortTerm" => 1,
|
||||
_ => 2,
|
||||
};
|
||||
(
|
||||
group_order,
|
||||
window.window_seconds.unwrap_or(0),
|
||||
window.key.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
fn quota_window_group(window_seconds: Option<u64>) -> &'static str {
|
||||
match window_seconds {
|
||||
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "monthly",
|
||||
Some(seconds)
|
||||
if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS)
|
||||
|| quota_window_matches(seconds, WEEK_WINDOW_SECONDS) =>
|
||||
{
|
||||
"shortTerm"
|
||||
}
|
||||
_ => "other",
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_matches(actual: u64, expected: u64) -> bool {
|
||||
actual > 0 && actual.abs_diff(expected) <= expected / 20
|
||||
}
|
||||
|
||||
fn quota_window_label_display(window_seconds: Option<u64>) -> String {
|
||||
let Some(window_seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
|
||||
return "额度".to_string();
|
||||
};
|
||||
match window_seconds {
|
||||
seconds if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
|
||||
"5小时限额".to_string()
|
||||
}
|
||||
seconds if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => "周限额".to_string(),
|
||||
seconds if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "月限额".to_string(),
|
||||
seconds if seconds % 86_400 == 0 => format!("{}天限额", seconds / 86_400),
|
||||
seconds if seconds % 3_600 == 0 => format!("{}小时限额", seconds / 3_600),
|
||||
seconds => format!("{}分钟限额", seconds.div_ceil(60)),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_used_display(
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
window_seconds: Option<u64>,
|
||||
) -> String {
|
||||
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
|
||||
return "-".to_string();
|
||||
};
|
||||
let remaining = reset_at
|
||||
.signed_duration_since(Utc::now())
|
||||
.num_seconds()
|
||||
.max(0)
|
||||
.cast_unsigned();
|
||||
let used = window_seconds.saturating_sub(remaining);
|
||||
format!(
|
||||
"{} / {}",
|
||||
format_duration_days(used),
|
||||
format_duration_days(window_seconds)
|
||||
)
|
||||
}
|
||||
|
||||
fn format_duration_days(seconds: u64) -> String {
|
||||
let days = seconds as f64 / 86_400.0;
|
||||
if days >= 1.0 {
|
||||
format!("{days:.1}d")
|
||||
} else {
|
||||
format!("{:.1}h", seconds as f64 / 3_600.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn number_value(value: &Value) -> Option<f64> {
|
||||
value
|
||||
.as_f64()
|
||||
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn is_review_limit_label(value: Option<&str>) -> bool {
|
||||
let normalized = value
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.replace(['-', ' '], "_");
|
||||
normalized == "review"
|
||||
|| normalized == "code_review"
|
||||
|| normalized == "codex_review"
|
||||
|| normalized == "codex_code_review"
|
||||
|| normalized.contains("code_review")
|
||||
|| normalized.contains("codex_review")
|
||||
}
|
||||
@@ -9,14 +9,14 @@ use super::oauth;
|
||||
|
||||
use crate::{
|
||||
fleet::{
|
||||
account_gateway::AccountUpstreamGateway,
|
||||
cookies::PgCookieStore,
|
||||
pool::AccountPoolService,
|
||||
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy},
|
||||
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy, TokenRefresher},
|
||||
store::PgAccountStore,
|
||||
},
|
||||
infra::identity::AccountPseudonymizer,
|
||||
models::service::ModelService,
|
||||
upstream::openai::{token_client::TokenRefresher, transport::CodexBackendClient},
|
||||
};
|
||||
|
||||
use super::types::{AccountManageError, RefreshedAccountTokens};
|
||||
@@ -25,7 +25,7 @@ use super::types::{AccountManageError, RefreshedAccountTokens};
|
||||
pub struct AccountManageService {
|
||||
pub store: PgAccountStore,
|
||||
pub(crate) cookies: PgCookieStore,
|
||||
pub(crate) codex: Arc<CodexBackendClient>,
|
||||
pub(crate) upstream: Arc<dyn AccountUpstreamGateway>,
|
||||
pub(crate) models: Arc<ModelService>,
|
||||
pub(crate) account_pool: Arc<AccountPoolService>,
|
||||
pub(crate) token_refresher: Arc<dyn TokenRefresher>,
|
||||
@@ -39,7 +39,7 @@ pub struct AccountManageService {
|
||||
pub struct AccountManageServiceParts {
|
||||
pub store: PgAccountStore,
|
||||
pub cookies: PgCookieStore,
|
||||
pub codex: Arc<CodexBackendClient>,
|
||||
pub upstream: Arc<dyn AccountUpstreamGateway>,
|
||||
pub models: Arc<ModelService>,
|
||||
pub account_pool: Arc<AccountPoolService>,
|
||||
pub token_refresher: Arc<dyn TokenRefresher>,
|
||||
@@ -54,7 +54,7 @@ impl AccountManageService {
|
||||
Self {
|
||||
store: parts.store,
|
||||
cookies: parts.cookies,
|
||||
codex: parts.codex,
|
||||
upstream: parts.upstream,
|
||||
models: parts.models,
|
||||
account_pool: parts.account_pool,
|
||||
token_refresher: parts.token_refresher,
|
||||
@@ -127,7 +127,7 @@ impl AccountManageService {
|
||||
}
|
||||
|
||||
pub(crate) async fn evict_account_websocket_pool(&self, account_id: &str) {
|
||||
self.codex.evict_websocket_account(account_id).await;
|
||||
self.upstream.evict_account_connections(account_id).await;
|
||||
match self.store.get(account_id).await {
|
||||
Ok(Some(account)) => {
|
||||
if let Some(upstream_account_id) = account
|
||||
@@ -135,8 +135,8 @@ impl AccountManageService {
|
||||
.as_deref()
|
||||
.filter(|value| *value != account_id)
|
||||
{
|
||||
self.codex
|
||||
.evict_websocket_account(upstream_account_id)
|
||||
self.upstream
|
||||
.evict_account_connections(upstream_account_id)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,13 +4,12 @@ use chrono::{DateTime, Utc};
|
||||
use serde::Serialize;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::{
|
||||
fleet::{
|
||||
account::{AccountStatus, AccountStatus as AcctStatus},
|
||||
account_failure::classify_client_failure,
|
||||
store::StoredAccountMetadata,
|
||||
},
|
||||
upstream::openai::{token_client::RefreshFailure, transport::CodexClientError},
|
||||
use crate::fleet::{
|
||||
account::{AccountStatus, AccountStatus as AcctStatus},
|
||||
account_failure::classify_account_failure,
|
||||
account_gateway::AccountGatewayError,
|
||||
refresh::RefreshFailure,
|
||||
store::StoredAccountMetadata,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
@@ -86,6 +85,8 @@ pub enum AccountManageError {
|
||||
FetchQuota(String),
|
||||
#[error("no available models returned")]
|
||||
NoModels,
|
||||
#[error("failed to refresh account models: {0}")]
|
||||
RefreshModels(String),
|
||||
#[error("invalid status: {0}")]
|
||||
InvalidStatus(String),
|
||||
#[error("label must be 64 characters or fewer")]
|
||||
@@ -296,21 +297,6 @@ pub(super) fn stored_to_admin_metadata(s: crate::fleet::store::StoredAccount) ->
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn import_usage_plan_type(usage: &serde_json::Value) -> Option<String> {
|
||||
usage
|
||||
.get("plan_type")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.and_then(normalized_plan_type)
|
||||
}
|
||||
|
||||
pub(super) fn import_usage_string(usage: &serde_json::Value, key: &str) -> Option<String> {
|
||||
usage
|
||||
.get(key)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.and_then(|value| crate::fleet::import::normalize_nonempty_str(Some(value)))
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
pub(super) fn import_quota_plan_type(quota: &serde_json::Value) -> Option<String> {
|
||||
quota
|
||||
.get("plan_type")
|
||||
@@ -323,8 +309,8 @@ fn normalized_plan_type(value: &str) -> Option<String> {
|
||||
(!value.is_empty() && !matches!(value.as_str(), "unknown" | "null")).then_some(value)
|
||||
}
|
||||
|
||||
pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option<AccountStatus> {
|
||||
classify_client_failure(error)?
|
||||
pub(super) fn import_status_from_usage_error(error: &AccountGatewayError) -> Option<AccountStatus> {
|
||||
classify_account_failure(error.failure()?)?
|
||||
.account_status()
|
||||
.filter(|status| *status == AccountStatus::Banned)
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
pub mod account;
|
||||
pub mod account_failure;
|
||||
pub mod account_gateway;
|
||||
pub mod cookies;
|
||||
pub mod import;
|
||||
pub mod manage;
|
||||
@@ -10,4 +11,5 @@ pub mod quota;
|
||||
pub mod refresh;
|
||||
pub mod scheduler;
|
||||
pub mod store;
|
||||
pub mod usage;
|
||||
mod window;
|
||||
|
||||
@@ -12,23 +12,24 @@ use chrono::{DateTime, Duration, Utc};
|
||||
use indexmap::IndexMap;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::fleet::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
|
||||
use crate::fleet::quota::{
|
||||
QuotaObservation, QuotaSnapshot, quota_from_observation, quota_snapshot_limit_reached,
|
||||
quota_snapshot_reset_at,
|
||||
};
|
||||
use crate::fleet::refresh::{JwtExpiry, jwt_expiry};
|
||||
use crate::fleet::scheduler::{
|
||||
AccountScheduler, ScoreWeights,
|
||||
AccountScheduler, AttemptFeedback, ScoreWeights,
|
||||
candidates::{self, CandidateFilter, CandidateRequest},
|
||||
};
|
||||
use crate::fleet::store::AccountStore;
|
||||
use crate::fleet::{
|
||||
account::{Account, AccountStatus},
|
||||
usage::{
|
||||
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageWindow,
|
||||
ResponseUsage,
|
||||
},
|
||||
window::should_reset_usage_window,
|
||||
};
|
||||
use crate::telemetry::account_usage::store::{
|
||||
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageWindow,
|
||||
};
|
||||
use crate::upstream::openai::protocol::events::{
|
||||
TokenUsage as CodexTokenUsage, parse_rate_limit_headers, rate_limit_quota,
|
||||
};
|
||||
|
||||
pub use crate::fleet::scheduler::RotationStrategy;
|
||||
|
||||
@@ -48,23 +49,10 @@ pub struct AccountPoolStaticSettings {
|
||||
pub tier_priority: Vec<String>,
|
||||
}
|
||||
|
||||
impl AccountPoolStaticSettings {
|
||||
/// 将设置快照转换为账号池选项。
|
||||
pub fn pool_options(&self, settings: &crate::settings::SettingsSnapshot) -> AccountPoolOptions {
|
||||
AccountPoolOptions {
|
||||
max_concurrent_per_account: settings.max_concurrent_per_account,
|
||||
rotation_strategy: match settings.rotation_strategy.as_str() {
|
||||
"smart" => RotationStrategy::Smart,
|
||||
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
|
||||
"round_robin" => RotationStrategy::RoundRobin,
|
||||
"sticky" => RotationStrategy::Sticky,
|
||||
_ => RotationStrategy::Smart,
|
||||
},
|
||||
skip_quota_limited: self.skip_quota_limited,
|
||||
tier_priority: self.tier_priority.clone(),
|
||||
..AccountPoolOptions::default()
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountPoolRuntimeOptions {
|
||||
pub options: AccountPoolOptions,
|
||||
pub request_interval_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -223,19 +211,20 @@ impl AccountPoolService {
|
||||
StdDuration::from_millis(request_interval_ms);
|
||||
}
|
||||
|
||||
/// 回灌一个已经启动的上游 attempt,供 Smart 调度器更新运行时健康反馈。
|
||||
pub async fn report_attempt(&self, account_id: &str, feedback: AttemptFeedback) {
|
||||
self.pool.lock().await.report_attempt(account_id, feedback);
|
||||
}
|
||||
|
||||
/// 持续接收运行时设置并更新账号池参数。
|
||||
pub async fn subscribe_settings(
|
||||
pub async fn subscribe_options(
|
||||
self: Arc<Self>,
|
||||
mut receiver: tokio::sync::watch::Receiver<crate::settings::SettingsSnapshot>,
|
||||
static_settings: AccountPoolStaticSettings,
|
||||
mut receiver: tokio::sync::watch::Receiver<AccountPoolRuntimeOptions>,
|
||||
) {
|
||||
while receiver.changed().await.is_ok() {
|
||||
let settings = receiver.borrow_and_update().clone();
|
||||
self.apply_options(
|
||||
static_settings.pool_options(&settings),
|
||||
settings.request_interval_ms,
|
||||
)
|
||||
.await;
|
||||
let runtime = receiver.borrow_and_update().clone();
|
||||
self.apply_options(runtime.options, runtime.request_interval_ms)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -541,16 +530,11 @@ impl AccountPoolService {
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录上游 token 用量。
|
||||
pub async fn record_token_usage(&self, account_id: &str, usage: &CodexTokenUsage) {
|
||||
self.record_response_usage(account_id, *usage, false).await;
|
||||
}
|
||||
|
||||
/// 记录 Responses 请求用量。
|
||||
pub async fn record_response_usage(
|
||||
&self,
|
||||
account_id: &str,
|
||||
usage: CodexTokenUsage,
|
||||
usage: ResponseUsage,
|
||||
image_generation_requested: bool,
|
||||
) {
|
||||
let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0;
|
||||
@@ -621,32 +605,29 @@ impl AccountPoolService {
|
||||
}
|
||||
}
|
||||
|
||||
/// 根据上游返回的 rate-limit header 被动同步 quota 状态。
|
||||
pub async fn sync_passive_rate_limit_headers(
|
||||
/// 被动同步上游 adapter 已归一化的 quota observation。
|
||||
pub async fn sync_passive_quota_observation(
|
||||
&self,
|
||||
account: &Account,
|
||||
headers: &[(String, String)],
|
||||
observation: &QuotaObservation,
|
||||
) {
|
||||
self.sync_passive_rate_limit_headers_for_account(
|
||||
self.sync_passive_quota_observation_for_account(
|
||||
&account.id,
|
||||
account.plan_type.as_deref(),
|
||||
headers,
|
||||
observation,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// 根据上游返回的 rate-limit header 被动同步指定账号 quota 状态。
|
||||
pub async fn sync_passive_rate_limit_headers_for_account(
|
||||
/// 被动同步指定账号的 typed quota observation。
|
||||
pub async fn sync_passive_quota_observation_for_account(
|
||||
&self,
|
||||
account_id: &str,
|
||||
plan_type: Option<&str>,
|
||||
headers: &[(String, String)],
|
||||
observation: &QuotaObservation,
|
||||
) {
|
||||
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
|
||||
return;
|
||||
};
|
||||
let existing_quota = match self.store.get_quota_json(account_id).await {
|
||||
Ok(Some(quota_json)) => serde_json::from_str::<serde_json::Value>("a_json).ok(),
|
||||
Ok(Some(quota_json)) => QuotaSnapshot::from_json("a_json).ok(),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
tracing::error!(
|
||||
@@ -657,18 +638,28 @@ impl AccountPoolService {
|
||||
None
|
||||
}
|
||||
};
|
||||
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
|
||||
let quota = quota_from_observation(observation, plan_type, existing_quota.as_ref());
|
||||
self.apply_quota_snapshot(account_id, "a).await;
|
||||
}
|
||||
|
||||
/// 应用 quota 快照到持久化存储和运行时账号池。
|
||||
pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &serde_json::Value) -> bool {
|
||||
pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &QuotaSnapshot) -> bool {
|
||||
let limit_reached = quota_snapshot_limit_reached(quota);
|
||||
let reset_at = quota_snapshot_reset_at(quota);
|
||||
let cooldown_until = limit_reached.then_some(reset_at).flatten();
|
||||
let limit_window_seconds =
|
||||
reset_at.and_then(|_| crate::fleet::quota::quota_snapshot_limit_window_seconds(quota));
|
||||
let quota_json = quota.to_string();
|
||||
let quota_json = match quota.to_json() {
|
||||
Ok(quota_json) => quota_json,
|
||||
Err(error) => {
|
||||
tracing::error!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"Failed to serialize verified quota snapshot"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
};
|
||||
let persisted = match self
|
||||
.store
|
||||
.apply_quota_snapshot(account_id, "a_json, limit_reached, cooldown_until)
|
||||
|
||||
@@ -206,10 +206,13 @@ impl AccountPool {
|
||||
self.options = options;
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果到调度器的运行时反馈(供 Smart 打分)。
|
||||
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
|
||||
self.scheduler
|
||||
.report_feedback(account_id, success, first_token_ms);
|
||||
/// 回灌一个已启动上游 attempt 的运行时反馈(供 Smart 打分)。
|
||||
pub fn report_attempt(
|
||||
&self,
|
||||
account_id: &str,
|
||||
feedback: crate::fleet::scheduler::AttemptFeedback,
|
||||
) {
|
||||
self.scheduler.report_attempt(account_id, feedback);
|
||||
}
|
||||
|
||||
/// 更新账号状态。
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
//! 账号配额解析与刷新。
|
||||
|
||||
mod model;
|
||||
mod observation;
|
||||
mod runtime;
|
||||
mod service;
|
||||
|
||||
pub use model::*;
|
||||
pub use observation::*;
|
||||
pub use runtime::*;
|
||||
pub use service::*;
|
||||
|
||||
@@ -0,0 +1,673 @@
|
||||
//! fleet 唯一拥有的 typed quota 快照。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
pub const FIVE_HOUR_WINDOW_MINUTES: u64 = 300;
|
||||
pub const WEEK_WINDOW_MINUTES: u64 = 10_080;
|
||||
pub const MONTH_WINDOW_MINUTES: u64 = 43_200;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum QuotaWindowKind {
|
||||
ShortTerm,
|
||||
Weekly,
|
||||
Monthly,
|
||||
#[default]
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum QuotaWindowRole {
|
||||
Primary,
|
||||
Secondary,
|
||||
Monthly,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub struct QuotaWindow {
|
||||
#[serde(skip)]
|
||||
pub kind: QuotaWindowKind,
|
||||
pub used_percent: Option<f64>,
|
||||
pub remaining_percent: Option<i64>,
|
||||
pub reset_at: Option<i64>,
|
||||
pub window_minutes: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub limit_reached: bool,
|
||||
pub used: Option<Value>,
|
||||
pub limit: Option<Value>,
|
||||
pub used_credits: Option<Value>,
|
||||
pub limit_credits: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct QuotaWindowWire {
|
||||
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
|
||||
used_percent: Option<f64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
|
||||
remaining_percent: Option<f64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_positive_i64")]
|
||||
reset_at: Option<i64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_positive_u64")]
|
||||
window_minutes: Option<u64>,
|
||||
#[serde(default)]
|
||||
limit_reached: bool,
|
||||
#[serde(default)]
|
||||
used: Option<Value>,
|
||||
#[serde(default)]
|
||||
limit: Option<Value>,
|
||||
#[serde(default)]
|
||||
used_credits: Option<Value>,
|
||||
#[serde(default)]
|
||||
limit_credits: Option<Value>,
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for QuotaWindow {
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let wire = QuotaWindowWire::deserialize(deserializer)?;
|
||||
Ok(Self {
|
||||
kind: quota_window_kind(wire.window_minutes),
|
||||
used_percent: wire.used_percent.map(clamp_percent),
|
||||
remaining_percent: wire
|
||||
.remaining_percent
|
||||
.map(clamp_percent)
|
||||
.map(|value| value.round() as i64),
|
||||
reset_at: wire.reset_at,
|
||||
window_minutes: wire.window_minutes,
|
||||
limit_reached: wire.limit_reached,
|
||||
used: wire.used,
|
||||
limit: wire.limit,
|
||||
used_credits: wire.used_credits,
|
||||
limit_credits: wire.limit_credits,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl QuotaWindow {
|
||||
pub fn from_usage(
|
||||
used_percent: f64,
|
||||
reset_at: Option<i64>,
|
||||
window_minutes: Option<u64>,
|
||||
used: Option<Value>,
|
||||
limit: Option<Value>,
|
||||
) -> Self {
|
||||
let used_percent = clamp_percent(used_percent);
|
||||
Self {
|
||||
kind: quota_window_kind(window_minutes),
|
||||
used_percent: Some(used_percent),
|
||||
remaining_percent: Some(remaining_percent(used_percent)),
|
||||
reset_at,
|
||||
window_minutes,
|
||||
limit_reached: used_percent >= 100.0,
|
||||
used,
|
||||
limit,
|
||||
used_credits: None,
|
||||
limit_credits: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_observation(
|
||||
used_percent: f64,
|
||||
reset_at: Option<i64>,
|
||||
window_minutes: Option<u64>,
|
||||
) -> Self {
|
||||
Self::from_usage(used_percent, reset_at, window_minutes, None, None)
|
||||
}
|
||||
|
||||
pub fn window_seconds(&self) -> Option<u64> {
|
||||
self.window_minutes?.checked_mul(60)
|
||||
}
|
||||
|
||||
pub fn reset_datetime(&self) -> Option<DateTime<Utc>> {
|
||||
DateTime::<Utc>::from_timestamp(self.reset_at?, 0)
|
||||
}
|
||||
|
||||
pub fn is_limit_reached(&self) -> bool {
|
||||
self.limit_reached || self.used_percent.is_some_and(|used| used >= 100.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct QuotaLimitSnapshot {
|
||||
pub source: String,
|
||||
#[serde(default)]
|
||||
pub limit_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub metered_feature: Option<String>,
|
||||
#[serde(default)]
|
||||
pub allowed: Option<bool>,
|
||||
#[serde(default)]
|
||||
pub limit_reached: Option<bool>,
|
||||
#[serde(default)]
|
||||
pub blocked: bool,
|
||||
#[serde(default)]
|
||||
pub primary: Option<QuotaWindow>,
|
||||
#[serde(default)]
|
||||
pub secondary: Option<QuotaWindow>,
|
||||
}
|
||||
|
||||
impl QuotaLimitSnapshot {
|
||||
pub fn is_blocked(&self) -> bool {
|
||||
self.blocked
|
||||
|| self.limit_reached.unwrap_or(false)
|
||||
|| self.allowed.is_some_and(|allowed| !allowed)
|
||||
|| self.windows().any(QuotaWindow::is_limit_reached)
|
||||
}
|
||||
|
||||
pub fn windows(&self) -> impl Iterator<Item = &QuotaWindow> {
|
||||
[self.primary.as_ref(), self.secondary.as_ref()]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MonthlyQuotaLimit {
|
||||
#[serde(default)]
|
||||
pub key: Option<String>,
|
||||
#[serde(default)]
|
||||
pub source: Option<String>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
|
||||
pub used_percent: Option<f64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
|
||||
pub remaining_percent: Option<f64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_positive_i64")]
|
||||
pub reset_at: Option<i64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_positive_u64")]
|
||||
pub window_minutes: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub limit_reached: bool,
|
||||
#[serde(default)]
|
||||
pub used: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub limit: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub used_credits: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub limit_credits: Option<Value>,
|
||||
}
|
||||
|
||||
impl MonthlyQuotaLimit {
|
||||
pub fn as_window(&self) -> QuotaWindow {
|
||||
QuotaWindow {
|
||||
kind: QuotaWindowKind::Monthly,
|
||||
used_percent: self.used_percent.map(clamp_percent),
|
||||
remaining_percent: self
|
||||
.remaining_percent
|
||||
.map(clamp_percent)
|
||||
.map(|value| value.round() as i64),
|
||||
reset_at: self.reset_at,
|
||||
window_minutes: self.window_minutes.or(Some(MONTH_WINDOW_MINUTES)),
|
||||
limit_reached: self.limit_reached,
|
||||
used: self.used.clone(),
|
||||
limit: self.limit.clone(),
|
||||
used_credits: self.used_credits.clone(),
|
||||
limit_credits: self.limit_credits.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_limit_reached(&self) -> bool {
|
||||
self.limit_reached || self.used_percent.is_some_and(|used| used >= 100.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct QuotaCredits {
|
||||
#[serde(default)]
|
||||
pub has_credits: bool,
|
||||
#[serde(default)]
|
||||
pub unlimited: bool,
|
||||
#[serde(default)]
|
||||
pub overage_limit_reached: bool,
|
||||
#[serde(default)]
|
||||
pub balance: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct QuotaSpendLimit {
|
||||
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
|
||||
pub used_percent: Option<f64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
|
||||
pub remaining_percent: Option<f64>,
|
||||
#[serde(default, deserialize_with = "deserialize_optional_positive_i64")]
|
||||
pub reset_at: Option<i64>,
|
||||
#[serde(default)]
|
||||
pub used: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub limit: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct QuotaSpendControl {
|
||||
#[serde(default)]
|
||||
pub reached: bool,
|
||||
#[serde(default)]
|
||||
pub individual_limit: Option<QuotaSpendLimit>,
|
||||
#[serde(flatten)]
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct QuotaSnapshot {
|
||||
#[serde(default)]
|
||||
pub plan_type: Option<String>,
|
||||
#[serde(default)]
|
||||
pub active_limit: Option<String>,
|
||||
#[serde(default, rename = "snapshots")]
|
||||
pub limits: Vec<QuotaLimitSnapshot>,
|
||||
#[serde(default)]
|
||||
pub monthly_limit: Option<MonthlyQuotaLimit>,
|
||||
#[serde(default)]
|
||||
pub credits: Option<QuotaCredits>,
|
||||
#[serde(default)]
|
||||
pub spend_control: Option<QuotaSpendControl>,
|
||||
#[serde(default)]
|
||||
pub promo_message: Option<String>,
|
||||
#[serde(default)]
|
||||
pub rate_limit_reached_type: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extensions: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
impl QuotaSnapshot {
|
||||
pub fn from_json(value: &str) -> serde_json::Result<Self> {
|
||||
serde_json::from_str(value)
|
||||
}
|
||||
|
||||
pub fn from_value(value: Value) -> serde_json::Result<Self> {
|
||||
serde_json::from_value(value)
|
||||
}
|
||||
|
||||
pub fn to_json(&self) -> serde_json::Result<String> {
|
||||
serde_json::to_string(self)
|
||||
}
|
||||
|
||||
pub fn is_limit_reached(&self) -> bool {
|
||||
self.limits.iter().any(QuotaLimitSnapshot::is_blocked)
|
||||
|| self
|
||||
.monthly_limit
|
||||
.as_ref()
|
||||
.is_some_and(MonthlyQuotaLimit::is_limit_reached)
|
||||
|| self
|
||||
.spend_control
|
||||
.as_ref()
|
||||
.is_some_and(quota_spend_control_limit_reached)
|
||||
|| self
|
||||
.credits
|
||||
.as_ref()
|
||||
.is_some_and(|credits| credits.overage_limit_reached)
|
||||
}
|
||||
|
||||
pub fn reset_at(&self) -> Option<DateTime<Utc>> {
|
||||
let reset_at = if self.monthly_or_spend_limit_reached() {
|
||||
self.monthly_reset_at()
|
||||
.or_else(|| self.blocking_window().and_then(|window| window.reset_at))
|
||||
.or_else(|| {
|
||||
self.core_window_with_reset()
|
||||
.and_then(|window| window.reset_at)
|
||||
})
|
||||
} else if let Some(reset_at) = self.blocking_window().and_then(|window| window.reset_at) {
|
||||
Some(reset_at)
|
||||
} else {
|
||||
self.core_window_with_reset()
|
||||
.and_then(|window| window.reset_at)
|
||||
.or_else(|| self.monthly_reset_at())
|
||||
}?;
|
||||
DateTime::<Utc>::from_timestamp(reset_at, 0)
|
||||
}
|
||||
|
||||
pub fn limit_window_seconds(&self) -> Option<u64> {
|
||||
if self.monthly_or_spend_limit_reached() {
|
||||
return self
|
||||
.monthly_window_seconds()
|
||||
.or(MONTH_WINDOW_MINUTES.checked_mul(60));
|
||||
}
|
||||
self.blocking_window()
|
||||
.and_then(QuotaWindow::window_seconds)
|
||||
.or_else(|| self.core_windows().find_map(QuotaWindow::window_seconds))
|
||||
.or_else(|| self.monthly_window_seconds())
|
||||
}
|
||||
|
||||
pub fn windows(&self) -> impl Iterator<Item = QuotaWindowEntry<'_>> {
|
||||
let monthly = self.monthly_limit.as_ref().map(|limit| QuotaWindowEntry {
|
||||
source: limit.source.as_deref().unwrap_or("monthly"),
|
||||
limit_name: limit.key.as_deref(),
|
||||
metered_feature: None,
|
||||
role: QuotaWindowRole::Monthly,
|
||||
window: QuotaWindowRef::Monthly(limit),
|
||||
});
|
||||
monthly
|
||||
.into_iter()
|
||||
.chain(self.limits.iter().flat_map(|limit| {
|
||||
let primary = limit.primary.as_ref().map(|window| QuotaWindowEntry {
|
||||
source: &limit.source,
|
||||
limit_name: limit.limit_name.as_deref(),
|
||||
metered_feature: limit.metered_feature.as_deref(),
|
||||
role: QuotaWindowRole::Primary,
|
||||
window: QuotaWindowRef::Regular(window),
|
||||
});
|
||||
let secondary = limit.secondary.as_ref().map(|window| QuotaWindowEntry {
|
||||
source: &limit.source,
|
||||
limit_name: limit.limit_name.as_deref(),
|
||||
metered_feature: limit.metered_feature.as_deref(),
|
||||
role: QuotaWindowRole::Secondary,
|
||||
window: QuotaWindowRef::Regular(window),
|
||||
});
|
||||
primary.into_iter().chain(secondary)
|
||||
}))
|
||||
}
|
||||
|
||||
fn monthly_or_spend_limit_reached(&self) -> bool {
|
||||
self.monthly_limit
|
||||
.as_ref()
|
||||
.is_some_and(MonthlyQuotaLimit::is_limit_reached)
|
||||
|| self
|
||||
.spend_control
|
||||
.as_ref()
|
||||
.is_some_and(quota_spend_control_limit_reached)
|
||||
}
|
||||
|
||||
fn monthly_reset_at(&self) -> Option<i64> {
|
||||
self.monthly_limit
|
||||
.as_ref()?
|
||||
.reset_at
|
||||
.filter(|value| *value > 0)
|
||||
}
|
||||
|
||||
fn monthly_window_seconds(&self) -> Option<u64> {
|
||||
self.monthly_limit
|
||||
.as_ref()?
|
||||
.window_minutes
|
||||
.filter(|value| *value > 0)?
|
||||
.checked_mul(60)
|
||||
}
|
||||
|
||||
fn blocking_window(&self) -> Option<&QuotaWindow> {
|
||||
self.limits
|
||||
.iter()
|
||||
.filter(|limit| limit.is_blocked())
|
||||
.flat_map(QuotaLimitSnapshot::windows)
|
||||
.find(|window| window.is_limit_reached())
|
||||
.or_else(|| {
|
||||
self.limits
|
||||
.iter()
|
||||
.filter(|limit| limit.is_blocked())
|
||||
.flat_map(QuotaLimitSnapshot::windows)
|
||||
.find(|window| window.reset_at.is_some())
|
||||
})
|
||||
}
|
||||
|
||||
fn core_window_with_reset(&self) -> Option<&QuotaWindow> {
|
||||
self.core_windows().find(|window| window.reset_at.is_some())
|
||||
}
|
||||
|
||||
fn core_windows(&self) -> impl Iterator<Item = &QuotaWindow> {
|
||||
self.limits
|
||||
.iter()
|
||||
.filter(|limit| limit.source == "core")
|
||||
.flat_map(QuotaLimitSnapshot::windows)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum QuotaWindowRef<'a> {
|
||||
Regular(&'a QuotaWindow),
|
||||
Monthly(&'a MonthlyQuotaLimit),
|
||||
}
|
||||
|
||||
impl QuotaWindowRef<'_> {
|
||||
pub fn to_owned(self) -> QuotaWindow {
|
||||
match self {
|
||||
Self::Regular(window) => window.clone(),
|
||||
Self::Monthly(limit) => limit.as_window(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct QuotaWindowEntry<'a> {
|
||||
pub source: &'a str,
|
||||
pub limit_name: Option<&'a str>,
|
||||
pub metered_feature: Option<&'a str>,
|
||||
pub role: QuotaWindowRole,
|
||||
pub window: QuotaWindowRef<'a>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ClassifiedQuotaWindow {
|
||||
pub key: String,
|
||||
pub source: String,
|
||||
pub limit_name: Option<String>,
|
||||
pub metered_feature: Option<String>,
|
||||
pub role: QuotaWindowRole,
|
||||
pub window: QuotaWindow,
|
||||
pub order: u8,
|
||||
}
|
||||
|
||||
impl QuotaSnapshot {
|
||||
pub fn classified_windows(&self) -> Vec<ClassifiedQuotaWindow> {
|
||||
let has_monthly_limit = self.monthly_limit.is_some();
|
||||
let mut windows = Vec::new();
|
||||
for entry in self.windows() {
|
||||
let window = entry.window.to_owned();
|
||||
if has_monthly_limit
|
||||
&& entry.source == "core"
|
||||
&& entry.role != QuotaWindowRole::Monthly
|
||||
&& window.kind == QuotaWindowKind::Monthly
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let has_semantic_label = entry
|
||||
.limit_name
|
||||
.or(entry.metered_feature)
|
||||
.is_some_and(|label| !label.trim().is_empty());
|
||||
if window.reset_at.is_none()
|
||||
&& window.window_minutes.is_none()
|
||||
&& (window.used_percent.is_none() || !has_semantic_label)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let source_key =
|
||||
quota_source_key(entry.source, entry.limit_name, entry.metered_feature);
|
||||
let bucket = quota_window_key_part(window.kind)
|
||||
.unwrap_or_else(|| quota_window_role_name(entry.role));
|
||||
let mut key = format!("{source_key}-{bucket}");
|
||||
if windows
|
||||
.iter()
|
||||
.any(|existing: &ClassifiedQuotaWindow| existing.key == key)
|
||||
{
|
||||
key = format!("{key}-{}", quota_window_role_name(entry.role));
|
||||
}
|
||||
windows.push(ClassifiedQuotaWindow {
|
||||
key,
|
||||
source: entry.source.to_string(),
|
||||
limit_name: entry.limit_name.map(ToString::to_string),
|
||||
metered_feature: entry.metered_feature.map(ToString::to_string),
|
||||
role: entry.role,
|
||||
order: quota_window_order(window.kind),
|
||||
window,
|
||||
});
|
||||
}
|
||||
windows.sort_by(|left, right| {
|
||||
left.order
|
||||
.cmp(&right.order)
|
||||
.then_with(|| {
|
||||
left.window
|
||||
.window_minutes
|
||||
.unwrap_or_default()
|
||||
.cmp(&right.window.window_minutes.unwrap_or_default())
|
||||
})
|
||||
.then_with(|| left.key.cmp(&right.key))
|
||||
});
|
||||
windows
|
||||
}
|
||||
|
||||
pub fn representative_used_percent(&self) -> Option<f64> {
|
||||
self.classified_windows()
|
||||
.into_iter()
|
||||
.filter_map(|window| {
|
||||
window
|
||||
.window
|
||||
.used_percent
|
||||
.map(|used| (quota_usage_priority(window.window.kind), used))
|
||||
})
|
||||
.fold(None::<(u8, f64)>, |selected, candidate| match selected {
|
||||
Some(current)
|
||||
if current.0 < candidate.0
|
||||
|| (current.0 == candidate.0 && current.1 >= candidate.1) =>
|
||||
{
|
||||
Some(current)
|
||||
}
|
||||
_ => Some(candidate),
|
||||
})
|
||||
.map(|(_, used)| used)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn quota_window_kind(window_minutes: Option<u64>) -> QuotaWindowKind {
|
||||
match window_minutes {
|
||||
Some(minutes) if window_minutes_match(minutes, FIVE_HOUR_WINDOW_MINUTES) => {
|
||||
QuotaWindowKind::ShortTerm
|
||||
}
|
||||
Some(minutes) if window_minutes_match(minutes, WEEK_WINDOW_MINUTES) => {
|
||||
QuotaWindowKind::Weekly
|
||||
}
|
||||
Some(minutes) if window_minutes_match(minutes, MONTH_WINDOW_MINUTES) => {
|
||||
QuotaWindowKind::Monthly
|
||||
}
|
||||
_ => QuotaWindowKind::Other,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn window_minutes_match(actual: u64, expected: u64) -> bool {
|
||||
actual > 0 && actual.abs_diff(expected) <= expected / 20
|
||||
}
|
||||
|
||||
fn quota_window_order(kind: QuotaWindowKind) -> u8 {
|
||||
match kind {
|
||||
QuotaWindowKind::Monthly => 0,
|
||||
QuotaWindowKind::ShortTerm => 1,
|
||||
QuotaWindowKind::Weekly => 2,
|
||||
QuotaWindowKind::Other => 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_usage_priority(kind: QuotaWindowKind) -> u8 {
|
||||
match kind {
|
||||
QuotaWindowKind::ShortTerm => 0,
|
||||
QuotaWindowKind::Weekly => 1,
|
||||
QuotaWindowKind::Monthly => 2,
|
||||
QuotaWindowKind::Other => 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_key_part(kind: QuotaWindowKind) -> Option<&'static str> {
|
||||
match kind {
|
||||
QuotaWindowKind::ShortTerm => Some("five-hour"),
|
||||
QuotaWindowKind::Weekly => Some("weekly"),
|
||||
QuotaWindowKind::Monthly => Some("monthly"),
|
||||
QuotaWindowKind::Other => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_role_name(role: QuotaWindowRole) -> &'static str {
|
||||
match role {
|
||||
QuotaWindowRole::Primary => "primary",
|
||||
QuotaWindowRole::Secondary => "secondary",
|
||||
QuotaWindowRole::Monthly => "monthly",
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_source_key(
|
||||
source: &str,
|
||||
limit_name: Option<&str>,
|
||||
metered_feature: Option<&str>,
|
||||
) -> String {
|
||||
let label = limit_name.or(metered_feature).unwrap_or(source);
|
||||
format!("{}-{}", quota_key_segment(source), quota_key_segment(label))
|
||||
}
|
||||
|
||||
fn quota_key_segment(value: &str) -> String {
|
||||
let mut segment = String::with_capacity(value.len());
|
||||
for ch in value.chars() {
|
||||
if ch.is_ascii_alphanumeric() {
|
||||
segment.push(ch.to_ascii_lowercase());
|
||||
} else if !segment.ends_with('-') {
|
||||
segment.push('-');
|
||||
}
|
||||
}
|
||||
let segment = segment.trim_matches('-');
|
||||
if segment.is_empty() {
|
||||
"quota".to_string()
|
||||
} else {
|
||||
segment.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn remaining_percent(used_percent: f64) -> i64 {
|
||||
(100.0 - clamp_percent(used_percent)).round() as i64
|
||||
}
|
||||
|
||||
fn quota_spend_control_limit_reached(spend_control: &QuotaSpendControl) -> bool {
|
||||
spend_control.reached
|
||||
|| spend_control
|
||||
.individual_limit
|
||||
.as_ref()
|
||||
.and_then(|limit| limit.used_percent)
|
||||
.is_some_and(|used| used >= 100.0)
|
||||
}
|
||||
|
||||
fn clamp_percent(value: f64) -> f64 {
|
||||
value.clamp(0.0, 100.0)
|
||||
}
|
||||
|
||||
fn deserialize_optional_finite_number<'de, D>(deserializer: D) -> Result<Option<f64>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<Value>::deserialize(deserializer)?;
|
||||
Ok(value.and_then(|value| {
|
||||
value
|
||||
.as_f64()
|
||||
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
|
||||
.filter(|value| value.is_finite())
|
||||
}))
|
||||
}
|
||||
|
||||
fn deserialize_optional_positive_u64<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<Value>::deserialize(deserializer)?;
|
||||
Ok(value
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok()))
|
||||
.or_else(|| value.as_str().and_then(|value| value.parse::<u64>().ok()))
|
||||
})
|
||||
.filter(|value| *value > 0))
|
||||
}
|
||||
|
||||
fn deserialize_optional_positive_i64<'de, D>(deserializer: D) -> Result<Option<i64>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<Value>::deserialize(deserializer)?;
|
||||
Ok(value
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok()))
|
||||
.or_else(|| value.as_str().and_then(|value| value.parse::<i64>().ok()))
|
||||
})
|
||||
.filter(|value| *value > 0))
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
//! 上游 adapter 交给 fleet 的 typed quota observation。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
MONTH_WINDOW_MINUTES, MonthlyQuotaLimit, QuotaCredits, QuotaLimitSnapshot, QuotaSnapshot,
|
||||
QuotaSpendControl, QuotaWindow, QuotaWindowKind, remaining_percent,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct QuotaWindowObservation {
|
||||
pub used_percent: f64,
|
||||
pub window_minutes: Option<u64>,
|
||||
pub reset_at: Option<i64>,
|
||||
pub used: Option<Value>,
|
||||
pub limit: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct QuotaLimitObservation {
|
||||
pub limit_id: String,
|
||||
pub limit_name: Option<String>,
|
||||
pub metered_feature: Option<String>,
|
||||
pub allowed: Option<bool>,
|
||||
pub limit_reached: Option<bool>,
|
||||
pub primary: Option<QuotaWindowObservation>,
|
||||
pub secondary: Option<QuotaWindowObservation>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct QuotaCreditsObservation {
|
||||
pub has_credits: bool,
|
||||
pub unlimited: bool,
|
||||
pub overage_limit_reached: bool,
|
||||
pub balance: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct QuotaObservation {
|
||||
pub limits: BTreeMap<String, QuotaLimitObservation>,
|
||||
pub active_limit: Option<String>,
|
||||
pub credits: Option<QuotaCreditsObservation>,
|
||||
pub spend_control: Option<QuotaSpendControl>,
|
||||
pub plan_type: Option<String>,
|
||||
pub promo_message: Option<String>,
|
||||
pub rate_limit_reached_type: Option<String>,
|
||||
}
|
||||
|
||||
/// 将 typed observation 合并为代理内部 quota 快照。
|
||||
pub fn quota_from_observation(
|
||||
observation: &QuotaObservation,
|
||||
plan_type: Option<&str>,
|
||||
existing: Option<&QuotaSnapshot>,
|
||||
) -> QuotaSnapshot {
|
||||
let mut limits = observation
|
||||
.limits
|
||||
.values()
|
||||
.filter_map(quota_limit_from_observation)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if let Some(existing) = existing {
|
||||
for limit in &existing.limits {
|
||||
if !observation
|
||||
.limits
|
||||
.keys()
|
||||
.any(|limit_id| quota_limit_matches(limit, limit_id))
|
||||
{
|
||||
limits.push(limit.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let spend_control = observation
|
||||
.spend_control
|
||||
.clone()
|
||||
.or_else(|| existing.and_then(|quota| quota.spend_control.clone()));
|
||||
let monthly_limit = spend_control
|
||||
.as_ref()
|
||||
.and_then(monthly_limit_from_spend_control)
|
||||
.or_else(|| monthly_limit_from_limits(&limits))
|
||||
.or_else(|| existing.and_then(|quota| quota.monthly_limit.clone()));
|
||||
let credits = observation
|
||||
.credits
|
||||
.as_ref()
|
||||
.map(|credits| QuotaCredits {
|
||||
has_credits: credits.has_credits,
|
||||
unlimited: credits.unlimited,
|
||||
overage_limit_reached: credits.overage_limit_reached,
|
||||
balance: credits.balance.clone(),
|
||||
})
|
||||
.or_else(|| existing.and_then(|quota| quota.credits.clone()));
|
||||
|
||||
QuotaSnapshot {
|
||||
plan_type: observation
|
||||
.plan_type
|
||||
.clone()
|
||||
.or_else(|| plan_type.map(ToString::to_string))
|
||||
.or_else(|| existing.and_then(|quota| quota.plan_type.clone()))
|
||||
.or_else(|| Some("unknown".to_string())),
|
||||
active_limit: observation
|
||||
.active_limit
|
||||
.clone()
|
||||
.or_else(|| existing.and_then(|quota| quota.active_limit.clone())),
|
||||
limits,
|
||||
monthly_limit,
|
||||
credits,
|
||||
spend_control,
|
||||
promo_message: observation
|
||||
.promo_message
|
||||
.clone()
|
||||
.or_else(|| existing.and_then(|quota| quota.promo_message.clone())),
|
||||
rate_limit_reached_type: observation
|
||||
.rate_limit_reached_type
|
||||
.clone()
|
||||
.or_else(|| existing.and_then(|quota| quota.rate_limit_reached_type.clone())),
|
||||
extensions: existing
|
||||
.map(|quota| quota.extensions.clone())
|
||||
.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_limit_from_observation(details: &QuotaLimitObservation) -> Option<QuotaLimitSnapshot> {
|
||||
let (source, metered_feature) = quota_identity(&details.limit_id);
|
||||
let limit_name = details
|
||||
.limit_name
|
||||
.clone()
|
||||
.or_else(|| (source != "core").then(|| details.limit_id.clone()));
|
||||
let primary = details.primary.as_ref().map(quota_window_from_observation);
|
||||
let secondary = details
|
||||
.secondary
|
||||
.as_ref()
|
||||
.map(quota_window_from_observation);
|
||||
if primary.is_none() && secondary.is_none() {
|
||||
return None;
|
||||
}
|
||||
let blocked = details.limit_reached.unwrap_or(false)
|
||||
|| details.allowed.is_some_and(|allowed| !allowed)
|
||||
|| primary.as_ref().is_some_and(QuotaWindow::is_limit_reached)
|
||||
|| secondary
|
||||
.as_ref()
|
||||
.is_some_and(QuotaWindow::is_limit_reached);
|
||||
Some(QuotaLimitSnapshot {
|
||||
source: source.to_string(),
|
||||
limit_name,
|
||||
metered_feature: details
|
||||
.metered_feature
|
||||
.clone()
|
||||
.or_else(|| metered_feature.map(ToString::to_string)),
|
||||
allowed: details.allowed,
|
||||
limit_reached: details.limit_reached,
|
||||
blocked,
|
||||
primary,
|
||||
secondary,
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_window_from_observation(window: &QuotaWindowObservation) -> QuotaWindow {
|
||||
QuotaWindow::from_usage(
|
||||
window.used_percent,
|
||||
window.reset_at,
|
||||
window.window_minutes,
|
||||
window.used.clone(),
|
||||
window.limit.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
fn quota_identity(limit_id: &str) -> (&'static str, Option<&str>) {
|
||||
if limit_id == "codex" {
|
||||
("core", None)
|
||||
} else if is_review_limit_name(limit_id) {
|
||||
("code_review", Some(limit_id))
|
||||
} else {
|
||||
("additional", Some(limit_id))
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_limit_matches(limit: &QuotaLimitSnapshot, limit_id: &str) -> bool {
|
||||
let (source, metered_feature) = quota_identity(limit_id);
|
||||
if limit.source != source {
|
||||
return false;
|
||||
}
|
||||
source == "core"
|
||||
|| limit.metered_feature.as_deref() == metered_feature
|
||||
|| limit.limit_name.as_deref() == Some(limit_id)
|
||||
}
|
||||
|
||||
fn monthly_limit_from_limits(limits: &[QuotaLimitSnapshot]) -> Option<MonthlyQuotaLimit> {
|
||||
let window = limits
|
||||
.iter()
|
||||
.filter(|limit| limit.source == "core")
|
||||
.flat_map(QuotaLimitSnapshot::windows)
|
||||
.find(|window| window.kind == QuotaWindowKind::Monthly)?;
|
||||
Some(MonthlyQuotaLimit {
|
||||
key: Some("core-monthly".to_string()),
|
||||
source: Some("rate_limit".to_string()),
|
||||
used_percent: window.used_percent,
|
||||
remaining_percent: window.remaining_percent.map(|value| value as f64),
|
||||
reset_at: window.reset_at,
|
||||
window_minutes: window.window_minutes,
|
||||
limit_reached: window.is_limit_reached(),
|
||||
used: window.used.clone(),
|
||||
limit: window.limit.clone(),
|
||||
used_credits: window.used_credits.clone(),
|
||||
limit_credits: window.limit_credits.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn monthly_limit_from_spend_control(
|
||||
spend_control: &QuotaSpendControl,
|
||||
) -> Option<MonthlyQuotaLimit> {
|
||||
let individual = spend_control.individual_limit.as_ref()?;
|
||||
let used_percent = individual.used_percent.unwrap_or(0.0).clamp(0.0, 100.0);
|
||||
Some(MonthlyQuotaLimit {
|
||||
key: Some("spend-control-monthly".to_string()),
|
||||
source: Some("spend_control".to_string()),
|
||||
used_percent: Some(used_percent),
|
||||
remaining_percent: individual
|
||||
.remaining_percent
|
||||
.or_else(|| Some(remaining_percent(used_percent) as f64)),
|
||||
reset_at: individual.reset_at,
|
||||
window_minutes: Some(MONTH_WINDOW_MINUTES),
|
||||
limit_reached: spend_control.reached || used_percent >= 100.0,
|
||||
used: individual.used.clone(),
|
||||
limit: individual.limit.clone(),
|
||||
used_credits: individual.used.clone(),
|
||||
limit_credits: individual.limit.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_review_limit_name(value: &str) -> bool {
|
||||
let normalized = value.trim().to_ascii_lowercase().replace('-', "_");
|
||||
matches!(
|
||||
normalized.as_str(),
|
||||
"review" | "code_review" | "codex_review" | "codex_code_review"
|
||||
)
|
||||
}
|
||||
@@ -7,26 +7,22 @@ use std::{
|
||||
};
|
||||
|
||||
use chrono::Utc;
|
||||
use secrecy::SecretString;
|
||||
use thiserror::Error;
|
||||
use tokio::time::{Instant, sleep};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::fleet::cookies::PgCookieStore;
|
||||
use crate::fleet::{
|
||||
account::{Account, AccountStatus},
|
||||
account_failure::apply_client_failure_effect,
|
||||
account_failure::{apply_account_state_effect, classify_account_failure},
|
||||
account_gateway::{AccountGatewayError, AccountUpstreamContext, AccountUpstreamGateway},
|
||||
pool::AccountPoolService,
|
||||
quota::{
|
||||
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
|
||||
quota_snapshot_reset_at,
|
||||
quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at,
|
||||
},
|
||||
store::{AccountStore, AccountStoreError, PgAccountStore, PgAccountStoreError},
|
||||
};
|
||||
use crate::{
|
||||
fleet::cookies::PgCookieStore,
|
||||
telemetry::account_usage::store::{
|
||||
AccountUsageStore, AccountUsageStoreError, PgAccountUsageStore,
|
||||
},
|
||||
upstream::openai::transport::{CodexBackendClient, CodexClientError, CodexRequestContext},
|
||||
store::{AccountStore, AccountStoreError},
|
||||
usage::{AccountUsageStore, AccountUsageStoreError},
|
||||
};
|
||||
|
||||
const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60;
|
||||
@@ -35,9 +31,9 @@ const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30;
|
||||
|
||||
/// 运行时 quota refresh 服务。
|
||||
pub struct QuotaRefreshService {
|
||||
store: PgAccountStore,
|
||||
usage_store: PgAccountUsageStore,
|
||||
codex: Arc<CodexBackendClient>,
|
||||
store: Arc<dyn AccountStore>,
|
||||
usage_store: Arc<dyn AccountUsageStore>,
|
||||
upstream: Arc<dyn AccountUpstreamGateway>,
|
||||
min_refresh_interval_secs: u64,
|
||||
request_spacing: Duration,
|
||||
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
|
||||
@@ -48,16 +44,16 @@ pub struct QuotaRefreshService {
|
||||
impl QuotaRefreshService {
|
||||
/// 构造默认 quota refresh 服务。
|
||||
pub fn new(
|
||||
store: PgAccountStore,
|
||||
usage_store: PgAccountUsageStore,
|
||||
codex: Arc<CodexBackendClient>,
|
||||
store: Arc<dyn AccountStore>,
|
||||
usage_store: Arc<dyn AccountUsageStore>,
|
||||
upstream: Arc<dyn AccountUpstreamGateway>,
|
||||
account_pool: Arc<AccountPoolService>,
|
||||
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
store,
|
||||
usage_store,
|
||||
codex,
|
||||
upstream,
|
||||
min_refresh_interval_secs: MIN_REFRESH_INTERVAL_SECS,
|
||||
request_spacing: Self::default_request_spacing(),
|
||||
account_pseudonymizer,
|
||||
@@ -72,9 +68,9 @@ impl QuotaRefreshService {
|
||||
|
||||
/// 使用自定义最小刷新间隔构造 quota refresh 服务。
|
||||
pub fn with_min_refresh_interval_secs(
|
||||
store: PgAccountStore,
|
||||
usage_store: PgAccountUsageStore,
|
||||
codex: Arc<CodexBackendClient>,
|
||||
store: Arc<dyn AccountStore>,
|
||||
usage_store: Arc<dyn AccountUsageStore>,
|
||||
upstream: Arc<dyn AccountUpstreamGateway>,
|
||||
account_pool: Arc<AccountPoolService>,
|
||||
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
|
||||
min_refresh_interval_secs: u64,
|
||||
@@ -82,7 +78,7 @@ impl QuotaRefreshService {
|
||||
Self {
|
||||
store,
|
||||
usage_store,
|
||||
codex,
|
||||
upstream,
|
||||
min_refresh_interval_secs,
|
||||
request_spacing: Self::default_request_spacing(),
|
||||
account_pseudonymizer,
|
||||
@@ -153,40 +149,29 @@ impl QuotaRefreshService {
|
||||
}
|
||||
last_refreshed.insert(account.id.clone(), now);
|
||||
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let cookie_header = self.usage_cookie_header(&account.id).await;
|
||||
let installation_id = self.account_pseudonymizer.installation_id(&account.id);
|
||||
match self
|
||||
.codex
|
||||
.fetch_usage(CodexRequestContext {
|
||||
access_token: &account.access_token,
|
||||
account_id: account.account_id.as_deref(),
|
||||
request_id: &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(&installation_id),
|
||||
session_id: None,
|
||||
thread_id: None,
|
||||
client_request_id: None,
|
||||
turn_id: None,
|
||||
.upstream
|
||||
.fetch_usage(AccountUpstreamContext {
|
||||
access_token: SecretString::new(account.access_token.clone().into()),
|
||||
account_id: account.account_id.clone(),
|
||||
request_id: uuid::Uuid::new_v4().to_string(),
|
||||
cookie_header,
|
||||
installation_id: Some(installation_id),
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(raw) => {
|
||||
let quota = quota_from_usage(&raw);
|
||||
Ok(result) => {
|
||||
let quota = result.quota;
|
||||
let limit_reached = quota_snapshot_limit_reached("a);
|
||||
let reset_at = quota_snapshot_reset_at("a);
|
||||
let quota_json = quota.to_json()?;
|
||||
if self
|
||||
.store
|
||||
.apply_quota_snapshot(
|
||||
&account.id,
|
||||
"a.to_string(),
|
||||
"a_json,
|
||||
limit_reached,
|
||||
limit_reached.then_some(reset_at).flatten(),
|
||||
)
|
||||
@@ -208,13 +193,17 @@ impl QuotaRefreshService {
|
||||
self.sync_runtime_pool_account(&account.id).await;
|
||||
}
|
||||
Err(error) => {
|
||||
apply_client_failure_effect(
|
||||
&self.account_pool,
|
||||
&self.codex,
|
||||
&account.id,
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
if let Some(classified) = error.failure().and_then(classify_account_failure)
|
||||
&& let Some(effect) = &classified.effect
|
||||
{
|
||||
apply_account_state_effect(
|
||||
&self.account_pool,
|
||||
self.upstream.as_ref(),
|
||||
&account.id,
|
||||
effect,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
warn!(
|
||||
account_id = %account.id,
|
||||
error = %error,
|
||||
@@ -289,15 +278,15 @@ pub enum QuotaRefreshServiceError {
|
||||
/// 账号读取失败。
|
||||
#[error("failed to list accounts for quota refresh: {0}")]
|
||||
AccountStore(#[from] AccountStoreError),
|
||||
/// 账号写入失败。
|
||||
#[error("failed to persist quota refresh result: {0}")]
|
||||
Store(#[from] PgAccountStoreError),
|
||||
/// 用量窗口写入失败。
|
||||
#[error("failed to persist quota usage window: {0}")]
|
||||
UsageStore(#[from] AccountUsageStoreError),
|
||||
/// typed quota 序列化失败。
|
||||
#[error("failed to serialize quota snapshot: {0}")]
|
||||
SerializeQuota(#[from] serde_json::Error),
|
||||
/// 上游 usage 请求失败。
|
||||
#[error("failed to fetch quota usage: {0}")]
|
||||
Codex(#[from] CodexClientError),
|
||||
Upstream(#[from] AccountGatewayError),
|
||||
}
|
||||
|
||||
/// quota refresh 服务结果。
|
||||
|
||||
@@ -1,412 +1,17 @@
|
||||
//! 配额检查策略。
|
||||
|
||||
use std::time::Duration;
|
||||
//! 账号调度读取 typed quota 的稳定策略入口。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
const MONTH_WINDOW_MINUTES: u64 = 43_200;
|
||||
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
|
||||
use super::QuotaSnapshot;
|
||||
|
||||
/// 判断配额快照是否已经触顶。
|
||||
pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
|
||||
snapshot_blocked(quota)
|
||||
|| spend_control_limit_reached(quota)
|
||||
|| monthly_limit_reached(quota)
|
||||
|| credits_overage_limit_reached(quota)
|
||||
pub fn quota_snapshot_limit_reached(quota: &QuotaSnapshot) -> bool {
|
||||
quota.is_limit_reached()
|
||||
}
|
||||
|
||||
/// 从配额快照读取 UTC 重置时间。
|
||||
pub fn quota_snapshot_reset_at(quota: &Value) -> Option<DateTime<Utc>> {
|
||||
let reset_at = if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
|
||||
monthly_limit_reset_at(quota)
|
||||
.or_else(|| blocking_bucket(quota).and_then(bucket_reset_at))
|
||||
.or_else(|| core_bucket_with_reset_at(quota).and_then(bucket_reset_at))
|
||||
} else if let Some(reset_at) = blocking_bucket(quota).and_then(bucket_reset_at) {
|
||||
Some(reset_at)
|
||||
} else {
|
||||
core_bucket_with_reset_at(quota)
|
||||
.and_then(bucket_reset_at)
|
||||
.or_else(|| monthly_limit_reset_at(quota))
|
||||
}?;
|
||||
DateTime::<Utc>::from_timestamp(reset_at, 0)
|
||||
pub fn quota_snapshot_reset_at(quota: &QuotaSnapshot) -> Option<DateTime<Utc>> {
|
||||
quota.reset_at()
|
||||
}
|
||||
|
||||
/// 从配额快照读取窗口大小。
|
||||
pub fn quota_snapshot_limit_window_seconds(quota: &Value) -> Option<u64> {
|
||||
if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
|
||||
return monthly_limit_window_seconds(quota).or(Some(MONTH_WINDOW_SECONDS));
|
||||
}
|
||||
blocking_bucket(quota)
|
||||
.and_then(bucket_window_minutes)
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
.or_else(|| {
|
||||
core_bucket_with_window_minutes(quota)
|
||||
.and_then(bucket_window_minutes)
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
})
|
||||
.or_else(|| monthly_limit_window_seconds(quota))
|
||||
}
|
||||
|
||||
/// 从 Codex usage 响应中提取持久化配额快照。
|
||||
pub fn quota_from_usage(usage: &Value) -> Value {
|
||||
let mut snapshots = Vec::new();
|
||||
if let Some(snapshot) = quota_snapshot_from_rate_limit(
|
||||
"core",
|
||||
usage.get("rate_limit_name").and_then(trimmed_str),
|
||||
usage.get("metered_feature").and_then(trimmed_str),
|
||||
usage.get("rate_limit"),
|
||||
) {
|
||||
snapshots.push(snapshot);
|
||||
}
|
||||
|
||||
if let Some(additional) = usage
|
||||
.get("additional_rate_limits")
|
||||
.and_then(Value::as_array)
|
||||
{
|
||||
for item in additional {
|
||||
let Some(limit_name) = item.get("limit_name").and_then(trimmed_str) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(snapshot) = quota_snapshot_from_rate_limit(
|
||||
"additional",
|
||||
Some(limit_name),
|
||||
item.get("metered_feature").and_then(trimmed_str),
|
||||
item.get("rate_limit"),
|
||||
) {
|
||||
snapshots.push(snapshot);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let monthly_limit = monthly_limit_from_spend_control(usage.get("spend_control"))
|
||||
.or_else(|| monthly_limit_from_snapshots(&snapshots))
|
||||
.unwrap_or(Value::Null);
|
||||
|
||||
json!({
|
||||
"plan_type": usage.get("plan_type").cloned().unwrap_or(Value::Null),
|
||||
"snapshots": snapshots,
|
||||
"monthly_limit": monthly_limit,
|
||||
"credits": normalize_quota_credits(usage.get("credits")),
|
||||
"spend_control": normalize_spend_control(usage.get("spend_control")),
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_snapshot_from_rate_limit(
|
||||
source: &str,
|
||||
limit_name: Option<&str>,
|
||||
metered_feature: Option<&str>,
|
||||
rate_limit: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let rate_limit = rate_limit.filter(|value| !value.is_null())?;
|
||||
let primary = window_from_rate_limit(rate_limit.get("primary_window"));
|
||||
let secondary = window_from_rate_limit(rate_limit.get("secondary_window"));
|
||||
if primary.is_null() && secondary.is_null() {
|
||||
return None;
|
||||
}
|
||||
Some(json!({
|
||||
"source": source,
|
||||
"limit_name": limit_name,
|
||||
"metered_feature": metered_feature,
|
||||
"allowed": rate_limit.get("allowed").cloned().unwrap_or(Value::Null),
|
||||
"limit_reached": rate_limit.get("limit_reached").cloned().unwrap_or(Value::Null),
|
||||
"blocked": rate_limit_blocked(rate_limit)
|
||||
|| window_limit_reached(&primary)
|
||||
|| window_limit_reached(&secondary),
|
||||
"primary": primary,
|
||||
"secondary": secondary,
|
||||
}))
|
||||
}
|
||||
|
||||
fn window_from_rate_limit(window: Option<&Value>) -> Value {
|
||||
let Some(window) = window.filter(|value| !value.is_null()) else {
|
||||
return Value::Null;
|
||||
};
|
||||
let used_percent = window
|
||||
.get("used_percent")
|
||||
.and_then(number_value)
|
||||
.unwrap_or(0.0)
|
||||
.clamp(0.0, 100.0);
|
||||
let window_minutes = window
|
||||
.get("limit_window_seconds")
|
||||
.and_then(number_value)
|
||||
.filter(|seconds| *seconds > 0.0)
|
||||
.and_then(|seconds| Duration::try_from_secs_f64((seconds / 60.0).round()).ok())
|
||||
.map(|duration| duration.as_secs());
|
||||
json!({
|
||||
"used_percent": used_percent,
|
||||
"remaining_percent": remaining_percent(used_percent),
|
||||
"reset_at": window.get("reset_at").and_then(positive_i64),
|
||||
"window_minutes": window_minutes,
|
||||
"limit_reached": used_percent >= 100.0,
|
||||
"used": first_present(window, &["used_tokens", "used", "used_credits", "usage", "consumed_tokens", "consumed"]),
|
||||
"limit": first_present(window, &["limit_tokens", "limit", "limit_credits", "quota", "total_tokens", "total"]),
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_quota_credits(raw: Option<&Value>) -> Value {
|
||||
let Some(raw) = raw.filter(|value| !value.is_null()) else {
|
||||
return Value::Null;
|
||||
};
|
||||
json!({
|
||||
"has_credits": raw.get("has_credits").and_then(Value::as_bool).unwrap_or(false),
|
||||
"unlimited": raw.get("unlimited").and_then(Value::as_bool).unwrap_or(false),
|
||||
"overage_limit_reached": raw.get("overage_limit_reached").and_then(Value::as_bool).unwrap_or(false),
|
||||
"balance": raw.get("balance").cloned().unwrap_or(Value::Null),
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_spend_control(raw: Option<&Value>) -> Value {
|
||||
let Some(raw) = raw.filter(|value| !value.is_null()) else {
|
||||
return Value::Null;
|
||||
};
|
||||
raw.clone()
|
||||
}
|
||||
|
||||
fn monthly_limit_from_spend_control(spend_control: Option<&Value>) -> Option<Value> {
|
||||
let individual_limit = spend_control?
|
||||
.get("individual_limit")
|
||||
.filter(|value| !value.is_null())?;
|
||||
let used_percent = individual_limit
|
||||
.get("used_percent")
|
||||
.and_then(number_value)
|
||||
.unwrap_or(0.0)
|
||||
.clamp(0.0, 100.0);
|
||||
Some(json!({
|
||||
"key": "spend-control-monthly",
|
||||
"source": "spend_control",
|
||||
"used_percent": used_percent,
|
||||
"remaining_percent": individual_limit
|
||||
.get("remaining_percent")
|
||||
.and_then(number_value)
|
||||
.unwrap_or_else(|| remaining_percent(used_percent) as f64),
|
||||
"reset_at": individual_limit.get("reset_at").and_then(positive_i64),
|
||||
"window_minutes": MONTH_WINDOW_MINUTES,
|
||||
"limit_reached": spend_control_limit_reached_value(spend_control.unwrap_or(&Value::Null), used_percent),
|
||||
"used_credits": individual_limit.get("used").cloned().unwrap_or(Value::Null),
|
||||
"limit_credits": individual_limit.get("limit").cloned().unwrap_or(Value::Null),
|
||||
}))
|
||||
}
|
||||
|
||||
fn monthly_limit_from_snapshots(snapshots: &[Value]) -> Option<Value> {
|
||||
for snapshot in snapshots {
|
||||
if snapshot.get("source").and_then(Value::as_str) != Some("core") {
|
||||
continue;
|
||||
}
|
||||
for key in ["primary", "secondary"] {
|
||||
let Some(bucket) = snapshot.get(key).filter(|value| !value.is_null()) else {
|
||||
continue;
|
||||
};
|
||||
if bucket
|
||||
.get("window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.is_some_and(|minutes| window_minutes_match(minutes, MONTH_WINDOW_MINUTES))
|
||||
{
|
||||
return Some(json!({
|
||||
"key": "core-monthly",
|
||||
"source": "rate_limit",
|
||||
"used_percent": bucket.get("used_percent").cloned().unwrap_or(Value::Null),
|
||||
"remaining_percent": bucket.get("remaining_percent").cloned().unwrap_or(Value::Null),
|
||||
"reset_at": bucket.get("reset_at").cloned().unwrap_or(Value::Null),
|
||||
"window_minutes": bucket.get("window_minutes").cloned().unwrap_or(Value::Null),
|
||||
"limit_reached": bucket.get("limit_reached").cloned().unwrap_or(Value::Bool(false)),
|
||||
"used": bucket.get("used").cloned().unwrap_or(Value::Null),
|
||||
"limit": bucket.get("limit").cloned().unwrap_or(Value::Null),
|
||||
"used_credits": bucket.get("used_credits").cloned().unwrap_or(Value::Null),
|
||||
"limit_credits": bucket.get("limit_credits").cloned().unwrap_or(Value::Null),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn remaining_percent(used_percent: f64) -> i64 {
|
||||
(100.0 - used_percent.clamp(0.0, 100.0)).round() as i64
|
||||
}
|
||||
|
||||
fn first_present(value: &Value, keys: &[&str]) -> Value {
|
||||
keys.iter()
|
||||
.find_map(|key| value.get(*key).filter(|value| !value.is_null()).cloned())
|
||||
.unwrap_or(Value::Null)
|
||||
}
|
||||
|
||||
fn rate_limit_blocked(rate_limit: &Value) -> bool {
|
||||
rate_limit
|
||||
.get("limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| rate_limit
|
||||
.get("allowed")
|
||||
.and_then(Value::as_bool)
|
||||
.is_some_and(|allowed| !allowed)
|
||||
}
|
||||
|
||||
fn window_limit_reached(window: &Value) -> bool {
|
||||
window
|
||||
.get("limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn spend_control_limit_reached(quota: &Value) -> bool {
|
||||
quota
|
||||
.pointer("/spend_control/reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| quota
|
||||
.pointer("/spend_control/individual_limit/used_percent")
|
||||
.and_then(number_value)
|
||||
.is_some_and(|used_percent| used_percent >= 100.0)
|
||||
}
|
||||
|
||||
fn spend_control_limit_reached_value(spend_control: &Value, used_percent: f64) -> bool {
|
||||
spend_control
|
||||
.get("reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| used_percent >= 100.0
|
||||
}
|
||||
|
||||
fn monthly_limit_reached(quota: &Value) -> bool {
|
||||
quota
|
||||
.pointer("/monthly_limit/limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| quota
|
||||
.pointer("/monthly_limit/used_percent")
|
||||
.and_then(number_value)
|
||||
.is_some_and(|used_percent| used_percent >= 100.0)
|
||||
}
|
||||
|
||||
fn quota_bucket_limit_reached(bucket: &Value) -> bool {
|
||||
!bucket.is_null()
|
||||
&& (bucket
|
||||
.get("limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| bucket
|
||||
.get("used_percent")
|
||||
.and_then(number_value)
|
||||
.is_some_and(|used_percent| used_percent >= 100.0))
|
||||
}
|
||||
|
||||
fn quota_snapshot_blocked(snapshot: &Value) -> bool {
|
||||
snapshot
|
||||
.get("blocked")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| snapshot
|
||||
.get("allowed")
|
||||
.and_then(Value::as_bool)
|
||||
.is_some_and(|allowed| !allowed)
|
||||
|| snapshot_buckets(snapshot).any(quota_bucket_limit_reached)
|
||||
}
|
||||
|
||||
fn snapshot_blocked(quota: &Value) -> bool {
|
||||
quota
|
||||
.get("snapshots")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.any(quota_snapshot_blocked)
|
||||
}
|
||||
|
||||
fn blocking_bucket(quota: &Value) -> Option<&Value> {
|
||||
let snapshots = quota.get("snapshots").and_then(Value::as_array)?;
|
||||
snapshots
|
||||
.iter()
|
||||
.filter(|snapshot| quota_snapshot_blocked(snapshot))
|
||||
.flat_map(snapshot_buckets)
|
||||
.find(|bucket| quota_bucket_limit_reached(bucket))
|
||||
.or_else(|| {
|
||||
snapshots
|
||||
.iter()
|
||||
.filter(|snapshot| quota_snapshot_blocked(snapshot))
|
||||
.flat_map(snapshot_buckets)
|
||||
.find(|bucket| bucket_reset_at(bucket).is_some())
|
||||
})
|
||||
}
|
||||
|
||||
fn snapshot_buckets(snapshot: &Value) -> impl Iterator<Item = &Value> {
|
||||
["primary", "secondary"]
|
||||
.into_iter()
|
||||
.filter_map(|key| snapshot.get(key))
|
||||
.filter(|bucket| !bucket.is_null())
|
||||
}
|
||||
|
||||
fn core_bucket_with_reset_at(quota: &Value) -> Option<&Value> {
|
||||
core_buckets(quota).find(|bucket| bucket_reset_at(bucket).is_some())
|
||||
}
|
||||
|
||||
fn core_bucket_with_window_minutes(quota: &Value) -> Option<&Value> {
|
||||
core_buckets(quota).find(|bucket| bucket_window_minutes(bucket).is_some())
|
||||
}
|
||||
|
||||
fn core_buckets<'a>(quota: &'a Value) -> impl Iterator<Item = &'a Value> + 'a {
|
||||
quota
|
||||
.get("snapshots")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
|
||||
.flat_map(snapshot_buckets)
|
||||
}
|
||||
|
||||
fn bucket_reset_at(bucket: &Value) -> Option<i64> {
|
||||
bucket.get("reset_at").and_then(positive_i64)
|
||||
}
|
||||
|
||||
fn bucket_window_minutes(bucket: &Value) -> Option<u64> {
|
||||
bucket
|
||||
.get("window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.filter(|value| *value > 0)
|
||||
}
|
||||
|
||||
fn credits_overage_limit_reached(quota: &Value) -> bool {
|
||||
quota
|
||||
.pointer("/credits/overage_limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn monthly_limit_reset_at(quota: &Value) -> Option<i64> {
|
||||
quota
|
||||
.pointer("/monthly_limit/reset_at")
|
||||
.and_then(positive_i64)
|
||||
}
|
||||
|
||||
fn monthly_limit_window_seconds(quota: &Value) -> Option<u64> {
|
||||
quota
|
||||
.pointer("/monthly_limit/window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.filter(|value| *value > 0)
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
}
|
||||
|
||||
fn number_value(value: &Value) -> Option<f64> {
|
||||
value
|
||||
.as_f64()
|
||||
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn positive_i64(value: &Value) -> Option<i64> {
|
||||
value
|
||||
.as_i64()
|
||||
.filter(|value| *value > 0)
|
||||
.or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok()))
|
||||
.filter(|value| *value > 0)
|
||||
}
|
||||
|
||||
fn trimmed_str(value: &Value) -> Option<&str> {
|
||||
value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn window_minutes_match(actual: u64, expected: u64) -> bool {
|
||||
actual > 0 && actual.abs_diff(expected) <= expected / 20
|
||||
pub fn quota_snapshot_limit_window_seconds(quota: &QuotaSnapshot) -> Option<u64> {
|
||||
quota.limit_window_seconds()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
//! token 刷新用例拥有的上游端口与稳定值。
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TokenPair {
|
||||
pub access_token: String,
|
||||
pub refresh_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum RefreshFailure {
|
||||
#[error("refresh token is invalid or expired")]
|
||||
InvalidGrant,
|
||||
#[error("account is banned")]
|
||||
Banned,
|
||||
#[error("refresh transport failed before server processing")]
|
||||
RetryableTransport,
|
||||
#[error("refresh transport failed after possible server processing")]
|
||||
Transport,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait TokenRefresher: Send + Sync + 'static {
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure>;
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
//! Token 刷新策略、租约与运行时服务。
|
||||
|
||||
mod gateway;
|
||||
mod lease;
|
||||
mod policy;
|
||||
mod service;
|
||||
|
||||
pub use gateway::*;
|
||||
pub use lease::*;
|
||||
pub use policy::*;
|
||||
pub use service::*;
|
||||
|
||||
@@ -12,8 +12,8 @@ use chrono::{DateTime, Duration, TimeZone, Utc};
|
||||
use serde_json::{Map, Value};
|
||||
use tokio::sync::{Semaphore, watch};
|
||||
|
||||
use super::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
use crate::fleet::account::{Account, AccountStatus};
|
||||
use crate::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
|
||||
pub(crate) const PERMANENT_FAILURE_CONFIRMATION_THRESHOLD: usize = 2;
|
||||
pub(crate) const RECOVERY_DELAY_SECONDS: i64 = 10 * 60;
|
||||
@@ -238,16 +238,9 @@ impl RuntimeRefreshPolicy {
|
||||
}
|
||||
|
||||
/// 持续接收运行时设置并更新刷新策略。
|
||||
pub async fn subscribe_settings(
|
||||
self,
|
||||
mut receiver: watch::Receiver<crate::settings::SettingsSnapshot>,
|
||||
) {
|
||||
pub async fn subscribe_policy(self, mut receiver: watch::Receiver<RefreshPolicy>) {
|
||||
while receiver.changed().await.is_ok() {
|
||||
let settings = receiver.borrow_and_update();
|
||||
self.update(RefreshPolicy {
|
||||
refresh_margin_seconds: settings.refresh_margin_seconds,
|
||||
refresh_concurrency: settings.refresh_concurrency,
|
||||
});
|
||||
self.update(*receiver.borrow_and_update());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -19,13 +19,14 @@ use tokio_util::task::TaskTracker;
|
||||
use tracing::{debug, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::TokenRefresher;
|
||||
use crate::fleet::{
|
||||
account::{Account, AccountStatus},
|
||||
pool::AccountPoolService,
|
||||
store::{
|
||||
AccountClaimsUpdate, AccountStore, PgAccountStore, PgAccountStoreError, StoredAccount,
|
||||
},
|
||||
};
|
||||
use crate::upstream::openai::token_client::TokenRefresher;
|
||||
|
||||
use super::{
|
||||
lease::{RedisRefreshLeaseStore, RedisRefreshLeaseStoreError},
|
||||
@@ -48,6 +49,7 @@ where
|
||||
C: TokenRefresher,
|
||||
{
|
||||
store: PgAccountStore,
|
||||
account_pool: Option<Arc<AccountPoolService>>,
|
||||
scheduler: Arc<RefreshScheduler<C>>,
|
||||
policy: RuntimeRefreshPolicy,
|
||||
refresh_leases: Option<RedisRefreshLeaseStore>,
|
||||
@@ -193,6 +195,7 @@ where
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
store: self.store.clone(),
|
||||
account_pool: self.account_pool.clone(),
|
||||
scheduler: self.scheduler.clone(),
|
||||
policy: self.policy.clone(),
|
||||
refresh_leases: self.refresh_leases.clone(),
|
||||
@@ -215,6 +218,7 @@ where
|
||||
let policy = policy.into();
|
||||
Self {
|
||||
store,
|
||||
account_pool: None,
|
||||
scheduler: Arc::new(RefreshScheduler::new(policy.clone(), client)),
|
||||
policy,
|
||||
refresh_leases: None,
|
||||
@@ -233,6 +237,12 @@ where
|
||||
self
|
||||
}
|
||||
|
||||
/// 刷新后同步持久化账号到运行时池,避免旧 token 快照覆盖新状态。
|
||||
pub fn with_account_pool_sync(mut self, account_pool: Arc<AccountPoolService>) -> Self {
|
||||
self.account_pool = Some(account_pool);
|
||||
self
|
||||
}
|
||||
|
||||
/// 更新刷新策略。
|
||||
pub fn update_policy(&self, policy: RefreshPolicy) {
|
||||
self.policy.update(policy);
|
||||
@@ -587,6 +597,7 @@ where
|
||||
if attempt.refresh_token.is_none() {
|
||||
persist_status(&self.store, &attempt.id, AccountStatus::Expired).await?;
|
||||
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
|
||||
self.sync_account_pool_after_update(&attempt.id).await;
|
||||
return Ok(TokenRefreshOutcome::StatusUpdated);
|
||||
}
|
||||
|
||||
@@ -601,6 +612,7 @@ where
|
||||
|
||||
persist_status(&self.store, &attempt.id, updated.status).await?;
|
||||
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
|
||||
self.sync_account_pool_after_update(&attempt.id).await;
|
||||
return Ok(TokenRefreshOutcome::StatusUpdated);
|
||||
}
|
||||
Ok(mut updated) => {
|
||||
@@ -609,6 +621,7 @@ where
|
||||
updated.next_refresh_at = self.computed_refresh_at(&updated);
|
||||
let refresh_token_rotated = updated.refresh_token != attempt.refresh_token;
|
||||
persist_token_update(&self.store, &updated).await?;
|
||||
self.sync_account_pool_after_update(&updated.id).await;
|
||||
info!(
|
||||
account_id = %updated.id,
|
||||
attempt = attempt_index + 1,
|
||||
@@ -643,6 +656,19 @@ where
|
||||
Ok(TokenRefreshOutcome::Skipped)
|
||||
}
|
||||
|
||||
async fn sync_account_pool_after_update(&self, account_id: &str) {
|
||||
let Some(account_pool) = self.account_pool.as_ref() else {
|
||||
return;
|
||||
};
|
||||
if let Err(error) = account_pool.sync_account_from_store(account_id).await {
|
||||
warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"Failed to sync runtime account pool after token refresh"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn prepare_refresh_attempt(
|
||||
&self,
|
||||
account: &Account,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! 账号运行时反馈:错误率与 TTFT 的 EWMA 统计。
|
||||
//! 账号运行时 feedback:首字前失败率、首字后中断率与 TTFT 的 EWMA 统计。
|
||||
//!
|
||||
//! 每个账号维护两个指数加权移动平均(EWMA):错误率与首个有效输出事件延迟(TTFT,
|
||||
//! 毫秒)。请求结束后经 [`FeedbackStats::report`] 回灌,直接进 [`super::score`]
|
||||
//! 每个账号维护三个指数加权移动平均(EWMA)。每个已启动 attempt 结束后经
|
||||
//! [`FeedbackStats::report_attempt`] 回灌,直接进 [`super::score`]
|
||||
//! 打分并可驱动 sticky escape。统计以 `AtomicU64`(存 `f64` 位模式)+ CAS 无锁更新,
|
||||
//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。
|
||||
|
||||
@@ -12,13 +12,32 @@ use std::sync::atomic::{AtomicU64, Ordering};
|
||||
/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。
|
||||
const EWMA_ALPHA: f64 = 0.2;
|
||||
|
||||
/// 单个上游 attempt 对账号调度器的反馈。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AttemptFeedback {
|
||||
Completed { first_token_ms: Option<u64> },
|
||||
Incomplete { first_token_ms: Option<u64> },
|
||||
FailedBeforeFirstToken,
|
||||
FailedAfterFirstToken { first_token_ms: u64 },
|
||||
Cancelled,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
impl AttemptFeedback {
|
||||
fn affects_account_health(self) -> bool {
|
||||
!matches!(self, Self::Cancelled | Self::Shutdown)
|
||||
}
|
||||
}
|
||||
|
||||
/// 单账号的运行时 EWMA 统计。
|
||||
///
|
||||
/// `*_bits` 存 `f64` 的位模式;`NaN`(全 1 的哨兵)表示"尚无样本"。
|
||||
#[derive(Debug)]
|
||||
struct AccountFeedback {
|
||||
/// 错误率 EWMA,范围 `[0, 1]`。
|
||||
error_rate_bits: AtomicU64,
|
||||
/// 首字前失败率 EWMA,范围 `[0, 1]`。
|
||||
pre_token_failure_rate_bits: AtomicU64,
|
||||
/// 首字后中断率 EWMA,范围 `[0, 1]`。
|
||||
post_token_abort_rate_bits: AtomicU64,
|
||||
/// TTFT EWMA,单位毫秒。
|
||||
ttft_ms_bits: AtomicU64,
|
||||
}
|
||||
@@ -29,16 +48,22 @@ const EMPTY_BITS: u64 = f64::NAN.to_bits();
|
||||
impl Default for AccountFeedback {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
error_rate_bits: AtomicU64::new(EMPTY_BITS),
|
||||
pre_token_failure_rate_bits: AtomicU64::new(EMPTY_BITS),
|
||||
post_token_abort_rate_bits: AtomicU64::new(EMPTY_BITS),
|
||||
ttft_ms_bits: AtomicU64::new(EMPTY_BITS),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AccountFeedback {
|
||||
/// 读取当前错误率 EWMA;无样本时返回 `None`。
|
||||
fn error_rate(&self) -> Option<f64> {
|
||||
load_sample(&self.error_rate_bits)
|
||||
/// 读取当前首字前失败率 EWMA;无样本时返回 `None`。
|
||||
fn pre_token_failure_rate(&self) -> Option<f64> {
|
||||
load_sample(&self.pre_token_failure_rate_bits)
|
||||
}
|
||||
|
||||
/// 读取当前首字后中断率 EWMA;无样本时返回 `None`。
|
||||
fn post_token_abort_rate(&self) -> Option<f64> {
|
||||
load_sample(&self.post_token_abort_rate_bits)
|
||||
}
|
||||
|
||||
/// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。
|
||||
@@ -46,12 +71,26 @@ impl AccountFeedback {
|
||||
load_sample(&self.ttft_ms_bits)
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果。
|
||||
fn report(&self, success: bool, first_token_ms: Option<u64>) {
|
||||
let error_sample = if success { 0.0 } else { 1.0 };
|
||||
update_ewma(&self.error_rate_bits, error_sample);
|
||||
if let Some(ms) = first_token_ms {
|
||||
update_ewma(&self.ttft_ms_bits, ms as f64);
|
||||
/// 回灌一次已经开始的 attempt。
|
||||
fn report_attempt(&self, feedback: AttemptFeedback) {
|
||||
match feedback {
|
||||
AttemptFeedback::Completed { first_token_ms }
|
||||
| AttemptFeedback::Incomplete { first_token_ms } => {
|
||||
update_ewma(&self.pre_token_failure_rate_bits, 0.0);
|
||||
update_ewma(&self.post_token_abort_rate_bits, 0.0);
|
||||
if let Some(ms) = first_token_ms {
|
||||
update_ewma(&self.ttft_ms_bits, ms as f64);
|
||||
}
|
||||
}
|
||||
AttemptFeedback::FailedBeforeFirstToken => {
|
||||
update_ewma(&self.pre_token_failure_rate_bits, 1.0);
|
||||
}
|
||||
AttemptFeedback::FailedAfterFirstToken { first_token_ms } => {
|
||||
update_ewma(&self.pre_token_failure_rate_bits, 0.0);
|
||||
update_ewma(&self.post_token_abort_rate_bits, 1.0);
|
||||
update_ewma(&self.ttft_ms_bits, first_token_ms as f64);
|
||||
}
|
||||
AttemptFeedback::Cancelled | AttemptFeedback::Shutdown => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -87,8 +126,10 @@ fn update_ewma(bits: &AtomicU64, sample: f64) {
|
||||
/// 账号 EWMA 反馈的读数快照,供打分使用。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq)]
|
||||
pub struct FeedbackSample {
|
||||
/// 错误率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
|
||||
pub error_rate: Option<f64>,
|
||||
/// 首字前失败率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
|
||||
pub pre_token_failure_rate: Option<f64>,
|
||||
/// 首字后中断率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
|
||||
pub post_token_abort_rate: Option<f64>,
|
||||
/// TTFT EWMA(毫秒);无样本时为 `None`。
|
||||
pub ttft_ms: Option<f64>,
|
||||
}
|
||||
@@ -116,23 +157,27 @@ impl FeedbackStats {
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
match accounts.get(account_id) {
|
||||
Some(feedback) => FeedbackSample {
|
||||
error_rate: feedback.error_rate(),
|
||||
pre_token_failure_rate: feedback.pre_token_failure_rate(),
|
||||
post_token_abort_rate: feedback.post_token_abort_rate(),
|
||||
ttft_ms: feedback.ttft_ms(),
|
||||
},
|
||||
None => FeedbackSample::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果:成功/失败 + 首个有效输出事件延迟(毫秒)。
|
||||
pub fn report(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
|
||||
/// 回灌一个已启动上游 attempt 的结果。
|
||||
pub fn report_attempt(&self, account_id: &str, sample: AttemptFeedback) {
|
||||
if !sample.affects_account_health() {
|
||||
return;
|
||||
}
|
||||
// 快路径:账号已存在时只需读锁 + 原子更新。
|
||||
{
|
||||
let accounts = self
|
||||
.accounts
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(feedback) = accounts.get(account_id) {
|
||||
feedback.report(success, first_token_ms);
|
||||
if let Some(account_feedback) = accounts.get(account_id) {
|
||||
account_feedback.report_attempt(sample);
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -144,7 +189,7 @@ impl FeedbackStats {
|
||||
accounts
|
||||
.entry(account_id.to_string())
|
||||
.or_default()
|
||||
.report(success, first_token_ms);
|
||||
.report_attempt(sample);
|
||||
}
|
||||
|
||||
/// 移除指定账号的反馈记录(账号被删除时调用)。
|
||||
|
||||
@@ -23,7 +23,7 @@ use crate::fleet::account::Account;
|
||||
use crate::infra::time::parse_rfc3339_utc;
|
||||
|
||||
pub use candidates::{CandidateFilter, CandidateRequest};
|
||||
pub use feedback::{FeedbackSample, FeedbackStats};
|
||||
pub use feedback::{AttemptFeedback, FeedbackSample, FeedbackStats};
|
||||
pub use strategy::smart::rank_candidates;
|
||||
pub use strategy::{ScoreBreakdown, ScoreWeights};
|
||||
|
||||
@@ -69,7 +69,7 @@ pub struct SelectionInput<'a> {
|
||||
/// 账号调度器:账号选择的对外门面。
|
||||
///
|
||||
/// 持有轮转游标(并列破并列用)、打分权重(Smart 用)与 EWMA 反馈存储。选择热路径
|
||||
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_feedback`] 回灌。游标用
|
||||
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_attempt`] 回灌。游标用
|
||||
/// [`Mutex`] 保护(选择需可变,但 `AccountScheduler` 以共享引用参与并发请求)。
|
||||
#[derive(Debug)]
|
||||
pub struct AccountScheduler {
|
||||
@@ -150,9 +150,9 @@ impl AccountScheduler {
|
||||
}
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果到运行时 EWMA 反馈(错误率 / TTFT),供 Smart 打分。
|
||||
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
|
||||
self.feedback.report(account_id, success, first_token_ms);
|
||||
/// 回灌一个已启动上游 attempt 的运行时反馈,供 Smart 打分。
|
||||
pub fn report_attempt(&self, account_id: &str, feedback: AttemptFeedback) {
|
||||
self.feedback.report_attempt(account_id, feedback);
|
||||
}
|
||||
|
||||
/// 移除指定账号的运行时反馈(账号被删除时调用)。
|
||||
|
||||
@@ -284,7 +284,7 @@ fn score_account(
|
||||
);
|
||||
|
||||
// 错误率 EWMA:本身即 [0,1],越低越健康。无样本时视为 0(满分健康)。
|
||||
let error_rate = 1.0 - sample.error_rate.unwrap_or(0.0).clamp(0.0, 1.0);
|
||||
let error_rate = 1.0 - sample.pre_token_failure_rate.unwrap_or(0.0).clamp(0.0, 1.0);
|
||||
|
||||
// TTFT:延迟越低越好。无样本的账号视为中性满分,避免冷启动被压。
|
||||
let ttft = match sample.ttft_ms {
|
||||
|
||||
@@ -174,6 +174,7 @@ update accounts
|
||||
set
|
||||
status = case
|
||||
when status in ('disabled', 'banned') then status
|
||||
when $4 = 'expired' and access_token_expires_at > $2 then status
|
||||
when (
|
||||
case
|
||||
when not $1 and quota_cooldown_until is not null and quota_cooldown_until > $2 then quota_limit_reached
|
||||
|
||||
@@ -7,7 +7,7 @@ use sqlx::{PgPool, Postgres, QueryBuilder, Row, postgres::PgRow};
|
||||
use crate::{
|
||||
fleet::{
|
||||
account::{Account, AccountStatus},
|
||||
quota::{quota_snapshot_limit_window_seconds, quota_snapshot_reset_at},
|
||||
quota::{QuotaSnapshot, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at},
|
||||
},
|
||||
infra::format::nonnegative_i64_to_u64,
|
||||
};
|
||||
@@ -45,13 +45,11 @@ pub(super) async fn get_pool_account(
|
||||
}
|
||||
|
||||
pub(super) fn pool_account_from_row(row: &PgRow) -> PgAccountStoreResult<Account> {
|
||||
let quota_json = row
|
||||
let quota = row
|
||||
.get::<Option<sqlx::types::Json<Value>>, _>("quota_json")
|
||||
.map(|quota_json| quota_json.0);
|
||||
let quota_window_reset_at = quota_json.as_ref().and_then(quota_snapshot_reset_at);
|
||||
let quota_limit_window_seconds = quota_json
|
||||
.as_ref()
|
||||
.and_then(quota_snapshot_limit_window_seconds);
|
||||
.and_then(|quota_json| QuotaSnapshot::from_value(quota_json.0).ok());
|
||||
let quota_window_reset_at = quota.as_ref().and_then(quota_snapshot_reset_at);
|
||||
let quota_limit_window_seconds = quota.as_ref().and_then(quota_snapshot_limit_window_seconds);
|
||||
|
||||
Ok(Account {
|
||||
id: row.get("id"),
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
//! 账号调度消费的用量值与持久化端口。
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use thiserror::Error;
|
||||
|
||||
/// 账号用量端口错误。
|
||||
#[derive(Debug, Error)]
|
||||
#[error("account usage store operation failed: {message}")]
|
||||
pub struct AccountUsageStoreError {
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl AccountUsageStoreError {
|
||||
/// 将 adapter 错误收窄为消费方稳定错误。
|
||||
pub fn adapter(error: impl std::fmt::Display) -> Self {
|
||||
Self {
|
||||
message: error.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageDelta {
|
||||
pub requests: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub reasoning_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub empty_responses: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
pub image_requests: u64,
|
||||
pub image_request_failures: u64,
|
||||
}
|
||||
|
||||
/// fleet 记录一次 Responses 完成结果所需的稳定用量事实。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct ResponseUsage {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub reasoning_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageSnapshot {
|
||||
pub request_count: u64,
|
||||
pub empty_response_count: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
pub image_request_count: u64,
|
||||
pub image_request_failed_count: u64,
|
||||
pub window_request_count: u64,
|
||||
pub window_input_tokens: u64,
|
||||
pub window_output_tokens: u64,
|
||||
pub window_cached_tokens: u64,
|
||||
pub window_image_input_tokens: u64,
|
||||
pub window_image_output_tokens: u64,
|
||||
pub window_image_request_count: u64,
|
||||
pub window_image_request_failed_count: u64,
|
||||
pub window_started_at: Option<DateTime<Utc>>,
|
||||
pub window_reset_at: Option<DateTime<Utc>>,
|
||||
pub limit_window_seconds: Option<u64>,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageWindow {
|
||||
pub request_count: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
pub image_request_count: u64,
|
||||
pub image_request_failed_count: u64,
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
pub reset_at: Option<DateTime<Utc>>,
|
||||
pub limit_window_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AccountUsageStore: Send + Sync + 'static {
|
||||
async fn snapshots(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
) -> Result<HashMap<String, AccountUsageSnapshot>, AccountUsageStoreError>;
|
||||
|
||||
async fn record_usage_delta(
|
||||
&self,
|
||||
account_id: &str,
|
||||
usage: AccountUsageDelta,
|
||||
) -> Result<(), AccountUsageStoreError>;
|
||||
|
||||
async fn sync_runtime_window(
|
||||
&self,
|
||||
account_id: &str,
|
||||
window: AccountUsageWindow,
|
||||
) -> Result<(), AccountUsageStoreError>;
|
||||
|
||||
async fn sync_rate_limit_window(
|
||||
&self,
|
||||
account_id: &str,
|
||||
reset_at: DateTime<Utc>,
|
||||
limit_window_seconds: Option<u64>,
|
||||
) -> Result<(), AccountUsageStoreError>;
|
||||
|
||||
async fn record_request(&self, account_id: &str) -> Result<(), AccountUsageStoreError> {
|
||||
self.record_usage_delta(
|
||||
account_id,
|
||||
AccountUsageDelta {
|
||||
requests: 1,
|
||||
..AccountUsageDelta::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,32 @@
|
||||
//! 账号配额窗口规则。
|
||||
|
||||
pub(crate) use crate::telemetry::account_usage::should_reset_usage_window;
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
/// 根据额度窗口边界漂移判断是否重置本地窗口计数。
|
||||
pub(crate) fn should_reset_usage_window(
|
||||
existing_reset_at: Option<DateTime<Utc>>,
|
||||
existing_limit_window_seconds: Option<u64>,
|
||||
new_reset_at: DateTime<Utc>,
|
||||
new_limit_window_seconds: Option<u64>,
|
||||
) -> bool {
|
||||
let Some(existing_reset_at) = existing_reset_at else {
|
||||
return false;
|
||||
};
|
||||
if existing_reset_at == new_reset_at {
|
||||
return false;
|
||||
}
|
||||
|
||||
let drift = existing_reset_at
|
||||
.signed_duration_since(new_reset_at)
|
||||
.num_seconds()
|
||||
.unsigned_abs();
|
||||
let window_seconds = new_limit_window_seconds
|
||||
.or(existing_limit_window_seconds)
|
||||
.unwrap_or(0);
|
||||
let threshold = if window_seconds > 0 {
|
||||
window_seconds / 2
|
||||
} else {
|
||||
3_600
|
||||
};
|
||||
drift >= threshold
|
||||
}
|
||||
|
||||
@@ -8,11 +8,18 @@ use tracing::info;
|
||||
|
||||
const MIGRATION_LOCK_KEY: i64 = i64::from_be_bytes(*b"CPR-MIGR");
|
||||
|
||||
const MIGRATIONS: &[Migration] = &[Migration {
|
||||
version: 1,
|
||||
name: "initial",
|
||||
sql: include_str!("migrations/0001_initial.sql"),
|
||||
}];
|
||||
const MIGRATIONS: &[Migration] = &[
|
||||
Migration {
|
||||
version: 1,
|
||||
name: "initial",
|
||||
sql: include_str!("migrations/0001_initial.sql"),
|
||||
},
|
||||
Migration {
|
||||
version: 2,
|
||||
name: "cache_write_tokens",
|
||||
sql: include_str!("migrations/0002_cache_write_tokens.sql"),
|
||||
},
|
||||
];
|
||||
|
||||
struct Migration {
|
||||
version: i64,
|
||||
|
||||
@@ -231,35 +231,3 @@ create index idx_account_cookies_account_domain
|
||||
on account_cookies(account_id, domain);
|
||||
create index idx_account_cookies_expires
|
||||
on account_cookies(expires_at) where expires_at is not null;
|
||||
|
||||
create table fingerprints (
|
||||
id text primary key,
|
||||
originator text not null,
|
||||
app_version text not null,
|
||||
build_number text not null,
|
||||
platform text not null,
|
||||
arch text not null,
|
||||
chromium_version text not null,
|
||||
user_agent_template text not null,
|
||||
default_headers_json jsonb not null,
|
||||
header_order_json jsonb not null,
|
||||
source text not null,
|
||||
created_at timestamptz not null,
|
||||
updated_at timestamptz not null
|
||||
);
|
||||
|
||||
create table fingerprint_update_history (
|
||||
id text primary key,
|
||||
current_fingerprint_id text not null references fingerprints(id) on delete cascade,
|
||||
app_version text not null,
|
||||
build_number text not null,
|
||||
chromium_version text,
|
||||
source text not null,
|
||||
manifest_json jsonb,
|
||||
created_at timestamptz not null
|
||||
);
|
||||
|
||||
create index idx_fingerprint_update_history_created_id
|
||||
on fingerprint_update_history(created_at desc, id desc);
|
||||
create index idx_fingerprint_update_history_fingerprint
|
||||
on fingerprint_update_history(current_fingerprint_id);
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
alter table usage_records
|
||||
add column cache_write_tokens bigint
|
||||
check (cache_write_tokens is null or cache_write_tokens >= 0);
|
||||
|
||||
alter table request_time_buckets
|
||||
add column cache_write_tokens bigint not null default 0
|
||||
check (cache_write_tokens >= 0);
|
||||
@@ -0,0 +1,36 @@
|
||||
//! 模型目录用例拥有的上游端口。
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::BackendModelEntry;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ModelCatalogRequest<'a> {
|
||||
pub access_token: &'a str,
|
||||
pub account_id: Option<&'a str>,
|
||||
pub request_id: &'a str,
|
||||
pub installation_id: Option<&'a str>,
|
||||
pub plan_type: &'a str,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("model catalog request failed: {message}")]
|
||||
pub struct ModelCatalogSourceError {
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl ModelCatalogSourceError {
|
||||
pub fn new(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
message: message.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ModelCatalogSource: Send + Sync + 'static {
|
||||
async fn fetch_models(
|
||||
&self,
|
||||
request: &ModelCatalogRequest<'_>,
|
||||
) -> Result<Vec<BackendModelEntry>, ModelCatalogSourceError>;
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
//! 模型目录领域。
|
||||
|
||||
pub mod catalog;
|
||||
pub mod gateway;
|
||||
pub mod service;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
|
||||
@@ -8,13 +8,9 @@ use std::{
|
||||
use thiserror::Error;
|
||||
use tokio::sync::watch;
|
||||
|
||||
use crate::{
|
||||
settings::SettingsSnapshot,
|
||||
upstream::openai::transport::{CodexModelCatalogClient, CodexModelCatalogRequest},
|
||||
};
|
||||
|
||||
use super::{
|
||||
catalog::ModelCatalog,
|
||||
gateway::{ModelCatalogRequest, ModelCatalogSource},
|
||||
store::{ModelSnapshotStore, ModelSnapshotStoreError},
|
||||
types::{ModelConfig, ModelPlanSnapshot},
|
||||
};
|
||||
@@ -91,7 +87,7 @@ pub struct ModelService {
|
||||
snapshots: Arc<RwLock<Vec<ModelPlanSnapshot>>>,
|
||||
catalog: Arc<RwLock<ModelCatalog>>,
|
||||
store: Option<Arc<dyn ModelSnapshotStore>>,
|
||||
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
|
||||
upstream_client: Option<Arc<dyn ModelCatalogSource>>,
|
||||
models_etag: Arc<RwLock<Option<String>>>,
|
||||
etag_change_tx: watch::Sender<u64>,
|
||||
refresh_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
@@ -102,7 +98,7 @@ impl ModelService {
|
||||
pub fn new(
|
||||
config: ModelConfig,
|
||||
store: Option<Arc<dyn ModelSnapshotStore>>,
|
||||
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
|
||||
upstream_client: Option<Arc<dyn ModelCatalogSource>>,
|
||||
) -> Self {
|
||||
let catalog = ModelCatalog::from_config(&config);
|
||||
let (etag_change_tx, _) = watch::channel(0);
|
||||
@@ -153,15 +149,9 @@ impl ModelService {
|
||||
}
|
||||
|
||||
/// 持续接收运行时设置并更新模型别名。
|
||||
pub async fn subscribe_settings(
|
||||
self: Arc<Self>,
|
||||
mut receiver: watch::Receiver<SettingsSnapshot>,
|
||||
) {
|
||||
pub async fn subscribe_config(self: Arc<Self>, mut receiver: watch::Receiver<ModelConfig>) {
|
||||
while receiver.changed().await.is_ok() {
|
||||
let settings = receiver.borrow_and_update().clone();
|
||||
self.update_config(ModelConfig {
|
||||
model_aliases: settings.model_aliases,
|
||||
});
|
||||
self.update_config(receiver.borrow_and_update().clone());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -189,6 +179,26 @@ impl ModelService {
|
||||
&self,
|
||||
plan_accounts: &[ModelRefreshPlanAccount],
|
||||
request_id: &str,
|
||||
) -> Result<ModelRefreshResult, ModelServiceError> {
|
||||
self.refresh_backend_models_inner(plan_accounts, request_id, false)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 刷新指定计划,并保留未参与本次刷新的计划快照。
|
||||
pub async fn refresh_selected_plan_models(
|
||||
&self,
|
||||
plan_accounts: &[ModelRefreshPlanAccount],
|
||||
request_id: &str,
|
||||
) -> Result<ModelRefreshResult, ModelServiceError> {
|
||||
self.refresh_backend_models_inner(plan_accounts, request_id, true)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn refresh_backend_models_inner(
|
||||
&self,
|
||||
plan_accounts: &[ModelRefreshPlanAccount],
|
||||
request_id: &str,
|
||||
preserve_other_plans: bool,
|
||||
) -> Result<ModelRefreshResult, ModelServiceError> {
|
||||
let _refresh_guard = self.refresh_lock.lock().await;
|
||||
let store = self
|
||||
@@ -211,7 +221,11 @@ impl ModelService {
|
||||
.into_iter()
|
||||
.map(|snapshot| (snapshot.plan_type.clone(), snapshot))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mut next_snapshots = BTreeMap::new();
|
||||
let mut next_snapshots = if preserve_other_plans {
|
||||
previous_snapshots.clone()
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
};
|
||||
|
||||
let mut result = ModelRefreshResult {
|
||||
refreshed_plans: 0,
|
||||
@@ -220,7 +234,7 @@ impl ModelService {
|
||||
};
|
||||
|
||||
for plan_account in plan_accounts {
|
||||
let request = CodexModelCatalogRequest {
|
||||
let request = ModelCatalogRequest {
|
||||
access_token: &plan_account.access_token,
|
||||
account_id: plan_account.account_id.as_deref(),
|
||||
request_id,
|
||||
@@ -248,7 +262,7 @@ impl ModelService {
|
||||
};
|
||||
|
||||
let snapshot =
|
||||
ModelPlanSnapshot::from_backend_values(plan_account.plan_type.clone(), entries);
|
||||
ModelPlanSnapshot::from_backend_entries(plan_account.plan_type.clone(), entries);
|
||||
if snapshot.models.is_empty() {
|
||||
result.failed_plans += 1;
|
||||
if let Some(snapshot) = previous_snapshots.get(&plan_account.plan_type) {
|
||||
|
||||
@@ -1,35 +1,4 @@
|
||||
//! 账号累计用量与时间桶查询。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
pub mod query;
|
||||
pub mod store;
|
||||
|
||||
/// 根据额度窗口边界漂移判断是否重置本地窗口计数。
|
||||
pub fn should_reset_usage_window(
|
||||
existing_reset_at: Option<DateTime<Utc>>,
|
||||
existing_limit_window_seconds: Option<u64>,
|
||||
new_reset_at: DateTime<Utc>,
|
||||
new_limit_window_seconds: Option<u64>,
|
||||
) -> bool {
|
||||
let Some(existing_reset_at) = existing_reset_at else {
|
||||
return false;
|
||||
};
|
||||
if existing_reset_at == new_reset_at {
|
||||
return false;
|
||||
}
|
||||
|
||||
let drift = existing_reset_at
|
||||
.signed_duration_since(new_reset_at)
|
||||
.num_seconds()
|
||||
.unsigned_abs();
|
||||
let window_seconds = new_limit_window_seconds
|
||||
.or(existing_limit_window_seconds)
|
||||
.unwrap_or(0);
|
||||
let threshold = if window_seconds > 0 {
|
||||
window_seconds / 2
|
||||
} else {
|
||||
3_600
|
||||
};
|
||||
drift >= threshold
|
||||
}
|
||||
|
||||
@@ -55,6 +55,7 @@ impl AccountUsageQueryService {
|
||||
nonnegative_i64_to_u64(bucket.totals.input_tokens),
|
||||
nonnegative_i64_to_u64(bucket.totals.output_tokens),
|
||||
nonnegative_i64_to_u64(bucket.totals.cached_tokens),
|
||||
nonnegative_i64_to_u64(bucket.totals.cache_write_tokens),
|
||||
&bucket.model,
|
||||
bucket.service_tier.as_deref(),
|
||||
);
|
||||
@@ -159,6 +160,7 @@ pub struct AccountUsageTimeBucket {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub cache_write_tokens: i64,
|
||||
pub first_token_latency_sum: i64,
|
||||
pub first_token_latency_count: i64,
|
||||
pub latency_sum: i64,
|
||||
@@ -178,6 +180,7 @@ impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
cached_tokens: record.cached_tokens,
|
||||
cache_write_tokens: record.cache_write_tokens,
|
||||
first_token_latency_sum: record.first_token_latency_sum,
|
||||
first_token_latency_count: record.first_token_latency_count,
|
||||
latency_sum: record.latency_sum,
|
||||
|
||||
@@ -7,7 +7,13 @@ use chrono::{DateTime, Utc};
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row, postgres::PgRow};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::format::nonnegative_i64_to_u64;
|
||||
use crate::{
|
||||
fleet::usage::{
|
||||
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageStoreError,
|
||||
AccountUsageWindow,
|
||||
},
|
||||
infra::format::nonnegative_i64_to_u64,
|
||||
};
|
||||
|
||||
const RECORD_USAGE_SQL: &str = r"
|
||||
insert into account_usage (
|
||||
@@ -140,108 +146,9 @@ pub enum PgAccountUsageStoreError {
|
||||
/// PostgreSQL 用量存储结果。
|
||||
pub type PgAccountUsageStoreResult<T> = Result<T, PgAccountUsageStoreError>;
|
||||
|
||||
/// 账号用量端口错误。
|
||||
#[derive(Debug, Error)]
|
||||
#[error("account usage store operation failed: {message}")]
|
||||
pub struct AccountUsageStoreError {
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl From<PgAccountUsageStoreError> for AccountUsageStoreError {
|
||||
fn from(error: PgAccountUsageStoreError) -> Self {
|
||||
Self {
|
||||
message: error.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageDelta {
|
||||
pub requests: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub reasoning_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub empty_responses: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
pub image_requests: u64,
|
||||
pub image_request_failures: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageSnapshot {
|
||||
pub request_count: u64,
|
||||
pub empty_response_count: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
pub image_request_count: u64,
|
||||
pub image_request_failed_count: u64,
|
||||
pub window_request_count: u64,
|
||||
pub window_input_tokens: u64,
|
||||
pub window_output_tokens: u64,
|
||||
pub window_cached_tokens: u64,
|
||||
pub window_image_input_tokens: u64,
|
||||
pub window_image_output_tokens: u64,
|
||||
pub window_image_request_count: u64,
|
||||
pub window_image_request_failed_count: u64,
|
||||
pub window_started_at: Option<DateTime<Utc>>,
|
||||
pub window_reset_at: Option<DateTime<Utc>>,
|
||||
pub limit_window_seconds: Option<u64>,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageWindow {
|
||||
pub request_count: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub image_input_tokens: u64,
|
||||
pub image_output_tokens: u64,
|
||||
pub image_request_count: u64,
|
||||
pub image_request_failed_count: u64,
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
pub reset_at: Option<DateTime<Utc>>,
|
||||
pub limit_window_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AccountUsageStore: Send + Sync + 'static {
|
||||
async fn snapshots(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
) -> Result<HashMap<String, AccountUsageSnapshot>, AccountUsageStoreError>;
|
||||
|
||||
async fn record_usage_delta(
|
||||
&self,
|
||||
account_id: &str,
|
||||
usage: AccountUsageDelta,
|
||||
) -> Result<(), AccountUsageStoreError>;
|
||||
|
||||
async fn sync_runtime_window(
|
||||
&self,
|
||||
account_id: &str,
|
||||
window: AccountUsageWindow,
|
||||
) -> Result<(), AccountUsageStoreError>;
|
||||
|
||||
async fn sync_rate_limit_window(
|
||||
&self,
|
||||
account_id: &str,
|
||||
reset_at: DateTime<Utc>,
|
||||
limit_window_seconds: Option<u64>,
|
||||
) -> Result<(), AccountUsageStoreError>;
|
||||
|
||||
async fn record_request(&self, account_id: &str) -> Result<(), AccountUsageStoreError> {
|
||||
self.record_usage_delta(
|
||||
account_id,
|
||||
AccountUsageDelta {
|
||||
requests: 1,
|
||||
..AccountUsageDelta::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
Self::adapter(error)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ struct ModelPricing {
|
||||
output_price_per_mtoken_priority: f64,
|
||||
cache_read_price_per_mtoken: f64,
|
||||
cache_read_price_per_mtoken_priority: f64,
|
||||
cache_write_multiplier: f64,
|
||||
long_input_price_per_mtoken: f64,
|
||||
long_input_price_per_mtoken_priority: f64,
|
||||
long_output_price_per_mtoken: f64,
|
||||
@@ -27,6 +28,7 @@ impl ModelPricing {
|
||||
output_price_per_mtoken_priority: 0.0,
|
||||
cache_read_price_per_mtoken: 0.0,
|
||||
cache_read_price_per_mtoken_priority: 0.0,
|
||||
cache_write_multiplier: 0.0,
|
||||
long_input_price_per_mtoken: 0.0,
|
||||
long_input_price_per_mtoken_priority: 0.0,
|
||||
long_output_price_per_mtoken: 0.0,
|
||||
@@ -41,6 +43,11 @@ impl ModelPricing {
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_cache_write(mut self, multiplier: f64) -> Self {
|
||||
self.cache_write_multiplier = multiplier;
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_priority(mut self, input: f64, output: f64, cache_read: f64) -> Self {
|
||||
self.input_price_per_mtoken_priority = input;
|
||||
self.output_price_per_mtoken_priority = output;
|
||||
@@ -78,24 +85,28 @@ const MODEL_PRICING_RULES: &[ModelPricingRule] = &[
|
||||
model: "gpt-5.6-sol",
|
||||
pricing: ModelPricing::new(5.0, 30.0)
|
||||
.with_cache(0.5)
|
||||
.with_cache_write(1.25)
|
||||
.with_priority(10.0, 60.0, 1.0),
|
||||
},
|
||||
ModelPricingRule {
|
||||
model: "gpt-5.6-terra",
|
||||
pricing: ModelPricing::new(2.5, 15.0)
|
||||
.with_cache(0.25)
|
||||
.with_cache_write(1.25)
|
||||
.with_priority(5.0, 30.0, 0.5),
|
||||
},
|
||||
ModelPricingRule {
|
||||
model: "gpt-5.6-luna",
|
||||
pricing: ModelPricing::new(1.0, 6.0)
|
||||
.with_cache(0.1)
|
||||
.with_cache_write(1.25)
|
||||
.with_priority(2.0, 12.0, 0.2),
|
||||
},
|
||||
ModelPricingRule {
|
||||
model: "gpt-5.6",
|
||||
pricing: ModelPricing::new(5.0, 30.0)
|
||||
.with_cache(0.5)
|
||||
.with_cache_write(1.25)
|
||||
.with_priority(10.0, 60.0, 1.0),
|
||||
},
|
||||
ModelPricingRule {
|
||||
@@ -181,6 +192,8 @@ pub(crate) struct BillingBreakdown {
|
||||
pub output_amount: f64,
|
||||
/// 缓存读取金额。
|
||||
pub cache_read_amount: f64,
|
||||
/// 缓存写入金额。
|
||||
pub cache_write_amount: f64,
|
||||
/// 应用服务档位倍率前的标准金额。
|
||||
pub standard_amount: f64,
|
||||
/// 最终计费金额。
|
||||
@@ -191,6 +204,8 @@ pub(crate) struct BillingBreakdown {
|
||||
pub output_price_per_mtoken: f64,
|
||||
/// 缓存读取单价,美元 / 1M token。
|
||||
pub cache_read_price_per_mtoken: f64,
|
||||
/// 缓存写入单价,美元 / 1M token。
|
||||
pub cache_write_price_per_mtoken: f64,
|
||||
/// 服务档位。
|
||||
pub service_tier: Option<String>,
|
||||
/// 服务档位倍率。
|
||||
@@ -202,16 +217,24 @@ pub(crate) fn calculate_billing(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
model: &str,
|
||||
service_tier: Option<&str>,
|
||||
) -> BillingBreakdown {
|
||||
let pricing = model_pricing(model).unwrap_or_else(|| {
|
||||
highest_pricing_for_request(input_tokens, output_tokens, cached_tokens, service_tier)
|
||||
highest_pricing_for_request(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
service_tier,
|
||||
)
|
||||
});
|
||||
calculate_billing_with_pricing(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
pricing,
|
||||
service_tier,
|
||||
)
|
||||
@@ -221,6 +244,7 @@ fn calculate_billing_with_pricing(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
pricing: ModelPricing,
|
||||
service_tier: Option<&str>,
|
||||
) -> BillingBreakdown {
|
||||
@@ -230,6 +254,7 @@ fn calculate_billing_with_pricing(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
pricing,
|
||||
service_tier,
|
||||
is_long,
|
||||
@@ -240,6 +265,7 @@ fn calculate_billing_with_pricing_for_context(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
pricing: ModelPricing,
|
||||
service_tier: Option<&str>,
|
||||
is_long: bool,
|
||||
@@ -277,26 +303,38 @@ fn calculate_billing_with_pricing_for_context(
|
||||
}
|
||||
|
||||
let cached_tokens = cached_tokens.min(input_tokens);
|
||||
let uncached_input_tokens = if cache_read_price > 0.0 {
|
||||
input_tokens - cached_tokens
|
||||
let billed_cache_read_tokens = if cache_read_price > 0.0 {
|
||||
cached_tokens
|
||||
} else {
|
||||
input_tokens
|
||||
0
|
||||
};
|
||||
let cache_write_price = input_price * pricing.cache_write_multiplier;
|
||||
let billed_cache_write_tokens = if cache_write_price > 0.0 {
|
||||
cache_write_tokens.min(input_tokens - billed_cache_read_tokens)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let uncached_input_tokens = input_tokens
|
||||
.saturating_sub(billed_cache_read_tokens)
|
||||
.saturating_sub(billed_cache_write_tokens);
|
||||
let input_amount = uncached_input_tokens as f64 / 1_000_000.0 * input_price;
|
||||
let cache_read_amount = cached_tokens as f64 / 1_000_000.0 * cache_read_price;
|
||||
let cache_read_amount = billed_cache_read_tokens as f64 / 1_000_000.0 * cache_read_price;
|
||||
let cache_write_amount = billed_cache_write_tokens as f64 / 1_000_000.0 * cache_write_price;
|
||||
let output_amount = output_tokens as f64 / 1_000_000.0 * output_price;
|
||||
let standard_amount = input_amount + cache_read_amount + output_amount;
|
||||
let standard_amount = input_amount + cache_read_amount + cache_write_amount + output_amount;
|
||||
let total_amount = standard_amount * tier_multiplier;
|
||||
|
||||
BillingBreakdown {
|
||||
input_amount,
|
||||
output_amount,
|
||||
cache_read_amount,
|
||||
cache_write_amount,
|
||||
standard_amount,
|
||||
total_amount,
|
||||
input_price_per_mtoken: input_price,
|
||||
output_price_per_mtoken: output_price,
|
||||
cache_read_price_per_mtoken: cache_read_price,
|
||||
cache_write_price_per_mtoken: cache_write_price,
|
||||
service_tier: normalize_service_tier(service_tier),
|
||||
tier_multiplier,
|
||||
}
|
||||
@@ -309,6 +347,7 @@ pub(crate) fn calculate_aggregate_billing(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
model: &str,
|
||||
service_tier: Option<&str>,
|
||||
is_long: bool,
|
||||
@@ -318,6 +357,7 @@ pub(crate) fn calculate_aggregate_billing(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
service_tier,
|
||||
is_long,
|
||||
)
|
||||
@@ -326,6 +366,7 @@ pub(crate) fn calculate_aggregate_billing(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
pricing,
|
||||
service_tier,
|
||||
is_long,
|
||||
@@ -337,6 +378,7 @@ pub(crate) fn calculate_billing_amount(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
model: &str,
|
||||
service_tier: Option<&str>,
|
||||
) -> f64 {
|
||||
@@ -344,6 +386,7 @@ pub(crate) fn calculate_billing_amount(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
model,
|
||||
service_tier,
|
||||
)
|
||||
@@ -365,6 +408,7 @@ fn highest_pricing_for_request(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
service_tier: Option<&str>,
|
||||
) -> ModelPricing {
|
||||
MODEL_PRICING_RULES
|
||||
@@ -374,6 +418,7 @@ fn highest_pricing_for_request(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
highest,
|
||||
service_tier,
|
||||
)
|
||||
@@ -382,6 +427,7 @@ fn highest_pricing_for_request(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
rule.pricing,
|
||||
service_tier,
|
||||
)
|
||||
@@ -398,6 +444,7 @@ fn highest_pricing_for_context(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
service_tier: Option<&str>,
|
||||
is_long: bool,
|
||||
) -> ModelPricing {
|
||||
@@ -408,6 +455,7 @@ fn highest_pricing_for_context(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
highest,
|
||||
service_tier,
|
||||
is_long,
|
||||
@@ -417,6 +465,7 @@ fn highest_pricing_for_context(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
rule.pricing,
|
||||
service_tier,
|
||||
is_long,
|
||||
|
||||
@@ -18,6 +18,7 @@ select
|
||||
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
|
||||
coalesce(sum(first_token_latency_sum), 0)::bigint as first_token_latency_sum,
|
||||
coalesce(sum(first_token_latency_count), 0)::bigint as first_token_latency_count,
|
||||
coalesce(sum(latency_sum), 0)::bigint as latency_sum,
|
||||
@@ -36,7 +37,8 @@ select
|
||||
coalesce(sum(success_count + error_count), 0)::bigint as request_count,
|
||||
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens
|
||||
from request_time_buckets
|
||||
group by bucket_start, model, service_tier";
|
||||
|
||||
@@ -59,6 +61,7 @@ pub struct UsageTimeBucketRecord {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub cache_write_tokens: i64,
|
||||
pub first_token_latency_sum: i64,
|
||||
pub first_token_latency_count: i64,
|
||||
pub latency_sum: i64,
|
||||
@@ -81,6 +84,7 @@ pub struct UsageBucketTotals {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub cache_write_tokens: i64,
|
||||
}
|
||||
|
||||
/// 保留期内单个计费桶的用量事实。
|
||||
@@ -97,6 +101,9 @@ impl UsageBucketTotals {
|
||||
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
|
||||
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
|
||||
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
|
||||
self.cache_write_tokens = self
|
||||
.cache_write_tokens
|
||||
.saturating_add(other.cache_write_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -117,6 +124,7 @@ pub struct ModelBucketUsage {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub cache_write_tokens: i64,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
@@ -163,6 +171,7 @@ impl PgRequestBucketQuery {
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
cache_write_tokens: row.get("cache_write_tokens"),
|
||||
},
|
||||
})
|
||||
.collect())
|
||||
@@ -198,7 +207,8 @@ impl PgRequestBucketQuery {
|
||||
coalesce(sum(success_count + error_count), 0)::bigint as request_count,
|
||||
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens
|
||||
from request_time_buckets
|
||||
where account_id in (",
|
||||
);
|
||||
@@ -223,6 +233,7 @@ impl PgRequestBucketQuery {
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
cache_write_tokens: row.get("cache_write_tokens"),
|
||||
};
|
||||
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
|
||||
continue;
|
||||
@@ -258,6 +269,7 @@ impl PgRequestBucketQuery {
|
||||
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
|
||||
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
|
||||
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
|
||||
max(bucket_start) as last_used_at
|
||||
from request_time_buckets
|
||||
where model != '__unknown__' and (",
|
||||
@@ -290,6 +302,7 @@ impl PgRequestBucketQuery {
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
cache_write_tokens: row.get("cache_write_tokens"),
|
||||
last_used_at: row.get("last_used_at"),
|
||||
}
|
||||
})
|
||||
@@ -307,6 +320,7 @@ fn usage_time_bucket_from_row(row: &sqlx::postgres::PgRow) -> UsageTimeBucketRec
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
cache_write_tokens: row.get("cache_write_tokens"),
|
||||
first_token_latency_sum: row.get("first_token_latency_sum"),
|
||||
first_token_latency_count: row.get("first_token_latency_count"),
|
||||
latency_sum: row.get("latency_sum"),
|
||||
|
||||
@@ -53,15 +53,16 @@ impl PgRequestBucketStore {
|
||||
r"
|
||||
insert into request_time_buckets (
|
||||
bucket_start, provider, account_id, model, service_tier, success_count,
|
||||
input_tokens, output_tokens, cached_tokens, first_token_latency_sum,
|
||||
input_tokens, output_tokens, cached_tokens, cache_write_tokens, first_token_latency_sum,
|
||||
first_token_latency_count, latency_sum, latency_count, max_latency_ms,
|
||||
min_latency_ms, updated_at
|
||||
) values ($1, $2, $3, $4, $5, 1, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
|
||||
) values ($1, $2, $3, $4, $5, 1, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
|
||||
on conflict (bucket_start, provider, account_id, model, service_tier) do update set
|
||||
success_count = request_time_buckets.success_count + 1,
|
||||
input_tokens = request_time_buckets.input_tokens + excluded.input_tokens,
|
||||
output_tokens = request_time_buckets.output_tokens + excluded.output_tokens,
|
||||
cached_tokens = request_time_buckets.cached_tokens + excluded.cached_tokens,
|
||||
cache_write_tokens = request_time_buckets.cache_write_tokens + excluded.cache_write_tokens,
|
||||
first_token_latency_sum = request_time_buckets.first_token_latency_sum + excluded.first_token_latency_sum,
|
||||
first_token_latency_count = request_time_buckets.first_token_latency_count + excluded.first_token_latency_count,
|
||||
latency_sum = request_time_buckets.latency_sum + excluded.latency_sum,
|
||||
@@ -82,6 +83,7 @@ on conflict (bucket_start, provider, account_id, model, service_tier) do update
|
||||
.bind(event.input_tokens.unwrap_or(0))
|
||||
.bind(event.output_tokens.unwrap_or(0))
|
||||
.bind(event.cached_tokens.unwrap_or(0))
|
||||
.bind(event.cache_write_tokens.unwrap_or(0))
|
||||
.bind(first_token.unwrap_or(0))
|
||||
.bind(i64::from(first_token.is_some()))
|
||||
.bind(latency.unwrap_or(0))
|
||||
@@ -188,6 +190,7 @@ with facts as (
|
||||
coalesce(input_tokens, 0)::bigint as input_tokens,
|
||||
coalesce(output_tokens, 0)::bigint as output_tokens,
|
||||
coalesce(cached_tokens, 0)::bigint as cached_tokens,
|
||||
coalesce(cache_write_tokens, 0)::bigint as cache_write_tokens,
|
||||
coalesce(first_token_ms, 0)::bigint as first_token_latency_sum,
|
||||
(first_token_ms is not null)::integer::bigint as first_token_latency_count,
|
||||
coalesce(latency_ms, 0)::bigint as latency_sum,
|
||||
@@ -215,6 +218,7 @@ with facts as (
|
||||
0::bigint,
|
||||
0::bigint,
|
||||
0::bigint,
|
||||
0::bigint,
|
||||
null::bigint
|
||||
from ops_error_logs
|
||||
where created_at >= $1
|
||||
@@ -226,6 +230,7 @@ with facts as (
|
||||
sum(input_tokens) as input_tokens,
|
||||
sum(output_tokens) as output_tokens,
|
||||
sum(cached_tokens) as cached_tokens,
|
||||
sum(cache_write_tokens) as cache_write_tokens,
|
||||
sum(first_token_latency_sum) as first_token_latency_sum,
|
||||
sum(first_token_latency_count) as first_token_latency_count,
|
||||
sum(latency_sum) as latency_sum,
|
||||
@@ -237,13 +242,13 @@ with facts as (
|
||||
)
|
||||
insert into request_time_buckets (
|
||||
bucket_start, provider, account_id, model, service_tier,
|
||||
success_count, error_count, input_tokens, output_tokens, cached_tokens,
|
||||
success_count, error_count, input_tokens, output_tokens, cached_tokens, cache_write_tokens,
|
||||
first_token_latency_sum, first_token_latency_count, latency_sum, latency_count,
|
||||
max_latency_ms, min_latency_ms, updated_at
|
||||
)
|
||||
select
|
||||
bucket_start, provider, account_id, model, service_tier,
|
||||
success_count, error_count, input_tokens, output_tokens, cached_tokens,
|
||||
success_count, error_count, input_tokens, output_tokens, cached_tokens, cache_write_tokens,
|
||||
first_token_latency_sum, first_token_latency_count, latency_sum, latency_count,
|
||||
max_latency_ms, min_latency_ms, now()
|
||||
from aggregated",
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
pub mod account_usage;
|
||||
pub(crate) mod billing;
|
||||
pub mod buckets;
|
||||
pub mod dashboard;
|
||||
pub mod ops;
|
||||
pub mod rebuild;
|
||||
pub mod recorder;
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::telemetry::{
|
||||
ops::store::{OpsErrorFilter, PgOpsErrorLogStore},
|
||||
ops::types::OpsErrorLog,
|
||||
ops::store::PgOpsErrorLogStore,
|
||||
ops::types::{OpsErrorFilter, OpsErrorLog},
|
||||
};
|
||||
|
||||
/// 管理端运维错误事件错误。
|
||||
|
||||
@@ -6,7 +6,10 @@ use thiserror::Error;
|
||||
|
||||
use crate::{
|
||||
infra::json::{NumberedPage, page_offset},
|
||||
telemetry::{buckets::store::PgRequestBucketStore, ops::types::OpsErrorLog},
|
||||
telemetry::{
|
||||
buckets::store::PgRequestBucketStore,
|
||||
ops::types::{OpsErrorFilter, OpsErrorLog},
|
||||
},
|
||||
};
|
||||
|
||||
const OPS_ERROR_SELECT_SQL: &str = r"select
|
||||
@@ -143,28 +146,6 @@ insert into ops_error_logs (
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct OpsErrorFilter {
|
||||
pub kind: Option<String>,
|
||||
pub client_api_key_id: Option<String>,
|
||||
pub provider: Option<String>,
|
||||
pub request_id: Option<String>,
|
||||
pub account_id: Option<String>,
|
||||
pub route: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub client_status_code: Option<i64>,
|
||||
pub upstream_status_code: Option<i64>,
|
||||
pub transport: Option<String>,
|
||||
pub attempt_index: Option<i64>,
|
||||
pub failure_class: Option<String>,
|
||||
pub response_id: Option<String>,
|
||||
pub upstream_request_id: Option<String>,
|
||||
pub search: Option<String>,
|
||||
pub start_time: Option<DateTime<Utc>>,
|
||||
pub end_time: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
fn push_filter(builder: &mut QueryBuilder<Postgres>, filter: &OpsErrorFilter) {
|
||||
builder.push(" where true");
|
||||
push_optional_text(builder, "kind", filter.kind.as_deref());
|
||||
|
||||
@@ -55,6 +55,29 @@ pub struct OpsErrorLog {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 运维错误查询条件。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct OpsErrorFilter {
|
||||
pub kind: Option<String>,
|
||||
pub client_api_key_id: Option<String>,
|
||||
pub provider: Option<String>,
|
||||
pub request_id: Option<String>,
|
||||
pub account_id: Option<String>,
|
||||
pub route: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub client_status_code: Option<i64>,
|
||||
pub upstream_status_code: Option<i64>,
|
||||
pub transport: Option<String>,
|
||||
pub attempt_index: Option<i64>,
|
||||
pub failure_class: Option<String>,
|
||||
pub response_id: Option<String>,
|
||||
pub upstream_request_id: Option<String>,
|
||||
pub search: Option<String>,
|
||||
pub start_time: Option<DateTime<Utc>>,
|
||||
pub end_time: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl OpsErrorLog {
|
||||
/// 构造运维错误事件。
|
||||
pub fn new(kind: impl Into<String>, message: impl Into<String>) -> Self {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user