Compare commits

...
38 Commits
Author SHA1 Message Date
ZyyandCodex 3920624862 chore: bump version to 3.3.3
Co-authored-by: Codex <noreply@openai.com>
2026-08-27 01:53:05 +08:00
Zyy d804bca495 fix(openai): strip unsupported Codex request fields 2026-08-27 01:51:30 +08:00
ZyyandCodex ca527d5833 fix(accounts): use personal usage reports only
Remove the workspace token-breakdown probe and its dedicated statistics mode so usage aggregation always uses the personal model-credit and token-total reports.

Drop the invariant mode field from provider projections, the admin response, and the account usage UI.

Co-authored-by: Codex <noreply@openai.com>
2026-08-26 11:31:55 +08:00
ZyyandCodex ae4eb49afc chore: bump version to 3.3.2
Co-authored-by: Codex <noreply@openai.com>
2026-08-26 00:51:31 +08:00
ZyyandCodex 5d97086511 feat(accounts): add official usage statistics
Fetch Codex cycle usage in the provider, calculate model and daily Token costs with the shared OpenAI pricing rules, and expose the result through the admin API.

Add the account usage statistics modal, structured loading state, and feature-scoped quota components.

Co-authored-by: Codex <noreply@openai.com>
2026-08-26 00:46:53 +08:00
ZyyandCodex 0e6da1e731 refactor: consolidate migration baseline and remove design source
Fold the nullable requested-model change into the initial migration and keep the frozen manifest and project guidance aligned.

Remove the staged Pen design source.

Co-authored-by: Codex <noreply@openai.com>
2026-08-25 17:53:29 +08:00
ZyyandCodex 5d0922b9d1 fix(theme): align functional colors with Ant Design
Keep functional base tokens on Ant Design P6 so light themes preserve the configured seed while dark themes retain palette adaptation.

Co-authored-by: Codex <noreply@openai.com>
2026-08-25 17:06:05 +08:00
ZyyandCodex ed3d622554 chore: bump version to 3.3.1
Co-authored-by: Codex <noreply@openai.com>
2026-08-25 15:53:32 +08:00
dependabot[bot] 0f71cf6df2 chore(deps): bump rust in /deploy in the minor-and-patch group (#46)
Bumps the minor-and-patch group in /deploy with 1 update: rust.


Updates `rust` from 1.97-bookworm to 1.98-bookworm

---
updated-dependencies:
- dependency-name: rust
  dependency-version: 1.98-bookworm
  dependency-type: direct:production
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-08-25 15:37:44 +08:00
dependabot[bot] a1d3eba874 chore(deps): bump the minor-and-patch group in /frontend with 8 updates (#45)
Bumps the minor-and-patch group in /frontend with 8 updates:

| Package | From | To |
| --- | --- | --- |
| [@lucide/vue](https://github.com/lucide-icons/lucide/tree/HEAD/packages/vue) | `1.31.0` | `1.33.0` |
| [dayjs](https://github.com/iamkun/dayjs) | `1.11.21` | `1.11.23` |
| [dompurify](https://github.com/cure53/DOMPurify) | `3.4.13` | `3.4.14` |
| [es-toolkit](https://github.com/toss/es-toolkit) | `1.50.0` | `1.51.0` |
| [marked](https://github.com/markedjs/marked) | `18.0.9` | `18.0.10` |
| [eslint](https://github.com/eslint/eslint) | `10.8.1` | `10.9.0` |
| [vite](https://github.com/vitejs/vite/tree/HEAD/packages/vite) | `8.2.1` | `8.2.2` |
| [vue-tsc](https://github.com/vuejs/language-tools/tree/HEAD/packages/tsc) | `3.3.10` | `3.3.11` |


Updates `@lucide/vue` from 1.31.0 to 1.33.0
- [Release notes](https://github.com/lucide-icons/lucide/releases)
- [Commits](https://github.com/lucide-icons/lucide/commits/1.33.0/packages/vue)

Updates `dayjs` from 1.11.21 to 1.11.23
- [Release notes](https://github.com/iamkun/dayjs/releases)
- [Changelog](https://github.com/iamkun/dayjs/blob/v1.11.23/CHANGELOG.md)
- [Commits](https://github.com/iamkun/dayjs/compare/v1.11.21...v1.11.23)

Updates `dompurify` from 3.4.13 to 3.4.14
- [Release notes](https://github.com/cure53/DOMPurify/releases)
- [Commits](https://github.com/cure53/DOMPurify/compare/3.4.13...3.4.14)

Updates `es-toolkit` from 1.50.0 to 1.51.0
- [Release notes](https://github.com/toss/es-toolkit/releases)
- [Changelog](https://github.com/toss/es-toolkit/blob/main/CHANGELOG.md)
- [Commits](https://github.com/toss/es-toolkit/compare/v1.50.0...v1.51.0)

Updates `marked` from 18.0.9 to 18.0.10
- [Release notes](https://github.com/markedjs/marked/releases)
- [Commits](https://github.com/markedjs/marked/compare/v18.0.9...v18.0.10)

Updates `eslint` from 10.8.1 to 10.9.0
- [Release notes](https://github.com/eslint/eslint/releases)
- [Commits](https://github.com/eslint/eslint/compare/v10.8.1...v10.9.0)

Updates `vite` from 8.2.1 to 8.2.2
- [Release notes](https://github.com/vitejs/vite/releases)
- [Changelog](https://github.com/vitejs/vite/blob/main/packages/vite/CHANGELOG.md)
- [Commits](https://github.com/vitejs/vite/commits/v8.2.2/packages/vite)

Updates `vue-tsc` from 3.3.10 to 3.3.11
- [Release notes](https://github.com/vuejs/language-tools/releases)
- [Changelog](https://github.com/vuejs/language-tools/blob/master/CHANGELOG.md)
- [Commits](https://github.com/vuejs/language-tools/commits/v3.3.11/packages/tsc)

---
updated-dependencies:
- dependency-name: "@lucide/vue"
  dependency-version: 1.33.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: dayjs
  dependency-version: 1.11.23
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: dompurify
  dependency-version: 3.4.14
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: es-toolkit
  dependency-version: 1.51.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: marked
  dependency-version: 18.0.10
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: eslint
  dependency-version: 10.9.0
  dependency-type: direct:development
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: vite
  dependency-version: 8.2.2
  dependency-type: direct:development
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: vue-tsc
  dependency-version: 3.3.11
  dependency-type: direct:development
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-08-25 15:37:10 +08:00
dependabot[bot] d22c1d4ab9 chore(deps): bump the minor-and-patch group in /backend with 4 updates (#44)
Bumps the minor-and-patch group in /backend with 4 updates: [uuid](https://github.com/uuid-rs/uuid), [aws-sdk-s3](https://github.com/awslabs/aws-sdk-rust), [aws-smithy-runtime-api](https://github.com/smithy-lang/smithy-rs) and [redis](https://github.com/redis-rs/redis-rs).


Updates `uuid` from 1.24.0 to 1.24.1
- [Release notes](https://github.com/uuid-rs/uuid/releases)
- [Commits](https://github.com/uuid-rs/uuid/compare/v1.24.0...v1.24.1)

Updates `aws-sdk-s3` from 1.141.0 to 1.142.0
- [Release notes](https://github.com/awslabs/aws-sdk-rust/releases)
- [Commits](https://github.com/awslabs/aws-sdk-rust/commits)

Updates `aws-smithy-runtime-api` from 1.14.0 to 1.15.0
- [Release notes](https://github.com/smithy-lang/smithy-rs/releases)
- [Changelog](https://github.com/smithy-lang/smithy-rs/blob/main/CHANGELOG.md)
- [Commits](https://github.com/smithy-lang/smithy-rs/commits)

Updates `redis` from 1.5.0 to 1.6.0
- [Release notes](https://github.com/redis-rs/redis-rs/releases)
- [Commits](https://github.com/redis-rs/redis-rs/compare/redis-1.5.0...redis-1.6.0)

---
updated-dependencies:
- dependency-name: uuid
  dependency-version: 1.24.1
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: aws-sdk-s3
  dependency-version: 1.142.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: aws-smithy-runtime-api
  dependency-version: 1.15.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: redis
  dependency-version: 1.6.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-08-25 15:36:39 +08:00
dependabot[bot] cc8dfe3dd8 chore(deps): bump the minor-and-patch group with 2 updates (#43)
Bumps the minor-and-patch group with 2 updates: [docker/setup-buildx-action](https://github.com/docker/setup-buildx-action) and [github/codeql-action/upload-sarif](https://github.com/github/codeql-action).


Updates `docker/setup-buildx-action` from 4.2.0 to 4.3.0
- [Release notes](https://github.com/docker/setup-buildx-action/releases)
- [Commits](https://github.com/docker/setup-buildx-action/compare/bb05f3f5519dd87d3ba754cc423b652a5edd6d2c...37fe631027851001ddb9b187196cc803df7f5f0e)

Updates `github/codeql-action/upload-sarif` from 4.37.7 to 4.37.8
- [Release notes](https://github.com/github/codeql-action/releases)
- [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md)
- [Commits](https://github.com/github/codeql-action/compare/ff2f1c621b7f889edc0d3c761ac2e6a3f8cdb0dd...db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28)

---
updated-dependencies:
- dependency-name: docker/setup-buildx-action
  dependency-version: 4.3.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: github/codeql-action/upload-sarif
  dependency-version: 4.37.8
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-08-25 15:36:01 +08:00
ZyyandCodex f25edadd5f fix(ci): keep OpenAI regression tests out of sources
Move the thread-spawn transport identity regression behind the public initialized-provider boundary so production source files remain free of embedded test modules.

Co-authored-by: Codex <noreply@openai.com>
2026-08-25 15:27:28 +08:00
ZyyandCodex 4f50c67332 fix(openai): preserve root session account affinity
Share the preferred account across derived child executions while keeping their continuation and WebSocket transport identities isolated.

Claim and migrate affinity atomically, clear account-bound state on fallback, reject cross-account previous-response continuations, and preserve complete replay transcripts.

Co-authored-by: Codex <noreply@openai.com>
2026-08-25 14:22:36 +08:00
ZyyandCodex 75f9d2a290 fix(openai): remove websocket pool slot caps
Co-authored-by: Codex <noreply@openai.com>
2026-08-25 11:22:07 +08:00
ZyyandCodex b93b8859f9 fix(frontend): polish theme controls and page headers
Remove the redundant preview mode switch, round the color picker input, and restore inline right-aligned page header actions.

Co-authored-by: Codex <noreply@openai.com>
2026-08-24 20:21:16 +08:00
ZyyandCodex 3510a45742 chore: bump version to 3.3.0
Co-authored-by: Codex <noreply@openai.com>
2026-08-24 04:07:37 +08:00
ZyyandCodex 85eb283b40 docs(readme): streamline project introduction
Co-authored-by: Codex <noreply@openai.com>
2026-08-24 04:06:49 +08:00
ZyyandCodex b8e1c5077d docs(theme): document runtime theme architecture
Describe the token model, preset themes, runtime stylesheet, isolated editor preview, naming conventions, and extension workflow.

Co-authored-by: Codex <noreply@openai.com>
2026-08-24 04:04:59 +08:00
ZyyandCodex 8422240967 feat(theme): add runtime token theme editor
Introduce Seed-to-Component token derivation, isolated live previews, semantic base component theming, and responsive theme controls.

Co-authored-by: Codex <noreply@openai.com>
2026-08-24 04:04:22 +08:00
ZyyandCodex fbd922117e fix(accounts): return persisted OAuth account ID
Co-authored-by: Codex <noreply@openai.com>
2026-08-22 19:37:03 +08:00
ZyyandCodex 941c4d57fa chore: bump version to 3.2.16
Co-authored-by: Codex <noreply@openai.com>
2026-08-22 09:28:34 +08:00
ZyyandCodex cf046ad9a4 fix(admin): close quota reset modal on success
Co-authored-by: Codex <noreply@openai.com>
2026-08-22 09:25:05 +08:00
ZyyandCodex 7699a954ac fix(admin): format account counts with separators
Co-authored-by: Codex <noreply@openai.com>
2026-08-22 09:21:11 +08:00
ZyyandCodex 2c31708d1b feat(admin): allow selecting quota reset credits
Co-authored-by: Codex <noreply@openai.com>
2026-08-22 09:19:20 +08:00
ZyyandCodex 80f338ce01 fix(admin): show raw account counts
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 23:12:54 +08:00
ZyyandCodex 299088a6e7 chore: bump version to 3.2.15
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 20:00:59 +08:00
ZyyandCodex a23bfca848 fix(admin): normalize OAuth credential formats
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 14:57:16 +08:00
ZyyandCodex 8951a64a05 fix(admin): raise billing abbreviation threshold
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 14:15:22 +08:00
ZyyandCodex eb20561893 fix(admin): align xAI quota usage displays
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 13:58:02 +08:00
ZyyandCodex a33637a99e fix(openai): use Codex image endpoint paths
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 13:02:40 +08:00
ZyyandCodex 97cd7a7097 chore: bump version to 3.2.14
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 01:11:57 +08:00
ZyyandCodex b0fe132736 docs: refresh architecture and project guides
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 01:01:07 +08:00
ZyyandCodex e2c75d1054 feat(admin): add direct token import and polish displays
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 00:33:26 +08:00
ZyyandCodex 94f298927d feat(openai): add active quota reset credits
Co-authored-by: Codex <noreply@openai.com>
2026-08-21 00:33:21 +08:00
ZyyandCodex 3be827204f fix(openai): keep image routes model-agnostic
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 19:05:14 +08:00
ZyyandCodex 227a1174b9 feat(openai): proxy image generation and edits
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 17:41:15 +08:00
ZyyandCodex c0defcb6df fix(admin): prefer account-wide dashboard quota
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 16:06:05 +08:00
292 changed files with 16690 additions and 36484 deletions
+2 -2
View File
@@ -45,7 +45,7 @@ jobs:
ref: ${{ inputs.ref }}
- name: Setup Docker Buildx
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
- name: Build runtime image
uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0
@@ -79,7 +79,7 @@ jobs:
${{ hashFiles('trivy-results.sarif') != '' &&
(github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository) }}
uses: github/codeql-action/upload-sarif@ff2f1c621b7f889edc0d3c761ac2e6a3f8cdb0dd # v4.37.7
uses: github/codeql-action/upload-sarif@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4.37.8
with:
sarif_file: trivy-results.sarif
category: runtime-container
+2 -2
View File
@@ -448,7 +448,7 @@ jobs:
PY
- name: Setup Docker Buildx
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
- name: Login to GHCR
uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0
@@ -519,7 +519,7 @@ jobs:
path: dist/release
- name: Setup Docker Buildx
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
- name: Install Cosign
uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2
+73 -70
View File
@@ -12,31 +12,32 @@
[![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) · [接口](docs/api.md) · [架构](docs/architecture.md)
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运行维护](#运行维护) · [部署](deploy/README.md) · [接口](docs/api.md) · [架构](docs/architecture.md)
</div>
> [!NOTE]
> 只支持 OpenAI Responses API,不提供 `/v1/chat/completions`。
> 数据面支持 OpenAI Responses、Images 和模型目录协议,不提供 `/v1/chat/completions`。
## 能力
## 功能
| 领域 | 实现 |
| -------- | ----------------------------------------------------------------------------------------------------- |
| 协议 | OpenAI Responses JSON、SSE、WebSocket 与模型目录 |
| Provider | 固定 OpenAI、xAI 两个 Provider;各自管理 OAuth、账号、额度与模型目录 |
| 路由 | Client Key 绑定 Provider、全局模型映射、会话亲和与同 Provider 安全 fallback |
| 透明边界 | OpenAI 请求保留未知字段和字段顺序,SSE 与 WebSocket 原样转发业务字节;xAI 在 Provider 内做协议转换 |
| 延续 | OpenAI native/replay continuation;xAI 使用客户端提交的完整历史 |
| 账号 | 导入、OAuth、credential 刷新、主动额度刷新,以及正常响应中的被动额度/状态观测 |
| 管理 | Client Key、账号、模型目录、设置、观测、系统更新与 OAuth |
| 计量 | 模型请求 Token、费用、延迟、账号与 Provider 归因,并记录响应确认的实际 `service_tier` |
| 领域 | 当前能力 |
| --- | --- |
| 数据面 | Responses JSON、SSE、WebSocket、review 子代理、Images 生成/编辑与模型目录 |
| Provider | 固定 `openai`、`xai` 两个 Provider,各自拥有 OAuth、credential、额度、目录与 transport |
| 路由 | Client Key 限定账号分组;按模型能力编译 Provider 候选;支持全局精确模型映射、会话亲和和安全 fallback |
| 透明边界 | OpenAI 保留未知请求字段及字段顺序,并原样转发 SSE、WebSocket 与 Images 业务字节;xAI 在 Provider 内转换协议 |
| 账号 | OAuth、OpenAI AT/RT、账号文件与 Agent Identity 导入,credential/额度刷新、账号恢复及主动额度重置卡 |
| 管理 | 账号与分组、Client Key、运行设置、用量/错误观测、S3/R2 备份、系统更新与回滚 |
| 计量 | Token、图片、费用、延迟、账号与 Provider 归因,并记录上游响应确认的实际 `service_tier` |
架构边界、状态所有权和请求不变量见 [架构文档](docs/architecture.md)。
## 快速开始
需要 Docker Engine 与 Docker Compose Plugin。
### 1. 准备
### 1. 准备配置和目录
```bash
git clone https://github.com/zyycn/codex-proxy-rs.git
@@ -56,16 +57,16 @@ openssl rand -hex 24
openssl rand -hex 24
```
把两个结果写入 `deploy/config.yaml`,并单独设置管理员初始密码:
把两个结果分别写入 `deploy/config.yaml`,再设置管理员初始密码:
| 配置 | 约束 |
| -------------------------- | ---------------------------- |
| `store.database.password` | 48 位十六进制 |
| `store.redis.password` | 48 位十六进制 |
| `admin.default_password` | 至少 12 位,不能包含 `$` |
| 配置 | 约束 |
| --- | --- |
| `store.database.password` | 48 位十六进制 |
| `store.redis.password` | 48 位十六进制 |
| `admin.default_password` | 至少 12 位,不能包含 `$` |
PostgreSQL、Redis 容器和应用通过同一份 `config.yaml` 使用这两个基础设施密码,不需要额外
导出环境变量。管理员密码只在首次创建管理员时生效。
PostgreSQL、Redis 和应用从同一份 `config.yaml` 读取基础设施密码,不需要额外导出环境变量。
管理员密码只在首次创建管理员时生效。
Linux 需要允许容器组写入运行目录:
@@ -83,37 +84,36 @@ docker compose -f deploy/compose.yaml up -d --no-build
curl -i http://127.0.0.1:8080/healthz
```
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址:
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址为
`http://127.0.0.1:8080`。
### 3. 初始化
### 3. 初始化网关
1. 使用 `admin@cpr.local` 与初始密码登录。
2. 在账号页选择对应 Provider:OpenAI 支持 OAuth、账号文件与 Agent Identity;xAI 支持 OAuth
和兼容账号文件导入。
3. 按需配置客户端模型到上游模型的全局映射;未命中映射的模型名原样透传。
4. 创建 `sk_...` Client Key,绑定 `openai` 或 `xai`,并设置速率与并发限制。
2. 导入或授权 Provider 账号:OpenAI 支持 OAuth、逐行 AT/RT、OAuth JSON 与 Agent Identity;
xAI 支持 OAuth 和账号 JSON。
3. 创建账号分组并按需配置全局模型映射。未命中的模型名原样交给候选 Provider。
4. 创建 `sk_...` Client Key,选择可用账号分组,并设置速率和并发限制;空分组范围表示全部账号。
OpenAI 账号导入和首次 OAuth 创建在 credential 提交后会立即尝试一次额度观测;额度接口临时失败不会
回滚已提交的账号。Free、K12 等套餐共用同一套额度与状态转换逻辑,套餐字段只作为账号事实并隔离
对应的模型目录 cache。credential 刷新、额度刷新和正常请求中的被动额度观测是三条独立链路,详见
[账号接口](docs/api.md#5-账号)。
OpenAI RT 导入会先向官方 token endpoint 换取 AT;仅 AT 的账号不具备自动续期能力,
能否参与调度仍取决于 token 中可验证的身份与后续上游事实。账号导入和首次 OAuth 创建会尽力读取一次额度,
额度上游暂时失败不会回滚已提交的账号。
OpenAI 账号的主动重置卡只在用户打开相应弹窗或手工刷新时查询。消费成功后管理端重新查询卡片与额度;
系统不会通过修改本地 `reset_at` 伪造重置。完整账号合同见 [账号接口](docs/api.md#5-账号)。
> [!IMPORTANT]
> xAI 使用 OAuth session,不支持把 xAI API Key 作为上游 credential。
> [!NOTE]
> `admin.default_password` 只在首次创建管理员时生效。
完整部署、权限和密码轮换规则见 [deploy/README.md](deploy/README.md)。
完整的权限、密码轮换、备份和升级规则见 [部署文档](deploy/README.md)。
## 客户端接入
| 配置 | 值 |
| -------- | ---------------------------------------- |
| Base URL | `http://127.0.0.1:8080/v1` |
| API Key | 管理端创建的 `sk_...` |
| 鉴权 | `Authorization: Bearer <client-api-key>` |
| 配置 | 值 |
| --- | --- |
| Base URL | `http://127.0.0.1:8080/v1` |
| API Key | 管理端创建的 `sk_...` |
| 鉴权 | `Authorization: Bearer <client-api-key>` |
```bash
curl http://127.0.0.1:8080/v1/models \
@@ -130,67 +130,70 @@ curl http://127.0.0.1:8080/v1/responses \
```
<details>
<summary>API 路由</summary>
<summary>公开数据面路由</summary>
| 路由 | 用途 |
| ---------------------------------- | ------------------------------ |
| `POST /v1/responses` | JSON 响应或 SSE Responses 流 |
| `GET /v1/responses` | Responses WebSocket 升级 |
| `POST /v1/responses/review` | review 子代理请求 |
| `GET /v1/models` | 当前 Client Key 可用的模型列表 |
| `GET /v1/models/catalog` | Codex 客户端展示用模型目录 |
| `GET /v1/models/{model_id}/info` | Codex 客户端展示用模型详情 |
| `GET /v1/models/{model_id}` | OpenAI 兼容模型详情 |
| 路由 | 用途 |
| --- | --- |
| `POST /v1/responses` | JSON 响应或 SSE Responses 流 |
| `GET /v1/responses` | Responses WebSocket 升级 |
| `POST /v1/responses/review` | review 子代理请求 |
| `POST /v1/images/generations` | OpenAI Images 生成 JSON 透明转发 |
| `POST /v1/images/edits` | OpenAI Images 编辑 JSON 透明转发 |
| `GET /v1/models` | 当前 Client Key 账号范围内的模型列表 |
| `GET /v1/models/catalog` | Codex 客户端展示用模型目录 |
| `GET /v1/models/{model_id}/info` | Codex 客户端展示用模型详情 |
| `GET /v1/models/{model_id}` | OpenAI 兼容模型详情 |
所有 `/v1/*` 路由都需要客户端 API Key。
所有 `/v1/*` 路由都需要 Client Key。Images 是 OpenAI Provider 自有端点,不读取模型字段,也不参与
文本模型映射。
</details>
完整的客户端与管理端路由、鉴权和 mutation 规则见 [接口文档](docs/api.md)。
完整的客户端与管理端路由、鉴权、请求字段和 mutation 语义见 [接口文档](docs/api.md)。
## 运维
## 运行维护
```bash
# 实例升级
# 同一大版本内升级发布镜像
docker compose -f deploy/compose.yaml pull codex-proxy-rs
docker compose -f deploy/compose.yaml up -d --no-build
# 日志
# 查看日志
docker compose -f deploy/compose.yaml logs -f codex-proxy-rs
# 从源码构建
# 从当前源码构建
docker compose -f deploy/compose.yaml build codex-proxy-rs
docker compose -f deploy/compose.yaml up -d
```
> [!IMPORTANT]
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点密钥和自更新状态。删除该目录会永久清除
> 运行状态。
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点和在线更新状态。删除该目录会永久清除
> 本地运行状态。
> [!WARNING]
> 不支持跨大版本在线升级。升级到新的大版本时,请使用全新的 `.runtime/` 数据目录重新部署,
> 并重新导入或重新授权 Provider 账号与客户端 Key。
> 不支持跨大版本在线升级。升级到新的大版本时请使用全新的 `.runtime/` 数据目录,并重新导入或授权
> Provider 账号与 Client Key。
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket
upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket upgrade 与真实客户端
IP;不要暴露 PostgreSQL 或 Redis。当前运行拓扑为单副本,不能通过复制应用容器扩容。
## 维护与发布
发布仓库版本必须从干净、已同步上游的发布分支运行:
仓库版本只能从干净、已同步上游的发布分支发布:
```bash
release/publish <version>
```
脚本依赖已登录的 GitHub CLI(`gh auth login`)。它会更新 `release/version.yaml`、创建约定提交与带
注释的 `v<version>` tag,原子推送分支和 tag,再从发布分支触发 Release 构建,使构建缓存可以跨版本
复用。不要手工修改版本后单独打 tag。该脚本不会登录服务器,也不会改变任何运行实例;源码提交、
仓库发版和实例升级是三个独立状态。
脚本依赖已登录的 GitHub CLI(`gh auth login`),会更新 `release/version.yaml`、创建约定提交和带注释的
`v<version>` tag,原子推送分支与 tag,再触发 Release 构建。源码提交、仓库发版和运行实例升级是三个
独立状态;脚本不会登录服务器或改变任何运行实例。
## 文档
- [部署](deploy/README.md)
- [部署与运维](deploy/README.md)
- [客户端与管理端接口](docs/api.md)
- [架构与数据边界](docs/architecture.md)
- [系统架构](docs/architecture.md)
- [数据库迁移规则](backend/migrations/README.md)
- [Release](https://github.com/zyycn/codex-proxy-rs/releases)
- [容器镜像](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
+18 -18
View File
@@ -216,9 +216,9 @@ dependencies = [
[[package]]
name = "aws-sdk-s3"
version = "1.141.0"
version = "1.142.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d9f9420d3a2467eed22ed3635ca653653162c386a0b0f65c78189f9bd3c1379e"
checksum = "f9e15a5c55e05f4b0b7e483160b3c85cccdf77cff02c95504f3e71d460855cd2"
dependencies = [
"arc-swap",
"aws-credential-types",
@@ -308,9 +308,9 @@ dependencies = [
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
version = "0.61.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
checksum = "6de526c7b567420a31bc283657a7921b45c4cafe0827fdf2490713dcc770c28f"
dependencies = [
"aws-smithy-types",
"bytes",
@@ -341,9 +341,9 @@ dependencies = [
[[package]]
name = "aws-smithy-http-client"
version = "1.2.0"
version = "1.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e"
checksum = "3c1c8a04cb31ba74d0115af5a890bb8c0d48fba64b52812fa13929a6ef0cc83c"
dependencies = [
"aws-smithy-async",
"aws-smithy-runtime-api",
@@ -385,9 +385,9 @@ dependencies = [
[[package]]
name = "aws-smithy-runtime"
version = "1.12.1"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "07505b34e8f4b3591a4fa69e9792b52289b95488dbbc68c3c0075b7bedb245e1"
checksum = "483b858ff67522011c4786310c5cd8fd88d0be7ea3d5f1a48328446300c4269e"
dependencies = [
"aws-smithy-async",
"aws-smithy-http",
@@ -411,9 +411,9 @@ dependencies = [
[[package]]
name = "aws-smithy-runtime-api"
version = "1.14.0"
version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3b98f2e1fd67ec06618f9c291e5e495a468e60519e44c9c1979cd0521f3affdb"
checksum = "954c563ce84507722d2679f07a35d21b9c6466b3872d513020d0281fc8112ac9"
dependencies = [
"aws-smithy-async",
"aws-smithy-runtime-api-macros",
@@ -1246,7 +1246,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.52.0",
"windows-sys 0.61.2",
]
[[package]]
@@ -2714,7 +2714,7 @@ dependencies = [
"once_cell",
"socket2",
"tracing",
"windows-sys 0.52.0",
"windows-sys 0.61.2",
]
[[package]]
@@ -2813,9 +2813,9 @@ dependencies = [
[[package]]
name = "redis"
version = "1.5.0"
version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3257df217f7eab0044627a268c9cc6cdb60c0c421c88f83ac41c4e31520b6b84"
checksum = "e37a4ca5c6ca42aa3e6df2fd32b987a65d32a4c2159a6f3fe0fd1df306a2658f"
dependencies = [
"arc-swap",
"arcstr",
@@ -2984,7 +2984,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.52.0",
"windows-sys 0.61.2",
]
[[package]]
@@ -3628,7 +3628,7 @@ dependencies = [
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.52.0",
"windows-sys 0.61.2",
]
[[package]]
@@ -4100,9 +4100,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
[[package]]
name = "uuid"
version = "1.24.0"
version = "1.24.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239"
checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9"
dependencies = [
"getrandom 0.4.3",
"js-sys",
@@ -119,11 +119,11 @@ pub struct RetentionDecision {
/// 从按完成时间倒序的 completed 计划备份中决定纳入删除流程的记录。
///
/// 规则(`docs/s3-backup-design-audit.md` §4.4):
/// 自动保留规则:
/// - `retentionDays = 0` 时不按天数清理;`retentionCount = 0` 时不按份数清理。
/// - 自动计划备份满足任一启用阈值即可进入删除。
/// - 无论阈值如何,至少保留最近一个成功的计划备份(倒序下标 0)。
/// - 手动备份不受自动阈值影响;无 TTL,仅由管理员手动删除。
/// - 手动备份不进入本规则;手工/计划记录各自的 `expires_at` 由到期扫描处理。
#[must_use]
pub fn decide_retention(
retention_days: u32,
+4 -1
View File
@@ -347,7 +347,10 @@ fn map_provider_registry_error(error: ProviderAdminError) -> AdminError {
}
ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound,
ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict,
ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::Unavailable | ProviderAdminErrorKind::CredentialRefreshRequired => {
AdminErrorKind::Unavailable
}
ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway,
ProviderAdminErrorKind::Internal => AdminErrorKind::Internal,
};
AdminError::new(kind, "Provider registry initialization failed")
@@ -288,7 +288,7 @@ pub struct BackupRecord {
pub error_message: Option<String>,
pub started_at: Option<DateTime<Utc>>,
pub completed_at: Option<DateTime<Utc>>,
/// 手动过期时间;创建时确定,到期自动清理。
/// 创建时确定的过期时间;手工和计划备份均可设置,到期自动清理。
pub expires_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
@@ -557,7 +557,7 @@ pub struct UsageRecord {
pub operation: String,
pub endpoint: String,
pub client_transport: String,
pub requested_model_id: String,
pub requested_model_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub provider_account_name: Option<String>,
@@ -3,15 +3,17 @@
use std::fmt;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use chrono::{DateTime, NaiveDate, Utc};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
routing::{ProviderKind, UpstreamModelId},
};
use uuid::Uuid;
use super::{
AdminError, MutationActor, MutationContext, PageSize, Revision,
accounts::{AccountRecord, AccountSummary, AccountUsage, CredentialState},
observability::CurrencyCost,
};
/// Provider-owned JSON;公共层只搬运且 Debug 不输出值。
@@ -626,6 +628,136 @@ pub struct ProviderQuota {
pub provider_data: Option<ProviderDocument>,
}
/// Provider 官方用量统计的目标额度周期。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderUsageStatisticsCycle {
pub offset: i8,
pub start_at: DateTime<Utc>,
pub end_at: DateTime<Utc>,
pub window_seconds: u64,
pub used_percent: Option<f64>,
pub is_current: bool,
pub can_go_previous: bool,
pub can_go_next: bool,
}
/// Provider 官方用量统计查询。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderUsageStatisticsRequest {
pub account_id: ProviderAccountId,
/// `0` 为当前周期,负数为历史周期。
pub cycle_offset: i8,
/// 客户端相对 UTC 的分钟偏移,例如 UTC+8 为 `480`。
pub utc_offset_minutes: i16,
}
/// 官方用量统计中的 Token 分类。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ProviderUsageStatisticsTokens {
pub uncached_input: u64,
pub cached_input: u64,
pub output: u64,
pub total: u64,
}
/// 官方模型报表中的服务档位。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderUsageStatisticsServiceTier {
Standard,
Fast,
}
impl ProviderUsageStatisticsServiceTier {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Standard => "standard",
Self::Fast => "fast",
}
}
}
/// 一个模型与服务档位在目标周期内的统计行。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderUsageStatisticsModel {
pub key: String,
pub model: String,
pub service_tier: ProviderUsageStatisticsServiceTier,
pub credit_share: Option<f64>,
pub quota_share: Option<f64>,
pub tokens: ProviderUsageStatisticsTokens,
pub estimated_cost: Option<CurrencyCost>,
pub has_unknown_pricing: bool,
pub has_estimated_allocation: bool,
pub has_rate_fallback: bool,
pub has_missing_token_data: bool,
}
/// 目标周期内的一天;官方日报只能提供本地日粒度。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderUsageStatisticsDay {
pub date: NaiveDate,
pub credit_share: Option<f64>,
pub tokens: ProviderUsageStatisticsTokens,
pub estimated_cost: Option<CurrencyCost>,
pub has_unknown_pricing: bool,
pub has_missing_token_data: bool,
pub is_boundary_day: bool,
}
/// 目标周期汇总。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderUsageStatisticsSummary {
pub tokens: ProviderUsageStatisticsTokens,
pub estimated_cost: Option<CurrencyCost>,
pub projected_tokens: Option<u64>,
pub projected_cost: Option<CurrencyCost>,
pub day_count: u32,
pub has_unknown_pricing: bool,
pub has_missing_token_data: bool,
}
/// Provider 已解释、计算并排序的官方用量统计。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderUsageStatistics {
pub cycle: ProviderUsageStatisticsCycle,
pub summary: ProviderUsageStatisticsSummary,
pub models: Vec<ProviderUsageStatisticsModel>,
pub daily: Vec<ProviderUsageStatisticsDay>,
}
/// Provider 返回的一张安全主动额度重置卡。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderResetCredit {
pub id: String,
pub status: Option<String>,
pub title: Option<String>,
pub expires_at: Option<DateTime<Utc>>,
pub reset_type: Option<String>,
}
/// Provider 主动额度重置卡列表。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderResetCredits {
pub available_count: u64,
pub credits: Vec<ProviderResetCredit>,
}
/// 一次主动额度重置卡消费命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConsumeProviderResetCredit {
pub account_id: ProviderAccountId,
pub credit_id: Option<String>,
pub redeem_request_id: Uuid,
}
/// Provider 返回的主动额度重置卡消费结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderResetCreditResult {
pub code: String,
pub credit: Option<ProviderResetCredit>,
}
impl ProviderQuota {
/// 返回账号展开面板使用的当前额度窗口用量。
///
@@ -648,8 +780,9 @@ impl ProviderQuota {
/// 返回 Dashboard 使用的代表性额度比例。
///
/// 优先使用短周期窗口,再依次使用周、月和其它窗口;同一优先级取较高的已用比例。
/// 这里只解释跨 Provider 共享的窗口语义,绝不读取 Provider 私有 JSON。
/// 优先使用覆盖账号全部请求的窗口;同一归属范围内再依次使用短周期、周、月
/// 和其它窗口,同一优先级取较高的已用比例。这里只解释跨 Provider 共享的窗口
/// 语义,绝不读取 Provider 私有 JSON。
#[must_use]
pub fn representative_used_percent(&self) -> Option<f64> {
if self.limit_reached {
@@ -708,7 +841,7 @@ impl ProviderQuota {
Some((index, quota_usage_priority(window), used_percent))
})
.fold(
None::<(usize, u8, f64)>,
None::<(usize, (u8, u8), f64)>,
|selected, candidate| match selected {
Some(current)
if current.1 < candidate.1
@@ -723,13 +856,18 @@ impl ProviderQuota {
}
}
fn quota_usage_priority(window: &ProviderQuotaWindow) -> u8 {
match window.group.as_str() {
fn quota_usage_priority(window: &ProviderQuotaWindow) -> (u8, u8) {
let attribution = match window.local_usage_attribution {
QuotaLocalUsageAttribution::AccountWide => 0,
QuotaLocalUsageAttribution::Unavailable => 1,
};
let duration = match window.group.as_str() {
"shortTerm" if window.window_seconds.is_some_and(is_week_window) => 1,
"shortTerm" => 0,
"monthly" => 2,
_ => 3,
}
};
(attribution, duration)
}
fn is_week_window(seconds: u64) -> bool {
@@ -13,11 +13,12 @@ use crate::model::observability::{
CalculatedBillingBreakdown, DashboardWireProfile, ProviderBillingInput,
};
use crate::model::provider_credentials::{
AuthorizationStarted, CompleteAuthorization, PendingAuthorizationMutation,
PrepareCredentialImport, PrepareCredentialRefresh, PrepareCredentialRotation,
PreparedAuthorizationCommit, PreparedCredentialImport, PreparedCredentialRotation,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
ProviderQuotaRequest,
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, ProviderModels,
ProviderQuota, ProviderQuotaRequest, ProviderResetCreditResult, ProviderResetCredits,
ProviderUsageStatistics, ProviderUsageStatisticsRequest,
};
/// Provider 管理失败的稳定分类。
@@ -28,6 +29,8 @@ pub enum ProviderAdminErrorKind {
NotFound,
Conflict,
Unavailable,
CredentialRefreshRequired,
BadGateway,
Internal,
}
@@ -144,6 +147,30 @@ pub trait ProviderAdmin: Send + Sync {
request: ProviderQuotaRequest,
) -> Result<ProviderQuota, ProviderAdminError>;
/// 查询 Provider 官方用量统计;不支持该能力的 Provider 使用默认拒绝。
async fn usage_statistics(
&self,
_request: ProviderUsageStatisticsRequest,
) -> Result<ProviderUsageStatistics, ProviderAdminError> {
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
}
/// 查询 Provider 主动额度重置卡;不支持该能力的 Provider 使用默认拒绝。
async fn reset_credits(
&self,
_account_id: &ProviderAccountId,
) -> Result<ProviderResetCredits, ProviderAdminError> {
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
}
/// 消费 Provider 主动额度重置卡;不支持该能力的 Provider 使用默认拒绝。
async fn consume_reset_credit(
&self,
_command: ConsumeProviderResetCredit,
) -> Result<ProviderResetCreditResult, ProviderAdminError> {
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
}
async fn models(
&self,
account_id: &ProviderAccountId,
@@ -24,8 +24,10 @@ use crate::{
observability::TimeRange,
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
PrepareCredentialRefresh, ProviderModels, ProviderQuota, ProviderQuotaRequest,
ProviderQuotaWindow, QuotaLocalUsageAttribution,
ConsumeProviderResetCredit, PrepareCredentialRefresh, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow, ProviderResetCreditResult,
ProviderResetCredits, ProviderUsageStatistics, ProviderUsageStatisticsRequest,
QuotaLocalUsageAttribution,
},
},
ports::{provider::ProviderAdminRegistry, store::AccountStore},
@@ -79,6 +81,35 @@ pub trait AccountsService: Send + Sync {
refresh: bool,
) -> Result<AccountDirectoryItem, AdminError>;
async fn usage_statistics(
&self,
_request: ProviderUsageStatisticsRequest,
) -> Result<ProviderUsageStatistics, AdminError> {
Err(AdminError::invalid(
"Provider usage statistics are not supported",
))
}
async fn reset_credits(
&self,
_context: &MutationContext,
_account_id: ProviderAccountId,
) -> Result<ProviderResetCredits, AdminError> {
Err(AdminError::invalid(
"Provider reset credits are not supported",
))
}
async fn consume_reset_credit(
&self,
_context: &MutationContext,
_command: ConsumeProviderResetCredit,
) -> Result<ProviderResetCreditResult, AdminError> {
Err(AdminError::invalid(
"Provider reset credits are not supported",
))
}
async fn models(
&self,
account_id: &ProviderAccountId,
@@ -97,11 +128,13 @@ pub(crate) struct DefaultAccountsService {
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
reset_credit_locks:
Arc<futures::lock::Mutex<BTreeMap<ProviderAccountId, Arc<futures::lock::Mutex<()>>>>>,
}
impl DefaultAccountsService {
#[must_use]
pub(crate) const fn new(
pub(crate) fn new(
accounts: Arc<dyn AccountStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
@@ -112,9 +145,22 @@ impl DefaultAccountsService {
providers,
snapshot,
probe,
reset_credit_locks: Arc::new(futures::lock::Mutex::new(BTreeMap::new())),
}
}
async fn reset_credit_lock(
&self,
account_id: &ProviderAccountId,
) -> Arc<futures::lock::Mutex<()>> {
let mut locks = self.reset_credit_locks.lock().await;
Arc::clone(
locks
.entry(account_id.clone())
.or_insert_with(|| Arc::new(futures::lock::Mutex::new(()))),
)
}
async fn load_account(
&self,
account_id: &ProviderAccountId,
@@ -502,6 +548,68 @@ impl AccountsService for DefaultAccountsService {
self.load_directory_item(account_id, refresh).await
}
async fn usage_statistics(
&self,
request: ProviderUsageStatisticsRequest,
) -> Result<ProviderUsageStatistics, AdminError> {
let (_, provider) = self.provider_for_account(&request.account_id).await?;
provider
.usage_statistics(request)
.await
.map_err(|error| map_provider_error(error, "provider usage statistics"))
}
async fn reset_credits(
&self,
context: &MutationContext,
account_id: ProviderAccountId,
) -> Result<ProviderResetCredits, AdminError> {
let (_, provider) = self.provider_for_account(&account_id).await?;
match provider.reset_credits(&account_id).await {
Ok(credits) => Ok(credits),
Err(error)
if error.kind()
== crate::ports::provider::ProviderAdminErrorKind::CredentialRefreshRequired =>
{
self.refresh(context, account_id.clone()).await?;
let (_, provider) = self.provider_for_account(&account_id).await?;
provider
.reset_credits(&account_id)
.await
.map_err(map_reset_credits_error_after_refresh)
}
Err(error) => Err(map_provider_error(error, "provider reset credits")),
}
}
async fn consume_reset_credit(
&self,
context: &MutationContext,
command: ConsumeProviderResetCredit,
) -> Result<ProviderResetCreditResult, AdminError> {
let account_id = command.account_id.clone();
// 覆盖 credential refresh + 同键重试的完整账号级临界区,避免 401 两次调用
// 之间插入另一笔不可逆消费。
let lock = self.reset_credit_lock(&account_id).await;
let _guard = lock.lock().await;
let (_, provider) = self.provider_for_account(&account_id).await?;
match provider.consume_reset_credit(command.clone()).await {
Ok(result) => Ok(result),
Err(error)
if error.kind()
== crate::ports::provider::ProviderAdminErrorKind::CredentialRefreshRequired =>
{
self.refresh(context, account_id.clone()).await?;
let (_, provider) = self.provider_for_account(&account_id).await?;
provider
.consume_reset_credit(command)
.await
.map_err(map_reset_credits_error_after_refresh)
}
Err(error) => Err(map_provider_error(error, "provider reset-credit consume")),
}
}
async fn models(
&self,
account_id: &ProviderAccountId,
@@ -582,6 +690,19 @@ impl AccountsService for DefaultAccountsService {
}
}
fn map_reset_credits_error_after_refresh(
error: crate::ports::provider::ProviderAdminError,
) -> AdminError {
if error.kind() == crate::ports::provider::ProviderAdminErrorKind::CredentialRefreshRequired {
return AdminError::bad_gateway(
error
.message()
.unwrap_or("OpenAI reset-credit upstream rejected the refreshed credential"),
);
}
map_provider_error(error, "provider reset credits")
}
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影。
fn empty_quota() -> ProviderQuota {
ProviderQuota {
@@ -57,6 +57,8 @@ fn map_provider_error(
ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound,
ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict,
ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::CredentialRefreshRequired => AdminErrorKind::Unavailable,
ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway,
ProviderAdminErrorKind::Internal => AdminErrorKind::Internal,
};
let message = error
@@ -136,6 +136,24 @@ fn representative_quota_should_prefer_short_window_and_highest_usage() {
assert_eq!(quota.representative_used_percent(), Some(45.0));
}
#[test]
fn representative_quota_should_prefer_account_wide_window_over_model_specific_window() {
let mut model_specific = quota_window("shortTerm", Some(18_000), Some(64.0));
model_specific.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
let quota = ProviderQuota {
observed_at: None,
refresh_token_expires_at: None,
windows: vec![
model_specific,
quota_window("shortTerm", Some(604_800), Some(5.0)),
],
limit_reached: false,
provider_data: None,
};
assert_eq!(quota.representative_used_percent(), Some(5.0));
}
#[test]
fn representative_window_usage_should_accept_account_wide_rolling_window_without_reset() {
let usage = AccountUsage {
@@ -30,14 +30,16 @@ use gateway_admin::{
provider_credentials::{
AuthorizationCommit, AuthorizationCommitGuard, AuthorizationCredentialCommit,
AuthorizationMutationTarget, AuthorizationStarted, CompleteAuthorization,
CredentialCommitGuard, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
ConsumeProviderResetCredit, CredentialCommitGuard, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow, QuotaLocalUsageAttribution,
ProviderQuotaRequest, ProviderQuotaWindow, ProviderResetCreditResult,
QuotaLocalUsageAttribution,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
@@ -69,6 +71,7 @@ pub(super) struct FakeProviderAdmin {
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
quota: Mutex<ProviderQuota>,
current_credential_revision: Mutex<Revision>,
reset_credit_commands: Mutex<Vec<ConsumeProviderResetCredit>>,
}
impl FakeProviderAdmin {
@@ -85,6 +88,7 @@ impl FakeProviderAdmin {
quota_requests: Mutex::new(Vec::new()),
quota: Mutex::new(empty_quota()),
current_credential_revision: Mutex::new(revision(1)),
reset_credit_commands: Mutex::new(Vec::new()),
})
}
@@ -118,6 +122,13 @@ impl FakeProviderAdmin {
.clone()
}
fn reset_credit_commands(&self) -> Vec<ConsumeProviderResetCredit> {
self.reset_credit_commands
.lock()
.expect("provider reset-credit commands")
.clone()
}
pub(super) fn set_quota(&self, quota: ProviderQuota) {
*self.quota.lock().expect("provider quota") = quota;
}
@@ -349,6 +360,22 @@ impl ProviderAdmin for FakeProviderAdmin {
Ok(self.quota.lock().expect("provider quota").clone())
}
async fn consume_reset_credit(
&self,
command: ConsumeProviderResetCredit,
) -> Result<ProviderResetCreditResult, ProviderAdminError> {
self.record("provider.consume_reset_credit");
self.reset_credit_commands
.lock()
.expect("provider reset-credit commands")
.push(command);
self.require_available()?;
Ok(ProviderResetCreditResult {
code: "reset".to_owned(),
credit: None,
})
}
async fn models(
&self,
_: &ProviderAccountId,
@@ -1667,6 +1694,33 @@ async fn accounts_refresh_provider_failure_should_not_call_store_commit() {
);
}
#[tokio::test]
async fn reset_credit_oauth_refresh_should_reuse_the_exact_consume_command() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
provider.fail_next(ProviderAdminErrorKind::CredentialRefreshRequired);
let store = FakeAccountStore::new("openai", events.clone());
let services = accounts_service(provider.clone(), store).await;
let redeem_request_id =
uuid::Uuid::parse_str("8fbf302d-11df-4bd5-82e4-08e4b3df7874").expect("UUID v4");
let command = ConsumeProviderResetCredit {
account_id: ProviderAccountId::new("acct_test").expect("account ID"),
credit_id: Some("credit_1".to_owned()),
redeem_request_id,
};
let result = services
.accounts()
.consume_reset_credit(&context("reset-credit-refresh"), command.clone())
.await
.expect("consume after OAuth refresh");
assert_eq!(result.code, "reset");
assert_eq!(provider.reset_credit_commands(), [command.clone(), command]);
assert!(recorded(&events).contains(&"provider.prepare_refresh"));
assert!(recorded(&events).contains(&"store.commit_refresh"));
}
#[tokio::test]
async fn accounts_refresh_should_preserve_the_provider_failure_message() {
let events = events();
@@ -662,7 +662,7 @@ fn total_record(
operation: "generate".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "gpt-5.5".to_owned(),
requested_model_id: Some("gpt-5.5".to_owned()),
provider_kind: provider_kind.map(str::to_owned),
provider_account_ref: None,
provider_account_name: None,
@@ -25,9 +25,13 @@ use gateway_admin::model::{
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion, CredentialDeletionResult,
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow, RotateCredential,
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation,
CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels,
ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult,
ProviderResetCredits, ProviderUsageStatistics, ProviderUsageStatisticsDay,
ProviderUsageStatisticsModel, ProviderUsageStatisticsRequest,
ProviderUsageStatisticsSummary, ProviderUsageStatisticsTokens, RotateCredential,
StartAuthorization,
},
};
@@ -37,6 +41,7 @@ use gateway_core::{
};
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
use uuid::{Uuid, Version};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
use super::{
@@ -56,6 +61,8 @@ const MAX_ID_TOKEN_BYTES: usize = 16 * 1024;
const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024;
const MAX_ACCOUNT_DELETE_BATCH: usize = 200;
const MAX_ACCOUNT_GROUP_BATCH: usize = 1000;
const MAX_USAGE_CYCLE_OFFSET: i8 = 8;
const MAX_UTC_OFFSET_MINUTES: i16 = 14 * 60;
/// 账号列表查询参数。
#[derive(Debug, Clone, Default, Deserialize)]
@@ -418,6 +425,35 @@ impl AccountIdQuery {
}
}
/// 官方每日用量统计查询。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountUsageStatisticsQuery {
pub account_id: String,
pub cycle_offset: Option<i8>,
pub utc_offset_minutes: Option<i16>,
}
impl AccountUsageStatisticsQuery {
pub fn into_request(self) -> Result<ProviderUsageStatisticsRequest, WireValidationError> {
require_account_id(&self.account_id, "accountId")?;
let cycle_offset = self.cycle_offset.unwrap_or_default();
if !(-MAX_USAGE_CYCLE_OFFSET..=0).contains(&cycle_offset) {
return Err(WireValidationError::new("cycleOffset"));
}
let utc_offset_minutes = self.utc_offset_minutes.unwrap_or_default();
if !(-MAX_UTC_OFFSET_MINUTES..=MAX_UTC_OFFSET_MINUTES).contains(&utc_offset_minutes) {
return Err(WireValidationError::new("utcOffsetMinutes"));
}
Ok(ProviderUsageStatisticsRequest {
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
cycle_offset,
utc_offset_minutes,
})
}
}
/// 敏感导出的固定 query;IDs 使用逗号分隔,禁止隐式导出全部账号。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
@@ -476,6 +512,48 @@ impl AccountActionRequest {
}
}
/// 主动额度重置卡消费请求。幂等键由 UI 生成并在不确定重试时复用,与官方一致。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountResetCreditConsumeRequest {
pub account_id: String,
#[serde(default)]
pub credit_id: Option<String>,
pub redeem_request_id: String,
}
impl AccountResetCreditConsumeRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_account_id(&self.account_id, "accountId")?;
if self.credit_id.as_deref().is_some_and(|value| {
let value = value.trim();
value.is_empty() || value.len() > 512 || value.chars().any(char::is_control)
}) {
return Err(WireValidationError::new("creditId"));
}
let value = self.redeem_request_id.trim();
let redeem_request_id =
Uuid::parse_str(value).map_err(|_| WireValidationError::new("redeemRequestId"))?;
if redeem_request_id.get_version() != Some(Version::Random)
|| redeem_request_id.hyphenated().to_string() != value
{
return Err(WireValidationError::new("redeemRequestId"));
}
Ok(())
}
fn into_command(self) -> Result<ConsumeProviderResetCredit, WireValidationError> {
self.validate()?;
Ok(ConsumeProviderResetCredit {
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
credit_id: self.credit_id.map(|value| value.trim().to_owned()),
redeem_request_id: Uuid::parse_str(&self.redeem_request_id)
.map_err(|_| WireValidationError::new("redeemRequestId"))?,
})
}
}
/// 手工 OAuth 刷新会变更持久 credential。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
@@ -543,6 +621,244 @@ pub struct AccountQuotaData {
pub account: AccountView,
}
/// Provider 官方每日用量统计响应。
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsData {
pub cycle: AccountUsageStatisticsCycleView,
pub summary: AccountUsageStatisticsSummaryView,
pub models: Vec<AccountUsageStatisticsModelView>,
pub daily: Vec<AccountUsageStatisticsDayView>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsCycleView {
pub offset: i8,
pub start_at: String,
pub end_at: String,
pub window_seconds: u64,
pub used_percent: Option<f64>,
pub is_current: bool,
pub can_go_previous: bool,
pub can_go_next: bool,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsTokensView {
pub uncached_input: u64,
pub cached_input: u64,
pub output: u64,
pub total: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsCostView {
pub currency: String,
pub amount: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsSummaryView {
pub tokens: AccountUsageStatisticsTokensView,
pub estimated_cost: Option<AccountUsageStatisticsCostView>,
pub projected_tokens: Option<u64>,
pub projected_cost: Option<AccountUsageStatisticsCostView>,
pub day_count: u32,
pub has_unknown_pricing: bool,
pub has_missing_token_data: bool,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsModelView {
pub key: String,
pub model: String,
pub service_tier: String,
pub credit_share: Option<f64>,
pub quota_share: Option<f64>,
pub tokens: AccountUsageStatisticsTokensView,
pub estimated_cost: Option<AccountUsageStatisticsCostView>,
pub has_unknown_pricing: bool,
pub has_estimated_allocation: bool,
pub has_rate_fallback: bool,
pub has_missing_token_data: bool,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUsageStatisticsDayView {
pub date: String,
pub credit_share: Option<f64>,
pub tokens: AccountUsageStatisticsTokensView,
pub estimated_cost: Option<AccountUsageStatisticsCostView>,
pub has_unknown_pricing: bool,
pub has_missing_token_data: bool,
pub is_boundary_day: bool,
}
impl From<ProviderUsageStatistics> for AccountUsageStatisticsData {
fn from(statistics: ProviderUsageStatistics) -> Self {
Self {
cycle: AccountUsageStatisticsCycleView {
offset: statistics.cycle.offset,
start_at: statistics.cycle.start_at.to_rfc3339(),
end_at: statistics.cycle.end_at.to_rfc3339(),
window_seconds: statistics.cycle.window_seconds,
used_percent: statistics.cycle.used_percent,
is_current: statistics.cycle.is_current,
can_go_previous: statistics.cycle.can_go_previous,
can_go_next: statistics.cycle.can_go_next,
},
summary: account_usage_statistics_summary_view(statistics.summary),
models: statistics
.models
.into_iter()
.map(account_usage_statistics_model_view)
.collect(),
daily: statistics
.daily
.into_iter()
.map(account_usage_statistics_day_view)
.collect(),
}
}
}
fn account_usage_statistics_tokens_view(
tokens: ProviderUsageStatisticsTokens,
) -> AccountUsageStatisticsTokensView {
AccountUsageStatisticsTokensView {
uncached_input: tokens.uncached_input,
cached_input: tokens.cached_input,
output: tokens.output,
total: tokens.total,
}
}
fn account_usage_statistics_cost_view(
cost: gateway_admin::model::observability::CurrencyCost,
) -> AccountUsageStatisticsCostView {
AccountUsageStatisticsCostView {
currency: cost.currency,
amount: cost.amount.to_string(),
}
}
fn account_usage_statistics_summary_view(
summary: ProviderUsageStatisticsSummary,
) -> AccountUsageStatisticsSummaryView {
AccountUsageStatisticsSummaryView {
tokens: account_usage_statistics_tokens_view(summary.tokens),
estimated_cost: summary
.estimated_cost
.map(account_usage_statistics_cost_view),
projected_tokens: summary.projected_tokens,
projected_cost: summary
.projected_cost
.map(account_usage_statistics_cost_view),
day_count: summary.day_count,
has_unknown_pricing: summary.has_unknown_pricing,
has_missing_token_data: summary.has_missing_token_data,
}
}
fn account_usage_statistics_model_view(
model: ProviderUsageStatisticsModel,
) -> AccountUsageStatisticsModelView {
AccountUsageStatisticsModelView {
key: model.key,
model: model.model,
service_tier: model.service_tier.as_str().to_owned(),
credit_share: model.credit_share,
quota_share: model.quota_share,
tokens: account_usage_statistics_tokens_view(model.tokens),
estimated_cost: model.estimated_cost.map(account_usage_statistics_cost_view),
has_unknown_pricing: model.has_unknown_pricing,
has_estimated_allocation: model.has_estimated_allocation,
has_rate_fallback: model.has_rate_fallback,
has_missing_token_data: model.has_missing_token_data,
}
}
fn account_usage_statistics_day_view(
day: ProviderUsageStatisticsDay,
) -> AccountUsageStatisticsDayView {
AccountUsageStatisticsDayView {
date: day.date.to_string(),
credit_share: day.credit_share,
tokens: account_usage_statistics_tokens_view(day.tokens),
estimated_cost: day.estimated_cost.map(account_usage_statistics_cost_view),
has_unknown_pricing: day.has_unknown_pricing,
has_missing_token_data: day.has_missing_token_data,
is_boundary_day: day.is_boundary_day,
}
}
/// 主动额度重置卡列表响应。
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountResetCreditsData {
pub available_count: u64,
pub credits: Vec<AccountResetCreditView>,
}
/// 一张安全主动额度重置卡视图。
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountResetCreditView {
pub id: String,
pub status: Option<String>,
pub title: Option<String>,
pub expires_at: Option<String>,
pub reset_type: Option<String>,
}
/// 主动额度重置卡消费响应。
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountResetCreditResultData {
pub code: String,
pub credit: Option<AccountResetCreditView>,
}
impl From<ProviderResetCredit> for AccountResetCreditView {
fn from(credit: ProviderResetCredit) -> Self {
Self {
id: credit.id,
status: credit.status,
title: credit.title,
expires_at: credit.expires_at.map(|value| value.to_rfc3339()),
reset_type: credit.reset_type,
}
}
}
impl From<ProviderResetCredits> for AccountResetCreditsData {
fn from(credits: ProviderResetCredits) -> Self {
Self {
available_count: credits.available_count,
credits: credits
.credits
.into_iter()
.map(AccountResetCreditView::from)
.collect(),
}
}
}
impl From<ProviderResetCreditResult> for AccountResetCreditResultData {
fn from(result: ProviderResetCreditResult) -> Self {
Self {
code: result.code,
credit: result.credit.map(AccountResetCreditView::from),
}
}
}
/// Provider-owned 明文导出文档;Debug 永远不输出内部 JSON。
#[derive(Serialize)]
#[serde(transparent)]
@@ -1055,6 +1371,14 @@ where
post(batch_update_accounts::<S>),
)
.route("/api/admin/accounts/quota", get(account_quota::<S>))
.route(
"/api/admin/accounts/usage-statistics",
get(account_usage_statistics::<S>),
)
.route(
"/api/admin/accounts/reset-credits",
get(account_reset_credits::<S>).post(consume_account_reset_credit::<S>),
)
.route(
"/api/admin/accounts/quota/refresh",
post(refresh_account_quota::<S>),
@@ -1385,6 +1709,27 @@ where
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn account_usage_statistics<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<AccountUsageStatisticsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let request = query.into_request().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.usage_statistics(request)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountUsageStatisticsData::from(result)),
))
}
async fn refresh_account_quota<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -1406,6 +1751,48 @@ where
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn account_reset_credits<S>(
auth: AdminAuth,
State(state): State<S>,
Query(query): Query<AccountIdQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let account_id = query.into_id().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.reset_credits(&auth.context().mutation_context(), account_id)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountResetCreditsData::from(result)),
))
}
async fn consume_account_reset_credit<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountResetCreditConsumeRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request.into_command().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.consume_reset_credit(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountResetCreditResultData::from(result)),
))
}
async fn account_models<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -142,11 +142,19 @@ pub(crate) fn format_service_tier(service_tier: Option<&str>) -> String {
match service_tier {
Some("priority" | "fast") => "Fast".to_owned(),
Some("flex") => "Flex".to_owned(),
Some("default") | None => "Default".to_owned(),
Some(other) => other.to_owned(),
Some("default" | "standard") | None => "Standard".to_owned(),
Some(other) => capitalize_first(other),
}
}
fn capitalize_first(value: &str) -> String {
let mut chars = value.chars();
chars
.next()
.map(|first| first.to_uppercase().chain(chars).collect())
.unwrap_or_default()
}
pub(crate) fn billing_view(billing: Option<&domain::UsageBilling>) -> Option<BillingView> {
match billing? {
domain::UsageBilling::Total { total, .. } => Some(BillingView {
@@ -222,7 +230,7 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
let model = record
.upstream_model_id
.clone()
.unwrap_or_else(|| record.requested_model_id.clone());
.or_else(|| record.requested_model_id.clone());
let transport = record
.upstream_transport
.clone()
@@ -243,7 +251,7 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
account_name: record.provider_account_name,
route: record.endpoint,
model,
requested_model: Some(record.requested_model_id),
requested_model: record.requested_model_id,
upstream_model: record.upstream_model_id,
service_tier: record.service_tier,
status_code,
@@ -351,9 +359,7 @@ pub(crate) fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptV
provider: attempt
.provider_kind
.unwrap_or_else(|| "unknown".to_owned()),
model: attempt
.upstream_model_id
.unwrap_or_else(|| "unknown".to_owned()),
model: attempt.upstream_model_id,
transport: attempt
.upstream_transport
.unwrap_or_else(|| "unknown".to_owned()),
@@ -731,21 +737,19 @@ pub(crate) fn dashboard_view(
last_used: relative_time(credential.last_used_at, range.end),
});
}
let unavailable_accounts = provider_accounts
.total
.saturating_sub(provider_accounts.normal);
DashboardDataView {
cards: DashboardCardsView {
credentials: DashboardCredentialsCardView {
total: format_compact_number(provider_accounts.total),
total: provider_accounts.total.to_string(),
total_value: provider_accounts.total,
available: format_compact_number(provider_accounts.normal),
available: provider_accounts.normal.to_string(),
available_value: provider_accounts.normal,
unavailable: format_compact_number(
provider_accounts
.total
.saturating_sub(provider_accounts.normal),
),
unavailable_value: provider_accounts
.total
.saturating_sub(provider_accounts.normal),
unavailable: unavailable_accounts.to_string(),
unavailable_value: unavailable_accounts,
},
traffic: DashboardTrafficCardView {
today_requests: format_compact_number(today.request_count),
@@ -94,7 +94,7 @@ pub struct UsageRecordView {
pub account_email: Option<String>,
pub account_name: Option<String>,
pub route: String,
pub model: String,
pub model: Option<String>,
pub requested_model: Option<String>,
pub upstream_model: Option<String>,
pub service_tier: Option<String>,
@@ -183,7 +183,7 @@ pub struct UsageAttemptView {
pub attempt_index: u32,
pub trigger: String,
pub provider: String,
pub model: String,
pub model: Option<String>,
pub transport: String,
pub send_state: String,
pub outcome: String,
@@ -1,7 +1,7 @@
//! Admin API 的展示层辅助:数值格式化。
//!
//! 按 docs/project-redundancy-boundary-audit.md BE-02 的目标,display 字段最终
//! 迁移到前端;过渡期统一由本模块提供,避免 accounts/observability 各自维护。
//! 当前 wire 仍保留 display 字段;统一由本模块提供,避免 accounts/observability
//! 各自维护格式化规则。
use gateway_core::accounting::Decimal;
@@ -0,0 +1,233 @@
//! Codex Images 非流式 HTTP adapter。
use std::net::SocketAddr;
use axum::{
body::{Body, Bytes},
extract::{Extension, State, connect_info::ConnectInfo},
http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header::CONTENT_TYPE},
response::{IntoResponse, Response},
};
use gateway_core::engine::execution::StartedExecution;
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
use gateway_core::operation::{ImageRequest, ImageRequestKind, Operation, RawJsonPayload};
use serde_json::{Map, Value};
use crate::ApiState;
use crate::openai::{
auth::{authenticate_client, authentication_error_response},
error::{engine_error_response, gateway_error_response, protocol_error_response},
responses::{PendingExecution, ProtocolError, ProtocolErrorBody, request_client_context},
};
const OPENAI_PROTOCOL: &str = "openai";
const IMAGE_TURN_ID_CONTEXT_KEY: &str = "image_turn_id";
/// `POST /v1/images/generations`。
pub(crate) async fn image_generations(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_image_request(
state,
connect_info,
headers,
body,
ImageRequestKind::Generation,
"/v1/images/generations",
)
.await
}
/// `POST /v1/images/edits`。
pub(crate) async fn image_edits(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_image_request(
state,
connect_info,
headers,
body,
ImageRequestKind::Edit,
"/v1/images/edits",
)
.await
}
async fn handle_image_request(
state: ApiState,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
headers: HeaderMap,
body: Bytes,
kind: ImageRequestKind,
endpoint: &'static str,
) -> Response {
let service = state.openai();
let client = match authenticate_client(service, &headers) {
Ok(client) => client,
Err(error) => return authentication_error_response(error),
};
let (client_ip, user_agent) = request_client_context(
&headers,
connect_info.map(|Extension(ConnectInfo(address))| address),
);
let operation = match image_operation(body, &headers, kind) {
Ok(operation) => operation,
Err(error) => return gateway_error_response(&error),
};
let started = match service
.start_image(client, operation, client_ip, user_agent, endpoint)
.await
{
Ok(started) => started,
Err(error) => return gateway_error_response(&error),
};
collect_image_response(started).await
}
fn image_operation(
body: Bytes,
headers: &HeaderMap,
kind: ImageRequestKind,
) -> Result<Operation, GatewayError> {
let mut context = Map::new();
if let Some(turn_id) = headers
.get("x-codex-image-turn-id")
.and_then(|value| value.to_str().ok())
{
context.insert(
IMAGE_TURN_ID_CONTEXT_KEY.to_owned(),
Value::String(turn_id.to_owned()),
);
}
let payload = RawJsonPayload::new(OPENAI_PROTOCOL, body)
.map_err(|_| {
GatewayError::new(
GatewayErrorKind::Internal,
"OpenAI protocol identifier is invalid",
)
})?
.with_context(context);
Ok(Operation::GenerateImage(ImageRequest::from_raw_json(
kind, payload,
)))
}
async fn collect_image_response(started: StartedExecution) -> Response {
let mut execution = PendingExecution::new(started.session);
let Some(session) = execution.session_mut() else {
return invalid_upstream_response();
};
let events = match session.collect_uncommitted().await {
Ok(events) => events,
Err(error) => {
let response = engine_error_response(&error);
return execution.record_response_status(response).await;
}
};
let body = match raw_json_body(events) {
Some(body) => body,
None => {
let response = invalid_upstream_response();
let response = execution.record_response_status(response).await;
execution.cancel_and_finalize().await;
return response;
}
};
let status = session
.response_status_code()
.and_then(|status| StatusCode::from_u16(status).ok())
.filter(StatusCode::is_success)
.unwrap_or(StatusCode::OK);
let response_headers = session.response_headers().to_vec();
let response = json_body_response(body, status, &response_headers);
let Some(session) = execution.session_mut() else {
return invalid_upstream_response();
};
if let Err(error) = session.commit_downstream(Some(status.as_u16())).await {
execution.cancel_and_finalize().await;
let response = engine_error_response(&error);
return execution.record_response_status(response).await;
}
if !session.is_finalized() {
execution.cancel_and_finalize().await;
return invalid_upstream_response();
}
execution.disarm();
response
}
fn raw_json_body(events: Vec<ProviderEvent>) -> Option<Bytes> {
let mut body = None;
for event in events {
let (_, wire) = event.into_parts();
let Some(wire) = wire.filter(|wire| wire.protocol() == "openai") else {
continue;
};
let Some(raw) = wire.into_raw_json_body() else {
continue;
};
if body.replace(raw).is_some() {
return None;
}
}
body
}
fn json_body_response(
body: Bytes,
status: StatusCode,
response_headers: &[ProviderResponseHeader],
) -> Response {
let mut response = Response::new(Body::from(body));
*response.status_mut() = status;
response
.headers_mut()
.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
apply_response_headers(response, response_headers)
}
fn apply_response_headers(
mut response: Response,
response_headers: &[ProviderResponseHeader],
) -> Response {
let connection_options =
crate::openai::responses::response_connection_options(response_headers);
for header in response_headers {
if !crate::openai::responses::response_header_is_forwardable(
header.name(),
&connection_options,
) {
continue;
}
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
continue;
};
let Ok(value) = HeaderValue::from_bytes(header.value()) else {
continue;
};
response.headers_mut().append(name, value);
}
response
}
fn invalid_upstream_response() -> Response {
protocol_error_response(
StatusCode::BAD_GATEWAY,
ProtocolErrorBody {
error: ProtocolError {
kind: "server_error",
code: "invalid_upstream_response",
message: "The gateway could not forward the upstream image response.".to_owned(),
param: None,
},
},
)
.into_response()
}
@@ -0,0 +1,5 @@
//! OpenAI/Codex Images JSON 透明 adapter。
mod http;
pub(crate) use http::{image_edits, image_generations};
@@ -2,6 +2,7 @@
pub mod auth;
pub mod error;
pub mod images;
pub mod models;
pub mod responses;
pub mod router;
@@ -123,7 +123,7 @@ async fn handle_responses(
}
/// 从 socket 与标准转发头提取旧 Usage 页面使用的诊断事实。
pub(super) fn request_client_context(
pub(in crate::openai) fn request_client_context(
headers: &HeaderMap,
peer_address: Option<SocketAddr>,
) -> (Option<IpAddr>, Option<String>) {
@@ -367,22 +367,24 @@ fn internal_gateway_response(message: &'static str) -> Response {
gateway_error_response(&GatewayError::new(GatewayErrorKind::Internal, message))
}
pub(super) struct PendingExecution {
pub(in crate::openai) struct PendingExecution {
session: Option<Box<dyn ExecutionSession>>,
}
impl PendingExecution {
pub(super) fn new(session: Box<dyn ExecutionSession>) -> Self {
pub(in crate::openai) fn new(session: Box<dyn ExecutionSession>) -> Self {
Self {
session: Some(session),
}
}
pub(super) fn session_mut(&mut self) -> Option<&mut (dyn ExecutionSession + 'static)> {
pub(in crate::openai) fn session_mut(
&mut self,
) -> Option<&mut (dyn ExecutionSession + 'static)> {
self.session.as_deref_mut()
}
async fn cancel_and_finalize(&mut self) {
pub(in crate::openai) async fn cancel_and_finalize(&mut self) {
if let Some(session) = self.session.as_mut() {
session.cancel();
if !session.is_finalized() {
@@ -391,7 +393,10 @@ impl PendingExecution {
}
}
async fn record_response_status(&mut self, response: Response) -> Response {
pub(in crate::openai) async fn record_response_status(
&mut self,
response: Response,
) -> Response {
if let Some(session) = self.session.as_mut() {
let _ = session
.record_client_status(response.status().as_u16())
@@ -400,7 +405,7 @@ impl PendingExecution {
response
}
pub(super) fn disarm(&mut self) {
pub(in crate::openai) fn disarm(&mut self) {
self.session = None;
}
@@ -6,7 +6,8 @@ mod request;
mod response;
mod websocket;
pub use error::{ProtocolErrorBody, RequestDecodeError, ResponseEncodeError};
pub use error::{ProtocolError, ProtocolErrorBody, RequestDecodeError, ResponseEncodeError};
pub(super) use http::{PendingExecution, request_client_context};
pub use http::{collect_execution_response, stream_execution_response};
pub(crate) use http::{responses, review_responses};
pub use request::{
@@ -165,7 +165,6 @@ impl fmt::Debug for ContinuationIntent {
#[derive(Clone, PartialEq, Eq)]
pub struct ResponsesRequestMetadata {
requested_model: String,
public_model: String,
stream: bool,
store: bool,
continuation: ContinuationIntent,
@@ -180,12 +179,6 @@ impl ResponsesRequestMetadata {
&self.requested_model
}
/// 返回协议层公开模型候选。
#[must_use]
pub fn public_model(&self) -> &str {
&self.public_model
}
/// 返回客户端是否请求 SSE。
#[must_use]
pub const fn stream(&self) -> bool {
@@ -222,7 +215,6 @@ impl fmt::Debug for ResponsesRequestMetadata {
formatter
.debug_struct("ResponsesRequestMetadata")
.field("requested_model", &self.requested_model)
.field("public_model", &self.public_model)
.field("stream", &self.stream)
.field("store", &self.store)
.field("continuation", &self.continuation)
@@ -366,8 +358,7 @@ pub(super) fn decode_request_object(
Ok(DecodedResponsesRequest {
operation,
metadata: ResponsesRequestMetadata {
requested_model: model.clone(),
public_model: model,
requested_model: model,
stream,
store,
continuation,
@@ -7,6 +7,7 @@ use axum::{
};
use super::{
images::{image_edits, image_generations},
models::{model_catalog, model_detail, model_info, models},
responses::{responses, responses_websocket, review_responses},
};
@@ -16,6 +17,8 @@ use crate::ApiState;
/// 构造 OpenAI 客户端协议路由。
pub(crate) fn router() -> Router<ApiState> {
Router::new()
.route("/v1/images/generations", post(image_generations))
.route("/v1/images/edits", post(image_edits))
.route("/v1/responses", get(responses_websocket).post(responses))
.route("/v1/responses/review", post(review_responses))
.route("/v1/models", get(models))
@@ -23,7 +26,7 @@ pub(crate) fn router() -> Router<ApiState> {
.route("/v1/models/{model_id}/info", get(model_info))
// 官方 OpenAI 模型详情合同使用 path ID;它不属于 Admin API 约束。
.route("/v1/models/{model_id}", get(model_detail))
// Responses 正文属于 OpenAI/Codex 协议;代理不能用私有大小上限提前拒绝
// OpenAI 数据面正文属于客户端/上游协议;代理不能用私有大小上限提前拒绝
// 上游本可接受的未来 payload。
.layer(DefaultBodyLimit::disable())
}
@@ -1,11 +1,12 @@
//! OpenAI wire adapter 到 Core 执行用例的唯一映射。
use std::net::IpAddr;
use std::sync::Arc;
use gateway_core::engine::continuation::PreviousResponseId;
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ClientTransport, ExecutionRequestMetadata,
ExecutionService, StartExecution, StartedExecution,
ExecutionService, StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle};
@@ -74,13 +75,13 @@ impl OpenAiService {
endpoint: &'static str,
) -> Result<StartedExecution, GatewayError> {
let (operation, metadata) = request.into_parts();
let public_model = PublicModelId::from_client_wire(metadata.public_model().to_owned())
let public_model = PublicModelId::from_client_wire(metadata.requested_model().to_owned())
.map_err(|_| {
GatewayError::new(
GatewayErrorKind::ModelNotFound,
"requested model was not found",
)
})?;
GatewayError::new(
GatewayErrorKind::ModelNotFound,
"requested model was not found",
)
})?;
let previous_response_id = match metadata.continuation() {
ContinuationIntent::None => None,
ContinuationIntent::PreviousResponseId(value) => {
@@ -105,6 +106,38 @@ impl OpenAiService {
.await
}
pub(crate) async fn start_image(
&self,
client: AuthenticatedClient,
operation: gateway_core::operation::Operation,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
endpoint: &'static str,
) -> Result<StartedExecution, GatewayError> {
let provider = gateway_core::routing::ProviderKind::new("openai").map_err(|_| {
GatewayError::new(
GatewayErrorKind::Internal,
"OpenAI provider identifier is invalid",
)
})?;
self.execution
.start_provider_endpoint(StartProviderExecution {
client,
provider,
operation,
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: endpoint.to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip,
user_agent,
previous_response_id: None,
},
})
.await
}
pub(crate) fn try_register_connection(
&self,
) -> Result<Box<dyn ConnectionGuard>, ConnectionDraining> {
@@ -70,6 +70,147 @@ mod query {
}
}
mod usage_statistics {
use std::str::FromStr as _;
use chrono::{NaiveDate, TimeZone as _, Utc};
use gateway_admin::model::{
observability::{CurrencyCost, DecimalAmount},
provider_credentials::{
ProviderUsageStatistics, ProviderUsageStatisticsCycle, ProviderUsageStatisticsDay,
ProviderUsageStatisticsModel, ProviderUsageStatisticsServiceTier,
ProviderUsageStatisticsSummary, ProviderUsageStatisticsTokens,
},
};
use gateway_api::admin::accounts::{AccountUsageStatisticsData, AccountUsageStatisticsQuery};
use serde_json::json;
fn tokens(total: u64) -> ProviderUsageStatisticsTokens {
ProviderUsageStatisticsTokens {
uncached_input: total / 4,
cached_input: total / 2,
output: total / 4,
total,
}
}
fn usd(amount: &str) -> CurrencyCost {
CurrencyCost {
currency: "USD".to_owned(),
amount: DecimalAmount::from_str(amount).expect("USD amount"),
}
}
#[test]
fn usage_statistics_query_accepts_only_current_or_bounded_history_offsets() {
let query: AccountUsageStatisticsQuery = serde_json::from_value(json!({
"accountId": "acct_openai",
"cycleOffset": -3,
"utcOffsetMinutes": 480
}))
.expect("decode usage-statistics query");
let request = query.into_request().expect("valid cycle offset");
assert_eq!(request.cycle_offset, -3);
assert_eq!(request.utc_offset_minutes, 480);
for (cycle_offset, field) in [(-9, "cycleOffset"), (1, "cycleOffset")] {
let query: AccountUsageStatisticsQuery = serde_json::from_value(json!({
"accountId": "acct_openai",
"cycleOffset": cycle_offset
}))
.expect("decode bounded integer");
assert_eq!(
query
.into_request()
.expect_err("reject cycle offset")
.field(),
field
);
}
assert!(
serde_json::from_value::<AccountUsageStatisticsQuery>(json!({
"accountId": "acct_openai",
"historyCycles": 8
}))
.is_err()
);
}
#[test]
fn usage_statistics_response_serializes_calculated_rows_without_raw_documents() {
let start_at = Utc
.with_ymd_and_hms(2026, 8, 24, 7, 50, 0)
.single()
.expect("cycle start");
let end_at = Utc
.with_ymd_and_hms(2026, 8, 31, 7, 50, 0)
.single()
.expect("cycle end");
let response = AccountUsageStatisticsData::from(ProviderUsageStatistics {
cycle: ProviderUsageStatisticsCycle {
offset: 0,
start_at,
end_at,
window_seconds: 604_800,
used_percent: Some(20.0),
is_current: true,
can_go_previous: true,
can_go_next: false,
},
summary: ProviderUsageStatisticsSummary {
tokens: tokens(100),
estimated_cost: Some(usd("1.25")),
projected_tokens: Some(500),
projected_cost: Some(usd("6.25")),
day_count: 1,
has_unknown_pricing: false,
has_missing_token_data: false,
},
models: vec![ProviderUsageStatisticsModel {
key: "gpt-5.6-sol::standard".to_owned(),
model: "gpt-5.6-sol".to_owned(),
service_tier: ProviderUsageStatisticsServiceTier::Standard,
credit_share: Some(1.0),
quota_share: Some(0.2),
tokens: tokens(100),
estimated_cost: Some(usd("1.25")),
has_unknown_pricing: false,
has_estimated_allocation: true,
has_rate_fallback: false,
has_missing_token_data: false,
}],
daily: vec![ProviderUsageStatisticsDay {
date: NaiveDate::from_ymd_opt(2026, 8, 25).expect("report date"),
credit_share: Some(1.0),
tokens: tokens(100),
estimated_cost: Some(usd("1.25")),
has_unknown_pricing: false,
has_missing_token_data: false,
is_boundary_day: true,
}],
});
let value = serde_json::to_value(response).expect("serialize usage statistics");
assert!(value.get("mode").is_none());
assert_eq!(value["cycle"]["offset"], 0);
assert_eq!(
value["summary"]["estimatedCost"],
json!({
"currency": "USD",
"amount": "1.25"
})
);
assert_eq!(value["models"][0]["serviceTier"], "standard");
assert!(value["models"][0].get("mode").is_none());
assert_eq!(value["daily"][0]["tokens"]["total"], 100);
assert_eq!(value["daily"][0]["isBoundaryDay"], true);
assert!(value["cycle"].get("deltaPercent").is_none());
assert!(value.get("modelBreakdown").is_none());
assert!(value.get("dailyTotals").is_none());
assert!(value.get("usage").is_none());
}
}
mod batch_update {
use gateway_api::admin::accounts::{BatchUpdateAccountsRequest, UpdateAccountRequest};
use serde_json::json;
@@ -266,8 +407,8 @@ mod actions {
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
AccountImportData, AccountImportRequest, AccountMutationData, AccountRefreshRequest,
AccountTestQuery, CompleteAccountAuthorizationRequest, RotateAccountRequest,
StartAccountAuthorizationRequest,
AccountResetCreditConsumeRequest, AccountTestQuery, CompleteAccountAuthorizationRequest,
RotateAccountRequest, StartAccountAuthorizationRequest,
};
use gateway_core::engine::credential::ProviderAccountId;
use serde_json::json;
@@ -325,6 +466,44 @@ mod actions {
);
}
#[test]
fn reset_credit_consume_should_require_canonical_v4_idempotency_key() {
let valid: AccountResetCreditConsumeRequest = serde_json::from_value(json!({
"accountId": "acct_1",
"creditId": "credit_1",
"redeemRequestId": "8fbf302d-11df-4bd5-82e4-08e4b3df7874"
}))
.expect("decode reset-credit consume");
valid.validate().expect("validate reset-credit consume");
for (redeem_request_id, field) in [
("019c0000-0000-7000-8000-000000000000", "redeemRequestId"),
("8FBF302D-11DF-4BD5-82E4-08E4B3DF7874", "redeemRequestId"),
("invalid", "redeemRequestId"),
] {
let request: AccountResetCreditConsumeRequest = serde_json::from_value(json!({
"accountId": "acct_1",
"redeemRequestId": redeem_request_id
}))
.expect("decode invalid reset-credit consume");
assert_eq!(request.validate().expect_err("reject UUID").field(), field);
}
let invalid_credit: AccountResetCreditConsumeRequest = serde_json::from_value(json!({
"accountId": "acct_1",
"creditId": " ",
"redeemRequestId": "8fbf302d-11df-4bd5-82e4-08e4b3df7874"
}))
.expect("decode invalid credit");
assert_eq!(
invalid_credit
.validate()
.expect_err("reject credit")
.field(),
"creditId"
);
}
#[test]
fn credential_recovery_requests_should_not_accept_client_revision_fences() {
let authorization: StartAccountAuthorizationRequest = serde_json::from_value(json!({
@@ -221,9 +221,9 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
attempts: AttemptMetrics::default(),
totals: Default::default(),
provider_accounts: AccountPoolMetrics {
total: 807,
normal: 15,
quota_exhausted: 785,
total: 1002,
normal: 943,
quota_exhausted: 52,
rate_limited: 0,
disabled: 3,
error: 4,
@@ -259,9 +259,21 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
&value["data"]["cards"]["credentials"]["unavailableValue"],
),
(
&serde_json::json!(807),
&serde_json::json!(15),
&serde_json::json!(792)
&serde_json::json!(1002),
&serde_json::json!(943),
&serde_json::json!(59)
),
);
assert_eq!(
(
&value["data"]["cards"]["credentials"]["total"],
&value["data"]["cards"]["credentials"]["available"],
&value["data"]["cards"]["credentials"]["unavailable"],
),
(
&serde_json::json!("1002"),
&serde_json::json!("943"),
&serde_json::json!("59")
),
);
}
@@ -681,7 +693,7 @@ fn usage_record_with_account(
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
requested_model_id: Some("coding".to_owned()),
provider_kind: Some("openai".to_owned()),
provider_account_ref: Some(account_id.to_owned()),
provider_account_name: Some(name.to_owned()),
@@ -810,7 +822,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
operation: "generate".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "grok-4.5".to_owned(),
requested_model_id: Some("grok-4.5".to_owned()),
provider_kind: Some("xai".to_owned()),
provider_account_ref: Some("acct_snapshot".to_owned()),
provider_account_name: Some("Snapshot Alpha".to_owned()),
@@ -820,7 +832,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
upstream_transport: Some("http_sse".to_owned()),
http_version: Some("h2".to_owned()),
websocket_pool: Some("reuse".to_owned()),
service_tier: Some("priority".to_owned()),
service_tier: Some("default".to_owned()),
provider_metadata_json: Some(
serde_json::json!({
"effectiveModel": "grok-4.5",
@@ -865,7 +877,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
output_price_per_million: usd("60.0000"),
cache_read_price_per_million: usd("1.0000"),
cache_write_price_per_million: usd("12.5000"),
service_tier: Some("priority".to_owned()),
service_tier: Some("default".to_owned()),
multiplier_percent: 100,
},
))),
@@ -921,6 +933,10 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
value["data"]["items"][0]["billing"]["cacheWritePriceDisplay"],
"$12.5 / 1M Token"
);
assert_eq!(
value["data"]["items"][0]["billing"]["serviceTierDisplay"],
"Standard"
);
assert_eq!(
serde_json::json!({
@@ -961,7 +977,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
"routingScope": "groups",
"routingGroupRefs": ["grp_usage"],
"routingGroupNamesSnapshot": ["Usage group"],
"serviceTier": "priority",
"serviceTier": "default",
"accountId": "acct_snapshot",
"accountName": "Snapshot Alpha",
"accountEmail": "alpha@example.invalid",
@@ -52,6 +52,8 @@ fn source_tree_should_match_frozen_machine_manifest() {
"src/lib.rs",
"src/openai/auth.rs",
"src/openai/error.rs",
"src/openai/images/http.rs",
"src/openai/images/mod.rs",
"src/openai/mod.rs",
"src/openai/models.rs",
"src/openai/responses/error.rs",
@@ -90,6 +92,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
"tests/main.rs",
"tests/openai/auth.rs",
"tests/openai/error.rs",
"tests/openai/images.rs",
"tests/openai/mod.rs",
"tests/openai/models.rs",
"tests/openai/responses/http.rs",
+8 -1
View File
@@ -8,7 +8,7 @@ use futures::future::BoxFuture;
use gateway_core::{
engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, StartExecution,
StartedExecution,
StartProviderExecution, StartedExecution,
},
error::GatewayError,
health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState},
@@ -120,4 +120,11 @@ impl ExecutionService for UnusedExecution {
fn start(&self, _: StartExecution) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async { unreachable!("health check does not execute requests") })
}
fn start_provider_endpoint(
&self,
_: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async { unreachable!("health check does not execute provider endpoints") })
}
}
@@ -0,0 +1,331 @@
use std::sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
};
use std::time::SystemTime;
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header::AUTHORIZATION},
};
use bytes::Bytes;
use futures::future::BoxFuture;
use gateway_core::engine::UpstreamSendState;
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ClientTransport, ExecutionService,
ExecutionSession, StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::engine::{CoordinatedEvent, EngineError, ModelRequestId};
use gateway_core::error::{
ClientVisibleUpstreamResponse, GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind,
};
use gateway_core::event::{ProtocolWireEvent, ProviderEvent, ProviderResponseHeader};
use gateway_core::operation::{ImageRequestKind, Operation};
use gateway_core::routing::PublicModelId;
use serde_json::{Value, json};
use tower::ServiceExt;
use super::{api_router, authenticated_client};
const IMAGE_RESPONSE: &[u8] =
br#"{ "created": 1787212800, "data": [{"b64_json":"AAEC"}], "future": 9007199254740993 }"#;
#[derive(Debug, Clone, PartialEq)]
struct CapturedImageRequest {
provider: String,
endpoint: String,
transport: ClientTransport,
kind: ImageRequestKind,
body: Bytes,
context: Value,
}
struct ImageExecution {
client: AuthenticatedClient,
captured: Arc<Mutex<Vec<CapturedImageRequest>>>,
committed_statuses: Arc<Mutex<Vec<u16>>>,
fail_with_upstream_response: bool,
}
impl ImageExecution {
fn new() -> Arc<Self> {
Arc::new(Self {
client: authenticated_client("sk_images_test"),
captured: Arc::new(Mutex::new(Vec::new())),
committed_statuses: Arc::new(Mutex::new(Vec::new())),
fail_with_upstream_response: false,
})
}
fn failing() -> Arc<Self> {
Arc::new(Self {
client: authenticated_client("sk_images_test"),
captured: Arc::new(Mutex::new(Vec::new())),
committed_statuses: Arc::new(Mutex::new(Vec::new())),
fail_with_upstream_response: true,
})
}
fn captured(&self) -> Vec<CapturedImageRequest> {
self.captured.lock().expect("capture lock").clone()
}
fn committed_statuses(&self) -> Vec<u16> {
self.committed_statuses
.lock()
.expect("committed status lock")
.clone()
}
}
impl ExecutionService for ImageExecution {
fn authenticate(
&self,
plaintext: &str,
) -> Result<AuthenticatedClient, ClientAuthenticationError> {
(plaintext == "sk_images_test")
.then(|| self.client.clone())
.ok_or(ClientAuthenticationError::InvalidKey)
}
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
Vec::new()
}
fn contains_public_model(&self, _: &AuthenticatedClient, _: &PublicModelId) -> bool {
false
}
fn start(&self, _: StartExecution) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async {
Err(GatewayError::new(
GatewayErrorKind::Internal,
"images test must use provider endpoint execution",
))
})
}
fn start_provider_endpoint(
&self,
request: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
let Operation::GenerateImage(image) = &request.operation else {
return Err(GatewayError::new(
GatewayErrorKind::Internal,
"images test received a non-image operation",
));
};
let payload = image.payload();
self.captured
.lock()
.expect("capture lock")
.push(CapturedImageRequest {
provider: request.provider.as_str().to_owned(),
endpoint: request.metadata.endpoint,
transport: request.metadata.transport,
kind: image.kind(),
body: payload.body().clone(),
context: Value::Object(payload.context().clone()),
});
Ok(StartedExecution {
request_id: ModelRequestId::new("req_images_test").expect("request ID"),
created_at: SystemTime::now(),
stream: false,
session: Box::new(ImageSession {
response: Some(Bytes::from_static(IMAGE_RESPONSE)),
response_headers: vec![ProviderResponseHeader::new(
"x-image-rate-limit",
Bytes::from_static(b"42"),
)],
committed_statuses: Arc::clone(&self.committed_statuses),
finalized: Arc::new(AtomicBool::new(false)),
fail_with_upstream_response: self.fail_with_upstream_response,
}),
})
})
}
}
struct ImageSession {
response: Option<Bytes>,
response_headers: Vec<ProviderResponseHeader>,
committed_statuses: Arc<Mutex<Vec<u16>>>,
finalized: Arc<AtomicBool>,
fail_with_upstream_response: bool,
}
impl ExecutionSession for ImageSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async { unreachable!("buffered image delivery does not stream") })
}
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
Box::pin(async move {
if self.fail_with_upstream_response {
return Err(EngineError::Provider(
ProviderError::new(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::Sent,
)
.with_status(StatusCode::UNPROCESSABLE_ENTITY.as_u16())
.with_client_visible_upstream_response(
ClientVisibleUpstreamResponse::new(
StatusCode::UNPROCESSABLE_ENTITY.as_u16(),
Some(b"application/problem+json".to_vec()),
Bytes::from_static(
br#"{ "error":{"message":"future image validation"}, "future":9007199254740993 }"#,
),
)
.with_headers(vec![ProviderResponseHeader::new(
"x-future-image-error",
Bytes::from_static(b"preserved"),
)]),
),
));
}
Ok(vec![ProviderEvent::wire(
ProtocolWireEvent::raw_json(
"openai",
self.response.take().expect("single image response"),
)
.expect("OpenAI protocol"),
)])
})
}
fn response_headers(&self) -> &[ProviderResponseHeader] {
&self.response_headers
}
fn response_status_code(&self) -> Option<u16> {
Some(StatusCode::CREATED.as_u16())
}
fn commit_downstream(
&mut self,
client_status_code: Option<u16>,
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
self.committed_statuses
.lock()
.expect("committed status lock")
.push(client_status_code.expect("image response status"));
self.finalized.store(true, Ordering::Release);
Ok(())
})
}
fn record_client_status(&mut self, _: u16) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async { Ok(()) })
}
fn is_finalized(&self) -> bool {
self.finalized.load(Ordering::Acquire)
}
fn cancel(&self) {
self.finalized.store(true, Ordering::Release);
}
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move {
self.finalized.store(true, Ordering::Release);
})
}
}
#[tokio::test]
async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directions() {
let execution = ImageExecution::new();
let router = api_router(execution.clone()).await;
let cases = [
(
"/v1/images/generations",
ImageRequestKind::Generation,
br#"{ "model":"gpt-image-future", "prompt":"a lighthouse", "future":9007199254740993 }"#.as_slice(),
),
(
"/v1/images/edits",
ImageRequestKind::Edit,
br#"{"model":"gpt-image-2","images":[{"image_url":"data:image/png;base64,AAEC"}],"prompt":"add fog","prompt":"duplicate stays raw"}"#.as_slice(),
),
];
for (path, _, body) in &cases {
let response = router
.clone()
.oneshot(
Request::post(*path)
.header(AUTHORIZATION, "Bearer sk_images_test")
.header("content-type", "application/json")
.header("x-codex-image-turn-id", "turn_image_route")
.body(Body::from(body.to_vec()))
.expect("image request"),
)
.await
.expect("image response");
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(
response.headers().get("x-image-rate-limit"),
Some(&"42".parse().expect("header value"))
);
let response_body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("read image response");
assert_eq!(response_body.as_ref(), IMAGE_RESPONSE);
}
let captured = execution.captured();
assert_eq!(captured.len(), cases.len());
for (captured, (endpoint, kind, body)) in captured.iter().zip(cases.iter()) {
assert_eq!(captured.provider, "openai");
assert_eq!(captured.endpoint, *endpoint);
assert_eq!(captured.transport, ClientTransport::HttpJson);
assert_eq!(captured.kind, *kind);
assert_eq!(captured.body.as_ref(), *body);
assert_eq!(
captured.context,
json!({"image_turn_id": "turn_image_route"})
);
}
assert_eq!(
execution.committed_statuses(),
vec![StatusCode::CREATED.as_u16(); cases.len()]
);
}
#[tokio::test]
async fn image_route_should_return_the_exact_upstream_error_response() {
let execution = ImageExecution::failing();
let response = api_router(execution)
.await
.oneshot(
Request::post("/v1/images/edits")
.header(AUTHORIZATION, "Bearer sk_images_test")
.header("content-type", "application/json")
.body(Body::from(
br#"{ "model":"gpt-image-2", "future_invalid":9007199254740993 }"#.to_vec(),
))
.expect("image request"),
)
.await
.expect("image response");
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
assert_eq!(
response.headers().get("content-type"),
Some(&"application/problem+json".parse().expect("content type"))
);
assert_eq!(
response.headers().get("x-future-image-error"),
Some(&"preserved".parse().expect("header value"))
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("read image response");
assert_eq!(
body.as_ref(),
br#"{ "error":{"message":"future image validation"}, "future":9007199254740993 }"#
);
}
@@ -1,5 +1,6 @@
mod auth;
mod error;
mod images;
mod models;
mod responses;
mod router;
@@ -7,7 +7,7 @@ use axum::{
use futures::future::BoxFuture;
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, StartExecution,
StartedExecution,
StartProviderExecution, StartedExecution,
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::routing::{
@@ -99,6 +99,18 @@ impl ExecutionService for ModelsExecution {
))
})
}
fn start_provider_endpoint(
&self,
_: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async {
Err(GatewayError::new(
GatewayErrorKind::Internal,
"models test must not start a provider endpoint",
))
})
}
}
fn authorized_request(path: &str) -> Request<Body> {
@@ -17,7 +17,7 @@ use bytes::Bytes;
use futures::future::{BoxFuture, pending};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
StartExecution, StartedExecution,
StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, UpstreamSendState};
use gateway_core::error::{
@@ -190,6 +190,18 @@ impl ExecutionService for ContextCaptureExecution {
))
})
}
fn start_provider_endpoint(
&self,
_: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async {
Err(GatewayError::new(
GatewayErrorKind::Internal,
"responses context test must not start a provider endpoint",
))
})
}
}
impl FakeSession {
@@ -130,7 +130,7 @@ fn decoder_should_preserve_opaque_client_model_values() {
"input": "hello"
}));
assert_eq!(decoded.metadata().public_model(), model);
assert_eq!(decoded.metadata().requested_model(), model);
assert_eq!(
openai_wire_body(&decoded)
.get("model")
@@ -454,14 +454,20 @@ fn decoder_should_not_reject_large_bodies_using_catalog_context_limits() {
let snapshot = super::snapshot("sk_context_test", "openai");
let plan = snapshot
.plan(
&PublicModelId::new(decoded.metadata().public_model()).expect("public model"),
&PublicModelId::new(decoded.metadata().requested_model()).expect("public model"),
decoded.operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("large body must not be locally context-gated");
assert_eq!(plan.candidates()[0].upstream_model().as_str(), "model-a");
assert_eq!(
plan.candidates()[0]
.upstream_model()
.expect("model route candidate")
.as_str(),
"model-a"
);
}
#[test]
@@ -13,7 +13,7 @@ use futures::{SinkExt, StreamExt, future::BoxFuture};
use gateway_api::openai::responses::ResponseCreateFrameError;
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
StartExecution, StartedExecution,
StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, UpstreamSendState};
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind};
@@ -271,6 +271,13 @@ impl ExecutionService for AtomicFailureExecution {
})
})
}
fn start_provider_endpoint(
&self,
_: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async { unreachable!("WebSocket test does not execute provider endpoints") })
}
}
fn atomic_failure_batch() -> CoordinatedEvent {
@@ -455,6 +455,15 @@ where
.unwrap_or_default()
}
/// 返回最终选中 attempt 观察到的上游 HTTP 状态码。
#[must_use]
pub fn response_status_code(&self) -> Option<u16> {
self.current
.as_ref()
.and_then(|current| current.response_observation.as_ref())
.and_then(ProviderResponseObservation::status_code)
}
/// 将已完成响应的账号事实与 Provider 私有状态封装为可丢失的亲和记录。
///
/// Core 不读取 `state` 内容;Provider 将在后续同账号 continuation 时自行解释。
@@ -828,7 +837,7 @@ where
provider_kind: metadata.provider().clone(),
provider_account_id: metadata.provider_account_id().cloned(),
provider_account_ref: metadata.provider_account_id().cloned(),
upstream_model_id: metadata.upstream_model().clone(),
upstream_model_id: metadata.upstream_model().cloned(),
upstream_transport: metadata.transport().as_str().to_owned(),
http_version: None,
};
@@ -1074,7 +1083,7 @@ where
trigger: current.trigger,
provider_kind: current.metadata.provider().clone(),
account_id: current.metadata.provider_account_id().cloned(),
upstream_model_id: current.metadata.upstream_model().clone(),
upstream_model_id: current.metadata.upstream_model().cloned(),
upstream_status_code: current
.response_observation
.as_ref()
@@ -1482,8 +1491,13 @@ where
}
fn image_generation_succeeded(&self) -> Option<bool> {
self.image_generation_requested
.then(|| self.usage.image_output_tokens.unwrap_or_default() > 0)
self.image_generation_requested.then(|| {
if matches!(self.operation, Operation::GenerateImage(_)) {
self.upstream_complete
} else {
self.usage.image_output_tokens.unwrap_or_default() > 0
}
})
}
fn current_send_state(&self) -> UpstreamSendState {
@@ -118,6 +118,35 @@ pub struct StartExecution {
pub metadata: ExecutionRequestMetadata,
}
/// 启动一个由协议 adapter 明确绑定到 Provider 自有端点的请求。
pub struct StartProviderExecution {
pub client: AuthenticatedClient,
pub provider: ProviderKind,
pub operation: Operation,
pub metadata: ExecutionRequestMetadata,
}
enum ExecutionTarget {
Model(PublicModelId),
ProviderEndpoint(ProviderKind),
}
impl ExecutionTarget {
fn into_public_model(self) -> Option<PublicModelId> {
match self {
Self::Model(model) => Some(model),
Self::ProviderEndpoint(_) => None,
}
}
}
struct PendingStartExecution {
client: AuthenticatedClient,
target: ExecutionTarget,
operation: Operation,
metadata: ExecutionRequestMetadata,
}
pub struct StartedExecution {
pub request_id: ModelRequestId,
pub created_at: SystemTime,
@@ -129,6 +158,9 @@ pub trait ExecutionSession: Send {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>>;
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>>;
fn response_headers(&self) -> &[ProviderResponseHeader];
fn response_status_code(&self) -> Option<u16> {
None
}
fn commit_downstream(
&mut self,
client_status_code: Option<u16>,
@@ -156,6 +188,10 @@ pub trait ExecutionService: Send + Sync {
&self,
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>>;
fn start_provider_endpoint(
&self,
request: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>>;
}
/// 成功认证后的 API Key 使用事实接收器。
@@ -247,9 +283,44 @@ impl DefaultExecutionService {
}
}
async fn start_inner(
async fn start_inner(&self, request: StartExecution) -> Result<StartedExecution, GatewayError> {
let StartExecution {
client,
public_model,
operation,
metadata,
} = request;
self.start_inner_with_target(PendingStartExecution {
client,
target: ExecutionTarget::Model(public_model),
operation,
metadata,
})
.await
}
async fn start_provider_endpoint_inner(
&self,
mut request: StartExecution,
request: StartProviderExecution,
) -> Result<StartedExecution, GatewayError> {
let StartProviderExecution {
client,
provider,
operation,
metadata,
} = request;
self.start_inner_with_target(PendingStartExecution {
client,
target: ExecutionTarget::ProviderEndpoint(provider),
operation,
metadata,
})
.await
}
async fn start_inner_with_target(
&self,
mut request: PendingStartExecution,
) -> Result<StartedExecution, GatewayError> {
request.client.policy.authorize().map_err(|_| {
GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled")
@@ -264,16 +335,24 @@ impl DefaultExecutionService {
let routing_context = self
.route_context(request.client.policy.account_scope().provider_kinds())
.await?;
let plan = request
.client
.snapshot
.plan(
&request.public_model,
let account_scope = Arc::clone(request.client.policy.account_scope());
let plan = match &request.target {
ExecutionTarget::ProviderEndpoint(provider) => {
request.client.snapshot.plan_provider_endpoint(
provider,
&request.operation,
account_scope,
&routing_context,
)
}
ExecutionTarget::Model(public_model) => request.client.snapshot.plan(
public_model,
&request.operation,
Arc::clone(request.client.policy.account_scope()),
account_scope,
&routing_context,
)
.map_err(map_routing_error)?;
),
}
.map_err(map_routing_error)?;
let continuation = match request.metadata.previous_response_id.as_ref() {
Some(previous) => {
let resolve = self
@@ -375,18 +454,24 @@ impl DefaultExecutionService {
async fn start_without_continuation(
&self,
request: StartExecution,
request: PendingStartExecution,
request_id: ModelRequestId,
started_at: SystemTime,
deadline_at: SystemTime,
plan: crate::routing::RoutingPlan,
continuation: Option<ContinuationBinding>,
) -> Result<StartedExecution, GatewayError> {
let PendingStartExecution {
client,
target,
operation,
metadata,
} = request;
let admission_request = ClientAdmissionRequest {
model_request_id: request_id.clone(),
client_api_key_id: request.client.policy.key_id().clone(),
client_api_key_id: client.policy.key_id().clone(),
lease_ttl: MODEL_REQUEST_DEADLINE,
limits: request.client.policy.limits(),
limits: client.policy.limits(),
};
match self
.admissions
@@ -411,7 +496,7 @@ impl DefaultExecutionService {
}
let admission = AdmissionLease {
port: Arc::clone(&self.admissions),
client_api_key_id: request.client.policy.key_id().clone(),
client_api_key_id: client.policy.key_id().clone(),
model_request_id: request_id.clone(),
};
let observation = plan
@@ -419,27 +504,27 @@ impl DefaultExecutionService {
.first()
.map_or_else(Default::default, |candidate| {
self.providers
.request_observation(candidate.provider(), &request.operation)
.request_observation(candidate.provider(), &operation)
});
let new_request = NewModelRequest {
id: request_id.clone(),
client_api_key_id: Some(request.client.policy.key_id().clone()),
client_api_key_ref: request.client.policy.key_id().clone(),
client_api_key_id: Some(client.policy.key_id().clone()),
client_api_key_ref: client.policy.key_id().clone(),
config_revision: plan.config_revision(),
routing: request.client.policy.account_scope().routing_snapshot(),
protocol: request.metadata.protocol,
operation: request.operation.kind(),
endpoint: request.metadata.endpoint,
client_transport: request.metadata.transport.as_str().to_owned(),
requested_model: request.public_model,
client_ip: request.metadata.client_ip,
user_agent: request.metadata.user_agent,
routing: client.policy.account_scope().routing_snapshot(),
protocol: metadata.protocol,
operation: operation.kind(),
endpoint: metadata.endpoint,
client_transport: metadata.transport.as_str().to_owned(),
requested_model: target.into_public_model(),
client_ip: metadata.client_ip,
user_agent: metadata.user_agent,
reasoning_effort: observation.reasoning_effort,
reasoning_preset: observation.reasoning_preset,
request_kind: observation.request_kind,
subagent_kind: observation.subagent_kind,
compact: observation.compact,
image_generation_requested: request.operation.image_generation_requested(),
image_generation_requested: operation.image_generation_requested(),
started_at,
deadline_at,
};
@@ -447,7 +532,7 @@ impl DefaultExecutionService {
.coordinator
.start(
new_request,
request.operation,
operation,
plan,
None,
continuation,
@@ -464,7 +549,7 @@ impl DefaultExecutionService {
Ok(StartedExecution {
request_id,
created_at: started_at,
stream: request.metadata.stream,
stream: metadata.stream,
session: Box::new(DefaultExecutionSession::new(
core,
admission,
@@ -572,7 +657,7 @@ impl DefaultExecutionService {
operation: operation.kind(),
endpoint: "/api/admin/accounts/connection-test".to_owned(),
client_transport: ClientTransport::InternalProbe.as_str().to_owned(),
requested_model: public_model,
requested_model: Some(public_model),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -783,6 +868,13 @@ impl ExecutionService for DefaultExecutionService {
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move { self.start_inner(request).await })
}
fn start_provider_endpoint(
&self,
request: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move { self.start_provider_endpoint_inner(request).await })
}
}
impl AccountProbe for DefaultExecutionService {
@@ -948,6 +1040,12 @@ impl ExecutionSession for DefaultExecutionSession {
.unwrap_or_default()
}
fn response_status_code(&self) -> Option<u16> {
self.core
.as_ref()
.and_then(ResponseExecutionSession::response_status_code)
}
fn commit_downstream(
&mut self,
client_status_code: Option<u16>,
@@ -1076,6 +1174,7 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
match error {
crate::error::RoutingError::NoCapableProvider { .. }
| crate::error::RoutingError::NoCapableProviderEndpoint { .. }
| crate::error::RoutingError::EmptyAccountScope => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider can execute this request",
@@ -560,7 +560,7 @@ pub struct NewModelRequest {
pub operation: OperationKind,
pub endpoint: String,
pub client_transport: String,
pub requested_model: PublicModelId,
pub requested_model: Option<PublicModelId>,
pub client_ip: Option<IpAddr>,
pub user_agent: Option<String>,
pub reasoning_effort: Option<String>,
@@ -582,7 +582,7 @@ pub struct AttemptRecord {
pub provider_kind: ProviderKind,
pub provider_account_id: Option<ProviderAccountId>,
pub provider_account_ref: Option<ProviderAccountId>,
pub upstream_model_id: UpstreamModelId,
pub upstream_model_id: Option<UpstreamModelId>,
pub upstream_transport: String,
pub http_version: Option<String>,
}
@@ -595,7 +595,7 @@ pub struct IntermediateFailure {
pub trigger: AttemptTrigger,
pub provider_kind: ProviderKind,
pub account_id: Option<ProviderAccountId>,
pub upstream_model_id: UpstreamModelId,
pub upstream_model_id: Option<UpstreamModelId>,
pub upstream_status_code: Option<u16>,
pub upstream_request_id: Option<String>,
pub error: ProviderError,
@@ -123,7 +123,7 @@ impl UpstreamTransport {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCallMetadata {
provider: ProviderKind,
upstream_model: UpstreamModelId,
upstream_model: Option<UpstreamModelId>,
resource: ProviderResource,
upstream_request_id: Option<OpaqueUpstreamValue>,
transport: UpstreamTransport,
@@ -140,7 +140,23 @@ impl ProviderCallMetadata {
) -> Self {
Self {
provider,
upstream_model,
upstream_model: Some(upstream_model),
resource,
upstream_request_id: None,
transport,
}
}
/// 创建一次不声明模型的 Provider 原生端点调用事实。
#[must_use]
pub const fn for_provider_endpoint(
provider: ProviderKind,
resource: ProviderResource,
transport: UpstreamTransport,
) -> Self {
Self {
provider,
upstream_model: None,
resource,
upstream_request_id: None,
transport,
@@ -160,10 +176,10 @@ impl ProviderCallMetadata {
&self.provider
}
/// 返回实际模型;必须与冻结 target 一致。
/// 返回实际模型;Provider 原生端点没有模型事实。
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
pub const fn upstream_model(&self) -> Option<&UpstreamModelId> {
self.upstream_model.as_ref()
}
/// 返回匿名资源;账号调用返回 `None`。
@@ -199,7 +215,8 @@ impl ProviderCallMetadata {
/// 确认 metadata 没有替换请求计划中冻结的 Provider 候选。
#[must_use]
pub fn confirms(&self, candidate: &ProviderCandidate) -> bool {
candidate.provider() == &self.provider && candidate.upstream_model() == &self.upstream_model
candidate.provider() == &self.provider
&& candidate.upstream_model() == self.upstream_model.as_ref()
}
}
@@ -517,7 +534,7 @@ impl fmt::Debug for ProviderRequest {
.debug_struct("ProviderRequest")
.field("operation", &self.operation)
.field("provider", self.candidate.provider())
.field("upstream_model", self.candidate.upstream_model())
.field("upstream_model", &self.candidate.upstream_model())
.finish()
}
}
+6
View File
@@ -98,6 +98,12 @@ pub enum RoutingError {
/// 客户端提交的模型名称。
model: String,
},
/// 固定 Provider 的原生端点当前不可执行。
#[error("provider endpoint `{provider}` is unavailable")]
NoCapableProviderEndpoint {
/// adapter 已绑定的 Provider。
provider: String,
},
}
/// 调用方策略不满足约束或拒绝请求。
+55 -4
View File
@@ -16,7 +16,7 @@ use crate::operation::ProviderSessionState;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResponseMeta {
response_id: String,
model: String,
model: Option<String>,
finish_reason: Option<FinishReason>,
}
@@ -26,7 +26,17 @@ impl ResponseMeta {
pub fn new(response_id: impl Into<String>, model: impl Into<String>) -> Self {
Self {
response_id: response_id.into(),
model: model.into(),
model: Some(model.into()),
finish_reason: None,
}
}
/// 创建不声明模型的 Provider 原生端点响应元数据。
#[must_use]
pub fn for_provider_endpoint(response_id: impl Into<String>) -> Self {
Self {
response_id: response_id.into(),
model: None,
finish_reason: None,
}
}
@@ -46,8 +56,8 @@ impl ResponseMeta {
/// 返回对客户端公开的模型名。
#[must_use]
pub fn model(&self) -> &str {
&self.model
pub fn model(&self) -> Option<&str> {
self.model.as_deref()
}
/// 返回规范化终止原因。
@@ -190,6 +200,7 @@ pub struct ProtocolWireEvent {
data: Value,
has_json_data: bool,
raw_sse_frame: Option<Bytes>,
raw_json_body: Option<Bytes>,
sse_id: Option<String>,
sse_retry: Option<u64>,
}
@@ -228,6 +239,7 @@ impl ProtocolWireEvent {
data,
has_json_data: true,
raw_sse_frame: None,
raw_json_body: None,
sse_id,
sse_retry,
})
@@ -274,6 +286,32 @@ impl ProtocolWireEvent {
data: Value::Null,
has_json_data: false,
raw_sse_frame: Some(raw_sse_frame),
raw_json_body: None,
sse_id: None,
sse_retry: None,
})
}
/// 创建未经改写的完整 JSON 响应正文。
///
/// 该载荷用于非流式协议端点;Core 不解析或重编码其中的值。
///
/// # Errors
///
/// 协议名不满足内部路由标识约束时返回错误。
pub fn raw_json(
protocol: impl Into<String>,
raw_json_body: Bytes,
) -> Result<Self, IdentifierError> {
let protocol = protocol.into();
validate_text(&protocol, 64, true, None)?;
Ok(Self {
protocol,
event_type: None,
data: Value::Null,
has_json_data: false,
raw_sse_frame: None,
raw_json_body: Some(raw_json_body),
sse_id: None,
sse_retry: None,
})
@@ -321,6 +359,18 @@ impl ProtocolWireEvent {
self.raw_sse_frame.as_ref()
}
/// 返回可直接写给 JSON 客户端的未经改写响应正文。
#[must_use]
pub const fn raw_json_body(&self) -> Option<&Bytes> {
self.raw_json_body.as_ref()
}
/// 拆出未经改写的完整 JSON 响应正文。
#[must_use]
pub fn into_raw_json_body(self) -> Option<Bytes> {
self.raw_json_body
}
/// 拆出协议原生 JSON 数据。
#[must_use]
pub fn into_data(self) -> Value {
@@ -336,6 +386,7 @@ impl fmt::Debug for ProtocolWireEvent {
.field("event_type", &self.event_type)
.field("has_json_data", &self.has_json_data)
.field("has_raw_sse_frame", &self.raw_sse_frame.is_some())
.field("has_raw_json_body", &self.raw_json_body.is_some())
.field("has_sse_id", &self.sse_id.is_some())
.field("sse_retry", &self.sse_retry)
.field("data", &"<not included in Debug>")
@@ -1,13 +1,14 @@
//! 协议无关的业务 operation。
//!
//! 这里只保留网关需要解释、路由或结算的稳定语义。协议 adapter 无法写入
//! wire body 的连接级事实随 [`ProtocolPayload`] 不透明传递,由对应 Provider
//! 自己解释。
//! wire body 与连接级事实随 [`ProtocolPayload`] 或 [`RawJsonPayload`] 不透明传递,
//! 由对应 Provider 自己解释。
use std::collections::BTreeSet;
use std::fmt;
use std::sync::Arc;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
@@ -239,6 +240,72 @@ impl fmt::Debug for ProtocolPayload {
}
}
/// 客户端协议交给 Provider 的未经重编码 JSON 正文。
///
/// 协议 adapter 只确定路由端点;Core 保留原始字节与非 wire 上下文,
/// 不验证、解析或改写正文。
#[derive(Clone, PartialEq, Eq)]
pub struct RawJsonPayload {
protocol: String,
body: Bytes,
context: Map<String, Value>,
}
impl RawJsonPayload {
/// 创建未经重编码的协议 JSON 正文。
///
/// # Errors
///
/// 协议名称为空、过长或含控制字符时返回错误。
pub fn new(protocol: impl Into<String>, body: Bytes) -> Result<Self, OperationError> {
let protocol = protocol.into();
validate_text(&protocol, 64, true, None).map_err(|_| OperationError::EmptyField {
field: "raw_json_payload protocol",
})?;
Ok(Self {
protocol,
body,
context: Map::new(),
})
}
/// 返回协议名称。
#[must_use]
pub fn protocol(&self) -> &str {
&self.protocol
}
/// 返回 adapter 收到的原始 JSON 字节。
#[must_use]
pub const fn body(&self) -> &Bytes {
&self.body
}
/// 附着不写入 wire body 的协议连接上下文。
#[must_use]
pub fn with_context(mut self, context: Map<String, Value>) -> Self {
self.context = context;
self
}
/// 返回仅供对应 Provider 解释的非 wire 上下文。
#[must_use]
pub const fn context(&self) -> &Map<String, Value> {
&self.context
}
}
impl fmt::Debug for RawJsonPayload {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("RawJsonPayload")
.field("protocol", &self.protocol)
.field("body", &"<not included in Debug>")
.field("context", &"<not included in Debug>")
.finish()
}
}
/// 通用生成请求。
#[derive(Clone, PartialEq)]
pub struct GenerateRequest {
@@ -501,32 +568,39 @@ impl fmt::Debug for RerankRequest {
}
}
/// 图像生成请求。
#[derive(Clone, PartialEq, Eq)]
/// 图像 API 的稳定端点语义。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ImageRequestKind {
/// 从文本生成图像。
Generation,
/// 使用输入图像执行编辑。
Edit,
}
/// 图像生成或编辑请求。
#[derive(Clone, PartialEq)]
pub struct ImageRequest {
prompt: String,
count: u32,
kind: ImageRequestKind,
payload: RawJsonPayload,
}
impl ImageRequest {
/// 创建图像生成请求。
///
/// # Errors
///
/// Prompt 为空时返回错误。
pub fn new(prompt: impl Into<String>) -> Result<Self, OperationError> {
let prompt = prompt.into();
if prompt.is_empty() {
return Err(OperationError::EmptyField { field: "prompt" });
}
Ok(Self { prompt, count: 1 })
/// 创建携带原始协议 JSON 正文的图像请求。
#[must_use]
pub const fn from_raw_json(kind: ImageRequestKind, payload: RawJsonPayload) -> Self {
Self { kind, payload }
}
/// 设置图像数量。
/// 返回图像 API 端点语义。
#[must_use]
pub const fn with_count(mut self, count: u32) -> Self {
self.count = count;
self
pub const fn kind(&self) -> ImageRequestKind {
self.kind
}
/// 返回未经重编码的协议 JSON 正文。
#[must_use]
pub const fn payload(&self) -> &RawJsonPayload {
&self.payload
}
}
@@ -534,8 +608,8 @@ impl fmt::Debug for ImageRequest {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ImageRequest")
.field("prompt", &"<not included in Debug>")
.field("count", &self.count)
.field("kind", &self.kind)
.field("payload", &"<not included in Debug>")
.finish()
}
}
@@ -649,7 +723,8 @@ impl Operation {
pub fn image_generation_requested(&self) -> bool {
match self {
Self::Generate(request) => request.image_generation_requested(),
Self::Embed(_) | Self::Rerank(_) | Self::GenerateImage(_) | Self::Speech(_) => false,
Self::GenerateImage(_) => true,
Self::Embed(_) | Self::Rerank(_) | Self::Speech(_) => false,
}
}
@@ -231,6 +231,32 @@ pub trait ProviderSessionAffinityPort: Send + Sync {
ttl: Duration,
) -> BoxFuture<'a, Result<(), ProviderStoreError>>;
/// 仅在亲和键尚未绑定时写入候选账号,并返回原子操作后的实际绑定。
///
/// 已存在的绑定绝不会被候选账号覆盖;同一根会话的并发首次请求据此收敛到
/// 单一账号。TTL 只在首次写入时设置,已有绑定由成功反馈负责刷新。
fn claim_or_load<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
candidate_account_id: &'a ProviderAccountId,
ttl: Duration,
) -> BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>>;
/// 仅当当前绑定等于 `expected_account_id`(或键已过期)时写入新账号,
/// 并返回原子操作后的实际绑定。
///
/// Provider 用它迁移不可用账号,以及在成功后以 `expected == replacement`
/// 刷新 TTL;迟到的旧账号成功不能覆盖较新的会话 winner。
fn compare_and_bind<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected_account_id: &'a ProviderAccountId,
replacement_account_id: &'a ProviderAccountId,
ttl: Duration,
) -> BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>>;
fn clear<'a>(
&'a self,
provider_kind: &'a ProviderKind,
+4 -10
View File
@@ -780,11 +780,11 @@ pub struct RoutingContext {
pub blocked_providers: BTreeSet<ProviderKind>,
}
/// 已绑定 Provider 与真实上游模型的请求候选。
/// 已绑定 Provider 的请求候选;模型端点携带真实上游模型,原生端点不虚构模型。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCandidate {
provider: ProviderKind,
upstream_model: UpstreamModelId,
upstream_model: Option<UpstreamModelId>,
emulated_features: BTreeSet<Feature>,
account_scope: Arc<FrozenAccountScope>,
}
@@ -796,8 +796,8 @@ impl ProviderCandidate {
}
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
pub const fn upstream_model(&self) -> Option<&UpstreamModelId> {
self.upstream_model.as_ref()
}
#[must_use]
@@ -816,7 +816,6 @@ impl ProviderCandidate {
pub struct RoutingPlan {
config_revision: ConfigRevision,
account_selection_policy: AccountSelectionPolicy,
public_model: PublicModelId,
operation: OperationKind,
max_attempts: NonZeroU32,
account_scope: Arc<FrozenAccountScope>,
@@ -834,11 +833,6 @@ impl RoutingPlan {
self.account_selection_policy
}
#[must_use]
pub const fn public_model(&self) -> &PublicModelId {
&self.public_model
}
#[must_use]
pub const fn operation(&self) -> OperationKind {
self.operation
@@ -765,7 +765,7 @@ impl RuntimeSnapshot {
};
candidates.push(ProviderCandidate {
provider: provider.clone(),
upstream_model,
upstream_model: Some(upstream_model),
emulated_features,
account_scope: Arc::clone(&account_scope),
});
@@ -780,7 +780,6 @@ impl RuntimeSnapshot {
Ok(RoutingPlan {
config_revision: self.revision,
account_selection_policy: self.account_selection_policy,
public_model: public_model.clone(),
operation: operation.kind(),
max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS)
.expect("constant request attempt limit is non-zero"),
@@ -788,6 +787,47 @@ impl RuntimeSnapshot {
candidates: Arc::from(candidates),
})
}
/// 为 Provider 自有、且不属于文本模型目录的端点冻结请求计划。
///
/// 端点 adapter 已经确定 Provider,因此这里只执行账号范围、circuit 和注册
/// 状态检查;不会读取业务正文、构造模型或查询 Provider 文本模型目录。
pub fn plan_provider_endpoint(
&self,
provider: &ProviderKind,
operation: &Operation,
account_scope: Arc<FrozenAccountScope>,
context: &RoutingContext,
) -> Result<RoutingPlan, RoutingError> {
let allowed = !account_scope.provider_kinds().is_empty()
&& account_scope.provider_kinds().contains(provider)
&& self.providers.contains(provider)
&& context
.required_provider
.as_ref()
.is_none_or(|required| required == provider)
&& !context.blocked_providers.contains(provider);
if !allowed {
return Err(RoutingError::NoCapableProviderEndpoint {
provider: provider.as_str().to_owned(),
});
}
let candidate = ProviderCandidate {
provider: provider.clone(),
upstream_model: None,
emulated_features: BTreeSet::new(),
account_scope: Arc::clone(&account_scope),
};
Ok(RoutingPlan {
config_revision: self.revision,
account_selection_policy: self.account_selection_policy,
operation: operation.kind(),
max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS)
.expect("constant request attempt limit is non-zero"),
account_scope,
candidates: Arc::from([candidate]),
})
}
}
/// 请求级冻结失败;此状态必须 fail closed。
@@ -317,7 +317,10 @@ impl Provider for ScriptedProvider {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().clone(),
candidate
.upstream_model()
.cloned()
.expect("model route candidate"),
ProviderResource::Account {
id: ProviderAccountId::new(account_id).expect("account id"),
revision: CredentialRevision::new(1).expect("revision"),
@@ -337,7 +340,10 @@ impl Provider for ScriptedProvider {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().clone(),
candidate
.upstream_model()
.cloned()
.expect("model route candidate"),
ProviderResource::Account {
id: ProviderAccountId::new(account_id).expect("account id"),
revision: CredentialRevision::new(1).expect("revision"),
@@ -357,7 +363,10 @@ impl Provider for ScriptedProvider {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().clone(),
candidate
.upstream_model()
.cloned()
.expect("model route candidate"),
ProviderResource::Account {
id: ProviderAccountId::new(account_id).expect("account id"),
revision: CredentialRevision::new(1).expect("revision"),
@@ -623,7 +632,7 @@ fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest
operation: operation.kind(),
endpoint: "responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model: PublicModelId::new("gpt-5").expect("model"),
requested_model: Some(PublicModelId::new("gpt-5").expect("model")),
client_ip: Some("127.0.0.1".parse().expect("client IP")),
user_agent: Some("gateway-core-test".to_owned()),
reasoning_effort: Some("medium".to_owned()),
@@ -18,17 +18,17 @@ use gateway_core::engine::continuation::{
NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
AccountSelectionPolicy, CredentialRevision, ProviderAccountId, RotationStrategy,
};
use gateway_core::engine::execution::{
ClientApiKeyUsageSink, ClientTransport, DefaultExecutionService, ExecutionRequestMetadata,
ExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
StartExecution, provider_failure_affects_circuit,
StartExecution, StartProviderExecution, provider_failure_affects_circuit,
};
use gateway_core::engine::probe::{AccountProbe, AccountProbeRequest};
use gateway_core::engine::provider::{
Provider, ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRegistry,
ProviderRequest, ProviderStream,
Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities,
ProviderRegistry, ProviderRequest, ProviderResource, ProviderStream, UpstreamTransport,
};
use gateway_core::engine::{
AttemptContext, AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
@@ -38,7 +38,10 @@ use gateway_core::error::{
ClientVisibleUpstreamResponse, GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError,
StoreErrorKind,
};
use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload};
use gateway_core::operation::{
GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProtocolPayload,
RawJsonPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
@@ -193,6 +196,50 @@ impl Provider for FailingProvider {
}
}
struct ColdFailingProvider;
#[async_trait]
impl Provider for ColdFailingProvider {
fn name(&self) -> &'static str {
"openai"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::default()
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
async fn execute(
&self,
request: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
let metadata = ProviderCallMetadata::for_provider_endpoint(
request.candidate().provider().clone(),
ProviderResource::Account {
id: ProviderAccountId::new("acct_usage").expect("account"),
revision: CredentialRevision::new(1).expect("revision"),
},
UpstreamTransport::new("http_json").expect("transport"),
);
Ok(ProviderStream::new(
metadata,
futures::stream::once(async {
Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
))
}),
(),
))
}
}
#[test]
fn successful_authentication_should_record_client_key_usage() {
let usage = Arc::new(RecordingClientApiKeyUsage::default());
@@ -213,6 +260,60 @@ fn successful_authentication_should_record_client_key_usage() {
assert_eq!(usage.recorded(), vec!["key_usage_test".to_owned()]);
}
#[test]
fn provider_endpoint_should_persist_its_real_v1_endpoint() {
let store = Arc::new(TrackingExecutionStore::default());
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(client_snapshot()),
store.clone(),
ProviderRegistry::new([Arc::new(ColdFailingProvider) as Arc<dyn Provider>])
.expect("provider registry"),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_usage_test")
.expect("authenticated client");
let operation = Operation::GenerateImage(ImageRequest::from_raw_json(
ImageRequestKind::Generation,
RawJsonPayload::new(
"openai",
Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"hello"}"#),
)
.expect("image payload"),
));
let mut started = block_on(service.start_provider_endpoint(StartProviderExecution {
client,
provider: ProviderKind::new("openai").expect("provider"),
operation,
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/images/generations".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("provider endpoint request should start without a text catalog entry");
let error = block_on(started.session.collect_uncommitted())
.expect_err("the cold provider stops execution after persistence");
assert_eq!(
store.model_request_endpoints(),
vec!["/v1/images/generations".to_owned()],
"execution error: {error:?}"
);
assert_eq!(store.requested_models(), vec![None]);
let upstream_models = store.upstream_models();
assert!(!upstream_models.is_empty());
assert!(upstream_models.iter().all(Option::is_none));
}
#[test]
fn circuit_store_failure_should_fail_open_during_request_start() {
let service = DefaultExecutionService::new(
@@ -412,6 +513,9 @@ impl ClientApiKeyUsageSink for RecordingClientApiKeyUsage {
struct TrackingExecutionStore {
touched: AtomicBool,
probe_failures: Mutex<Vec<String>>,
model_request_endpoints: Mutex<Vec<String>>,
requested_models: Mutex<Vec<Option<String>>>,
upstream_models: Mutex<Vec<Option<String>>>,
fail_probe_observation: AtomicBool,
}
@@ -426,17 +530,60 @@ impl TrackingExecutionStore {
.expect("probe failures lock")
.clone()
}
fn model_request_endpoints(&self) -> Vec<String> {
self.model_request_endpoints
.lock()
.expect("model request endpoints lock")
.clone()
}
fn requested_models(&self) -> Vec<Option<String>> {
self.requested_models
.lock()
.expect("requested models lock")
.clone()
}
fn upstream_models(&self) -> Vec<Option<String>> {
self.upstream_models
.lock()
.expect("upstream models lock")
.clone()
}
}
#[async_trait]
impl ExecutionStore for TrackingExecutionStore {
async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> {
async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> {
self.touch();
self.requested_models
.lock()
.expect("requested models lock")
.push(
request
.requested_model
.as_ref()
.map(|model| model.as_str().to_owned()),
);
self.model_request_endpoints
.lock()
.expect("model request endpoints lock")
.push(request.endpoint);
Ok(())
}
async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> {
async fn record_attempt(&self, attempt: AttemptRecord) -> Result<(), StoreError> {
self.touch();
self.upstream_models
.lock()
.expect("upstream models lock")
.push(
attempt
.upstream_model_id
.as_ref()
.map(|model| model.as_str().to_owned()),
);
Ok(())
}
@@ -77,6 +77,17 @@ fn protocol_wire_event_should_preserve_opaque_sse_metadata() {
assert_eq!(wire.sse_id(), Some(sse_id.as_str()));
}
#[test]
fn protocol_wire_event_should_preserve_raw_json_body_without_parsing_it() {
let raw = Bytes::from_static(
br#"{ "created": 1, "data": [{"b64_json":"AAAA"}], "future": 9007199254740993 }"#,
);
let wire = ProtocolWireEvent::raw_json("openai", raw.clone()).expect("raw JSON event");
assert_eq!(wire.raw_json_body(), Some(&raw));
assert_eq!(wire.into_raw_json_body(), Some(raw));
}
#[test]
fn validator_should_accept_started_content_delta_completed_sequence() {
let mut validator = EventSequenceValidator::new();
@@ -1,7 +1,9 @@
use bytes::Bytes;
use gateway_core::error::OperationError;
use gateway_core::operation::{
EmbedRequest, Feature, GenerateRequest, ImageRequest, Operation, OperationKind,
ProtocolPayload, ProviderSessionState, RerankRequest, RetrySafety, SpeechRequest,
EmbedRequest, Feature, GenerateRequest, ImageRequest, ImageRequestKind, Operation,
OperationKind, ProtocolPayload, ProviderSessionState, RawJsonPayload, RerankRequest,
RetrySafety, SpeechRequest,
};
use serde_json::{Map, Value, json};
@@ -12,6 +14,17 @@ fn generate(body: Value) -> GenerateRequest {
)
}
fn image(kind: ImageRequestKind, body: Value) -> ImageRequest {
ImageRequest::from_raw_json(
kind,
RawJsonPayload::new(
"openai",
Bytes::from(serde_json::to_vec(&body).expect("image JSON")),
)
.expect("OpenAI payload"),
)
}
#[test]
fn generate_request_should_keep_client_body_opaque_and_redacted() {
let secret = "private prompt body";
@@ -113,10 +126,6 @@ fn non_generate_requests_should_reject_missing_required_fields() {
RerankRequest::new("", vec!["document".to_owned()]).expect_err("query is required"),
OperationError::EmptyField { field: "query" }
);
assert_eq!(
ImageRequest::new("").expect_err("prompt is required"),
OperationError::EmptyField { field: "prompt" }
);
assert_eq!(
SpeechRequest::new("hello", "").expect_err("voice is required"),
OperationError::EmptyField { field: "voice" }
@@ -132,7 +141,10 @@ fn operation_kind_and_retry_safety_should_remain_stable() {
let rerank = Operation::Rerank(
RerankRequest::new("query", vec!["document".to_owned()]).expect("rerank request is valid"),
);
let image = Operation::GenerateImage(ImageRequest::new("draw").expect("image request"));
let image = Operation::GenerateImage(image(
ImageRequestKind::Generation,
json!({"model": "gpt-image-2", "prompt": "draw"}),
));
let speech = Operation::Speech(SpeechRequest::new("hello", "alloy").expect("speech request"));
assert_eq!(generate.kind(), OperationKind::Generate);
@@ -140,5 +152,35 @@ fn operation_kind_and_retry_safety_should_remain_stable() {
assert_eq!(embed.retry_safety(), RetrySafety::Idempotent);
assert_eq!(rerank.retry_safety(), RetrySafety::Idempotent);
assert_eq!(image.kind(), OperationKind::GenerateImage);
assert!(image.image_generation_requested());
assert_eq!(speech.kind(), OperationKind::Speech);
}
#[test]
fn image_request_should_preserve_generation_and_edit_payloads_opaque() {
let generation_body = json!({
"model": "gpt-image-2",
"prompt": "draw",
"future_official_field": {"keep": true},
});
let edit_body = json!({
"model": "gpt-image-2",
"images": [{"image_url": "data:image/png;base64,AAAA"}],
"prompt": "edit",
"future_official_field": [1, 2, 3],
});
let generation = image(ImageRequestKind::Generation, generation_body.clone());
let edit = image(ImageRequestKind::Edit, edit_body.clone());
assert_eq!(generation.kind(), ImageRequestKind::Generation);
assert_eq!(edit.kind(), ImageRequestKind::Edit);
assert_eq!(
serde_json::from_slice::<Value>(generation.payload().body()).expect("generation JSON"),
generation_body
);
assert_eq!(
serde_json::from_slice::<Value>(edit.payload().body()).expect("edit JSON"),
edit_body
);
assert!(!format!("{edit:?}").contains("data:image/png"));
}
@@ -3,11 +3,13 @@ use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::operation::{
CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload,
CapabilityRequirements, Feature, GenerateRequest, ImageRequest, ImageRequestKind, Operation,
OperationKind, ProtocolPayload, RawJsonPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
@@ -50,6 +52,21 @@ fn operation() -> Operation {
))
}
fn image_operation() -> Operation {
let body = serde_json::json!({
"model": "gpt-image-2",
"prompt": "draw",
});
Operation::GenerateImage(ImageRequest::from_raw_json(
ImageRequestKind::Generation,
RawJsonPayload::new(
"openai",
Bytes::from(serde_json::to_vec(&body).expect("image JSON")),
)
.expect("OpenAI payload"),
))
}
fn account_directory() -> Arc<RuntimeAccountDirectory> {
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
@@ -267,7 +284,13 @@ fn selected_provider_should_use_global_model_mapping() {
assert_eq!(plan.candidates().len(), 1);
assert_eq!(plan.candidates()[0].provider().as_str(), "openai");
assert_eq!(plan.candidates()[0].upstream_model().as_str(), "gpt-5.5");
assert_eq!(
plan.candidates()[0]
.upstream_model()
.expect("model route candidate")
.as_str(),
"gpt-5.5"
);
}
#[test]
@@ -311,6 +334,51 @@ fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
));
}
#[test]
fn provider_endpoint_plan_should_not_consult_or_publish_the_text_model_catalog() {
let snapshot = snapshot();
let provider = ProviderKind::new("openai").expect("provider");
let public_model = PublicModelId::new("gpt-image-2").expect("model");
assert!(!snapshot.contains_public_model_for_provider(&public_model, &provider));
let plan = snapshot
.plan_provider_endpoint(
&provider,
&image_operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("provider-owned endpoint remains routable");
assert_eq!(plan.operation(), OperationKind::GenerateImage);
assert_eq!(plan.candidates().len(), 1);
assert_eq!(plan.candidates()[0].provider(), &provider);
assert_eq!(plan.candidates()[0].upstream_model(), None);
assert!(!snapshot.contains_public_model_for_provider(&public_model, &provider));
}
#[test]
fn provider_endpoint_plan_should_still_respect_circuit_filtering() {
let snapshot = snapshot();
let provider = ProviderKind::new("openai").expect("provider");
let error = snapshot
.plan_provider_endpoint(
&provider,
&image_operation(),
snapshot.all_account_scope(),
&RoutingContext {
blocked_providers: BTreeSet::from([provider.clone()]),
..RoutingContext::default()
},
)
.expect_err("blocked provider has no endpoint candidate");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProviderEndpoint { .. }
));
}
#[test]
fn unmapped_model_should_pass_through_unchanged() {
let snapshot = RuntimeSnapshot::new(
@@ -340,7 +408,13 @@ fn unmapped_model_should_pass_through_unchanged() {
)
.expect("unknown model remains transparent");
assert_eq!(plan.candidates()[0].upstream_model().as_str(), requested);
assert_eq!(
plan.candidates()[0]
.upstream_model()
.expect("model route candidate")
.as_str(),
requested
);
}
}
+1 -1
View File
@@ -1,6 +1,6 @@
//! 多 Provider 网关的 PostgreSQL 持久化与 Redis 协调 adapter。
//!
//! 业务规则与 port 由 `gateway-core` 拥有。本 crate 只负责把终态的七张业务表
//! 业务规则与 port 由 `gateway-core` / `gateway-admin` 拥有。本 crate 只负责把当前十二张业务表
//! 和可丢失 Redis 状态映射为明确的基础设施操作。
use std::sync::Arc;
@@ -91,7 +91,7 @@ pub struct NewModelRequest {
pub operation: String,
pub endpoint: String,
pub client_transport: String,
pub requested_model_id: String,
pub requested_model_id: Option<String>,
pub client_ip: Option<String>,
pub user_agent: Option<String>,
pub reasoning_effort: Option<String>,
@@ -112,7 +112,9 @@ impl NewModelRequest {
require_nonempty(ENTITY, "operation", &self.operation)?;
require_nonempty(ENTITY, "endpoint", &self.endpoint)?;
require_nonempty(ENTITY, "client_transport", &self.client_transport)?;
require_nonempty(ENTITY, "requested_model_id", &self.requested_model_id)?;
if let Some(requested_model_id) = self.requested_model_id.as_deref() {
require_nonempty(ENTITY, "requested_model_id", requested_model_id)?;
}
validate_routing_snapshot(
&self.routing_scope,
&self.routing_group_refs,
@@ -141,7 +143,7 @@ pub struct ModelRequestAttemptStart {
pub provider_kind: String,
pub provider_account_id: Option<String>,
pub provider_account_ref: Option<String>,
pub upstream_model_id: String,
pub upstream_model_id: Option<String>,
pub upstream_transport: String,
pub http_version: Option<String>,
}
@@ -151,11 +153,13 @@ impl ModelRequestAttemptStart {
for (field, value) in [
("model_request_id", self.model_request_id.as_str()),
("provider_kind", self.provider_kind.as_str()),
("upstream_model_id", self.upstream_model_id.as_str()),
("upstream_transport", self.upstream_transport.as_str()),
] {
require_nonempty(ENTITY, field, value)?;
}
if let Some(upstream_model_id) = self.upstream_model_id.as_deref() {
require_nonempty(ENTITY, "upstream_model_id", upstream_model_id)?;
}
if self.attempt_count == 0 {
return Err(invalid("attempt_count must be positive"));
}
@@ -841,7 +845,10 @@ impl ExecutionStore for PgExecutionStore {
provider_kind: Some(failure.provider_kind.as_str().to_owned()),
provider_account_id: failure.account_id.as_ref().map(|id| id.as_str().to_owned()),
provider_account_ref: failure.account_id.as_ref().map(|id| id.as_str().to_owned()),
upstream_model_id: Some(failure.upstream_model_id.as_str().to_owned()),
upstream_model_id: failure
.upstream_model_id
.as_ref()
.map(|model| model.as_str().to_owned()),
failure_kind: error.kind().as_str().to_owned(),
status_code: error.upstream_status().or(failure.upstream_status_code),
provider_error_code: error.upstream_code().map(|code| code.as_str().to_owned()),
@@ -1081,7 +1088,9 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
operation: request.operation.as_str().to_owned(),
endpoint: request.endpoint,
client_transport: request.client_transport,
requested_model_id: request.requested_model.as_str().to_owned(),
requested_model_id: request
.requested_model
.map(|model| model.as_str().to_owned()),
client_ip: request.client_ip.map(|address| address.to_string()),
user_agent: request.user_agent,
reasoning_effort: request.reasoning_effort,
@@ -1150,7 +1159,9 @@ fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart {
.provider_account_ref
.as_ref()
.map(|id| id.as_str().to_owned()),
upstream_model_id: attempt.upstream_model_id.as_str().to_owned(),
upstream_model_id: attempt
.upstream_model_id
.map(|model| model.as_str().to_owned()),
upstream_transport: attempt.upstream_transport,
http_version: attempt.http_version,
}
@@ -551,7 +551,10 @@ impl ExecutionObservationWrite {
Some(failure.request_id.as_str()),
Some(failure.provider_kind.as_str()),
failure.account_id.as_ref().map(|value| value.as_str()),
Some(failure.upstream_model_id.as_str()),
failure
.upstream_model_id
.as_ref()
.map(|model| model.as_str()),
failure.upstream_request_id.as_deref(),
])
.saturating_add(provider_error_bytes(&failure.error)),
@@ -638,7 +641,7 @@ fn new_request_bytes(request: &NewModelRequest) -> usize {
Some(request.protocol.as_str()),
Some(request.endpoint.as_str()),
Some(request.client_transport.as_str()),
Some(request.requested_model.as_str()),
request.requested_model.as_ref().map(|model| model.as_str()),
request.user_agent.as_deref(),
request.reasoning_effort.as_deref(),
request.reasoning_preset.as_deref(),
@@ -659,7 +662,10 @@ fn attempt_bytes(attempt: &AttemptRecord) -> usize {
.provider_account_ref
.as_ref()
.map(|value| value.as_str()),
Some(attempt.upstream_model_id.as_str()),
attempt
.upstream_model_id
.as_ref()
.map(|model| model.as_str()),
Some(attempt.upstream_transport.as_str()),
attempt.http_version.as_deref(),
]))
@@ -1,4 +1,4 @@
//! 九张终态业务表的 PostgreSQL adapters。
//! 当前十二张业务表的 PostgreSQL adapters。
use async_trait::async_trait;
use sqlx::{PgPool, postgres::PgPoolOptions};
@@ -504,7 +504,7 @@ pub struct UsageRecord {
pub operation: String,
pub endpoint: String,
pub client_transport: String,
pub requested_model_id: String,
pub requested_model_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub provider_account_name: Option<String>,
@@ -259,6 +259,7 @@ pub(crate) async fn account_model_rows(
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id)
order by mr.provider_account_ref, request_count desc, model",
)
@@ -286,6 +287,7 @@ pub(crate) async fn account_model_costs(
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and mr.cost_amount is not null and mr.cost_currency is not null
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id),
mr.cost_currency
order by mr.provider_account_ref, model, mr.cost_currency",
@@ -682,8 +682,10 @@ impl AccountStore for PgAdminAccountStore {
) -> AdminStoreResult<CredentialMutationResult> {
match command.credential {
AuthorizationCredentialCommit::Create(credential) => {
let account_id = credential.account_id.clone();
let result = self
let CredentialImportResult {
config_revision,
credential_ids,
} = self
.commit_prepared_import(
PreparedCredentialImport {
provider_kind: credential.provider_kind.clone(),
@@ -693,6 +695,14 @@ impl AccountStore for PgAdminAccountStore {
"authorize",
)
.await?;
let [account_id]: [CoreProviderAccountId; 1] =
credential_ids.try_into().map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
ENTITY,
"authorization import returned an unexpected account count",
)
})?;
let details = self
.accounts
.load_provider_account(account_id.as_str())
@@ -706,7 +716,7 @@ impl AccountStore for PgAdminAccountStore {
)
})?;
Ok(CredentialMutationResult {
config_revision: result.config_revision,
config_revision,
account_id,
credential_revision: Some(admin_revision(details.summary.credential_revision)?),
})
@@ -91,6 +91,7 @@ pub(crate) const ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL: &str = "with requested_win
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
group by requested.account_id, requested.window_key,
coalesce(mr.upstream_model_id, mr.requested_model_id)
order by requested.account_id, requested.window_key, request_count desc, model";
@@ -115,6 +116,7 @@ pub(crate) const ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL: &str = "with requeste
and mr.client_status_code between 200 and 399
and mr.cost_amount is not null
and mr.cost_currency is not null
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
group by requested.account_id, requested.window_key,
coalesce(mr.upstream_model_id, mr.requested_model_id), mr.cost_currency
order by requested.account_id, requested.window_key, model, mr.cost_currency";
@@ -16,6 +16,24 @@ use super::{namespace, resource_fingerprint};
const MAX_SESSION_AFFINITY_TTL: Duration = Duration::from_secs(24 * 60 * 60);
const CLAIM_OR_LOAD_SCRIPT: &str = r#"
local current = redis.call('GET', KEYS[1])
if current then
return current
end
redis.call('PSETEX', KEYS[1], tonumber(ARGV[2]), ARGV[1])
return ARGV[1]
"#;
const COMPARE_AND_BIND_SCRIPT: &str = r#"
local current = redis.call('GET', KEYS[1])
if not current or current == ARGV[1] then
redis.call('PSETEX', KEYS[1], tonumber(ARGV[3]), ARGV[2])
return ARGV[2]
end
return current
"#;
#[derive(Clone)]
pub struct RedisProviderSessionAffinityRepository {
connection: ConnectionManager,
@@ -79,11 +97,7 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<(), ProviderStoreError>> {
Box::pin(async move {
if ttl.is_zero() || ttl > MAX_SESSION_AFFINITY_TTL {
return Err(provider_invalid("validate provider session affinity TTL"));
}
let ttl_millis = u64::try_from(ttl.as_millis())
.map_err(|_| provider_invalid("validate provider session affinity TTL"))?;
let ttl_millis = session_affinity_ttl_millis(ttl)?;
let mut connection = self.connection.clone();
redis::cmd("PSETEX")
.arg(self.key(provider_kind, key)?)
@@ -95,6 +109,50 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
})
}
fn claim_or_load<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
candidate_account_id: &'a ProviderAccountId,
ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>> {
Box::pin(async move {
let ttl_millis = session_affinity_ttl_millis(ttl)?;
let mut connection = self.connection.clone();
let effective_account_id = redis::Script::new(CLAIM_OR_LOAD_SCRIPT)
.key(self.key(provider_kind, key)?)
.arg(candidate_account_id.as_str())
.arg(ttl_millis)
.invoke_async::<String>(&mut connection)
.await
.map_err(|_| provider_unavailable("claim provider session affinity"))?;
decode_account_id(effective_account_id)
})
}
fn compare_and_bind<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected_account_id: &'a ProviderAccountId,
replacement_account_id: &'a ProviderAccountId,
ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>> {
Box::pin(async move {
let ttl_millis = session_affinity_ttl_millis(ttl)?;
let mut connection = self.connection.clone();
let effective_account_id = redis::Script::new(COMPARE_AND_BIND_SCRIPT)
.key(self.key(provider_kind, key)?)
.arg(expected_account_id.as_str())
.arg(replacement_account_id.as_str())
.arg(ttl_millis)
.invoke_async::<String>(&mut connection)
.await
.map_err(|_| provider_unavailable("compare provider session affinity"))?;
decode_account_id(effective_account_id)
})
}
fn clear<'a>(
&'a self,
provider_kind: &'a ProviderKind,
@@ -112,6 +170,19 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
}
}
fn session_affinity_ttl_millis(ttl: Duration) -> Result<u64, ProviderStoreError> {
if ttl.is_zero() || ttl > MAX_SESSION_AFFINITY_TTL {
return Err(provider_invalid("validate provider session affinity TTL"));
}
u64::try_from(ttl.as_millis())
.map_err(|_| provider_invalid("validate provider session affinity TTL"))
}
fn decode_account_id(account_id: String) -> Result<ProviderAccountId, ProviderStoreError> {
ProviderAccountId::new(account_id)
.map_err(|_| provider_invalid("decode provider session affinity"))
}
fn provider_unavailable(operation: &'static str) -> ProviderStoreError {
ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation)
}
@@ -28,7 +28,7 @@ fn model_request_rejects_mismatched_client_key_live_id() {
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
requested_model_id: Some("coding".to_owned()),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
@@ -47,7 +47,7 @@ fn model_request_rejects_mismatched_client_key_live_id() {
}
#[tokio::test]
async fn merged_first_attempt_insert_should_match_sequential_semantics() {
async fn merged_model_less_first_attempt_should_match_sequential_semantics() {
let Some(database) = TestDatabase::create("execution_merged_insert").await else {
return;
};
@@ -75,10 +75,10 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
client_api_key_ref: "key_merged".to_owned(),
config_revision: 1,
protocol: "openai".to_owned(),
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
operation: "generate_image".to_owned(),
endpoint: "/v1/images/generations".to_owned(),
client_transport: "http_json".to_owned(),
requested_model_id: None,
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
@@ -99,8 +99,8 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
provider_kind: "openai".to_owned(),
provider_account_id: Some("acct_merged".to_owned()),
provider_account_ref: Some("acct_merged".to_owned()),
upstream_model_id: "gpt-test".to_owned(),
upstream_transport: "http_sse".to_owned(),
upstream_model_id: None,
upstream_transport: "http_json".to_owned(),
http_version: None,
};
@@ -109,22 +109,31 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
.await
.expect("merged insert");
let (attempt_count, send_state, provider_kind, outcome): (i32, String, String, String) =
sqlx::query_as(
"select attempt_count, upstream_send_state, provider_kind, outcome
let (attempt_count, send_state, provider_kind, outcome, requested_model, upstream_model): (
i32,
String,
String,
String,
Option<String>,
Option<String>,
) = sqlx::query_as(
"select attempt_count, upstream_send_state, provider_kind, outcome,
requested_model_id, upstream_model_id
from model_requests where id = 'req_merged'",
)
.fetch_one(&database.pool)
.await
.expect("load merged request");
)
.fetch_one(&database.pool)
.await
.expect("load merged request");
assert_eq!(
(
attempt_count,
send_state.as_str(),
provider_kind.as_str(),
outcome.as_str()
outcome.as_str(),
requested_model,
upstream_model,
),
(1, "not_sent", "openai", "running")
(1, "not_sent", "openai", "running", None, None)
);
// 后续 attempt 沿用常规 CAS 递增路径;已持久化的 sent 水位不被重试重置。
@@ -142,8 +151,8 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
provider_kind: "openai".to_owned(),
provider_account_id: Some("acct_merged".to_owned()),
provider_account_ref: Some("acct_merged".to_owned()),
upstream_model_id: "gpt-test".to_owned(),
upstream_transport: "http_sse".to_owned(),
upstream_model_id: None,
upstream_transport: "http_json".to_owned(),
http_version: None,
})
.await
@@ -182,7 +191,7 @@ async fn model_request_persists_group_routing_snapshot_without_live_group_foreig
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
requested_model_id: Some("coding".to_owned()),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -10,7 +10,12 @@ use gateway_admin::{
UpdateAccount,
},
observability::TimeRange,
provider_credentials::{CredentialListQuery, CredentialListWindow, CredentialStateFilter},
provider_credentials::{
AuthorizationCommit, AuthorizationCredentialCommit, AuthorizationMutationTarget,
AuthorizationOwnerBinding, CredentialListQuery, CredentialListWindow,
CredentialStateFilter, PendingAuthorizationMutation, PreparedCredentialCreate,
ProviderDocument,
},
},
ports::store::AccountStore,
};
@@ -1326,6 +1331,75 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
database.close().await;
}
#[tokio::test]
async fn authorization_create_returns_existing_account_id_when_identity_is_upserted() {
let Some(database) = TestDatabase::create("provider_account_authorization_upsert").await else {
return;
};
PgProviderAccountRepository::new(database.pool.clone())
.insert_provider_account(account(
"acct_authorization_existing",
"user-authorization-upsert",
))
.await
.expect("seed existing authorized identity");
let context = MutationContext {
actor: MutationActor::System,
request_id: "request_authorization_upsert".to_owned(),
};
let provider_kind = ProviderKind::new("openai").expect("OpenAI Provider kind");
let serde_json::Value::Object(provider_material) =
json!({ "access_token": "authorized-secret" })
else {
unreachable!("credential fixture must be an object");
};
let result = PgAdminAccountStore::new(database.pool.clone(), None)
.commit_authorization(
AuthorizationCommit {
pending: PendingAuthorizationMutation::new(
provider_kind.clone(),
AuthorizationMutationTarget::Create {
name: "authorized account".to_owned(),
},
AuthorizationOwnerBinding::from_context(&context),
),
credential: AuthorizationCredentialCommit::Create(PreparedCredentialCreate {
account_id: ProviderAccountId::new("acct_authorization_candidate")
.expect("candidate account ID"),
provider_kind,
name: "authorized account".to_owned(),
email: Some("authorized@example.invalid".to_owned()),
upstream_user_id: Some("user-authorization-upsert".to_owned()),
upstream_account_id: None,
plan_type: Some("free".to_owned()),
authentication_kind: "oauth".to_owned(),
provider_material: ProviderDocument::new(OpaqueProviderData::new(
provider_material,
)),
has_refresh_token: true,
access_token_expires_at: Some(Utc::now() + TimeDelta::hours(1)),
next_refresh_at: None,
enabled: true,
credential_state: CredentialState::Ready,
credential_observed_at: Utc::now(),
}),
},
&context,
)
.await
.expect("authorize existing identity");
assert_eq!(
(
result.account_id.as_str(),
result.credential_revision.map(|revision| revision.get()),
),
("acct_authorization_existing", Some(2)),
);
database.close().await;
}
#[tokio::test]
async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
let Some(database) = TestDatabase::create("provider_account_core_refresh").await else {
@@ -436,7 +436,7 @@ async fn diagnostics_should_fallback_to_name_then_ref_for_missing_snapshots() {
provider_kind: "openai".to_owned(),
provider_account_id: None,
provider_account_ref: Some("acct_legacy_deleted".to_owned()),
upstream_model_id: "legacy-model".to_owned(),
upstream_model_id: Some("legacy-model".to_owned()),
upstream_transport: "http_sse".to_owned(),
http_version: None,
},
@@ -547,7 +547,7 @@ async fn api_key_diagnostics_should_display_key_name_and_fallback_to_ref() {
provider_kind: "openai".to_owned(),
provider_account_id: None,
provider_account_ref: Some("acct_key".to_owned()),
upstream_model_id: "upstream-model".to_owned(),
upstream_model_id: Some("upstream-model".to_owned()),
upstream_transport: "http_sse".to_owned(),
http_version: None,
};
@@ -664,7 +664,7 @@ fn new_request(id: &str, started_at: DateTime<Utc>) -> NewModelRequest {
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
requested_model_id: Some("coding".to_owned()),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
@@ -688,7 +688,7 @@ fn attempt(id: &str, count: u32, account_id: &str) -> ModelRequestAttemptStart {
provider_kind: "openai".to_owned(),
provider_account_id: Some(account_id.to_owned()),
provider_account_ref: Some(account_id.to_owned()),
upstream_model_id: "upstream-model".to_owned(),
upstream_model_id: Some("upstream-model".to_owned()),
upstream_transport: "http_sse".to_owned(),
http_version: None,
}
@@ -65,6 +65,55 @@ async fn session_affinity_should_overwrite_the_previous_account() {
);
}
#[tokio::test]
async fn session_affinity_claim_should_keep_the_first_account() {
let Some((repository, _connection, _namespace)) = affinity_repository().await else {
return;
};
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("claim-session").expect("affinity key");
let first = ProviderAccountId::new("acct_first").expect("first account");
let second = ProviderAccountId::new("acct_second").expect("second account");
let first_winner = repository
.claim_or_load(&provider, &key, &first, Duration::from_secs(60))
.await
.expect("claim first affinity");
let second_winner = repository
.claim_or_load(&provider, &key, &second, Duration::from_secs(60))
.await
.expect("load existing affinity");
assert_eq!((first_winner, second_winner), (first.clone(), first));
}
#[tokio::test]
async fn session_affinity_compare_should_not_replace_a_newer_winner() {
let Some((repository, _connection, _namespace)) = affinity_repository().await else {
return;
};
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("compare-session").expect("affinity key");
let first = ProviderAccountId::new("acct_first").expect("first account");
let second = ProviderAccountId::new("acct_second").expect("second account");
let stale = ProviderAccountId::new("acct_stale").expect("stale account");
repository
.bind(&provider, &key, &first, Duration::from_secs(60))
.await
.expect("seed affinity");
let migrated = repository
.compare_and_bind(&provider, &key, &first, &second, Duration::from_secs(60))
.await
.expect("migrate expected affinity");
let stale_result = repository
.compare_and_bind(&provider, &key, &first, &stale, Duration::from_secs(60))
.await
.expect("reject stale affinity migration");
assert_eq!((migrated, stale_result), (second.clone(), second));
}
#[tokio::test]
async fn session_affinity_should_apply_ttl_and_support_explicit_clear() {
let Some((repository, mut connection, namespace)) = affinity_repository().await else {
+264 -4
View File
@@ -13,13 +13,17 @@ use gateway_admin::model::observability::{
};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwner, AuthorizationOwnerBinding,
AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard,
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit, CredentialCommitGuard,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedAuthorizationCredential,
PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation,
PreparedCredentialRotationFacts, ProviderDocument, ProviderExport,
ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
ProviderQuotaRequest, ProviderQuotaWindow, ProviderQuotaWindowRole, ProviderResetCredit,
ProviderResetCreditResult, ProviderResetCredits, ProviderUsageStatistics,
ProviderUsageStatisticsCycle, ProviderUsageStatisticsDay, ProviderUsageStatisticsModel,
ProviderUsageStatisticsRequest, ProviderUsageStatisticsServiceTier,
ProviderUsageStatisticsSummary, ProviderUsageStatisticsTokens, QuotaLocalUsageAttribution,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
@@ -47,8 +51,10 @@ use crate::credential::{
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexPendingAuthorization, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole,
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
CodexResetCreditsError, CodexUsageStatistics, CodexUsageStatisticsError,
CodexUsageStatisticsServiceTier, CodexUsageStatisticsTokens, CompleteCodexOAuthAuthorization,
CompletedCodexOAuthCredential, ExportManagedCodexCredential, StartCodexOAuthAuthorization,
StoredCodexPendingAuthorization,
};
use crate::credential::{
CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref, parse_access_token_expiration,
@@ -442,6 +448,54 @@ impl ProviderAdmin for OpenAiAdminProvider {
Ok(project_quota(snapshot, &account))
}
async fn usage_statistics(
&self,
request: ProviderUsageStatisticsRequest,
) -> Result<ProviderUsageStatistics, ProviderAdminError> {
self.account(&request.account_id).await?;
let statistics = self
.quota
.usage_statistics(
&request.account_id,
request.cycle_offset,
request.utc_offset_minutes,
)
.await
.map_err(map_usage_statistics_error)?;
project_usage_statistics(statistics)
}
async fn reset_credits(
&self,
account_id: &ProviderAccountId,
) -> Result<ProviderResetCredits, ProviderAdminError> {
self.account(account_id).await?;
self.quota
.list_reset_credits(account_id)
.await
.map(project_reset_credits)
.map_err(map_reset_credits_error)
}
async fn consume_reset_credit(
&self,
command: ConsumeProviderResetCredit,
) -> Result<ProviderResetCreditResult, ProviderAdminError> {
self.account(&command.account_id).await?;
self.quota
.consume_reset_credit(
&command.account_id,
command.credit_id.as_deref(),
command.redeem_request_id,
)
.await
.map(|result| ProviderResetCreditResult {
code: result.code,
credit: result.credit.map(project_reset_credit),
})
.map_err(map_reset_credits_error)
}
async fn models(
&self,
account_id: &ProviderAccountId,
@@ -768,6 +822,119 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
}
}
fn project_usage_statistics(
statistics: CodexUsageStatistics,
) -> Result<ProviderUsageStatistics, ProviderAdminError> {
let models = statistics
.models
.into_iter()
.map(|model| {
Ok(ProviderUsageStatisticsModel {
key: model.key,
model: model.model,
service_tier: project_usage_service_tier(model.service_tier),
credit_share: model.credit_share,
quota_share: model.quota_share,
tokens: project_usage_tokens(model.tokens),
estimated_cost: model.estimated_cost.map(currency_cost).transpose()?,
has_unknown_pricing: model.has_unknown_pricing,
has_estimated_allocation: model.has_estimated_allocation,
has_rate_fallback: model.has_rate_fallback,
has_missing_token_data: model.has_missing_token_data,
})
})
.collect::<Result<Vec<_>, ProviderAdminError>>()?;
let daily = statistics
.daily
.into_iter()
.map(|day| {
Ok(ProviderUsageStatisticsDay {
date: day.date,
credit_share: day.credit_share,
tokens: project_usage_tokens(day.tokens),
estimated_cost: day.estimated_cost.map(currency_cost).transpose()?,
has_unknown_pricing: day.has_unknown_pricing,
has_missing_token_data: day.has_missing_token_data,
is_boundary_day: day.is_boundary_day,
})
})
.collect::<Result<Vec<_>, ProviderAdminError>>()?;
Ok(ProviderUsageStatistics {
cycle: ProviderUsageStatisticsCycle {
offset: statistics.cycle.offset,
start_at: statistics.cycle.start_at,
end_at: statistics.cycle.end_at,
window_seconds: statistics.cycle.window_seconds,
used_percent: statistics.cycle.used_percent,
is_current: statistics.cycle.is_current,
can_go_previous: statistics.cycle.can_go_previous,
can_go_next: statistics.cycle.can_go_next,
},
summary: ProviderUsageStatisticsSummary {
tokens: project_usage_tokens(statistics.summary.tokens),
estimated_cost: statistics
.summary
.estimated_cost
.map(currency_cost)
.transpose()?,
projected_tokens: statistics.summary.projected_tokens,
projected_cost: statistics
.summary
.projected_cost
.map(currency_cost)
.transpose()?,
day_count: statistics.summary.day_count,
has_unknown_pricing: statistics.summary.has_unknown_pricing,
has_missing_token_data: statistics.summary.has_missing_token_data,
},
models,
daily,
})
}
const fn project_usage_tokens(tokens: CodexUsageStatisticsTokens) -> ProviderUsageStatisticsTokens {
ProviderUsageStatisticsTokens {
uncached_input: tokens.uncached_input,
cached_input: tokens.cached_input,
output: tokens.output,
total: tokens.total,
}
}
const fn project_usage_service_tier(
service_tier: CodexUsageStatisticsServiceTier,
) -> ProviderUsageStatisticsServiceTier {
match service_tier {
CodexUsageStatisticsServiceTier::Standard => ProviderUsageStatisticsServiceTier::Standard,
CodexUsageStatisticsServiceTier::Fast => ProviderUsageStatisticsServiceTier::Fast,
}
}
fn project_reset_credits(
credits: crate::transport::CodexRateLimitResetCredits,
) -> ProviderResetCredits {
ProviderResetCredits {
available_count: credits.available_count,
credits: credits
.credits
.into_iter()
.map(project_reset_credit)
.collect(),
}
}
fn project_reset_credit(
credit: crate::transport::CodexRateLimitResetCredit,
) -> ProviderResetCredit {
ProviderResetCredit {
id: credit.id,
status: credit.status,
title: credit.title,
expires_at: credit.expires_at,
reset_type: credit.reset_type,
}
}
/// `secondary_window` 有时只是上游的空占位:没有时长、重置时间,也没有用量。
/// 它不能被可靠翻译成 5 小时或周额度,因此不应占用账号面板;带有实际事实的
/// 次级窗口(时长、重置、非零用量或触顶)仍完整保留。
@@ -1283,6 +1450,8 @@ const fn provider_admin_error_code(kind: ProviderAdminErrorKind) -> &'static str
ProviderAdminErrorKind::NotFound => "not_found",
ProviderAdminErrorKind::Conflict => "conflict",
ProviderAdminErrorKind::Unavailable => "unavailable",
ProviderAdminErrorKind::CredentialRefreshRequired => "credential_refresh_required",
ProviderAdminErrorKind::BadGateway => "bad_gateway",
ProviderAdminErrorKind::Internal => "internal",
}
}
@@ -1338,6 +1507,97 @@ fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
})
}
fn map_usage_statistics_error(error: CodexUsageStatisticsError) -> ProviderAdminError {
use CodexUsageStatisticsError as Error;
match error {
Error::InvalidRequest | Error::InvalidCredentialData => {
provider_admin_error(ProviderAdminErrorKind::Invalid)
}
Error::NotFound => provider_admin_error(ProviderAdminErrorKind::NotFound),
Error::Store { .. } | Error::TransportUnavailable => {
provider_admin_error(ProviderAdminErrorKind::Unavailable)
}
Error::QuotaWindowUnavailable => provider_admin_error(ProviderAdminErrorKind::Conflict)
.with_message("OpenAI did not expose a resettable quota window for usage statistics"),
Error::CredentialRefreshRequired { upstream_body } => {
let mut error = provider_admin_error(ProviderAdminErrorKind::CredentialRefreshRequired);
if let Some(body) = upstream_body {
error = error.with_message(format!(
"OpenAI usage-statistics upstream returned HTTP 401: {}",
bounded_upstream_body(&body)
));
}
error
}
Error::Upstream {
status,
body,
retry_after_seconds,
} => {
let retry_after = retry_after_seconds
.map(|seconds| format!("; retry-after={seconds}s"))
.unwrap_or_default();
provider_admin_error(ProviderAdminErrorKind::BadGateway).with_message(format!(
"OpenAI usage-statistics upstream returned HTTP {status}{retry_after}: {}",
bounded_upstream_body(&body)
))
}
}
}
fn map_reset_credits_error(error: CodexResetCreditsError) -> ProviderAdminError {
use CodexResetCreditsError as Error;
match error {
Error::InvalidCredentialData => provider_admin_error(ProviderAdminErrorKind::Invalid),
Error::NotFound => provider_admin_error(ProviderAdminErrorKind::NotFound),
Error::Store { .. } | Error::TransportUnavailable => {
provider_admin_error(ProviderAdminErrorKind::Unavailable)
}
Error::CredentialRefreshRequired { upstream_body } => {
let mut error =
provider_admin_error(ProviderAdminErrorKind::CredentialRefreshRequired);
if let Some(body) = upstream_body {
error = error.with_message(format!(
"OpenAI reset-credit upstream returned HTTP 401: {}",
bounded_upstream_body(&body)
));
}
error
}
Error::Upstream {
status,
body,
retry_after_seconds,
} => {
let retry_after = retry_after_seconds
.map(|seconds| format!("; retry-after={seconds}s"))
.unwrap_or_default();
provider_admin_error(ProviderAdminErrorKind::BadGateway).with_message(format!(
"OpenAI reset-credit upstream returned HTTP {status}{retry_after}: {}",
bounded_upstream_body(&body)
))
}
Error::ConsumeResultUnknown => provider_admin_error(ProviderAdminErrorKind::BadGateway)
.with_message(
"OpenAI reset-credit consume result is unknown; refresh the credit list before retrying",
),
}
}
fn bounded_upstream_body(body: &str) -> String {
const MAX_ADMIN_UPSTREAM_BODY_BYTES: usize = 8 * 1024;
if body.len() <= MAX_ADMIN_UPSTREAM_BODY_BYTES {
return body.to_owned();
}
let mut end = MAX_ADMIN_UPSTREAM_BODY_BYTES;
while !body.is_char_boundary(end) {
end -= 1;
}
body[..end].to_owned()
}
fn map_catalog_error(error: CodexCredentialCatalogError) -> ProviderAdminError {
use CodexCredentialCatalogError as Error;
provider_admin_error(match error {
+1 -12
View File
@@ -155,8 +155,6 @@ impl CodexApiConfig {
pub struct CodexWebSocketPoolSettings {
pub enabled: bool,
pub max_age_ms: u64,
pub max_per_account: usize,
pub max_total: usize,
pub max_connecting: usize,
pub initial_event_timeout_ms: u64,
}
@@ -166,8 +164,6 @@ impl Default for CodexWebSocketPoolSettings {
Self {
enabled: true,
max_age_ms: 55 * 60 * 1000,
max_per_account: 8,
max_total: 64,
max_connecting: 8,
initial_event_timeout_ms: 20_000,
}
@@ -176,12 +172,7 @@ impl Default for CodexWebSocketPoolSettings {
impl CodexWebSocketPoolSettings {
fn validate(&self) -> Result<(), OpenAiConfigError> {
if self.max_age_ms == 0
|| self.max_per_account == 0
|| self.max_total == 0
|| self.max_connecting == 0
|| self.max_connecting > self.max_total
{
if self.max_age_ms == 0 || self.max_connecting == 0 {
return Err(OpenAiConfigError::InvalidField("openai.ws_pool"));
}
Ok(())
@@ -191,8 +182,6 @@ impl CodexWebSocketPoolSettings {
CodexWebSocketPoolConfig {
enabled: self.enabled,
max_age: Duration::from_millis(self.max_age_ms),
max_per_account: self.max_per_account,
max_total: self.max_total,
max_connecting: self.max_connecting,
initial_event_timeout: (self.initial_event_timeout_ms != 0)
.then(|| Duration::from_millis(self.initial_event_timeout_ms)),
@@ -1288,9 +1288,9 @@ fn parse_oauth_import_tokens(
Value::Object(object) => {
for (key, value) in object {
let token = match key.as_str() {
"accessToken" => &mut access_token,
"refreshToken" => &mut refresh_token,
"idToken" => &mut id_token,
"accessToken" | "access_token" => &mut access_token,
"refreshToken" | "refresh_token" => &mut refresh_token,
"idToken" | "id_token" => &mut id_token,
_ => {
pending.push(value);
continue;
@@ -1,4 +1,4 @@
//! OpenAI 会话锚点到 Store 不透明亲和键的单向派生。
//! OpenAI 根会话锚点到 Store 不透明账号亲和键的单向派生。
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::ProviderSessionAffinityKey;
@@ -11,34 +11,8 @@ pub(crate) fn derive_codex_session_affinity_key(
request: &CodexResponsesRequest,
client_api_key_id: &ClientApiKeyId,
) -> Option<ProviderSessionAffinityKey> {
let session_key = if let Some(conversation_id) = request
.local_conversation_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if let Ok(key) = ProviderSessionAffinityKey::try_new(conversation_id.to_owned()) {
Some(key)
} else {
opaque_affinity_key("local-conversation", conversation_id)
}
} else {
let (domain, value) = derive_conversation_anchor(request)?;
opaque_affinity_key(domain, &value)
};
let session_discriminator = request.subagent_kind();
let session_key =
session_key.and_then(|session_key| match session_discriminator.as_deref() {
None => Some(session_key),
Some(discriminator) => {
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
opaque_affinity_key(
"subagent-session",
&format!("{discriminator}\0{}", session_key.expose_to_store()),
)
}
})?;
let (domain, value) = derive_account_affinity_anchor(request)?;
let session_key = opaque_affinity_key(domain, &value)?;
opaque_affinity_key(
"client-session",
&format!(
@@ -49,6 +23,33 @@ pub(crate) fn derive_codex_session_affinity_key(
)
}
/// 账号选择优先使用根会话事实;`local_conversation_id` 仍可按 child thread 隔离
/// WebSocket 与 continuation 状态,但 child 身份不再拆分账号首选项。账号不可调度时
/// selector 仍可按既有策略切换,并由 request scope 清理账号绑定状态。
fn derive_account_affinity_anchor(
request: &CodexResponsesRequest,
) -> Option<(&'static str, String)> {
non_empty(request.client_session_id.as_deref())
.map(|value| ("root-session", value.to_owned()))
.or_else(|| {
non_empty(request.client_conversation_id.as_deref())
.map(|value| ("root-conversation", value.to_owned()))
})
.or_else(|| {
request
.explicit_prompt_cache_key
.then(|| request.prompt_cache_key())
.flatten()
.and_then(|value| non_empty(Some(value)))
.map(|value| ("root-prompt-cache", value.to_owned()))
})
.or_else(|| {
non_empty(request.local_conversation_id.as_deref())
.map(|value| ("local-conversation", value.to_owned()))
})
.or_else(|| derive_conversation_anchor(request))
}
fn opaque_affinity_key(domain: &str, value: &str) -> Option<ProviderSessionAffinityKey> {
let mut hasher = Sha256::new();
hasher.update(b"codex-session-affinity-v1\0");
@@ -47,7 +47,10 @@ pub use oauth::{
pub use quota::{
CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService,
CodexQuotaFact, CodexQuotaRefreshPolicy, CodexQuotaSyncSummary, CodexQuotaWindow,
CodexQuotaWindowKind, CodexQuotaWindowRole, parse_codex_quota_usage,
CodexQuotaWindowKind, CodexQuotaWindowRole, CodexResetCreditsError, CodexUsageStatistics,
CodexUsageStatisticsCycle, CodexUsageStatisticsDay, CodexUsageStatisticsError,
CodexUsageStatisticsModel, CodexUsageStatisticsServiceTier, CodexUsageStatisticsSummary,
CodexUsageStatisticsTokens, parse_codex_quota_usage,
};
pub use refresh::{
CodexCredentialRefreshError, CodexCredentialRefreshOutcome, CodexCredentialRefreshService,
@@ -58,11 +61,11 @@ pub use security::{
CodexCredentialCodec, CodexCredentialDataError, CodexRuntimeAuthentication,
CodexRuntimeCredential,
};
pub(crate) use selector::CodexCyberPolicyScope;
pub use selector::{
CodexAccountFailure, CodexCredentialLease, CodexCredentialSelector, CredentialSelectionError,
SelectCodexCredential,
};
pub(crate) use selector::{CodexCyberPolicyScope, SelectCodexProviderEndpointCredential};
pub use types::{
CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY, CODEX_AUTHENTICATION_KIND_OAUTH, CodexAccountProfile,
CodexAgentIdentityAuthMode, CodexAgentIdentityCredentialData, CodexCookie,
@@ -7,17 +7,23 @@
mod document;
pub(crate) mod evidence;
pub(crate) mod snapshot;
mod usage_statistics;
pub use snapshot::{
CodexAccountQuotaSnapshot, CodexQuotaFact, CodexQuotaWindow, CodexQuotaWindowKind,
CodexQuotaWindowRole, parse_codex_quota_usage,
};
pub use usage_statistics::{
CodexUsageStatistics, CodexUsageStatisticsCycle, CodexUsageStatisticsDay,
CodexUsageStatisticsModel, CodexUsageStatisticsServiceTier, CodexUsageStatisticsSummary,
CodexUsageStatisticsTokens,
};
use std::collections::{BTreeMap, BTreeSet};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant, SystemTime};
use chrono::Utc;
use chrono::{DateTime, FixedOffset, Utc};
use gateway_core::engine::credential::{
AccountErrorReason, AccountQuotaSignals, CredentialRevision, CredentialState,
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
@@ -36,10 +42,15 @@ use tokio::sync::Mutex;
use uuid::Uuid;
use crate::transport::profile::CodexWireProfileState;
use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext};
use crate::transport::usage_statistics::CodexDailyUsageReport;
use crate::transport::{
CodexBackendClient, CodexClientError, CodexRateLimitResetCredits,
CodexRateLimitResetCreditsConsumeResult, CodexRequestContext,
};
use super::agent_identity::{CodexAgentIdentityTaskService, PreparedCodexRuntimeCredential};
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use super::types::CODEX_AUTHENTICATION_KIND_OAUTH;
use document::{
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId,
canonicalize_rate_limit_document,
@@ -118,6 +129,105 @@ pub enum CodexCredentialQuotaError {
Upstream { detail: String },
}
/// 主动额度重置卡查询/消费失败。
#[derive(Error)]
pub enum CodexResetCreditsError {
#[error("Codex reset-credit credential data is invalid")]
InvalidCredentialData,
#[error("Codex OAuth access token must be refreshed before using reset credits")]
CredentialRefreshRequired { upstream_body: Option<String> },
#[error("Codex reset-credit account was not found")]
NotFound,
#[error("provider account store is unavailable: {detail}")]
Store { detail: String },
#[error("Codex reset-credit upstream returned HTTP {status}")]
Upstream {
status: u16,
body: String,
retry_after_seconds: Option<u64>,
},
#[error("Codex reset-credit query transport is unavailable")]
TransportUnavailable,
#[error("Codex reset-credit consume result is unknown")]
ConsumeResultUnknown,
}
/// 官方每日用量统计查询失败。
#[derive(Error)]
pub enum CodexUsageStatisticsError {
#[error("Codex usage-statistics request is invalid")]
InvalidRequest,
#[error("Codex usage-statistics credential data is invalid")]
InvalidCredentialData,
#[error("Codex OAuth access token must be refreshed before querying usage statistics")]
CredentialRefreshRequired { upstream_body: Option<String> },
#[error("Codex usage-statistics account was not found")]
NotFound,
#[error("provider account store is unavailable: {detail}")]
Store { detail: String },
#[error("Codex usage-statistics quota window is unavailable")]
QuotaWindowUnavailable,
#[error("Codex usage-statistics upstream returned HTTP {status}")]
Upstream {
status: u16,
body: String,
retry_after_seconds: Option<u64>,
},
#[error("Codex usage-statistics query transport is unavailable")]
TransportUnavailable,
}
impl std::fmt::Debug for CodexUsageStatisticsError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::InvalidRequest => formatter.write_str("InvalidRequest"),
Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"),
Self::CredentialRefreshRequired { .. } => {
formatter.write_str("CredentialRefreshRequired { upstream_body: <redacted> }")
}
Self::NotFound => formatter.write_str("NotFound"),
Self::Store { .. } => formatter.write_str("Store { detail: <redacted> }"),
Self::QuotaWindowUnavailable => formatter.write_str("QuotaWindowUnavailable"),
Self::Upstream {
status,
retry_after_seconds,
..
} => formatter
.debug_struct("Upstream")
.field("status", status)
.field("body", &"<redacted>")
.field("retry_after_seconds", retry_after_seconds)
.finish(),
Self::TransportUnavailable => formatter.write_str("TransportUnavailable"),
}
}
}
impl std::fmt::Debug for CodexResetCreditsError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"),
Self::CredentialRefreshRequired { .. } => {
formatter.write_str("CredentialRefreshRequired { upstream_body: <redacted> }")
}
Self::NotFound => formatter.write_str("NotFound"),
Self::Store { .. } => formatter.write_str("Store { detail: <redacted> }"),
Self::Upstream {
status,
retry_after_seconds,
..
} => formatter
.debug_struct("Upstream")
.field("status", status)
.field("body", &"<redacted>")
.field("retry_after_seconds", retry_after_seconds)
.finish(),
Self::TransportUnavailable => formatter.write_str("TransportUnavailable"),
Self::ConsumeResultUnknown => formatter.write_str("ConsumeResultUnknown"),
}
}
}
impl From<gateway_core::error::StoreError> for CodexCredentialQuotaError {
fn from(error: gateway_core::error::StoreError) -> Self {
Self::Store {
@@ -135,6 +245,7 @@ pub struct CodexCredentialQuotaService {
agent_identity: Arc<CodexAgentIdentityTaskService>,
cooldowns: Arc<dyn ProviderCooldownPort>,
scheduling: CodexQuotaSchedulingProjection,
reset_consume_locks: Mutex<HashMap<ProviderAccountId, Arc<Mutex<()>>>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -189,6 +300,30 @@ enum CodexQuotaFetchError {
},
}
enum ResetCreditAttemptError {
InvalidCredential,
Upstream(CodexClientError),
}
enum DailyUsageAttemptError {
InvalidCredential,
Upstream(CodexClientError),
}
enum DailyUsageFetchError {
InvalidCredential,
Recovery,
Upstream {
prepared: Box<PreparedCodexRuntimeCredential>,
error: CodexClientError,
},
}
struct FetchedDailyUsage {
model_breakdown: Value,
daily_totals: Value,
}
impl CodexQuotaSchedulingProjection {
fn invalidate(&self, account_ids: &[ProviderAccountId]) {
let mut state = self
@@ -413,9 +548,236 @@ impl CodexCredentialQuotaService {
agent_identity,
cooldowns,
scheduling: CodexQuotaSchedulingProjection::default(),
reset_consume_locks: Mutex::new(HashMap::new()),
}
}
/// 查询当前账号由 Codex Desktop 暴露的主动额度重置卡。
pub async fn list_reset_credits(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexRateLimitResetCredits, CodexResetCreditsError> {
let account = self.reset_credit_account(account_id).await?;
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let mut prepared = self
.agent_identity
.prepare(&account)
.await
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?;
let request_id = format!("reset_credits_{}", Uuid::now_v7().simple());
let mut result = list_reset_credits_once(&client, &prepared, &request_id).await;
if let Err(ResetCreditAttemptError::Upstream(error)) = &result
&& let Some(recovered) = self
.agent_identity
.recover_after_rejected_task(
prepared.account.id(),
&prepared.credential.authentication,
error,
)
.await
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?
{
prepared = recovered;
result = list_reset_credits_once(&client, &prepared, &request_id).await;
}
result.map_err(|error| map_reset_credit_attempt_error(error, &prepared, false))
}
/// 查询当前账号的官方每日用量报表,并把报表范围锚定到最接近七天的额度窗口。
pub async fn usage_statistics(
&self,
account_id: &ProviderAccountId,
cycle_offset: i8,
utc_offset_minutes: i16,
) -> Result<CodexUsageStatistics, CodexUsageStatisticsError> {
const MAX_CYCLE_OFFSET: u8 = 8;
const MAX_UTC_OFFSET_MINUTES: i16 = 14 * 60;
const TARGET_WINDOW_SECONDS: u64 = 7 * 24 * 60 * 60;
if !(-(MAX_CYCLE_OFFSET as i8)..=0).contains(&cycle_offset)
|| !(-MAX_UTC_OFFSET_MINUTES..=MAX_UTC_OFFSET_MINUTES).contains(&utc_offset_minutes)
{
return Err(CodexUsageStatisticsError::InvalidRequest);
}
let account = self
.store
.get_account(account_id)
.await
.map_err(|error| CodexUsageStatisticsError::Store {
detail: error.to_string(),
})?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexUsageStatisticsError::NotFound)?;
let observed_at = SystemTime::now();
if !access_token_is_current(&account, observed_at) {
return Err(CodexUsageStatisticsError::CredentialRefreshRequired {
upstream_body: None,
});
}
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let FetchedCodexQuota { account, value } = self
.fetch_usage_with_recovery(&client, &account)
.await
.map_err(map_statistics_quota_fetch_error)?;
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&value,
)
.map_err(|_| CodexUsageStatisticsError::InvalidCredentialData)?;
let window = snapshot
.windows()
.iter()
.filter(|window| window.is_account_wide())
.filter_map(|window| {
Some((
window,
window.window_seconds()?.abs_diff(TARGET_WINDOW_SECONDS),
))
})
.min_by_key(|(_, distance)| *distance)
.map(|(window, _)| window)
.ok_or(CodexUsageStatisticsError::QuotaWindowUnavailable)?;
let window_seconds = window
.window_seconds()
.filter(|seconds| *seconds > 0)
.ok_or(CodexUsageStatisticsError::QuotaWindowUnavailable)?;
let end_at = window
.reset_at()
.ok_or(CodexUsageStatisticsError::QuotaWindowUnavailable)?;
let window_delta = chrono::TimeDelta::seconds(
i64::try_from(window_seconds)
.map_err(|_| CodexUsageStatisticsError::QuotaWindowUnavailable)?,
);
let current_start_at = end_at - window_delta;
let timezone = FixedOffset::east_opt(i32::from(utc_offset_minutes) * 60)
.ok_or(CodexUsageStatisticsError::InvalidRequest)?;
let now = DateTime::<Utc>::from(observed_at);
let period = usage_statistics::statistics_period(
current_start_at,
end_at,
window_seconds,
cycle_offset,
timezone,
now,
)
.ok_or(CodexUsageStatisticsError::InvalidRequest)?;
let fetched = self
.fetch_daily_usage_with_recovery(
&client,
&account,
period.query_start_date,
period.query_end_date,
)
.await
.map_err(map_daily_usage_fetch_error)?;
Ok(usage_statistics::build_usage_statistics(
usage_statistics::BuildUsageStatistics {
period,
timezone,
current_used_percent: window.used_percent(),
now,
model_breakdown: &fetched.model_breakdown,
daily_totals: &fetched.daily_totals,
max_cycle_offset: MAX_CYCLE_OFFSET,
},
))
}
/// 消费一张主动额度重置卡。相同账号在本进程内串行,且不做传输重试。
pub async fn consume_reset_credit(
&self,
account_id: &ProviderAccountId,
credit_id: Option<&str>,
redeem_request_id: Uuid,
) -> Result<CodexRateLimitResetCreditsConsumeResult, CodexResetCreditsError> {
let lock = {
let mut locks = self.reset_consume_locks.lock().await;
Arc::clone(
locks
.entry(account_id.clone())
.or_insert_with(|| Arc::new(Mutex::new(()))),
)
};
let _guard = lock.lock().await;
let account = self.reset_credit_account(account_id).await?;
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let mut prepared = self
.agent_identity
.prepare(&account)
.await
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?;
let request_id = format!("reset_credit_consume_{}", Uuid::now_v7().simple());
let mut result = consume_reset_credit_once(
&client,
&prepared,
&request_id,
credit_id,
redeem_request_id,
)
.await;
// Agent Identity task recovery is the only automatic retry below. It keeps the exact
// official credit selection and redeem_request_id; generic HTTP/transport failures are
// never retried because the debit result may already be committed upstream.
if let Err(ResetCreditAttemptError::Upstream(error)) = &result
&& let Some(recovered) = self
.agent_identity
.recover_after_rejected_task(
prepared.account.id(),
&prepared.credential.authentication,
error,
)
.await
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?
{
prepared = recovered;
result = consume_reset_credit_once(
&client,
&prepared,
&request_id,
credit_id,
redeem_request_id,
)
.await;
}
result.map_err(|error| map_reset_credit_attempt_error(error, &prepared, true))
}
async fn reset_credit_account(
&self,
account_id: &ProviderAccountId,
) -> Result<ProviderAccount, CodexResetCreditsError> {
let account = self
.store
.get_account(account_id)
.await
.map_err(|error| CodexResetCreditsError::Store {
detail: error.to_string(),
})?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexResetCreditsError::NotFound)?;
if !access_token_is_current(&account, SystemTime::now()) {
return Err(CodexResetCreditsError::CredentialRefreshRequired {
upstream_body: None,
});
}
Ok(account)
}
/// 真实推理错误只更新额度访问事实,不伪造 Provider JSON 或展示百分比。
pub(crate) async fn record_confirmed_exhaustion(
&self,
@@ -1085,6 +1447,247 @@ impl CodexCredentialQuotaService {
}
}
}
async fn fetch_daily_usage_with_recovery(
&self,
client: &CodexBackendClient,
account: &ProviderAccount,
start_date: chrono::NaiveDate,
end_date: chrono::NaiveDate,
) -> Result<FetchedDailyUsage, DailyUsageFetchError> {
let mut prepared = self
.agent_identity
.prepare(account)
.await
.map_err(|_| DailyUsageFetchError::InvalidCredential)?;
let mut result = fetch_daily_usage_once(client, &prepared, start_date, end_date).await;
if let Err(DailyUsageAttemptError::Upstream(error)) = &result
&& let Some(recovered) = self
.agent_identity
.recover_after_rejected_task(
prepared.account.id(),
&prepared.credential.authentication,
error,
)
.await
.map_err(|_| DailyUsageFetchError::Recovery)?
{
prepared = recovered;
result = fetch_daily_usage_once(client, &prepared, start_date, end_date).await;
}
match result {
Ok(value) => Ok(value),
Err(DailyUsageAttemptError::InvalidCredential) => {
Err(DailyUsageFetchError::InvalidCredential)
}
Err(DailyUsageAttemptError::Upstream(error)) => Err(DailyUsageFetchError::Upstream {
prepared: Box::new(prepared),
error,
}),
}
}
}
async fn fetch_daily_usage_once(
client: &CodexBackendClient,
prepared: &PreparedCodexRuntimeCredential,
start_date: chrono::NaiveDate,
end_date: chrono::NaiveDate,
) -> Result<FetchedDailyUsage, DailyUsageAttemptError> {
let authorization = prepared
.credential
.authentication
.authorization_header(Utc::now())
.map_err(|_| DailyUsageAttemptError::InvalidCredential)?;
let model_request_id = format!("usage_statistics_{}", Uuid::now_v7().simple());
let totals_request_id = format!("usage_statistics_{}", Uuid::now_v7().simple());
let model_breakdown = client.fetch_daily_usage(
CodexRequestContext::auxiliary(
authorization.expose_secret(),
prepared.account.upstream_account_id(),
&model_request_id,
None,
),
CodexDailyUsageReport::PersonalModelCredits,
start_date,
end_date,
);
let daily_totals = client.fetch_daily_usage(
CodexRequestContext::auxiliary(
authorization.expose_secret(),
prepared.account.upstream_account_id(),
&totals_request_id,
None,
),
CodexDailyUsageReport::PersonalTokenTotals,
start_date,
end_date,
);
let (model_breakdown, daily_totals) = tokio::try_join!(model_breakdown, daily_totals)
.map_err(DailyUsageAttemptError::Upstream)?;
Ok(FetchedDailyUsage {
model_breakdown,
daily_totals,
})
}
fn map_statistics_quota_fetch_error(error: CodexQuotaFetchError) -> CodexUsageStatisticsError {
match error {
CodexQuotaFetchError::InvalidCredential => CodexUsageStatisticsError::InvalidCredentialData,
CodexQuotaFetchError::Recovery => CodexUsageStatisticsError::TransportUnavailable,
CodexQuotaFetchError::Upstream { account, error } => map_statistics_client_error(
error,
account.authentication_kind() == CODEX_AUTHENTICATION_KIND_OAUTH,
),
}
}
fn map_daily_usage_fetch_error(error: DailyUsageFetchError) -> CodexUsageStatisticsError {
match error {
DailyUsageFetchError::InvalidCredential => CodexUsageStatisticsError::InvalidCredentialData,
DailyUsageFetchError::Recovery => CodexUsageStatisticsError::TransportUnavailable,
DailyUsageFetchError::Upstream { prepared, error } => {
map_statistics_client_error(error, prepared.credential.authentication.oauth().is_some())
}
}
}
fn map_statistics_client_error(
error: CodexClientError,
credential_can_refresh: bool,
) -> CodexUsageStatisticsError {
match error {
CodexClientError::Upstream { status, body, .. }
if status == reqwest::StatusCode::UNAUTHORIZED && credential_can_refresh =>
{
CodexUsageStatisticsError::CredentialRefreshRequired {
upstream_body: Some(body),
}
}
CodexClientError::Upstream {
status,
body,
retry_after_seconds,
..
} => CodexUsageStatisticsError::Upstream {
status: status.as_u16(),
body,
retry_after_seconds,
},
_ => CodexUsageStatisticsError::TransportUnavailable,
}
}
async fn list_reset_credits_once(
client: &CodexBackendClient,
prepared: &PreparedCodexRuntimeCredential,
request_id: &str,
) -> Result<CodexRateLimitResetCredits, ResetCreditAttemptError> {
let authorization = prepared
.credential
.authentication
.authorization_header(Utc::now())
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
client
.list_rate_limit_reset_credits(CodexRequestContext::auxiliary(
authorization.expose_secret(),
prepared.account.upstream_account_id(),
request_id,
None,
))
.await
.map_err(ResetCreditAttemptError::Upstream)
}
async fn consume_reset_credit_once(
client: &CodexBackendClient,
prepared: &PreparedCodexRuntimeCredential,
request_id: &str,
credit_id: Option<&str>,
redeem_request_id: Uuid,
) -> Result<CodexRateLimitResetCreditsConsumeResult, ResetCreditAttemptError> {
let authorization = prepared
.credential
.authentication
.authorization_header(Utc::now())
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
client
.consume_rate_limit_reset_credit(
CodexRequestContext::auxiliary(
authorization.expose_secret(),
prepared.account.upstream_account_id(),
request_id,
None,
),
credit_id,
redeem_request_id,
)
.await
.map_err(ResetCreditAttemptError::Upstream)
}
fn map_reset_credit_attempt_error(
error: ResetCreditAttemptError,
prepared: &PreparedCodexRuntimeCredential,
consume: bool,
) -> CodexResetCreditsError {
match error {
ResetCreditAttemptError::InvalidCredential => CodexResetCreditsError::InvalidCredentialData,
ResetCreditAttemptError::Upstream(error) => {
map_reset_credit_client_error(error, prepared, consume)
}
}
}
fn map_reset_credit_client_error(
error: CodexClientError,
prepared: &PreparedCodexRuntimeCredential,
consume: bool,
) -> CodexResetCreditsError {
match error {
CodexClientError::Upstream {
status,
body,
diagnostics,
..
} if status == reqwest::StatusCode::UNAUTHORIZED
&& reset_credit_response_was_explicit_rejection(status, &diagnostics)
&& prepared.credential.authentication.oauth().is_some() =>
{
CodexResetCreditsError::CredentialRefreshRequired {
upstream_body: Some(body),
}
}
CodexClientError::Upstream {
status,
body,
retry_after_seconds,
diagnostics,
..
} => {
// 2xx 后发生的解码/响应体上限错误使用 synthetic 502 表示,但额度卡
// 可能已经消费。只有 transport 记录的真实非成功状态与错误状态一致时,
// 才能把它当作确定的上游拒绝并允许前端清除 pending 幂等键。
if consume && !reset_credit_response_was_explicit_rejection(status, &diagnostics) {
CodexResetCreditsError::ConsumeResultUnknown
} else {
CodexResetCreditsError::Upstream {
status: status.as_u16(),
body,
retry_after_seconds,
}
}
}
_ if consume => CodexResetCreditsError::ConsumeResultUnknown,
_ => CodexResetCreditsError::TransportUnavailable,
}
}
fn reset_credit_response_was_explicit_rejection(
status: reqwest::StatusCode,
diagnostics: &crate::transport::CodexUpstreamDiagnostics,
) -> bool {
!status.is_success() && diagnostics.status_code == Some(status.as_u16())
}
fn quota_refresh_allows_recovery(
@@ -0,0 +1,874 @@
use std::{
cmp::{Ordering, Reverse},
collections::BTreeMap,
};
use chrono::{DateTime, FixedOffset, NaiveDate, NaiveTime, TimeDelta, Utc};
use gateway_core::accounting::{Decimal, Money};
use serde_json::Value;
use crate::transport::{OpenAiBillingUsage, openai_aggregate_billing_breakdown};
const EPSILON: f64 = 1e-9;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum CodexUsageStatisticsServiceTier {
Standard,
Fast,
}
impl CodexUsageStatisticsServiceTier {
pub const fn as_str(self) -> &'static str {
match self {
Self::Standard => "standard",
Self::Fast => "fast",
}
}
const fn billing_service_tier(self) -> &'static str {
match self {
Self::Standard => "default",
Self::Fast => "fast",
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CodexUsageStatisticsTokens {
pub uncached_input: u64,
pub cached_input: u64,
pub output: u64,
pub total: u64,
}
impl CodexUsageStatisticsTokens {
fn checked_add(self, other: Self) -> Option<Self> {
Some(Self {
uncached_input: self.uncached_input.checked_add(other.uncached_input)?,
cached_input: self.cached_input.checked_add(other.cached_input)?,
output: self.output.checked_add(other.output)?,
total: self.total.checked_add(other.total)?,
})
}
const fn is_empty(self) -> bool {
self.total == 0 && self.uncached_input == 0 && self.cached_input == 0 && self.output == 0
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct CodexUsageStatisticsCycle {
pub offset: i8,
pub start_at: DateTime<Utc>,
pub end_at: DateTime<Utc>,
pub window_seconds: u64,
pub used_percent: Option<f64>,
pub is_current: bool,
pub can_go_previous: bool,
pub can_go_next: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct CodexUsageStatisticsModel {
pub key: String,
pub model: String,
pub service_tier: CodexUsageStatisticsServiceTier,
pub credit_share: Option<f64>,
pub quota_share: Option<f64>,
pub tokens: CodexUsageStatisticsTokens,
pub estimated_cost: Option<Money>,
pub has_unknown_pricing: bool,
pub has_estimated_allocation: bool,
pub has_rate_fallback: bool,
pub has_missing_token_data: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct CodexUsageStatisticsDay {
pub date: NaiveDate,
pub credit_share: Option<f64>,
pub tokens: CodexUsageStatisticsTokens,
pub estimated_cost: Option<Money>,
pub has_unknown_pricing: bool,
pub has_missing_token_data: bool,
pub is_boundary_day: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct CodexUsageStatisticsSummary {
pub tokens: CodexUsageStatisticsTokens,
pub estimated_cost: Option<Money>,
pub projected_tokens: Option<u64>,
pub projected_cost: Option<Money>,
pub day_count: u32,
pub has_unknown_pricing: bool,
pub has_missing_token_data: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct CodexUsageStatistics {
pub cycle: CodexUsageStatisticsCycle,
pub summary: CodexUsageStatisticsSummary,
pub models: Vec<CodexUsageStatisticsModel>,
pub daily: Vec<CodexUsageStatisticsDay>,
}
#[derive(Debug, Clone, Copy)]
pub struct CodexUsageStatisticsPeriod {
pub cycle_offset: i8,
pub start_at: DateTime<Utc>,
pub end_at: DateTime<Utc>,
pub window_seconds: u64,
pub query_start_date: NaiveDate,
pub query_end_date: NaiveDate,
}
pub fn statistics_period(
current_start_at: DateTime<Utc>,
current_end_at: DateTime<Utc>,
window_seconds: u64,
cycle_offset: i8,
timezone: FixedOffset,
now: DateTime<Utc>,
) -> Option<CodexUsageStatisticsPeriod> {
let window_delta = TimeDelta::seconds(i64::try_from(window_seconds).ok()?);
let shift = window_delta.checked_mul(i32::from(cycle_offset))?;
let start_at = current_start_at.checked_add_signed(shift)?;
let end_at = current_end_at.checked_add_signed(shift)?;
let query_start_date = start_at
.checked_sub_signed(TimeDelta::days(1))?
.with_timezone(&timezone)
.date_naive();
let effective_end = end_at.min(now);
let query_end_date = effective_end
.checked_add_signed(TimeDelta::days(1))?
.with_timezone(&timezone)
.date_naive();
Some(CodexUsageStatisticsPeriod {
cycle_offset,
start_at,
end_at,
window_seconds,
query_start_date,
query_end_date,
})
}
pub(super) struct BuildUsageStatistics<'a> {
pub period: CodexUsageStatisticsPeriod,
pub timezone: FixedOffset,
pub current_used_percent: Option<f64>,
pub now: DateTime<Utc>,
pub model_breakdown: &'a Value,
pub daily_totals: &'a Value,
pub max_cycle_offset: u8,
}
pub(super) fn build_usage_statistics(input: BuildUsageStatistics<'_>) -> CodexUsageStatistics {
let BuildUsageStatistics {
period,
timezone,
current_used_percent,
now,
model_breakdown,
daily_totals,
max_cycle_offset,
} = input;
let all_days = build_personal_days(daily_totals, model_breakdown);
let selected_days = select_cycle_days(
&all_days,
period.start_at,
period.end_at,
timezone,
now,
true,
);
let is_current = period.cycle_offset == 0;
let used_percent = is_current
.then_some(current_used_percent)
.flatten()
.filter(|value| value.is_finite())
.map(|value| value.clamp(0.0, 100.0));
let used_ratio = used_percent
.map(|value| value.clamp(0.0, 100.0) / 100.0)
.unwrap_or_default();
let summary = summarize(&selected_days, used_ratio, is_current);
let models = aggregate_models(&selected_days, used_ratio, is_current);
let total_credits = selected_days
.iter()
.map(|day| day.credit_weight)
.sum::<f64>();
let mut daily = selected_days
.into_iter()
.map(|day| CodexUsageStatisticsDay {
date: day.date,
credit_share: (total_credits > EPSILON).then_some(day.credit_weight / total_credits),
tokens: day.tokens,
estimated_cost: day.estimated_cost,
has_unknown_pricing: day.has_unknown_pricing,
has_missing_token_data: day.has_missing_token_data,
is_boundary_day: day.is_boundary_day,
})
.collect::<Vec<_>>();
daily.sort_by_key(|day| Reverse(day.date));
CodexUsageStatistics {
cycle: CodexUsageStatisticsCycle {
offset: period.cycle_offset,
start_at: period.start_at,
end_at: period.end_at,
window_seconds: period.window_seconds,
used_percent,
is_current,
can_go_previous: period.cycle_offset > -(max_cycle_offset as i8),
can_go_next: period.cycle_offset < 0,
},
summary,
models,
daily,
}
}
#[derive(Debug, Clone)]
struct DailyModelUsage {
model: String,
service_tier: CodexUsageStatisticsServiceTier,
credits: f64,
tokens: CodexUsageStatisticsTokens,
estimated_cost: Option<Money>,
has_unknown_pricing: bool,
has_estimated_allocation: bool,
has_rate_fallback: bool,
has_missing_token_data: bool,
}
#[derive(Debug, Clone)]
struct DailyUsage {
date: NaiveDate,
credit_weight: f64,
tokens: CodexUsageStatisticsTokens,
estimated_cost: Option<Money>,
has_unknown_pricing: bool,
has_missing_token_data: bool,
is_boundary_day: bool,
models: Vec<DailyModelUsage>,
}
#[derive(Default)]
struct MoneySum {
value: Option<Money>,
overflowed: bool,
}
impl MoneySum {
fn push(&mut self, value: Option<Money>) {
let Some(value) = value else {
return;
};
if self.overflowed {
return;
}
self.value = match self.value {
Some(current) => match current.checked_add(value) {
Some(total) => Some(total),
None => {
self.overflowed = true;
None
}
},
None => Some(value),
};
}
fn finish(self) -> Option<Money> {
(!self.overflowed).then_some(self.value).flatten()
}
}
fn build_personal_days(totals_payload: &Value, breakdown_payload: &Value) -> Vec<DailyUsage> {
let totals_by_date = rows_by_date(totals_payload);
let breakdown_by_date = rows_by_date(breakdown_payload);
let mut dates = totals_by_date
.keys()
.chain(breakdown_by_date.keys())
.copied()
.collect::<Vec<_>>();
dates.sort_unstable();
dates.dedup();
dates
.into_iter()
.filter_map(|date| {
let totals_row = totals_by_date.get(&date).unwrap_or(&Value::Null);
let breakdown_row = breakdown_by_date.get(&date).unwrap_or(&Value::Null);
let totals = token_parts(totals_row.get("totals").unwrap_or(totals_row));
let credit_weight = day_credit_weight(breakdown_row);
let weighted = weighted_models(breakdown_row);
if totals.is_empty() {
if credit_weight <= EPSILON {
return None;
}
let models = weighted
.into_iter()
.map(|model| DailyModelUsage {
model: model.model,
service_tier: model.service_tier,
credits: model.weight,
tokens: CodexUsageStatisticsTokens::default(),
estimated_cost: None,
has_unknown_pricing: false,
has_estimated_allocation: false,
has_rate_fallback: false,
has_missing_token_data: true,
})
.collect();
return Some(finalize_day(date, credit_weight, totals, models, true));
}
let has_model_breakdown = !weighted.is_empty();
let weighted = if has_model_breakdown {
weighted
} else {
vec![WeightedModel {
model: "unknown".to_owned(),
service_tier: CodexUsageStatisticsServiceTier::Standard,
weight: 1.0,
}]
};
let allocation = rate_adjusted_shares(&weighted, totals);
let uncached = allocate_integer_total(totals.uncached_input, &allocation.token_shares);
let cached = allocate_integer_total(totals.cached_input, &allocation.token_shares);
let output = allocate_integer_total(totals.output, &allocation.token_shares);
let total = allocate_integer_total(totals.total, &allocation.token_shares);
let estimated_allocation = weighted.len() > 1 || !has_model_breakdown;
let rate_fallback = allocation.has_rate_fallback || !has_model_breakdown;
let models = weighted
.into_iter()
.enumerate()
.map(|(index, model)| {
let tokens = CodexUsageStatisticsTokens {
uncached_input: uncached.get(index).copied().unwrap_or_default(),
cached_input: cached.get(index).copied().unwrap_or_default(),
output: output.get(index).copied().unwrap_or_default(),
total: total.get(index).copied().unwrap_or_default(),
};
let estimated_cost = aggregate_cost(&model.model, model.service_tier, tokens);
DailyModelUsage {
model: model.model,
service_tier: model.service_tier,
credits: if has_model_breakdown {
model.weight
} else {
0.0
},
tokens,
has_unknown_pricing: estimated_cost.is_none() && !tokens.is_empty(),
estimated_cost,
has_estimated_allocation: estimated_allocation,
has_rate_fallback: rate_fallback,
has_missing_token_data: false,
}
})
.collect();
Some(finalize_day(date, credit_weight, totals, models, false))
})
.collect()
}
fn finalize_day(
date: NaiveDate,
credit_weight: f64,
tokens: CodexUsageStatisticsTokens,
models: Vec<DailyModelUsage>,
has_missing_token_data: bool,
) -> DailyUsage {
let mut costs = MoneySum::default();
for model in &models {
costs.push(model.estimated_cost);
}
DailyUsage {
date,
credit_weight,
tokens,
estimated_cost: costs.finish(),
has_unknown_pricing: models.iter().any(|model| model.has_unknown_pricing),
has_missing_token_data,
is_boundary_day: false,
models,
}
}
#[derive(Debug)]
struct WeightedModel {
model: String,
service_tier: CodexUsageStatisticsServiceTier,
weight: f64,
}
struct Allocation {
token_shares: Vec<f64>,
has_rate_fallback: bool,
}
fn weighted_models(row: &Value) -> Vec<WeightedModel> {
value_array(row.get("models"))
.iter()
.filter_map(|value| {
let weight = non_negative_f64(value.get("credits"));
(weight > EPSILON).then(|| WeightedModel {
model: model_name(value),
service_tier: model_service_tier(value.get("speed").or_else(|| value.get("mode"))),
weight,
})
})
.collect()
}
fn rate_adjusted_shares(
models: &[WeightedModel],
totals: CodexUsageStatisticsTokens,
) -> Allocation {
let total_weight = models.iter().map(|model| model.weight).sum::<f64>();
let credit_shares = models
.iter()
.map(|model| model.weight / total_weight)
.collect::<Vec<_>>();
if models.len() == 1 {
return Allocation {
token_shares: vec![1.0],
has_rate_fallback: false,
};
}
let classified_tokens = totals
.uncached_input
.saturating_add(totals.cached_input)
.saturating_add(totals.output);
if classified_tokens == 0 {
return Allocation {
token_shares: credit_shares,
has_rate_fallback: true,
};
}
let rates = models
.iter()
.map(|model| blended_rate(&model.model, model.service_tier, totals, classified_tokens))
.collect::<Vec<_>>();
let known = rates
.iter()
.enumerate()
.filter_map(|(index, rate)| rate.map(|rate| (index, rate)))
.collect::<Vec<_>>();
if known.is_empty() {
return Allocation {
token_shares: credit_shares,
has_rate_fallback: true,
};
}
let known_weight = known
.iter()
.map(|(index, _)| models[*index].weight)
.sum::<f64>();
let fallback_rate = known
.iter()
.map(|(index, rate)| rate * models[*index].weight / known_weight)
.sum::<f64>();
let inverse = models
.iter()
.enumerate()
.map(|(index, model)| model.weight / rates[index].unwrap_or(fallback_rate))
.collect::<Vec<_>>();
let inverse_total = inverse.iter().sum::<f64>();
Allocation {
token_shares: inverse
.into_iter()
.map(|value| value / inverse_total)
.collect(),
has_rate_fallback: known.len() != models.len(),
}
}
#[allow(clippy::cast_precision_loss)]
fn blended_rate(
model: &str,
service_tier: CodexUsageStatisticsServiceTier,
totals: CodexUsageStatisticsTokens,
classified_tokens: u64,
) -> Option<f64> {
let cost = aggregate_cost(model, service_tier, totals)?;
let rate = cost.amount().scaled() as f64 / classified_tokens as f64;
(rate > EPSILON).then_some(rate)
}
fn aggregate_cost(
model: &str,
service_tier: CodexUsageStatisticsServiceTier,
tokens: CodexUsageStatisticsTokens,
) -> Option<Money> {
let input = tokens.uncached_input.checked_add(tokens.cached_input)?;
openai_aggregate_billing_breakdown(
model,
OpenAiBillingUsage::new(input, tokens.output, tokens.cached_input, 0),
Some(service_tier.billing_service_tier()),
)
.map(|breakdown| breakdown.total_amount())
}
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
fn allocate_integer_total(total: u64, shares: &[f64]) -> Vec<u64> {
if shares.is_empty() {
return Vec::new();
}
let share_total = shares.iter().copied().sum::<f64>();
if !share_total.is_finite() || share_total <= EPSILON {
return vec![0; shares.len()];
}
let raw = shares
.iter()
.map(|share| total as f64 * (share.max(0.0) / share_total))
.collect::<Vec<_>>();
let mut allocated = raw
.iter()
.map(|value| value.floor() as u64)
.collect::<Vec<_>>();
let allocated_total = allocated.iter().copied().sum::<u64>();
let mut remainder = total.saturating_sub(allocated_total);
let mut order = raw
.iter()
.enumerate()
.map(|(index, value)| (index, value - value.floor()))
.collect::<Vec<_>>();
order.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| left.0.cmp(&right.0))
});
let mut index = 0_usize;
while remainder > 0 {
let target = order[index % order.len()].0;
allocated[target] = allocated[target].saturating_add(1);
remainder -= 1;
index += 1;
}
allocated
}
fn select_cycle_days(
days: &[DailyUsage],
start_at: DateTime<Utc>,
end_at: DateTime<Utc>,
timezone: FixedOffset,
now: DateTime<Utc>,
mark_boundary: bool,
) -> Vec<DailyUsage> {
let local_start = start_at.with_timezone(&timezone);
let start_date = local_start.date_naive();
let effective_end = end_at.min(now);
let end_date = effective_end.with_timezone(&timezone).date_naive();
let exclude_end_date = end_at <= now && end_date > start_date;
let is_boundary = local_start.time() != NaiveTime::MIN;
days.iter()
.filter(|day| {
day.date >= start_date
&& if exclude_end_date {
day.date < end_date
} else {
day.date <= end_date
}
})
.cloned()
.map(|mut day| {
day.is_boundary_day = mark_boundary && is_boundary && day.date == start_date;
day
})
.collect()
}
fn summarize(
days: &[DailyUsage],
used_ratio: f64,
is_current: bool,
) -> CodexUsageStatisticsSummary {
let tokens = sum_tokens(days.iter().map(|day| day.tokens));
let mut costs = MoneySum::default();
for day in days {
costs.push(day.estimated_cost);
}
let estimated_cost = costs.finish();
let has_unknown_pricing = days.iter().any(|day| day.has_unknown_pricing);
let projected_tokens = (is_current && used_ratio > EPSILON && tokens.total > 0)
.then(|| project_u64(tokens.total, used_ratio))
.flatten();
let projected_cost = (is_current && !has_unknown_pricing && used_ratio > EPSILON)
.then(|| estimated_cost.and_then(|cost| project_money(cost, used_ratio)))
.flatten();
CodexUsageStatisticsSummary {
tokens,
estimated_cost,
projected_tokens,
projected_cost,
day_count: u32::try_from(days.len()).unwrap_or(u32::MAX),
has_unknown_pricing,
has_missing_token_data: days.iter().any(|day| day.has_missing_token_data),
}
}
#[derive(Default)]
struct ModelAggregate {
credits: f64,
tokens: CodexUsageStatisticsTokens,
costs: MoneySum,
has_unknown_pricing: bool,
has_estimated_allocation: bool,
has_rate_fallback: bool,
has_missing_token_data: bool,
}
fn aggregate_models(
days: &[DailyUsage],
used_ratio: f64,
is_current: bool,
) -> Vec<CodexUsageStatisticsModel> {
let mut aggregates =
BTreeMap::<(String, CodexUsageStatisticsServiceTier), ModelAggregate>::new();
for model in days.iter().flat_map(|day| &day.models) {
if model.tokens.is_empty() && model.credits <= EPSILON {
continue;
}
let aggregate = aggregates
.entry((model.model.clone(), model.service_tier))
.or_default();
aggregate.credits += model.credits;
aggregate.tokens =
aggregate
.tokens
.checked_add(model.tokens)
.unwrap_or(CodexUsageStatisticsTokens {
uncached_input: u64::MAX,
cached_input: u64::MAX,
output: u64::MAX,
total: u64::MAX,
});
aggregate.costs.push(model.estimated_cost);
aggregate.has_unknown_pricing |= model.has_unknown_pricing;
aggregate.has_estimated_allocation |= model.has_estimated_allocation;
aggregate.has_rate_fallback |= model.has_rate_fallback;
aggregate.has_missing_token_data |= model.has_missing_token_data;
}
let total_credits = aggregates
.values()
.map(|aggregate| aggregate.credits)
.sum::<f64>();
let mut rows = aggregates
.into_iter()
.map(|((model, service_tier), aggregate)| {
let credit_share =
(total_credits > EPSILON).then_some(aggregate.credits / total_credits);
let quota_share = (is_current && used_ratio > EPSILON)
.then(|| credit_share.map(|share| share * used_ratio))
.flatten();
CodexUsageStatisticsModel {
key: format!("{model}::{}", service_tier.as_str()),
model,
service_tier,
credit_share,
quota_share,
tokens: aggregate.tokens,
estimated_cost: aggregate.costs.finish(),
has_unknown_pricing: aggregate.has_unknown_pricing,
has_estimated_allocation: aggregate.has_estimated_allocation,
has_rate_fallback: aggregate.has_rate_fallback,
has_missing_token_data: aggregate.has_missing_token_data,
}
})
.collect::<Vec<_>>();
rows.sort_by(compare_model_rows);
rows
}
fn compare_model_rows(
left: &CodexUsageStatisticsModel,
right: &CodexUsageStatisticsModel,
) -> Ordering {
right
.credit_share
.unwrap_or_default()
.total_cmp(&left.credit_share.unwrap_or_default())
.then_with(|| cost_ticks(right.estimated_cost).cmp(&cost_ticks(left.estimated_cost)))
.then_with(|| right.tokens.total.cmp(&left.tokens.total))
.then_with(|| left.model.cmp(&right.model))
.then_with(|| left.service_tier.cmp(&right.service_tier))
}
fn cost_ticks(cost: Option<Money>) -> u128 {
cost.map(|value| value.amount().scaled())
.unwrap_or_default()
}
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
fn project_u64(value: u64, used_ratio: f64) -> Option<u64> {
let projected = value as f64 / used_ratio;
(projected.is_finite() && projected >= 0.0 && projected <= u64::MAX as f64)
.then(|| projected.round() as u64)
}
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
fn project_money(cost: Money, used_ratio: f64) -> Option<Money> {
let projected = cost.amount().scaled() as f64 / used_ratio;
if !projected.is_finite() || projected < 0.0 || projected > u128::MAX as f64 {
return None;
}
let amount = Decimal::from_scaled(projected.round() as u128).ok()?;
Some(Money::new(amount, cost.currency()))
}
fn sum_tokens(
values: impl IntoIterator<Item = CodexUsageStatisticsTokens>,
) -> CodexUsageStatisticsTokens {
values
.into_iter()
.fold(CodexUsageStatisticsTokens::default(), |total, value| {
total
.checked_add(value)
.unwrap_or(CodexUsageStatisticsTokens {
uncached_input: u64::MAX,
cached_input: u64::MAX,
output: u64::MAX,
total: u64::MAX,
})
})
}
fn rows_by_date(payload: &Value) -> BTreeMap<NaiveDate, Value> {
extract_daily_list(payload)
.into_iter()
.filter_map(|row| Some((date_value(row.get("date")?)?, row.clone())))
.collect()
}
fn extract_daily_list(payload: &Value) -> Vec<&Value> {
if let Some(rows) = payload.as_array() {
return rows.iter().collect();
}
[
"data",
"items",
"results",
"daily",
"daily_usage",
"dailyWorkspaceUsageCounts",
"daily_workspace_usage_counts",
"workspace_usage_counts",
]
.into_iter()
.find_map(|key| payload.get(key).and_then(Value::as_array))
.map(|rows| rows.iter().collect())
.unwrap_or_default()
}
fn value_array(value: Option<&Value>) -> &[Value] {
value.and_then(Value::as_array).map_or(&[], Vec::as_slice)
}
fn token_parts(value: &Value) -> CodexUsageStatisticsTokens {
let uncached_input = non_negative_u64(value.get("uncached_text_input_tokens"));
let cached_input = non_negative_u64(value.get("cached_text_input_tokens"));
let output = non_negative_u64(value.get("text_output_tokens"));
let derived = uncached_input
.saturating_add(cached_input)
.saturating_add(output);
let reported = non_negative_u64(value.get("text_total_tokens"));
CodexUsageStatisticsTokens {
uncached_input,
cached_input,
output,
total: if reported > 0 { reported } else { derived },
}
}
fn day_credit_weight(value: &Value) -> f64 {
let model_credits = value_array(value.get("models"))
.iter()
.map(|model| non_negative_f64(model.get("credits")))
.sum::<f64>();
if model_credits > EPSILON {
return model_credits;
}
value
.get("product_surface_usage_values")
.and_then(Value::as_object)
.map(|values| {
values
.values()
.map(|value| non_negative_f64(Some(value)))
.sum()
})
.unwrap_or_default()
}
fn model_name(value: &Value) -> String {
["model", "model_id", "model_name", "name", "id"]
.into_iter()
.find_map(|key| value.get(key).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("unknown")
.to_ascii_lowercase()
}
fn model_service_tier(value: Option<&Value>) -> CodexUsageStatisticsServiceTier {
match value
.and_then(Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.as_deref()
{
Some("fast" | "priority") => CodexUsageStatisticsServiceTier::Fast,
Some(_) | None => CodexUsageStatisticsServiceTier::Standard,
}
}
fn date_value(value: &Value) -> Option<NaiveDate> {
NaiveDate::parse_from_str(value.as_str()?.get(..10)?, "%Y-%m-%d").ok()
}
fn non_negative_u64(value: Option<&Value>) -> u64 {
let Some(value) = value else {
return 0;
};
if let Some(number) = value.as_u64() {
return number;
}
let normalized = value
.as_str()
.map(|raw| raw.replace(',', ""))
.unwrap_or_else(|| value.to_string());
normalized
.trim()
.parse::<u64>()
.ok()
.or_else(|| {
normalized
.trim()
.parse::<f64>()
.ok()
.filter(|number| number.is_finite() && *number >= 0.0)
.map(|number| number.round() as u64)
})
.unwrap_or_default()
}
fn non_negative_f64(value: Option<&Value>) -> f64 {
let Some(value) = value else {
return 0.0;
};
let parsed = value.as_f64().or_else(|| {
value
.as_str()
.map(|raw| raw.replace(',', ""))
.and_then(|raw| raw.trim().parse::<f64>().ok())
});
parsed
.filter(|number| number.is_finite() && *number >= 0.0)
.unwrap_or_default()
}
@@ -95,6 +95,17 @@ pub struct SelectCodexCredential<'a> {
pub session_affinity_key: Option<&'a ProviderSessionAffinityKey>,
}
pub(crate) struct SelectCodexProviderEndpointCredential<'a> {
pub request_url: &'a Url,
pub attempt: &'a AttemptContext,
}
struct CredentialSelectionInput<'a> {
request_url: &'a Url,
attempt: &'a AttemptContext,
session_affinity_key: Option<&'a ProviderSessionAffinityKey>,
}
#[derive(Clone)]
pub(crate) struct CodexCyberPolicyScope {
key: ProviderSessionAffinityKey,
@@ -240,6 +251,12 @@ struct AffinityTelemetry {
account_switch: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ModelCatalogEligibility<'a> {
Required(&'a str),
NotApplicable,
}
impl CodexCredentialSelector {
#[must_use]
// 选择器显式持有各能力边界,避免把 Provider 私有服务重新包装成通用容器。
@@ -275,13 +292,59 @@ impl CodexCredentialSelector {
&self,
request: &SelectCodexCredential<'_>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
self.select_with_cyber_policy(request, None).await
let input = CredentialSelectionInput {
request_url: request.request_url,
attempt: request.attempt,
session_affinity_key: request.session_affinity_key,
};
self.select_inner(
&input,
None,
ModelCatalogEligibility::Required(request.upstream_model),
)
.await
}
pub(crate) async fn select_with_cyber_policy(
&self,
request: &SelectCodexCredential<'_>,
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let input = CredentialSelectionInput {
request_url: request.request_url,
attempt: request.attempt,
session_affinity_key: request.session_affinity_key,
};
self.select_inner(
&input,
cyber_policy_session_key,
ModelCatalogEligibility::Required(request.upstream_model),
)
.await
}
/// 为不属于 Responses 文本模型目录的 Provider 原生端点选择账号。
///
/// 账号范围、健康度、配额、并发租约、cookie 与认证准备仍走同一套选择链路;
/// 唯一差异是不能用文本模型目录成员关系淘汰账号。
pub(crate) async fn select_for_provider_endpoint(
&self,
request: &SelectCodexProviderEndpointCredential<'_>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let input = CredentialSelectionInput {
request_url: request.request_url,
attempt: request.attempt,
session_affinity_key: None,
};
self.select_inner(&input, None, ModelCatalogEligibility::NotApplicable)
.await
}
async fn select_inner(
&self,
request: &CredentialSelectionInput<'_>,
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
model_catalog_eligibility: ModelCatalogEligibility<'_>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let diagnostic = request.attempt.is_diagnostic_required_account();
let accounts = self.repository.list_for_provider().await?;
@@ -294,12 +357,15 @@ impl CodexCredentialSelector {
.attempt
.account_scope()
.is_some_and(|scope| scope.allows(account.id())))
&& (diagnostic || {
let observed_support = self
.catalog
.observed_model_support(account, request.upstream_model);
matches!(observed_support, Ok(None | Some(true)))
})
&& (diagnostic
|| match model_catalog_eligibility {
ModelCatalogEligibility::NotApplicable => true,
ModelCatalogEligibility::Required(upstream_model) => {
let observed_support =
self.catalog.observed_model_support(account, upstream_model);
matches!(observed_support, Ok(None | Some(true)))
}
})
})
.collect::<Vec<_>>();
if !diagnostic {
@@ -410,14 +476,16 @@ impl CodexCredentialSelector {
{
affinity.escape(AffinityEscapeReason::PinnedAccount);
}
let preferred = pinned_account
.clone()
.or_else(|| affinity.preferred_account().cloned());
let mut observed_affinity_account = affinity.bound_account().cloned();
let mut shortest_retry = None;
let policy = request.attempt.account_selection_policy();
loop {
let preferred = pinned_account
.clone()
.or_else(|| affinity.preferred_account().cloned());
let context = AccountSelectionContext {
policy: request.attempt.account_selection_policy(),
policy,
now: SystemTime::now(),
excluded_accounts: excluded.clone(),
preferred_account: preferred.clone(),
@@ -445,7 +513,6 @@ impl CodexCredentialSelector {
.map(|candidate| candidate.account.clone())
.ok_or(CredentialSelectionError::InvalidCredential)?;
let allows_account_state_mutation = !diagnostic || account.enabled();
let policy = request.attempt.account_selection_policy();
match self
.leases
.try_acquire(ProviderLeaseRequest::Scheduling(
@@ -466,6 +533,40 @@ impl CodexCredentialSelector {
excluded.insert(account.id().clone());
}
ProviderLeaseAcquisition::Acquired(guard) => {
let initial_affinity_claim = if !diagnostic
&& observed_affinity_account.is_none()
&& let Some(key) = request.session_affinity_key
{
self.claim_initial_session_affinity(key, account.id()).await
} else {
None
};
// 原生 continuation/required account 比根亲和绑定更严格;它可以
// 使用 owner 账号,但不能把已经迁移的根会话拉回旧号。
if pinned_account.is_none()
&& let Some(effective_account) = initial_affinity_claim
&& &effective_account != account.id()
{
drop(guard);
observed_affinity_account = Some(effective_account.clone());
affinity = affinity_selection_for_bound_account(
effective_account,
&candidates,
SystemTime::now(),
);
continue;
}
let affinity_expected_account_id = if pinned_account
.as_ref()
.zip(observed_affinity_account.as_ref())
.is_some_and(|(pinned, bound)| pinned != bound)
{
account.id().clone()
} else {
observed_affinity_account
.clone()
.unwrap_or_else(|| account.id().clone())
};
let affinity_telemetry = affinity.telemetry(account.id());
tracing::info!(
request_id = %request.attempt.request_id(),
@@ -517,6 +618,7 @@ impl CodexCredentialSelector {
cyber_policy_scope,
allows_account_state_mutation,
affinity_telemetry,
affinity_expected_account_id,
_guard: guard,
});
}
@@ -537,24 +639,7 @@ impl CodexCredentialSelector {
else {
return AffinitySelection::default();
};
let Some(candidate) = candidates
.iter()
.find(|candidate| candidate.account.id() == &account_id)
else {
self.clear_session_affinity(key).await;
return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable);
};
if candidate
.account
.status_projection(now, candidate.signals.rate_limited_until)
.status
== AccountStatus::Normal
{
return AffinitySelection::preferred(account_id);
}
let reason = affinity_unavailable_reason(candidate, now);
self.clear_session_affinity(key).await;
AffinitySelection::escaped(account_id, reason)
affinity_selection_for_bound_account(account_id, candidates, now)
}
async fn lookup_session_affinity(
@@ -583,22 +668,38 @@ impl CodexCredentialSelector {
}
}
async fn clear_session_affinity(&self, key: &ProviderSessionAffinityKey) {
async fn claim_initial_session_affinity(
&self,
key: &ProviderSessionAffinityKey,
selected_account_id: &ProviderAccountId,
) -> Option<ProviderAccountId> {
match tokio::time::timeout(
SESSION_AFFINITY_TIMEOUT,
self.session_affinity.clear(&self.provider_kind, key),
self.session_affinity.claim_or_load(
&self.provider_kind,
key,
selected_account_id,
SESSION_AFFINITY_TTL,
),
)
.await
{
Ok(Ok(_)) => {}
Ok(Ok(account_id)) => Some(account_id),
Ok(Err(error)) => {
tracing::warn!(error = %error, "OpenAI stale session affinity clear failed open");
tracing::warn!(
account_id = %selected_account_id,
error = %error,
"OpenAI initial session affinity claim failed open"
);
None
}
Err(_) => {
tracing::warn!(
account_id = %selected_account_id,
timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(),
"OpenAI stale session affinity clear timed out"
"OpenAI initial session affinity claim timed out"
);
None
}
}
}
@@ -835,6 +936,7 @@ impl CodexCredentialSelector {
&self,
account: &ProviderAccount,
session_affinity_key: Option<&ProviderSessionAffinityKey>,
expected_affinity_account_id: &ProviderAccountId,
) {
self.restore_recoverable_account_state(account).await;
self.risk_recovery
@@ -844,24 +946,26 @@ impl CodexCredentialSelector {
let Some(key) = session_affinity_key else {
return;
};
if matches!(
self.lookup_session_affinity(key).await,
SessionAffinityLookup::Unavailable
) {
return;
}
match tokio::time::timeout(
SESSION_AFFINITY_TIMEOUT,
self.session_affinity.bind(
self.session_affinity.compare_and_bind(
&self.provider_kind,
key,
expected_affinity_account_id,
account.id(),
SESSION_AFFINITY_TTL,
),
)
.await
{
Ok(Ok(())) => {}
Ok(Ok(effective_account)) if &effective_account == account.id() => {}
Ok(Ok(effective_account)) => {
tracing::debug!(
account_id = %account.id(),
effective_account_id = %effective_account,
"OpenAI late success left the newer session affinity binding unchanged"
);
}
Ok(Err(error)) => {
tracing::warn!(
account_id = %account.id(),
@@ -1038,6 +1142,32 @@ impl CodexCredentialSelector {
}
}
fn affinity_selection_for_bound_account(
account_id: ProviderAccountId,
candidates: &[AccountCandidate],
now: SystemTime,
) -> AffinitySelection {
let Some(candidate) = candidates
.iter()
.find(|candidate| candidate.account.id() == &account_id)
else {
return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable);
};
match candidate
.account
.status_projection(now, candidate.signals.rate_limited_until)
.status
{
AccountStatus::Normal => AffinitySelection::preferred(account_id),
AccountStatus::QuotaExhausted
| AccountStatus::RateLimited
| AccountStatus::Disabled
| AccountStatus::Error => {
AffinitySelection::escaped(account_id, affinity_unavailable_reason(candidate, now))
}
}
}
fn affinity_unavailable_reason(
candidate: &AccountCandidate,
now: SystemTime,
@@ -1093,6 +1223,7 @@ pub struct CodexCredentialLease {
cyber_policy_scope: Option<CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
affinity_telemetry: AffinityTelemetry,
affinity_expected_account_id: ProviderAccountId,
_guard: Box<dyn ProviderLeaseGuard>,
}
@@ -1150,6 +1281,12 @@ impl CodexCredentialLease {
pub const fn account_switch(&self) -> bool {
self.affinity_telemetry.account_switch
}
/// 成功反馈只能从选择时观察到的绑定迁移,避免迟到请求覆盖较新 winner。
#[must_use]
pub(crate) const fn affinity_expected_account_id(&self) -> &ProviderAccountId {
&self.affinity_expected_account_id
}
}
impl fmt::Debug for CodexCredentialLease {
+337 -26
View File
@@ -23,12 +23,17 @@ use gateway_core::error::{
OpaqueUpstreamValue, ProviderError, ProviderErrorKind,
};
use gateway_core::event::{
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseMetadata,
ProviderResponseObservation, ProviderResponseTimings, UpstreamHttpVersion, WebSocketPoolKind,
FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader,
ProviderResponseMetadata, ProviderResponseObservation, ProviderResponseTimings, ResponseMeta,
UpstreamHttpVersion, WebSocketPoolKind,
};
use gateway_core::operation::{
GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProviderSessionState,
};
use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProviderSessionState};
use gateway_core::provider_ports::ProviderSessionAffinityKey;
use gateway_core::routing::{ModelCapabilities, ModelPresentation, ProviderKind, UpstreamModelId};
use gateway_core::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId,
};
use gateway_core::task::{
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration,
@@ -48,8 +53,8 @@ use crate::credential::{
CodexCredentialCatalogService, CodexCredentialLease, CodexCredentialQuotaService,
CodexCredentialRefreshOutcome, CodexCredentialRefreshService, CodexCredentialSelector,
CodexCyberPolicyScope, CodexQuotaRefreshPolicy, CredentialSelectionError, RuntimeCodexCookie,
SelectCodexCredential, derive_codex_cyber_policy_session_key,
derive_codex_session_affinity_key,
SelectCodexCredential, SelectCodexProviderEndpointCredential,
derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key,
};
use crate::transport::canonical::{
CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome,
@@ -66,6 +71,9 @@ use crate::transport::profile::{
use crate::transport::protocol::responses::{
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
};
use crate::transport::protocol::websocket::{
PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
};
use crate::transport::request::{
CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account,
};
@@ -73,7 +81,8 @@ use crate::transport::session::CodexSessionIdentity;
use crate::transport::usage::normalize_service_tier;
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
use crate::transport::{
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
CODEX_IMAGE_EDITS_PATH, CODEX_IMAGE_GENERATIONS_PATH, CODEX_RESPONSES_PATH,
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates,
CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexUpstreamDiagnostics,
CodexWebSocketPool, endpoint_url,
@@ -81,6 +90,7 @@ use crate::transport::{
const PROVIDER_NAME: &str = "openai";
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const HTTP_JSON_TRANSPORT: &str = "http_json";
const WEBSOCKET_TRANSPORT: &str = "websocket";
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
/// 提交边界前最多保留 64 KiB 原始上游 chunk;达到阈值后结束无感换号窗口,
@@ -115,6 +125,8 @@ pub struct CodexProvider {
account_feedback: Arc<AccountFeedbackStats>,
client: CodexBackendClient,
responses_url: Url,
image_generations_url: Url,
image_edits_url: Url,
session_identity: Option<CodexSessionIdentity>,
}
@@ -134,6 +146,11 @@ impl CodexProvider {
) -> Result<Self, CodexProviderConfigError> {
let responses_url = Url::parse(&endpoint_url(&base_url, CODEX_RESPONSES_PATH))
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
let image_generations_url =
Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_GENERATIONS_PATH))
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
let image_edits_url = Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_EDITS_PATH))
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
let client =
CodexBackendClient::new(http, base_url, profile).with_websocket_pool(websocket_pool);
Ok(Self {
@@ -144,6 +161,8 @@ impl CodexProvider {
account_feedback,
client,
responses_url,
image_generations_url,
image_edits_url,
session_identity: None,
})
}
@@ -231,17 +250,25 @@ impl Provider for CodexProvider {
UpstreamSendState::NotSent,
));
}
if let Operation::GenerateImage(image) = request.operation() {
return self.execute_image(image, candidate, context).await;
}
let Operation::Generate(generate) = request.operation() else {
return Err(provider_error(
ProviderErrorKind::Unsupported,
UpstreamSendState::NotSent,
));
};
let Some(upstream_model) = candidate.upstream_model() else {
return Err(provider_error(
ProviderErrorKind::Protocol,
UpstreamSendState::NotSent,
));
};
let previous_session = decode_openai_session_state(generate);
let continuation_requested = generate.native_continuation_requested();
let mut upstream_request =
encode_generate_request(generate, candidate.upstream_model().as_str())
.map_err(map_request_error)?;
let mut upstream_request = encode_generate_request(generate, upstream_model.as_str())
.map_err(map_request_error)?;
if let Some(conversation_id) = previous_session
.as_ref()
.and_then(|state| state.conversation_id.as_ref())
@@ -272,7 +299,7 @@ impl Provider for CodexProvider {
.selector
.select_with_cyber_policy(
&SelectCodexCredential {
upstream_model: candidate.upstream_model().as_str(),
upstream_model: upstream_model.as_str(),
request_url: &self.responses_url,
attempt: &context,
session_affinity_key: session_affinity_key.as_ref(),
@@ -294,8 +321,14 @@ impl Provider for CodexProvider {
.is_some_and(|state| state.account_id != lease.account_id().as_str());
let has_explicit_state_owner =
context.account_state_owner().is_some() || previous_session.is_some();
let cross_account =
state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch());
let account_scope =
if state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch()) {
RequestAccountScope::Different
} else if has_explicit_state_owner || lease.affinity_hit() {
RequestAccountScope::Same
} else {
RequestAccountScope::Unknown
};
if context.continuation_attempt() != ContinuationAttempt::None
&& let Some(continuation) = context.continuation()
{
@@ -342,13 +375,11 @@ impl Provider for CodexProvider {
}
}
}
let account_scope = if !cross_account {
RequestAccountScope::SameAccount
} else if upstream_request.previous_response_id().is_some() {
RequestAccountScope::CrossAccountContinuationProbe
} else {
RequestAccountScope::CrossAccountFullRequest
};
if upstream_request.previous_response_id().is_some()
&& !account_scope.can_reuse_account_state()
{
return Err(continuation_replay_required_error());
}
scope_request_to_account(
&mut upstream_request,
lease.installation_id(),
@@ -356,7 +387,7 @@ impl Provider for CodexProvider {
);
let metadata = ProviderCallMetadata::new(
provider_kind,
candidate.upstream_model().clone(),
upstream_model.clone(),
ProviderResource::Account {
id: lease.account_id().clone(),
revision: lease.account().revision(),
@@ -380,7 +411,7 @@ impl Provider for CodexProvider {
client: self.client.clone(),
response_origin: self.responses_url.clone(),
request: upstream_request,
upstream_model: candidate.upstream_model().clone(),
upstream_model: upstream_model.clone(),
transport_policy: transport,
context,
selector: Arc::clone(&self.selector),
@@ -401,6 +432,76 @@ impl Provider for CodexProvider {
}
}
impl CodexProvider {
async fn execute_image(
&self,
image: &ImageRequest,
candidate: &ProviderCandidate,
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
if image.payload().protocol() != PROVIDER_NAME || candidate.upstream_model().is_some() {
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
));
}
let (endpoint_path, response_origin) = match image.kind() {
ImageRequestKind::Generation => (
CODEX_IMAGE_GENERATIONS_PATH,
self.image_generations_url.clone(),
),
ImageRequestKind::Edit => (CODEX_IMAGE_EDITS_PATH, self.image_edits_url.clone()),
};
let lease = self
.selector
.select_for_provider_endpoint(&SelectCodexProviderEndpointCredential {
request_url: &response_origin,
attempt: &context,
})
.await
.map_err(map_selection_error)?;
let lease = Arc::new(lease);
let allows_account_state_mutation = lease.allows_account_state_mutation();
let provider_kind = ProviderKind::new(PROVIDER_NAME)
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
let metadata = ProviderCallMetadata::for_provider_endpoint(
provider_kind,
ProviderResource::Account {
id: lease.account_id().clone(),
revision: lease.account().revision(),
},
UpstreamTransport::new(HTTP_JSON_TRANSPORT).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
})?,
);
let image_turn_id = image
.payload()
.context()
.get("image_turn_id")
.and_then(Value::as_str)
.map(str::to_owned);
let events = cold_json_response_stream(ColdJsonResponse {
client: self.client.clone(),
response_origin,
endpoint_path,
body: image.payload().body().clone(),
image_turn_id,
context,
selector: Arc::clone(&self.selector),
quota: Arc::clone(&self.quota),
agent_identity: Arc::clone(&self.agent_identity),
lease: Arc::clone(&lease),
output_started_at: Instant::now(),
});
let stream = ProviderStream::new(metadata, events, lease);
Ok(if allows_account_state_mutation {
stream.with_account_feedback(Arc::clone(&self.account_feedback))
} else {
stream
})
}
}
struct ColdResponse {
client: CodexBackendClient,
response_origin: Url,
@@ -418,6 +519,20 @@ struct ColdResponse {
session_capture: Option<OpenAiSessionCapture>,
}
struct ColdJsonResponse {
client: CodexBackendClient,
response_origin: Url,
endpoint_path: &'static str,
body: Bytes,
image_turn_id: Option<String>,
context: AttemptContext,
selector: Arc<CodexCredentialSelector>,
quota: Arc<CodexCredentialQuotaService>,
agent_identity: Arc<CodexAgentIdentityTaskService>,
lease: Arc<CodexCredentialLease>,
output_started_at: Instant,
}
#[derive(Clone, Serialize, Deserialize)]
struct OpenAiSessionState {
account_id: String,
@@ -569,6 +684,178 @@ fn map_handshake_attempt_error(error: CodexHandshakeAttemptError) -> MappedProvi
}
}
async fn create_json_attempt(
request: &ColdJsonResponse,
account: &ProviderAccount,
installation_id: &str,
authorization: &SecretString,
cookie_header: Option<&SecretString>,
account_selection: CodexAccountSelectionTelemetry<'_>,
) -> Result<CodexBackendJsonResponse, CodexHandshakeAttemptError> {
let Some(handshake_deadline) = remaining(request.context.deadline()) else {
return Err(CodexHandshakeAttemptError::Timeout);
};
let request_id = request.context.request_id().as_str();
let mut request_context = CodexRequestContext::auxiliary(
authorization.expose_secret(),
account.upstream_account_id(),
request_id,
Some(installation_id),
);
request_context.cookie_header = cookie_header.map(ExposeSecret::expose_secret);
request_context.account_selection = account_selection;
tokio::select! {
biased;
_ = request.context.cancellation().cancelled() => Err(CodexHandshakeAttemptError::Cancelled),
_ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout),
response = request.client.post_raw_json(
request.endpoint_path,
request.body.clone(),
request.image_turn_id.as_deref(),
request_context,
) => response.map_err(CodexHandshakeAttemptError::Client),
}
}
fn cold_json_response_stream(request: ColdJsonResponse) -> EventStream {
Box::pin(async_stream::try_stream! {
let allows_account_state_mutation = request.lease.allows_account_state_mutation();
let failure_context = OpenAiFailureContext {
client: &request.client,
selector: &request.selector,
quota: &request.quota,
response_origin: &request.response_origin,
cyber_policy_scope: None,
allows_account_state_mutation,
};
let mut active_account = request.lease.account().clone();
let cookie_header = build_cookie_header(request.lease.cookies())?;
let mut authorization = request
.lease
.authentication()
.authorization_header(chrono::Utc::now())
.map_err(|_| {
provider_error(
ProviderErrorKind::Unauthorized,
UpstreamSendState::NotSent,
)
})?;
let account_selection = CodexAccountSelectionTelemetry::new(
request.lease.affinity_hit(),
request.lease.escape_reason(),
request.lease.account_switch(),
);
let mut response = create_json_attempt(
&request,
&active_account,
request.lease.installation_id(),
&authorization,
cookie_header.as_ref(),
account_selection,
)
.await;
if allows_account_state_mutation
&& let Err(CodexHandshakeAttemptError::Client(error)) = &response
{
match request
.agent_identity
.recover_after_rejected_task(
active_account.id(),
request.lease.authentication(),
error,
)
.await
{
Ok(Some(recovered)) => {
authorization = recovered
.credential
.authentication
.authorization_header(chrono::Utc::now())
.map_err(|_| {
provider_error(
ProviderErrorKind::Unauthorized,
UpstreamSendState::NotSent,
)
})?;
active_account = recovered.account;
response = create_json_attempt(
&request,
&active_account,
&recovered.credential.installation_id,
&authorization,
cookie_header.as_ref(),
account_selection,
)
.await;
}
Ok(None) => {}
Err(_) => response = Err(CodexHandshakeAttemptError::AgentIdentityRecovery),
}
}
let response = match response.map_err(map_handshake_attempt_error) {
Ok(response) => response,
Err(mut failure) => {
if let Some(observation) = failure.observation.take() {
yield ProviderEvent::observation(observation);
}
apply_failure(&failure_context, &active_account, &failure).await;
Err(failure.error)?;
return;
}
};
let mut metrics = response.transport_metrics.clone();
metrics.first_event_ms = Some(
i64::try_from(request.output_started_at.elapsed().as_millis()).unwrap_or(i64::MAX),
);
if let Some(observation) = codex_response_observation(
CodexBackendTransport::HttpJson,
&response.diagnostics,
&response.response_metadata,
&metrics,
None,
openai_response_timings(&metrics, &response.response_metadata),
) {
yield ProviderEvent::observation(observation);
}
if allows_account_state_mutation {
synchronize_passive_quota_headers(
&request.quota,
&active_account,
&response.rate_limit_headers,
)
.await;
if !response.set_cookie_headers.is_empty()
&& let Err(error) = request
.selector
.capture_response_cookies(
&active_account,
&request.response_origin,
&response.set_cookie_headers,
)
.await
{
tracing::warn!(
account_id = %active_account.id(),
error = %error,
"Failed to persist OpenAI image response cookies"
);
}
}
let response_meta =
ResponseMeta::for_provider_endpoint(request.context.request_id().as_str());
yield ProviderEvent::canonical(GatewayEvent::Started(response_meta.clone()));
let wire = ProtocolWireEvent::raw_json(PROVIDER_NAME, response.body).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
})?;
yield ProviderEvent::wire(wire);
yield ProviderEvent::canonical(GatewayEvent::Completed(
response_meta.with_finish_reason(FinishReason::Stop),
));
})
}
fn cold_response_stream(response: ColdResponse) -> EventStream {
let ColdResponse {
client,
@@ -898,7 +1185,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
// 完成事件一旦交给下游,Core 可以立刻停止轮询 Provider stream;
// 在此之前持久化亲和关系,保证成功请求不会因流被提前 drop 而丢失绑定。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.record_success(
&active_account,
session_affinity_key.as_ref(),
lease.affinity_expected_account_id(),
)
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
@@ -1002,7 +1293,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 同上:尾部 finish() 也可能产出 completed,亲和记录必须先于任何下游 yield。
selector
.record_success(&active_account, session_affinity_key.as_ref())
.record_success(
&active_account,
session_affinity_key.as_ref(),
lease.affinity_expected_account_id(),
)
.await;
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
@@ -1603,6 +1898,7 @@ async fn synchronize_passive_quota_headers(
const fn actual_transport_name(transport: CodexBackendTransport) -> &'static str {
match transport {
CodexBackendTransport::HttpSse => HTTP_SSE_TRANSPORT,
CodexBackendTransport::HttpJson => HTTP_JSON_TRANSPORT,
CodexBackendTransport::WebSocket => WEBSOCKET_TRANSPORT,
}
}
@@ -2083,6 +2379,19 @@ fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure {
map_client_error(error, UpstreamSendState::Ambiguous, true)
}
fn continuation_replay_required_error() -> ProviderError {
ProviderError::new(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
)
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
Some("invalid_request_error".to_owned()),
))
}
fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error);
let mut failure = map_client_error(error, UpstreamSendState::Sent, false);
@@ -2094,7 +2403,9 @@ fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool {
match error {
CodexClientError::Http(_) | CodexClientError::StreamIdleTimeout { .. } => true,
CodexClientError::Http(_)
| CodexClientError::HttpJson(_)
| CodexClientError::StreamIdleTimeout { .. } => true,
CodexClientError::WebSocket(error) => matches!(
error,
CodexWebSocketExchangeError::Transport(_)
@@ -2175,7 +2486,7 @@ fn map_client_error(
ProviderErrorKind::Protocol,
UpstreamSendState::Sent,
)),
CodexClientError::Http(error) => {
CodexClientError::Http(error) | CodexClientError::HttpJson(error) => {
let send_state = if error.is_connect() {
UpstreamSendState::NotSent
} else {
@@ -159,6 +159,9 @@ pub enum CodexClientError {
/// Reqwest 传输失败。
#[error("http transport error: {0}")]
Http(#[from] reqwest::Error),
/// 非流式 JSON 请求的 Reqwest 传输失败。
#[error("HTTP JSON transport error: {0}")]
HttpJson(#[source] reqwest::Error),
/// 自定义 CA 构建失败。
#[error("custom CA transport error: {0}")]
CustomCa(#[from] CustomCaError),
@@ -222,6 +225,7 @@ impl fmt::Debug for CodexClientError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Http(_) => formatter.write_str("CodexClientError::Http([REDACTED])"),
Self::HttpJson(_) => formatter.write_str("CodexClientError::HttpJson([REDACTED])"),
Self::CustomCa(_) => formatter.write_str("CodexClientError::CustomCa([REDACTED])"),
Self::InvalidHeaderName(_) => {
formatter.write_str("CodexClientError::InvalidHeaderName([REDACTED])")
@@ -274,6 +278,7 @@ impl CodexClientError {
| Self::StreamIdleTimeout { .. }
| Self::InvalidSse(_)
| Self::ModelCatalog(_) => Some(CodexBackendTransport::HttpSse),
Self::HttpJson(_) => Some(CodexBackendTransport::HttpJson),
Self::WebSocket(_) => Some(CodexBackendTransport::WebSocket),
Self::Upstream { transport, .. } => Some(*transport),
Self::CustomCa(_)
@@ -472,6 +477,8 @@ impl fmt::Debug for CodexRequestContext<'_> {
pub enum CodexBackendTransport {
/// HTTP SSE 传输。
HttpSse,
/// 非流式 HTTP JSON 传输。
HttpJson,
/// WebSocket 传输。
WebSocket,
}
@@ -553,6 +560,22 @@ pub struct CodexBackendStreamingResponse {
pub connection_local_continuation: bool,
}
/// Codex 上游非流式 JSON 响应。
pub struct CodexBackendJsonResponse {
/// 未经解析或重编码的完整响应正文。
pub body: Bytes,
/// 上游透传的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游透传的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 上游诊断元数据。
pub diagnostics: CodexUpstreamDiagnostics,
/// 已筛选、可交给客户端的响应头。
pub response_metadata: CodexResponseMetadata,
/// HTTP 阶段耗时与版本。
pub transport_metrics: CodexTransportMetrics,
}
// ---------------------------------------------------------------------------
// CodexBackendClient
// ---------------------------------------------------------------------------
@@ -0,0 +1,109 @@
//! Codex 非流式 JSON 上游透明传输。
use std::time::Instant;
use bytes::Bytes;
use reqwest::header::{ACCEPT, CONTENT_TYPE, HeaderName, HeaderValue};
use super::{
client::{
CodexBackendClient, CodexBackendJsonResponse, CodexBackendTransport, CodexClientError,
CodexClientResult, CodexClientVisibleUpstreamResponse, CodexRequestContext,
CodexTransportMetrics, elapsed_duration_millis, http_version_name,
read_error_response_body, retry_after_seconds,
},
diagnostics::CodexUpstreamSendPhase,
endpoints::endpoint_url,
headers::{build_codex_base_headers, insert_optional_header, insert_optional_protocol_header},
response_meta,
};
impl CodexBackendClient {
/// 向固定 Provider 端点发送一次原始 JSON;请求与成功响应正文均不经过 serde。
pub(crate) async fn post_raw_json(
&self,
endpoint_path: &'static str,
body: Bytes,
image_turn_id: Option<&str>,
context: CodexRequestContext<'_>,
) -> CodexClientResult<CodexBackendJsonResponse> {
let profile = self.profile.snapshot();
let mut headers =
build_codex_base_headers(&profile, context.authorization, context.account_id)?;
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
insert_optional_header(
&mut headers,
"x-codex-installation-id",
context.installation_id,
)?;
insert_optional_protocol_header(&mut headers, "x-codex-image-turn-id", image_turn_id);
headers.insert(
HeaderName::from_static("x-client-request-id"),
HeaderValue::from_str(context.request_id)?,
);
let headers_started_at = Instant::now();
let response = self
.client
.post(endpoint_url(&self.base_url, endpoint_path))
.headers(headers)
.body(body)
.send()
.await
.map_err(CodexClientError::HttpJson)?;
let upstream_headers_ms = elapsed_duration_millis(headers_started_at.elapsed());
let http_version = http_version_name(response.version()).to_owned();
let status = response.status();
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
let response_metadata = response_meta::response_metadata(response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
let transport_metrics = CodexTransportMetrics {
upstream_headers_ms: Some(upstream_headers_ms),
http_version: Some(http_version),
..CodexTransportMetrics::default()
};
if !status.is_success() {
let content_type = response
.headers()
.get(CONTENT_TYPE)
.map(|value| value.as_bytes().to_vec());
let client_headers = response_meta::client_headers(response.headers());
let raw_body = read_error_response_body(response)
.await
.map_err(CodexClientError::HttpJson)?;
let error_body = String::from_utf8_lossy(&raw_body).into_owned();
return Err(CodexClientError::Upstream {
status,
body: error_body,
client_response: Some(Box::new(CodexClientVisibleUpstreamResponse::new(
status,
content_type,
client_headers,
raw_body,
))),
retry_after_seconds,
diagnostics: Box::new(diagnostics),
set_cookie_headers,
rate_limit_headers,
transport: CodexBackendTransport::HttpJson,
transport_metrics: Box::new(transport_metrics),
send_phase: CodexUpstreamSendPhase::AfterPayload,
});
}
let body = response.bytes().await.map_err(CodexClientError::HttpJson)?;
Ok(CodexBackendJsonResponse {
body,
set_cookie_headers,
rate_limit_headers,
diagnostics,
response_metadata,
transport_metrics,
})
}
}
@@ -1,9 +1,23 @@
/// `/codex/responses`
pub const CODEX_RESPONSES_PATH: &str = "/codex/responses";
/// `/codex/images/generations`
pub const CODEX_IMAGE_GENERATIONS_PATH: &str = "/codex/images/generations";
/// `/codex/images/edits`
pub const CODEX_IMAGE_EDITS_PATH: &str = "/codex/images/edits";
/// `/api/codex/usage`
pub const CODEX_USAGE_API_PATH: &str = "/api/codex/usage";
/// `/wham/usage`
pub const WHAM_USAGE_PATH: &str = "/wham/usage";
/// 个人账号每日模型额度权重。
pub const WHAM_DAILY_TOKEN_USAGE_PATH: &str = "/wham/usage/daily-token-usage-breakdown";
/// 个人账号每日 Token 总量。
pub const WHAM_DAILY_WORKSPACE_USAGE_COUNTS_PATH: &str =
"/wham/analytics/daily-workspace-usage-counts";
/// Codex Desktop 主动额度重置卡列表。
pub const WHAM_RATE_LIMIT_RESET_CREDITS_PATH: &str = "/wham/rate-limit-reset-credits";
/// Codex Desktop 主动消费额度重置卡。
pub const WHAM_RATE_LIMIT_RESET_CREDITS_CONSUME_PATH: &str =
"/wham/rate-limit-reset-credits/consume";
/// 拼接完整 endpoint URL。
pub fn endpoint_url(base_url: &str, endpoint_path: &str) -> String {
@@ -3,6 +3,7 @@
pub mod canonical;
pub mod catalog;
pub mod client;
mod client_json;
mod client_sse;
pub mod diagnostics;
pub mod endpoints;
@@ -10,11 +11,13 @@ pub mod headers;
pub mod profile;
pub mod protocol;
pub mod request;
pub mod reset_credits;
mod response_meta;
pub(crate) mod session;
mod time;
pub mod tls;
pub mod usage;
pub(crate) mod usage_statistics;
pub mod websocket;
pub use self::{
@@ -25,21 +28,29 @@ pub use self::{
CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, parse_codex_model_catalog,
},
client::{
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendSseStream,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult,
CodexRateLimitUpdates, CodexRequestContext, CodexTransportDecision, CodexTransportMetrics,
CodexTurnStateUpdate, build_reqwest_client,
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse,
CodexBackendSseStream, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexClientResult, CodexRateLimitUpdates, CodexRequestContext,
CodexTransportDecision, CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client,
},
diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase},
endpoints::{
CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, WHAM_USAGE_PATH, endpoint_url,
usage_endpoint_url,
CODEX_IMAGE_EDITS_PATH, CODEX_IMAGE_GENERATIONS_PATH, CODEX_RESPONSES_PATH,
CODEX_USAGE_API_PATH, WHAM_DAILY_TOKEN_USAGE_PATH, WHAM_DAILY_WORKSPACE_USAGE_COUNTS_PATH,
WHAM_RATE_LIMIT_RESET_CREDITS_CONSUME_PATH, WHAM_RATE_LIMIT_RESET_CREDITS_PATH,
WHAM_USAGE_PATH, endpoint_url, usage_endpoint_url,
},
headers::build_codex_base_headers,
request::{CodexRequestEncodeError, encode_generate_request},
reset_credits::{
CodexRateLimitResetCredit, CodexRateLimitResetCredits,
CodexRateLimitResetCreditsConsumeResult, MAX_CODEX_RESET_CREDITS_BODY_BYTES,
},
response_meta::CodexResponseMetadata,
usage::{MAX_CODEX_USAGE_BODY_BYTES, OpenAiBillingUsage, openai_billing_breakdown},
websocket::{
CodexWebSocketPool, CodexWebSocketPoolConfig, CodexWebSocketPoolKey, WebSocketPoolDecision,
},
};
pub(crate) use self::usage::openai_aggregate_billing_breakdown;
@@ -73,6 +73,19 @@ impl CodexWireProfile {
self.desktop_version,
)
}
/// 按 Codex Desktop 暴露的稳定桌面 surface 格式构造 User-Agent。
///
/// Electron 网络栈会为 renderer 请求附加运行时 User-Agent;非 Electron transport
/// 使用同一官方包中 `getDesktopUserAgent` 的稳定格式。它与 bundled Core 请求使用的
/// 复合 User-Agent 是两个独立 surface。
#[must_use]
pub fn desktop_user_agent(&self) -> String {
format!(
"{}/{} ({}; {})",
self.originator, self.desktop_version, self.os_type, self.arch
)
}
}
/// 跨 Codex 上游请求共享的运行时请求画像。
@@ -525,11 +525,6 @@ impl CodexResponsesRequest {
.map_or(&[], Vec::as_slice)
}
/// 替换输入条目(仅用于代理持有完整历史时的换号重放)。
pub fn set_input(&mut self, input: Vec<Value>) {
self.body.insert("input".to_string(), Value::Array(input));
}
/// 是否流式返回。
pub fn stream(&self) -> bool {
self.body
@@ -617,11 +612,12 @@ impl CodexResponsesRequest {
)
}
/// 返回会话亲和所用的子代理区分值。
/// 返回请求语义与 child transport 隔离所用的子代理区分值。
///
/// Codex 原生请求在 turn metadata 中声明 `subagent_kind`;网关的
/// `/responses/review` 等入口则通过 `client_metadata.x-openai-subagent`
/// 传递给上游。该可选值仅参与会话亲和键派生。
/// 传递给上游。它不拆分根会话的账号首选项;`thread_spawn` 仍用它派生独立
/// WebSocket/continuation transport identity。
pub fn subagent_kind(&self) -> Option<String> {
self.semantics()
.subagent_kind
@@ -5,7 +5,9 @@ use serde_json::{Map, Value};
use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement};
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found";
pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found";
pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE: &str =
"Previous response was not found. Retrying the full request.";
const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`.";
/// WebSocket 握手审计快照。
@@ -12,6 +12,7 @@ use crate::transport::protocol::responses::CodexResponsesRequest;
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
const THREAD_SPAWN_SUBAGENT_KIND: &str = "thread_spawn";
const THREAD_SPAWN_CONVERSATION_PREFIX: &str = "thread-spawn:";
const UNSUPPORTED_CODEX_RESPONSES_FIELDS: &[&str] = &["max_output_tokens", "temperature"];
const CROSS_ACCOUNT_IDENTITY_KEYS: &[&str] = &[
"authorization",
@@ -63,18 +64,14 @@ const INSTALLATION_ID_KEYS: &[&str] = &[
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestAccountScope {
SameAccount,
CrossAccountFullRequest,
CrossAccountContinuationProbe,
Same,
Different,
Unknown,
}
impl RequestAccountScope {
const fn is_cross_account(self) -> bool {
!matches!(self, Self::SameAccount)
}
const fn preserves_continuation(self) -> bool {
matches!(self, Self::CrossAccountContinuationProbe)
pub(crate) const fn can_reuse_account_state(self) -> bool {
matches!(self, Self::Same)
}
}
@@ -94,7 +91,7 @@ pub fn encode_generate_request(
return Err(CodexRequestEncodeError::InvalidProtocolPayload);
}
let mut body = payload.body().clone();
body.insert("model".to_owned(), Value::String(upstream_model.to_owned()));
adapt_codex_responses_body(&mut body, upstream_model);
let mut encoded = CodexResponsesRequest::from_body(body);
encoded.explicit_prompt_cache_key = encoded.prompt_cache_key().is_some();
@@ -103,6 +100,13 @@ pub fn encode_generate_request(
Ok(encoded)
}
fn adapt_codex_responses_body(body: &mut Map<String, Value>, upstream_model: &str) {
body.insert("model".to_owned(), Value::String(upstream_model.to_owned()));
for field in UNSUPPORTED_CODEX_RESPONSES_FIELDS {
body.remove(*field);
}
}
fn extract_request_context(request: &mut CodexResponsesRequest) {
let context = ExtractedRequestContext::from_body(request.body());
request.turn_state = context.turn_state;
@@ -205,8 +209,9 @@ pub(crate) fn derive_conversation_anchor(
})
}
/// A `thread_spawn` shares its parent session but is a separate child execution.
/// Its child thread must therefore own the affinity and WebSocket pool identity.
/// `thread_spawn` 与父任务共享根会话,但它本身是独立的子任务执行。
/// 因此子线程必须拥有独立的 continuation/WebSocket 传输身份;账号亲和则从共享的
/// 根会话单独派生。
fn thread_spawn_conversation_anchor(
request: &CodexResponsesRequest,
) -> Option<(&'static str, String)> {
@@ -306,21 +311,17 @@ fn normalize_conversation_anchor_text(text: &str) -> String {
///
/// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入
/// 请求头;正文只替换客户端原本声明过的 installation 字段,绝不接受客户端
/// 提供的 token、cookie 或账号身份。Continuation probe 保留 canonical
/// previous-response 和原始增量 input,不能把它们误当成完整请求清理。
/// 提供的 token、cookie 或账号身份。`input` 是 Responses 的可回放会话正文,
/// item ID、encrypted content 和 compaction 都必须原样保留。
pub(crate) fn scope_request_to_account(
request: &mut CodexResponsesRequest,
installation_id: &str,
account_scope: RequestAccountScope,
) {
let cross_account = account_scope.is_cross_account();
let continuation = account_scope
.preserves_continuation()
.then(|| request.previous_response_id().map(ToOwned::to_owned))
.flatten();
let reset_account_state = !account_scope.can_reuse_account_state();
let client_metadata_turn_state = metadata_string(request, "x-codex-turn-state");
let preserve_turn_state =
!cross_account && (request.turn_state.is_some() || client_metadata_turn_state.is_some());
let preserve_turn_state = !reset_account_state
&& (request.turn_state.is_some() || client_metadata_turn_state.is_some());
let turn_state = preserve_turn_state
.then(|| request.turn_state.clone())
.flatten();
@@ -329,38 +330,22 @@ pub(crate) fn scope_request_to_account(
} else {
None
};
let turn_metadata = (!account_scope.preserves_continuation())
.then(|| {
request
.turn_metadata
.as_deref()
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account))
})
.flatten();
let client_metadata_turn_metadata = (!account_scope.preserves_continuation())
.then(|| {
metadata_string(request, "x-codex-turn-metadata")
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account))
})
.flatten();
let turn_metadata = request
.turn_metadata
.as_deref()
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, reset_account_state));
let client_metadata_turn_metadata = metadata_string(request, "x-codex-turn-metadata")
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, reset_account_state));
if cross_account {
if reset_account_state {
request.passthrough_headers.remove("x-codex-turn-state");
request.passthrough_headers.remove("x-codex-turn-metadata");
if !account_scope.preserves_continuation() {
sanitize_cross_account_input(request);
}
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.iter()
.chain(ACCOUNT_BOUND_STATE_KEYS)
{
request.body_mut().remove(*key);
}
if account_scope.preserves_continuation() {
for key in TURN_METADATA_KEYS {
request.body_mut().remove(*key);
}
}
}
for key in INSTALLATION_ID_KEYS {
@@ -371,7 +356,7 @@ pub(crate) fn scope_request_to_account(
.body()
.get(*key)
.and_then(Value::as_str)
.and_then(|value| scope_turn_metadata(value, installation_id, cross_account));
.and_then(|value| scope_turn_metadata(value, installation_id, reset_account_state));
replace_existing_body_string(request, key, scoped.as_deref());
}
@@ -383,22 +368,17 @@ pub(crate) fn scope_request_to_account(
(
key,
metadata.get(key).and_then(Value::as_str).and_then(|value| {
scope_turn_metadata(value, installation_id, cross_account)
scope_turn_metadata(value, installation_id, reset_account_state)
}),
)
});
if cross_account {
if reset_account_state {
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.iter()
.chain(ACCOUNT_BOUND_STATE_KEYS)
{
metadata.remove(*key);
}
if account_scope.preserves_continuation() {
for key in TURN_METADATA_KEYS {
metadata.remove(*key);
}
}
}
metadata.insert(
"x-codex-installation-id".to_owned(),
@@ -436,36 +416,6 @@ pub(crate) fn scope_request_to_account(
request.turn_state = turn_state;
request.turn_metadata = turn_metadata;
if let Some(continuation) = continuation {
request.set_previous_response_id(Some(continuation));
}
}
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);
}
pub(crate) 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 metadata_string(request: &CodexResponsesRequest, key: &str) -> Option<String> {
@@ -0,0 +1,285 @@
//! Codex Desktop 主动额度重置卡 HTTP contract。
use chrono::{DateTime, Utc};
use reqwest::{
StatusCode,
header::{ACCEPT, AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue, USER_AGENT},
};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use super::{
CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
client::{
CodexBackendTransport, CodexTransportMetrics, read_capped_response_body,
retry_after_seconds,
},
diagnostics::CodexUpstreamSendPhase,
endpoints::{
WHAM_RATE_LIMIT_RESET_CREDITS_CONSUME_PATH, WHAM_RATE_LIMIT_RESET_CREDITS_PATH,
endpoint_url,
},
headers::insert_optional_header,
response_meta,
};
/// 单次 reset-credit 响应允许保留和解析的最大字节数。
pub const MAX_CODEX_RESET_CREDITS_BODY_BYTES: usize = 1024 * 1024;
const MAX_RESET_CREDITS: usize = 1024;
const MAX_IDENTIFIER_BYTES: usize = 512;
const MAX_TEXT_BYTES: usize = 2048;
/// 上游返回的一张安全重置卡投影。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexRateLimitResetCredit {
pub id: String,
pub status: Option<String>,
pub title: Option<String>,
pub expires_at: Option<DateTime<Utc>>,
pub reset_type: Option<String>,
}
/// 上游重置卡列表。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexRateLimitResetCredits {
pub available_count: u64,
pub credits: Vec<CodexRateLimitResetCredit>,
}
/// 上游消费结果。`code` 的业务含义由调用方按官方状态机解释。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexRateLimitResetCreditsConsumeResult {
pub code: String,
pub credit: Option<CodexRateLimitResetCredit>,
}
#[derive(Deserialize)]
struct ResetCreditWire {
id: String,
#[serde(default)]
status: Option<String>,
#[serde(default)]
title: Option<String>,
#[serde(default)]
expires_at: Option<String>,
#[serde(default)]
reset_type: Option<String>,
}
#[derive(Deserialize)]
struct ResetCreditsWire {
available_count: u64,
credits: Vec<ResetCreditWire>,
}
#[derive(Deserialize)]
struct ConsumeResultWire {
code: String,
#[serde(default)]
credit: Option<ResetCreditWire>,
}
#[derive(Serialize)]
struct ConsumeRequest<'a> {
#[serde(skip_serializing_if = "Option::is_none")]
credit_id: Option<&'a str>,
redeem_request_id: String,
}
impl CodexBackendClient {
/// 查询当前账号的 Codex Desktop 主动额度重置卡。
pub async fn list_rate_limit_reset_credits(
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<CodexRateLimitResetCredits> {
let response = self
.client
.get(endpoint_url(
&self.base_url,
WHAM_RATE_LIMIT_RESET_CREDITS_PATH,
))
.headers(self.reset_credits_headers(context, false)?)
.send()
.await
.map_err(CodexClientError::HttpJson)?;
let body = read_reset_credits_response(response).await?;
let wire = serde_json::from_str::<ResetCreditsWire>(&body)
.map_err(|_| invalid_response("invalid reset-credit list response"))?;
if wire.credits.len() > MAX_RESET_CREDITS {
return Err(invalid_response(
"reset-credit list exceeded the item limit",
));
}
Ok(CodexRateLimitResetCredits {
available_count: wire.available_count,
credits: wire
.credits
.into_iter()
.map(parse_credit)
.collect::<CodexClientResult<Vec<_>>>()?,
})
}
/// 消费一张 Codex Desktop 主动额度重置卡。
///
/// 此方法不做任何自动重试;请求发出后的传输失败必须由上层按不确定结果处理。
pub async fn consume_rate_limit_reset_credit(
&self,
context: CodexRequestContext<'_>,
credit_id: Option<&str>,
redeem_request_id: Uuid,
) -> CodexClientResult<CodexRateLimitResetCreditsConsumeResult> {
let response = self
.client
.post(endpoint_url(
&self.base_url,
WHAM_RATE_LIMIT_RESET_CREDITS_CONSUME_PATH,
))
.headers(self.reset_credits_headers(context, true)?)
.json(&ConsumeRequest {
credit_id,
redeem_request_id: redeem_request_id.to_string(),
})
.send()
.await
.map_err(CodexClientError::HttpJson)?;
let body = read_reset_credits_response(response).await?;
let wire = serde_json::from_str::<ConsumeResultWire>(&body)
.map_err(|_| invalid_response("invalid reset-credit consume response"))?;
Ok(CodexRateLimitResetCreditsConsumeResult {
code: required_text(wire.code, MAX_IDENTIFIER_BYTES, "reset-credit result code")?,
credit: wire.credit.map(parse_credit).transpose()?,
})
}
fn reset_credits_headers(
&self,
context: CodexRequestContext<'_>,
json_body: bool,
) -> CodexClientResult<HeaderMap> {
let profile = self.profile.snapshot();
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, HeaderValue::from_str(context.authorization)?);
insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?;
headers.insert(
HeaderName::from_static("originator"),
HeaderValue::from_str(&profile.originator)?,
);
headers.insert(
HeaderName::from_static("oai-language"),
HeaderValue::from_static("en"),
);
headers.insert(
USER_AGENT,
HeaderValue::from_str(&profile.desktop_user_agent())?,
);
// Electron `net.fetch` 为 renderer 请求附加的默认 Accept;该 endpoint 的
// renderer 并没有显式改写为 Core auxiliary 请求的 application/json 画像。
headers.insert(ACCEPT, HeaderValue::from_static("*/*"));
if json_body {
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
}
Ok(headers)
}
}
async fn read_reset_credits_response(response: reqwest::Response) -> CodexClientResult<String> {
let status = response.status();
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
let body = read_capped_response_body(response, MAX_CODEX_RESET_CREDITS_BODY_BYTES)
.await
.map_err(CodexClientError::HttpJson)?;
if body.limit_exceeded() {
return Err(upstream_error(
if status.is_success() {
StatusCode::BAD_GATEWAY
} else {
status
},
"upstream reset-credit response exceeded the body limit".to_owned(),
retry_after_seconds,
diagnostics,
));
}
let body = body.into_string();
if status.is_success() {
Ok(body)
} else {
Err(upstream_error(
status,
body,
retry_after_seconds,
diagnostics,
))
}
}
fn upstream_error(
status: StatusCode,
body: String,
retry_after_seconds: Option<u64>,
diagnostics: super::CodexUpstreamDiagnostics,
) -> CodexClientError {
CodexClientError::Upstream {
status,
body,
client_response: None,
retry_after_seconds,
diagnostics: Box::new(diagnostics),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
transport: CodexBackendTransport::HttpJson,
transport_metrics: Box::<CodexTransportMetrics>::default(),
send_phase: CodexUpstreamSendPhase::AfterPayload,
}
}
fn invalid_response(message: &'static str) -> CodexClientError {
upstream_error(
StatusCode::BAD_GATEWAY,
message.to_owned(),
None,
super::CodexUpstreamDiagnostics::default(),
)
}
fn parse_credit(wire: ResetCreditWire) -> CodexClientResult<CodexRateLimitResetCredit> {
let expires_at = wire
.expires_at
.map(|value| {
DateTime::parse_from_rfc3339(value.trim())
.map(|value| value.with_timezone(&Utc))
.map_err(|_| invalid_response("invalid reset-credit expiry"))
})
.transpose()?;
Ok(CodexRateLimitResetCredit {
id: required_text(wire.id, MAX_IDENTIFIER_BYTES, "reset-credit ID")?,
status: optional_text(wire.status, MAX_IDENTIFIER_BYTES, "reset-credit status")?,
title: optional_text(wire.title, MAX_TEXT_BYTES, "reset-credit title")?,
expires_at,
reset_type: optional_text(wire.reset_type, MAX_IDENTIFIER_BYTES, "reset-credit type")?,
})
}
fn required_text(
value: String,
max_bytes: usize,
field: &'static str,
) -> CodexClientResult<String> {
let value = value.trim();
if value.is_empty() || value.len() > max_bytes || value.chars().any(char::is_control) {
return Err(invalid_response(field));
}
Ok(value.to_owned())
}
fn optional_text(
value: Option<String>,
max_bytes: usize,
field: &'static str,
) -> CodexClientResult<Option<String>> {
value
.map(|value| required_text(value, max_bytes, field))
.transpose()
}
@@ -412,6 +412,30 @@ pub fn openai_billing_breakdown(
model: &str,
usage: OpenAiBillingUsage,
service_tier: Option<&str>,
) -> Option<CalculatedCostBreakdown> {
openai_billing_breakdown_with_context(
model,
usage,
service_tier,
usage.input_tokens > LONG_CONTEXT_THRESHOLD,
)
}
/// 官方日报没有单次请求边界,不能把整日输入量套用长上下文阶梯。
/// 聚合统计固定使用公开的短上下文价格,其余价格规则与请求级计费共用。
pub(crate) fn openai_aggregate_billing_breakdown(
model: &str,
usage: OpenAiBillingUsage,
service_tier: Option<&str>,
) -> Option<CalculatedCostBreakdown> {
openai_billing_breakdown_with_context(model, usage, service_tier, false)
}
fn openai_billing_breakdown_with_context(
model: &str,
usage: OpenAiBillingUsage,
service_tier: Option<&str>,
long_context: bool,
) -> Option<CalculatedCostBreakdown> {
if usage.cached_tokens > usage.input_tokens
|| usage.image_input_tokens > 0
@@ -421,7 +445,6 @@ pub fn openai_billing_breakdown(
}
let web_search_ticks = web_search_amount_ticks(usage)?;
let pricing = model_pricing(model)?;
let long_context = usage.input_tokens > LONG_CONTEXT_THRESHOLD;
let normalized_tier = normalize_service_tier(service_tier);
let tier = pricing_tier(normalized_tier.as_deref())?;
let standard_rates = pricing.rates(PricingTier::Standard, long_context)?;
@@ -0,0 +1,119 @@
//! Codex 官方每日用量报表 transport。
use chrono::NaiveDate;
use gateway_protocol::openai::events::retry_after_seconds_from_body;
use reqwest::StatusCode;
use serde_json::Value;
use super::{
CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
client::{
CodexBackendTransport, CodexTransportMetrics, read_capped_response_body,
retry_after_seconds, truncate_for_error,
},
diagnostics::CodexUpstreamSendPhase,
endpoints::{
WHAM_DAILY_TOKEN_USAGE_PATH, WHAM_DAILY_WORKSPACE_USAGE_COUNTS_PATH, endpoint_url,
},
response_meta,
};
const MAX_CODEX_DAILY_USAGE_BODY_BYTES: usize = 4 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CodexDailyUsageReport {
PersonalModelCredits,
PersonalTokenTotals,
}
impl CodexDailyUsageReport {
const fn path(self) -> &'static str {
match self {
Self::PersonalModelCredits => WHAM_DAILY_TOKEN_USAGE_PATH,
Self::PersonalTokenTotals => WHAM_DAILY_WORKSPACE_USAGE_COUNTS_PATH,
}
}
const fn workspace_user(self) -> bool {
matches!(self, Self::PersonalTokenTotals)
}
}
impl CodexBackendClient {
pub(crate) async fn fetch_daily_usage(
&self,
context: CodexRequestContext<'_>,
report: CodexDailyUsageReport,
start_date: NaiveDate,
end_date: NaiveDate,
) -> CodexClientResult<Value> {
let headers = self.usage_request_headers(context)?;
let workspace_user = if report.workspace_user() {
"&workspace_user=true"
} else {
""
};
let url = format!(
"{}?start_date={start_date}&end_date={end_date}&group_by=day{workspace_user}",
endpoint_url(&self.base_url, report.path()),
);
let response = self.client.get(url).headers(headers).send().await?;
let status = response.status();
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
let body = read_capped_response_body(response, MAX_CODEX_DAILY_USAGE_BODY_BYTES).await?;
if body.limit_exceeded() {
return Err(upstream_error(
if status.is_success() {
StatusCode::BAD_GATEWAY
} else {
status
},
"upstream daily usage response exceeded the body limit".to_owned(),
retry_after_seconds,
diagnostics,
));
}
let body = body.into_string();
if !status.is_success() {
return Err(upstream_error(
status,
body.clone(),
retry_after_seconds.or_else(|| retry_after_seconds_from_body(&body)),
diagnostics,
));
}
serde_json::from_str::<Value>(&body).map_err(|_| {
upstream_error(
StatusCode::BAD_GATEWAY,
format!(
"invalid daily usage response: {}",
truncate_for_error(&body)
),
None,
diagnostics,
)
})
}
}
fn upstream_error(
status: StatusCode,
body: String,
retry_after_seconds: Option<u64>,
diagnostics: super::CodexUpstreamDiagnostics,
) -> CodexClientError {
CodexClientError::Upstream {
status,
body,
client_response: None,
retry_after_seconds,
diagnostics: Box::new(diagnostics),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
transport: CodexBackendTransport::HttpSse,
transport_metrics: Box::<CodexTransportMetrics>::default(),
send_phase: CodexUpstreamSendPhase::AfterPayload,
}
}
@@ -10,13 +10,13 @@ use std::{
time::Duration,
};
use tokio::{sync::Semaphore, time::Instant};
use tokio::sync::Semaphore;
use tokio_util::{sync::CancellationToken, task::TaskTracker};
use uuid::Uuid;
use self::state::{
WebSocketPoolConnecting, WebSocketPoolSlot, WebSocketPoolState, account_slot_count,
close_pooled_connection, close_pooled_connections, take_lru_idle_connection,
WebSocketPoolConnecting, WebSocketPoolSlot, WebSocketPoolState, close_pooled_connection,
close_pooled_connections,
};
use super::pump::PumpKeepalive;
@@ -31,8 +31,6 @@ pub(crate) use self::{
},
};
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAX_TOTAL: usize = 64;
const DEFAULT_MAX_CONNECTING: usize = 16;
const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
@@ -64,10 +62,6 @@ pub struct CodexWebSocketPoolConfig {
pub enabled: bool,
/// 单个 socket 的最大生命周期。
pub max_age: Duration,
/// 单个账号允许占用的最大池 slot 数。
pub max_per_account: usize,
/// 所有账号合计允许占用的最大池 slot 数。
pub max_total: usize,
/// 所有账号合计允许并发执行的 opening 数。
pub max_connecting: usize,
/// 后台维护间隔;`None` 表示不启动后台任务。
@@ -87,8 +81,6 @@ impl Default for CodexWebSocketPoolConfig {
Self {
enabled: true,
max_age: DEFAULT_MAX_AGE,
max_per_account: DEFAULT_MAX_PER_ACCOUNT,
max_total: DEFAULT_MAX_TOTAL,
max_connecting: DEFAULT_MAX_CONNECTING,
maintenance_interval: Some(DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(DEFAULT_PING_INTERVAL),
@@ -113,10 +105,9 @@ impl CodexWebSocketPoolConfig {
}
impl CodexWebSocketPool {
/// 构造连接池策略和状态。
pub fn new(max_per_account: usize, max_age: Duration) -> Self {
/// 构造不限制累计 slot 数的连接池策略和状态。
pub fn new(max_age: Duration) -> Self {
Self::with_config(CodexWebSocketPoolConfig {
max_per_account: max_per_account.max(1),
max_age,
maintenance_interval: None,
ping_interval: None,
@@ -218,43 +209,25 @@ impl CodexWebSocketPool {
let acquire = if required_response_id.is_some() {
WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::ContinuationNotFound)
} else if self.config.max_per_account == 0
|| account_slot_count(&state.slots, key.account_id()) >= self.config.max_per_account
{
WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap)
} else {
let connect_permit = self.connect_semaphore.clone().try_acquire_owned().ok();
match connect_permit {
Some(connect_permit) => {
let has_total_capacity = if self.config.max_total == 0 {
false
} else if state.slots.len() < self.config.max_total {
true
} else if let Some(connection) = take_lru_idle_connection(&mut state) {
connections_to_close.push(connection);
true
} else {
false
};
if has_total_capacity {
let lease = WebSocketPoolConnectLease::reserve(
self.clone(),
key.clone(),
connect_permit,
);
state.slots.insert(
key,
WebSocketPoolSlot::Connecting(WebSocketPoolConnecting {
id: lease.id,
started_at: lease.started_at,
outcome: lease.outcome.clone(),
cancellation: lease.cancellation.clone(),
}),
);
WebSocketPoolAcquire::Connect(lease)
} else {
WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap)
}
let lease = WebSocketPoolConnectLease::reserve(
self.clone(),
key.clone(),
connect_permit,
);
state.slots.insert(
key,
WebSocketPoolSlot::Connecting(WebSocketPoolConnecting {
id: lease.id,
started_at: lease.started_at,
outcome: lease.outcome.clone(),
cancellation: lease.cancellation.clone(),
}),
);
WebSocketPoolAcquire::Connect(lease)
}
None => WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap),
}
@@ -292,7 +265,6 @@ impl CodexWebSocketPool {
key.clone(),
WebSocketPoolSlot::Idle {
connection: Box::new(connection),
last_used_at: Instant::now(),
},
);
}
@@ -108,7 +108,6 @@ impl WebSocketContinuationState {
pub(super) enum WebSocketPoolSlot {
Idle {
connection: Box<PooledWebSocketConnection>,
last_used_at: Instant,
},
Busy(WebSocketPoolReservation),
Connecting(WebSocketPoolConnecting),
@@ -138,34 +137,6 @@ pub(super) struct WebSocketPoolReservation {
pub(super) latest_response_id: Option<String>,
}
pub(super) fn account_slot_count(
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
account_id: &str,
) -> usize {
slots
.keys()
.filter(|key| key.account_id() == account_id)
.count()
}
pub(super) fn take_lru_idle_connection(
state: &mut WebSocketPoolState,
) -> Option<PooledWebSocketConnection> {
let key = state
.slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle { last_used_at, .. } => Some((key, *last_used_at)),
WebSocketPoolSlot::Busy(_) | WebSocketPoolSlot::Connecting(_) => None,
})
.min_by_key(|(_, last_used_at)| *last_used_at)
.map(|(key, _)| key.clone())?;
let Some(WebSocketPoolSlot::Idle { connection, .. }) = state.slots.remove(&key) else {
return None;
};
Some(*connection)
}
pub(super) async fn close_pooled_connection(connection: PooledWebSocketConnection) {
connection.websocket.close().await;
}
+257 -7
View File
@@ -1,20 +1,21 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, TimeZone as _, Utc};
use futures::future::BoxFuture;
use futures::{StreamExt, future::BoxFuture};
use gateway_admin::model::accounts::AccountRecord;
use gateway_admin::model::observability::{
CurrencyCost, DesktopReleaseStatus, ProviderBillingInput,
};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwnerBinding, CompleteAuthorization,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, ProviderDocument, ProviderExportCredentialInput,
ProviderQuotaRequest, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
ConsumeProviderResetCredit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, ProviderDocument,
ProviderExportCredentialInput, ProviderQuotaRequest, ProviderQuotaWindowRole,
QuotaLocalUsageAttribution,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::ProviderAdminErrorKind;
@@ -22,7 +23,12 @@ use gateway_core::engine::credential::{
CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
ProviderAccountStore, QuotaAccessChange, QuotaEvidence, QuotaObservation, QuotaState,
};
use gateway_core::engine::provider::ProviderRequest;
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
};
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
@@ -31,7 +37,11 @@ use gateway_core::provider_ports::{
ProviderCredentialState, ProviderCredentialStatePort, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderStoreError, ProviderStorePorts,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use gateway_core::routing::{
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory,
RuntimeSnapshot, UpstreamModelId,
};
use gateway_core::task::{WorkerContribution, WorkerKind, WorkerRunnable};
use provider_openai::config::{CodexWireProfileConfig, OpenAiConfig};
use provider_openai::credential::{CodexCredentialCodec, ImportCodexOAuthCredential};
@@ -39,12 +49,20 @@ use provider_openai::transport::profile::APPCAST_POLL_INTERVAL;
use secrecy::SecretString;
use serde_json::{Map, Value, json};
use tempfile::TempDir;
use uuid::Uuid;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{
MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator,
profile, secret,
account_policy, profile, secret,
};
const COMPLETED_SESSION_SSE: &str = concat!(
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_initialized_session\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
);
#[tokio::test]
async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions_once() {
let config = valid_config();
@@ -93,6 +111,84 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
assert!(bundle.take_worker_contributions().is_empty());
}
#[tokio::test]
async fn initialized_provider_keeps_thread_spawn_transport_conversations_distinct() {
let account_id = "acct_initialized_thread_spawn";
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: account_id.to_owned(),
secret: secret("at-initialized-thread-spawn"),
verified_account: profile("chatgpt-initialized-thread-spawn"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(COMPLETED_SESSION_SSE),
)
.expect(2)
.mount(&server)
.await;
let mut config = valid_config();
config.config.api.base_url = server.uri();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(store, Arc::new(TestOAuthPending::default())),
)
.await
.expect("initialized OpenAI provider");
let provider = bundle.core_provider();
let thread_spawn = r#"{"subagent_kind":"thread_spawn"}"#;
let mut conversation_ids = Vec::new();
for (request_id, thread_id) in [
("req_initialized_thread_spawn_first", "child-one"),
("req_initialized_thread_spawn_second", "child-two"),
] {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("child task")),
("session_id".to_owned(), json!("parent-session")),
("thread_id".to_owned(), json!(thread_id)),
("turnMetadata".to_owned(), json!(thread_spawn)),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(payload));
let mut stream = provider
.execute(
initialized_provider_request(operation, account_id),
initialized_attempt_context(request_id, account_id),
)
.await
.expect("prepare child provider stream");
let mut conversation_id = None;
while let Some(event) = stream.next().await {
let event = event.expect("child provider response");
if let Some(update) = event.session_update() {
conversation_id = update
.payload()
.get("conversation_id")
.and_then(Value::as_str)
.map(ToOwned::to_owned);
}
}
conversation_ids.push(conversation_id.expect("child transport conversation id"));
}
assert_ne!(conversation_ids[0], conversation_ids[1]);
}
#[tokio::test]
async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing() {
let config = valid_config();
@@ -160,6 +256,64 @@ async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing()
assert_eq!(fast_billing.total_amount.amount.as_str(), "4.25");
}
#[tokio::test]
async fn reset_credit_success_with_invalid_body_should_remain_an_unknown_consume_result() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/wham/rate-limit-reset-credits/consume"))
.respond_with(ResponseTemplate::new(200).set_body_raw("{}", "application/json"))
.expect(1)
.mount(&server)
.await;
let (bundle, account_id, _config) = reset_credit_admin(&server).await;
let error = bundle
.admin_provider()
.consume_reset_credit(reset_credit_command(account_id))
.await
.expect_err("invalid success body must be ambiguous");
assert_eq!(error.kind(), ProviderAdminErrorKind::BadGateway);
assert_eq!(
error.message(),
Some(
"OpenAI reset-credit consume result is unknown; refresh the credit list before retrying"
)
);
}
#[tokio::test]
async fn reset_credit_explicit_http_rejection_should_preserve_the_raw_body() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/wham/rate-limit-reset-credits/consume"))
.respond_with(ResponseTemplate::new(409).set_body_raw(
r#"{"code":"nothing_to_reset","detail":"window is fresh"}"#,
"application/json",
))
.expect(1)
.mount(&server)
.await;
let (bundle, account_id, _config) = reset_credit_admin(&server).await;
let error = bundle
.admin_provider()
.consume_reset_credit(reset_credit_command(account_id))
.await
.expect_err("explicit upstream rejection");
assert_eq!(error.kind(), ProviderAdminErrorKind::BadGateway);
assert_eq!(
error.message(),
Some(
r#"OpenAI reset-credit upstream returned HTTP 409: {"code":"nothing_to_reset","detail":"window is fresh"}"#
)
);
let debug = format!("{error:?}");
assert!(debug.contains("<redacted>"));
assert!(!debug.contains("window is fresh"));
}
#[tokio::test]
async fn openai_core_provider_projects_codex_request_observation_without_routing_side_effects() {
let config = valid_config();
@@ -820,6 +974,102 @@ async fn openai_rotation_preserves_the_new_access_token_jwt_expiration() {
assert_eq!(prepared.facts().access_token_expires_at, Some(expires_at));
}
async fn reset_credit_admin(
server: &MockServer,
) -> (
provider_openai::ProviderBundle,
ProviderAccountId,
TestOpenAiConfig,
) {
let account_id = ProviderAccountId::new("acct_reset_credit").expect("account ID");
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_string(),
name: "reset credit".to_owned(),
secret: secret("reset-credit-access"),
verified_account: profile("chatgpt-reset-credit"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let mut config = valid_config();
config.config.api.base_url = server.uri();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(store, Arc::new(TestOAuthPending::default())),
)
.await
.expect("OpenAI reset-credit bundle");
(bundle, account_id, config)
}
fn reset_credit_command(account_id: ProviderAccountId) -> ConsumeProviderResetCredit {
ConsumeProviderResetCredit {
account_id,
credit_id: Some("credit_1".to_owned()),
redeem_request_id: Uuid::parse_str("8fbf302d-11df-4bd5-82e4-08e4b3df7874")
.expect("UUID v4"),
}
}
fn initialized_provider_request(operation: Operation, account_id: &str) -> ProviderRequest {
let provider = ProviderKind::new("openai").expect("provider");
let upstream_model = UpstreamModelId::new("gpt-5.4").expect("upstream model");
let public_model = PublicModelId::new(upstream_model.as_str()).expect("public model");
let account_scope = initialized_account_scope(account_id);
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
account_policy(),
vec![provider.clone()],
vec![ProviderModel::new(
provider,
upstream_model,
ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
.with_upstream_feature_validation(),
)],
Vec::new(),
)
.expect("runtime snapshot");
let plan = snapshot
.plan(
&public_model,
&operation,
account_scope,
&RoutingContext::default(),
)
.expect("routing plan");
ProviderRequest::new(operation, plan.candidates()[0].clone())
}
fn initialized_attempt_context(request_id: &str, account_id: &str) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_initialized").expect("client key id"),
),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(initialized_account_scope(account_id)),
None,
CancellationToken::new(),
)
}
fn initialized_account_scope(account_id: &str) -> Arc<FrozenAccountScope> {
let provider = ProviderKind::new("openai").expect("provider");
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new(account_id).expect("account id"),
RuntimeAccount::new(provider, BTreeSet::new()),
)]))),
ClientRoutingScope::all_accounts(),
))
}
fn provider_ports() -> ProviderStorePorts {
provider_ports_with(
Arc::new(MemoryAccountStore::default()),
@@ -101,8 +101,6 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
config.api.base_url.as_str(),
config.ws_pool.enabled,
config.ws_pool.max_age_ms,
config.ws_pool.max_per_account,
config.ws_pool.max_total,
config.ws_pool.max_connecting,
config.ws_pool.initial_event_timeout_ms,
config.quota.refresh_interval_minutes,
@@ -115,8 +113,6 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
true,
3_300_000,
8,
64,
8,
20_000,
15,
true,
@@ -59,6 +59,48 @@ async fn direct_import_persists_opaque_tokens_without_profile_or_token_validatio
assert!(account.account.next_refresh_at().is_none());
}
#[tokio::test]
async fn direct_import_accepts_snake_case_oauth_token_aliases() {
let service = CodexCredentialAdminService::new(
Arc::new(UnusedRefresher),
Arc::new(TestLeaseCoordinator::default()),
runtime_policy(),
);
let prepared = service
.prepare_import_document(serde_json::json!({
"accounts": [{
"platform": "openai",
"type": "oauth",
"credentials": {
"access_token": "snake-access-token",
"refresh_token": "snake-refresh-token",
"id_token": "snake-id-token"
}
}]
}))
.await
.expect("snake_case OAuth token aliases must be accepted");
let account = prepared.accounts().first().expect("one prepared account");
let runtime = CodexCredentialCodec::decode(&account.credential).expect("stored credential");
let oauth = runtime.authentication.oauth().expect("OAuth credential");
assert_eq!(
(
oauth.access_token.expose_secret(),
oauth
.refresh_token
.as_ref()
.map(|token| token.expose_secret()),
oauth.id_token.as_ref().map(|token| token.expose_secret()),
),
(
"snake-access-token",
Some("snake-refresh-token"),
Some("snake-id-token"),
)
);
}
#[tokio::test]
async fn direct_import_projects_access_token_jwt_expiry_without_persisting_refresh_margin() {
let expires_at = Utc
@@ -445,6 +445,74 @@ fn selector_round_robin_cursor_advances_across_requests() {
);
}
#[tokio::test]
async fn selector_should_claim_the_initial_session_account_before_upstream_send() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let affinity = Arc::new(MemorySessionAffinity::default());
let selector = selector_with_affinity(
&store,
Arc::new(TestLeaseCoordinator::default()),
Arc::clone(&affinity),
);
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("initial-claim").expect("affinity key");
let request_attempt = attempt(BTreeSet::new());
let selected = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
attempt: &request_attempt,
session_affinity_key: Some(&key),
})
.await
.expect("select initial account");
assert_eq!(
affinity
.load(&provider, &key)
.await
.expect("load claimed affinity"),
Some(selected.account_id().clone())
);
}
#[tokio::test]
async fn record_success_should_not_overwrite_a_newer_session_winner() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let affinity = Arc::new(MemorySessionAffinity::default());
let selector = selector_with_affinity(
&store,
Arc::new(TestLeaseCoordinator::default()),
Arc::clone(&affinity),
);
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("late-success").expect("affinity key");
let first = store.account("acct_first").expect("first account");
let second = ProviderAccountId::new("acct_second").expect("second account");
affinity
.bind(&provider, &key, &second, Duration::from_secs(60))
.await
.expect("seed newer affinity");
selector
.record_success(&first, Some(&key), first.id())
.await;
assert_eq!(
affinity
.load(&provider, &key)
.await
.expect("load preserved affinity"),
Some(second)
);
}
#[tokio::test]
async fn selector_should_reuse_the_account_bound_to_the_same_session() {
let store = Arc::new(MemoryAccountStore::default());
@@ -469,7 +537,9 @@ async fn selector_should_reuse_the_account_bound_to_the_same_session() {
})
.await
.expect("select first account");
selector.record_success(first.account(), Some(&key)).await;
selector
.record_success(first.account(), Some(&key), first.account_id())
.await;
let first_account = first.account_id().clone();
let second_attempt = attempt(BTreeSet::new());
@@ -535,7 +605,7 @@ async fn selector_should_replace_a_busy_affinity_binding_after_the_fallback_succ
.await
.expect("select fallback account");
selector
.record_success(selected.account(), Some(&key))
.record_success(selected.account(), Some(&key), &bound)
.await;
assert_eq!(
@@ -1005,7 +1075,7 @@ fn successful_upstream_response_recovers_non_quota_terminal_states() {
let current = store.account("acct_primary").expect("stale account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_success(&current, None));
block_on(selector.record_success(&current, None, current.id()));
assert_eq!(
store
@@ -1289,7 +1359,7 @@ fn cloudflare_challenge_does_not_change_persisted_account_facts() {
.credential_state(),
CredentialState::Ready
);
block_on(selector.record_success(lease.account(), None));
block_on(selector.record_success(lease.account(), None, lease.account_id()));
}
#[test]
@@ -2,6 +2,7 @@
mod recovery;
mod snapshot;
mod usage_statistics;
use std::sync::Arc;
use std::time::SystemTime;
@@ -0,0 +1,217 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use chrono::{DateTime, Days, NaiveTime, TimeDelta, Utc};
use gateway_core::engine::credential::ProviderAccountId;
use provider_openai::credential::{CodexUsageStatisticsError, CodexUsageStatisticsServiceTier};
use serde_json::json;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::{create_account, quota_service_with_base_url};
use crate::support::MemoryAccountStore;
const PERSONAL_MODEL_REPORT_PATH: &str = "/wham/usage/daily-token-usage-breakdown";
const PERSONAL_TOTALS_REPORT_PATH: &str = "/wham/analytics/daily-workspace-usage-counts";
fn cycle_anchor() -> (DateTime<Utc>, DateTime<Utc>) {
let start_date = Utc::now()
.date_naive()
.checked_sub_days(Days::new(1))
.expect("cycle start date");
let start_at = DateTime::<Utc>::from_naive_utc_and_offset(
start_date.and_time(NaiveTime::from_hms_opt(15, 50, 0).expect("cycle time")),
Utc,
);
(
start_at,
start_at
.checked_add_signed(TimeDelta::days(7))
.expect("cycle end"),
)
}
async fn mount_quota(server: &MockServer, reset_at: DateTime<Utc>) {
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"primary_window": {
"used_percent": 20,
"reset_at": reset_at.timestamp(),
"limit_window_seconds": 604800
}
}
})))
.mount(server)
.await;
}
async fn service(
account_id: &str,
server: &MockServer,
reset_at: DateTime<Utc>,
) -> (
Arc<MemoryAccountStore>,
provider_openai::credential::CodexCredentialQuotaService,
) {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, account_id).await;
mount_quota(server, reset_at).await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
(store, service)
}
#[tokio::test]
async fn personal_report_allocates_tokens_on_the_server_and_preserves_official_totals() {
let server = MockServer::start().await;
let (start_at, end_at) = cycle_anchor();
let report_date = Utc::now().date_naive().to_string();
Mock::given(method("GET"))
.and(path(PERSONAL_MODEL_REPORT_PATH))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": [{
"date": report_date,
"models": [
{"model": "gpt-5.6-sol", "speed": "fast", "credits": 75},
{"model": "gpt-5.6-terra", "speed": "standard", "credits": 25}
]
}]})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path(PERSONAL_TOTALS_REPORT_PATH))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": [{
"date": report_date,
"uncached_text_input_tokens": 101,
"cached_text_input_tokens": 203,
"text_output_tokens": 17,
"text_total_tokens": 321
}]})))
.mount(&server)
.await;
let (store, service) = service("acct_personal_statistics", &server, end_at).await;
let account = store
.account("acct_personal_statistics")
.expect("personal account");
let statistics = service
.usage_statistics(account.id(), 0, 0)
.await
.expect("personal statistics");
assert_eq!(statistics.cycle.start_at, start_at);
assert_eq!(statistics.summary.tokens.uncached_input, 101);
assert_eq!(statistics.summary.tokens.cached_input, 203);
assert_eq!(statistics.summary.tokens.output, 17);
assert_eq!(statistics.summary.tokens.total, 321);
assert_eq!(
statistics
.models
.iter()
.map(|model| model.tokens.uncached_input)
.sum::<u64>(),
101
);
assert_eq!(
statistics
.models
.iter()
.map(|model| model.tokens.total)
.sum::<u64>(),
321
);
assert!(
statistics
.models
.iter()
.all(|model| model.has_estimated_allocation)
);
assert!(statistics.models.iter().any(|model| {
model.model == "gpt-5.6-sol" && model.service_tier == CodexUsageStatisticsServiceTier::Fast
}));
assert!(statistics.models.iter().any(|model| {
model.model == "gpt-5.6-terra"
&& model.service_tier == CodexUsageStatisticsServiceTier::Standard
}));
}
#[tokio::test]
async fn cycle_offset_fetches_only_the_selected_cycle_with_boundary_padding() {
let server = MockServer::start().await;
let (current_start_at, current_end_at) = cycle_anchor();
let selected_date = current_start_at
.date_naive()
.checked_sub_days(Days::new(14))
.expect("selected date");
Mock::given(method("GET"))
.and(path(PERSONAL_MODEL_REPORT_PATH))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": [{
"date": selected_date.to_string(),
"models": [{
"model": "gpt-5.6-terra",
"speed": "standard",
"credits": 10
}]
}]})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path(PERSONAL_TOTALS_REPORT_PATH))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": [{
"date": selected_date.to_string(),
"uncached_text_input_tokens": 10,
"text_total_tokens": 10
}]})))
.mount(&server)
.await;
let (store, service) = service("acct_offset_statistics", &server, current_end_at).await;
let account = store
.account("acct_offset_statistics")
.expect("offset account");
let statistics = service
.usage_statistics(account.id(), -2, 0)
.await
.expect("historical statistics");
assert_eq!(statistics.cycle.offset, -2);
assert!(!statistics.cycle.is_current);
assert!(statistics.cycle.can_go_next);
assert!(statistics.cycle.used_percent.is_none());
assert_eq!(statistics.daily[0].date, selected_date);
let requests = server.received_requests().await.expect("recorded requests");
let report = requests
.iter()
.find(|request| request.url.path() == PERSONAL_TOTALS_REPORT_PATH)
.expect("personal totals report request");
let query = report
.url
.query_pairs()
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect::<BTreeMap<_, _>>();
let expected_start = current_start_at
.date_naive()
.checked_sub_days(Days::new(15))
.expect("query start");
let expected_end = current_start_at
.date_naive()
.checked_sub_days(Days::new(6))
.expect("query end");
assert_eq!(query.get("start_date"), Some(&expected_start.to_string()));
assert_eq!(query.get("end_date"), Some(&expected_end.to_string()));
assert_eq!(query.get("group_by").map(String::as_str), Some("day"));
let invalid_account = ProviderAccountId::new("acct_invalid_offset").expect("account ID");
assert!(matches!(
service.usage_statistics(&invalid_account, -9, 0).await,
Err(CodexUsageStatisticsError::InvalidRequest)
));
}
@@ -4,6 +4,7 @@ use std::sync::Arc;
use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use bytes::Bytes;
use chrono::Utc;
use futures::{SinkExt, StreamExt};
use gateway_core::engine::continuation::{
@@ -22,8 +23,8 @@ use gateway_core::engine::{
use gateway_core::error::{ContinuationFailure, ProviderErrorKind};
use gateway_core::event::GatewayEvent;
use gateway_core::operation::{
CapabilityRequirements, GenerateRequest, Operation, OperationKind, ProtocolPayload,
ProviderSessionState,
CapabilityRequirements, GenerateRequest, ImageRequest, ImageRequestKind, Operation,
OperationKind, ProtocolPayload, ProviderSessionState, RawJsonPayload,
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
@@ -47,7 +48,7 @@ use tokio::{
};
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::protocol::{CloseFrame, frame::coding::CloseCode};
use wiremock::matchers::{header, method, path};
use wiremock::matchers::{body_bytes, header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{
@@ -291,6 +292,35 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
ProviderRequest::new(operation, plan.candidates()[0].clone())
}
fn planned_provider_endpoint_request(provider_name: &str, operation: Operation) -> ProviderRequest {
let provider = ProviderKind::new(provider_name).expect("provider");
let account_scope = Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_provider_contract").expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)]))),
ClientRoutingScope::all_accounts(),
));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
account_policy(),
vec![provider.clone()],
Vec::new(),
Vec::new(),
)
.expect("snapshot");
let plan = snapshot
.plan_provider_endpoint(
&provider,
&operation,
account_scope,
&RoutingContext::default(),
)
.expect("provider endpoint routing plan");
ProviderRequest::new(operation, plan.candidates()[0].clone())
}
fn contract_account_scope() -> Arc<FrozenAccountScope> {
let provider = ProviderKind::new("openai").expect("provider");
let accounts = [
@@ -314,6 +344,8 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_local_affinity",
"acct_metadata_new",
"acct_metadata_old",
"acct_unknown_continuation",
"acct_unknown_turn_state",
"acct_prefetch_limit",
"acct_presentation",
"acct_provider_contract",
@@ -322,7 +354,8 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_scope_same",
"acct_semantic_failure",
"acct_session_affinity",
"acct_subagent",
"acct_subagent_a",
"acct_subagent_b",
"acct_success_exhausted",
"acct_thread_spawn_affinity",
"acct_truncated_stream",
@@ -426,6 +459,24 @@ fn replay_any_context(request_id: &str, owner_account_id: &str) -> AttemptContex
.with_continuation_attempt(ContinuationAttempt::ReplayAny)
}
fn external_continuation_context(request_id: &str) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
Some(ContinuationBinding::External(PreviousResponseId::new(
"external-previous-response",
))),
CancellationToken::new(),
)
}
async fn capture_scoped_http_request(
request_id: &str,
selected_account_id: &str,
@@ -772,6 +823,169 @@ async fn generate_without_an_eligible_openai_account_fails_before_network_io() {
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn image_endpoints_bypass_only_the_text_catalog_and_preserve_the_current_codex_wire() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_raw(
br#"{"models":[{"slug":"gpt-5.4","display_name":"GPT-5.4","supported_in_api":true}]}"#,
"application/json",
))
.expect(1)
.mount(&server)
.await;
let cases = [
(
ImageRequestKind::Generation,
"/codex/images/generations",
br#"{ "model":"gpt-image-future", "prompt":"a lighthouse", "background":"transparent", "future_option":{"schema":2}, "future_integer":9007199254740993 }"#.as_slice(),
br#"{ "created": 1787212800, "data": [{"b64_json":"AAEC"}], "future": 9007199254740993 }"#.as_slice(),
),
(
ImageRequestKind::Edit,
"/codex/images/edits",
br#"{"model":"gpt-image-2","images":[{"image_url":"data:image/png;base64,AAEC"}],"prompt":"add fog","prompt":"duplicate remains opaque"}"#.as_slice(),
br#"{"created":1787212801,"data":[{"b64_json":"AwQF"}],"quality":"high"}"#.as_slice(),
),
];
for (_, endpoint, body, response_body) in &cases {
Mock::given(method("POST"))
.and(path(*endpoint))
.and(header("originator", "codex_cli_rs"))
.and(header("x-codex-image-turn-id", "turn_image_contract"))
.and(body_bytes(body.to_vec()))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "application/json")
.insert_header("x-request-id", "upstream_image_request")
.insert_header("x-future-image-header", "preserved")
.set_body_raw(response_body.to_vec(), "application/json"),
)
.expect(1)
.mount(&server)
.await;
}
let provider = provider_with_base_url(&store, server.uri());
let catalog = provider
.query_model_capabilities()
.await
.expect("text model catalog");
assert!(
catalog
.iter()
.all(|model| model.upstream_model().as_str() != "gpt-image-2"),
"the image model must not be published as a text model"
);
for (index, (kind, _, body, expected_response)) in cases.iter().enumerate() {
let payload = RawJsonPayload::new("openai", Bytes::copy_from_slice(body))
.expect("image payload")
.with_context(Map::from_iter([(
"image_turn_id".to_owned(),
json!("turn_image_contract"),
)]));
let operation = Operation::GenerateImage(ImageRequest::from_raw_json(*kind, payload));
let request = planned_provider_endpoint_request("openai", operation);
let mut stream = provider
.execute(
request,
context(
&format!("req_image_contract_{index}"),
CancellationToken::new(),
),
)
.await
.expect("prepare image provider stream");
let mut raw_response = None;
let mut completed = false;
let mut observed_http_json = false;
while let Some(event) = stream.next().await {
let event = event.expect("image provider event");
completed |= event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)));
if let Some(raw) = event.wire_event().and_then(|wire| wire.raw_json_body()) {
assert!(raw_response.replace(raw.clone()).is_none());
}
if let Some(observation) = event.response_observation() {
observed_http_json |= observation.transport().as_str() == "http_json"
&& observation.status_code() == Some(200)
&& observation.client_headers().iter().any(|header| {
header.name() == "x-future-image-header"
&& header.value().as_ref() == b"preserved"
});
}
}
assert!(completed);
assert!(observed_http_json);
assert_eq!(raw_response.as_deref(), Some(*expected_response));
}
server.verify().await;
}
#[tokio::test]
async fn image_endpoint_returns_the_exact_upstream_error_response() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
let request_body =
br#"{ "model":"gpt-image-2", "images":[], "future_invalid":9007199254740993 }"#;
let response_body = br#"{ "error":{"message":"future image validation","type":"image_error","code":"future_code"}, "future":9007199254740993 }"#;
Mock::given(method("POST"))
.and(path("/codex/images/edits"))
.and(body_bytes(request_body.to_vec()))
.respond_with(
ResponseTemplate::new(422)
.insert_header("content-type", "application/problem+json")
.insert_header("x-future-image-error", "preserved")
.set_body_bytes(response_body.to_vec()),
)
.expect(1)
.mount(&server)
.await;
let payload =
RawJsonPayload::new("openai", Bytes::from_static(request_body)).expect("image payload");
let operation =
Operation::GenerateImage(ImageRequest::from_raw_json(ImageRequestKind::Edit, payload));
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_provider_endpoint_request("openai", operation),
context("req_image_error", CancellationToken::new()),
)
.await
.expect("prepare image provider stream");
let error = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("image rejection must surface an upstream response"),
}
};
assert_eq!(error.upstream_status(), Some(422));
let response = error
.client_visible_upstream_response()
.expect("raw upstream image error response");
assert_eq!(response.status(), 422);
assert_eq!(
response.content_type(),
Some(b"application/problem+json".as_slice())
);
assert_eq!(response.body().as_ref(), response_body);
assert!(response.headers().iter().any(|header| {
header.name() == "x-future-image-error" && header.value().as_ref() == b"preserved"
}));
server.verify().await;
}
#[tokio::test]
async fn capacity_selection_error_preserves_classification_and_retry_after() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1126,6 +1340,53 @@ async fn cross_account_scope_removes_only_account_bound_body_fields() {
);
}
#[tokio::test]
async fn cross_account_full_replay_should_preserve_the_complete_transcript() {
let transcript = json!([
{
"type": "reasoning",
"id": "reasoning-item-id",
"encrypted_content": "reasoning-ciphertext",
"summary": []
},
{
"type": "compaction",
"id": "compaction-item-id",
"encrypted_content": "compaction-ciphertext"
},
{
"type": "agent_message",
"id": "agent-message-id",
"content": [
{"type": "output_text", "text": "visible"},
{"type": "encrypted_content", "encrypted_content": "nested-ciphertext"}
]
},
{
"type": "function_call",
"id": "function-call-id",
"call_id": "call-id",
"name": "tool",
"arguments": "{}",
"encrypted_function_args": "function-args-ciphertext"
}
]);
let request = capture_scoped_http_request(
"req_scope_transcript",
"acct_scope_new",
"acct_scope_old",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), transcript.clone()),
]),
Map::new(),
)
.await;
let body: Value = captured_request_body(&request);
assert_eq!(body.get("input"), Some(&transcript));
}
#[tokio::test]
async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() {
let request = capture_scoped_http_request(
@@ -1166,40 +1427,19 @@ async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() {
}
#[tokio::test]
async fn cross_account_continuation_probe_should_preserve_delta_input_and_previous_response() {
async fn cross_account_continuation_should_require_client_replay_without_an_upstream_probe() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_scope_new").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let delta_input = json!([
{
"type": "message",
"id": "delta-message-id",
"encrypted_content": "opaque-delta-ciphertext",
"role": "user",
"content": "current turn only"
},
{
"type": "compaction",
"id": "delta-compaction-id",
"encrypted_content": "opaque-compaction-ciphertext"
}
]);
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), delta_input.clone()),
(
"input".to_owned(),
json!([{"role": "user", "content": "delta"}]),
),
(
"previous_response_id".to_owned(),
json!("client-previous-response"),
@@ -1239,41 +1479,137 @@ async fn cross_account_continuation_probe_should_preserve_delta_input_and_previo
)
.expect("provider session state"),
);
let mut stream = provider_with_base_url(&store, server.uri())
let result = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
replay_any_context("req_cross_account_probe", "acct_scope_old"),
)
.await;
let Err(error) = result else {
panic!("cross-account continuation must require a full client replay");
};
let detail = error
.client_visible_upstream_error()
.expect("client replay error detail");
let requests = server
.received_requests()
.await
.expect("prepare continuation probe stream");
.expect("captured upstream requests");
assert_eq!(
(
error.kind(),
error.send_state(),
error.continuation_failure(),
detail.code(),
detail.error_type(),
),
(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
Some(ContinuationFailure::HistoryUnavailable),
Some("previous_response_not_found"),
Some("invalid_request_error"),
)
);
assert!(requests.is_empty());
}
#[tokio::test]
async fn missing_affinity_continuation_should_require_client_replay_without_an_upstream_probe() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_unknown_continuation").await;
let server = MockServer::start().await;
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
(
"input".to_owned(),
json!([{"role": "user", "content": "delta"}]),
),
(
"previous_response_id".to_owned(),
json!("external-previous-response"),
),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([
("use_websocket".to_owned(), json!(false)),
("session_id".to_owned(), json!("expired-affinity-session")),
])),
);
let result = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
external_continuation_context("req_unknown_continuation"),
)
.await;
let Err(error) = result else {
panic!("unowned continuation must require a full client replay");
};
let requests = server
.received_requests()
.await
.expect("captured upstream requests");
assert_eq!(
error
.client_visible_upstream_error()
.and_then(|detail| detail.code()),
Some("previous_response_not_found")
);
assert!(requests.is_empty());
}
#[tokio::test]
async fn missing_affinity_full_request_should_clear_unowned_turn_state() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_unknown_turn_state").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("full transcript")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([
("use_websocket".to_owned(), json!(false)),
("session_id".to_owned(), json!("expired-turn-state-session")),
("turn_id".to_owned(), json!("turn-after-affinity-expiry")),
("turn_state".to_owned(), json!("unowned-turn-state")),
])),
);
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
context("req_unknown_turn_state", CancellationToken::new()),
)
.await
.expect("prepare full replay stream");
while let Some(event) = stream.next().await {
event.expect("continuation probe response");
event.expect("full replay response");
}
let requests = server
.received_requests()
.await
.expect("captured continuation probe request");
let body: Value = captured_request_body(&requests[0]);
.expect("captured full replay request");
assert_eq!(body.get("input"), Some(&delta_input));
assert_eq!(
body.get("previous_response_id"),
Some(&json!("upstream-previous-response"))
);
assert!(body.get("authorization").is_none());
assert!(body.get("turnMetadata").is_none());
assert!(
body.pointer("/client_metadata/x-codex-turn-state")
.is_none()
);
assert!(
body.pointer("/client_metadata/x-codex-turn-metadata")
.is_none()
);
assert_eq!(
body.pointer("/client_metadata/future"),
Some(&json!("keep"))
);
assert!(captured_header_values(&requests[0], "x-codex-turn-state").is_empty());
}
@@ -1444,9 +1780,10 @@ async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key()
}
#[tokio::test]
async fn subagent_requests_should_use_an_affinity_key_isolated_from_the_root_session() {
async fn subagent_requests_should_share_the_root_session_account_affinity_key() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent").await;
create_account(&store, "acct_subagent_a").await;
create_account(&store, "acct_subagent_b").await;
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
Mock::given(method("POST"))
@@ -1496,10 +1833,25 @@ async fn subagent_requests_should_use_an_affinity_key_isolated_from_the_root_ses
let keys = affinity.lookup_keys();
assert_eq!(keys.len(), 3);
assert_ne!(keys[0], keys[1], "subagent must not inherit the root key");
assert_eq!(
keys[1], keys[2],
"the same subagent should retain its own key"
assert!(
keys.iter().all(|key| key == &keys[0]),
"root and derived subagent requests must prefer the same account"
);
let requests = server
.received_requests()
.await
.expect("captured root and subagent requests");
let selected_accounts = requests
.iter()
.filter_map(|request| request.headers.get("chatgpt-account-id"))
.filter_map(|value| value.to_str().ok())
.collect::<Vec<_>>();
assert_eq!(selected_accounts.len(), 3);
assert!(
selected_accounts
.iter()
.all(|account| account == &selected_accounts[0]),
"root and subagents should route to the same preferred account"
);
}
@@ -1674,7 +2026,7 @@ async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_st
}
#[tokio::test]
async fn thread_spawn_children_should_have_independent_session_affinity_keys() {
async fn thread_spawn_children_should_share_the_parent_account_affinity_key() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_thread_spawn_affinity").await;
let affinity = Arc::new(MemorySessionAffinity::default());
@@ -1719,10 +2071,10 @@ async fn thread_spawn_children_should_have_independent_session_affinity_keys() {
let keys = affinity.lookup_keys();
assert_eq!(keys.len(), 5);
assert_ne!(keys[0], keys[1]);
assert_eq!(keys[0], keys[2]);
assert_ne!(keys[3], keys[4]);
assert_ne!(keys[0], keys[3]);
assert!(
keys.iter().all(|key| key == &keys[0]),
"child thread transport identities must not split parent account affinity"
);
}
#[tokio::test]
@@ -2346,7 +2698,7 @@ async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() {
}
#[tokio::test]
async fn previous_response_not_found_should_remain_client_visible_without_internal_replay_proof() {
async fn same_account_previous_response_not_found_should_remain_client_visible() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_client_history").await;
let server = MockServer::start().await;
@@ -2385,7 +2737,7 @@ async fn previous_response_not_found_should_remain_client_visible_without_intern
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
),
context("req_history_missing", CancellationToken::new()),
context_with_state_owner("req_history_missing", "acct_client_history"),
)
.await
.expect("prepare missing-history stream");
@@ -678,6 +678,55 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity {
})
}
fn claim_or_load<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
candidate_account_id: &'a ProviderAccountId,
_ttl: Duration,
) -> BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>> {
Box::pin(async move {
let mut bindings = self.bindings.lock().expect("session affinity lock");
Ok(bindings
.entry((
provider_kind.as_str().to_owned(),
key.expose_to_store().to_owned(),
))
.or_insert_with(|| candidate_account_id.clone())
.clone())
})
}
fn compare_and_bind<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected_account_id: &'a ProviderAccountId,
replacement_account_id: &'a ProviderAccountId,
_ttl: Duration,
) -> BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>> {
Box::pin(async move {
let binding_key = (
provider_kind.as_str().to_owned(),
key.expose_to_store().to_owned(),
);
let mut bindings = self.bindings.lock().expect("session affinity lock");
if bindings
.get(&binding_key)
.is_none_or(|current| current == expected_account_id)
{
bindings.insert(binding_key, replacement_account_id.clone());
return Ok(replacement_account_id.clone());
}
bindings.get(&binding_key).cloned().ok_or_else(|| {
ProviderStoreError::new(
gateway_core::provider_ports::ProviderStoreErrorKind::Unavailable,
"resolve in-memory provider session affinity",
)
})
})
}
fn clear<'a>(
&'a self,
provider_kind: &'a ProviderKind,
@@ -743,7 +743,7 @@ fn bare_response_failed_should_project_started_identity_with_the_fallback_model(
canonical_facts(failure.events()).as_slice(),
[GatewayEvent::Started(metadata)]
if metadata.response_id() == "resp_bare_failed"
&& metadata.model() == "fallback-model"
&& metadata.model() == Some("fallback-model")
));
}
@@ -195,7 +195,7 @@ fn bare_http_429_remains_a_temporary_rate_limit() {
}
#[test]
fn sub2api_openai_account_failure_matrix_is_preserved() {
fn openai_account_failure_matrix_is_preserved() {
let cases = [
(
400,
@@ -265,7 +265,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
format!("http://{address}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))));
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
let response = backend
.create_response(
@@ -480,7 +480,7 @@ async fn backend_websocket_should_drop_only_unrepresentable_opaque_header_values
format!("http://{address}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))));
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
client
.create_response(
@@ -932,7 +932,7 @@ async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_la
});
let profile = test_wire_profile();
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let client = CodexBackendClient::new(
reqwest::Client::builder()
.no_proxy()
@@ -300,7 +300,7 @@ async fn account_eviction_should_cancel_a_background_websocket_opening() {
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -354,7 +354,7 @@ async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() {
.unwrap();
assert_eq!(read, 0);
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -733,7 +733,7 @@ async fn concurrent_multi_conversation_continuations_should_select_each_exact_so
tokio::join!(first_continuation, second_continuation);
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -1190,7 +1190,7 @@ async fn cancelled_half_open_opening_should_allow_another_probe() {
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -1325,7 +1325,7 @@ async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() {
write_completed_sse_response(&mut second_http).await;
}
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -56,6 +56,7 @@ mod latency;
mod profile;
mod protocol;
mod request;
mod reset_credits;
mod usage;
mod websocket;
mod websocket_pool;
@@ -494,8 +495,6 @@ fn websocket_pool_config_for_tests(
CodexWebSocketPoolConfig {
enabled: true,
max_age: Duration::from_mins(1),
max_per_account: 8,
max_total: 64,
max_connecting: 16,
maintenance_interval,
ping_interval,

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