Compare commits

..
34 Commits
Author SHA1 Message Date
Zyy 521201c926 fix(telemetry): preserve timed-out record writes 2026-07-17 14:52:19 +08:00
ZyyandCodex be2123ea30 chore: bump version to 2.5.0
Co-authored-by: Codex <noreply@openai.com>
2026-07-17 13:40:14 +08:00
ZyyandCodex 4b6d369225 refactor(backend): simplify admin query structure
Move account and dashboard read orchestration into their existing admin API modules, and remove the top-level admin_queries package and service wrapper.

Update composition wiring, dependency guards, and architecture documentation to reflect feature-local query ownership.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 13:09:17 +08:00
ZyyandCodex 76b81a6ceb refactor(frontend): simplify client architecture
Remove frontend schemas and tests, keep API modules as thin GET/POST adapters, and replace manual SSE and browser helpers with existing VueUse and utility primitives.

Reorganize feature helpers, move unauthorized handling into an app plugin, align the admin key removal route, and satisfy the enabled accessibility rules.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 11:44:04 +08:00
zyycnandCodex 29c8b679ba refactor(architecture): complete dependency boundary audit
Harden WebSocket ownership and capacity semantics, split the transport into focused modules, and move admin queries, quota observations, and upstream adapters behind explicit consumer-owned ports.

Parse frontend payloads at the API boundary, centralize SSE and store ownership, add latest-request-wins queries with inferred view types, and enforce dependency direction through lint and AST tests.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 09:03:33 +08:00
zyycnandCodex 4091a0c2d8 feat(proxy): prefer passive upstream state refresh
Remove request-path quota probes and consume account, request, and rate-limit facts from normal upstream responses, including 429 active-limit metadata.

Refresh model catalogs only at startup, on x-models-etag changes, or through the explicit account model refresh control in the connection test dialog.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 02:12:16 +08:00
zyycnandCodex f0ccc981fb feat(proxy): align latest Responses semantics
Preserve Responses Lite and Memgen controls across HTTP and WebSocket transports without rewriting request bodies, and standardize continuation recovery signals.

Track cache-write usage and billing while applying appcast profile updates to new WebSocket generations without breaking exact continuations.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 00:50:27 +08:00
zyycnandCodex d01ea7f775 fix(dispatch): generalize cross-account replay sanitization
Sanitize every replay item envelope without maintaining an output-type allowlist, while preserving nested client and tool semantics.

Build connection snapshots from the actual successful attempt input so cleaned source-account state cannot return during target-account recovery.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 23:59:10 +08:00
zyycnandCodex fd749ea8f8 fix(dispatch): isolate cross-account replay state
Track connection transcripts by the actual response account and sanitize typed Responses output state before replaying across accounts.

Rebuild target auth, Cookie, installation metadata, and turn state boundaries, fail closed on malformed turn metadata, and persist only target-account state after replay.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 23:18:26 +08:00
zyycnandCodex f0543b637e fix(frontend): refine analytics interactions
Align usage chart bounds and diagnostic loading states, simplify dashboard hover feedback, and improve User-Agent layout.

Restore the login light palette to the primary accent color.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 20:03:51 +08:00
ZyyandCodex 549bb9f62b feat(telemetry): expose latency breakdowns
Capture first semantic output and transport timing, then surface concise latency values with safe phase details in dashboard and usage records.

Keep the Codex Desktop request profile synchronized with successful appcast checks at runtime, without rewriting configuration or evicting existing WebSocket connections.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 18:07:08 +08:00
ZyyandCodex 9e57f78bab ci: speed up validation pipelines
Run affected CI checks in parallel and skip unrelated backend or frontend work while retaining workflow lint.

Start release binary builds after release metadata is available, while keeping image assembly, packaging, and publication gated by quality checks.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 16:07:55 +08:00
ZyyandCodex 73c0ecdc4c feat(telemetry): add first-token latency observability
Measure TTFT from the first semantic response event across streaming, retry, scheduler, usage telemetry, and dashboard paths.

Expose first-token percentiles and concise UI labels, and isolate concurrent Redis-backed test runs.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 15:57:19 +08:00
zyycnandCodex bfd05722d5 chore: bump version to 2.4.0
Co-authored-by: Codex <noreply@openai.com>
2026-07-16 02:22:13 +08:00
dependabot[bot]andCodex 9dc2ea4a2f chore(deps): update backend dependency group
Update redis to 1.4.0, rustls to 0.23.42, and uuid to 1.23.5 after full backend, security, container, and Trivy validation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 02:12:01 +08:00
dependabot[bot]andCodex 6fb27e869e chore(deps): refresh Rust 1.97 builder digest
Refresh the pinned rust:1.97-bookworm digest after a successful container build and Trivy scan.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 02:01:08 +08:00
dependabot[bot]andCodex 81c58b1cdc chore(deps): bump setup-node to 7.0.0
Update the pinned setup-node action to 7.0.0 across quality and security workflows.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 01:51:31 +08:00
zyycnandCodex c8dc2e293a fix(transport): eliminate websocket fallback race
Notify foreground and shared waiters before awaiting pool cleanup so an observed WebSocket opening error cannot be replaced by a fast-path timeout and HTTP fallback.

Drive the half-open breaker test through explicit state transitions instead of timing sleeps.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 01:46:33 +08:00
dependabot[bot] 25081de913 chore(deps): bump action-gh-release to 3.0.2
Update the pinned GitHub Release action from 3.0.1 to 3.0.2.\n\nCo-authored-by: Codex <noreply@openai.com>
2026-07-16 01:23:07 +08:00
zyycnandCodex 6eb0917ff1 feat(frontend): enrich interface micro-interactions
Add responsive hover reactions across dashboard, account, usage, settings, login, and navigation surfaces. Keep brand motion distinct, pair enter and leave states, and respect reduced-motion preferences.

Refine sidebar transitions, dashboard heartbeat, trend feedback, and 15-minute metric buckets without animating page content.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 01:07:33 +08:00
zyycnandCodex 72d892d5e3 feat(dashboard): improve request health visibility
Classify final request outcomes into success, failure, cancellation, and caller-error buckets so availability reflects only eligible requests. Expose structured 15-minute health points with focused coverage.

Redesign the health timeline, trend, wire-profile, and diagnostic surfaces with responsive shared interactions, and standardize configuration schema version 1.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 23:52:35 +08:00
ZyyandCodex 0b2da3087b refactor: replace fingerprint with wire profile
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 18:23:59 +08:00
zyycnandCodex 5391fdb31c fix(lifecycle): eliminate async finalization races
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:57:54 +08:00
zyycnandCodex 24036f1f84 docs: refine project README
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:27:27 +08:00
zyycnandCodex b043fced0c docs: streamline project guidance
Keep the public README terse and operational, and align the repository skills with the current lifecycle, transport, validation, and release boundaries.

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

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

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

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

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 03:09:39 +08:00
378 changed files with 24219 additions and 14442 deletions
+2
View File
@@ -34,6 +34,8 @@ updates:
day: monday
time: '09:20'
timezone: Asia/Shanghai
cooldown:
default-days: 1
groups:
minor-and-patch:
update-types:
+13 -1
View File
@@ -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
+17 -3
View File
@@ -3,8 +3,10 @@ name: CI
on:
push:
branches:
- '**'
- main
pull_request:
branches:
- main
permissions:
contents: read
@@ -18,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
@@ -26,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/**'
@@ -41,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
+3 -3
View File
@@ -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
+4 -2
View File
@@ -3,8 +3,10 @@ name: Security Scan
on:
push:
branches:
- '**'
- main
pull_request:
branches:
- main
schedule:
- cron: '0 3 * * 1'
@@ -55,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
+67 -122
View File
@@ -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 多账号网关代理**
[![CI](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
[![Release](https://img.shields.io/github/v/release/zyycn/codex-proxy-rs?display_name=tag&sort=semver)](https://github.com/zyycn/codex-proxy-rs/releases)
[![Rust 1.97](https://img.shields.io/badge/Rust-1.97-000?logo=rust&logoColor=white)](https://www.rust-lang.org/)
[![Vue 3.5](https://img.shields.io/badge/Vue-3.5-42b883?logo=vuedotjs&logoColor=white)](https://vuejs.org/)
[![Vite 8](https://img.shields.io/badge/Vite-8-646cff?logo=vite&logoColor=white)](https://vite.dev/)
[![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
[![MIT](https://img.shields.io/badge/License-MIT-blue.svg)](#开源许可)
[![Release](https://img.shields.io/github/v/release/zyycn/codex-proxy-rs?display_name=tag&sort=semver&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/releases)
[![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
[![MIT](https://img.shields.io/badge/License-MIT-blue.svg?style=flat-square)](https://opensource.org/license/mit)
[快速开始](#快速开始) · [首次使用](#首次使用) · [客户端接入](#客户端接入) · [常见问题](#常见问题) · [部署文档](deploy/README.md)
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运维](#运维) · [部署文档](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) 开源。
+17 -6
View File
@@ -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
View File
@@ -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"
+41 -106
View File
@@ -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, &quota_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(&quota_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), &quota_by_account).await?;
Ok(stats.data_for(account))
}
@@ -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::{AccountUsageSummary, AccountUsageTimeBucket};
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)]
@@ -79,7 +110,7 @@ struct DashboardTokenCardData {
today_tokens_value: u64,
yesterday_tokens_value: u64,
total_tokens: String,
today_billing_amount_usd: String,
total_billing_amount_usd: String,
}
#[derive(Debug, Serialize)]
@@ -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
@@ -174,25 +213,25 @@ impl UsageWindow {
pub fn dashboard_cards(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
retained_usage: &RetainedUsageSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, buckets, lifetime_usage, Utc::now())
dashboard_cards_at(accounts, buckets, retained_usage, Utc::now())
}
pub fn dashboard_cards_at(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
retained_usage: &RetainedUsageSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
let yesterday_start = today_start - Duration::days(1);
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_billing_amount = billing_window(buckets, today_start, now).unwrap_or(0.0);
let total_requests = nonnegative_i64_to_u64(lifetime_usage.request_count);
let total_tokens = nonnegative_i64_to_u64(lifetime_usage.total_tokens);
let total_cached_tokens = nonnegative_i64_to_u64(lifetime_usage.cached_tokens);
let total_requests = nonnegative_i64_to_u64(retained_usage.totals.request_count);
let total_tokens = nonnegative_i64_to_u64(retained_usage.totals.input_tokens)
.saturating_add(nonnegative_i64_to_u64(retained_usage.totals.output_tokens));
let total_cached_tokens = nonnegative_i64_to_u64(retained_usage.totals.cached_tokens);
DashboardCardsData {
accounts: DashboardAccountsCardData {
@@ -214,13 +253,13 @@ pub fn dashboard_cards_at(
today_tokens_value: today.tokens(),
yesterday_tokens_value: yesterday.tokens(),
total_tokens: format_tokens(total_tokens),
today_billing_amount_usd: format_billing_amount(today_billing_amount),
total_billing_amount_usd: format_billing_amount(retained_usage.billing_amount_usd),
},
cache: DashboardCacheCardData {
today_hit_rate: format_rate(today.cache_hit_rate()),
today_hit_rate_value: today.cache_hit_rate(),
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
total_hit_rate: format_rate(summary_cache_hit_rate(lifetime_usage)),
total_hit_rate: format_rate(summary_cache_hit_rate(retained_usage)),
total_cached_tokens: format_tokens(total_cached_tokens),
average_first_token_latency_ms: format_optional_duration_ms(
today.avg_first_token_latency(),
@@ -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(),
}
@@ -486,19 +551,6 @@ fn usage_window(
})
}
fn billing_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Option<f64> {
let billing_amounts = records
.iter()
.filter(|record| record.bucket_start >= start && record.bucket_start < end)
.filter_map(|record| record.billing_amount_usd)
.collect::<Vec<_>>();
(!billing_amounts.is_empty()).then(|| billing_amounts.into_iter().sum())
}
fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
@@ -523,27 +575,47 @@ 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 summary_cache_hit_rate(summary: &AccountUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
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)
.then(|| nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
.then(|| nonnegative_i64_to_u64(summary.totals.cached_tokens) as f64 / input_tokens as f64)
}
fn format_optional_duration_ms(value: Option<u64>) -> String {
File diff suppressed because it is too large Load Diff
+1
View File
@@ -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;
+4 -1
View File
@@ -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)]
+7 -3
View File
@@ -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))
+1 -1
View File
@@ -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,
+25
View File
@@ -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),
+1 -1
View File
@@ -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",
+25 -1
View File
@@ -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"))
+33 -4
View File
@@ -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)
+9 -14
View File
@@ -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,
+67 -71
View File
@@ -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 {
+131 -48
View File
@@ -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 {
@@ -268,7 +323,7 @@ impl Services {
)
.with_refresh_lease_store(stores.refresh_leases.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 +335,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 +368,8 @@ impl Services {
}));
let system_update = Arc::new(SystemUpdateService::from_env());
let desktop_release = DesktopReleaseStatus::default();
Ok(Self {
database,
redis,
@@ -333,7 +390,8 @@ impl Services {
session_affinity,
codex,
websocket_pool,
fingerprint,
wire_profile,
desktop_release,
account_pseudonymizer,
system_update,
connection_drain,
@@ -351,13 +409,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 +440,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 +491,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>> {
@@ -560,14 +649,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 +667,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),
)?;
+51 -39
View File
@@ -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(),
);
@@ -82,6 +75,7 @@ impl TaskCoordinator {
stores.accounts.clone(),
stores.account_usage.clone(),
services.codex.clone(),
services.account_pool.clone(),
services.account_pseudonymizer.clone(),
config
.quota
@@ -91,20 +85,16 @@ impl TaskCoordinator {
30 * 60,
)
.with_cookie_store(stores.cookies.clone())
.with_account_pool(services.account_pool.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 定期版本检查失败");
}
}
}
})
}
}
+1 -1
View File
@@ -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;
+45 -85
View File
@@ -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 或手动操作再次触发");
}
}
-32
View File
@@ -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! {
+6 -9
View File
@@ -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,
@@ -35,15 +36,17 @@ impl QuotaRefreshTask {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<AccountPseudonymizer>,
interval_secs: u64,
min_refresh_interval_secs: u64,
) -> 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,
min_refresh_interval_secs,
),
@@ -58,12 +61,6 @@ impl QuotaRefreshTask {
self
}
/// 设置运行时账号池,用于刷新后同步内存调度状态。
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.service = self.service.with_account_pool(account_pool);
self
}
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let config = PeriodicTaskConfig::new(
+3 -21
View File
@@ -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, "清理超限指纹更新历史失败"),
}
}
}
+1 -1
View File
@@ -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,
+114 -38
View File
@@ -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,
@@ -2,19 +2,25 @@
use crate::{
dispatch::{
controllers::{
ControllerFailureFact,
account_state::{AccountStateEffect, AccountStateEffects},
},
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
},
transport::observation::{UpstreamFailureFacts, UpstreamFailureKind},
},
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{protocol::responses::ResponsesSseFailure, transport::CodexBackendClient},
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_account_failure,
},
pool::AccountPoolService,
},
upstream::openai::{
failure::{UpstreamFailureFacts, UpstreamFailureKind},
protocol::responses::ResponsesSseFailure,
transport::CodexBackendClient,
},
};
pub(super) struct AccountFailureController;
@@ -46,7 +52,7 @@ impl AccountFailureController {
return;
};
if let Some(effect) = &classified.effect {
AccountStateEffects::apply(account_pool, codex, account_id, effect).await;
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
}
}
@@ -110,33 +116,34 @@ impl AccountFailureController {
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let error_type = failure
.upstream_type
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let message = failure.message.to_ascii_lowercase();
if is_model_unsupported(&code)
|| is_model_unsupported(&message)
|| code.contains("invalid_request")
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 is_auth_failure(&code, &error_type, &message) {
return Some(ClientFailure::new(failure.clone(), 401, false));
}
if code.contains("forbidden") || code.contains("banned") || error_type == "permission_error"
if matches!(code.as_str(), "server_is_overloaded" | "slow_down")
|| code.contains("server_overloaded")
{
return Some(ClientFailure::new(failure.clone(), 403, false));
}
if code.contains("server_overloaded") {
return Some(ClientFailure::new(failure.clone(), 503, true));
}
None
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)
},
)
}
}
@@ -168,167 +175,42 @@ fn classify_failure(
}
fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<ClassifiedFailure> {
let status = facts.status_code;
let body = facts.body.as_str();
let code = facts.code.as_deref().unwrap_or_default();
let error_type = facts.error_type.as_deref().unwrap_or_default();
let message = facts.message.as_str();
if is_model_unsupported(code) || is_model_unsupported(message) || is_model_unsupported(body) {
return Some(model_unsupported_failure(account_id, body));
}
if let Some(account_status) = explicit_account_status(code, error_type, body) {
return Some(account_status_failure(
account_id,
account_status,
body,
status,
));
}
if status == Some(403) && !is_html(body) {
return Some(account_status_failure(
account_id,
AccountStatus::Banned,
body,
status,
));
}
if status == Some(401) {
return Some(account_status_failure(
account_id,
AccountStatus::Expired,
body,
status,
));
}
None
classified_failure(
account_id,
&facts.body,
facts.status_code,
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
)
}
fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<ClassifiedFailure> {
let code = failure.upstream_code.as_deref().unwrap_or_default();
let error_type = failure.upstream_type.as_deref().unwrap_or_default();
if is_model_unsupported(code) || is_model_unsupported(&failure.message) {
return Some(model_unsupported_failure(
account_id,
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
));
}
let account_status =
explicit_account_status(code, error_type, &failure.message).or_else(|| {
(code.eq_ignore_ascii_case("forbidden") || error_type == "permission_error")
.then_some(AccountStatus::Banned)
})?;
Some(account_status_failure(
classified_failure(
account_id,
account_status,
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
failure.explicit_status_code,
))
classify_account_failure(&super::failure_observation::from_response(failure))?,
)
}
fn model_unsupported_failure(account_id: &str, message: &str) -> ClassifiedFailure {
ClassifiedFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::ModelUnsupported,
message,
),
effect: None,
}
}
fn account_status_failure(
fn classified_failure(
account_id: &str,
status: AccountStatus,
message: &str,
status_code: Option<u16>,
) -> ClassifiedFailure {
let kind = match status {
AccountStatus::Expired => ExhaustedAccountKind::Expired,
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
AccountStatus::Banned => ExhaustedAccountKind::Banned,
AccountStatus::Active | AccountStatus::QuotaExhausted => {
unreachable!("account failure controller only invalidates unavailable statuses")
}
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
) -> Option<ClassifiedFailure> {
let kind = match classified.kind {
AccountFailureKind::ModelUnsupported => ExhaustedAccountKind::ModelUnsupported,
AccountFailureKind::Expired => ExhaustedAccountKind::Expired,
AccountFailureKind::Disabled => ExhaustedAccountKind::Disabled,
AccountFailureKind::Banned => ExhaustedAccountKind::Banned,
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => return None,
};
let mut exhaustion = AccountExhaustionRecord::new(account_id, kind, message);
if let Some(status_code) = status_code {
exhaustion = exhaustion.with_status_code(status_code);
}
ClassifiedFailure {
Some(ClassifiedFailure {
exhaustion,
effect: Some(AccountStateEffect::SetStatus(status)),
}
}
fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Option<AccountStatus> {
let code = code.trim().to_ascii_lowercase();
let error_type = error_type.trim().to_ascii_lowercase();
let message = message.to_ascii_lowercase();
if matches!(
code.as_str(),
"identity_verification_required" | "verification_required"
) || message.contains("identity verification is required")
{
return Some(AccountStatus::Disabled);
}
if matches!(
code.as_str(),
"account_banned"
| "account_deactivated"
| "account_disabled"
| "account_suspended"
| "deactivated_workspace"
| "organization_disabled"
| "workspace_deactivated"
) || message.contains("account is banned")
|| message.contains("account has been banned")
|| message.contains("account deactivated")
|| message.contains("account has been deactivated")
|| message.contains("account disabled")
|| message.contains("account has been disabled")
|| message.contains("account suspended")
|| message.contains("organization has been disabled")
|| message.contains("workspace has been deactivated")
|| message.contains("deactivated_workspace")
{
return Some(AccountStatus::Banned);
}
is_auth_failure(&code, &error_type, &message).then_some(AccountStatus::Expired)
}
fn is_auth_failure(code: &str, error_type: &str, message: &str) -> bool {
matches!(
code,
"token_invalid"
| "token_invalidated"
| "token_expired"
| "token_revoked"
| "refresh_token_invalidated"
| "unauthorized"
| "invalid_api_key"
| "authentication_error"
) || error_type == "authentication_error"
|| message.contains("token revoked")
|| message.contains("token invalidated")
|| message.contains("token invalid")
|| message.contains("token expired")
|| message.contains("unauthorized")
|| message.contains("invalid api key")
}
fn is_model_unsupported(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("model_not_supported")
|| value.contains("model_not_available")
|| (value.contains("model")
&& (value.contains("not supported")
|| value.contains("not available")
|| value.contains("not_supported")
|| value.contains("not_available")))
}
fn is_html(value: &str) -> bool {
let value = value.trim_start().to_ascii_lowercase();
value.starts_with("<!doctype") || value.starts_with("<html") || value.contains("<html")
effect: classified.effect,
})
}
@@ -1,38 +0,0 @@
//! 明确账号事实触发的运行时失效与连接驱逐。
use chrono::{DateTime, Utc};
use crate::{
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::transport::CodexBackendClient,
};
pub(super) enum AccountStateEffect {
SetStatus(AccountStatus),
MarkQuotaLimitedUntil(DateTime<Utc>),
}
pub(super) struct AccountStateEffects;
impl AccountStateEffects {
pub(super) async fn apply(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
effect: &AccountStateEffect,
) {
match effect {
AccountStateEffect::SetStatus(status) => {
account_pool
.set_status_immediately(account_id, *status)
.await;
}
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until_immediately(account_id, *until)
.await;
}
}
codex.evict_websocket_account(account_id).await;
}
}
@@ -9,13 +9,13 @@ use crate::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
lifecycle::stream::StreamTerminal,
transport::observation::UpstreamFailureFacts,
},
fleet::{
account::AccountStatus,
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
pool::AccountPoolService,
},
upstream::openai::failure::UpstreamFailureFacts,
};
pub(super) struct CloudflareController;
@@ -15,9 +15,11 @@ use crate::{
controllers::ControllerFailureFact,
errors::ClientFailure,
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
transport::observation::UpstreamFailureFacts,
},
upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
},
};
use self::{store::Store, types::SessionState};
@@ -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()
}
}
+267 -79
View File
@@ -13,10 +13,13 @@ use crate::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
AttemptRoutingFacts, CompleteResponseFacts, PinnedCandidateAcquireFailureKind,
},
transport::observation::{UpstreamFailureFacts, UpstreamFailureKind},
},
upstream::openai::{
protocol::responses::{CodexResponsesRequest, PreviousResponseScope, StreamCommitPolicy},
failure::{UpstreamFailureFacts, UpstreamFailureKind},
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,
+128 -55
View File
@@ -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},
@@ -49,17 +52,20 @@ use crate::{
},
stream::StreamTerminal,
},
transport::observation::UpstreamFailureFacts,
},
fleet::pool::AccountPoolService,
upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
fleet::scheduler::AttemptFeedback,
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
},
};
mod account_failure;
mod account_state;
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;
@@ -269,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,
@@ -298,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(
@@ -338,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(),
)))
}
@@ -383,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;
}
@@ -393,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(
@@ -452,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(
@@ -656,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(),
@@ -725,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)
@@ -749,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,
@@ -917,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(),
@@ -952,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(
+38 -254
View File
@@ -1,140 +1,29 @@
//! 请求前配额验证规则 owner。
use chrono::{Duration, Utc};
use serde_json::Value;
//! 配额失败规则 owner。
use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::{
ControllerFailureFact,
account_state::{AccountStateEffect, AccountStateEffects},
},
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
transport::observation::UpstreamFailureFacts,
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
},
fleet::{
pool::{AccountLease, AccountPoolService},
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_account_failure,
},
pool::AccountPoolService,
},
upstream::openai::{
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";
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
pub(super) enum QuotaEnterOutcome {
Ready(Box<AccountLease>),
LimitReached,
}
pub(super) struct QuotaEnterContext<'a> {
pub account_pool: &'a AccountPoolService,
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
pub cookie_header: Option<&'a str>,
}
pub(super) struct QuotaController;
impl QuotaController {
pub(super) async fn enter(
context: QuotaEnterContext<'_>,
acquired: AccountLease,
) -> QuotaEnterOutcome {
if !acquired.account.quota_verify_required {
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
let account_id = acquired.account.id.clone();
let identity = context.account_identity.scope_auxiliary(&account_id);
let usage = context
.codex
.fetch_usage(CodexRequestContext {
access_token: &acquired.account.access_token,
account_id: acquired.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: None,
parent_thread_id: None,
cookie_header: context.cookie_header,
installation_id: Some(&identity.installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
})
.await;
let raw = match usage {
Ok(raw) => raw,
Err(error) => {
tracing::warn!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "upstream_error",
retry_with_another_account = false,
error = %error,
"Failed to verify stale quota state before upstream request"
);
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
};
let quota = quota_from_usage(&raw);
context
.account_pool
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
acquired.release_without_usage().await;
tracing::info!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "limit_reached",
retry_with_another_account = true,
"Quota verification reported exhausted account before upstream request"
);
return QuotaEnterOutcome::LimitReached;
}
QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, &quota)))
}
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
if matches!(
observation.kind,
AttemptObservationKind::CandidatePreparationRejected
) {
return observation
.account
.as_ref()
.map(|account| ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account.id.clone(),
ExhaustedAccountKind::RateLimited,
LIMIT_REACHED_MESSAGE,
),
effect: None,
});
}
classify(observation)
}
@@ -167,7 +56,7 @@ impl QuotaController {
return;
};
if let Some(effect) = &classified.effect {
AccountStateEffects::apply(account_pool, codex, account_id, effect).await;
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
}
}
@@ -201,145 +90,40 @@ fn classify_upstream(
account_id: &str,
facts: &UpstreamFailureFacts,
) -> Option<ClassifiedQuotaFailure> {
let signaled_kind = facts
.code
.as_deref()
.and_then(classify_quota_signal)
.or_else(|| facts.error_type.as_deref().and_then(classify_quota_signal))
.or_else(|| classify_quota_message(&facts.message));
match signaled_kind.or(match facts.status_code {
Some(429) => Some(StreamQuotaKind::RateLimited),
Some(402) => Some(StreamQuotaKind::QuotaExhausted),
_ => None,
})? {
StreamQuotaKind::RateLimited => Some(rate_limited_failure(
account_id,
facts.body.clone(),
facts.retry_after_seconds,
)),
StreamQuotaKind::QuotaExhausted => {
Some(quota_exhausted_failure(account_id, facts.body.clone()))
}
}
classified_quota_failure(
account_id,
&facts.body,
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
)
}
fn classify_stream(
account_id: &str,
failure: &ResponsesSseFailure,
) -> Option<ClassifiedQuotaFailure> {
let kind = classify_stream_signal(failure).or(match failure.explicit_status_code {
Some(429) => Some(StreamQuotaKind::RateLimited),
Some(402) => Some(StreamQuotaKind::QuotaExhausted),
_ => None,
})?;
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
Some(match kind {
StreamQuotaKind::RateLimited => {
rate_limited_failure(account_id, body, failure.retry_after_seconds)
}
StreamQuotaKind::QuotaExhausted => quota_exhausted_failure(account_id, body),
classified_quota_failure(
account_id,
&body,
classify_account_failure(&super::failure_observation::from_response(failure))?,
)
}
fn classified_quota_failure(
account_id: &str,
message: &str,
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
) -> Option<ClassifiedQuotaFailure> {
let kind = match classified.kind {
AccountFailureKind::RateLimited => ExhaustedAccountKind::RateLimited,
AccountFailureKind::QuotaExhausted => ExhaustedAccountKind::QuotaExhausted,
AccountFailureKind::ModelUnsupported
| AccountFailureKind::Expired
| AccountFailureKind::Disabled
| AccountFailureKind::Banned => return None,
};
Some(ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(account_id, kind, message),
effect: classified.effect,
})
}
#[derive(Clone, Copy)]
enum StreamQuotaKind {
RateLimited,
QuotaExhausted,
}
fn classify_stream_signal(failure: &ResponsesSseFailure) -> Option<StreamQuotaKind> {
failure
.upstream_code
.as_deref()
.and_then(classify_quota_signal)
.or_else(|| {
failure
.upstream_type
.as_deref()
.and_then(classify_quota_signal)
})
.or_else(|| classify_quota_message(&failure.message))
}
fn classify_quota_signal(signal: &str) -> Option<StreamQuotaKind> {
let signal = signal.trim().to_ascii_lowercase();
match signal.as_str() {
"usage_limit_reached"
| "rate_limit_exceeded"
| "rate_limit_reached"
| "rate_limit_error"
| "workspace_owner_usage_limit_reached"
| "workspace_member_usage_limit_reached" => Some(StreamQuotaKind::RateLimited),
"quota_exhausted"
| "quota_exceeded"
| "payment_required"
| "insufficient_quota"
| "workspace_owner_credits_depleted"
| "workspace_member_credits_depleted" => Some(StreamQuotaKind::QuotaExhausted),
signal if signal.starts_with("billing_limit") => Some(StreamQuotaKind::QuotaExhausted),
_ => None,
}
}
fn classify_quota_message(message: &str) -> Option<StreamQuotaKind> {
let message = message.to_ascii_lowercase();
if message.contains("rate limit") || message.contains("usage limit") {
return Some(StreamQuotaKind::RateLimited);
}
if message.contains("quota")
|| message.contains("payment required")
|| message.contains("billing limit")
{
return Some(StreamQuotaKind::QuotaExhausted);
}
None
}
fn rate_limited_failure(
account_id: &str,
message: String,
retry_after_seconds: Option<u64>,
) -> ClassifiedQuotaFailure {
let seconds = retry_after_seconds
.unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS)
.min(i64::MAX as u64) as i64;
ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::RateLimited,
message,
),
effect: Some(AccountStateEffect::MarkQuotaLimitedUntil(
Utc::now() + Duration::seconds(seconds),
)),
}
}
fn quota_exhausted_failure(account_id: &str, message: String) -> ClassifiedQuotaFailure {
ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::QuotaExhausted,
message,
),
effect: Some(AccountStateEffect::SetStatus(
crate::fleet::account::AccountStatus::QuotaExhausted,
)),
}
}
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
let limit_reached = quota_snapshot_limit_reached(quota);
acquired.account.quota_verify_required = false;
acquired.account.quota_limit_reached = limit_reached;
acquired.account.quota_cooldown_until = limit_reached
.then_some(quota_snapshot_reset_at(quota))
.flatten();
if let Some(reset_at) = quota_snapshot_reset_at(quota) {
acquired.account.window_reset_at = Some(reset_at);
if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) {
acquired.account.limit_window_seconds = Some(limit_window_seconds);
}
}
acquired
}
@@ -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
@@ -546,6 +864,9 @@ pub(super) async fn record_live_response_stream_event(
status_code,
body,
);
if !is_success_usage_event(level, status_code) {
ensure_stream_failure_class(&mut metadata);
}
log_live_response_stream_finalized(context, status_code, level, message, &metadata);
enrich_usage_record_identity(
&mut metadata,
@@ -559,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());
@@ -591,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");
}
@@ -601,6 +922,19 @@ fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
level != UsageRecordLevel::Error && status_code < 400
}
fn ensure_stream_failure_class(metadata: &mut Value) {
let failure_class = ["failureClass", "upstreamCode", "terminal"]
.into_iter()
.find_map(|field| metadata_string_field(metadata, field))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string);
let (Some(failure_class), Some(object)) = (failure_class, metadata.as_object_mut()) else {
return;
};
object.insert("failureClass".to_string(), Value::String(failure_class));
}
fn log_live_response_stream_finalized(
context: &StreamContext<'_>,
status_code: i64,
@@ -667,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;
}
+98 -11
View File
@@ -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,
}
}
+9
View File
@@ -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(),
+23 -54
View File
@@ -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,
@@ -21,7 +20,6 @@ use crate::{
create_response_with_account, prepare_response_transport_with_account,
},
canonical::{CanonicalStreamDecoder, normalize_complete_response},
observation::upstream_failure_facts,
prefetch::{PrefetchedStream, StreamPrefetchError, prefetch_until_commit},
},
},
@@ -30,6 +28,7 @@ use crate::{
pool::{AccountCandidateLease, AccountPoolService},
},
upstream::openai::{
failure::upstream_failure_facts,
protocol::responses::{
CodexResponsesRequest, CollectedResponse, TransportRequirement, transport_requirement,
},
@@ -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",
+3 -8
View File
@@ -6,13 +6,11 @@ use crate::{
dispatch::{
errors::{ClientFailure, ResponseDispatchError},
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccount},
transport::{
canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
observation::UpstreamFailureFacts,
},
transport::canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
},
fleet::{account::Account, pool::AccountLease},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::{
responses::{CodexResponsesRequest, CollectedResponse, ResponsesSseFailure},
sse::SseError,
@@ -95,7 +93,6 @@ pub(in crate::dispatch) enum AttemptObservationKind {
ledger: CandidateLedgerFacts,
last_exhausted: Option<ExhaustedAccount>,
},
CandidatePreparationRejected,
RoutePreparationRejected {
message: String,
},
@@ -175,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,
@@ -258,9 +256,6 @@ pub(super) enum PendingAttempt {
account_id: String,
kind: PinnedCandidateAcquireFailureKind,
},
CandidatePreparationRejected {
account: Account,
},
RoutePreparationRejected {
account: Account,
message: String,
+24 -2
View File
@@ -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;
}
+1 -1
View File
@@ -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();
+2
View File
@@ -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>,
+8
View File
@@ -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
}
}
+20 -5
View File
@@ -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,
+5
View File
@@ -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(
+54 -9
View File
@@ -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),
+26 -37
View File
@@ -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());
}
+1 -2
View File
@@ -1,6 +1,5 @@
//! Attempt 阶段的账号作用域上游调用与事实规范化。
//! Attempt 阶段的账号作用域上游调用与流规范化。
pub(in crate::dispatch) mod account;
pub(crate) mod canonical;
pub(in crate::dispatch) mod observation;
pub(in crate::dispatch) mod prefetch;
+284
View File
@@ -0,0 +1,284 @@
//! 上游失败到账号状态 effect 的唯一领域分类。
use chrono::{Duration, Utc};
use crate::fleet::{
account::AccountStatus,
account_gateway::{AccountFailureObservation, AccountUpstreamGateway},
pool::AccountPoolService,
};
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountFailureKind {
ModelUnsupported,
Expired,
Disabled,
Banned,
QuotaExhausted,
RateLimited,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AccountStateEffect {
SetStatus(AccountStatus),
MarkQuotaLimitedUntil(chrono::DateTime<Utc>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClassifiedAccountFailure {
pub kind: AccountFailureKind,
pub effect: Option<AccountStateEffect>,
}
impl ClassifiedAccountFailure {
pub fn account_status(&self) -> Option<AccountStatus> {
match self.effect.as_ref()? {
AccountStateEffect::SetStatus(status) => Some(*status),
AccountStateEffect::MarkQuotaLimitedUntil(_) => None,
}
}
}
pub fn classify_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 async fn apply_account_state_effect_immediately(
account_pool: &AccountPoolService,
upstream: &dyn AccountUpstreamGateway,
account_id: &str,
effect: &AccountStateEffect,
) {
match effect {
AccountStateEffect::SetStatus(status) => {
account_pool
.set_status_immediately(account_id, *status)
.await;
}
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until_immediately(account_id, *until)
.await;
}
}
upstream.evict_account_connections(account_id).await;
}
pub async fn apply_account_state_effect(
account_pool: &AccountPoolService,
upstream: &dyn AccountUpstreamGateway,
account_id: &str,
effect: &AccountStateEffect,
) -> bool {
let updated = match effect {
AccountStateEffect::SetStatus(status) => account_pool.set_status(account_id, *status).await,
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until(account_id, *until)
.await
}
};
upstream.evict_account_connections(account_id).await;
updated
}
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>,
}
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()
.any(is_model_unsupported)
{
return Some(classified(AccountFailureKind::ModelUnsupported, None));
}
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(401) {
return Some(status_failure(AccountStatus::Expired));
}
let quota_kind = [code, error_type]
.into_iter()
.find_map(classify_quota_signal)
.or_else(|| classify_quota_message(fields.message))
.or_else(|| classify_quota_message(fields.body))
.or(match fields.status_code {
Some(429) => Some(AccountFailureKind::RateLimited),
Some(402) => Some(AccountFailureKind::QuotaExhausted),
_ => None,
})?;
Some(match quota_kind {
AccountFailureKind::RateLimited => {
let seconds = fields
.retry_after_seconds
.unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS)
.min(i64::MAX as u64) as i64;
classified(
AccountFailureKind::RateLimited,
Some(AccountStateEffect::MarkQuotaLimitedUntil(
Utc::now() + Duration::seconds(seconds),
)),
)
}
AccountFailureKind::QuotaExhausted => classified(
AccountFailureKind::QuotaExhausted,
Some(AccountStateEffect::SetStatus(AccountStatus::QuotaExhausted)),
),
AccountFailureKind::ModelUnsupported
| AccountFailureKind::Expired
| AccountFailureKind::Disabled
| AccountFailureKind::Banned => unreachable!("quota classifier returned account status"),
})
}
fn status_failure(status: AccountStatus) -> ClassifiedAccountFailure {
let kind = match status {
AccountStatus::Expired => AccountFailureKind::Expired,
AccountStatus::Disabled => AccountFailureKind::Disabled,
AccountStatus::Banned => AccountFailureKind::Banned,
AccountStatus::Active | AccountStatus::QuotaExhausted => {
unreachable!("explicit account failure must make the account unavailable")
}
};
classified(kind, Some(AccountStateEffect::SetStatus(status)))
}
fn classified(
kind: AccountFailureKind,
effect: Option<AccountStateEffect>,
) -> ClassifiedAccountFailure {
ClassifiedAccountFailure { kind, effect }
}
fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Option<AccountStatus> {
let code = code.trim().to_ascii_lowercase();
let error_type = error_type.trim().to_ascii_lowercase();
let message = message.to_ascii_lowercase();
if matches!(
code.as_str(),
"identity_verification_required" | "verification_required"
) || message.contains("identity verification is required")
{
return Some(AccountStatus::Disabled);
}
if matches!(
code.as_str(),
"account_banned"
| "account_deactivated"
| "account_disabled"
| "account_suspended"
| "deactivated_workspace"
| "organization_disabled"
| "workspace_deactivated"
) || message.contains("account is banned")
|| message.contains("account has been banned")
|| message.contains("account deactivated")
|| message.contains("account has been deactivated")
|| message.contains("account disabled")
|| message.contains("account has been disabled")
|| message.contains("account suspended")
|| message.contains("organization has been disabled")
|| message.contains("workspace has been deactivated")
|| message.contains("deactivated_workspace")
{
return Some(AccountStatus::Banned);
}
is_auth_failure(&code, &error_type, &message).then_some(AccountStatus::Expired)
}
fn is_auth_failure(code: &str, error_type: &str, message: &str) -> bool {
matches!(
code,
"token_invalid"
| "token_invalidated"
| "token_expired"
| "token_revoked"
| "refresh_token_invalidated"
| "unauthorized"
| "invalid_api_key"
| "authentication_error"
) || error_type == "authentication_error"
|| message.contains("token revoked")
|| message.contains("token invalidated")
|| message.contains("token invalid")
|| message.contains("token expired")
|| message.contains("unauthorized")
|| message.contains("invalid api key")
}
fn classify_quota_signal(signal: &str) -> Option<AccountFailureKind> {
let signal = signal.trim().to_ascii_lowercase();
match signal.as_str() {
"usage_limit_reached"
| "rate_limit_exceeded"
| "rate_limit_reached"
| "rate_limit_error"
| "workspace_owner_usage_limit_reached"
| "workspace_member_usage_limit_reached" => Some(AccountFailureKind::RateLimited),
"quota_exhausted"
| "quota_exceeded"
| "payment_required"
| "insufficient_quota"
| "workspace_owner_credits_depleted"
| "workspace_member_credits_depleted" => Some(AccountFailureKind::QuotaExhausted),
signal if signal.starts_with("billing_limit") => Some(AccountFailureKind::QuotaExhausted),
_ => None,
}
}
fn classify_quota_message(message: &str) -> Option<AccountFailureKind> {
let message = message.to_ascii_lowercase();
if message.contains("rate limit") || message.contains("usage limit") {
return Some(AccountFailureKind::RateLimited);
}
if message.contains("quota")
|| message.contains("payment required")
|| message.contains("billing limit")
{
return Some(AccountFailureKind::QuotaExhausted);
}
None
}
fn is_model_unsupported(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("model_not_supported")
|| value.contains("model_not_available")
|| (value.contains("model")
&& (value.contains("not supported")
|| value.contains("not available")
|| value.contains("not_supported")
|| value.contains("not_available")))
}
+98
View File
@@ -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);
}
+18 -31
View File
@@ -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,
+1 -1
View File
@@ -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::{
+1 -1
View File
@@ -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::*;
+151 -267
View File
@@ -1,64 +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, store::StoredAccount},
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{SseEvent, encode_sse_event, parse_sse_events},
},
transport::{
CodexClientError, CodexRequestContext, is_banned_auth_signal, is_banned_upstream_error,
fleet::{
account::AccountStatus,
account_failure::classify_account_failure,
account_gateway::{
AccountGatewayError, AccountProbeEvent, AccountProbeRequest, AccountProbeSession,
AccountUpstreamContext,
},
store::StoredAccount,
},
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,
},
}
}
}
@@ -74,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);
@@ -86,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(),
@@ -94,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(
@@ -119,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))
}
@@ -221,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,
@@ -235,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
{
@@ -251,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(),
@@ -268,168 +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> {
if is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else {
match error {
CodexClientError::Upstream { status, .. } if status.as_u16() == 402 => {
Some(AccountStatus::QuotaExhausted)
}
CodexClientError::Upstream { status, body, .. } if status.as_u16() == 401 => {
Some(if is_banned_auth_signal(body) {
AccountStatus::Banned
} else {
AccountStatus::Expired
})
}
_ => None,
}
}
}
fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
let code = failure
.upstream_code
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let message = failure.message.to_ascii_lowercase();
if matches!(code.as_str(), "quota_exceeded" | "insufficient_quota") || message.contains("quota")
{
return Some(AccountStatus::QuotaExhausted);
}
let authentication_failed = matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
) || message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired");
authentication_failed.then(|| {
if is_banned_auth_signal(&failure.message) {
AccountStatus::Banned
} else {
AccountStatus::Expired
}
})
}
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;
}
+40 -31
View File
@@ -1,6 +1,16 @@
use secrecy::ExposeSecret;
use serde_json::Value;
use super::{AccountManageService, types::AccountManageError};
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> {
@@ -23,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)
@@ -31,39 +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 result = match self.upstream.fetch_usage(context).await {
Ok(result) => result,
Err(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 raw = self
.codex
.fetch_usage(context)
.await
.map_err(|e| AccountManageError::FetchQuota(e.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,
})
}
}
-471
View File
@@ -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(&quota);
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")
}
+8 -8
View File
@@ -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;
}
}
+12 -27
View File
@@ -4,12 +4,12 @@ use chrono::{DateTime, Utc};
use serde::Serialize;
use thiserror::Error;
use crate::{
fleet::{
account::{AccountStatus, AccountStatus as AcctStatus},
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)]
@@ -85,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")]
@@ -295,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")
@@ -322,12 +309,10 @@ 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> {
if crate::upstream::openai::transport::is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else {
None
}
pub(super) fn import_status_from_usage_error(error: &AccountGatewayError) -> Option<AccountStatus> {
classify_account_failure(error.failure()?)?
.account_status()
.filter(|status| *status == AccountStatus::Banned)
}
pub(super) fn parse_account_status(status: &str) -> Result<AcctStatus, AccountManageError> {
+3
View File
@@ -1,6 +1,8 @@
//! 账号、调度、配额与刷新领域。
pub mod account;
pub mod account_failure;
pub mod account_gateway;
pub mod cookies;
pub mod import;
pub mod manage;
@@ -9,4 +11,5 @@ pub mod quota;
pub mod refresh;
pub mod scheduler;
pub mod store;
pub mod usage;
mod window;
+46 -55
View File
@@ -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>(&quota_json).ok(),
Ok(Some(quota_json)) => QuotaSnapshot::from_json(&quota_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, &quota).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, &quota_json, limit_reached, cooldown_until)
+7 -4
View File
@@ -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);
}
/// 更新账号状态。
+4
View File
@@ -1,7 +1,11 @@
//! 账号配额解析与刷新。
mod model;
mod observation;
mod runtime;
mod service;
pub use model::*;
pub use observation::*;
pub use runtime::*;
pub use service::*;
+673
View File
@@ -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))
}
+238
View File
@@ -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"
)
}
+50 -60
View File
@@ -7,25 +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_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;
@@ -34,33 +31,34 @@ 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>,
cookie_store: Option<PgCookieStore>,
account_pool: Option<Arc<AccountPoolService>>,
account_pool: Arc<AccountPoolService>,
}
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,
cookie_store: None,
account_pool: None,
account_pool,
}
}
@@ -70,21 +68,22 @@ 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,
) -> Self {
Self {
store,
usage_store,
codex,
upstream,
min_refresh_interval_secs,
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
cookie_store: None,
account_pool: None,
account_pool,
}
}
@@ -100,12 +99,6 @@ impl QuotaRefreshService {
self
}
/// 设置运行时账号池,用于刷新后同步内存调度状态。
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.account_pool = Some(account_pool);
self
}
/// 执行一次配额锁定或待验证账号刷新,复用调用方传入的最近刷新记录。
pub async fn refresh_locked_accounts(
&self,
@@ -156,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(&quota);
let reset_at = quota_snapshot_reset_at(&quota);
let quota_json = quota.to_json()?;
if self
.store
.apply_quota_snapshot(
&account.id,
&quota.to_string(),
&quota_json,
limit_reached,
limit_reached.then_some(reset_at).flatten(),
)
@@ -211,6 +193,17 @@ impl QuotaRefreshService {
self.sync_runtime_pool_account(&account.id).await;
}
Err(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;
}
warn!(
account_id = %account.id,
error = %error,
@@ -238,10 +231,7 @@ impl QuotaRefreshService {
}
async fn sync_runtime_pool_account(&self, account_id: &str) {
let Some(account_pool) = &self.account_pool else {
return;
};
if let Err(error) = account_pool.sync_account_from_store(account_id).await {
if let Err(error) = self.account_pool.sync_account_from_store(account_id).await {
warn!(
account_id = %account_id,
error = %error,
@@ -288,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 服务结果。
+8 -403
View File
@@ -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()
}
+26
View File
@@ -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>;
}
+2
View File
@@ -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::*;
+3 -10
View File
@@ -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());
}
}
+1 -1
View File
@@ -19,13 +19,13 @@ use tokio_util::task::TaskTracker;
use tracing::{debug, info, warn};
use uuid::Uuid;
use super::TokenRefresher;
use crate::fleet::{
account::{Account, AccountStatus},
store::{
AccountClaimsUpdate, AccountStore, PgAccountStore, PgAccountStoreError, StoredAccount,
},
};
use crate::upstream::openai::token_client::TokenRefresher;
use super::{
lease::{RedisRefreshLeaseStore, RedisRefreshLeaseStoreError},
+68 -23
View File
@@ -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);
}
/// 移除指定账号的反馈记录(账号被删除时调用)。
+5 -5
View File
@@ -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 {
+5 -7
View File
@@ -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"),
+125
View File
@@ -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
}
}
+30 -1
View File
@@ -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
}
+12 -5
View File
@@ -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);
+36
View File
@@ -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
View File
@@ -1,6 +1,7 @@
//! 模型目录领域。
pub mod catalog;
pub mod gateway;
pub mod service;
pub mod store;
pub mod types;
+32 -18
View File
@@ -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
}
+31 -70
View File
@@ -6,7 +6,7 @@ use thiserror::Error;
use crate::infra::format::nonnegative_i64_to_u64;
use crate::telemetry::{
account_usage::store::{PgAccountUsageStore, UsageListRecord, UsageSummary},
account_usage::store::{PgAccountUsageStore, UsageListRecord},
billing,
buckets::query::{
ModelBucketUsage, ModelUsageWindow, PgRequestBucketQuery, UsageBucketTotals,
@@ -42,13 +42,26 @@ impl AccountUsageQueryService {
.map_err(|_| AccountUsageQueryError::List)
}
/// 汇总账号用量。
pub async fn summary(&self) -> Result<AccountUsageSummary, AccountUsageQueryError> {
self.store
.usage_summary()
/// 汇总请求时间桶保留期内的全局用量与计费。
pub async fn retained_summary(&self) -> Result<RetainedUsageSummary, AccountUsageQueryError> {
let buckets = self
.buckets
.retained_usage_buckets()
.await
.map(AccountUsageSummary::from)
.map_err(|_| AccountUsageQueryError::Summary)
.map_err(|_| AccountUsageQueryError::RetainedSummary)?;
let mut summary = RetainedUsageSummary::default();
for bucket in buckets {
summary.billing_amount_usd += billing::calculate_billing_amount(
nonnegative_i64_to_u64(bucket.totals.input_tokens),
nonnegative_i64_to_u64(bucket.totals.output_tokens),
nonnegative_i64_to_u64(bucket.totals.cached_tokens),
nonnegative_i64_to_u64(bucket.totals.cache_write_tokens),
&bucket.model,
bucket.service_tier.as_deref(),
);
summary.totals.add(bucket.totals);
}
Ok(summary)
}
/// 列出指定时间范围内的聚合时间桶。
@@ -95,12 +108,19 @@ impl AccountUsageQueryService {
pub enum AccountUsageQueryError {
#[error("failed to list account usage")]
List,
#[error("failed to summarize account usage")]
Summary,
#[error("failed to summarize retained usage")]
RetainedSummary,
#[error("failed to list usage time buckets")]
TimeBuckets,
}
/// 请求时间桶保留期内的全局用量与计费汇总。
#[derive(Debug, Clone, Default, PartialEq)]
pub struct RetainedUsageSummary {
pub totals: UsageBucketTotals,
pub billing_amount_usd: f64,
}
/// 管理端账号用量记录。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUsageRecord {
@@ -129,23 +149,6 @@ pub struct AccountUsageRecord {
pub last_used_at: Option<DateTime<Utc>>,
}
/// 管理端账号用量汇总。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageSummary {
pub account_count: i64,
pub request_count: i64,
pub empty_response_count: i64,
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub reasoning_tokens: i64,
pub total_tokens: i64,
pub image_input_tokens: i64,
pub image_output_tokens: i64,
pub image_request_count: i64,
pub image_request_failed_count: i64,
}
/// 管理端时间桶用量记录。
#[derive(Debug, Clone, PartialEq)]
pub struct AccountUsageTimeBucket {
@@ -157,37 +160,17 @@ 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,
pub latency_count: i64,
pub max_latency_ms: i64,
pub min_latency_ms: i64,
pub billing_amount_usd: Option<f64>,
}
impl From<UsageSummary> for AccountUsageSummary {
fn from(s: UsageSummary) -> Self {
Self {
account_count: s.account_count,
request_count: s.request_count,
empty_response_count: s.empty_response_count,
input_tokens: s.input_tokens,
output_tokens: s.output_tokens,
cached_tokens: s.cached_tokens,
reasoning_tokens: s.reasoning_tokens,
total_tokens: s.total_tokens,
image_input_tokens: s.image_input_tokens,
image_output_tokens: s.image_output_tokens,
image_request_count: s.image_request_count,
image_request_failed_count: s.image_request_failed_count,
}
}
}
impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
fn from(record: UsageTimeBucketRecord) -> Self {
let billing_amount_usd = time_bucket_billing_amount_usd(&record);
Self {
bucket_start: record.bucket_start,
model: record.model,
@@ -197,39 +180,17 @@ 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,
latency_count: record.latency_count,
max_latency_ms: record.max_latency_ms,
min_latency_ms: record.min_latency_ms,
billing_amount_usd,
}
}
}
fn time_bucket_billing_amount_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(record.input_tokens);
let output_tokens = nonnegative_i64_to_u64(record.output_tokens);
let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens);
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
return None;
}
let model = record.model.trim();
if model.is_empty() {
return None;
}
Some(billing::calculate_billing_amount(
input_tokens,
output_tokens,
cached_tokens,
model,
record.service_tier.as_deref(),
))
}
impl From<UsageListRecord> for AccountUsageRecord {
fn from(usage: UsageListRecord) -> Self {
Self {
+8 -165
View File
@@ -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 (
@@ -129,22 +135,6 @@ select
from account_usage au
left join accounts a on a.id = au.account_id";
const USAGE_SUMMARY_SQL: &str = r"
select
count(*) as account_count,
coalesce(sum(request_count), 0)::bigint as request_count,
coalesce(sum(empty_response_count), 0)::bigint as empty_response_count,
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(reasoning_tokens), 0)::bigint as reasoning_tokens,
coalesce(sum(total_tokens), 0)::bigint as total_tokens,
coalesce(sum(image_input_tokens), 0)::bigint as image_input_tokens,
coalesce(sum(image_output_tokens), 0)::bigint as image_output_tokens,
coalesce(sum(image_request_count), 0)::bigint as image_request_count,
coalesce(sum(image_request_failed_count), 0)::bigint as image_request_failed_count
from account_usage";
/// PostgreSQL 用量存储错误。
#[derive(Debug, Error)]
pub enum PgAccountUsageStoreError {
@@ -156,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)
}
}
@@ -312,35 +203,6 @@ pub struct UsageListRecord {
pub last_used_at: Option<DateTime<Utc>>,
}
/// 账号用量汇总。
#[derive(Debug, Clone)]
pub struct UsageSummary {
/// 账号数。
pub account_count: i64,
/// 总请求数。
pub request_count: i64,
/// 总空响应数。
pub empty_response_count: i64,
/// 总输入 token。
pub input_tokens: i64,
/// 总输出 token。
pub output_tokens: i64,
/// 总缓存 token。
pub cached_tokens: i64,
/// 总 reasoning token。
pub reasoning_tokens: i64,
/// 总 token。
pub total_tokens: i64,
/// 总图片输入 token。
pub image_input_tokens: i64,
/// 总图片输出 token。
pub image_output_tokens: i64,
/// 总图片请求成功数。
pub image_request_count: i64,
/// 总图片请求失败数。
pub image_request_failed_count: i64,
}
/// PostgreSQL 用量存储。
#[derive(Clone)]
pub struct PgAccountUsageStore {
@@ -492,25 +354,6 @@ where account_id in (",
let rows = builder.build().fetch_all(&self.pool).await?;
rows.iter().map(usage_list_from_row).collect()
}
/// 汇总账号用量。
pub async fn usage_summary(&self) -> PgAccountUsageStoreResult<UsageSummary> {
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
Ok(UsageSummary {
account_count: row.get("account_count"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
reasoning_tokens: row.get("reasoning_tokens"),
total_tokens: row.get("total_tokens"),
image_input_tokens: row.get("image_input_tokens"),
image_output_tokens: row.get("image_output_tokens"),
image_request_count: row.get("image_request_count"),
image_request_failed_count: row.get("image_request_failed_count"),
})
}
}
#[async_trait]
+55 -6
View File
@@ -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,
+57 -2
View File
@@ -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,
@@ -29,6 +30,18 @@ where bucket_start >= $1 and bucket_start <= $2
group by bucket_start, model, service_tier
order by bucket_start asc, model asc, service_tier asc";
const RETAINED_USAGE_BUCKETS_SQL: &str = r"
select
model,
nullif(service_tier, '__unknown__') as service_tier,
coalesce(sum(success_count + error_count), 0)::bigint as request_count,
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens
from request_time_buckets
group by bucket_start, model, service_tier";
/// 请求时间桶查询错误。
#[derive(Debug, Error)]
pub enum PgRequestBucketQueryError {
@@ -48,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,
@@ -70,14 +84,26 @@ pub struct UsageBucketTotals {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub cache_write_tokens: i64,
}
/// 保留期内单个计费桶的用量事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedUsageBucket {
pub model: String,
pub service_tier: Option<String>,
pub totals: UsageBucketTotals,
}
impl UsageBucketTotals {
fn add(&mut self, other: Self) {
pub(crate) fn add(&mut self, other: Self) {
self.request_count = self.request_count.saturating_add(other.request_count);
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
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);
}
}
@@ -98,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>>,
}
@@ -127,6 +154,29 @@ impl PgRequestBucketQuery {
Ok(rows.iter().map(usage_time_bucket_from_row).collect())
}
/// 列出请求时间桶保留期内的全局计费桶,不受当前账号集合影响。
pub async fn retained_usage_buckets(
&self,
) -> Result<Vec<RetainedUsageBucket>, PgRequestBucketQueryError> {
let rows = sqlx::query(RETAINED_USAGE_BUCKETS_SQL)
.fetch_all(&self.pool)
.await?;
Ok(rows
.iter()
.map(|row| RetainedUsageBucket {
model: row.get("model"),
service_tier: row.get("service_tier"),
totals: UsageBucketTotals {
request_count: row.get("request_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
cache_write_tokens: row.get("cache_write_tokens"),
},
})
.collect())
}
pub async fn usage_by_windows(
&self,
windows: &[UsageBucketWindow],
@@ -157,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 (",
);
@@ -182,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;
@@ -217,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 (",
@@ -249,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"),
}
})
@@ -266,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"),
+9 -4
View File
@@ -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",
-1
View File
@@ -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;

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