Compare commits

...
28 Commits
Author SHA1 Message Date
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
ZyyandCodex 53739ebe5c chore: bump version to 3.2.13
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 15:20:29 +08:00
ZyyandCodex 7f420257a6 fix(api): keep provider degradation out of healthz
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 15:14:28 +08:00
ZyyandCodex 4f85209dc9 fix(config): centralize runtime data paths
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 15:14:28 +08:00
ZyyandCodex cd89e2e8fc fix(admin): align quota metadata and usage layout
Co-authored-by: Codex <noreply@openai.com>
2026-08-20 15:14:28 +08:00
ZyyandCodex 0cb5a29699 chore: bump version to 3.2.12
Co-authored-by: Codex <noreply@openai.com>
2026-08-19 23:52:11 +08:00
ZyyandCodex 61788804b4 fix(openai): recover cross-account websocket continuations
Co-authored-by: Codex <noreply@openai.com>
2026-08-19 23:40:40 +08:00
ZyyandCodex 2751ea6aa4 fix(openai): preserve cross-account session recovery
Co-authored-by: Codex <noreply@openai.com>
2026-08-19 23:05:48 +08:00
ZyyandCodex 58fe66789c chore(deps): bump h2 to 0.4.16
Resolve RUSTSEC-2026-0258 reported by cargo audit.

Co-authored-by: Codex <noreply@openai.com>
2026-08-19 11:34:07 +08:00
ZyyandCodex 573b6b26db fix(openai): isolate quota buckets by limit id
Co-authored-by: Codex <noreply@openai.com>
2026-08-18 16:48:37 +08:00
ZyyandCodex d425428ebe chore: bump version to 3.2.11
Co-authored-by: Codex <noreply@openai.com>
2026-08-17 20:02:02 +08:00
ZyyandCodex b26d3b386a fix(admin): align quota diagnostics and usage layout
Co-authored-by: Codex <noreply@openai.com>
2026-08-17 20:00:29 +08:00
165 changed files with 8187 additions and 1450 deletions
+76 -70
View File
@@ -12,31 +12,35 @@
[![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>
Codex Proxy RS 在一个网关中管理 OpenAI 与 xAI 账号,为客户端提供统一鉴权、模型路由、额度状态、
请求观测和管理控制面。
> [!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 +60,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 +87,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 +133,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)
+6 -6
View File
@@ -1246,7 +1246,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -1686,9 +1686,9 @@ dependencies = [
[[package]]
name = "h2"
version = "0.4.15"
version = "0.4.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27"
dependencies = [
"atomic-waker",
"bytes",
@@ -2714,7 +2714,7 @@ dependencies = [
"once_cell",
"socket2",
"tracing",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -2984,7 +2984,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -3628,7 +3628,7 @@ dependencies = [
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
+3 -2
View File
@@ -27,8 +27,9 @@ impl LoadableConfig for GatewayConfig {
return Err(ConfigError::InvalidField("schema_version"));
}
self.host.resolve_and_validate(source_dir)?;
let runtime_data_dir = self.host.runtime_data_dir().to_path_buf();
self.store
.resolve_and_validate(source_dir)
.resolve_and_validate(&runtime_data_dir)
.map_err(|_| ConfigError::InvalidField("store"))?;
self.admin
.resolve_and_validate(source_dir)
@@ -37,7 +38,7 @@ impl LoadableConfig for GatewayConfig {
.resolve_and_validate(source_dir)
.map_err(|_| ConfigError::InvalidField("api"))?;
self.openai
.resolve_and_validate(source_dir)
.resolve_and_validate(&runtime_data_dir)
.map_err(|_| ConfigError::InvalidField("openai"))?;
self.xai
.resolve_and_validate(source_dir)
+8
View File
@@ -18,10 +18,18 @@ fn config_loader_should_load_complete_terminal_example() {
fn config_loader_should_resolve_paths_relative_to_config_file() {
let (config, _directory) = parse_config(&valid_config()).expect("resolved config");
let debug = format!("{config:?}");
assert!(debug.contains(".runtime/data"));
assert!(debug.contains(".runtime/logs"));
assert!(debug.contains("frontend/dist"));
}
#[test]
fn config_loader_should_reject_missing_runtime_data_dir() {
let config = valid_config().replace(" runtime_data_dir: '../.runtime/data'\n", "");
assert!(parse_config(&config).is_err());
}
#[test]
fn config_loader_should_inject_connection_passwords_into_urls() {
parse_config(&valid_config()).expect("Store validates password injection into both URLs");
@@ -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>,
@@ -8,6 +8,7 @@ use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
routing::{ProviderKind, UpstreamModelId},
};
use uuid::Uuid;
use super::{
AdminError, MutationActor, MutationContext, PageSize, Revision,
@@ -570,13 +571,34 @@ pub enum QuotaLocalUsageAttribution {
Unavailable,
}
/// Provider quota bucket 中窗口的官方位置语义。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderQuotaWindowRole {
Primary,
Secondary,
Monthly,
}
impl ProviderQuotaWindowRole {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Primary => "primary",
Self::Secondary => "secondary",
Self::Monthly => "monthly",
}
}
}
/// 一个 Provider quota 窗口的公共投影。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderQuotaWindow {
pub key: String,
pub group: String,
pub label: String,
pub source: Option<String>,
pub limit_id: Option<String>,
pub limit_name: Option<String>,
pub role: Option<ProviderQuotaWindowRole>,
pub local_usage_attribution: QuotaLocalUsageAttribution,
pub window_seconds: Option<u64>,
pub used_percent: Option<f64>,
@@ -605,6 +627,38 @@ pub struct ProviderQuota {
pub provider_data: Option<ProviderDocument>,
}
/// 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 {
/// 返回账号展开面板使用的当前额度窗口用量。
///
@@ -627,8 +681,9 @@ impl ProviderQuota {
/// 返回 Dashboard 使用的代表性额度比例。
///
/// 优先使用短周期窗口,再依次使用周、月和其它窗口;同一优先级取较高的已用比例。
/// 这里只解释跨 Provider 共享的窗口语义,绝不读取 Provider 私有 JSON。
/// 优先使用覆盖账号全部请求的窗口;同一归属范围内再依次使用短周期、周、月
/// 和其它窗口,同一优先级取较高的已用比例。这里只解释跨 Provider 共享的窗口
/// 语义,绝不读取 Provider 私有 JSON。
#[must_use]
pub fn representative_used_percent(&self) -> Option<f64> {
if self.limit_reached {
@@ -687,7 +742,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
@@ -702,13 +757,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,11 @@ 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,
};
/// Provider 管理失败的稳定分类。
@@ -28,6 +28,8 @@ pub enum ProviderAdminErrorKind {
NotFound,
Conflict,
Unavailable,
CredentialRefreshRequired,
BadGateway,
Internal,
}
@@ -144,6 +146,22 @@ pub trait ProviderAdmin: Send + Sync {
request: ProviderQuotaRequest,
) -> Result<ProviderQuota, ProviderAdminError>;
/// 查询 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,9 @@ use crate::{
observability::TimeRange,
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
PrepareCredentialRefresh, ProviderModels, ProviderQuota, ProviderQuotaRequest,
ProviderQuotaWindow, QuotaLocalUsageAttribution,
ConsumeProviderResetCredit, PrepareCredentialRefresh, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow, ProviderResetCreditResult,
ProviderResetCredits, QuotaLocalUsageAttribution,
},
},
ports::{provider::ProviderAdminRegistry, store::AccountStore},
@@ -79,6 +80,26 @@ pub trait AccountsService: Send + Sync {
refresh: bool,
) -> Result<AccountDirectoryItem, AdminError>;
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 +118,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 +135,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 +538,57 @@ impl AccountsService for DefaultAccountsService {
self.load_directory_item(account_id, refresh).await
}
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 +669,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
+24 -2
View File
@@ -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 {
@@ -165,7 +183,9 @@ fn representative_window_usage_should_accept_account_wide_rolling_window_without
key: "free-rolling-24h".to_owned(),
group: "shortTerm".to_owned(),
label: "日限额".to_owned(),
source: None,
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(86_400),
used_percent: None,
@@ -234,7 +254,9 @@ fn quota_window(
key: group.to_owned(),
group: group.to_owned(),
label: group.to_owned(),
source: None,
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds,
used_percent,
@@ -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,
@@ -1262,7 +1289,9 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
@@ -1391,7 +1420,9 @@ async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view()
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(100.0),
@@ -1436,7 +1467,9 @@ async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(100.0),
@@ -1482,7 +1515,9 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
@@ -1547,7 +1582,9 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
key: "core-primary".to_owned(),
group: "weekly".to_owned(),
label: "周额度".to_owned(),
source: Some("codex".to_owned()),
limit_id: Some("codex".to_owned()),
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(7 * 24 * 60 * 60),
used_percent: Some(1.0),
@@ -1559,8 +1596,10 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
ProviderQuotaWindow {
key: "codex-bengalfox-primary".to_owned(),
group: "weekly".to_owned(),
label: "GPT-5.3-Codex-Spark · 周额度".to_owned(),
source: Some("codex_bengalfox".to_owned()),
label: "周额度".to_owned(),
limit_id: Some("codex_bengalfox".to_owned()),
limit_name: Some("GPT-5.3-Codex-Spark".to_owned()),
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
window_seconds: Some(7 * 24 * 60 * 60),
used_percent: Some(0.0),
@@ -1655,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,10 +25,11 @@ use gateway_admin::model::{
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion, CredentialDeletionResult,
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow, RotateCredential,
StartAuthorization,
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation,
CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels,
ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult,
ProviderResetCredits, RotateCredential, StartAuthorization,
},
};
use gateway_core::{
@@ -37,6 +38,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::{
@@ -305,8 +307,12 @@ pub struct AccountQuotaView {
pub struct AccountQuotaWindowView {
pub key: String,
pub group: String,
pub limit_id: Option<String>,
pub limit_name: Option<String>,
pub role: Option<String>,
pub window_seconds: Option<u64>,
pub label_display: String,
pub window_label_display: String,
pub used_percent: Option<f64>,
pub used_percent_display: String,
pub limit_reached: bool,
@@ -472,6 +478,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)]
@@ -539,6 +587,67 @@ pub struct AccountQuotaData {
pub account: AccountView,
}
/// 主动额度重置卡列表响应。
#[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)]
@@ -1051,6 +1160,10 @@ where
post(batch_update_accounts::<S>),
)
.route("/api/admin/accounts/quota", get(account_quota::<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>),
@@ -1402,6 +1515,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>,
@@ -1603,7 +1758,9 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
key,
group,
label,
source: _,
limit_id,
limit_name,
role,
local_usage_attribution: _,
window_seconds,
used_percent,
@@ -1612,10 +1769,17 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
local_usage,
provider_data: _,
} = window;
let label_display = limit_name
.as_ref()
.map_or_else(|| label.clone(), |name| format!("{name} · {label}"));
AccountQuotaWindowView {
label_display: label,
label_display,
window_label_display: label,
key,
group,
limit_id,
limit_name,
role: role.map(|value| value.as_str().to_owned()),
window_seconds,
used_percent,
used_percent_display: used_percent
@@ -222,7 +222,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 +243,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 +351,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 +729,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;
+21 -3
View File
@@ -1,4 +1,4 @@
//! `/healthz` 对 Core、Store 与 Host worker 健康事实的聚合。
//! `/healthz` 对 Core、Store 与 Host 关键 worker 健康事实的聚合。
use std::sync::Arc;
use std::time::Duration;
@@ -6,8 +6,10 @@ use std::time::Duration;
use axum::{extract::State, http::StatusCode};
use futures::future::join_all;
use gateway_core::health::{
HealthProbe, HealthState, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState,
HealthProbe, HealthState, WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource,
WorkerRuntimeState,
};
use gateway_core::task::WorkerKind;
use crate::ApiState;
@@ -32,7 +34,13 @@ impl HealthStatus {
}
pub(crate) async fn healthy(&self) -> bool {
if !self.workers.snapshot().iter().all(worker_is_healthy) {
if !self
.workers
.snapshot()
.iter()
.filter(|worker| worker_affects_healthz(worker))
.all(worker_is_healthy)
{
return false;
}
let checks = self.probes.iter().map(|probe| probe.check());
@@ -54,6 +62,16 @@ pub(crate) async fn healthz(State(state): State<ApiState>) -> StatusCode {
}
}
fn worker_affects_healthz(worker: &WorkerHealthSnapshot) -> bool {
let WorkerHealthKey::Task(id) = &worker.key else {
return false;
};
!matches!(
id.kind(),
WorkerKind::OAuthRefresh | WorkerKind::QuotaCatalogHealth
)
}
fn worker_is_healthy(worker: &WorkerHealthSnapshot) -> bool {
match worker.state {
WorkerRuntimeState::Disabled | WorkerRuntimeState::Standby => true,
@@ -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::{
@@ -73,7 +74,6 @@ pub(super) fn response_header_is_forwardable(name: &str, connection_options: &[S
| "x-openai-organization"
| "x-openai-project"
| "x-codex-installation-id"
| "x-codex-turn-state"
| "x-codex-turn-metadata"
)
}
@@ -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> {
@@ -266,8 +266,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 +325,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()),
@@ -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",
+79 -1
View File
@@ -8,10 +8,12 @@ use futures::future::BoxFuture;
use gateway_core::{
engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, StartExecution,
StartedExecution,
StartProviderExecution, StartedExecution,
},
error::GatewayError,
health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState},
routing::PublicModelId,
task::{WorkerId, WorkerKind},
};
use tower::ServiceExt;
@@ -31,6 +33,75 @@ async fn healthz_should_return_no_content_when_all_inputs_are_healthy() {
assert_eq!(response.status(), StatusCode::NO_CONTENT);
}
#[tokio::test]
async fn healthz_should_ignore_provider_maintenance_worker_failure() {
let worker_health = StaticWorkerHealth(vec![worker_snapshot(
WorkerKind::QuotaCatalogHealth,
WorkerRuntimeState::BackingOff,
)]);
let response = crate::openai::api_router_with_worker_health(
Arc::new(UnusedExecution),
Arc::new(worker_health),
)
.await
.oneshot(
Request::builder()
.uri("/healthz")
.body(Body::empty())
.expect("health request"),
)
.await
.expect("health response");
assert_eq!(response.status(), StatusCode::NO_CONTENT);
}
#[tokio::test]
async fn healthz_should_reject_critical_worker_failure() {
let worker_health = StaticWorkerHealth(vec![worker_snapshot(
WorkerKind::RuntimeSnapshotReconciliation,
WorkerRuntimeState::BackingOff,
)]);
let response = crate::openai::api_router_with_worker_health(
Arc::new(UnusedExecution),
Arc::new(worker_health),
)
.await
.oneshot(
Request::builder()
.uri("/healthz")
.body(Body::empty())
.expect("health request"),
)
.await
.expect("health response");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}
#[derive(Clone)]
struct StaticWorkerHealth(Vec<WorkerHealthSnapshot>);
impl WorkerHealthSource for StaticWorkerHealth {
fn snapshot(&self) -> Vec<WorkerHealthSnapshot> {
self.0.clone()
}
}
fn worker_snapshot(kind: WorkerKind, state: WorkerRuntimeState) -> WorkerHealthSnapshot {
let id = WorkerId::try_new(kind, "health-test").expect("worker ID");
WorkerHealthSnapshot {
key: WorkerHealthKey::Task(id),
state,
consecutive_failures: 1,
completed_cycles: 0,
last_fencing_token: None,
last_success_at: None,
last_failure_at: None,
last_error: Some("test worker failure".to_owned()),
}
}
struct UnusedExecution;
impl ExecutionService for UnusedExecution {
@@ -49,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 }"#
);
}
+22 -1
View File
@@ -1,5 +1,6 @@
mod auth;
mod error;
mod images;
mod models;
mod responses;
mod router;
@@ -44,9 +45,29 @@ pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Ro
api_router_with_origins(execution, Vec::new()).await
}
pub(super) async fn api_router_with_worker_health(
execution: Arc<dyn ExecutionService>,
worker_health: Arc<dyn WorkerHealthSource>,
) -> axum::Router {
api_router_with_origins_and_worker_health(execution, Vec::new(), worker_health).await
}
pub(super) async fn api_router_with_origins(
execution: Arc<dyn ExecutionService>,
cors_allowed_origins: Vec<String>,
) -> axum::Router {
api_router_with_origins_and_worker_health(
execution,
cors_allowed_origins,
Arc::new(EmptyWorkerHealth),
)
.await
}
async fn api_router_with_origins_and_worker_health(
execution: Arc<dyn ExecutionService>,
cors_allowed_origins: Vec<String>,
worker_health: Arc<dyn WorkerHealthSource>,
) -> axum::Router {
let admin = crate::admin::AdminTestFixture::new().await;
gateway_api::initialize(
@@ -59,7 +80,7 @@ pub(super) async fn api_router_with_origins(
execution,
admin.services,
Vec::new(),
Arc::new(EmptyWorkerHealth),
worker_health,
Arc::new(TestLifecycle::default()),
)
.expect("API bundle")
@@ -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 {
@@ -788,6 +800,10 @@ async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresenta
let trace = Arc::new(Trace::default());
let headers = vec![
ProviderResponseHeader::new("x-models-etag", Bytes::from_static(b"models-v2")),
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-from-upstream"),
),
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"first")),
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"second")),
ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")),
@@ -834,6 +850,13 @@ async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresenta
.map(HeaderValue::as_bytes),
Some(b"\xffopaque".as_slice())
);
assert_eq!(
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok()),
Some("turn-state-from-upstream")
);
assert!(response.headers().get("authorization").is_none());
assert!(response.headers().get("connection").is_none());
assert!(response.headers().get("x-hop-secret").is_none());
@@ -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 {
@@ -411,6 +418,10 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
ProviderResponseHeader::new("x-future-header", Bytes::from_static(b"future-value")),
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"first")),
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"second")),
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-from-upstream"),
),
ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")),
ProviderResponseHeader::new("bad\0name", Bytes::from_static(b"unrepresentable")),
ProviderResponseHeader::new(
@@ -483,6 +494,10 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
// 官方 metadata 是 string map,同名多值只能保持既有的后值覆盖语义。
assert_eq!(headers.get("x-future-multi"), Some(&json!("second")));
assert_eq!(headers.get("x-request-id"), Some(&json!("req_upstream")));
assert_eq!(
headers.get("x-codex-turn-state"),
Some(&json!("turn-state-from-upstream"))
);
for omitted in [
"x-future-bytes",
"bad\0name",
@@ -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()
@@ -1162,10 +1171,11 @@ where
true
}
/// 原生续写的原账号在本地调度阶段已不可用时,改用完整输入重放。
/// 原生续写的原账号在本地调度阶段已不可用时,交给 Provider 执行跨账号恢复。
///
/// 这不是强制 Smart:下一次 attempt 仍使用请求配置的调度策略。只有网关持有
/// 对应 Provider 的可携带会话状态时,才可安全移除 native handle 并换号。
/// 对应 Provider 的会话状态时才进入恢复;是否保留 native handle、执行 probe
/// 或使用完整 transcript,由 Provider 自己的协议边界决定。
fn prepare_unavailable_native_continuation_replay(&mut self, error: &ProviderError) -> bool {
if self.account_selection.required_account().is_some()
|| self.continuation_attempt != ContinuationAttempt::Native
@@ -1481,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",
+10 -5
View File
@@ -362,9 +362,14 @@ pub enum ContinuationAttempt {
None,
/// 使用 Store 解析出的原生 handle 与账号绑定。
Native,
/// 在原账号上用完整连接内 transcript 重放。
/// 在原账号上执行 Provider 定义的恢复。
///
/// 只有持有可携带 transcript 的 Provider 才能清除 native handle 后重放;
/// 只持有 opaque state 的 Provider 必须保留 continuation 依赖或返回客户端错误。
ReplayOwner,
/// transcript 已可携带,允许选择其他账号。
/// 允许选择其他账号并执行 Provider 定义的恢复。
///
/// 该枚举本身不证明客户端 input 完整,也不授权 Provider 清除 previous-response。
ReplayAny,
}
@@ -555,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>,
@@ -577,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>,
}
@@ -590,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,
}
}
+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"),
@@ -493,6 +502,31 @@ fn bare_atomic_response_failed(response_id: &str) -> ProviderError {
.with_atomic_client_events(vec![failed])
}
fn atomic_history_unavailable(response_id: &str) -> ProviderError {
let failed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.failed".to_owned()),
json!({
"type": "response.failed",
"response": {
"id": response_id,
"status": "failed",
"error": {
"code": "previous_response_not_found",
"message": "Previous response was not found. Retrying the full request."
}
}
}),
)
.expect("history unavailable wire"),
);
ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
.with_status(400)
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
.with_atomic_client_events(vec![failed])
}
fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, ProviderError>> {
let mut usage = Usage::new();
usage.input_tokens = Some(12);
@@ -598,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()),
@@ -1481,7 +1515,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
}
#[test]
fn native_continuation_replays_owner_before_safely_switching_account() {
fn native_continuation_uses_provider_defined_owner_then_cross_account_recovery() {
let Operation::Generate(generate) = generate_operation() else {
panic!("generate operation");
};
@@ -1536,7 +1570,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
block_on(session.collect_uncommitted()).expect("provider-defined recovery succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let contexts = provider.contexts.lock().expect("contexts lock");
@@ -1561,7 +1595,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
}
#[test]
fn unavailable_native_continuation_replays_with_the_configured_account_policy() {
fn unavailable_native_continuation_uses_provider_recovery_with_the_configured_account_policy() {
let Operation::Generate(generate) = generate_operation() else {
panic!("generate operation");
};
@@ -1602,7 +1636,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
block_on(session.collect_uncommitted()).expect("provider-defined recovery succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let contexts = provider.contexts.lock().expect("contexts lock");
@@ -1625,6 +1659,76 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn continuation_history_failure_without_replay_proof_is_delivered_without_hidden_retry() {
let Operation::Generate(generate) = generate_operation() else {
panic!("generate operation");
};
let operation = Operation::Generate(generate.with_provider_session_state(
ProviderSessionState::new("openai", Map::new()).expect("provider session state"),
));
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![
Script::ObservedStream {
account_id: "acct_one",
items: vec![Err(atomic_history_unavailable("response-history-missing"))],
},
Script::Stream {
account_id: "acct_two",
items: complete_stream(None),
},
]);
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
Some(ContinuationBinding::Pinned(continuation)),
CancellationToken::new(),
))
.expect("start execution");
let terminal = block_on(session.next_event())
.expect("history failure remains deliverable")
.expect("terminal history failure batch");
assert_eq!(
terminal.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
let events = terminal.into_provider_events();
let failure = events
.iter()
.filter_map(ProviderEvent::wire_event)
.find(|wire| wire.event_type() == Some("response.failed"))
.expect("raw response.failed event");
assert_eq!(
failure
.data()
.pointer("/response/error/code")
.and_then(Value::as_str),
Some("previous_response_not_found")
);
block_on(session.commit_downstream(Some(200))).expect("commit terminal history failure");
assert!(matches!(
block_on(session.next_event()),
Err(EngineError::Provider(_))
));
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 1);
assert_eq!(
contexts[0].continuation_attempt(),
ContinuationAttempt::Native
);
}
#[test]
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
let operation = operation(RetrySafety::Idempotent);
@@ -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
);
}
}
+12 -1
View File
@@ -41,6 +41,7 @@ pub fn load_config<T: LoadableConfig>() -> Result<T, ConfigError> {
#[serde(deny_unknown_fields)]
pub struct HostConfig {
pub listen: ListenConfig,
pub runtime_data_dir: PathBuf,
pub logging: LoggingConfig,
#[serde(default)]
pub system_update: SystemUpdateConfig,
@@ -76,11 +77,21 @@ impl HostConfig {
"host.worker_shutdown_timeout_seconds",
));
}
if self.runtime_data_dir.as_os_str().is_empty() {
return Err(ConfigError::InvalidField("host.runtime_data_dir"));
}
resolve_relative_path(source_dir, &mut self.runtime_data_dir);
self.logging.resolve_and_validate(source_dir)?;
self.system_update.resolve_and_validate(source_dir)?;
self.system_update
.resolve_and_validate(source_dir, &self.runtime_data_dir)?;
Ok(())
}
#[must_use]
pub fn runtime_data_dir(&self) -> &Path {
&self.runtime_data_dir
}
#[must_use]
pub const fn drain_timeout(&self) -> Duration {
Duration::from_secs(self.drain_timeout_seconds)
@@ -81,10 +81,10 @@ impl Default for SystemUpdateConfig {
});
let update_state_file = environment_value("CPR_UPDATE_STATE_FILE")
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/app/.runtime/data/update-state.json"));
.unwrap_or_else(|| PathBuf::from("update-state.json"));
let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE")
.map(PathBuf::from)
.unwrap_or_else(|| update_state_file.with_extension("lock"));
.unwrap_or_else(|| update_state_file.with_file_name("update.lock"));
let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR")
.map(PathBuf::from)
.unwrap_or_else(|| state::default_temp_dir(&update_state_file));
@@ -120,21 +120,28 @@ impl Default for SystemUpdateConfig {
}
impl SystemUpdateConfig {
pub(crate) fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ConfigError> {
for path in [
self.executable_path.as_mut(),
Some(&mut self.web_dist_dir),
Some(&mut self.update_state_file),
Some(&mut self.update_lock_file),
Some(&mut self.update_temp_dir),
]
.into_iter()
.flatten()
pub(crate) fn resolve_and_validate(
&mut self,
source_dir: &Path,
runtime_data_dir: &Path,
) -> Result<(), ConfigError> {
for path in [self.executable_path.as_mut(), Some(&mut self.web_dist_dir)]
.into_iter()
.flatten()
{
if path.is_relative() {
*path = source_dir.join(&*path);
}
}
for path in [
&mut self.update_state_file,
&mut self.update_lock_file,
&mut self.update_temp_dir,
] {
if path.is_relative() {
*path = runtime_data_dir.join(&*path);
}
}
if self.version.trim().is_empty() {
return Err(ConfigError::InvalidField("host.system_update.version"));
}
+38 -3
View File
@@ -34,6 +34,40 @@ fn host_config_resolves_only_host_owned_relative_paths() {
);
}
#[test]
fn host_config_resolves_runtime_data_dir_relative_to_configuration() {
let mut config = valid_config();
config
.resolve_and_validate(std::path::Path::new("/srv/gateway"))
.expect("valid host config");
assert_eq!(
config.runtime_data_dir(),
PathBuf::from("/srv/gateway/runtime-data")
);
}
#[test]
fn host_config_derives_update_paths_from_runtime_data_dir() {
let mut config = valid_config();
config
.resolve_and_validate(std::path::Path::new("/srv/gateway"))
.expect("valid host config");
assert_eq!(
(
config.system_update.update_state_file,
config.system_update.update_lock_file,
config.system_update.update_temp_dir,
),
(
PathBuf::from("/srv/gateway/runtime-data/update-state.json"),
PathBuf::from("/srv/gateway/runtime-data/update.lock"),
PathBuf::from("/srv/gateway/runtime-data/update-tmp"),
)
);
}
#[test]
fn host_config_rejects_zero_drain_window() {
let mut config = valid_config();
@@ -48,9 +82,9 @@ fn host_config_rejects_zero_drain_window() {
fn valid_config() -> HostConfig {
let system_update = SystemUpdateConfig {
update_state_file: PathBuf::from(".runtime/update-state.json"),
update_lock_file: PathBuf::from(".runtime/update-state.lock"),
update_temp_dir: PathBuf::from(".runtime/update-tmp"),
update_state_file: PathBuf::from("update-state.json"),
update_lock_file: PathBuf::from("update.lock"),
update_temp_dir: PathBuf::from("update-tmp"),
..SystemUpdateConfig::default()
};
HostConfig {
@@ -58,6 +92,7 @@ fn valid_config() -> HostConfig {
host: "127.0.0.1".to_owned(),
port: 8080,
},
runtime_data_dir: PathBuf::from("runtime-data"),
logging: LoggingConfig {
level: "info".to_owned(),
stdout: true,
@@ -23,6 +23,7 @@ fn logging_requires_at_least_one_sink() {
host: "127.0.0.1".to_owned(),
port: 8080,
},
runtime_data_dir: PathBuf::from("/tmp/runtime-data"),
logging: LoggingConfig {
level: "info".to_owned(),
stdout: false,
@@ -98,6 +99,7 @@ fn write_oauth_recovery_log(directory: PathBuf) {
host: "127.0.0.1".to_owned(),
port: 8080,
},
runtime_data_dir: PathBuf::from("/tmp/runtime-data"),
logging: LoggingConfig {
level: "off".to_owned(),
stdout: false,
@@ -208,11 +208,24 @@ pub struct RateLimitWindow {
pub reset_at: Option<i64>,
}
/// 限流项键的来源。
///
/// HTTP headers 和事件里的 `metered_limit_name` 提供稳定 `limit_id`;
/// WebSocket `additional_rate_limits` 在部分协议版本里只以可读名称为 map key,
/// 此时必须由持久化层用既有名称元数据解析,不能把名称误当成稳定 ID。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RateLimitKeySource {
LimitId,
LimitName,
}
/// 单个计量项的限流信息。
#[derive(Debug, Clone, PartialEq)]
pub struct RateLimitDetails {
/// 稳定的计量项 ID。
pub limit_id: String,
/// [`limit_id`](Self::limit_id) 是稳定 ID,还是名称生成的临时解析键。
pub key_source: RateLimitKeySource,
/// 上游提供的可读名称。
pub limit_name: Option<String>,
/// 当前请求是否被允许。
@@ -341,20 +354,40 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
return None;
}
let limit_id = rate_limit_name(value)
.map(normalize_limit_id)
let explicit_limit_id = top_level_rate_limit_id(value).map(normalize_limit_id);
let explicit_limit_name = non_empty_string(value, "limit_name");
let default_limit_id = explicit_limit_id
.clone()
.unwrap_or_else(|| "codex".to_string());
let limit_name = value
.get("metered_limit_name")
.and_then(Value::as_str)
.and_then(|_| value.get("limit_name"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|name| !name.is_empty())
.map(ToString::to_string);
let details = value
.get("rate_limits")
.and_then(|details| parse_details_from_object(details, &limit_id, limit_name));
let default_details = value.get("rate_limits").and_then(|details| {
parse_details_from_object(
details,
&default_limit_id,
explicit_limit_name,
RateLimitKeySource::LimitId,
)
});
let mut limits = parse_additional_rate_limits(value);
let matching_additional_ids = default_details
.as_ref()
.map(|default| {
limits
.iter()
.filter(|(_, additional)| same_rate_limit_facts(default, additional))
.map(|(limit_id, _)| limit_id.clone())
.collect::<Vec<_>>()
})
.unwrap_or_default();
// 当前 WebSocket 协议会把活动具名桶同时镜像到顶层 `rate_limits`,但不附
// `metered_limit_name`。只要 additional 中存在相同事实,顶层就不是 core 桶。
let default_is_additional_alias =
explicit_limit_id.is_none() && !matching_additional_ids.is_empty();
if let Some(details) = default_details
&& !default_is_additional_alias
{
limits.insert(default_limit_id.clone(), details);
}
let credits = value.get("credits").and_then(parse_credits_from_object);
let plan_type = value
.get("plan_type")
@@ -363,17 +396,18 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
.filter(|plan| !plan.is_empty())
.map(ToString::to_string);
if details.is_none() && credits.is_none() && plan_type.is_none() {
if limits.is_empty() && credits.is_none() && plan_type.is_none() {
return None;
}
let mut limits = BTreeMap::new();
if let Some(details) = details {
limits.insert(limit_id.clone(), details);
}
let active_limit = explicit_limit_id.or_else(|| match matching_additional_ids.as_slice() {
[limit_id] => Some(limit_id.clone()),
[] if !limits.is_empty() => Some(default_limit_id),
_ => None,
});
Some(ParsedRateLimits {
limits,
active_limit: Some(limit_id),
active_limit,
credits,
plan_type,
promo_message: value
@@ -597,6 +631,7 @@ fn parse_details_from_lookup(
}
Some(RateLimitDetails {
limit_id: limit_id.to_string(),
key_source: RateLimitKeySource::LimitId,
limit_name,
allowed,
limit_reached,
@@ -687,6 +722,7 @@ fn parse_details_from_object(
value: &Value,
limit_id: &str,
limit_name: Option<String>,
key_source: RateLimitKeySource,
) -> Option<RateLimitDetails> {
let primary = value.get("primary").and_then(parse_window_from_object);
let secondary = value.get("secondary").and_then(parse_window_from_object);
@@ -697,6 +733,7 @@ fn parse_details_from_object(
}
Some(RateLimitDetails {
limit_id: limit_id.to_string(),
key_source,
limit_name,
allowed,
limit_reached,
@@ -752,11 +789,59 @@ fn value_as_positive_i64(value: &Value) -> Option<i64> {
.filter(|value| *value > 0)
}
fn rate_limit_name(value: &Value) -> Option<&str> {
fn rate_limit_id(value: &Value) -> Option<&str> {
value
.get("metered_limit_name")
.or_else(|| value.get("limit_name"))
.or_else(|| value.get("limit_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn top_level_rate_limit_id(value: &Value) -> Option<&str> {
rate_limit_id(value).or_else(|| {
value
.get("limit_name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
})
}
fn non_empty_string(value: &Value, field: &str) -> Option<String> {
value
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn parse_additional_rate_limits(value: &Value) -> BTreeMap<String, RateLimitDetails> {
value
.get("additional_rate_limits")
.and_then(Value::as_object)
.into_iter()
.flatten()
.filter_map(|(map_key, value)| {
let explicit_limit_id = rate_limit_id(value).map(normalize_limit_id);
let (limit_id, key_source) = explicit_limit_id.map_or_else(
|| (normalize_limit_id(map_key), RateLimitKeySource::LimitName),
|limit_id| (limit_id, RateLimitKeySource::LimitId),
);
let limit_name = non_empty_string(value, "limit_name")
.or_else(|| Some(map_key.trim().to_owned()).filter(|name| !name.is_empty()));
parse_details_from_object(value, &limit_id, limit_name, key_source)
.map(|details| (limit_id, details))
})
.collect()
}
fn same_rate_limit_facts(left: &RateLimitDetails, right: &RateLimitDetails) -> bool {
left.allowed == right.allowed
&& left.limit_reached == right.limit_reached
&& left.primary == right.primary
&& left.secondary == right.secondary
}
fn normalize_limit_id(value: impl AsRef<str>) -> String {
@@ -1,7 +1,7 @@
use gateway_protocol::openai::events::{
RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage, is_codex_quota_header_name,
is_rate_limit_header_name, parse_rate_limit_headers, parse_rate_limits_event,
rate_limits_to_header_pairs, retry_after_seconds_from_body,
RateLimitKeySource, RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage,
is_codex_quota_header_name, is_rate_limit_header_name, parse_rate_limit_headers,
parse_rate_limits_event, rate_limits_to_header_pairs, retry_after_seconds_from_body,
};
use serde_json::json;
@@ -483,6 +483,58 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
);
}
#[test]
fn parse_rate_limits_event_should_preserve_name_keyed_additional_limit() {
let event = json!({
"type": "codex.rate_limits",
"plan_type": "pro",
"rate_limits": {
"allowed": true,
"limit_reached": false,
"primary": {
"used_percent": 0,
"window_minutes": 10080,
"reset_after_seconds": 354916,
"reset_at": 1787396838
},
"secondary": null
},
"additional_rate_limits": {
"GPT-5.3-Codex-Spark": {
"allowed": true,
"limit_reached": false,
"primary": {
"used_percent": 0,
"window_minutes": 10080,
"reset_after_seconds": 354916,
"reset_at": 1787396838
},
"secondary": null
}
},
"credits": null
});
let parsed = parse_rate_limits_event(&event).expect("rate-limit event");
let details = parsed
.limits
.get("gpt_5.3_codex_spark")
.expect("name-keyed additional limit");
assert_eq!(parsed.active_limit.as_deref(), Some("gpt_5.3_codex_spark"));
assert_eq!(details.key_source, RateLimitKeySource::LimitName);
assert_eq!(details.limit_name.as_deref(), Some("GPT-5.3-Codex-Spark"));
assert_eq!(
details.primary,
Some(RateLimitWindow {
used_percent: 0.0,
window_minutes: Some(10080),
reset_at: Some(1787396838),
})
);
assert!(!parsed.limits.contains_key("codex"));
}
#[test]
fn parse_rate_limits_event_should_reject_unrelated_event_type() {
for event in [
+2 -2
View File
@@ -50,7 +50,7 @@ impl StoreBundle {
pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
const REDIS_NAMESPACE: &str = "codex-proxy-rs";
config.resolve_and_validate(std::path::Path::new("."))?;
config.validate_resolved()?;
let pool = postgres::connect_and_migrate(&config.database_url()?, config.pool).await?;
let redis_client = ::redis::Client::open(config.redis_url()?)
.map_err(|_| redis_unavailable("create Redis client"))?;
@@ -209,7 +209,7 @@ pub(crate) fn backup_ports(
config: &StoreConfig,
) -> StoreResult<BackupStorePorts> {
let staging = Arc::new(backup::staging::StagingArea::open(
std::path::PathBuf::from(".runtime/data/backup-staging"),
config.backup_staging_dir().to_path_buf(),
backup::staging::DEFAULT_MAX_ARCHIVE_BYTES,
)?);
let repository = Arc::new(postgres::PgBackupRepository::new(pool));
+28 -1
View File
@@ -1,5 +1,7 @@
//! Store 启动配置、环境变量解析与校验。
use std::path::{Path, PathBuf};
use super::*;
pub(crate) const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL";
@@ -15,6 +17,8 @@ pub struct StoreConfig {
pub(crate) redis: StoreConnectionConfig,
#[serde(default)]
pub(crate) pool: StorePoolConfig,
#[serde(skip)]
backup_staging_dir: PathBuf,
}
/// PostgreSQL 连接池预算;acquire 超时决定池耗尽时快速失败而非排队积压。
@@ -47,7 +51,18 @@ impl StorePoolConfig {
}
impl StoreConfig {
pub fn resolve_and_validate(&mut self, _source_dir: &std::path::Path) -> StoreResult<()> {
pub fn resolve_and_validate(&mut self, runtime_data_dir: &Path) -> StoreResult<()> {
if runtime_data_dir.as_os_str().is_empty() {
return Err(StoreError::InvalidData {
entity: "store config",
message: "runtime_data_dir must not be empty".to_owned(),
});
}
self.backup_staging_dir = runtime_data_dir.join("backup-staging");
self.validate_resolved()
}
pub(crate) fn validate_resolved(&mut self) -> StoreResult<()> {
if let Some(url) = optional_environment_value(DATABASE_URL_ENV)? {
self.database.url = url;
}
@@ -63,6 +78,12 @@ impl StoreConfig {
self.database.validate("database")?;
self.redis.validate("redis")?;
self.pool.validate()?;
if self.backup_staging_dir.as_os_str().is_empty() {
return Err(StoreError::InvalidData {
entity: "store config",
message: "runtime_data_dir was not resolved".to_owned(),
});
}
Ok(())
}
@@ -73,6 +94,12 @@ impl StoreConfig {
pub(crate) fn redis_url(&self) -> StoreResult<String> {
self.redis.connection_url("redis")
}
/// 返回由统一运行数据根目录派生的备份暂存目录。
#[must_use]
pub fn backup_staging_dir(&self) -> &Path {
&self.backup_staging_dir
}
}
pub(crate) fn optional_environment_value(name: &'static str) -> StoreResult<Option<String>> {
+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",
@@ -367,7 +367,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
require_nonempty(ENTITY, "account_id", account_id)?;
let result = sqlx::query(
"update provider_accounts
set quota_observed_at = $3
set quota_observed_at = $3, updated_at = greatest(now(), $3)
where id = $1 and credential_revision = $2
and provider_quota_json is not null
and (quota_observed_at is null or quota_observed_at <= $3)",
@@ -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";
@@ -0,0 +1,32 @@
use gateway_store::StoreConfig;
const PASSWORD: &str = "111111111111111111111111111111111111111111111111";
#[test]
fn store_config_should_derive_backup_staging_from_runtime_data_dir() {
let root = tempfile::tempdir().expect("runtime data root");
let mut config = valid_config();
config
.resolve_and_validate(root.path())
.expect("valid Store configuration");
assert_eq!(
config.backup_staging_dir(),
root.path().join("backup-staging")
);
}
fn valid_config() -> StoreConfig {
serde_json::from_value(serde_json::json!({
"database": {
"url": "postgres://codex_proxy@127.0.0.1:5432/codex_proxy",
"password": PASSWORD,
},
"redis": {
"url": "redis://127.0.0.1:6379/",
"password": PASSWORD,
},
}))
.expect("test Store configuration")
}
@@ -1,3 +1,4 @@
mod config;
mod postgres;
mod redis;
mod s3;
@@ -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,
@@ -173,7 +173,7 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
}
#[tokio::test]
async fn quota_observation_touch_only_advances_the_refresh_time() {
async fn quota_observation_touch_preserves_quota_state_and_advances_account_update_time() {
let Some(database) = TestDatabase::create("provider_account_quota_touch").await else {
return;
};
@@ -188,7 +188,7 @@ async fn quota_observation_touch_only_advances_the_refresh_time() {
.checked_sub(Duration::from_secs(60))
.expect("old observation time");
let touched_at = observed_at
.checked_add(Duration::from_secs(30))
.checked_add(Duration::from_secs(90))
.expect("new observation time");
let reset_at = observed_at
.checked_add(Duration::from_secs(3_600))
@@ -272,7 +272,11 @@ async fn quota_observation_touch_only_advances_the_refresh_time() {
stored.5.map(|value| value.timestamp_micros()),
Some(chrono::DateTime::<Utc>::from(reset_at).timestamp_micros())
);
assert_eq!(stored.6, account_updated_at);
assert_eq!(
stored.6.timestamp_micros(),
chrono::DateTime::<Utc>::from(touched_at).timestamp_micros()
);
assert!(stored.6 > account_updated_at);
database.close().await;
}
@@ -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,
}
@@ -21,12 +21,18 @@ async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() {
return;
};
let database = TestDatabase::create(&database_url).await;
let config = store_config(&database.url, &redis_url);
let runtime_data = tempfile::tempdir().expect("Store runtime data");
let mut config = store_config(&database.url, &redis_url);
config
.resolve_and_validate(runtime_data.path())
.expect("resolved Store config");
let mut first = initialize(config.clone())
.await
.expect("first Store bundle");
let second = initialize(config).await.expect("second Store bundle");
assert!(runtime_data.path().join("backup-staging").is_dir());
let kinds = first
.take_worker_contributions()
.into_iter()
+139 -67
View File
@@ -13,13 +13,14 @@ 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, QuotaLocalUsageAttribution,
ProviderQuotaRequest, ProviderQuotaWindow, ProviderQuotaWindowRole, ProviderResetCredit,
ProviderResetCreditResult, ProviderResetCredits, QuotaLocalUsageAttribution,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
@@ -47,8 +48,8 @@ use crate::credential::{
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError,
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
CodexPendingAuthorization, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole,
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
CodexResetCreditsError, CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential,
ExportManagedCodexCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
};
use crate::credential::{
CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref, parse_access_token_expiration,
@@ -442,6 +443,37 @@ impl ProviderAdmin for OpenAiAdminProvider {
Ok(project_quota(snapshot, &account))
}
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,
@@ -734,39 +766,24 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
.windows()
.iter()
.filter(|window| should_project_quota_window(window))
.map(|window| {
let mut data = Map::new();
data.insert(
"source".to_owned(),
Value::String(window.source().to_owned()),
);
data.insert(
"role".to_owned(),
Value::String(quota_role(window.role()).to_owned()),
);
ProviderQuotaWindow {
key: window.key().to_owned(),
group: quota_group(window.kind()).to_owned(),
label: codex_quota_window_label(
window.limit_name(),
window.kind(),
window.role(),
window.source(),
window.window_seconds(),
),
source: Some(window.source().to_owned()),
local_usage_attribution: if window.is_account_wide() {
QuotaLocalUsageAttribution::AccountWide
} else {
QuotaLocalUsageAttribution::Unavailable
},
window_seconds: window.window_seconds(),
used_percent: window.used_percent(),
reset_at: window.reset_at(),
limit_reached: window.limit_reached(),
local_usage: None,
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))),
}
.map(|window| ProviderQuotaWindow {
key: window.key().to_owned(),
group: quota_group(window.kind()).to_owned(),
label: codex_quota_window_label(window.kind(), window.role(), window.window_seconds()),
limit_id: Some(window.source().to_owned()),
limit_name: window.limit_name().map(str::to_owned),
role: Some(quota_role(window.role())),
local_usage_attribution: if window.is_account_wide() {
QuotaLocalUsageAttribution::AccountWide
} else {
QuotaLocalUsageAttribution::Unavailable
},
window_seconds: window.window_seconds(),
used_percent: window.used_percent(),
reset_at: window.reset_at(),
limit_reached: window.limit_reached(),
local_usage: None,
provider_data: None,
})
.collect();
// 快照级 limit_reached 只看滚动后的窗口触顶:顶层标记是观测事实,不能
@@ -783,6 +800,31 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
}
}
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 小时或周额度,因此不应占用账号面板;带有实际事实的
/// 次级窗口(时长、重置、非零用量或触顶)仍完整保留。
@@ -806,23 +848,21 @@ const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str {
}
}
const fn quota_role(role: CodexQuotaWindowRole) -> &'static str {
const fn quota_role(role: CodexQuotaWindowRole) -> ProviderQuotaWindowRole {
match role {
CodexQuotaWindowRole::Primary => "primary",
CodexQuotaWindowRole::Secondary => "secondary",
CodexQuotaWindowRole::Monthly => "monthly",
CodexQuotaWindowRole::Primary => ProviderQuotaWindowRole::Primary,
CodexQuotaWindowRole::Secondary => ProviderQuotaWindowRole::Secondary,
CodexQuotaWindowRole::Monthly => ProviderQuotaWindowRole::Monthly,
}
}
/// 按窗口时长显示汉化额度名;非核心桶补上上游桶名称作为前缀。
/// 按窗口时长显示汉化额度名;额度桶名称通过独立字段投影。
fn codex_quota_window_label(
limit_name: Option<&str>,
kind: CodexQuotaWindowKind,
role: CodexQuotaWindowRole,
source: &str,
window_seconds: Option<u64>,
) -> String {
let base = match kind {
match kind {
CodexQuotaWindowKind::Monthly => "月额度".to_owned(),
CodexQuotaWindowKind::Weekly => "周额度".to_owned(),
CodexQuotaWindowKind::ShortTerm => {
@@ -833,29 +873,7 @@ fn codex_quota_window_label(
}
}
CodexQuotaWindowKind::Other => custom_quota_window_label(window_seconds, role),
};
if let Some(prefix) = codex_quota_window_prefix(limit_name, source) {
return format!("{prefix} · {base}");
}
base
}
fn codex_quota_window_prefix(limit_name: Option<&str>, source: &str) -> Option<String> {
let source_identifier = quota_identifier(source);
if matches!(source_identifier.as_str(), "core" | "codex") {
return None;
}
let label = limit_name.unwrap_or(source);
match quota_identifier(label).as_str() {
"premium" => Some("Premium".to_owned()),
"code_review" => Some("代码审查".to_owned()),
"spend_control" => Some("消费控制".to_owned()),
_ => Some(label.replace('_', " ")),
}
}
fn quota_identifier(value: &str) -> String {
value.trim().to_ascii_lowercase().replace(['-', ' '], "_")
}
fn custom_quota_window_label(window_seconds: Option<u64>, role: CodexQuotaWindowRole) -> String {
@@ -1322,6 +1340,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",
}
}
@@ -1377,6 +1397,58 @@ fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
})
}
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 {
@@ -36,15 +36,17 @@ pub struct OpenAiConfig {
}
impl OpenAiConfig {
/// 校验 Provider-owned 字段,并定位旧版会话锚点持续使用的身份密钥。
pub fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), OpenAiConfigError> {
/// 校验 Provider-owned 字段,并从统一运行数据目录定位会话身份密钥。
pub fn resolve_and_validate(
&mut self,
runtime_data_dir: &Path,
) -> Result<(), OpenAiConfigError> {
self.api.validate()?;
self.ws_pool.validate()?;
self.quota.validate()?;
self.auth.validate()?;
self.wire_profile.validate()?;
let project_root = source_dir.parent().unwrap_or(source_dir);
self.identity_secret_path = project_root.join(".runtime/data/identity_hmac_secret");
self.identity_secret_path = runtime_data_dir.join("identity_hmac_secret");
Ok(())
}
@@ -101,7 +103,7 @@ impl Default for OpenAiConfig {
quota: CodexQuotaSettings::default(),
auth: CodexAuthSettings::default(),
wire_profile: CodexWireProfileConfig::default(),
identity_secret_path: PathBuf::from(".runtime/data/identity_hmac_secret"),
identity_secret_path: PathBuf::new(),
}
}
}
@@ -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;
@@ -47,7 +47,7 @@ pub use oauth::{
pub use quota::{
CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService,
CodexQuotaFact, CodexQuotaRefreshPolicy, CodexQuotaSyncSummary, CodexQuotaWindow,
CodexQuotaWindowKind, CodexQuotaWindowRole, parse_codex_quota_usage,
CodexQuotaWindowKind, CodexQuotaWindowRole, CodexResetCreditsError, parse_codex_quota_usage,
};
pub use refresh::{
CodexCredentialRefreshError, CodexCredentialRefreshOutcome, CodexCredentialRefreshService,
@@ -58,11 +58,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,
@@ -0,0 +1,173 @@
//! OpenAI quota 文档的规范化 map 与上游 `/usage` 输入适配。
use std::collections::BTreeMap;
use gateway_protocol::openai::events::{RateLimitDetails, RateLimitKeySource};
use serde_json::{Map, Value};
pub(super) const DEFAULT_CODEX_LIMIT_ID: &str = "codex";
pub(super) const RATE_LIMITS_BY_LIMIT_ID: &str = "rate_limits_by_limit_id";
/// 官方新协议的多桶视图:每个 `limit_id` 对应一个独立快照。
#[derive(Debug, Default)]
pub(super) struct RateLimitSnapshotsByLimitId {
snapshots: BTreeMap<String, Map<String, Value>>,
}
impl RateLimitSnapshotsByLimitId {
/// 接收上游 `/usage` 形状或已经规范化的 map;输出始终只保留 map。
pub(super) fn take_from_document(quota: &mut Map<String, Value>) -> Self {
let mut result = Self::default();
if let Some(by_limit_id) = quota
.remove(RATE_LIMITS_BY_LIMIT_ID)
.and_then(|value| value.as_object().cloned())
{
for (raw_limit_id, value) in by_limit_id {
let Some(limit_id) = normalize_limit_id(&raw_limit_id) else {
continue;
};
let mut snapshot = snapshot_object(value);
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.clone()));
result.snapshots.insert(limit_id, snapshot);
}
}
if let Some(rate_limit) = quota.remove("rate_limit") {
result.insert_usage_snapshot(DEFAULT_CODEX_LIMIT_ID, None, rate_limit);
}
if let Some(rate_limit) = quota.remove("code_review_rate_limit") {
result.insert_usage_snapshot("code_review", Some("code_review"), rate_limit);
}
let additional = quota
.remove("additional_rate_limits")
.and_then(|value| value.as_array().cloned())
.unwrap_or_default();
for (index, value) in additional.into_iter().enumerate() {
let Value::Object(mut item) = value else {
result.insert_usage_snapshot(&format!("additional_{index}"), None, value);
continue;
};
let limit_id = item_limit_id(&item).unwrap_or_else(|| format!("additional_{index}"));
let limit_name = item
.get("limit_name")
.and_then(Value::as_str)
.map(str::to_owned);
let rate_limit = item.remove("rate_limit").unwrap_or(Value::Null);
result.insert_usage_snapshot(&limit_id, limit_name.as_deref(), rate_limit);
}
result
}
fn insert_usage_snapshot(
&mut self,
limit_id: &str,
limit_name: Option<&str>,
rate_limit: Value,
) {
let Some(limit_id) = normalize_limit_id(limit_id) else {
return;
};
let mut snapshot = snapshot_object(rate_limit);
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.clone()));
if let Some(limit_name) = limit_name {
snapshot.insert(
"limit_name".to_owned(),
Value::String(limit_name.to_owned()),
);
}
// map 协议中一个 ID 只能有一个快照;规范化 map 与顶层账号桶优先。
self.snapshots.entry(limit_id).or_insert(snapshot);
}
pub(super) fn upsert(
&mut self,
limit_id: &str,
details: &RateLimitDetails,
rate_limit: Map<String, Value>,
) {
let mut snapshot = self.snapshots.remove(limit_id).unwrap_or_default();
for field in [
"allowed",
"limit_reached",
"primary_window",
"secondary_window",
] {
snapshot.remove(field);
}
snapshot.extend(rate_limit);
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.to_owned()));
if let Some(name) = details.limit_name.as_ref() {
snapshot.insert("limit_name".to_owned(), Value::String(name.clone()));
}
self.snapshots.insert(limit_id.to_owned(), snapshot);
}
/// 把 wire 观察解析到稳定 `limit_id`。
///
/// WebSocket 的 `additional_rate_limits` 可能只返回可读名称;这时只接受
/// `/usage` 或 HTTP headers 已建立的唯一名称映射。无法唯一解析就不落库,
/// 避免凭名称猜 ID 或污染默认 `codex` 桶。
pub(super) fn resolve_limit_id(&self, details: &RateLimitDetails) -> Option<String> {
if details.key_source == RateLimitKeySource::LimitId {
return normalize_limit_id(&details.limit_id);
}
let limit_name = details.limit_name.as_deref()?.trim();
let mut matches = self
.snapshots
.iter()
.filter(|(_, snapshot)| {
snapshot
.get("limit_name")
.and_then(Value::as_str)
.is_some_and(|stored| stored.trim().eq_ignore_ascii_case(limit_name))
})
.map(|(limit_id, _)| limit_id.clone());
let resolved = matches.next()?;
matches.next().is_none().then_some(resolved)
}
pub(super) fn write_to_document(self, quota: &mut Map<String, Value>) {
let by_limit_id = self
.snapshots
.into_iter()
.map(|(limit_id, snapshot)| (limit_id, Value::Object(snapshot)))
.collect::<Map<_, _>>();
quota.insert(
RATE_LIMITS_BY_LIMIT_ID.to_owned(),
Value::Object(by_limit_id),
);
}
}
pub(super) fn canonicalize_rate_limit_document(
mut quota: Map<String, Value>,
) -> Map<String, Value> {
let snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota);
snapshots.write_to_document(&mut quota);
quota
}
fn snapshot_object(value: Value) -> Map<String, Value> {
value.as_object().cloned().unwrap_or_else(|| {
let mut snapshot = Map::new();
snapshot.insert("rate_limit".to_owned(), value);
snapshot
})
}
fn item_limit_id(item: &Map<String, Value>) -> Option<String> {
item.get("limit_id")
.or_else(|| item.get("metered_feature"))
.or_else(|| item.get("limit_name"))
.and_then(Value::as_str)
.and_then(normalize_limit_id)
}
fn normalize_limit_id(value: &str) -> Option<String> {
let value = value.trim();
(!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control))
.then(|| value.to_ascii_lowercase().replace('-', "_"))
}
@@ -1,8 +1,10 @@
//! Codex quota 服务编排:权威额度事实、展示投影、主动/被动同步与 429 冷却。
//!
//! - [`document`]:上游 `/usage` 输入到多桶 map 的单向规范化。
//! - [`snapshot`]:快照/窗口解析、聚合、O(1) 滚动与调度信号。
//! - [`evidence`]:额度接口错误到凭据/额度事实的分类。
mod document;
pub(crate) mod evidence;
pub(crate) mod snapshot;
@@ -11,7 +13,7 @@ pub use snapshot::{
CodexQuotaWindowRole, parse_codex_quota_usage,
};
use std::collections::{BTreeMap, BTreeSet};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant, SystemTime};
@@ -34,10 +36,17 @@ use tokio::sync::Mutex;
use uuid::Uuid;
use crate::transport::profile::CodexWireProfileState;
use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext};
use crate::transport::{
CodexBackendClient, CodexClientError, CodexRateLimitResetCredits,
CodexRateLimitResetCreditsConsumeResult, CodexRequestContext,
};
use super::agent_identity::{CodexAgentIdentityTaskService, PreparedCodexRuntimeCredential};
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use document::{
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId,
canonicalize_rate_limit_document,
};
use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure};
use snapshot::{
parse_account_quota_snapshot, quota_projection_ttl, quota_snapshot_from_observation,
@@ -112,6 +121,54 @@ 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,
}
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 {
@@ -129,6 +186,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)]
@@ -183,6 +241,11 @@ enum CodexQuotaFetchError {
},
}
enum ResetCreditAttemptError {
InvalidCredential,
Upstream(CodexClientError),
}
impl CodexQuotaSchedulingProjection {
fn invalidate(&self, account_ids: &[ProviderAccountId]) {
let mut state = self
@@ -407,9 +470,129 @@ 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 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,
@@ -706,10 +889,25 @@ impl CodexCredentialQuotaService {
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
return Ok(false);
};
let has_quota_facts = rate_limits
.limits
.values()
.any(|details| passive_rate_limit_snapshot(details).is_some());
self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits))
.await
}
/// 把一次推理响应中采集的结构化限流观察合并后单次落库。
pub async fn synchronize_passive_rate_limits(
&self,
account: &ProviderAccount,
rate_limits: &[ParsedRateLimits],
) -> Result<bool, CodexCredentialQuotaError> {
if rate_limits.is_empty() {
return Ok(false);
}
let has_quota_facts = rate_limits.iter().any(|observation| {
observation
.limits
.values()
.any(|details| passive_rate_limit_snapshot(details).is_some())
});
let existing = self
.store
.get_quotas(std::slice::from_ref(account.id()))
@@ -734,14 +932,14 @@ impl CodexCredentialQuotaService {
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(merge_passive_quota(existing, &rate_limits)),
quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)),
observed_at: SystemTime::now(),
state,
})
.await?;
return Ok(outcome != QuotaWriteOutcome::Conflict);
}
let quota = merge_passive_quota(existing, &rate_limits);
let quota = merge_passive_quota(existing, rate_limits);
let observed_at = SystemTime::now();
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
@@ -1066,6 +1264,118 @@ impl CodexCredentialQuotaService {
}
}
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(
account: &ProviderAccount,
refreshed: &CodexAccountQuotaSnapshot,
@@ -1177,18 +1487,48 @@ fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool {
fn merge_passive_quota(
mut quota: Map<String, Value>,
rate_limits: &ParsedRateLimits,
observations: &[ParsedRateLimits],
) -> Map<String, Value> {
let active_limit = rate_limits
.active_limit
.as_deref()
.or_else(|| rate_limits.limits.contains_key("codex").then_some("codex"));
if let Some(active_limit) = active_limit {
quota.insert(
"active_limit".to_owned(),
Value::String(active_limit.to_owned()),
);
let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota);
for rate_limits in observations {
let default_is_named_alias = default_limit_is_named_alias(rate_limits);
let mut resolved_limit_ids = BTreeMap::new();
for (wire_limit_id, details) in &rate_limits.limits {
// HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口。
// 同一 wire 观察里存在相同具名事实时丢弃镜像,不触碰 core 桶。
if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias {
continue;
}
let Some(limit_id) = snapshots.resolve_limit_id(details) else {
continue;
};
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
continue;
};
snapshots.upsert(&limit_id, details, rate_limit);
resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id);
}
let active_limit = rate_limits
.active_limit
.as_deref()
.and_then(|wire_limit_id| {
resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| {
(!rate_limits.limits.contains_key(wire_limit_id))
.then(|| wire_limit_id.to_owned())
})
})
.or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned());
if let Some(active_limit) = active_limit {
quota.insert("active_limit".to_owned(), Value::String(active_limit));
}
merge_passive_metadata(&mut quota, rate_limits);
}
snapshots.write_to_document(&mut quota);
quota
}
fn merge_passive_metadata(quota: &mut Map<String, Value>, rate_limits: &ParsedRateLimits) {
if let Some(plan_type) = rate_limits.plan_type.as_ref() {
quota.insert("plan_type".to_owned(), Value::String(plan_type.clone()));
}
@@ -1216,48 +1556,19 @@ fn merge_passive_quota(
}
quota.insert("credits".to_owned(), Value::Object(value));
}
}
let mut additional = quota
.remove("additional_rate_limits")
.and_then(|value| value.as_array().cloned())
.unwrap_or_default();
for (limit_id, details) in &rate_limits.limits {
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
continue;
};
if Some(limit_id.as_str()) == active_limit {
quota.insert("rate_limit".to_owned(), Value::Object(rate_limit));
continue;
}
let index = additional.iter().position(|item| {
item.get("metered_feature")
.or_else(|| item.get("limit_name"))
.and_then(Value::as_str)
.is_some_and(|existing| existing == limit_id)
});
let mut item = index
.and_then(|index| additional.get(index))
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
item.insert(
"metered_feature".to_owned(),
Value::String(limit_id.clone()),
);
if let Some(name) = details.limit_name.as_ref() {
item.insert("limit_name".to_owned(), Value::String(name.clone()));
}
item.insert("rate_limit".to_owned(), Value::Object(rate_limit));
match index {
Some(index) => additional[index] = Value::Object(item),
None => additional.push(Value::Object(item)),
}
}
quota.insert(
"additional_rate_limits".to_owned(),
Value::Array(additional),
);
quota
fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool {
let Some(default_primary) = rate_limits
.limits
.get(DEFAULT_CODEX_LIMIT_ID)
.and_then(|details| details.primary)
else {
return false;
};
rate_limits.limits.iter().any(|(limit_id, details)| {
limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary)
})
}
/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位。
@@ -1269,7 +1580,13 @@ fn merge_passive_quota(
/// 无事实值。带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按
/// 原有额度逻辑保留。
fn normalize_quota_window_placeholders(mut quota: Map<String, Value>) -> Map<String, Value> {
if let Some(rate_limit) = quota.get_mut("rate_limit").and_then(Value::as_object_mut) {
quota = canonicalize_rate_limit_document(quota);
if let Some(rate_limit) = quota
.get_mut(RATE_LIMITS_BY_LIMIT_ID)
.and_then(Value::as_object_mut)
.and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID))
.and_then(Value::as_object_mut)
{
drop_secondary_window_placeholder(rate_limit);
}
quota
@@ -12,6 +12,9 @@ use gateway_core::engine::credential::{
};
use serde_json::{Map, Value};
use super::document::{
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, canonicalize_rate_limit_document,
};
use super::{CodexCredentialQuotaError, QUOTA_SCHEDULING_TTL};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -173,9 +176,18 @@ impl CodexQuotaFact {
}
pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
let object = usage
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
let object = canonicalize_rate_limit_document(
usage
.as_object()
.cloned()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
);
parse_codex_quota_object(&object)
}
fn parse_codex_quota_object(
object: &Map<String, Value>,
) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
let mut aggregate = QuotaAggregate::default();
for limit in canonical_rate_limits(object)? {
aggregate.observe_rate_limit(limit.rate_limit)?;
@@ -196,6 +208,15 @@ pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCre
})
}
fn canonical_quota_object(usage: &Value) -> Result<Map<String, Value>, CodexCredentialQuotaError> {
Ok(canonicalize_rate_limit_document(
usage
.as_object()
.cloned()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
))
}
pub(crate) fn quota_snapshot_from_observation(
observation: &QuotaObservation,
) -> Option<CodexAccountQuotaSnapshot> {
@@ -277,13 +298,11 @@ pub(crate) fn parse_account_quota_snapshot(
observed_at: SystemTime,
usage: &Value,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let fact = parse_codex_quota_usage(usage)?;
let object = usage
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
let quota = authoritative_quota_state(object, observed_at)?;
let object = canonical_quota_object(usage)?;
let fact = parse_codex_quota_object(&object)?;
let quota = authoritative_quota_state(&object, observed_at)?;
let mut windows = Vec::new();
for limit in canonical_rate_limits(object)? {
for limit in canonical_rate_limits(&object)? {
parse_rate_limit_windows(
&limit.key,
&limit.source,
@@ -316,12 +335,16 @@ fn quota_source_order(source: &str) -> u8 {
}
}
/// 顶层 account-wide Codex 桶是唯一可改变账号额度访问结论的 quota 响应字段。
/// 默认 `codex` 桶是唯一可改变账号额度访问结论的 quota 响应字段。
fn authoritative_quota_state(
usage: &Map<String, Value>,
observed_at: SystemTime,
) -> Result<QuotaState, CodexCredentialQuotaError> {
let Some(rate_limit) = usage.get("rate_limit") else {
let Some(rate_limit) = usage
.get(RATE_LIMITS_BY_LIMIT_ID)
.and_then(Value::as_object)
.and_then(|limits| limits.get(DEFAULT_CODEX_LIMIT_ID))
else {
return Ok(QuotaState::observed_unknown(observed_at));
};
if rate_limit.is_null() {
@@ -398,64 +421,32 @@ fn canonical_rate_limits(
object: &Map<String, Value>,
) -> Result<Vec<CanonicalRateLimit<'_>>, CodexCredentialQuotaError> {
let mut limits = Vec::new();
let has_primary_codex_rate_limit = object
.get("rate_limit")
.is_some_and(|value| !value.is_null());
if let Some(rate_limit) = object.get("rate_limit") {
// 官方 Codex 将顶层 `rate_limit` 固定识别为 `codex`,而不把
// `active_limit` 投影成另一个展示桶。后者只是当前活动套餐的元数据。
limits.push(CanonicalRateLimit {
key: "core".to_owned(),
source: "codex".to_owned(),
account_wide: true,
limit_name: None,
rate_limit,
});
}
if let Some(rate_limit) = object
.get("code_review_rate_limit")
.filter(|value| !value.is_null())
{
limits.push(CanonicalRateLimit {
key: "code-review".to_owned(),
source: "code_review".to_owned(),
account_wide: false,
limit_name: Some("code_review".to_owned()),
rate_limit,
});
}
let Some(additional) = object
.get("additional_rate_limits")
.filter(|value| !value.is_null())
else {
let Some(by_limit_id) = object.get(RATE_LIMITS_BY_LIMIT_ID) else {
return Ok(limits);
};
for (index, value) in additional
.as_array()
for (limit_id, rate_limit) in by_limit_id
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?
.iter()
.enumerate()
{
let item = value
let item = rate_limit
.as_object()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
let (source, limit_name) = canonical_rate_limit_source(item, index);
// 部分上游响应会在 `additional_rate_limits` 中重复默认的 `codex`
// 桶。它没有独立的 limit id,不能在账号面板再次展示。
if has_primary_codex_rate_limit && is_codex_limit_id(&source) {
continue;
}
if let Some(rate_limit) = item.get("rate_limit") {
// 非 Codex additional 桶保留为独立快照;额外用索引生成窗口键,
// 避免同名 additional 产生重复前端 key。
limits.push(CanonicalRateLimit {
key: format!("additional-{index}-{source}"),
account_wide: is_codex_limit_id(&source),
source,
limit_name,
rate_limit,
});
}
let source = item
.get("limit_id")
.and_then(Value::as_str)
.unwrap_or(limit_id)
.to_owned();
let limit_name = item
.get("limit_name")
.and_then(Value::as_str)
.map(str::to_owned);
limits.push(CanonicalRateLimit {
key: source.clone(),
account_wide: is_codex_limit_id(&source),
source,
limit_name,
rate_limit,
});
}
Ok(limits)
}
@@ -464,30 +455,6 @@ fn is_codex_limit_id(value: &str) -> bool {
value.trim().eq_ignore_ascii_case("codex")
}
fn canonical_rate_limit_source(
item: &Map<String, Value>,
index: usize,
) -> (String, Option<String>) {
let source = item
.get("metered_feature")
.or_else(|| item.get("limit_name"))
.and_then(Value::as_str)
.filter(|value| {
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
})
.map(str::to_owned)
.unwrap_or_else(|| format!("additional-{index}"));
let limit_name = item
.get("limit_name")
.or_else(|| item.get("metered_feature"))
.and_then(Value::as_str)
.filter(|value| {
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
})
.map(str::to_owned);
(source, limit_name)
}
fn parse_rate_limit_windows(
key_source: &str,
source: &str,
@@ -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 {
+496 -93
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use bytes::Bytes;
use futures::{StreamExt, future::BoxFuture};
use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope};
use gateway_core::engine::credential::{AccountFeedbackStats, ProviderAccount};
@@ -22,17 +23,25 @@ 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,
WorkerRunnable, WorkerSchedule, WorkerTaskError,
};
use gateway_protocol::openai::events::{
ParsedRateLimits, parse_rate_limit_headers, rate_limits_to_header_pairs,
};
use reqwest::Client;
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
@@ -44,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,
@@ -63,20 +72,22 @@ use crate::transport::protocol::responses::{
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
};
use crate::transport::request::{
CodexRequestEncodeError, encode_generate_request, scope_request_to_account,
CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account,
};
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,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError,
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics,
CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
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,
};
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;达到阈值后结束无感换号窗口,
@@ -111,6 +122,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>,
}
@@ -130,6 +143,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 {
@@ -140,6 +158,8 @@ impl CodexProvider {
account_feedback,
client,
responses_url,
image_generations_url,
image_edits_url,
session_identity: None,
})
}
@@ -227,17 +247,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())
@@ -247,7 +275,12 @@ impl Provider for CodexProvider {
if let Some(identity) = &self.session_identity {
identity.prepare_local_conversation(&mut upstream_request);
}
if upstream_request.turn_state.is_none() {
if previous_session.as_ref().is_some_and(|state| {
same_client_turn(
state.client_turn_id.as_deref(),
upstream_request.client_turn_id.as_deref(),
)
}) {
upstream_request.turn_state = previous_session
.as_ref()
.and_then(|state| state.turn_state.clone());
@@ -263,7 +296,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(),
@@ -277,57 +310,48 @@ impl Provider for CodexProvider {
let output_started_at = Instant::now();
let provider_kind = ProviderKind::new(PROVIDER_NAME)
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
let cross_account = context
let state_owner_cross_account = context
.account_state_owner()
.is_some_and(|owner| !owner.matches(&provider_kind, lease.account_id()))
|| previous_session
.as_ref()
.is_some_and(|state| state.account_id != lease.account_id().as_str());
let replay_previous_response = matches!(
context.continuation_attempt(),
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny
) || previous_session.as_ref().is_some_and(|state| {
state.continuation_scope == OpenAiContinuationScope::ReplayRequired
});
if replay_previous_response && previous_session.is_some() {
// Codex 客户端 input 是完整历史;恢复时只解除上游续链绑定,避免重复叠加代理副本。
upstream_request.set_previous_response_id(None);
upstream_request.previous_response_scope = None;
upstream_request.turn_state = None;
}
scope_request_to_account(
&mut upstream_request,
lease.installation_id(),
cross_account,
);
if context.continuation_attempt() == ContinuationAttempt::Native
&& !replay_previous_response
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());
if context.continuation_attempt() != ContinuationAttempt::None
&& let Some(continuation) = context.continuation()
{
match continuation {
ContinuationBinding::Pinned(continuation) => {
let previous_response_scope = match previous_session
.as_ref()
.map(|state| state.continuation_scope)
{
Some(OpenAiContinuationScope::Persisted) => {
PreviousResponseScope::Persisted
}
Some(OpenAiContinuationScope::ConnectionLocal) => {
PreviousResponseScope::ConnectionLocal
}
Some(OpenAiContinuationScope::ReplayRequired) => {
return Err(provider_error(
ProviderErrorKind::Protocol,
UpstreamSendState::NotSent,
));
}
None => match continuation.scope() {
NativeContinuationScope::Persisted => PreviousResponseScope::Persisted,
NativeContinuationScope::ConnectionLocal => {
let previous_response_scope = match context.continuation_attempt() {
ContinuationAttempt::Native => match previous_session
.as_ref()
.map(|state| state.continuation_scope)
{
Some(OpenAiContinuationScope::Persisted) => {
PreviousResponseScope::Persisted
}
Some(OpenAiContinuationScope::ConnectionLocal) => {
PreviousResponseScope::ConnectionLocal
}
Some(OpenAiContinuationScope::ReplayRequired) => {
PreviousResponseScope::ExternalUnknown
}
None => match continuation.scope() {
NativeContinuationScope::Persisted => {
PreviousResponseScope::Persisted
}
NativeContinuationScope::ConnectionLocal => {
PreviousResponseScope::ConnectionLocal
}
},
},
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => {
PreviousResponseScope::ExternalUnknown
}
ContinuationAttempt::None => PreviousResponseScope::ExternalUnknown,
};
upstream_request.set_previous_response_id(Some(
continuation.upstream_response_id().as_str().to_owned(),
@@ -342,9 +366,21 @@ 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
};
scope_request_to_account(
&mut upstream_request,
lease.installation_id(),
account_scope,
);
let metadata = ProviderCallMetadata::new(
provider_kind,
candidate.upstream_model().clone(),
upstream_model.clone(),
ProviderResource::Account {
id: lease.account_id().clone(),
revision: lease.account().revision(),
@@ -359,6 +395,7 @@ impl Provider for CodexProvider {
account_id: lease.account_id().as_str().to_owned(),
conversation_id: upstream_request.local_conversation_id.clone(),
turn_state: upstream_request.turn_state.clone(),
client_turn_id: upstream_request.client_turn_id.clone(),
response_store,
continuation_scope: None,
});
@@ -367,7 +404,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),
@@ -388,6 +425,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,
@@ -405,12 +512,28 @@ 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,
conversation_id: Option<String>,
#[serde(default)]
turn_state: Option<String>,
#[serde(default)]
client_turn_id: Option<String>,
continuation_scope: OpenAiContinuationScope,
}
@@ -426,10 +549,17 @@ struct OpenAiSessionCapture {
account_id: String,
conversation_id: Option<String>,
turn_state: Option<String>,
client_turn_id: Option<String>,
response_store: bool,
continuation_scope: Option<OpenAiContinuationScope>,
}
fn same_client_turn(previous: Option<&str>, current: Option<&str>) -> bool {
previous
.zip(current)
.is_some_and(|(previous, current)| !previous.is_empty() && previous == current)
}
fn decode_openai_session_state(request: &GenerateRequest) -> Option<OpenAiSessionState> {
request
.provider_session_state(PROVIDER_NAME)
@@ -471,6 +601,7 @@ fn attach_openai_session_update(
account_id: capture.account_id,
conversation_id: capture.conversation_id,
turn_state: capture.turn_state,
client_turn_id: capture.client_turn_id,
continuation_scope,
}) else {
return;
@@ -546,6 +677,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,
@@ -698,9 +1001,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if let Some(observation) = observation_state.observation() {
yield ProviderEvent::observation(observation);
}
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
}
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(etag)
{
@@ -728,7 +1028,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let failure_diagnostics = response.diagnostics.clone();
let failure_set_cookie_headers = response.set_cookie_headers.clone();
let failure_rate_limit_headers = response.rate_limit_headers.clone();
let rate_limit_updates = response.rate_limit_header_updates;
let mut passive_quota_observation =
OpenAiPassiveQuotaObservation::new(response.rate_limit_headers);
let rate_limit_updates = response.rate_limit_updates;
let turn_state_updates = response.turn_state_update;
// OpenAI 线路为透明代理:HTTP SSE 与 WebSocket 两条上游均启用 raw 透传,
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
@@ -741,6 +1043,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let mut pre_commit_events = PreCommitClientEvents::new();
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
if allows_account_state_mutation {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
return;
};
@@ -773,15 +1083,22 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure) => {
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
if observation_state.merge_rate_limit_headers(&updates)
passive_quota_observation.observe(&updates);
let update_headers = rate_limit_update_headers(&updates);
if observation_state.merge_rate_limit_headers(&update_headers)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
}
if allows_account_state_mutation {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
apply_failure(&failure_context, &active_account, &failure)
.await;
Err(failure.error)?;
@@ -793,18 +1110,19 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let rate_limits_changed = if updates.is_empty() {
false
} else {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
passive_quota_observation.observe(&updates);
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
};
if let (Some(capture), Some(updates)) =
(session_capture.as_mut(), turn_state_updates.as_ref())
// take:仅在上游写入过新 turn_state 时更新 capture,避免逐 chunk 克隆。
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
&& let Some(turn_state) = updates.lock().await.take()
{
capture.turn_state = Some(turn_state);
}
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
} else {
false
};
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, output_started_at);
let chunk_len = chunk.len();
@@ -843,6 +1161,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if allows_account_state_mutation && (completed || terminal_failure.is_some()) {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
@@ -861,6 +1187,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
if (rate_limits_changed
|| service_tier_changed
|| timing_changed
|| turn_state_changed
|| terminal_changed)
&& let Some(observation) = observation_state.observation()
{
@@ -920,21 +1247,31 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let rate_limits_changed = if updates.is_empty() {
false
} else {
if allows_account_state_mutation {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
passive_quota_observation.observe(&updates);
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
};
if allows_account_state_mutation {
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
}
if let Some((failure, _)) = terminal_failure.as_ref() {
apply_failure(&failure_context, &active_account, failure)
.await;
}
if let (Some(capture), Some(updates)) =
(session_capture.as_mut(), turn_state_updates.as_ref())
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
&& let Some(turn_state) = updates.lock().await.take()
{
capture.turn_state = Some(turn_state);
}
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
} else {
false
};
attach_openai_session_update(&mut events, &mut session_capture);
let completed = events
.iter()
@@ -951,7 +1288,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
}
if (service_tier_changed || timing_changed || rate_limits_changed || terminal_changed)
if (service_tier_changed
|| timing_changed
|| rate_limits_changed
|| turn_state_changed
|| terminal_changed)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
@@ -1001,6 +1342,26 @@ struct OpenAiResponseObservationState {
terminal: Option<OpenAiResponseTerminal>,
}
struct OpenAiPassiveQuotaObservation {
rate_limits: Vec<ParsedRateLimits>,
}
impl OpenAiPassiveQuotaObservation {
fn new(headers: Vec<(String, String)>) -> Self {
Self {
rate_limits: parse_rate_limit_headers(&headers).into_iter().collect(),
}
}
fn observe(&mut self, updates: &[ParsedRateLimits]) {
self.rate_limits.extend_from_slice(updates);
}
fn rate_limits(&self) -> &[ParsedRateLimits] {
&self.rate_limits
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum OpenAiResponseTerminal {
Completed,
@@ -1120,6 +1481,26 @@ impl OpenAiResponseObservationState {
changed
}
fn merge_client_header(&mut self, name: &str, value: &str) -> bool {
let value = Bytes::copy_from_slice(value.as_bytes());
if let Some((_, existing)) = self
.response_metadata
.client_headers
.iter_mut()
.find(|(existing_name, _)| existing_name.eq_ignore_ascii_case(name))
{
if *existing == value {
return false;
}
*existing = value;
return true;
}
self.response_metadata
.client_headers
.push((name.to_owned(), value));
true
}
fn mark_completed(&mut self, incomplete: bool) -> bool {
let terminal = if incomplete {
OpenAiResponseTerminal::Incomplete
@@ -1432,15 +1813,20 @@ fn insert_optional_metadata_millis(
}
}
async fn take_rate_limit_updates(
updates: Option<&CodexRateLimitHeaderUpdates>,
) -> Vec<(String, String)> {
async fn take_rate_limit_updates(updates: Option<&CodexRateLimitUpdates>) -> Vec<ParsedRateLimits> {
let Some(updates) = updates else {
return Vec::new();
};
std::mem::take(&mut *updates.lock().await)
}
fn rate_limit_update_headers(updates: &[ParsedRateLimits]) -> Vec<(String, String)> {
updates
.iter()
.flat_map(rate_limits_to_header_pairs)
.collect()
}
fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
events
.iter()
@@ -1466,12 +1852,15 @@ fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
async fn synchronize_passive_quota(
quota: &CodexCredentialQuotaService,
account: &ProviderAccount,
headers: &[(String, String)],
rate_limits: &[ParsedRateLimits],
) {
if headers.is_empty() {
if rate_limits.is_empty() {
return;
}
if let Err(error) = quota.synchronize_passive_headers(account, headers).await {
if let Err(error) = quota
.synchronize_passive_rate_limits(account, rate_limits)
.await
{
tracing::warn!(
account_id = %account.id(),
error = %error,
@@ -1480,9 +1869,21 @@ async fn synchronize_passive_quota(
}
}
async fn synchronize_passive_quota_headers(
quota: &CodexCredentialQuotaService,
account: &ProviderAccount,
headers: &[(String, String)],
) {
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
return;
};
synchronize_passive_quota(quota, account, std::slice::from_ref(&rate_limits)).await;
}
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,
}
}
@@ -1873,7 +2274,7 @@ async fn apply_failure(
if !context.allows_account_state_mutation {
return;
}
synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await;
synchronize_passive_quota_headers(context.quota, account, &failure.rate_limit_headers).await;
let needs_authoritative_quota_refresh = matches!(
failure.account_failure,
Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. })
@@ -1974,7 +2375,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(_)
@@ -2055,7 +2458,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 {
@@ -2152,7 +2555,7 @@ fn map_upstream_failure(
error = error.with_status(status.as_u16());
}
if replay_boundary.permits_provider_proof()
&& (failure.replay_is_safe() || continuation_failure.is_some() || cyber_policy_failure)
&& (failure.replay_is_safe() || cyber_policy_failure)
{
error = error.with_replay_safe();
}
@@ -30,7 +30,7 @@ use super::response_meta::CodexResponseMetadata;
use super::tls::{CustomCaError, build_reqwest_client_with_custom_ca, custom_ca_env_cache_key};
use super::websocket::{
CodexWebSocketExchangeError, CodexWebSocketPool, CodexWebSocketPoolKey,
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
CodexWebSocketRateLimitUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
PreparedWebSocket, WebSocketOriginBreaker, WebSocketPoolDecision,
};
@@ -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,
}
@@ -519,8 +526,8 @@ pub struct CodexTransportMetrics {
pub http_version: Option<String>,
}
/// 响应头之后在 live 流中采集的限流更新。
pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates;
/// 响应头之后在 live 流中采集的结构化限流更新。
pub type CodexRateLimitUpdates = CodexWebSocketRateLimitUpdates;
/// 响应头之后在 live 流中采集的 turn state 更新。
pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate;
@@ -537,8 +544,8 @@ pub struct CodexBackendStreamingResponse {
pub set_cookie_headers: Vec<String>,
/// 上游透传的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// live stream 期间捕获的限流头更新。
pub rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
/// live stream 期间捕获的结构化限流更新。
pub rate_limit_updates: Option<CodexRateLimitUpdates>,
/// live stream 期间捕获的 turn-state 更新。
pub turn_state_update: Option<CodexTurnStateUpdate>,
/// 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,
})
}
}
@@ -147,14 +147,14 @@ impl CodexBackendClient {
});
}
let rate_limit_header_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
let rate_limit_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
Ok(CodexBackendStreamingResponse {
body: http_sse_stream(response, Arc::clone(&rate_limit_header_updates)),
body: http_sse_stream(response, Arc::clone(&rate_limit_updates)),
transport: CodexBackendTransport::HttpSse,
turn_state,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates: Some(rate_limit_header_updates),
rate_limit_updates: Some(rate_limit_updates),
turn_state_update: None,
websocket_pool_decision: None,
diagnostics,
@@ -245,11 +245,12 @@ impl CodexBackendClient {
self.websocket_pool.as_deref().zip(pool_key)
};
let fast_path_budget = match requirement {
TransportRequirement::PersistedContinuation
| TransportRequirement::ExternalUnknown
| TransportRequirement::NewChain => Some(WEBSOCKET_FAST_PATH_BUDGET),
TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => {
Some(WEBSOCKET_FAST_PATH_BUDGET)
}
TransportRequirement::ExplicitWebSocketWarmup
| TransportRequirement::ExactWebSocketContinuation => None,
| TransportRequirement::ExactWebSocketContinuation
| TransportRequirement::ExternalUnknown => None,
TransportRequirement::HttpRequired => None,
};
let prepare_started_at = Instant::now();
@@ -515,7 +516,7 @@ impl CodexBackendClient {
turn_state: exchange.turn_state,
set_cookie_headers: exchange.set_cookie_headers,
rate_limit_headers: exchange.rate_limit_headers,
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
rate_limit_updates: Some(exchange.rate_limit_updates),
turn_state_update: Some(exchange.turn_state_update),
websocket_pool_decision: exchange.pool_decision,
diagnostics: exchange.diagnostics,
@@ -878,7 +879,7 @@ fn websocket_connection_profile(headers: &HeaderMap) -> String {
fn http_sse_stream(
response: ReqwestResponse,
rate_limit_updates: CodexRateLimitHeaderUpdates,
rate_limit_updates: CodexRateLimitUpdates,
) -> CodexBackendSseStream {
let stream: CodexBackendSseStream =
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
@@ -901,7 +902,7 @@ fn http_sse_stream(
fn observe_http_sse_rate_limits(
stream: CodexBackendSseStream,
updates: CodexRateLimitHeaderUpdates,
updates: CodexRateLimitUpdates,
) -> CodexBackendSseStream {
Box::pin(futures::stream::unfold(
(stream, SseEventDecoder::default(), updates),
@@ -925,9 +926,9 @@ fn observe_http_sse_rate_limits(
async fn append_http_sse_rate_limit_updates(
frames: Vec<SseFrame>,
updates: &CodexRateLimitHeaderUpdates,
updates: &CodexRateLimitUpdates,
) {
let mut headers = Vec::new();
let mut observations = Vec::new();
for frame in frames {
for event in frame.events() {
if event
@@ -940,10 +941,10 @@ async fn append_http_sse_rate_limit_updates(
let Some(rate_limits) = events::parse_rate_limits_event_raw(&event.data) else {
continue;
};
headers.extend(events::rate_limits_to_header_pairs(&rate_limits));
observations.push(rate_limits);
}
}
if !headers.is_empty() {
updates.lock().await.extend(headers);
if !observations.is_empty() {
updates.lock().await.extend(observations);
}
}
@@ -1,9 +1,18 @@
/// `/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";
/// 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,6 +11,7 @@ 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;
@@ -25,18 +27,23 @@ pub use self::{
CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, parse_codex_model_catalog,
},
client::{
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendSseStream,
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult,
CodexRateLimitHeaderUpdates, 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_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::{
@@ -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 上游请求共享的运行时请求画像。
@@ -5,6 +5,8 @@ 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";
const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`.";
/// WebSocket 握手审计快照。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
@@ -54,16 +56,49 @@ pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
websocket_event_frame(&value, raw)
}
/// 同上,复用已解析的事件 JSON;`data` 段逐字节内嵌 `raw` 原文。
/// 同上,复用已解析的事件 JSON;`data` 段通常逐字节内嵌 `raw` 原文。
///
/// 仅剥离传输层内部帧:`codex.rate_limits` 与承载 turn_state 的
/// `response.metadata` 不下发客户端;业务事件一律原样转发,不做 schema 校验。
/// `response.metadata` 不下发客户端。唯一的业务帧适配是为上游省略错误码的
/// `Invalid previous_response_id` 补齐官方客户端用于完整历史重放的稳定错误码;
/// 原始 status、type 与 message 均保持不变。
pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option<String> {
let event = websocket_event_type(value)?;
if is_internal_websocket_event(event) || event == "response.metadata" {
return None;
}
Some(encode_sse_event(event, raw))
let normalized = previous_response_not_found_event(value);
Some(encode_sse_event(
event,
normalized.as_deref().unwrap_or(raw),
))
}
fn previous_response_not_found_event(value: &Value) -> Option<String> {
if websocket_event_type(value) != Some("error")
|| value
.get("status")
.or_else(|| value.get("status_code"))
.and_then(Value::as_u64)
!= Some(400)
{
return None;
}
let error = value.get("error")?.as_object()?;
if error.get("code").is_some_and(|code| !code.is_null())
|| error.get("type").and_then(Value::as_str) != Some("invalid_request_error")
|| error.get("message").and_then(Value::as_str)
!= Some(INVALID_PREVIOUS_RESPONSE_ID_MESSAGE)
{
return None;
}
let mut normalized = value.clone();
normalized.get_mut("error")?.as_object_mut()?.insert(
"code".to_owned(),
Value::String(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
);
serde_json::to_string(&normalized).ok()
}
pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> {
@@ -61,6 +61,23 @@ const INSTALLATION_ID_KEYS: &[&str] = &[
"x-codex-installation-id",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestAccountScope {
SameAccount,
CrossAccountFullRequest,
CrossAccountContinuationProbe,
}
impl RequestAccountScope {
const fn is_cross_account(self) -> bool {
!matches!(self, Self::SameAccount)
}
const fn preserves_continuation(self) -> bool {
matches!(self, Self::CrossAccountContinuationProbe)
}
}
/// Provider 专属编码错误;不保存 prompt、schema 或 option 值。
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum CodexRequestEncodeError {
@@ -289,12 +306,18 @@ fn normalize_conversation_anchor_text(text: &str) -> String {
///
/// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入
/// 请求头;正文只替换客户端原本声明过的 installation 字段,绝不接受客户端
/// 提供的 token、cookie 或账号身份。
/// 提供的 token、cookie 或账号身份。Continuation probe 保留 canonical
/// previous-response 和原始增量 input,不能把它们误当成完整请求清理。
pub(crate) fn scope_request_to_account(
request: &mut CodexResponsesRequest,
installation_id: &str,
cross_account: bool,
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 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());
@@ -306,23 +329,38 @@ pub(crate) fn scope_request_to_account(
} else {
None
};
let turn_metadata = request
.turn_metadata
.as_deref()
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account));
let client_metadata_turn_metadata = metadata_string(request, "x-codex-turn-metadata")
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account));
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();
if cross_account {
request.passthrough_headers.remove("x-codex-turn-state");
request.passthrough_headers.remove("x-codex-turn-metadata");
sanitize_cross_account_input(request);
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 {
@@ -356,6 +394,11 @@ pub(crate) fn scope_request_to_account(
{
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(),
@@ -393,6 +436,9 @@ 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) {
@@ -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()
}
@@ -193,7 +193,6 @@ fn client_response_header_is_forwardable(name: &str, connection_options: &[Strin
| "x-openai-organization"
| "x-openai-project"
| "x-codex-installation-id"
| "x-codex-turn-state"
| "x-codex-turn-metadata"
)
}
@@ -1,6 +1,6 @@
//! Codex 本地会话锚点的稳定派生。
//!
//! 旧版 Redis 会话亲和与 WebSocket 复用以 `.runtime/data/identity_hmac_secret`
//! 旧版 Redis 会话亲和与 WebSocket 复用以运行数据目录下的 `identity_hmac_secret`
//! 生成的 `lc_` 值为键。该密钥是持久运行数据,不是账号凭据;保留它可避免
//! 架构迁移后同一客户端会话全部换键。
@@ -8,6 +8,7 @@ use std::{pin::Pin, sync::Arc, time::Duration};
use bytes::Bytes;
use futures::Stream;
use gateway_protocol::openai::events::ParsedRateLimits;
use tokio::sync::Mutex;
use uuid::Uuid;
@@ -36,8 +37,8 @@ pub struct CodexWebSocketStreamingExchange {
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 上游内部 rate-limit 事件里的动态更新。
pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
/// 上游内部 `codex.rate_limits` 事件里的结构化动态更新。
pub rate_limit_updates: CodexWebSocketRateLimitUpdates,
/// 上游内部 metadata 事件里的动态 turn state。
pub turn_state_update: CodexWebSocketTurnStateUpdate,
/// WebSocket 连接池决策。
@@ -53,8 +54,8 @@ pub struct CodexWebSocketStreamingExchange {
/// Responses WebSocket live SSE 字节流。
pub type CodexWebSocketSseStream =
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
/// live 流中的限流头动态更新。
pub type CodexWebSocketRateLimitHeaderUpdates = Arc<Mutex<Vec<(String, String)>>>;
/// live 流中的结构化限流动态更新。
pub type CodexWebSocketRateLimitUpdates = Arc<Mutex<Vec<ParsedRateLimits>>>;
/// live 流中的 turn state 动态更新。
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
@@ -20,7 +20,7 @@ pub(in crate::transport::websocket) enum WebSocketTerminalKind {
}
pub(super) enum ExchangeAction {
RateLimits(Vec<(String, String)>),
RateLimits(events::ParsedRateLimits),
TurnState(String),
Forward {
frame: String,
@@ -41,8 +41,8 @@ pub(super) fn reduce_websocket_event(
};
if let Some(parsed) = events::parse_rate_limits_event(&value) {
let headers = events::rate_limits_to_header_pairs(&parsed);
metadata.rate_limit_headers.extend(headers.iter().cloned());
return Ok(ExchangeAction::RateLimits(headers));
metadata.rate_limit_headers.extend(headers);
return Ok(ExchangeAction::RateLimits(parsed));
}
response_meta::merge_response_metadata(
@@ -16,9 +16,8 @@ use super::super::{
use super::io::{next_websocket_message, receive_idle_timeout, reused_stream_receive_error};
use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event};
use super::{
CodexWebSocketExchangeError, CodexWebSocketRateLimitHeaderUpdates,
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER,
reusable_websocket_metadata,
CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates, CodexWebSocketStreamingExchange,
CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER, reusable_websocket_metadata,
};
pub(in crate::transport::websocket) struct WebSocketStreamPoolReturn {
@@ -63,8 +62,8 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
) -> CodexWebSocketStreamingExchange {
let websocket_connection_id = websocket.connection_id();
let response_metadata = metadata.clone();
let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new()));
let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates);
let rate_limit_updates = Arc::new(Mutex::new(Vec::new()));
let rate_limit_updates_for_task = Arc::clone(&rate_limit_updates);
let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone()));
let turn_state_update_for_task = Arc::clone(&turn_state_update);
let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER);
@@ -83,7 +82,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
reused_connection,
initial_event_timeout,
shutdown,
rate_limit_header_updates: rate_limit_header_updates_for_task,
rate_limit_updates: rate_limit_updates_for_task,
turn_state_update: turn_state_update_for_task,
tx,
})
@@ -105,7 +104,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
turn_state: response_metadata.turn_state,
set_cookie_headers: response_metadata.set_cookie_headers,
rate_limit_headers: response_metadata.rate_limit_headers,
rate_limit_header_updates,
rate_limit_updates,
turn_state_update,
pool_decision: None,
connection_local_continuation: false,
@@ -121,7 +120,7 @@ struct WebSocketStreamForwardState {
reused_connection: bool,
initial_event_timeout: Option<Duration>,
shutdown: CancellationToken,
rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
rate_limit_updates: CodexWebSocketRateLimitUpdates,
turn_state_update: CodexWebSocketTurnStateUpdate,
tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
}
@@ -134,7 +133,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
reused_connection,
initial_event_timeout,
shutdown,
rate_limit_header_updates,
rate_limit_updates,
turn_state_update,
tx,
} = state;
@@ -231,8 +230,8 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
};
let (frame, terminal) = match reduce_websocket_event(&raw, &mut metadata, &mut continuation)
{
Ok(ExchangeAction::RateLimits(headers)) => {
rate_limit_header_updates.lock().await.extend(headers);
Ok(ExchangeAction::RateLimits(rate_limits)) => {
rate_limit_updates.lock().await.push(rate_limits);
continue;
}
Ok(ExchangeAction::TurnState(turn_state)) => {
@@ -21,8 +21,8 @@ pub use self::{
},
error::{CodexWebSocketCloseError, CodexWebSocketExchangeError, CodexWebSocketUpstreamError},
exchange::{
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream,
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate,
CodexWebSocketRateLimitUpdates, CodexWebSocketSseStream, CodexWebSocketStreamingExchange,
CodexWebSocketTurnStateUpdate,
},
handshake::responses_websocket_endpoint,
model::{
+111 -8
View File
@@ -12,9 +12,10 @@ use gateway_admin::model::observability::{
};
use gateway_admin::model::provider_credentials::{
AuthorizationMutationTarget, AuthorizationOwnerBinding, CompleteAuthorization,
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
PrepareCredentialRotation, ProviderDocument, ProviderExportCredentialInput,
ProviderQuotaRequest, QuotaLocalUsageAttribution,
ConsumeProviderResetCredit, PendingAuthorizationMutation, PrepareCredentialImport,
PrepareCredentialRefresh, PrepareCredentialRotation, ProviderDocument,
ProviderExportCredentialInput, ProviderQuotaRequest, ProviderQuotaWindowRole,
QuotaLocalUsageAttribution,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::ProviderAdminErrorKind;
@@ -39,6 +40,9 @@ 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,
@@ -160,6 +164,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();
@@ -566,7 +628,7 @@ async fn openai_admin_provider_projects_official_codex_quota_and_independent_buc
assert_eq!(monthly.len(), 1);
assert!(monthly.iter().any(|window| {
window.label == "月额度"
&& window.source.as_deref() == Some("codex")
&& window.limit_id.as_deref() == Some("codex")
&& window.used_percent == Some(91.0)
}));
assert!(
@@ -579,20 +641,22 @@ async fn openai_admin_provider_projects_official_codex_quota_and_independent_buc
let secondary = quota
.windows
.iter()
.find(|window| window.key == "core-secondary")
.find(|window| window.key == "codex-secondary")
.expect("core secondary quota");
assert_eq!(secondary.label, "次级额度");
assert_eq!(secondary.used_percent, Some(88.0));
let review = quota
.windows
.iter()
.find(|window| window.source.as_deref() == Some("code_review"))
.find(|window| window.limit_id.as_deref() == Some("code_review"))
.expect("code review quota");
assert_eq!(review.label, "代码审查 · 周额度");
assert_eq!(review.label, "周额度");
assert_eq!(review.limit_name.as_deref(), Some("code_review"));
assert_eq!(review.role, Some(ProviderQuotaWindowRole::Primary));
let spark = quota
.windows
.iter()
.find(|window| window.source.as_deref() == Some("codex_bengalfox"))
.find(|window| window.limit_id.as_deref() == Some("codex_bengalfox"))
.expect("Spark quota");
assert_eq!(
(
@@ -818,6 +882,45 @@ 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 provider_ports() -> ProviderStorePorts {
provider_ports_with(
Arc::new(MemoryAccountStore::default()),
@@ -18,6 +18,16 @@ fn openai_config_builds_the_audited_wire_profile() {
assert_eq!(profile.desktop_build, "19012345678");
}
#[test]
fn openai_config_derives_identity_secret_from_runtime_data_dir() {
let mut config = valid_config();
config
.resolve_and_validate(Path::new("/srv/gateway/runtime-data"))
.expect("valid OpenAI config");
assert!(format!("{config:?}").contains("/srv/gateway/runtime-data/identity_hmac_secret"));
}
#[test]
fn openai_config_rejects_noncanonical_versions_and_empty_fields() {
let mut config = valid_config();

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