mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ae4eb49afc | ||
|
|
5d97086511 | ||
|
|
0e6da1e731 | ||
|
|
5d0922b9d1 | ||
|
|
ed3d622554 | ||
|
|
0f71cf6df2 | ||
|
|
a1d3eba874 | ||
|
|
d22c1d4ab9 | ||
|
|
cc8dfe3dd8 | ||
|
|
f25edadd5f | ||
|
|
4f50c67332 | ||
|
|
75f9d2a290 | ||
|
|
b93b8859f9 | ||
|
|
3510a45742 | ||
|
|
85eb283b40 | ||
|
|
b8e1c5077d | ||
|
|
8422240967 | ||
|
|
fbd922117e | ||
|
|
941c4d57fa | ||
|
|
cf046ad9a4 | ||
|
|
7699a954ac | ||
|
|
2c31708d1b | ||
|
|
80f338ce01 | ||
|
|
299088a6e7 | ||
|
|
a23bfca848 | ||
|
|
8951a64a05 | ||
|
|
eb20561893 | ||
|
|
a33637a99e | ||
|
|
97cd7a7097 | ||
|
|
b0fe132736 | ||
|
|
e2c75d1054 | ||
|
|
94f298927d | ||
|
|
3be827204f | ||
|
|
227a1174b9 | ||
|
|
c0defcb6df | ||
|
|
53739ebe5c | ||
|
|
7f420257a6 | ||
|
|
4f85209dc9 | ||
|
|
cd89e2e8fc | ||
|
|
0cb5a29699 | ||
|
|
61788804b4 | ||
|
|
2751ea6aa4 | ||
|
|
58fe66789c | ||
|
|
573b6b26db | ||
|
|
d425428ebe | ||
|
|
b26d3b386a | ||
|
|
cf0b14c5fd | ||
|
|
056561de41 | ||
|
|
f2215df72e | ||
|
|
72f8957163 | ||
|
|
73f0c15790 | ||
|
|
0558693949 | ||
|
|
2a9cb02cc9 | ||
|
|
5ce36931a8 | ||
|
|
0d826d9eba | ||
|
|
304242957a | ||
|
|
ef37ae8d9e | ||
|
|
fc13d4f2b6 |
@@ -45,7 +45,7 @@ jobs:
|
||||
ref: ${{ inputs.ref }}
|
||||
|
||||
- name: Setup Docker Buildx
|
||||
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
|
||||
uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
|
||||
|
||||
- name: Build runtime image
|
||||
uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0
|
||||
@@ -79,7 +79,7 @@ jobs:
|
||||
${{ hashFiles('trivy-results.sarif') != '' &&
|
||||
(github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository) }}
|
||||
uses: github/codeql-action/upload-sarif@5595ccaf912efad79be6eef63a5619ff05969be3 # v4.37.6
|
||||
uses: github/codeql-action/upload-sarif@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4.37.8
|
||||
with:
|
||||
sarif_file: trivy-results.sarif
|
||||
category: runtime-container
|
||||
|
||||
@@ -448,7 +448,7 @@ jobs:
|
||||
PY
|
||||
|
||||
- name: Setup Docker Buildx
|
||||
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
|
||||
uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
|
||||
|
||||
- name: Login to GHCR
|
||||
uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0
|
||||
@@ -519,7 +519,7 @@ jobs:
|
||||
path: dist/release
|
||||
|
||||
- name: Setup Docker Buildx
|
||||
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
|
||||
uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
|
||||
|
||||
- name: Install Cosign
|
||||
uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2
|
||||
|
||||
@@ -12,31 +12,32 @@
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](https://opensource.org/license/mit)
|
||||
|
||||
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运维](#运维) · [部署](deploy/README.md) · [接口](docs/api.md) · [架构](docs/architecture.md)
|
||||
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运行维护](#运行维护) · [部署](deploy/README.md) · [接口](docs/api.md) · [架构](docs/architecture.md)
|
||||
|
||||
</div>
|
||||
|
||||
> [!NOTE]
|
||||
> 只支持 OpenAI Responses API,不提供 `/v1/chat/completions`。
|
||||
> 数据面支持 OpenAI Responses、Images 和模型目录协议,不提供 `/v1/chat/completions`。
|
||||
|
||||
## 能力
|
||||
## 功能
|
||||
|
||||
| 领域 | 实现 |
|
||||
| -------- | ----------------------------------------------------------------------------------------------------- |
|
||||
| 协议 | OpenAI Responses JSON、SSE、WebSocket 与模型目录 |
|
||||
| Provider | 固定 OpenAI、xAI 两个 Provider;各自管理 OAuth、账号、额度与模型目录 |
|
||||
| 路由 | Client Key 绑定 Provider、全局模型映射、会话亲和与同 Provider 安全 fallback |
|
||||
| 透明边界 | OpenAI 请求保留未知字段和字段顺序,SSE 与 WebSocket 原样转发业务字节;xAI 在 Provider 内做协议转换 |
|
||||
| 延续 | OpenAI native/replay continuation;xAI 使用客户端提交的完整历史 |
|
||||
| 账号 | 导入、OAuth、credential 刷新、主动额度刷新,以及正常响应中的被动额度/状态观测 |
|
||||
| 管理 | Client Key、账号、模型目录、设置、观测、系统更新与 OAuth |
|
||||
| 计量 | 模型请求 Token、费用、延迟、账号与 Provider 归因,并记录响应确认的实际 `service_tier` |
|
||||
| 领域 | 当前能力 |
|
||||
| --- | --- |
|
||||
| 数据面 | Responses JSON、SSE、WebSocket、review 子代理、Images 生成/编辑与模型目录 |
|
||||
| Provider | 固定 `openai`、`xai` 两个 Provider,各自拥有 OAuth、credential、额度、目录与 transport |
|
||||
| 路由 | Client Key 限定账号分组;按模型能力编译 Provider 候选;支持全局精确模型映射、会话亲和和安全 fallback |
|
||||
| 透明边界 | OpenAI 保留未知请求字段及字段顺序,并原样转发 SSE、WebSocket 与 Images 业务字节;xAI 在 Provider 内转换协议 |
|
||||
| 账号 | OAuth、OpenAI AT/RT、账号文件与 Agent Identity 导入,credential/额度刷新、账号恢复及主动额度重置卡 |
|
||||
| 管理 | 账号与分组、Client Key、运行设置、用量/错误观测、S3/R2 备份、系统更新与回滚 |
|
||||
| 计量 | Token、图片、费用、延迟、账号与 Provider 归因,并记录上游响应确认的实际 `service_tier` |
|
||||
|
||||
架构边界、状态所有权和请求不变量见 [架构文档](docs/architecture.md)。
|
||||
|
||||
## 快速开始
|
||||
|
||||
需要 Docker Engine 与 Docker Compose Plugin。
|
||||
|
||||
### 1. 准备
|
||||
### 1. 准备配置和目录
|
||||
|
||||
```bash
|
||||
git clone https://github.com/zyycn/codex-proxy-rs.git
|
||||
@@ -56,16 +57,16 @@ openssl rand -hex 24
|
||||
openssl rand -hex 24
|
||||
```
|
||||
|
||||
把两个结果写入 `deploy/config.yaml`,并单独设置管理员初始密码:
|
||||
把两个结果分别写入 `deploy/config.yaml`,再设置管理员初始密码:
|
||||
|
||||
| 配置 | 约束 |
|
||||
| -------------------------- | ---------------------------- |
|
||||
| `store.database.password` | 48 位十六进制 |
|
||||
| `store.redis.password` | 48 位十六进制 |
|
||||
| `admin.default_password` | 至少 12 位,不能包含 `$` |
|
||||
| 配置 | 约束 |
|
||||
| --- | --- |
|
||||
| `store.database.password` | 48 位十六进制 |
|
||||
| `store.redis.password` | 48 位十六进制 |
|
||||
| `admin.default_password` | 至少 12 位,不能包含 `$` |
|
||||
|
||||
PostgreSQL、Redis 容器和应用通过同一份 `config.yaml` 使用这两个基础设施密码,不需要额外
|
||||
导出环境变量。管理员密码只在首次创建管理员时生效。
|
||||
PostgreSQL、Redis 和应用从同一份 `config.yaml` 读取基础设施密码,不需要额外导出环境变量。
|
||||
管理员密码只在首次创建管理员时生效。
|
||||
|
||||
Linux 需要允许容器组写入运行目录:
|
||||
|
||||
@@ -83,37 +84,36 @@ docker compose -f deploy/compose.yaml up -d --no-build
|
||||
curl -i http://127.0.0.1:8080/healthz
|
||||
```
|
||||
|
||||
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址:
|
||||
`204 No Content` 表示应用、PostgreSQL 与 Redis 均可用。管理端地址为
|
||||
`http://127.0.0.1:8080`。
|
||||
|
||||
### 3. 初始化
|
||||
### 3. 初始化网关
|
||||
|
||||
1. 使用 `admin@cpr.local` 与初始密码登录。
|
||||
2. 在账号页选择对应 Provider:OpenAI 支持 OAuth、账号文件与 Agent Identity;xAI 支持 OAuth
|
||||
和兼容账号文件导入。
|
||||
3. 按需配置客户端模型到上游模型的全局映射;未命中映射的模型名原样透传。
|
||||
4. 创建 `sk_...` Client Key,绑定 `openai` 或 `xai`,并设置速率与并发限制。
|
||||
2. 导入或授权 Provider 账号:OpenAI 支持 OAuth、逐行 AT/RT、OAuth JSON 与 Agent Identity;
|
||||
xAI 支持 OAuth 和账号 JSON。
|
||||
3. 创建账号分组并按需配置全局模型映射。未命中的模型名原样交给候选 Provider。
|
||||
4. 创建 `sk_...` Client Key,选择可用账号分组,并设置速率和并发限制;空分组范围表示全部账号。
|
||||
|
||||
OpenAI 账号导入和首次 OAuth 创建在 credential 提交后会立即尝试一次额度观测;额度接口临时失败不会
|
||||
回滚已提交的账号。Free、K12 等套餐共用同一套额度与状态转换逻辑,套餐字段只作为账号事实并隔离
|
||||
对应的模型目录 cache。credential 刷新、额度刷新和正常请求中的被动额度观测是三条独立链路,详见
|
||||
[账号接口](docs/api.md#5-账号)。
|
||||
OpenAI RT 导入会先向官方 token endpoint 换取 AT;仅 AT 的账号不具备自动续期能力,
|
||||
能否参与调度仍取决于 token 中可验证的身份与后续上游事实。账号导入和首次 OAuth 创建会尽力读取一次额度,
|
||||
额度上游暂时失败不会回滚已提交的账号。
|
||||
|
||||
OpenAI 账号的主动重置卡只在用户打开相应弹窗或手工刷新时查询。消费成功后管理端重新查询卡片与额度;
|
||||
系统不会通过修改本地 `reset_at` 伪造重置。完整账号合同见 [账号接口](docs/api.md#5-账号)。
|
||||
|
||||
> [!IMPORTANT]
|
||||
> xAI 使用 OAuth session,不支持把 xAI API Key 作为上游 credential。
|
||||
|
||||
> [!NOTE]
|
||||
> `admin.default_password` 只在首次创建管理员时生效。
|
||||
|
||||
完整部署、权限和密码轮换规则见 [deploy/README.md](deploy/README.md)。
|
||||
完整的权限、密码轮换、备份和升级规则见 [部署文档](deploy/README.md)。
|
||||
|
||||
## 客户端接入
|
||||
|
||||
| 配置 | 值 |
|
||||
| -------- | ---------------------------------------- |
|
||||
| Base URL | `http://127.0.0.1:8080/v1` |
|
||||
| API Key | 管理端创建的 `sk_...` |
|
||||
| 鉴权 | `Authorization: Bearer <client-api-key>` |
|
||||
| 配置 | 值 |
|
||||
| --- | --- |
|
||||
| Base URL | `http://127.0.0.1:8080/v1` |
|
||||
| API Key | 管理端创建的 `sk_...` |
|
||||
| 鉴权 | `Authorization: Bearer <client-api-key>` |
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8080/v1/models \
|
||||
@@ -130,67 +130,70 @@ curl http://127.0.0.1:8080/v1/responses \
|
||||
```
|
||||
|
||||
<details>
|
||||
<summary>API 路由</summary>
|
||||
<summary>公开数据面路由</summary>
|
||||
|
||||
| 路由 | 用途 |
|
||||
| ---------------------------------- | ------------------------------ |
|
||||
| `POST /v1/responses` | JSON 响应或 SSE Responses 流 |
|
||||
| `GET /v1/responses` | Responses WebSocket 升级 |
|
||||
| `POST /v1/responses/review` | review 子代理请求 |
|
||||
| `GET /v1/models` | 当前 Client Key 可用的模型列表 |
|
||||
| `GET /v1/models/catalog` | Codex 客户端展示用模型目录 |
|
||||
| `GET /v1/models/{model_id}/info` | Codex 客户端展示用模型详情 |
|
||||
| `GET /v1/models/{model_id}` | OpenAI 兼容模型详情 |
|
||||
| 路由 | 用途 |
|
||||
| --- | --- |
|
||||
| `POST /v1/responses` | JSON 响应或 SSE Responses 流 |
|
||||
| `GET /v1/responses` | Responses WebSocket 升级 |
|
||||
| `POST /v1/responses/review` | review 子代理请求 |
|
||||
| `POST /v1/images/generations` | OpenAI Images 生成 JSON 透明转发 |
|
||||
| `POST /v1/images/edits` | OpenAI Images 编辑 JSON 透明转发 |
|
||||
| `GET /v1/models` | 当前 Client Key 账号范围内的模型列表 |
|
||||
| `GET /v1/models/catalog` | Codex 客户端展示用模型目录 |
|
||||
| `GET /v1/models/{model_id}/info` | Codex 客户端展示用模型详情 |
|
||||
| `GET /v1/models/{model_id}` | OpenAI 兼容模型详情 |
|
||||
|
||||
所有 `/v1/*` 路由都需要客户端 API Key。
|
||||
所有 `/v1/*` 路由都需要 Client Key。Images 是 OpenAI Provider 自有端点,不读取模型字段,也不参与
|
||||
文本模型映射。
|
||||
|
||||
</details>
|
||||
|
||||
完整的客户端与管理端路由、鉴权和 mutation 规则见 [接口文档](docs/api.md)。
|
||||
完整的客户端与管理端路由、鉴权、请求字段和 mutation 语义见 [接口文档](docs/api.md)。
|
||||
|
||||
## 运维
|
||||
## 运行维护
|
||||
|
||||
```bash
|
||||
# 实例升级
|
||||
# 同一大版本内升级发布镜像
|
||||
docker compose -f deploy/compose.yaml pull codex-proxy-rs
|
||||
docker compose -f deploy/compose.yaml up -d --no-build
|
||||
|
||||
# 日志
|
||||
# 查看日志
|
||||
docker compose -f deploy/compose.yaml logs -f codex-proxy-rs
|
||||
|
||||
# 从源码构建
|
||||
# 从当前源码构建
|
||||
docker compose -f deploy/compose.yaml build codex-proxy-rs
|
||||
docker compose -f deploy/compose.yaml up -d
|
||||
```
|
||||
|
||||
> [!IMPORTANT]
|
||||
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点密钥和自更新状态。删除该目录会永久清除
|
||||
> 运行状态。
|
||||
> `.runtime/` 保存 PostgreSQL、Redis、日志、OpenAI 会话锚点和在线更新状态。删除该目录会永久清除
|
||||
> 本地运行状态。
|
||||
|
||||
> [!WARNING]
|
||||
> 不支持跨大版本在线升级。升级到新的大版本时,请使用全新的 `.runtime/` 数据目录重新部署,
|
||||
> 并重新导入或重新授权 Provider 账号与客户端 Key。
|
||||
> 不支持跨大版本在线升级。升级到新的大版本时请使用全新的 `.runtime/` 数据目录,并重新导入或授权
|
||||
> Provider 账号与 Client Key。
|
||||
|
||||
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket
|
||||
upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
|
||||
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket upgrade 与真实客户端
|
||||
IP;不要暴露 PostgreSQL 或 Redis。当前运行拓扑为单副本,不能通过复制应用容器扩容。
|
||||
|
||||
## 维护与发布
|
||||
|
||||
发布仓库版本必须从干净、已同步上游的发布分支运行:
|
||||
仓库版本只能从干净、已同步上游的发布分支发布:
|
||||
|
||||
```bash
|
||||
release/publish <version>
|
||||
```
|
||||
|
||||
脚本依赖已登录的 GitHub CLI(`gh auth login`)。它会更新 `release/version.yaml`、创建约定提交与带
|
||||
注释的 `v<version>` tag,原子推送分支和 tag,再从发布分支触发 Release 构建,使构建缓存可以跨版本
|
||||
复用。不要手工修改版本后单独打 tag。该脚本不会登录服务器,也不会改变任何运行实例;源码提交、
|
||||
仓库发版和实例升级是三个独立状态。
|
||||
脚本依赖已登录的 GitHub CLI(`gh auth login`),会更新 `release/version.yaml`、创建约定提交和带注释的
|
||||
`v<version>` tag,原子推送分支与 tag,再触发 Release 构建。源码提交、仓库发版和运行实例升级是三个
|
||||
独立状态;脚本不会登录服务器或改变任何运行实例。
|
||||
|
||||
## 文档
|
||||
|
||||
- [部署](deploy/README.md)
|
||||
- [部署与运维](deploy/README.md)
|
||||
- [客户端与管理端接口](docs/api.md)
|
||||
- [架构与数据边界](docs/architecture.md)
|
||||
- [系统架构](docs/architecture.md)
|
||||
- [数据库迁移规则](backend/migrations/README.md)
|
||||
- [Release](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
- [容器镜像](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
|
||||
Generated
+45
-45
@@ -120,9 +120,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.91"
|
||||
version = "0.1.92"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec"
|
||||
checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -216,9 +216,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-sdk-s3"
|
||||
version = "1.141.0"
|
||||
version = "1.142.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d9f9420d3a2467eed22ed3635ca653653162c386a0b0f65c78189f9bd3c1379e"
|
||||
checksum = "f9e15a5c55e05f4b0b7e483160b3c85cccdf77cff02c95504f3e71d460855cd2"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"aws-credential-types",
|
||||
@@ -308,9 +308,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-eventstream"
|
||||
version = "0.61.1"
|
||||
version = "0.61.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
|
||||
checksum = "6de526c7b567420a31bc283657a7921b45c4cafe0827fdf2490713dcc770c28f"
|
||||
dependencies = [
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
@@ -341,9 +341,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-http-client"
|
||||
version = "1.2.0"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e"
|
||||
checksum = "3c1c8a04cb31ba74d0115af5a890bb8c0d48fba64b52812fa13929a6ef0cc83c"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api",
|
||||
@@ -385,9 +385,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-runtime"
|
||||
version = "1.12.1"
|
||||
version = "1.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "07505b34e8f4b3591a4fa69e9792b52289b95488dbbc68c3c0075b7bedb245e1"
|
||||
checksum = "483b858ff67522011c4786310c5cd8fd88d0be7ea3d5f1a48328446300c4269e"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-http",
|
||||
@@ -411,9 +411,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-runtime-api"
|
||||
version = "1.14.0"
|
||||
version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3b98f2e1fd67ec06618f9c291e5e495a468e60519e44c9c1979cd0521f3affdb"
|
||||
checksum = "954c563ce84507722d2679f07a35d21b9c6466b3872d513020d0281fc8112ac9"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api-macros",
|
||||
@@ -451,9 +451,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-types"
|
||||
version = "1.6.1"
|
||||
version = "1.6.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8"
|
||||
checksum = "fce83ce9abbb198d25bc7131e468d0f9fe1257125e58c39f3f9fc9f5098c9647"
|
||||
dependencies = [
|
||||
"base64-simd",
|
||||
"bytes",
|
||||
@@ -801,9 +801,9 @@ checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
|
||||
|
||||
[[package]]
|
||||
name = "cookie"
|
||||
version = "0.18.1"
|
||||
version = "0.18.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747"
|
||||
checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87"
|
||||
dependencies = [
|
||||
"percent-encoding",
|
||||
"time",
|
||||
@@ -1384,9 +1384,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218"
|
||||
checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
@@ -1399,9 +1399,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "futures-channel"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae"
|
||||
checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
@@ -1409,15 +1409,15 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "futures-core"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7"
|
||||
checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
|
||||
|
||||
[[package]]
|
||||
name = "futures-executor"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458"
|
||||
checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-task",
|
||||
@@ -1437,32 +1437,32 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "futures-io"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a"
|
||||
checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b"
|
||||
checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307"
|
||||
checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d"
|
||||
|
||||
[[package]]
|
||||
name = "futures-task"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109"
|
||||
checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd"
|
||||
|
||||
[[package]]
|
||||
name = "futures-timer"
|
||||
@@ -1472,9 +1472,9 @@ checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968"
|
||||
|
||||
[[package]]
|
||||
name = "futures-util"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa"
|
||||
checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
@@ -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",
|
||||
@@ -2813,9 +2813,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.5.0"
|
||||
version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3257df217f7eab0044627a268c9cc6cdb60c0c421c88f83ac41c4e31520b6b84"
|
||||
checksum = "e37a4ca5c6ca42aa3e6df2fd32b987a65d32a4c2159a6f3fe0fd1df306a2658f"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
@@ -3633,18 +3633,18 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.19"
|
||||
version = "2.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9"
|
||||
checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "2.0.19"
|
||||
version = "2.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd"
|
||||
checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -4100,9 +4100,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.24.0"
|
||||
version = "1.24.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239"
|
||||
checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"js-sys",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -347,7 +347,10 @@ fn map_provider_registry_error(error: ProviderAdminError) -> AdminError {
|
||||
}
|
||||
ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound,
|
||||
ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict,
|
||||
ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable,
|
||||
ProviderAdminErrorKind::Unavailable | ProviderAdminErrorKind::CredentialRefreshRequired => {
|
||||
AdminErrorKind::Unavailable
|
||||
}
|
||||
ProviderAdminErrorKind::BadGateway => AdminErrorKind::BadGateway,
|
||||
ProviderAdminErrorKind::Internal => AdminErrorKind::Internal,
|
||||
};
|
||||
AdminError::new(kind, "Provider registry initialization failed")
|
||||
|
||||
@@ -288,7 +288,7 @@ pub struct BackupRecord {
|
||||
pub error_message: Option<String>,
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
/// 手动过期时间;创建时确定,到期自动清理。
|
||||
/// 创建时确定的过期时间;手工和计划备份均可设置,到期自动清理。
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
|
||||
@@ -557,7 +557,7 @@ pub struct UsageRecord {
|
||||
pub operation: String,
|
||||
pub endpoint: String,
|
||||
pub client_transport: String,
|
||||
pub requested_model_id: String,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub provider_account_name: Option<String>,
|
||||
|
||||
@@ -3,15 +3,17 @@
|
||||
use std::fmt;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::{DateTime, NaiveDate, Utc};
|
||||
use gateway_core::{
|
||||
engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
|
||||
routing::{ProviderKind, UpstreamModelId},
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{
|
||||
AdminError, MutationActor, MutationContext, PageSize, Revision,
|
||||
accounts::{AccountRecord, AccountSummary, AccountUsage, CredentialState},
|
||||
observability::CurrencyCost,
|
||||
};
|
||||
|
||||
/// Provider-owned JSON;公共层只搬运且 Debug 不输出值。
|
||||
@@ -570,13 +572,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 +628,154 @@ pub struct ProviderQuota {
|
||||
pub provider_data: Option<ProviderDocument>,
|
||||
}
|
||||
|
||||
/// Provider 官方用量统计的账号形态。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ProviderUsageStatisticsMode {
|
||||
Workspace,
|
||||
Personal,
|
||||
}
|
||||
|
||||
impl ProviderUsageStatisticsMode {
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Workspace => "workspace",
|
||||
Self::Personal => "personal",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 官方用量统计的目标额度周期。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderUsageStatisticsCycle {
|
||||
pub offset: i8,
|
||||
pub start_at: DateTime<Utc>,
|
||||
pub end_at: DateTime<Utc>,
|
||||
pub window_seconds: u64,
|
||||
pub used_percent: Option<f64>,
|
||||
pub is_current: bool,
|
||||
pub can_go_previous: bool,
|
||||
pub can_go_next: bool,
|
||||
}
|
||||
|
||||
/// Provider 官方用量统计查询。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderUsageStatisticsRequest {
|
||||
pub account_id: ProviderAccountId,
|
||||
/// `0` 为当前周期,负数为历史周期。
|
||||
pub cycle_offset: i8,
|
||||
/// 客户端相对 UTC 的分钟偏移,例如 UTC+8 为 `480`。
|
||||
pub utc_offset_minutes: i16,
|
||||
}
|
||||
|
||||
/// 官方用量统计中的 Token 分类。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct ProviderUsageStatisticsTokens {
|
||||
pub uncached_input: u64,
|
||||
pub cached_input: u64,
|
||||
pub output: u64,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
/// 官方模型报表中的服务档位。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ProviderUsageStatisticsServiceTier {
|
||||
Standard,
|
||||
Fast,
|
||||
}
|
||||
|
||||
impl ProviderUsageStatisticsServiceTier {
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Standard => "standard",
|
||||
Self::Fast => "fast",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 一个模型与服务档位在目标周期内的统计行。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderUsageStatisticsModel {
|
||||
pub key: String,
|
||||
pub model: String,
|
||||
pub service_tier: ProviderUsageStatisticsServiceTier,
|
||||
pub credit_share: Option<f64>,
|
||||
pub quota_share: Option<f64>,
|
||||
pub tokens: ProviderUsageStatisticsTokens,
|
||||
pub estimated_cost: Option<CurrencyCost>,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_estimated_allocation: bool,
|
||||
pub has_rate_fallback: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
/// 目标周期内的一天;官方日报只能提供本地日粒度。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderUsageStatisticsDay {
|
||||
pub date: NaiveDate,
|
||||
pub credit_share: Option<f64>,
|
||||
pub tokens: ProviderUsageStatisticsTokens,
|
||||
pub estimated_cost: Option<CurrencyCost>,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
pub is_boundary_day: bool,
|
||||
}
|
||||
|
||||
/// 目标周期汇总。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderUsageStatisticsSummary {
|
||||
pub tokens: ProviderUsageStatisticsTokens,
|
||||
pub estimated_cost: Option<CurrencyCost>,
|
||||
pub projected_tokens: Option<u64>,
|
||||
pub projected_cost: Option<CurrencyCost>,
|
||||
pub day_count: u32,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
/// Provider 已解释、计算并排序的官方用量统计。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderUsageStatistics {
|
||||
pub mode: ProviderUsageStatisticsMode,
|
||||
pub cycle: ProviderUsageStatisticsCycle,
|
||||
pub summary: ProviderUsageStatisticsSummary,
|
||||
pub models: Vec<ProviderUsageStatisticsModel>,
|
||||
pub daily: Vec<ProviderUsageStatisticsDay>,
|
||||
}
|
||||
|
||||
/// Provider 返回的一张安全主动额度重置卡。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderResetCredit {
|
||||
pub id: String,
|
||||
pub status: Option<String>,
|
||||
pub title: Option<String>,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
pub reset_type: Option<String>,
|
||||
}
|
||||
|
||||
/// Provider 主动额度重置卡列表。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderResetCredits {
|
||||
pub available_count: u64,
|
||||
pub credits: Vec<ProviderResetCredit>,
|
||||
}
|
||||
|
||||
/// 一次主动额度重置卡消费命令。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ConsumeProviderResetCredit {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub credit_id: Option<String>,
|
||||
pub redeem_request_id: Uuid,
|
||||
}
|
||||
|
||||
/// Provider 返回的主动额度重置卡消费结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderResetCreditResult {
|
||||
pub code: String,
|
||||
pub credit: Option<ProviderResetCredit>,
|
||||
}
|
||||
|
||||
impl ProviderQuota {
|
||||
/// 返回账号展开面板使用的当前额度窗口用量。
|
||||
///
|
||||
@@ -627,8 +798,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 +859,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 +874,18 @@ impl ProviderQuota {
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_usage_priority(window: &ProviderQuotaWindow) -> u8 {
|
||||
match window.group.as_str() {
|
||||
fn quota_usage_priority(window: &ProviderQuotaWindow) -> (u8, u8) {
|
||||
let attribution = match window.local_usage_attribution {
|
||||
QuotaLocalUsageAttribution::AccountWide => 0,
|
||||
QuotaLocalUsageAttribution::Unavailable => 1,
|
||||
};
|
||||
let duration = match window.group.as_str() {
|
||||
"shortTerm" if window.window_seconds.is_some_and(is_week_window) => 1,
|
||||
"shortTerm" => 0,
|
||||
"monthly" => 2,
|
||||
_ => 3,
|
||||
}
|
||||
};
|
||||
(attribution, duration)
|
||||
}
|
||||
|
||||
fn is_week_window(seconds: u64) -> bool {
|
||||
|
||||
@@ -13,11 +13,12 @@ use crate::model::observability::{
|
||||
CalculatedBillingBreakdown, DashboardWireProfile, ProviderBillingInput,
|
||||
};
|
||||
use crate::model::provider_credentials::{
|
||||
AuthorizationStarted, CompleteAuthorization, PendingAuthorizationMutation,
|
||||
PrepareCredentialImport, PrepareCredentialRefresh, PrepareCredentialRotation,
|
||||
PreparedAuthorizationCommit, PreparedCredentialImport, PreparedCredentialRotation,
|
||||
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
|
||||
ProviderQuotaRequest,
|
||||
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
|
||||
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
|
||||
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
|
||||
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, ProviderModels,
|
||||
ProviderQuota, ProviderQuotaRequest, ProviderResetCreditResult, ProviderResetCredits,
|
||||
ProviderUsageStatistics, ProviderUsageStatisticsRequest,
|
||||
};
|
||||
|
||||
/// Provider 管理失败的稳定分类。
|
||||
@@ -28,26 +29,56 @@ pub enum ProviderAdminErrorKind {
|
||||
NotFound,
|
||||
Conflict,
|
||||
Unavailable,
|
||||
CredentialRefreshRequired,
|
||||
BadGateway,
|
||||
Internal,
|
||||
}
|
||||
|
||||
/// 不泄漏 OAuth material 与 Provider 响应体的管理错误。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
/// 不携带 OAuth 请求材料的管理错误。
|
||||
///
|
||||
/// `message` 用于当前认证管理请求展示;上游没有标准 `error.message` 时,
|
||||
/// Provider 可以按其官方契约回退为完整的非成功响应正文。
|
||||
#[derive(Clone, PartialEq, Eq, thiserror::Error)]
|
||||
#[error("provider admin operation failed: {kind:?}")]
|
||||
pub struct ProviderAdminError {
|
||||
kind: ProviderAdminErrorKind,
|
||||
message: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ProviderAdminError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ProviderAdminError")
|
||||
.field("kind", &self.kind)
|
||||
.field("message", &self.message.as_ref().map(|_| "<redacted>"))
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderAdminError {
|
||||
#[must_use]
|
||||
pub const fn new(kind: ProviderAdminErrorKind) -> Self {
|
||||
Self { kind }
|
||||
Self {
|
||||
kind,
|
||||
message: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_message(mut self, message: impl Into<String>) -> Self {
|
||||
self.message = Some(message.into());
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn kind(&self) -> ProviderAdminErrorKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn message(&self) -> Option<&str> {
|
||||
self.message.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// 一个具体 Provider 对管理控制面提供的解析、验证、上游交互与运行时资源回收能力。
|
||||
@@ -116,6 +147,30 @@ pub trait ProviderAdmin: Send + Sync {
|
||||
request: ProviderQuotaRequest,
|
||||
) -> Result<ProviderQuota, ProviderAdminError>;
|
||||
|
||||
/// 查询 Provider 官方用量统计;不支持该能力的 Provider 使用默认拒绝。
|
||||
async fn usage_statistics(
|
||||
&self,
|
||||
_request: ProviderUsageStatisticsRequest,
|
||||
) -> Result<ProviderUsageStatistics, ProviderAdminError> {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
/// 查询 Provider 主动额度重置卡;不支持该能力的 Provider 使用默认拒绝。
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderResetCredits, ProviderAdminError> {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
/// 消费 Provider 主动额度重置卡;不支持该能力的 Provider 使用默认拒绝。
|
||||
async fn consume_reset_credit(
|
||||
&self,
|
||||
_command: ConsumeProviderResetCredit,
|
||||
) -> Result<ProviderResetCreditResult, ProviderAdminError> {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
async fn models(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
|
||||
@@ -24,8 +24,10 @@ use crate::{
|
||||
observability::TimeRange,
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
PrepareCredentialRefresh, ProviderModels, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderQuotaWindow, QuotaLocalUsageAttribution,
|
||||
ConsumeProviderResetCredit, PrepareCredentialRefresh, ProviderModels, ProviderQuota,
|
||||
ProviderQuotaRequest, ProviderQuotaWindow, ProviderResetCreditResult,
|
||||
ProviderResetCredits, ProviderUsageStatistics, ProviderUsageStatisticsRequest,
|
||||
QuotaLocalUsageAttribution,
|
||||
},
|
||||
},
|
||||
ports::{provider::ProviderAdminRegistry, store::AccountStore},
|
||||
@@ -79,6 +81,35 @@ pub trait AccountsService: Send + Sync {
|
||||
refresh: bool,
|
||||
) -> Result<AccountDirectoryItem, AdminError>;
|
||||
|
||||
async fn usage_statistics(
|
||||
&self,
|
||||
_request: ProviderUsageStatisticsRequest,
|
||||
) -> Result<ProviderUsageStatistics, AdminError> {
|
||||
Err(AdminError::invalid(
|
||||
"Provider usage statistics are not supported",
|
||||
))
|
||||
}
|
||||
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
_context: &MutationContext,
|
||||
_account_id: ProviderAccountId,
|
||||
) -> Result<ProviderResetCredits, AdminError> {
|
||||
Err(AdminError::invalid(
|
||||
"Provider reset credits are not supported",
|
||||
))
|
||||
}
|
||||
|
||||
async fn consume_reset_credit(
|
||||
&self,
|
||||
_context: &MutationContext,
|
||||
_command: ConsumeProviderResetCredit,
|
||||
) -> Result<ProviderResetCreditResult, AdminError> {
|
||||
Err(AdminError::invalid(
|
||||
"Provider reset credits are not supported",
|
||||
))
|
||||
}
|
||||
|
||||
async fn models(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -97,11 +128,13 @@ pub(crate) struct DefaultAccountsService {
|
||||
providers: ProviderAdminRegistry,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
probe: Arc<dyn AccountProbe>,
|
||||
reset_credit_locks:
|
||||
Arc<futures::lock::Mutex<BTreeMap<ProviderAccountId, Arc<futures::lock::Mutex<()>>>>>,
|
||||
}
|
||||
|
||||
impl DefaultAccountsService {
|
||||
#[must_use]
|
||||
pub(crate) const fn new(
|
||||
pub(crate) fn new(
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
providers: ProviderAdminRegistry,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
@@ -112,9 +145,22 @@ impl DefaultAccountsService {
|
||||
providers,
|
||||
snapshot,
|
||||
probe,
|
||||
reset_credit_locks: Arc::new(futures::lock::Mutex::new(BTreeMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn reset_credit_lock(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Arc<futures::lock::Mutex<()>> {
|
||||
let mut locks = self.reset_credit_locks.lock().await;
|
||||
Arc::clone(
|
||||
locks
|
||||
.entry(account_id.clone())
|
||||
.or_insert_with(|| Arc::new(futures::lock::Mutex::new(()))),
|
||||
)
|
||||
}
|
||||
|
||||
async fn load_account(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -502,6 +548,68 @@ impl AccountsService for DefaultAccountsService {
|
||||
self.load_directory_item(account_id, refresh).await
|
||||
}
|
||||
|
||||
async fn usage_statistics(
|
||||
&self,
|
||||
request: ProviderUsageStatisticsRequest,
|
||||
) -> Result<ProviderUsageStatistics, AdminError> {
|
||||
let (_, provider) = self.provider_for_account(&request.account_id).await?;
|
||||
provider
|
||||
.usage_statistics(request)
|
||||
.await
|
||||
.map_err(|error| map_provider_error(error, "provider usage statistics"))
|
||||
}
|
||||
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
account_id: ProviderAccountId,
|
||||
) -> Result<ProviderResetCredits, AdminError> {
|
||||
let (_, provider) = self.provider_for_account(&account_id).await?;
|
||||
match provider.reset_credits(&account_id).await {
|
||||
Ok(credits) => Ok(credits),
|
||||
Err(error)
|
||||
if error.kind()
|
||||
== crate::ports::provider::ProviderAdminErrorKind::CredentialRefreshRequired =>
|
||||
{
|
||||
self.refresh(context, account_id.clone()).await?;
|
||||
let (_, provider) = self.provider_for_account(&account_id).await?;
|
||||
provider
|
||||
.reset_credits(&account_id)
|
||||
.await
|
||||
.map_err(map_reset_credits_error_after_refresh)
|
||||
}
|
||||
Err(error) => Err(map_provider_error(error, "provider reset credits")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn consume_reset_credit(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: ConsumeProviderResetCredit,
|
||||
) -> Result<ProviderResetCreditResult, AdminError> {
|
||||
let account_id = command.account_id.clone();
|
||||
// 覆盖 credential refresh + 同键重试的完整账号级临界区,避免 401 两次调用
|
||||
// 之间插入另一笔不可逆消费。
|
||||
let lock = self.reset_credit_lock(&account_id).await;
|
||||
let _guard = lock.lock().await;
|
||||
let (_, provider) = self.provider_for_account(&account_id).await?;
|
||||
match provider.consume_reset_credit(command.clone()).await {
|
||||
Ok(result) => Ok(result),
|
||||
Err(error)
|
||||
if error.kind()
|
||||
== crate::ports::provider::ProviderAdminErrorKind::CredentialRefreshRequired =>
|
||||
{
|
||||
self.refresh(context, account_id.clone()).await?;
|
||||
let (_, provider) = self.provider_for_account(&account_id).await?;
|
||||
provider
|
||||
.consume_reset_credit(command)
|
||||
.await
|
||||
.map_err(map_reset_credits_error_after_refresh)
|
||||
}
|
||||
Err(error) => Err(map_provider_error(error, "provider reset-credit consume")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn models(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -565,11 +673,7 @@ impl AccountsService for DefaultAccountsService {
|
||||
let upstream_body = error
|
||||
.upstream_response()
|
||||
.map(|response| String::from_utf8_lossy(response.body()).into_owned());
|
||||
let message = upstream_body
|
||||
.as_deref()
|
||||
.filter(|body| !body.is_empty())
|
||||
.unwrap_or_else(|| error.client_message())
|
||||
.to_owned();
|
||||
let message = error.client_message().to_owned();
|
||||
vec![AccountConnectionTestEvent::Failed {
|
||||
message,
|
||||
provider_error_code: error.client_error_code().map(ToOwned::to_owned),
|
||||
@@ -586,6 +690,19 @@ impl AccountsService for DefaultAccountsService {
|
||||
}
|
||||
}
|
||||
|
||||
fn map_reset_credits_error_after_refresh(
|
||||
error: crate::ports::provider::ProviderAdminError,
|
||||
) -> AdminError {
|
||||
if error.kind() == crate::ports::provider::ProviderAdminErrorKind::CredentialRefreshRequired {
|
||||
return AdminError::bad_gateway(
|
||||
error
|
||||
.message()
|
||||
.unwrap_or("OpenAI reset-credit upstream rejected the refreshed credential"),
|
||||
);
|
||||
}
|
||||
map_provider_error(error, "provider reset credits")
|
||||
}
|
||||
|
||||
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影。
|
||||
fn empty_quota() -> ProviderQuota {
|
||||
ProviderQuota {
|
||||
|
||||
@@ -57,9 +57,15 @@ 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,
|
||||
};
|
||||
AdminError::new(kind, format!("{resource} operation failed"))
|
||||
let message = error
|
||||
.message()
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| format!("{resource} operation failed"));
|
||||
AdminError::new(kind, message)
|
||||
}
|
||||
|
||||
async fn publish_committed(
|
||||
|
||||
@@ -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,
|
||||
@@ -60,7 +62,7 @@ pub(super) type EventLog = Arc<Mutex<Vec<&'static str>>>;
|
||||
pub(super) struct FakeProviderAdmin {
|
||||
kind: ProviderKind,
|
||||
events: EventLog,
|
||||
failure: Mutex<Option<ProviderAdminErrorKind>>,
|
||||
failure: Mutex<Option<ProviderAdminError>>,
|
||||
quota_failure: Mutex<Option<ProviderAdminErrorKind>>,
|
||||
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
|
||||
retry_authorization_after_abort: Mutex<bool>,
|
||||
@@ -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,11 +88,17 @@ 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()),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn fail_next(&self, kind: ProviderAdminErrorKind) {
|
||||
*self.failure.lock().expect("provider failure") = Some(kind);
|
||||
*self.failure.lock().expect("provider failure") = Some(ProviderAdminError::new(kind));
|
||||
}
|
||||
|
||||
fn fail_next_with_message(&self, kind: ProviderAdminErrorKind, message: &str) {
|
||||
*self.failure.lock().expect("provider failure") =
|
||||
Some(ProviderAdminError::new(kind).with_message(message));
|
||||
}
|
||||
|
||||
pub(super) fn fail_next_quota(&self, kind: ProviderAdminErrorKind) {
|
||||
@@ -113,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;
|
||||
}
|
||||
@@ -145,7 +161,7 @@ impl FakeProviderAdmin {
|
||||
|
||||
fn require_available(&self) -> Result<(), ProviderAdminError> {
|
||||
match self.failure.lock().expect("provider failure").take() {
|
||||
Some(kind) => Err(ProviderAdminError::new(kind)),
|
||||
Some(error) => Err(error),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -344,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,
|
||||
@@ -1257,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),
|
||||
@@ -1386,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),
|
||||
@@ -1431,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),
|
||||
@@ -1477,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),
|
||||
@@ -1542,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),
|
||||
@@ -1554,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),
|
||||
@@ -1650,6 +1694,59 @@ 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();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let upstream_message = "Your refresh token has already been used.";
|
||||
provider.fail_next_with_message(ProviderAdminErrorKind::Conflict, upstream_message);
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let services = accounts_service(provider, store).await;
|
||||
|
||||
let error = services
|
||||
.accounts()
|
||||
.refresh(
|
||||
&context("refresh-provider-message"),
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
)
|
||||
.await
|
||||
.expect_err("Provider preparation must fail");
|
||||
|
||||
assert_eq!(error.kind(), gateway_admin::model::AdminErrorKind::Conflict);
|
||||
assert_eq!(error.to_string(), upstream_message);
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
["store.load_account", "provider.prepare_refresh"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_refresh_store_failure_should_drop_guard_after_commit_attempt() {
|
||||
let events = events();
|
||||
|
||||
@@ -662,7 +662,7 @@ fn total_record(
|
||||
operation: "generate".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "gpt-5.5".to_owned(),
|
||||
requested_model_id: Some("gpt-5.5".to_owned()),
|
||||
provider_kind: provider_kind.map(str::to_owned),
|
||||
provider_account_ref: None,
|
||||
provider_account_name: None,
|
||||
|
||||
@@ -25,9 +25,13 @@ use gateway_admin::model::{
|
||||
},
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
AuthorizationStarted, CompleteAuthorization, CredentialDeletion, CredentialDeletionResult,
|
||||
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
|
||||
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow, RotateCredential,
|
||||
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
|
||||
CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation,
|
||||
CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels,
|
||||
ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult,
|
||||
ProviderResetCredits, ProviderUsageStatistics, ProviderUsageStatisticsDay,
|
||||
ProviderUsageStatisticsModel, ProviderUsageStatisticsRequest,
|
||||
ProviderUsageStatisticsSummary, ProviderUsageStatisticsTokens, RotateCredential,
|
||||
StartAuthorization,
|
||||
},
|
||||
};
|
||||
@@ -37,6 +41,7 @@ use gateway_core::{
|
||||
};
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
use uuid::{Uuid, Version};
|
||||
|
||||
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
|
||||
use super::{
|
||||
@@ -56,6 +61,8 @@ const MAX_ID_TOKEN_BYTES: usize = 16 * 1024;
|
||||
const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024;
|
||||
const MAX_ACCOUNT_DELETE_BATCH: usize = 200;
|
||||
const MAX_ACCOUNT_GROUP_BATCH: usize = 1000;
|
||||
const MAX_USAGE_CYCLE_OFFSET: i8 = 8;
|
||||
const MAX_UTC_OFFSET_MINUTES: i16 = 14 * 60;
|
||||
|
||||
/// 账号列表查询参数。
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
@@ -305,8 +312,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,
|
||||
@@ -414,6 +425,35 @@ impl AccountIdQuery {
|
||||
}
|
||||
}
|
||||
|
||||
/// 官方每日用量统计查询。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct AccountUsageStatisticsQuery {
|
||||
pub account_id: String,
|
||||
pub cycle_offset: Option<i8>,
|
||||
pub utc_offset_minutes: Option<i16>,
|
||||
}
|
||||
|
||||
impl AccountUsageStatisticsQuery {
|
||||
pub fn into_request(self) -> Result<ProviderUsageStatisticsRequest, WireValidationError> {
|
||||
require_account_id(&self.account_id, "accountId")?;
|
||||
let cycle_offset = self.cycle_offset.unwrap_or_default();
|
||||
if !(-MAX_USAGE_CYCLE_OFFSET..=0).contains(&cycle_offset) {
|
||||
return Err(WireValidationError::new("cycleOffset"));
|
||||
}
|
||||
let utc_offset_minutes = self.utc_offset_minutes.unwrap_or_default();
|
||||
if !(-MAX_UTC_OFFSET_MINUTES..=MAX_UTC_OFFSET_MINUTES).contains(&utc_offset_minutes) {
|
||||
return Err(WireValidationError::new("utcOffsetMinutes"));
|
||||
}
|
||||
Ok(ProviderUsageStatisticsRequest {
|
||||
account_id: ProviderAccountId::new(self.account_id)
|
||||
.map_err(|_| WireValidationError::new("accountId"))?,
|
||||
cycle_offset,
|
||||
utc_offset_minutes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 敏感导出的固定 query;IDs 使用逗号分隔,禁止隐式导出全部账号。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
@@ -472,6 +512,48 @@ impl AccountActionRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// 主动额度重置卡消费请求。幂等键由 UI 生成并在不确定重试时复用,与官方一致。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct AccountResetCreditConsumeRequest {
|
||||
pub account_id: String,
|
||||
#[serde(default)]
|
||||
pub credit_id: Option<String>,
|
||||
pub redeem_request_id: String,
|
||||
}
|
||||
|
||||
impl AccountResetCreditConsumeRequest {
|
||||
pub fn validate(&self) -> Result<(), WireValidationError> {
|
||||
require_account_id(&self.account_id, "accountId")?;
|
||||
if self.credit_id.as_deref().is_some_and(|value| {
|
||||
let value = value.trim();
|
||||
value.is_empty() || value.len() > 512 || value.chars().any(char::is_control)
|
||||
}) {
|
||||
return Err(WireValidationError::new("creditId"));
|
||||
}
|
||||
let value = self.redeem_request_id.trim();
|
||||
let redeem_request_id =
|
||||
Uuid::parse_str(value).map_err(|_| WireValidationError::new("redeemRequestId"))?;
|
||||
if redeem_request_id.get_version() != Some(Version::Random)
|
||||
|| redeem_request_id.hyphenated().to_string() != value
|
||||
{
|
||||
return Err(WireValidationError::new("redeemRequestId"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn into_command(self) -> Result<ConsumeProviderResetCredit, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(ConsumeProviderResetCredit {
|
||||
account_id: ProviderAccountId::new(self.account_id)
|
||||
.map_err(|_| WireValidationError::new("accountId"))?,
|
||||
credit_id: self.credit_id.map(|value| value.trim().to_owned()),
|
||||
redeem_request_id: Uuid::parse_str(&self.redeem_request_id)
|
||||
.map_err(|_| WireValidationError::new("redeemRequestId"))?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 手工 OAuth 刷新会变更持久 credential。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
@@ -539,6 +621,246 @@ pub struct AccountQuotaData {
|
||||
pub account: AccountView,
|
||||
}
|
||||
|
||||
/// Provider 官方每日用量统计响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsData {
|
||||
pub mode: String,
|
||||
pub cycle: AccountUsageStatisticsCycleView,
|
||||
pub summary: AccountUsageStatisticsSummaryView,
|
||||
pub models: Vec<AccountUsageStatisticsModelView>,
|
||||
pub daily: Vec<AccountUsageStatisticsDayView>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsCycleView {
|
||||
pub offset: i8,
|
||||
pub start_at: String,
|
||||
pub end_at: String,
|
||||
pub window_seconds: u64,
|
||||
pub used_percent: Option<f64>,
|
||||
pub is_current: bool,
|
||||
pub can_go_previous: bool,
|
||||
pub can_go_next: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsTokensView {
|
||||
pub uncached_input: u64,
|
||||
pub cached_input: u64,
|
||||
pub output: u64,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsCostView {
|
||||
pub currency: String,
|
||||
pub amount: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsSummaryView {
|
||||
pub tokens: AccountUsageStatisticsTokensView,
|
||||
pub estimated_cost: Option<AccountUsageStatisticsCostView>,
|
||||
pub projected_tokens: Option<u64>,
|
||||
pub projected_cost: Option<AccountUsageStatisticsCostView>,
|
||||
pub day_count: u32,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsModelView {
|
||||
pub key: String,
|
||||
pub model: String,
|
||||
pub service_tier: String,
|
||||
pub credit_share: Option<f64>,
|
||||
pub quota_share: Option<f64>,
|
||||
pub tokens: AccountUsageStatisticsTokensView,
|
||||
pub estimated_cost: Option<AccountUsageStatisticsCostView>,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_estimated_allocation: bool,
|
||||
pub has_rate_fallback: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageStatisticsDayView {
|
||||
pub date: String,
|
||||
pub credit_share: Option<f64>,
|
||||
pub tokens: AccountUsageStatisticsTokensView,
|
||||
pub estimated_cost: Option<AccountUsageStatisticsCostView>,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
pub is_boundary_day: bool,
|
||||
}
|
||||
|
||||
impl From<ProviderUsageStatistics> for AccountUsageStatisticsData {
|
||||
fn from(statistics: ProviderUsageStatistics) -> Self {
|
||||
Self {
|
||||
mode: statistics.mode.as_str().to_owned(),
|
||||
cycle: AccountUsageStatisticsCycleView {
|
||||
offset: statistics.cycle.offset,
|
||||
start_at: statistics.cycle.start_at.to_rfc3339(),
|
||||
end_at: statistics.cycle.end_at.to_rfc3339(),
|
||||
window_seconds: statistics.cycle.window_seconds,
|
||||
used_percent: statistics.cycle.used_percent,
|
||||
is_current: statistics.cycle.is_current,
|
||||
can_go_previous: statistics.cycle.can_go_previous,
|
||||
can_go_next: statistics.cycle.can_go_next,
|
||||
},
|
||||
summary: account_usage_statistics_summary_view(statistics.summary),
|
||||
models: statistics
|
||||
.models
|
||||
.into_iter()
|
||||
.map(account_usage_statistics_model_view)
|
||||
.collect(),
|
||||
daily: statistics
|
||||
.daily
|
||||
.into_iter()
|
||||
.map(account_usage_statistics_day_view)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn account_usage_statistics_tokens_view(
|
||||
tokens: ProviderUsageStatisticsTokens,
|
||||
) -> AccountUsageStatisticsTokensView {
|
||||
AccountUsageStatisticsTokensView {
|
||||
uncached_input: tokens.uncached_input,
|
||||
cached_input: tokens.cached_input,
|
||||
output: tokens.output,
|
||||
total: tokens.total,
|
||||
}
|
||||
}
|
||||
|
||||
fn account_usage_statistics_cost_view(
|
||||
cost: gateway_admin::model::observability::CurrencyCost,
|
||||
) -> AccountUsageStatisticsCostView {
|
||||
AccountUsageStatisticsCostView {
|
||||
currency: cost.currency,
|
||||
amount: cost.amount.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn account_usage_statistics_summary_view(
|
||||
summary: ProviderUsageStatisticsSummary,
|
||||
) -> AccountUsageStatisticsSummaryView {
|
||||
AccountUsageStatisticsSummaryView {
|
||||
tokens: account_usage_statistics_tokens_view(summary.tokens),
|
||||
estimated_cost: summary
|
||||
.estimated_cost
|
||||
.map(account_usage_statistics_cost_view),
|
||||
projected_tokens: summary.projected_tokens,
|
||||
projected_cost: summary
|
||||
.projected_cost
|
||||
.map(account_usage_statistics_cost_view),
|
||||
day_count: summary.day_count,
|
||||
has_unknown_pricing: summary.has_unknown_pricing,
|
||||
has_missing_token_data: summary.has_missing_token_data,
|
||||
}
|
||||
}
|
||||
|
||||
fn account_usage_statistics_model_view(
|
||||
model: ProviderUsageStatisticsModel,
|
||||
) -> AccountUsageStatisticsModelView {
|
||||
AccountUsageStatisticsModelView {
|
||||
key: model.key,
|
||||
model: model.model,
|
||||
service_tier: model.service_tier.as_str().to_owned(),
|
||||
credit_share: model.credit_share,
|
||||
quota_share: model.quota_share,
|
||||
tokens: account_usage_statistics_tokens_view(model.tokens),
|
||||
estimated_cost: model.estimated_cost.map(account_usage_statistics_cost_view),
|
||||
has_unknown_pricing: model.has_unknown_pricing,
|
||||
has_estimated_allocation: model.has_estimated_allocation,
|
||||
has_rate_fallback: model.has_rate_fallback,
|
||||
has_missing_token_data: model.has_missing_token_data,
|
||||
}
|
||||
}
|
||||
|
||||
fn account_usage_statistics_day_view(
|
||||
day: ProviderUsageStatisticsDay,
|
||||
) -> AccountUsageStatisticsDayView {
|
||||
AccountUsageStatisticsDayView {
|
||||
date: day.date.to_string(),
|
||||
credit_share: day.credit_share,
|
||||
tokens: account_usage_statistics_tokens_view(day.tokens),
|
||||
estimated_cost: day.estimated_cost.map(account_usage_statistics_cost_view),
|
||||
has_unknown_pricing: day.has_unknown_pricing,
|
||||
has_missing_token_data: day.has_missing_token_data,
|
||||
is_boundary_day: day.is_boundary_day,
|
||||
}
|
||||
}
|
||||
|
||||
/// 主动额度重置卡列表响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountResetCreditsData {
|
||||
pub available_count: u64,
|
||||
pub credits: Vec<AccountResetCreditView>,
|
||||
}
|
||||
|
||||
/// 一张安全主动额度重置卡视图。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountResetCreditView {
|
||||
pub id: String,
|
||||
pub status: Option<String>,
|
||||
pub title: Option<String>,
|
||||
pub expires_at: Option<String>,
|
||||
pub reset_type: Option<String>,
|
||||
}
|
||||
|
||||
/// 主动额度重置卡消费响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountResetCreditResultData {
|
||||
pub code: String,
|
||||
pub credit: Option<AccountResetCreditView>,
|
||||
}
|
||||
|
||||
impl From<ProviderResetCredit> for AccountResetCreditView {
|
||||
fn from(credit: ProviderResetCredit) -> Self {
|
||||
Self {
|
||||
id: credit.id,
|
||||
status: credit.status,
|
||||
title: credit.title,
|
||||
expires_at: credit.expires_at.map(|value| value.to_rfc3339()),
|
||||
reset_type: credit.reset_type,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ProviderResetCredits> for AccountResetCreditsData {
|
||||
fn from(credits: ProviderResetCredits) -> Self {
|
||||
Self {
|
||||
available_count: credits.available_count,
|
||||
credits: credits
|
||||
.credits
|
||||
.into_iter()
|
||||
.map(AccountResetCreditView::from)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ProviderResetCreditResult> for AccountResetCreditResultData {
|
||||
fn from(result: ProviderResetCreditResult) -> Self {
|
||||
Self {
|
||||
code: result.code,
|
||||
credit: result.credit.map(AccountResetCreditView::from),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider-owned 明文导出文档;Debug 永远不输出内部 JSON。
|
||||
#[derive(Serialize)]
|
||||
#[serde(transparent)]
|
||||
@@ -1051,6 +1373,14 @@ where
|
||||
post(batch_update_accounts::<S>),
|
||||
)
|
||||
.route("/api/admin/accounts/quota", get(account_quota::<S>))
|
||||
.route(
|
||||
"/api/admin/accounts/usage-statistics",
|
||||
get(account_usage_statistics::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/reset-credits",
|
||||
get(account_reset_credits::<S>).post(consume_account_reset_credit::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/quota/refresh",
|
||||
post(refresh_account_quota::<S>),
|
||||
@@ -1381,6 +1711,27 @@ where
|
||||
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
|
||||
}
|
||||
|
||||
async fn account_usage_statistics<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Query(query): Query<AccountUsageStatisticsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let request = query.into_request().map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.accounts()
|
||||
.usage_statistics(request)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountUsageStatisticsData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn refresh_account_quota<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
@@ -1402,6 +1753,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 +1996,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 +2007,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
|
||||
|
||||
@@ -142,11 +142,19 @@ pub(crate) fn format_service_tier(service_tier: Option<&str>) -> String {
|
||||
match service_tier {
|
||||
Some("priority" | "fast") => "Fast".to_owned(),
|
||||
Some("flex") => "Flex".to_owned(),
|
||||
Some("default") | None => "Default".to_owned(),
|
||||
Some(other) => other.to_owned(),
|
||||
Some("default" | "standard") | None => "Standard".to_owned(),
|
||||
Some(other) => capitalize_first(other),
|
||||
}
|
||||
}
|
||||
|
||||
fn capitalize_first(value: &str) -> String {
|
||||
let mut chars = value.chars();
|
||||
chars
|
||||
.next()
|
||||
.map(|first| first.to_uppercase().chain(chars).collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(crate) fn billing_view(billing: Option<&domain::UsageBilling>) -> Option<BillingView> {
|
||||
match billing? {
|
||||
domain::UsageBilling::Total { total, .. } => Some(BillingView {
|
||||
@@ -222,7 +230,7 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
|
||||
let model = record
|
||||
.upstream_model_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| record.requested_model_id.clone());
|
||||
.or_else(|| record.requested_model_id.clone());
|
||||
let transport = record
|
||||
.upstream_transport
|
||||
.clone()
|
||||
@@ -243,7 +251,7 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
|
||||
account_name: record.provider_account_name,
|
||||
route: record.endpoint,
|
||||
model,
|
||||
requested_model: Some(record.requested_model_id),
|
||||
requested_model: record.requested_model_id,
|
||||
upstream_model: record.upstream_model_id,
|
||||
service_tier: record.service_tier,
|
||||
status_code,
|
||||
@@ -351,9 +359,7 @@ pub(crate) fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptV
|
||||
provider: attempt
|
||||
.provider_kind
|
||||
.unwrap_or_else(|| "unknown".to_owned()),
|
||||
model: attempt
|
||||
.upstream_model_id
|
||||
.unwrap_or_else(|| "unknown".to_owned()),
|
||||
model: attempt.upstream_model_id,
|
||||
transport: attempt
|
||||
.upstream_transport
|
||||
.unwrap_or_else(|| "unknown".to_owned()),
|
||||
@@ -731,21 +737,19 @@ pub(crate) fn dashboard_view(
|
||||
last_used: relative_time(credential.last_used_at, range.end),
|
||||
});
|
||||
}
|
||||
let unavailable_accounts = provider_accounts
|
||||
.total
|
||||
.saturating_sub(provider_accounts.normal);
|
||||
|
||||
DashboardDataView {
|
||||
cards: DashboardCardsView {
|
||||
credentials: DashboardCredentialsCardView {
|
||||
total: format_compact_number(provider_accounts.total),
|
||||
total: provider_accounts.total.to_string(),
|
||||
total_value: provider_accounts.total,
|
||||
available: format_compact_number(provider_accounts.normal),
|
||||
available: provider_accounts.normal.to_string(),
|
||||
available_value: provider_accounts.normal,
|
||||
unavailable: format_compact_number(
|
||||
provider_accounts
|
||||
.total
|
||||
.saturating_sub(provider_accounts.normal),
|
||||
),
|
||||
unavailable_value: provider_accounts
|
||||
.total
|
||||
.saturating_sub(provider_accounts.normal),
|
||||
unavailable: unavailable_accounts.to_string(),
|
||||
unavailable_value: unavailable_accounts,
|
||||
},
|
||||
traffic: DashboardTrafficCardView {
|
||||
today_requests: format_compact_number(today.request_count),
|
||||
|
||||
@@ -94,7 +94,7 @@ pub struct UsageRecordView {
|
||||
pub account_email: Option<String>,
|
||||
pub account_name: Option<String>,
|
||||
pub route: String,
|
||||
pub model: String,
|
||||
pub model: Option<String>,
|
||||
pub requested_model: Option<String>,
|
||||
pub upstream_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
@@ -183,7 +183,7 @@ pub struct UsageAttemptView {
|
||||
pub attempt_index: u32,
|
||||
pub trigger: String,
|
||||
pub provider: String,
|
||||
pub model: String,
|
||||
pub model: Option<String>,
|
||||
pub transport: String,
|
||||
pub send_state: String,
|
||||
pub outcome: String,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Admin API 的展示层辅助:数值格式化。
|
||||
//!
|
||||
//! 按 docs/project-redundancy-boundary-audit.md BE-02 的目标,display 字段最终
|
||||
//! 迁移到前端;过渡期统一由本模块提供,避免 accounts/observability 各自维护。
|
||||
//! 当前 wire 仍保留 display 字段;统一由本模块提供,避免 accounts/observability
|
||||
//! 各自维护格式化规则。
|
||||
|
||||
use gateway_core::accounting::Decimal;
|
||||
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -70,6 +70,148 @@ mod query {
|
||||
}
|
||||
}
|
||||
|
||||
mod usage_statistics {
|
||||
use std::str::FromStr as _;
|
||||
|
||||
use chrono::{NaiveDate, TimeZone as _, Utc};
|
||||
use gateway_admin::model::{
|
||||
observability::{CurrencyCost, DecimalAmount},
|
||||
provider_credentials::{
|
||||
ProviderUsageStatistics, ProviderUsageStatisticsCycle, ProviderUsageStatisticsDay,
|
||||
ProviderUsageStatisticsMode, ProviderUsageStatisticsModel,
|
||||
ProviderUsageStatisticsServiceTier, ProviderUsageStatisticsSummary,
|
||||
ProviderUsageStatisticsTokens,
|
||||
},
|
||||
};
|
||||
use gateway_api::admin::accounts::{AccountUsageStatisticsData, AccountUsageStatisticsQuery};
|
||||
use serde_json::json;
|
||||
|
||||
fn tokens(total: u64) -> ProviderUsageStatisticsTokens {
|
||||
ProviderUsageStatisticsTokens {
|
||||
uncached_input: total / 4,
|
||||
cached_input: total / 2,
|
||||
output: total / 4,
|
||||
total,
|
||||
}
|
||||
}
|
||||
|
||||
fn usd(amount: &str) -> CurrencyCost {
|
||||
CurrencyCost {
|
||||
currency: "USD".to_owned(),
|
||||
amount: DecimalAmount::from_str(amount).expect("USD amount"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_statistics_query_accepts_only_current_or_bounded_history_offsets() {
|
||||
let query: AccountUsageStatisticsQuery = serde_json::from_value(json!({
|
||||
"accountId": "acct_openai",
|
||||
"cycleOffset": -3,
|
||||
"utcOffsetMinutes": 480
|
||||
}))
|
||||
.expect("decode usage-statistics query");
|
||||
let request = query.into_request().expect("valid cycle offset");
|
||||
assert_eq!(request.cycle_offset, -3);
|
||||
assert_eq!(request.utc_offset_minutes, 480);
|
||||
|
||||
for (cycle_offset, field) in [(-9, "cycleOffset"), (1, "cycleOffset")] {
|
||||
let query: AccountUsageStatisticsQuery = serde_json::from_value(json!({
|
||||
"accountId": "acct_openai",
|
||||
"cycleOffset": cycle_offset
|
||||
}))
|
||||
.expect("decode bounded integer");
|
||||
assert_eq!(
|
||||
query
|
||||
.into_request()
|
||||
.expect_err("reject cycle offset")
|
||||
.field(),
|
||||
field
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
serde_json::from_value::<AccountUsageStatisticsQuery>(json!({
|
||||
"accountId": "acct_openai",
|
||||
"historyCycles": 8
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_statistics_response_serializes_calculated_rows_without_raw_documents() {
|
||||
let start_at = Utc
|
||||
.with_ymd_and_hms(2026, 8, 24, 7, 50, 0)
|
||||
.single()
|
||||
.expect("cycle start");
|
||||
let end_at = Utc
|
||||
.with_ymd_and_hms(2026, 8, 31, 7, 50, 0)
|
||||
.single()
|
||||
.expect("cycle end");
|
||||
let response = AccountUsageStatisticsData::from(ProviderUsageStatistics {
|
||||
mode: ProviderUsageStatisticsMode::Personal,
|
||||
cycle: ProviderUsageStatisticsCycle {
|
||||
offset: 0,
|
||||
start_at,
|
||||
end_at,
|
||||
window_seconds: 604_800,
|
||||
used_percent: Some(20.0),
|
||||
is_current: true,
|
||||
can_go_previous: true,
|
||||
can_go_next: false,
|
||||
},
|
||||
summary: ProviderUsageStatisticsSummary {
|
||||
tokens: tokens(100),
|
||||
estimated_cost: Some(usd("1.25")),
|
||||
projected_tokens: Some(500),
|
||||
projected_cost: Some(usd("6.25")),
|
||||
day_count: 1,
|
||||
has_unknown_pricing: false,
|
||||
has_missing_token_data: false,
|
||||
},
|
||||
models: vec![ProviderUsageStatisticsModel {
|
||||
key: "gpt-5.6-sol::standard".to_owned(),
|
||||
model: "gpt-5.6-sol".to_owned(),
|
||||
service_tier: ProviderUsageStatisticsServiceTier::Standard,
|
||||
credit_share: Some(1.0),
|
||||
quota_share: Some(0.2),
|
||||
tokens: tokens(100),
|
||||
estimated_cost: Some(usd("1.25")),
|
||||
has_unknown_pricing: false,
|
||||
has_estimated_allocation: true,
|
||||
has_rate_fallback: false,
|
||||
has_missing_token_data: false,
|
||||
}],
|
||||
daily: vec![ProviderUsageStatisticsDay {
|
||||
date: NaiveDate::from_ymd_opt(2026, 8, 25).expect("report date"),
|
||||
credit_share: Some(1.0),
|
||||
tokens: tokens(100),
|
||||
estimated_cost: Some(usd("1.25")),
|
||||
has_unknown_pricing: false,
|
||||
has_missing_token_data: false,
|
||||
is_boundary_day: true,
|
||||
}],
|
||||
});
|
||||
let value = serde_json::to_value(response).expect("serialize usage statistics");
|
||||
|
||||
assert_eq!(value["cycle"]["offset"], 0);
|
||||
assert_eq!(
|
||||
value["summary"]["estimatedCost"],
|
||||
json!({
|
||||
"currency": "USD",
|
||||
"amount": "1.25"
|
||||
})
|
||||
);
|
||||
assert_eq!(value["models"][0]["serviceTier"], "standard");
|
||||
assert!(value["models"][0].get("mode").is_none());
|
||||
assert_eq!(value["daily"][0]["tokens"]["total"], 100);
|
||||
assert_eq!(value["daily"][0]["isBoundaryDay"], true);
|
||||
assert!(value["cycle"].get("deltaPercent").is_none());
|
||||
assert!(value.get("modelBreakdown").is_none());
|
||||
assert!(value.get("dailyTotals").is_none());
|
||||
assert!(value.get("usage").is_none());
|
||||
}
|
||||
}
|
||||
|
||||
mod batch_update {
|
||||
use gateway_api::admin::accounts::{BatchUpdateAccountsRequest, UpdateAccountRequest};
|
||||
use serde_json::json;
|
||||
@@ -266,8 +408,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 +467,44 @@ mod actions {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_credit_consume_should_require_canonical_v4_idempotency_key() {
|
||||
let valid: AccountResetCreditConsumeRequest = serde_json::from_value(json!({
|
||||
"accountId": "acct_1",
|
||||
"creditId": "credit_1",
|
||||
"redeemRequestId": "8fbf302d-11df-4bd5-82e4-08e4b3df7874"
|
||||
}))
|
||||
.expect("decode reset-credit consume");
|
||||
valid.validate().expect("validate reset-credit consume");
|
||||
|
||||
for (redeem_request_id, field) in [
|
||||
("019c0000-0000-7000-8000-000000000000", "redeemRequestId"),
|
||||
("8FBF302D-11DF-4BD5-82E4-08E4B3DF7874", "redeemRequestId"),
|
||||
("invalid", "redeemRequestId"),
|
||||
] {
|
||||
let request: AccountResetCreditConsumeRequest = serde_json::from_value(json!({
|
||||
"accountId": "acct_1",
|
||||
"redeemRequestId": redeem_request_id
|
||||
}))
|
||||
.expect("decode invalid reset-credit consume");
|
||||
assert_eq!(request.validate().expect_err("reject UUID").field(), field);
|
||||
}
|
||||
|
||||
let invalid_credit: AccountResetCreditConsumeRequest = serde_json::from_value(json!({
|
||||
"accountId": "acct_1",
|
||||
"creditId": " ",
|
||||
"redeemRequestId": "8fbf302d-11df-4bd5-82e4-08e4b3df7874"
|
||||
}))
|
||||
.expect("decode invalid credit");
|
||||
assert_eq!(
|
||||
invalid_credit
|
||||
.validate()
|
||||
.expect_err("reject credit")
|
||||
.field(),
|
||||
"creditId"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn credential_recovery_requests_should_not_accept_client_revision_fences() {
|
||||
let authorization: StartAccountAuthorizationRequest = serde_json::from_value(json!({
|
||||
|
||||
@@ -221,9 +221,9 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
|
||||
attempts: AttemptMetrics::default(),
|
||||
totals: Default::default(),
|
||||
provider_accounts: AccountPoolMetrics {
|
||||
total: 807,
|
||||
normal: 15,
|
||||
quota_exhausted: 785,
|
||||
total: 1002,
|
||||
normal: 943,
|
||||
quota_exhausted: 52,
|
||||
rate_limited: 0,
|
||||
disabled: 3,
|
||||
error: 4,
|
||||
@@ -259,9 +259,21 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
|
||||
&value["data"]["cards"]["credentials"]["unavailableValue"],
|
||||
),
|
||||
(
|
||||
&serde_json::json!(807),
|
||||
&serde_json::json!(15),
|
||||
&serde_json::json!(792)
|
||||
&serde_json::json!(1002),
|
||||
&serde_json::json!(943),
|
||||
&serde_json::json!(59)
|
||||
),
|
||||
);
|
||||
assert_eq!(
|
||||
(
|
||||
&value["data"]["cards"]["credentials"]["total"],
|
||||
&value["data"]["cards"]["credentials"]["available"],
|
||||
&value["data"]["cards"]["credentials"]["unavailable"],
|
||||
),
|
||||
(
|
||||
&serde_json::json!("1002"),
|
||||
&serde_json::json!("943"),
|
||||
&serde_json::json!("59")
|
||||
),
|
||||
);
|
||||
}
|
||||
@@ -681,7 +693,7 @@ fn usage_record_with_account(
|
||||
operation: "responses".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
requested_model_id: Some("coding".to_owned()),
|
||||
provider_kind: Some("openai".to_owned()),
|
||||
provider_account_ref: Some(account_id.to_owned()),
|
||||
provider_account_name: Some(name.to_owned()),
|
||||
@@ -810,7 +822,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
operation: "generate".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "grok-4.5".to_owned(),
|
||||
requested_model_id: Some("grok-4.5".to_owned()),
|
||||
provider_kind: Some("xai".to_owned()),
|
||||
provider_account_ref: Some("acct_snapshot".to_owned()),
|
||||
provider_account_name: Some("Snapshot Alpha".to_owned()),
|
||||
@@ -820,7 +832,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
upstream_transport: Some("http_sse".to_owned()),
|
||||
http_version: Some("h2".to_owned()),
|
||||
websocket_pool: Some("reuse".to_owned()),
|
||||
service_tier: Some("priority".to_owned()),
|
||||
service_tier: Some("default".to_owned()),
|
||||
provider_metadata_json: Some(
|
||||
serde_json::json!({
|
||||
"effectiveModel": "grok-4.5",
|
||||
@@ -865,7 +877,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
output_price_per_million: usd("60.0000"),
|
||||
cache_read_price_per_million: usd("1.0000"),
|
||||
cache_write_price_per_million: usd("12.5000"),
|
||||
service_tier: Some("priority".to_owned()),
|
||||
service_tier: Some("default".to_owned()),
|
||||
multiplier_percent: 100,
|
||||
},
|
||||
))),
|
||||
@@ -921,6 +933,10 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
value["data"]["items"][0]["billing"]["cacheWritePriceDisplay"],
|
||||
"$12.5 / 1M Token"
|
||||
);
|
||||
assert_eq!(
|
||||
value["data"]["items"][0]["billing"]["serviceTierDisplay"],
|
||||
"Standard"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
serde_json::json!({
|
||||
@@ -961,7 +977,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
"routingScope": "groups",
|
||||
"routingGroupRefs": ["grp_usage"],
|
||||
"routingGroupNamesSnapshot": ["Usage group"],
|
||||
"serviceTier": "priority",
|
||||
"serviceTier": "default",
|
||||
"accountId": "acct_snapshot",
|
||||
"accountName": "Snapshot Alpha",
|
||||
"accountEmail": "alpha@example.invalid",
|
||||
|
||||
@@ -52,6 +52,8 @@ fn source_tree_should_match_frozen_machine_manifest() {
|
||||
"src/lib.rs",
|
||||
"src/openai/auth.rs",
|
||||
"src/openai/error.rs",
|
||||
"src/openai/images/http.rs",
|
||||
"src/openai/images/mod.rs",
|
||||
"src/openai/mod.rs",
|
||||
"src/openai/models.rs",
|
||||
"src/openai/responses/error.rs",
|
||||
@@ -90,6 +92,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/main.rs",
|
||||
"tests/openai/auth.rs",
|
||||
"tests/openai/error.rs",
|
||||
"tests/openai/images.rs",
|
||||
"tests/openai/mod.rs",
|
||||
"tests/openai/models.rs",
|
||||
"tests/openai/responses/http.rs",
|
||||
|
||||
@@ -8,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 }"#
|
||||
);
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -946,7 +946,28 @@ impl fmt::Debug for LoadedCredential {
|
||||
}
|
||||
}
|
||||
|
||||
/// 刷新成功后的完整 CAS 写回。
|
||||
/// 与 credential revision CAS 同事务提交的账号错误事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialStateWrite {
|
||||
pub credential_state: CredentialState,
|
||||
pub observed_at: SystemTime,
|
||||
pub error_reason: Option<AccountErrorReason>,
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
/// [`CredentialCasUpdate`] 跨 store 边界的命名字段。
|
||||
pub struct CredentialCasUpdateParts {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub profile: ProviderAccountUpdate,
|
||||
pub credential: PlaintextCredential,
|
||||
pub has_refresh_token: bool,
|
||||
pub access_token_expires_at: Option<SystemTime>,
|
||||
pub next_refresh_at: Option<SystemTime>,
|
||||
pub account_state: Option<Box<CredentialStateWrite>>,
|
||||
}
|
||||
|
||||
/// 刷新后的完整 CAS 写回。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct CredentialCasUpdate {
|
||||
account_id: ProviderAccountId,
|
||||
@@ -956,6 +977,7 @@ pub struct CredentialCasUpdate {
|
||||
has_refresh_token: bool,
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
next_refresh_at: Option<SystemTime>,
|
||||
account_state: Option<Box<CredentialStateWrite>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for CredentialCasUpdate {
|
||||
@@ -969,6 +991,7 @@ impl fmt::Debug for CredentialCasUpdate {
|
||||
.field("has_refresh_token", &self.has_refresh_token)
|
||||
.field("access_token_expires_at", &self.access_token_expires_at)
|
||||
.field("next_refresh_at", &self.next_refresh_at)
|
||||
.field("account_state", &self.account_state)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -1002,9 +1025,34 @@ impl CredentialCasUpdate {
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
account_state: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// 将刷新调度与账号错误事实放入同一个 revision CAS。
|
||||
#[must_use]
|
||||
pub fn with_account_state(
|
||||
mut self,
|
||||
credential_state: CredentialState,
|
||||
observed_at: SystemTime,
|
||||
error_reason: Option<AccountErrorReason>,
|
||||
message: Option<String>,
|
||||
) -> Self {
|
||||
let message = message.filter(|value| !value.trim().is_empty());
|
||||
let error_reason = if credential_state == CredentialState::Ready {
|
||||
message.as_ref().and(error_reason)
|
||||
} else {
|
||||
error_reason.or_else(|| credential_state.error_reason())
|
||||
};
|
||||
self.account_state = Some(Box::new(CredentialStateWrite {
|
||||
credential_state,
|
||||
observed_at,
|
||||
error_reason,
|
||||
message,
|
||||
}));
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_id(&self) -> &ProviderAccountId {
|
||||
&self.account_id
|
||||
@@ -1041,26 +1089,17 @@ impl CredentialCasUpdate {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn into_parts(
|
||||
self,
|
||||
) -> (
|
||||
ProviderAccountId,
|
||||
CredentialRevision,
|
||||
ProviderAccountUpdate,
|
||||
PlaintextCredential,
|
||||
bool,
|
||||
Option<SystemTime>,
|
||||
Option<SystemTime>,
|
||||
) {
|
||||
(
|
||||
self.account_id,
|
||||
self.expected_revision,
|
||||
self.profile,
|
||||
self.credential,
|
||||
self.has_refresh_token,
|
||||
self.access_token_expires_at,
|
||||
self.next_refresh_at,
|
||||
)
|
||||
pub fn into_parts(self) -> CredentialCasUpdateParts {
|
||||
CredentialCasUpdateParts {
|
||||
account_id: self.account_id,
|
||||
expected_revision: self.expected_revision,
|
||||
profile: self.profile,
|
||||
credential: self.credential,
|
||||
has_refresh_token: self.has_refresh_token,
|
||||
access_token_expires_at: self.access_token_expires_at,
|
||||
next_refresh_at: self.next_refresh_at,
|
||||
account_state: self.account_state,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1102,6 +1141,14 @@ pub enum QuotaWriteOutcome {
|
||||
Conflict,
|
||||
}
|
||||
|
||||
/// 只推进 Provider quota 文档的最后成功查询时间。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct QuotaObservationTouch {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub observed_at: SystemTime,
|
||||
}
|
||||
|
||||
/// 不改 Provider 原始 quota JSON 的额度访问事实写入。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct QuotaAccessChange {
|
||||
@@ -1117,9 +1164,13 @@ pub struct AccountStateChange {
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub credential_state: CredentialState,
|
||||
pub observed_at: SystemTime,
|
||||
/// 受控错误原因;恢复 Ready 时必须清空。
|
||||
/// 受控错误原因;无失败事实的 Ready 写入必须清空。
|
||||
///
|
||||
/// Ready 凭据可以在 AT 过期后继续 RT 恢复,此时允许保留最近一次
|
||||
/// 刷新失败原因,供错误状态投影展示。
|
||||
pub error_reason: Option<AccountErrorReason>,
|
||||
/// 已脱敏的上游诊断信息;恢复 Ready 时必须清空。
|
||||
/// 供管理端展示的错误消息;结构化上游失败应保留原始 message,不能写入整个正文。
|
||||
/// 刷新成功或其他无失败事实的 Ready 写入必须清空。
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
@@ -1170,6 +1221,11 @@ pub trait ProviderAccountStore: Send + Sync {
|
||||
observation: QuotaObservation,
|
||||
) -> Result<QuotaWriteOutcome, StoreError>;
|
||||
|
||||
async fn touch_quota_observation(
|
||||
&self,
|
||||
touch: QuotaObservationTouch,
|
||||
) -> Result<QuotaWriteOutcome, StoreError>;
|
||||
|
||||
async fn apply_quota_access(
|
||||
&self,
|
||||
change: QuotaAccessChange,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,6 +98,12 @@ pub enum RoutingError {
|
||||
/// 客户端提交的模型名称。
|
||||
model: String,
|
||||
},
|
||||
/// 固定 Provider 的原生端点当前不可执行。
|
||||
#[error("provider endpoint `{provider}` is unavailable")]
|
||||
NoCapableProviderEndpoint {
|
||||
/// adapter 已绑定的 Provider。
|
||||
provider: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// 调用方策略不满足约束或拒绝请求。
|
||||
|
||||
@@ -16,7 +16,7 @@ use crate::operation::ProviderSessionState;
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ResponseMeta {
|
||||
response_id: String,
|
||||
model: String,
|
||||
model: Option<String>,
|
||||
finish_reason: Option<FinishReason>,
|
||||
}
|
||||
|
||||
@@ -26,7 +26,17 @@ impl ResponseMeta {
|
||||
pub fn new(response_id: impl Into<String>, model: impl Into<String>) -> Self {
|
||||
Self {
|
||||
response_id: response_id.into(),
|
||||
model: model.into(),
|
||||
model: Some(model.into()),
|
||||
finish_reason: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建不声明模型的 Provider 原生端点响应元数据。
|
||||
#[must_use]
|
||||
pub fn for_provider_endpoint(response_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
response_id: response_id.into(),
|
||||
model: None,
|
||||
finish_reason: None,
|
||||
}
|
||||
}
|
||||
@@ -46,8 +56,8 @@ impl ResponseMeta {
|
||||
|
||||
/// 返回对客户端公开的模型名。
|
||||
#[must_use]
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
self.model.as_deref()
|
||||
}
|
||||
|
||||
/// 返回规范化终止原因。
|
||||
@@ -190,6 +200,7 @@ pub struct ProtocolWireEvent {
|
||||
data: Value,
|
||||
has_json_data: bool,
|
||||
raw_sse_frame: Option<Bytes>,
|
||||
raw_json_body: Option<Bytes>,
|
||||
sse_id: Option<String>,
|
||||
sse_retry: Option<u64>,
|
||||
}
|
||||
@@ -228,6 +239,7 @@ impl ProtocolWireEvent {
|
||||
data,
|
||||
has_json_data: true,
|
||||
raw_sse_frame: None,
|
||||
raw_json_body: None,
|
||||
sse_id,
|
||||
sse_retry,
|
||||
})
|
||||
@@ -274,6 +286,32 @@ impl ProtocolWireEvent {
|
||||
data: Value::Null,
|
||||
has_json_data: false,
|
||||
raw_sse_frame: Some(raw_sse_frame),
|
||||
raw_json_body: None,
|
||||
sse_id: None,
|
||||
sse_retry: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// 创建未经改写的完整 JSON 响应正文。
|
||||
///
|
||||
/// 该载荷用于非流式协议端点;Core 不解析或重编码其中的值。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 协议名不满足内部路由标识约束时返回错误。
|
||||
pub fn raw_json(
|
||||
protocol: impl Into<String>,
|
||||
raw_json_body: Bytes,
|
||||
) -> Result<Self, IdentifierError> {
|
||||
let protocol = protocol.into();
|
||||
validate_text(&protocol, 64, true, None)?;
|
||||
Ok(Self {
|
||||
protocol,
|
||||
event_type: None,
|
||||
data: Value::Null,
|
||||
has_json_data: false,
|
||||
raw_sse_frame: None,
|
||||
raw_json_body: Some(raw_json_body),
|
||||
sse_id: None,
|
||||
sse_retry: None,
|
||||
})
|
||||
@@ -321,6 +359,18 @@ impl ProtocolWireEvent {
|
||||
self.raw_sse_frame.as_ref()
|
||||
}
|
||||
|
||||
/// 返回可直接写给 JSON 客户端的未经改写响应正文。
|
||||
#[must_use]
|
||||
pub const fn raw_json_body(&self) -> Option<&Bytes> {
|
||||
self.raw_json_body.as_ref()
|
||||
}
|
||||
|
||||
/// 拆出未经改写的完整 JSON 响应正文。
|
||||
#[must_use]
|
||||
pub fn into_raw_json_body(self) -> Option<Bytes> {
|
||||
self.raw_json_body
|
||||
}
|
||||
|
||||
/// 拆出协议原生 JSON 数据。
|
||||
#[must_use]
|
||||
pub fn into_data(self) -> Value {
|
||||
@@ -336,6 +386,7 @@ impl fmt::Debug for ProtocolWireEvent {
|
||||
.field("event_type", &self.event_type)
|
||||
.field("has_json_data", &self.has_json_data)
|
||||
.field("has_raw_sse_frame", &self.raw_sse_frame.is_some())
|
||||
.field("has_raw_json_body", &self.raw_json_body.is_some())
|
||||
.field("has_sse_id", &self.sse_id.is_some())
|
||||
.field("sse_retry", &self.sse_retry)
|
||||
.field("data", &"<not included in Debug>")
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
//! 协议无关的业务 operation。
|
||||
//!
|
||||
//! 这里只保留网关需要解释、路由或结算的稳定语义。协议 adapter 无法写入
|
||||
//! wire body 的连接级事实随 [`ProtocolPayload`] 不透明传递,由对应 Provider
|
||||
//! 自己解释。
|
||||
//! wire body 与连接级事实随 [`ProtocolPayload`] 或 [`RawJsonPayload`] 不透明传递,
|
||||
//! 由对应 Provider 自己解释。
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::Bytes;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
@@ -239,6 +240,72 @@ impl fmt::Debug for ProtocolPayload {
|
||||
}
|
||||
}
|
||||
|
||||
/// 客户端协议交给 Provider 的未经重编码 JSON 正文。
|
||||
///
|
||||
/// 协议 adapter 只确定路由端点;Core 保留原始字节与非 wire 上下文,
|
||||
/// 不验证、解析或改写正文。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct RawJsonPayload {
|
||||
protocol: String,
|
||||
body: Bytes,
|
||||
context: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl RawJsonPayload {
|
||||
/// 创建未经重编码的协议 JSON 正文。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 协议名称为空、过长或含控制字符时返回错误。
|
||||
pub fn new(protocol: impl Into<String>, body: Bytes) -> Result<Self, OperationError> {
|
||||
let protocol = protocol.into();
|
||||
validate_text(&protocol, 64, true, None).map_err(|_| OperationError::EmptyField {
|
||||
field: "raw_json_payload protocol",
|
||||
})?;
|
||||
Ok(Self {
|
||||
protocol,
|
||||
body,
|
||||
context: Map::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 返回协议名称。
|
||||
#[must_use]
|
||||
pub fn protocol(&self) -> &str {
|
||||
&self.protocol
|
||||
}
|
||||
|
||||
/// 返回 adapter 收到的原始 JSON 字节。
|
||||
#[must_use]
|
||||
pub const fn body(&self) -> &Bytes {
|
||||
&self.body
|
||||
}
|
||||
|
||||
/// 附着不写入 wire body 的协议连接上下文。
|
||||
#[must_use]
|
||||
pub fn with_context(mut self, context: Map<String, Value>) -> Self {
|
||||
self.context = context;
|
||||
self
|
||||
}
|
||||
|
||||
/// 返回仅供对应 Provider 解释的非 wire 上下文。
|
||||
#[must_use]
|
||||
pub const fn context(&self) -> &Map<String, Value> {
|
||||
&self.context
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for RawJsonPayload {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RawJsonPayload")
|
||||
.field("protocol", &self.protocol)
|
||||
.field("body", &"<not included in Debug>")
|
||||
.field("context", &"<not included in Debug>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// 通用生成请求。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct GenerateRequest {
|
||||
@@ -501,32 +568,39 @@ impl fmt::Debug for RerankRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// 图像生成请求。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
/// 图像 API 的稳定端点语义。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum ImageRequestKind {
|
||||
/// 从文本生成图像。
|
||||
Generation,
|
||||
/// 使用输入图像执行编辑。
|
||||
Edit,
|
||||
}
|
||||
|
||||
/// 图像生成或编辑请求。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct ImageRequest {
|
||||
prompt: String,
|
||||
count: u32,
|
||||
kind: ImageRequestKind,
|
||||
payload: RawJsonPayload,
|
||||
}
|
||||
|
||||
impl ImageRequest {
|
||||
/// 创建图像生成请求。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Prompt 为空时返回错误。
|
||||
pub fn new(prompt: impl Into<String>) -> Result<Self, OperationError> {
|
||||
let prompt = prompt.into();
|
||||
if prompt.is_empty() {
|
||||
return Err(OperationError::EmptyField { field: "prompt" });
|
||||
}
|
||||
Ok(Self { prompt, count: 1 })
|
||||
/// 创建携带原始协议 JSON 正文的图像请求。
|
||||
#[must_use]
|
||||
pub const fn from_raw_json(kind: ImageRequestKind, payload: RawJsonPayload) -> Self {
|
||||
Self { kind, payload }
|
||||
}
|
||||
|
||||
/// 设置图像数量。
|
||||
/// 返回图像 API 端点语义。
|
||||
#[must_use]
|
||||
pub const fn with_count(mut self, count: u32) -> Self {
|
||||
self.count = count;
|
||||
self
|
||||
pub const fn kind(&self) -> ImageRequestKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
/// 返回未经重编码的协议 JSON 正文。
|
||||
#[must_use]
|
||||
pub const fn payload(&self) -> &RawJsonPayload {
|
||||
&self.payload
|
||||
}
|
||||
}
|
||||
|
||||
@@ -534,8 +608,8 @@ impl fmt::Debug for ImageRequest {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ImageRequest")
|
||||
.field("prompt", &"<not included in Debug>")
|
||||
.field("count", &self.count)
|
||||
.field("kind", &self.kind)
|
||||
.field("payload", &"<not included in Debug>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -649,7 +723,8 @@ impl Operation {
|
||||
pub fn image_generation_requested(&self) -> bool {
|
||||
match self {
|
||||
Self::Generate(request) => request.image_generation_requested(),
|
||||
Self::Embed(_) | Self::Rerank(_) | Self::GenerateImage(_) | Self::Speech(_) => false,
|
||||
Self::GenerateImage(_) => true,
|
||||
Self::Embed(_) | Self::Rerank(_) | Self::Speech(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -231,6 +231,32 @@ pub trait ProviderSessionAffinityPort: Send + Sync {
|
||||
ttl: Duration,
|
||||
) -> BoxFuture<'a, Result<(), ProviderStoreError>>;
|
||||
|
||||
/// 仅在亲和键尚未绑定时写入候选账号,并返回原子操作后的实际绑定。
|
||||
///
|
||||
/// 已存在的绑定绝不会被候选账号覆盖;同一根会话的并发首次请求据此收敛到
|
||||
/// 单一账号。TTL 只在首次写入时设置,已有绑定由成功反馈负责刷新。
|
||||
fn claim_or_load<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
key: &'a ProviderSessionAffinityKey,
|
||||
candidate_account_id: &'a ProviderAccountId,
|
||||
ttl: Duration,
|
||||
) -> BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>>;
|
||||
|
||||
/// 仅当当前绑定等于 `expected_account_id`(或键已过期)时写入新账号,
|
||||
/// 并返回原子操作后的实际绑定。
|
||||
///
|
||||
/// Provider 用它迁移不可用账号,以及在成功后以 `expected == replacement`
|
||||
/// 刷新 TTL;迟到的旧账号成功不能覆盖较新的会话 winner。
|
||||
fn compare_and_bind<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
key: &'a ProviderSessionAffinityKey,
|
||||
expected_account_id: &'a ProviderAccountId,
|
||||
replacement_account_id: &'a ProviderAccountId,
|
||||
ttl: Duration,
|
||||
) -> BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>>;
|
||||
|
||||
fn clear<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
|
||||
@@ -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 [
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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,6 +1,6 @@
|
||||
//! 多 Provider 网关的 PostgreSQL 持久化与 Redis 协调 adapter。
|
||||
//!
|
||||
//! 业务规则与 port 由 `gateway-core` 拥有。本 crate 只负责把终态的七张业务表
|
||||
//! 业务规则与 port 由 `gateway-core` / `gateway-admin` 拥有。本 crate 只负责把当前十二张业务表
|
||||
//! 和可丢失 Redis 状态映射为明确的基础设施操作。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -91,7 +91,7 @@ pub struct NewModelRequest {
|
||||
pub operation: String,
|
||||
pub endpoint: String,
|
||||
pub client_transport: String,
|
||||
pub requested_model_id: String,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
@@ -112,7 +112,9 @@ impl NewModelRequest {
|
||||
require_nonempty(ENTITY, "operation", &self.operation)?;
|
||||
require_nonempty(ENTITY, "endpoint", &self.endpoint)?;
|
||||
require_nonempty(ENTITY, "client_transport", &self.client_transport)?;
|
||||
require_nonempty(ENTITY, "requested_model_id", &self.requested_model_id)?;
|
||||
if let Some(requested_model_id) = self.requested_model_id.as_deref() {
|
||||
require_nonempty(ENTITY, "requested_model_id", requested_model_id)?;
|
||||
}
|
||||
validate_routing_snapshot(
|
||||
&self.routing_scope,
|
||||
&self.routing_group_refs,
|
||||
@@ -141,7 +143,7 @@ pub struct ModelRequestAttemptStart {
|
||||
pub provider_kind: String,
|
||||
pub provider_account_id: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub upstream_model_id: String,
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: String,
|
||||
pub http_version: Option<String>,
|
||||
}
|
||||
@@ -151,11 +153,13 @@ impl ModelRequestAttemptStart {
|
||||
for (field, value) in [
|
||||
("model_request_id", self.model_request_id.as_str()),
|
||||
("provider_kind", self.provider_kind.as_str()),
|
||||
("upstream_model_id", self.upstream_model_id.as_str()),
|
||||
("upstream_transport", self.upstream_transport.as_str()),
|
||||
] {
|
||||
require_nonempty(ENTITY, field, value)?;
|
||||
}
|
||||
if let Some(upstream_model_id) = self.upstream_model_id.as_deref() {
|
||||
require_nonempty(ENTITY, "upstream_model_id", upstream_model_id)?;
|
||||
}
|
||||
if self.attempt_count == 0 {
|
||||
return Err(invalid("attempt_count must be positive"));
|
||||
}
|
||||
@@ -841,7 +845,10 @@ impl ExecutionStore for PgExecutionStore {
|
||||
provider_kind: Some(failure.provider_kind.as_str().to_owned()),
|
||||
provider_account_id: failure.account_id.as_ref().map(|id| id.as_str().to_owned()),
|
||||
provider_account_ref: failure.account_id.as_ref().map(|id| id.as_str().to_owned()),
|
||||
upstream_model_id: Some(failure.upstream_model_id.as_str().to_owned()),
|
||||
upstream_model_id: failure
|
||||
.upstream_model_id
|
||||
.as_ref()
|
||||
.map(|model| model.as_str().to_owned()),
|
||||
failure_kind: error.kind().as_str().to_owned(),
|
||||
status_code: error.upstream_status().or(failure.upstream_status_code),
|
||||
provider_error_code: error.upstream_code().map(|code| code.as_str().to_owned()),
|
||||
@@ -1081,7 +1088,9 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
|
||||
operation: request.operation.as_str().to_owned(),
|
||||
endpoint: request.endpoint,
|
||||
client_transport: request.client_transport,
|
||||
requested_model_id: request.requested_model.as_str().to_owned(),
|
||||
requested_model_id: request
|
||||
.requested_model
|
||||
.map(|model| model.as_str().to_owned()),
|
||||
client_ip: request.client_ip.map(|address| address.to_string()),
|
||||
user_agent: request.user_agent,
|
||||
reasoning_effort: request.reasoning_effort,
|
||||
@@ -1150,7 +1159,9 @@ fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart {
|
||||
.provider_account_ref
|
||||
.as_ref()
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
upstream_model_id: attempt.upstream_model_id.as_str().to_owned(),
|
||||
upstream_model_id: attempt
|
||||
.upstream_model_id
|
||||
.map(|model| model.as_str().to_owned()),
|
||||
upstream_transport: attempt.upstream_transport,
|
||||
http_version: attempt.http_version,
|
||||
}
|
||||
|
||||
@@ -551,7 +551,10 @@ impl ExecutionObservationWrite {
|
||||
Some(failure.request_id.as_str()),
|
||||
Some(failure.provider_kind.as_str()),
|
||||
failure.account_id.as_ref().map(|value| value.as_str()),
|
||||
Some(failure.upstream_model_id.as_str()),
|
||||
failure
|
||||
.upstream_model_id
|
||||
.as_ref()
|
||||
.map(|model| model.as_str()),
|
||||
failure.upstream_request_id.as_deref(),
|
||||
])
|
||||
.saturating_add(provider_error_bytes(&failure.error)),
|
||||
@@ -638,7 +641,7 @@ fn new_request_bytes(request: &NewModelRequest) -> usize {
|
||||
Some(request.protocol.as_str()),
|
||||
Some(request.endpoint.as_str()),
|
||||
Some(request.client_transport.as_str()),
|
||||
Some(request.requested_model.as_str()),
|
||||
request.requested_model.as_ref().map(|model| model.as_str()),
|
||||
request.user_agent.as_deref(),
|
||||
request.reasoning_effort.as_deref(),
|
||||
request.reasoning_preset.as_deref(),
|
||||
@@ -659,7 +662,10 @@ fn attempt_bytes(attempt: &AttemptRecord) -> usize {
|
||||
.provider_account_ref
|
||||
.as_ref()
|
||||
.map(|value| value.as_str()),
|
||||
Some(attempt.upstream_model_id.as_str()),
|
||||
attempt
|
||||
.upstream_model_id
|
||||
.as_ref()
|
||||
.map(|model| model.as_str()),
|
||||
Some(attempt.upstream_transport.as_str()),
|
||||
attempt.http_version.as_deref(),
|
||||
]))
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 九张终态业务表的 PostgreSQL adapters。
|
||||
//! 当前十二张业务表的 PostgreSQL adapters。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, postgres::PgPoolOptions};
|
||||
|
||||
@@ -504,7 +504,7 @@ pub struct UsageRecord {
|
||||
pub operation: String,
|
||||
pub endpoint: String,
|
||||
pub client_transport: String,
|
||||
pub requested_model_id: String,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub provider_account_name: Option<String>,
|
||||
|
||||
@@ -259,6 +259,7 @@ pub(crate) async fn account_model_rows(
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id)
|
||||
order by mr.provider_account_ref, request_count desc, model",
|
||||
)
|
||||
@@ -286,6 +287,7 @@ pub(crate) async fn account_model_costs(
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_amount is not null and mr.cost_currency is not null
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id),
|
||||
mr.cost_currency
|
||||
order by mr.provider_account_ref, model, mr.cost_currency",
|
||||
|
||||
@@ -682,8 +682,10 @@ impl AccountStore for PgAdminAccountStore {
|
||||
) -> AdminStoreResult<CredentialMutationResult> {
|
||||
match command.credential {
|
||||
AuthorizationCredentialCommit::Create(credential) => {
|
||||
let account_id = credential.account_id.clone();
|
||||
let result = self
|
||||
let CredentialImportResult {
|
||||
config_revision,
|
||||
credential_ids,
|
||||
} = self
|
||||
.commit_prepared_import(
|
||||
PreparedCredentialImport {
|
||||
provider_kind: credential.provider_kind.clone(),
|
||||
@@ -693,6 +695,14 @@ impl AccountStore for PgAdminAccountStore {
|
||||
"authorize",
|
||||
)
|
||||
.await?;
|
||||
let [account_id]: [CoreProviderAccountId; 1] =
|
||||
credential_ids.try_into().map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
ENTITY,
|
||||
"authorization import returned an unexpected account count",
|
||||
)
|
||||
})?;
|
||||
let details = self
|
||||
.accounts
|
||||
.load_provider_account(account_id.as_str())
|
||||
@@ -706,7 +716,7 @@ impl AccountStore for PgAdminAccountStore {
|
||||
)
|
||||
})?;
|
||||
Ok(CredentialMutationResult {
|
||||
config_revision: result.config_revision,
|
||||
config_revision,
|
||||
account_id,
|
||||
credential_revision: Some(admin_revision(details.summary.credential_revision)?),
|
||||
})
|
||||
|
||||
@@ -103,7 +103,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
&self,
|
||||
update: CredentialCasUpdate,
|
||||
) -> Result<CredentialCasOutcome, CoreStoreError> {
|
||||
let (
|
||||
let CredentialCasUpdateParts {
|
||||
account_id,
|
||||
expected_revision,
|
||||
profile,
|
||||
@@ -111,7 +111,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
) = update.into_parts();
|
||||
account_state,
|
||||
} = update.into_parts();
|
||||
if profile.account_id != account_id {
|
||||
return Err(CoreStoreError::new(CoreStoreErrorKind::InvalidData));
|
||||
}
|
||||
@@ -122,13 +123,38 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
CREDENTIALS_MAX_BYTES,
|
||||
)
|
||||
.map_err(core_store_error)?;
|
||||
let credential_state = account_state
|
||||
.as_ref()
|
||||
.map(|state| state.credential_state.as_str());
|
||||
let credential_observed_at = account_state
|
||||
.as_ref()
|
||||
.map(|state| DateTime::<Utc>::from(state.observed_at));
|
||||
let error_reason = account_state
|
||||
.as_ref()
|
||||
.and_then(|state| state.error_reason.map(AccountErrorReason::as_str));
|
||||
let message = account_state
|
||||
.as_ref()
|
||||
.and_then(|state| state.message.as_deref());
|
||||
let next = sqlx::query_scalar::<_, i64>(
|
||||
"update provider_accounts
|
||||
set name = $3, email = $4, plan_type = $5,
|
||||
provider_credentials_json = $6,
|
||||
credential_revision = credential_revision + 1,
|
||||
has_refresh_token = $7, access_token_expires_at = $8,
|
||||
next_refresh_at = $9, updated_at = now()
|
||||
next_refresh_at = $9,
|
||||
credential_state = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $10 else credential_state end,
|
||||
credential_observed_at = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $11 else credential_observed_at end,
|
||||
last_error_reason = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $12 else last_error_reason end,
|
||||
last_error_message = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $13 else last_error_message end,
|
||||
updated_at = greatest(now(), coalesce($11, now()))
|
||||
where id = $1 and credential_revision = $2
|
||||
returning credential_revision",
|
||||
)
|
||||
@@ -141,6 +167,10 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.bind(has_refresh_token)
|
||||
.bind(access_token_expires_at.map(DateTime::<Utc>::from))
|
||||
.bind(next_refresh_at.map(DateTime::<Utc>::from))
|
||||
.bind(credential_state)
|
||||
.bind(credential_observed_at)
|
||||
.bind(error_reason)
|
||||
.bind(message)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::Unavailable))?;
|
||||
@@ -280,6 +310,25 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
})
|
||||
}
|
||||
|
||||
async fn touch_quota_observation(
|
||||
&self,
|
||||
touch: QuotaObservationTouch,
|
||||
) -> Result<QuotaWriteOutcome, CoreStoreError> {
|
||||
let updated = self
|
||||
.touch_provider_quota_observation(
|
||||
touch.account_id.as_str(),
|
||||
Revision::new(touch.expected_revision.get()).map_err(core_store_error)?,
|
||||
DateTime::<Utc>::from(touch.observed_at),
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)?;
|
||||
Ok(if updated {
|
||||
QuotaWriteOutcome::Updated
|
||||
} else {
|
||||
QuotaWriteOutcome::Conflict
|
||||
})
|
||||
}
|
||||
|
||||
async fn apply_state_change(&self, change: AccountStateChange) -> Result<(), CoreStoreError> {
|
||||
let updated = self
|
||||
.apply_provider_account_state(ProviderAccountStateUpdate {
|
||||
|
||||
@@ -38,12 +38,13 @@ use sqlx::{PgPool, Postgres, Row, Transaction};
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision as CoreCredentialRevision,
|
||||
CredentialState, LoadedCredential, NewProviderAccount as CoreNewProviderAccount,
|
||||
OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount,
|
||||
ProviderAccountId as CoreProviderAccountId, ProviderAccountIdentity, ProviderAccountStore,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialCasUpdateParts,
|
||||
CredentialRevision as CoreCredentialRevision, CredentialState, LoadedCredential,
|
||||
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
|
||||
ProviderAccountIdentity, ProviderAccountStore,
|
||||
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState,
|
||||
QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
QuotaEvidence, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
|
||||
@@ -29,6 +29,12 @@ pub trait ProviderAccountRepository: Send + Sync {
|
||||
observed_at: DateTime<Utc>,
|
||||
state: QuotaState,
|
||||
) -> StoreResult<bool>;
|
||||
async fn touch_provider_quota_observation(
|
||||
&self,
|
||||
account_id: &str,
|
||||
expected_revision: Revision,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<bool>;
|
||||
async fn apply_provider_quota_access(
|
||||
&self,
|
||||
account_id: &str,
|
||||
@@ -237,13 +243,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
else credential_observed_at
|
||||
end,
|
||||
last_error_reason = case
|
||||
when enabled and upstream_user_id is not null then
|
||||
case when $3 = 'ready' then null else $5 end
|
||||
when enabled and upstream_user_id is not null then $5
|
||||
else last_error_reason
|
||||
end,
|
||||
last_error_message = case
|
||||
when enabled and upstream_user_id is not null then
|
||||
case when $3 = 'ready' then null else $6 end
|
||||
when enabled and upstream_user_id is not null then $6
|
||||
else last_error_message
|
||||
end,
|
||||
updated_at = case when enabled then greatest(now(), $4) else updated_at end
|
||||
@@ -354,6 +358,29 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
async fn touch_provider_quota_observation(
|
||||
&self,
|
||||
account_id: &str,
|
||||
expected_revision: Revision,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "account_id", account_id)?;
|
||||
let result = sqlx::query(
|
||||
"update provider_accounts
|
||||
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)",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(to_i64(expected_revision.get())?)
|
||||
.bind(observed_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("touch provider quota observation"))?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "id", id)?;
|
||||
let result = sqlx::query("delete from provider_accounts where id = $1 and not enabled")
|
||||
|
||||
@@ -91,6 +91,7 @@ pub(crate) const ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL: &str = "with requested_win
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by requested.account_id, requested.window_key,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id)
|
||||
order by requested.account_id, requested.window_key, request_count desc, model";
|
||||
@@ -115,6 +116,7 @@ pub(crate) const ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL: &str = "with requeste
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_amount is not null
|
||||
and mr.cost_currency is not null
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by requested.account_id, requested.window_key,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id), mr.cost_currency
|
||||
order by requested.account_id, requested.window_key, model, mr.cost_currency";
|
||||
|
||||
@@ -16,6 +16,24 @@ use super::{namespace, resource_fingerprint};
|
||||
|
||||
const MAX_SESSION_AFFINITY_TTL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
|
||||
const CLAIM_OR_LOAD_SCRIPT: &str = r#"
|
||||
local current = redis.call('GET', KEYS[1])
|
||||
if current then
|
||||
return current
|
||||
end
|
||||
redis.call('PSETEX', KEYS[1], tonumber(ARGV[2]), ARGV[1])
|
||||
return ARGV[1]
|
||||
"#;
|
||||
|
||||
const COMPARE_AND_BIND_SCRIPT: &str = r#"
|
||||
local current = redis.call('GET', KEYS[1])
|
||||
if not current or current == ARGV[1] then
|
||||
redis.call('PSETEX', KEYS[1], tonumber(ARGV[3]), ARGV[2])
|
||||
return ARGV[2]
|
||||
end
|
||||
return current
|
||||
"#;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RedisProviderSessionAffinityRepository {
|
||||
connection: ConnectionManager,
|
||||
@@ -79,11 +97,7 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
|
||||
ttl: Duration,
|
||||
) -> futures::future::BoxFuture<'a, Result<(), ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
if ttl.is_zero() || ttl > MAX_SESSION_AFFINITY_TTL {
|
||||
return Err(provider_invalid("validate provider session affinity TTL"));
|
||||
}
|
||||
let ttl_millis = u64::try_from(ttl.as_millis())
|
||||
.map_err(|_| provider_invalid("validate provider session affinity TTL"))?;
|
||||
let ttl_millis = session_affinity_ttl_millis(ttl)?;
|
||||
let mut connection = self.connection.clone();
|
||||
redis::cmd("PSETEX")
|
||||
.arg(self.key(provider_kind, key)?)
|
||||
@@ -95,6 +109,50 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
|
||||
})
|
||||
}
|
||||
|
||||
fn claim_or_load<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
key: &'a ProviderSessionAffinityKey,
|
||||
candidate_account_id: &'a ProviderAccountId,
|
||||
ttl: Duration,
|
||||
) -> futures::future::BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let ttl_millis = session_affinity_ttl_millis(ttl)?;
|
||||
let mut connection = self.connection.clone();
|
||||
let effective_account_id = redis::Script::new(CLAIM_OR_LOAD_SCRIPT)
|
||||
.key(self.key(provider_kind, key)?)
|
||||
.arg(candidate_account_id.as_str())
|
||||
.arg(ttl_millis)
|
||||
.invoke_async::<String>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("claim provider session affinity"))?;
|
||||
decode_account_id(effective_account_id)
|
||||
})
|
||||
}
|
||||
|
||||
fn compare_and_bind<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
key: &'a ProviderSessionAffinityKey,
|
||||
expected_account_id: &'a ProviderAccountId,
|
||||
replacement_account_id: &'a ProviderAccountId,
|
||||
ttl: Duration,
|
||||
) -> futures::future::BoxFuture<'a, Result<ProviderAccountId, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let ttl_millis = session_affinity_ttl_millis(ttl)?;
|
||||
let mut connection = self.connection.clone();
|
||||
let effective_account_id = redis::Script::new(COMPARE_AND_BIND_SCRIPT)
|
||||
.key(self.key(provider_kind, key)?)
|
||||
.arg(expected_account_id.as_str())
|
||||
.arg(replacement_account_id.as_str())
|
||||
.arg(ttl_millis)
|
||||
.invoke_async::<String>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("compare provider session affinity"))?;
|
||||
decode_account_id(effective_account_id)
|
||||
})
|
||||
}
|
||||
|
||||
fn clear<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
@@ -112,6 +170,19 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
|
||||
}
|
||||
}
|
||||
|
||||
fn session_affinity_ttl_millis(ttl: Duration) -> Result<u64, ProviderStoreError> {
|
||||
if ttl.is_zero() || ttl > MAX_SESSION_AFFINITY_TTL {
|
||||
return Err(provider_invalid("validate provider session affinity TTL"));
|
||||
}
|
||||
u64::try_from(ttl.as_millis())
|
||||
.map_err(|_| provider_invalid("validate provider session affinity TTL"))
|
||||
}
|
||||
|
||||
fn decode_account_id(account_id: String) -> Result<ProviderAccountId, ProviderStoreError> {
|
||||
ProviderAccountId::new(account_id)
|
||||
.map_err(|_| provider_invalid("decode provider session affinity"))
|
||||
}
|
||||
|
||||
fn provider_unavailable(operation: &'static str) -> ProviderStoreError {
|
||||
ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -10,7 +10,12 @@ use gateway_admin::{
|
||||
UpdateAccount,
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{CredentialListQuery, CredentialListWindow, CredentialStateFilter},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, AuthorizationCredentialCommit, AuthorizationMutationTarget,
|
||||
AuthorizationOwnerBinding, CredentialListQuery, CredentialListWindow,
|
||||
CredentialStateFilter, PendingAuthorizationMutation, PreparedCredentialCreate,
|
||||
ProviderDocument,
|
||||
},
|
||||
},
|
||||
ports::store::AccountStore,
|
||||
};
|
||||
@@ -18,8 +23,8 @@ use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, CredentialState, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccountId, ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch,
|
||||
QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
use gateway_store::{
|
||||
@@ -51,6 +56,22 @@ struct RecoveredAccountRow {
|
||||
access_token_expires_at: Option<chrono::DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(sqlx::FromRow)]
|
||||
struct CoreRefreshRow {
|
||||
name: String,
|
||||
email: Option<String>,
|
||||
plan_type: Option<String>,
|
||||
provider_credentials_json: serde_json::Value,
|
||||
credential_revision: i64,
|
||||
credential_state: String,
|
||||
last_error_reason: Option<String>,
|
||||
last_error_message: Option<String>,
|
||||
quota_access_state: String,
|
||||
quota_evidence: Option<String>,
|
||||
quota_access_observed_at: Option<chrono::DateTime<Utc>>,
|
||||
quota_reset_at: Option<chrono::DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn postgres_provider_account_adapter_implements_core_port() {
|
||||
fn assert_port<T: ProviderAccountStore>() {}
|
||||
@@ -156,6 +177,115 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
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;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let account_id = ProviderAccountId::new("acct_quota_touch").expect("account id");
|
||||
repository
|
||||
.insert_provider_account(account(account_id.as_str(), "user-quota-touch"))
|
||||
.await
|
||||
.expect("insert quota fixture");
|
||||
let revision = CredentialRevision::new(1).expect("revision");
|
||||
let observed_at = SystemTime::now()
|
||||
.checked_sub(Duration::from_secs(60))
|
||||
.expect("old observation time");
|
||||
let touched_at = observed_at
|
||||
.checked_add(Duration::from_secs(90))
|
||||
.expect("new observation time");
|
||||
let reset_at = observed_at
|
||||
.checked_add(Duration::from_secs(3_600))
|
||||
.expect("quota reset time");
|
||||
let quota_json = json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
}
|
||||
});
|
||||
assert_eq!(
|
||||
repository
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
quota: OpaqueProviderData::new(
|
||||
quota_json.as_object().expect("quota object").clone(),
|
||||
),
|
||||
observed_at,
|
||||
state: QuotaState::exhausted(
|
||||
QuotaEvidence::UsageLimitReached,
|
||||
observed_at,
|
||||
Some(reset_at),
|
||||
),
|
||||
})
|
||||
.await
|
||||
.expect("persist exhausted quota"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
let account_updated_at: chrono::DateTime<Utc> = sqlx::query_scalar(
|
||||
"select updated_at from provider_accounts where id = 'acct_quota_touch'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load account update time before quota touch");
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.touch_quota_observation(QuotaObservationTouch {
|
||||
account_id,
|
||||
expected_revision: revision,
|
||||
observed_at: touched_at,
|
||||
})
|
||||
.await
|
||||
.expect("touch quota refresh time"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
|
||||
let stored = sqlx::query_as::<
|
||||
_,
|
||||
(
|
||||
serde_json::Value,
|
||||
chrono::DateTime<Utc>,
|
||||
String,
|
||||
Option<String>,
|
||||
chrono::DateTime<Utc>,
|
||||
Option<chrono::DateTime<Utc>>,
|
||||
chrono::DateTime<Utc>,
|
||||
),
|
||||
>(
|
||||
"select provider_quota_json, quota_observed_at, quota_access_state,
|
||||
quota_evidence, quota_access_observed_at, quota_reset_at, updated_at
|
||||
from provider_accounts where id = 'acct_quota_touch'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load touched quota row");
|
||||
assert_eq!(stored.0, quota_json);
|
||||
assert_eq!(
|
||||
stored.1.timestamp_micros(),
|
||||
chrono::DateTime::<Utc>::from(touched_at).timestamp_micros()
|
||||
);
|
||||
assert_eq!(stored.2, "exhausted");
|
||||
assert_eq!(stored.3.as_deref(), Some("usage_limit_reached"));
|
||||
assert_eq!(
|
||||
stored.4.timestamp_micros(),
|
||||
chrono::DateTime::<Utc>::from(observed_at).timestamp_micros()
|
||||
);
|
||||
assert_eq!(
|
||||
stored.5.map(|value| value.timestamp_micros()),
|
||||
Some(chrono::DateTime::<Utc>::from(reset_at).timestamp_micros())
|
||||
);
|
||||
assert_eq!(
|
||||
stored.6.timestamp_micros(),
|
||||
chrono::DateTime::<Utc>::from(touched_at).timestamp_micros()
|
||||
);
|
||||
assert!(stored.6 > account_updated_at);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
|
||||
let Some(database) = TestDatabase::create("provider_account_observation_fence").await else {
|
||||
@@ -280,6 +410,27 @@ async fn credential_error_message_is_persisted_and_cleared_on_recovery() {
|
||||
Some("upstream returned 402 payment required")
|
||||
);
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: Some("Invalid refresh token.".to_owned()),
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
credential_state: CredentialState::Ready,
|
||||
observed_at: SystemTime::now(),
|
||||
error_reason: Some(AccountErrorReason::AccessTokenExpired),
|
||||
})
|
||||
.await
|
||||
.expect("persist recoverable refresh failure");
|
||||
let recoverable: (Option<String>, Option<String>) = sqlx::query_as(
|
||||
"select last_error_reason, last_error_message from provider_accounts where id = $1",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("read recoverable refresh failure");
|
||||
assert_eq!(recoverable.0.as_deref(), Some("access_token_expired"));
|
||||
assert_eq!(recoverable.1.as_deref(), Some("Invalid refresh token."));
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
@@ -1180,6 +1331,75 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorization_create_returns_existing_account_id_when_identity_is_upserted() {
|
||||
let Some(database) = TestDatabase::create("provider_account_authorization_upsert").await else {
|
||||
return;
|
||||
};
|
||||
PgProviderAccountRepository::new(database.pool.clone())
|
||||
.insert_provider_account(account(
|
||||
"acct_authorization_existing",
|
||||
"user-authorization-upsert",
|
||||
))
|
||||
.await
|
||||
.expect("seed existing authorized identity");
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "request_authorization_upsert".to_owned(),
|
||||
};
|
||||
let provider_kind = ProviderKind::new("openai").expect("OpenAI Provider kind");
|
||||
let serde_json::Value::Object(provider_material) =
|
||||
json!({ "access_token": "authorized-secret" })
|
||||
else {
|
||||
unreachable!("credential fixture must be an object");
|
||||
};
|
||||
|
||||
let result = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
.commit_authorization(
|
||||
AuthorizationCommit {
|
||||
pending: PendingAuthorizationMutation::new(
|
||||
provider_kind.clone(),
|
||||
AuthorizationMutationTarget::Create {
|
||||
name: "authorized account".to_owned(),
|
||||
},
|
||||
AuthorizationOwnerBinding::from_context(&context),
|
||||
),
|
||||
credential: AuthorizationCredentialCommit::Create(PreparedCredentialCreate {
|
||||
account_id: ProviderAccountId::new("acct_authorization_candidate")
|
||||
.expect("candidate account ID"),
|
||||
provider_kind,
|
||||
name: "authorized account".to_owned(),
|
||||
email: Some("authorized@example.invalid".to_owned()),
|
||||
upstream_user_id: Some("user-authorization-upsert".to_owned()),
|
||||
upstream_account_id: None,
|
||||
plan_type: Some("free".to_owned()),
|
||||
authentication_kind: "oauth".to_owned(),
|
||||
provider_material: ProviderDocument::new(OpaqueProviderData::new(
|
||||
provider_material,
|
||||
)),
|
||||
has_refresh_token: true,
|
||||
access_token_expires_at: Some(Utc::now() + TimeDelta::hours(1)),
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
credential_state: CredentialState::Ready,
|
||||
credential_observed_at: Utc::now(),
|
||||
}),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.expect("authorize existing identity");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
result.account_id.as_str(),
|
||||
result.credential_revision.map(|revision| revision.get()),
|
||||
),
|
||||
("acct_authorization_existing", Some(2)),
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
let Some(database) = TestDatabase::create("provider_account_core_refresh").await else {
|
||||
@@ -1208,8 +1428,37 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
})
|
||||
.await
|
||||
.expect("seed provider account");
|
||||
sqlx::query(
|
||||
"update provider_accounts
|
||||
set last_error_reason = 'access_token_expired',
|
||||
last_error_message = 'stale refresh failure'
|
||||
where id = 'acct_core_refresh'",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed refresh failure details");
|
||||
|
||||
let account_id = ProviderAccountId::new("acct_core_refresh").expect("account ID");
|
||||
let quota_observed_at = SystemTime::now();
|
||||
let quota_reset_at = quota_observed_at
|
||||
.checked_add(Duration::from_secs(3_600))
|
||||
.expect("quota reset time");
|
||||
assert_eq!(
|
||||
repository
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("credential revision"),
|
||||
state: QuotaState::exhausted(
|
||||
QuotaEvidence::UsageLimitReached,
|
||||
quota_observed_at,
|
||||
Some(quota_reset_at),
|
||||
),
|
||||
})
|
||||
.await
|
||||
.expect("seed quota exhaustion"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
let state_observed_at = SystemTime::now();
|
||||
let refreshed = CredentialCasUpdate::new(
|
||||
account_id.clone(),
|
||||
CredentialRevision::new(1).expect("credential revision"),
|
||||
@@ -1224,7 +1473,8 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
Some(SystemTime::now() + Duration::from_secs(3_600)),
|
||||
Some(SystemTime::now() + Duration::from_secs(1_800)),
|
||||
)
|
||||
.expect("valid refresh update");
|
||||
.expect("valid refresh update")
|
||||
.with_account_state(CredentialState::Ready, state_observed_at, None, None);
|
||||
assert_eq!(
|
||||
repository
|
||||
.compare_and_swap_credential(refreshed)
|
||||
@@ -1235,24 +1485,37 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
)
|
||||
);
|
||||
|
||||
let row: (
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
serde_json::Value,
|
||||
i64,
|
||||
) = sqlx::query_as(
|
||||
"select name, email, plan_type, provider_credentials_json, credential_revision
|
||||
let row = sqlx::query_as::<_, CoreRefreshRow>(
|
||||
"select name, email, plan_type, provider_credentials_json, credential_revision,
|
||||
credential_state, last_error_reason, last_error_message,
|
||||
quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at
|
||||
from provider_accounts where id = 'acct_core_refresh'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load refreshed account");
|
||||
assert_eq!(row.0, "after refresh");
|
||||
assert_eq!(row.1.as_deref(), Some("after@example.invalid"));
|
||||
assert_eq!(row.2.as_deref(), Some("premium"));
|
||||
assert_eq!(row.3["access_token"], "after-secret");
|
||||
assert_eq!(row.4, 2);
|
||||
assert_eq!(row.name, "after refresh");
|
||||
assert_eq!(row.email.as_deref(), Some("after@example.invalid"));
|
||||
assert_eq!(row.plan_type.as_deref(), Some("premium"));
|
||||
assert_eq!(
|
||||
row.provider_credentials_json["access_token"],
|
||||
"after-secret"
|
||||
);
|
||||
assert_eq!(row.credential_revision, 2);
|
||||
assert_eq!(row.credential_state, "ready");
|
||||
assert_eq!(row.last_error_reason, None);
|
||||
assert_eq!(row.last_error_message, None);
|
||||
assert_eq!(row.quota_access_state, "exhausted");
|
||||
assert_eq!(row.quota_evidence.as_deref(), Some("usage_limit_reached"));
|
||||
assert_eq!(
|
||||
row.quota_access_observed_at
|
||||
.map(|value| value.timestamp_micros()),
|
||||
Some(chrono::DateTime::<Utc>::from(quota_observed_at).timestamp_micros())
|
||||
);
|
||||
assert_eq!(
|
||||
row.quota_reset_at.map(|value| value.timestamp_micros()),
|
||||
Some(chrono::DateTime::<Utc>::from(quota_reset_at).timestamp_micros())
|
||||
);
|
||||
|
||||
assert!(
|
||||
repository
|
||||
|
||||
@@ -436,7 +436,7 @@ async fn diagnostics_should_fallback_to_name_then_ref_for_missing_snapshots() {
|
||||
provider_kind: "openai".to_owned(),
|
||||
provider_account_id: None,
|
||||
provider_account_ref: Some("acct_legacy_deleted".to_owned()),
|
||||
upstream_model_id: "legacy-model".to_owned(),
|
||||
upstream_model_id: Some("legacy-model".to_owned()),
|
||||
upstream_transport: "http_sse".to_owned(),
|
||||
http_version: None,
|
||||
},
|
||||
@@ -547,7 +547,7 @@ async fn api_key_diagnostics_should_display_key_name_and_fallback_to_ref() {
|
||||
provider_kind: "openai".to_owned(),
|
||||
provider_account_id: None,
|
||||
provider_account_ref: Some("acct_key".to_owned()),
|
||||
upstream_model_id: "upstream-model".to_owned(),
|
||||
upstream_model_id: Some("upstream-model".to_owned()),
|
||||
upstream_transport: "http_sse".to_owned(),
|
||||
http_version: None,
|
||||
};
|
||||
@@ -664,7 +664,7 @@ fn new_request(id: &str, started_at: DateTime<Utc>) -> NewModelRequest {
|
||||
operation: "responses".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
requested_model_id: Some("coding".to_owned()),
|
||||
routing_scope: "all".to_owned(),
|
||||
routing_group_refs: Vec::new(),
|
||||
routing_group_names_snapshot: serde_json::json!([]),
|
||||
@@ -688,7 +688,7 @@ fn attempt(id: &str, count: u32, account_id: &str) -> ModelRequestAttemptStart {
|
||||
provider_kind: "openai".to_owned(),
|
||||
provider_account_id: Some(account_id.to_owned()),
|
||||
provider_account_ref: Some(account_id.to_owned()),
|
||||
upstream_model_id: "upstream-model".to_owned(),
|
||||
upstream_model_id: Some("upstream-model".to_owned()),
|
||||
upstream_transport: "http_sse".to_owned(),
|
||||
http_version: None,
|
||||
}
|
||||
|
||||
@@ -65,6 +65,55 @@ async fn session_affinity_should_overwrite_the_previous_account() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_affinity_claim_should_keep_the_first_account() {
|
||||
let Some((repository, _connection, _namespace)) = affinity_repository().await else {
|
||||
return;
|
||||
};
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let key = ProviderSessionAffinityKey::try_new("claim-session").expect("affinity key");
|
||||
let first = ProviderAccountId::new("acct_first").expect("first account");
|
||||
let second = ProviderAccountId::new("acct_second").expect("second account");
|
||||
|
||||
let first_winner = repository
|
||||
.claim_or_load(&provider, &key, &first, Duration::from_secs(60))
|
||||
.await
|
||||
.expect("claim first affinity");
|
||||
let second_winner = repository
|
||||
.claim_or_load(&provider, &key, &second, Duration::from_secs(60))
|
||||
.await
|
||||
.expect("load existing affinity");
|
||||
|
||||
assert_eq!((first_winner, second_winner), (first.clone(), first));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_affinity_compare_should_not_replace_a_newer_winner() {
|
||||
let Some((repository, _connection, _namespace)) = affinity_repository().await else {
|
||||
return;
|
||||
};
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let key = ProviderSessionAffinityKey::try_new("compare-session").expect("affinity key");
|
||||
let first = ProviderAccountId::new("acct_first").expect("first account");
|
||||
let second = ProviderAccountId::new("acct_second").expect("second account");
|
||||
let stale = ProviderAccountId::new("acct_stale").expect("stale account");
|
||||
repository
|
||||
.bind(&provider, &key, &first, Duration::from_secs(60))
|
||||
.await
|
||||
.expect("seed affinity");
|
||||
|
||||
let migrated = repository
|
||||
.compare_and_bind(&provider, &key, &first, &second, Duration::from_secs(60))
|
||||
.await
|
||||
.expect("migrate expected affinity");
|
||||
let stale_result = repository
|
||||
.compare_and_bind(&provider, &key, &first, &stale, Duration::from_secs(60))
|
||||
.await
|
||||
.expect("reject stale affinity migration");
|
||||
|
||||
assert_eq!((migrated, stale_result), (second.clone(), second));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_affinity_should_apply_ttl_and_support_explicit_clear() {
|
||||
let Some((repository, mut connection, namespace)) = affinity_repository().await else {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -13,20 +13,26 @@ 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, ProviderUsageStatistics,
|
||||
ProviderUsageStatisticsCycle, ProviderUsageStatisticsDay, ProviderUsageStatisticsMode,
|
||||
ProviderUsageStatisticsModel, ProviderUsageStatisticsRequest,
|
||||
ProviderUsageStatisticsServiceTier, ProviderUsageStatisticsSummary,
|
||||
ProviderUsageStatisticsTokens, QuotaLocalUsageAttribution,
|
||||
};
|
||||
use gateway_admin::model::{MutationActor, MutationContext, Revision};
|
||||
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
|
||||
use gateway_core::accounting::Money;
|
||||
use gateway_core::engine::credential::{
|
||||
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
|
||||
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
|
||||
CredentialCasUpdateParts, CredentialRevision, LoadedCredential, NewProviderAccount,
|
||||
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountStore,
|
||||
};
|
||||
use gateway_core::error::StoreErrorKind;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
|
||||
@@ -46,10 +52,14 @@ use crate::credential::{
|
||||
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError,
|
||||
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
|
||||
CodexPendingAuthorization, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole,
|
||||
CodexResetCreditsError, CodexUsageStatistics, CodexUsageStatisticsError,
|
||||
CodexUsageStatisticsMode, CodexUsageStatisticsServiceTier, CodexUsageStatisticsTokens,
|
||||
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
|
||||
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
};
|
||||
use crate::credential::{CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref};
|
||||
use crate::credential::{
|
||||
CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref, parse_access_token_expiration,
|
||||
};
|
||||
use crate::transport::CodexWebSocketPool;
|
||||
use crate::transport::profile::{
|
||||
CodexDesktopReleaseSnapshot, CodexDesktopReleaseStatus, CodexWireProfile, CodexWireProfileState,
|
||||
@@ -258,7 +268,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
))
|
||||
.await
|
||||
.inspect_err(|error| {
|
||||
log_import_failure("prepare_document", credential_admin_error_code(*error));
|
||||
log_import_failure("prepare_document", credential_admin_error_code(error));
|
||||
})
|
||||
.map_err(map_credential_admin_error)?;
|
||||
let observed_at = Utc::now();
|
||||
@@ -380,8 +390,10 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
.oauth()
|
||||
.and_then(|oauth| oauth.id_token.clone());
|
||||
}
|
||||
let access_token_expires_at =
|
||||
parse_access_token_expiration(secret.access_token.expose_secret());
|
||||
let prepared = CodexCredentialAdmin
|
||||
.prepare_refreshed_oauth_rotation(current, secret, None, None)
|
||||
.prepare_refreshed_oauth_rotation(current, secret, access_token_expires_at, None)
|
||||
.map_err(map_credential_admin_error)?;
|
||||
prepared_rotation(prepared, command.account.provider_kind)
|
||||
}
|
||||
@@ -437,6 +449,54 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
Ok(project_quota(snapshot, &account))
|
||||
}
|
||||
|
||||
async fn usage_statistics(
|
||||
&self,
|
||||
request: ProviderUsageStatisticsRequest,
|
||||
) -> Result<ProviderUsageStatistics, ProviderAdminError> {
|
||||
self.account(&request.account_id).await?;
|
||||
let statistics = self
|
||||
.quota
|
||||
.usage_statistics(
|
||||
&request.account_id,
|
||||
request.cycle_offset,
|
||||
request.utc_offset_minutes,
|
||||
)
|
||||
.await
|
||||
.map_err(map_usage_statistics_error)?;
|
||||
project_usage_statistics(statistics)
|
||||
}
|
||||
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderResetCredits, ProviderAdminError> {
|
||||
self.account(account_id).await?;
|
||||
self.quota
|
||||
.list_reset_credits(account_id)
|
||||
.await
|
||||
.map(project_reset_credits)
|
||||
.map_err(map_reset_credits_error)
|
||||
}
|
||||
|
||||
async fn consume_reset_credit(
|
||||
&self,
|
||||
command: ConsumeProviderResetCredit,
|
||||
) -> Result<ProviderResetCreditResult, ProviderAdminError> {
|
||||
self.account(&command.account_id).await?;
|
||||
self.quota
|
||||
.consume_reset_credit(
|
||||
&command.account_id,
|
||||
command.credit_id.as_deref(),
|
||||
command.redeem_request_id,
|
||||
)
|
||||
.await
|
||||
.map(|result| ProviderResetCreditResult {
|
||||
code: result.code,
|
||||
credit: result.credit.map(project_reset_credit),
|
||||
})
|
||||
.map_err(map_reset_credits_error)
|
||||
}
|
||||
|
||||
async fn models(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -546,15 +606,16 @@ fn prepared_rotation(
|
||||
provider_kind: ProviderKind,
|
||||
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
|
||||
let (profile, credential, replacement_identity, guard) = prepared.into_parts();
|
||||
let (
|
||||
let CredentialCasUpdateParts {
|
||||
account_id,
|
||||
expected_revision,
|
||||
credential_profile,
|
||||
profile: credential_profile,
|
||||
credential,
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
) = credential.into_parts();
|
||||
account_state: _account_state,
|
||||
} = credential.into_parts();
|
||||
if profile != credential_profile || profile.account_id != account_id {
|
||||
return Err(provider_admin_error(ProviderAdminErrorKind::Internal));
|
||||
}
|
||||
@@ -728,39 +789,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 只看滚动后的窗口触顶:顶层标记是观测事实,不能
|
||||
@@ -777,6 +823,123 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
|
||||
}
|
||||
}
|
||||
|
||||
fn project_usage_statistics(
|
||||
statistics: CodexUsageStatistics,
|
||||
) -> Result<ProviderUsageStatistics, ProviderAdminError> {
|
||||
let models = statistics
|
||||
.models
|
||||
.into_iter()
|
||||
.map(|model| {
|
||||
Ok(ProviderUsageStatisticsModel {
|
||||
key: model.key,
|
||||
model: model.model,
|
||||
service_tier: project_usage_service_tier(model.service_tier),
|
||||
credit_share: model.credit_share,
|
||||
quota_share: model.quota_share,
|
||||
tokens: project_usage_tokens(model.tokens),
|
||||
estimated_cost: model.estimated_cost.map(currency_cost).transpose()?,
|
||||
has_unknown_pricing: model.has_unknown_pricing,
|
||||
has_estimated_allocation: model.has_estimated_allocation,
|
||||
has_rate_fallback: model.has_rate_fallback,
|
||||
has_missing_token_data: model.has_missing_token_data,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, ProviderAdminError>>()?;
|
||||
let daily = statistics
|
||||
.daily
|
||||
.into_iter()
|
||||
.map(|day| {
|
||||
Ok(ProviderUsageStatisticsDay {
|
||||
date: day.date,
|
||||
credit_share: day.credit_share,
|
||||
tokens: project_usage_tokens(day.tokens),
|
||||
estimated_cost: day.estimated_cost.map(currency_cost).transpose()?,
|
||||
has_unknown_pricing: day.has_unknown_pricing,
|
||||
has_missing_token_data: day.has_missing_token_data,
|
||||
is_boundary_day: day.is_boundary_day,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, ProviderAdminError>>()?;
|
||||
Ok(ProviderUsageStatistics {
|
||||
mode: match statistics.mode {
|
||||
CodexUsageStatisticsMode::Workspace => ProviderUsageStatisticsMode::Workspace,
|
||||
CodexUsageStatisticsMode::Personal => ProviderUsageStatisticsMode::Personal,
|
||||
},
|
||||
cycle: ProviderUsageStatisticsCycle {
|
||||
offset: statistics.cycle.offset,
|
||||
start_at: statistics.cycle.start_at,
|
||||
end_at: statistics.cycle.end_at,
|
||||
window_seconds: statistics.cycle.window_seconds,
|
||||
used_percent: statistics.cycle.used_percent,
|
||||
is_current: statistics.cycle.is_current,
|
||||
can_go_previous: statistics.cycle.can_go_previous,
|
||||
can_go_next: statistics.cycle.can_go_next,
|
||||
},
|
||||
summary: ProviderUsageStatisticsSummary {
|
||||
tokens: project_usage_tokens(statistics.summary.tokens),
|
||||
estimated_cost: statistics
|
||||
.summary
|
||||
.estimated_cost
|
||||
.map(currency_cost)
|
||||
.transpose()?,
|
||||
projected_tokens: statistics.summary.projected_tokens,
|
||||
projected_cost: statistics
|
||||
.summary
|
||||
.projected_cost
|
||||
.map(currency_cost)
|
||||
.transpose()?,
|
||||
day_count: statistics.summary.day_count,
|
||||
has_unknown_pricing: statistics.summary.has_unknown_pricing,
|
||||
has_missing_token_data: statistics.summary.has_missing_token_data,
|
||||
},
|
||||
models,
|
||||
daily,
|
||||
})
|
||||
}
|
||||
|
||||
const fn project_usage_tokens(tokens: CodexUsageStatisticsTokens) -> ProviderUsageStatisticsTokens {
|
||||
ProviderUsageStatisticsTokens {
|
||||
uncached_input: tokens.uncached_input,
|
||||
cached_input: tokens.cached_input,
|
||||
output: tokens.output,
|
||||
total: tokens.total,
|
||||
}
|
||||
}
|
||||
|
||||
const fn project_usage_service_tier(
|
||||
service_tier: CodexUsageStatisticsServiceTier,
|
||||
) -> ProviderUsageStatisticsServiceTier {
|
||||
match service_tier {
|
||||
CodexUsageStatisticsServiceTier::Standard => ProviderUsageStatisticsServiceTier::Standard,
|
||||
CodexUsageStatisticsServiceTier::Fast => ProviderUsageStatisticsServiceTier::Fast,
|
||||
}
|
||||
}
|
||||
|
||||
fn project_reset_credits(
|
||||
credits: crate::transport::CodexRateLimitResetCredits,
|
||||
) -> ProviderResetCredits {
|
||||
ProviderResetCredits {
|
||||
available_count: credits.available_count,
|
||||
credits: credits
|
||||
.credits
|
||||
.into_iter()
|
||||
.map(project_reset_credit)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn project_reset_credit(
|
||||
credit: crate::transport::CodexRateLimitResetCredit,
|
||||
) -> ProviderResetCredit {
|
||||
ProviderResetCredit {
|
||||
id: credit.id,
|
||||
status: credit.status,
|
||||
title: credit.title,
|
||||
expires_at: credit.expires_at,
|
||||
reset_type: credit.reset_type,
|
||||
}
|
||||
}
|
||||
|
||||
/// `secondary_window` 有时只是上游的空占位:没有时长、重置时间,也没有用量。
|
||||
/// 它不能被可靠翻译成 5 小时或周额度,因此不应占用账号面板;带有实际事实的
|
||||
/// 次级窗口(时长、重置、非零用量或触顶)仍完整保留。
|
||||
@@ -800,23 +963,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 => {
|
||||
@@ -827,29 +988,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 {
|
||||
@@ -1275,31 +1414,37 @@ fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError
|
||||
|
||||
fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdminError {
|
||||
use CodexCredentialAdminError as Error;
|
||||
provider_admin_error(match error {
|
||||
let upstream_message = error.upstream_message().map(ToOwned::to_owned);
|
||||
let kind = match error {
|
||||
Error::InvalidInput
|
||||
| Error::InvalidCredential
|
||||
| Error::MissingRefreshToken
|
||||
| Error::RefreshRejected
|
||||
| Error::AccountBanned => ProviderAdminErrorKind::Invalid,
|
||||
| Error::RefreshRejected { .. }
|
||||
| Error::AccountBanned { .. } => ProviderAdminErrorKind::Invalid,
|
||||
Error::NotFound => ProviderAdminErrorKind::NotFound,
|
||||
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => {
|
||||
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous { .. } => {
|
||||
ProviderAdminErrorKind::Conflict
|
||||
}
|
||||
Error::RefreshUnavailable => ProviderAdminErrorKind::Unavailable,
|
||||
})
|
||||
};
|
||||
let error = provider_admin_error(kind);
|
||||
match upstream_message {
|
||||
Some(message) => error.with_message(message),
|
||||
None => error,
|
||||
}
|
||||
}
|
||||
|
||||
const fn credential_admin_error_code(error: CodexCredentialAdminError) -> &'static str {
|
||||
const fn credential_admin_error_code(error: &CodexCredentialAdminError) -> &'static str {
|
||||
match error {
|
||||
CodexCredentialAdminError::InvalidInput => "invalid_input",
|
||||
CodexCredentialAdminError::InvalidCredential => "invalid_credential",
|
||||
CodexCredentialAdminError::NotFound => "not_found",
|
||||
CodexCredentialAdminError::MissingRefreshToken => "missing_refresh_token",
|
||||
CodexCredentialAdminError::RefreshLeaseUnavailable => "refresh_lease_unavailable",
|
||||
CodexCredentialAdminError::RefreshRejected => "refresh_rejected",
|
||||
CodexCredentialAdminError::AccountBanned => "account_banned",
|
||||
CodexCredentialAdminError::RefreshRejected { .. } => "refresh_rejected",
|
||||
CodexCredentialAdminError::AccountBanned { .. } => "account_banned",
|
||||
CodexCredentialAdminError::RefreshUnavailable => "refresh_unavailable",
|
||||
CodexCredentialAdminError::RefreshAmbiguous => "refresh_ambiguous",
|
||||
CodexCredentialAdminError::RefreshAmbiguous { .. } => "refresh_ambiguous",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1310,6 +1455,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",
|
||||
}
|
||||
}
|
||||
@@ -1365,6 +1512,97 @@ fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
|
||||
})
|
||||
}
|
||||
|
||||
fn map_usage_statistics_error(error: CodexUsageStatisticsError) -> ProviderAdminError {
|
||||
use CodexUsageStatisticsError as Error;
|
||||
|
||||
match error {
|
||||
Error::InvalidRequest | Error::InvalidCredentialData => {
|
||||
provider_admin_error(ProviderAdminErrorKind::Invalid)
|
||||
}
|
||||
Error::NotFound => provider_admin_error(ProviderAdminErrorKind::NotFound),
|
||||
Error::Store { .. } | Error::TransportUnavailable => {
|
||||
provider_admin_error(ProviderAdminErrorKind::Unavailable)
|
||||
}
|
||||
Error::QuotaWindowUnavailable => provider_admin_error(ProviderAdminErrorKind::Conflict)
|
||||
.with_message("OpenAI did not expose a resettable quota window for usage statistics"),
|
||||
Error::CredentialRefreshRequired { upstream_body } => {
|
||||
let mut error = provider_admin_error(ProviderAdminErrorKind::CredentialRefreshRequired);
|
||||
if let Some(body) = upstream_body {
|
||||
error = error.with_message(format!(
|
||||
"OpenAI usage-statistics upstream returned HTTP 401: {}",
|
||||
bounded_upstream_body(&body)
|
||||
));
|
||||
}
|
||||
error
|
||||
}
|
||||
Error::Upstream {
|
||||
status,
|
||||
body,
|
||||
retry_after_seconds,
|
||||
} => {
|
||||
let retry_after = retry_after_seconds
|
||||
.map(|seconds| format!("; retry-after={seconds}s"))
|
||||
.unwrap_or_default();
|
||||
provider_admin_error(ProviderAdminErrorKind::BadGateway).with_message(format!(
|
||||
"OpenAI usage-statistics upstream returned HTTP {status}{retry_after}: {}",
|
||||
bounded_upstream_body(&body)
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_reset_credits_error(error: CodexResetCreditsError) -> ProviderAdminError {
|
||||
use CodexResetCreditsError as Error;
|
||||
|
||||
match error {
|
||||
Error::InvalidCredentialData => provider_admin_error(ProviderAdminErrorKind::Invalid),
|
||||
Error::NotFound => provider_admin_error(ProviderAdminErrorKind::NotFound),
|
||||
Error::Store { .. } | Error::TransportUnavailable => {
|
||||
provider_admin_error(ProviderAdminErrorKind::Unavailable)
|
||||
}
|
||||
Error::CredentialRefreshRequired { upstream_body } => {
|
||||
let mut error =
|
||||
provider_admin_error(ProviderAdminErrorKind::CredentialRefreshRequired);
|
||||
if let Some(body) = upstream_body {
|
||||
error = error.with_message(format!(
|
||||
"OpenAI reset-credit upstream returned HTTP 401: {}",
|
||||
bounded_upstream_body(&body)
|
||||
));
|
||||
}
|
||||
error
|
||||
}
|
||||
Error::Upstream {
|
||||
status,
|
||||
body,
|
||||
retry_after_seconds,
|
||||
} => {
|
||||
let retry_after = retry_after_seconds
|
||||
.map(|seconds| format!("; retry-after={seconds}s"))
|
||||
.unwrap_or_default();
|
||||
provider_admin_error(ProviderAdminErrorKind::BadGateway).with_message(format!(
|
||||
"OpenAI reset-credit upstream returned HTTP {status}{retry_after}: {}",
|
||||
bounded_upstream_body(&body)
|
||||
))
|
||||
}
|
||||
Error::ConsumeResultUnknown => provider_admin_error(ProviderAdminErrorKind::BadGateway)
|
||||
.with_message(
|
||||
"OpenAI reset-credit consume result is unknown; refresh the credit list before retrying",
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn bounded_upstream_body(body: &str) -> String {
|
||||
const MAX_ADMIN_UPSTREAM_BODY_BYTES: usize = 8 * 1024;
|
||||
if body.len() <= MAX_ADMIN_UPSTREAM_BODY_BYTES {
|
||||
return body.to_owned();
|
||||
}
|
||||
let mut end = MAX_ADMIN_UPSTREAM_BODY_BYTES;
|
||||
while !body.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
body[..end].to_owned()
|
||||
}
|
||||
|
||||
fn map_catalog_error(error: CodexCredentialCatalogError) -> ProviderAdminError {
|
||||
use CodexCredentialCatalogError as Error;
|
||||
provider_admin_error(match error {
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -153,8 +155,6 @@ impl CodexApiConfig {
|
||||
pub struct CodexWebSocketPoolSettings {
|
||||
pub enabled: bool,
|
||||
pub max_age_ms: u64,
|
||||
pub max_per_account: usize,
|
||||
pub max_total: usize,
|
||||
pub max_connecting: usize,
|
||||
pub initial_event_timeout_ms: u64,
|
||||
}
|
||||
@@ -164,8 +164,6 @@ impl Default for CodexWebSocketPoolSettings {
|
||||
Self {
|
||||
enabled: true,
|
||||
max_age_ms: 55 * 60 * 1000,
|
||||
max_per_account: 8,
|
||||
max_total: 64,
|
||||
max_connecting: 8,
|
||||
initial_event_timeout_ms: 20_000,
|
||||
}
|
||||
@@ -174,12 +172,7 @@ impl Default for CodexWebSocketPoolSettings {
|
||||
|
||||
impl CodexWebSocketPoolSettings {
|
||||
fn validate(&self) -> Result<(), OpenAiConfigError> {
|
||||
if self.max_age_ms == 0
|
||||
|| self.max_per_account == 0
|
||||
|| self.max_total == 0
|
||||
|| self.max_connecting == 0
|
||||
|| self.max_connecting > self.max_total
|
||||
{
|
||||
if self.max_age_ms == 0 || self.max_connecting == 0 {
|
||||
return Err(OpenAiConfigError::InvalidField("openai.ws_pool"));
|
||||
}
|
||||
Ok(())
|
||||
@@ -189,8 +182,6 @@ impl CodexWebSocketPoolSettings {
|
||||
CodexWebSocketPoolConfig {
|
||||
enabled: self.enabled,
|
||||
max_age: Duration::from_millis(self.max_age_ms),
|
||||
max_per_account: self.max_per_account,
|
||||
max_total: self.max_total,
|
||||
max_connecting: self.max_connecting,
|
||||
initial_event_timeout: (self.initial_event_timeout_ms != 0)
|
||||
.then(|| Duration::from_millis(self.initial_event_timeout_ms)),
|
||||
|
||||
@@ -376,7 +376,7 @@ impl fmt::Debug for PreparedCodexCredentialRotationGuard {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum CodexCredentialAdminError {
|
||||
#[error("Codex account input is invalid")]
|
||||
InvalidInput,
|
||||
@@ -389,13 +389,30 @@ pub enum CodexCredentialAdminError {
|
||||
#[error("Codex refresh lease is unavailable")]
|
||||
RefreshLeaseUnavailable,
|
||||
#[error("Codex refresh token was rejected")]
|
||||
RefreshRejected,
|
||||
RefreshRejected { message: Option<String> },
|
||||
#[error("Codex account is banned")]
|
||||
AccountBanned,
|
||||
AccountBanned { message: Option<String> },
|
||||
#[error("Codex refresh service is unavailable")]
|
||||
RefreshUnavailable,
|
||||
#[error("Codex refresh send state is ambiguous")]
|
||||
RefreshAmbiguous,
|
||||
RefreshAmbiguous { message: Option<String> },
|
||||
}
|
||||
|
||||
impl CodexCredentialAdminError {
|
||||
#[must_use]
|
||||
pub fn upstream_message(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::RefreshRejected { message }
|
||||
| Self::AccountBanned { message }
|
||||
| Self::RefreshAmbiguous { message } => message.as_deref(),
|
||||
Self::InvalidInput
|
||||
| Self::InvalidCredential
|
||||
| Self::NotFound
|
||||
| Self::MissingRefreshToken
|
||||
| Self::RefreshLeaseUnavailable
|
||||
| Self::RefreshUnavailable => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 无状态的 Codex Admin command preparer。
|
||||
@@ -800,8 +817,15 @@ impl CodexCredentialAdminService {
|
||||
.refresher
|
||||
.refresh(refresh_token.expose_secret())
|
||||
.await
|
||||
.inspect_err(|error| log_manual_refresh_failure(&account_id, error))
|
||||
.map_err(map_refresh_failure)?;
|
||||
let access_token_rotated = tokens.access_token.is_some();
|
||||
let access_token_expires_at = match tokens.access_token.as_deref() {
|
||||
Some(access_token) => parse_access_token_expiration(access_token),
|
||||
None => current
|
||||
.account
|
||||
.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from),
|
||||
};
|
||||
let secret = CodexOAuthSecret {
|
||||
access_token: tokens
|
||||
.access_token
|
||||
@@ -821,17 +845,6 @@ impl CodexCredentialAdminService {
|
||||
Some(account_id.as_str()),
|
||||
&secret,
|
||||
);
|
||||
let access_token_expires_at = if access_token_rotated {
|
||||
tokens
|
||||
.expires_in
|
||||
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
|
||||
.map(DateTime::<Utc>::from)
|
||||
} else {
|
||||
current
|
||||
.account
|
||||
.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from)
|
||||
};
|
||||
// 正常预刷新由 worker 读取当时的 runtime policy 动态判断;这里仅清除
|
||||
// 之前瞬态失败留下的 retry-not-before。
|
||||
let next_refresh_at = None;
|
||||
@@ -998,6 +1011,7 @@ impl CodexCredentialAdminService {
|
||||
let access_token = tokens
|
||||
.access_token
|
||||
.ok_or(CodexCredentialAdminError::InvalidCredential)?;
|
||||
let access_token_expires_at = parse_access_token_expiration(&access_token);
|
||||
let secret = CodexOAuthSecret {
|
||||
access_token: SecretString::from(access_token),
|
||||
// RT 未轮换时仍保留导入提供的 RT,保证本次补全不会丢失后续刷新能力。
|
||||
@@ -1012,10 +1026,6 @@ impl CodexCredentialAdminService {
|
||||
Some(account_id.as_str()),
|
||||
&secret,
|
||||
);
|
||||
let access_token_expires_at = tokens
|
||||
.expires_in
|
||||
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
|
||||
.map(DateTime::<Utc>::from);
|
||||
Ok((secret, access_token_expires_at))
|
||||
}
|
||||
|
||||
@@ -1133,15 +1143,35 @@ fn china_rfc3339(value: DateTime<Utc>) -> String {
|
||||
.to_rfc3339()
|
||||
}
|
||||
|
||||
const fn map_refresh_failure(error: RefreshFailure) -> CodexCredentialAdminError {
|
||||
fn map_refresh_failure(error: RefreshFailure) -> CodexCredentialAdminError {
|
||||
match error {
|
||||
RefreshFailure::InvalidGrant => CodexCredentialAdminError::RefreshRejected,
|
||||
RefreshFailure::Banned => CodexCredentialAdminError::AccountBanned,
|
||||
RefreshFailure::RetryableTransport => CodexCredentialAdminError::RefreshUnavailable,
|
||||
RefreshFailure::Transport => CodexCredentialAdminError::RefreshAmbiguous,
|
||||
RefreshFailure::InvalidGrant { message, .. } => {
|
||||
CodexCredentialAdminError::RefreshRejected { message }
|
||||
}
|
||||
RefreshFailure::Banned { message, .. } => {
|
||||
CodexCredentialAdminError::AccountBanned { message }
|
||||
}
|
||||
RefreshFailure::RetryableTransport { .. } => CodexCredentialAdminError::RefreshUnavailable,
|
||||
RefreshFailure::Transport { message, .. } => {
|
||||
CodexCredentialAdminError::RefreshAmbiguous { message }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn log_manual_refresh_failure(account_id: &ProviderAccountId, error: &RefreshFailure) {
|
||||
let upstream = error.upstream();
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
failure_class = error.classification(),
|
||||
upstream_message = ?error.message(),
|
||||
upstream_status = ?upstream.map(super::token_client::RefreshUpstreamFailure::status),
|
||||
upstream_code = ?upstream.and_then(super::token_client::RefreshUpstreamFailure::code),
|
||||
upstream_type = ?upstream.and_then(super::token_client::RefreshUpstreamFailure::error_type),
|
||||
upstream_body = ?upstream.map(super::token_client::RefreshUpstreamFailure::body),
|
||||
"OpenAI OAuth manual refresh failed"
|
||||
);
|
||||
}
|
||||
|
||||
fn parse_import_document(
|
||||
payload: &Value,
|
||||
) -> Result<Vec<ParsedCodexImportAccount>, CodexCredentialAdminError> {
|
||||
@@ -1258,9 +1288,9 @@ fn parse_oauth_import_tokens(
|
||||
Value::Object(object) => {
|
||||
for (key, value) in object {
|
||||
let token = match key.as_str() {
|
||||
"accessToken" => &mut access_token,
|
||||
"refreshToken" => &mut refresh_token,
|
||||
"idToken" => &mut id_token,
|
||||
"accessToken" | "access_token" => &mut access_token,
|
||||
"refreshToken" | "refresh_token" => &mut refresh_token,
|
||||
"idToken" | "id_token" => &mut id_token,
|
||||
_ => {
|
||||
pending.push(value);
|
||||
continue;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! OpenAI 会话锚点到 Store 不透明亲和键的单向派生。
|
||||
//! OpenAI 根会话锚点到 Store 不透明账号亲和键的单向派生。
|
||||
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::ProviderSessionAffinityKey;
|
||||
@@ -11,34 +11,8 @@ pub(crate) fn derive_codex_session_affinity_key(
|
||||
request: &CodexResponsesRequest,
|
||||
client_api_key_id: &ClientApiKeyId,
|
||||
) -> Option<ProviderSessionAffinityKey> {
|
||||
let session_key = if let Some(conversation_id) = request
|
||||
.local_conversation_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
if let Ok(key) = ProviderSessionAffinityKey::try_new(conversation_id.to_owned()) {
|
||||
Some(key)
|
||||
} else {
|
||||
opaque_affinity_key("local-conversation", conversation_id)
|
||||
}
|
||||
} else {
|
||||
let (domain, value) = derive_conversation_anchor(request)?;
|
||||
opaque_affinity_key(domain, &value)
|
||||
};
|
||||
|
||||
let session_discriminator = request.subagent_kind();
|
||||
let session_key =
|
||||
session_key.and_then(|session_key| match session_discriminator.as_deref() {
|
||||
None => Some(session_key),
|
||||
Some(discriminator) => {
|
||||
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
|
||||
opaque_affinity_key(
|
||||
"subagent-session",
|
||||
&format!("{discriminator}\0{}", session_key.expose_to_store()),
|
||||
)
|
||||
}
|
||||
})?;
|
||||
let (domain, value) = derive_account_affinity_anchor(request)?;
|
||||
let session_key = opaque_affinity_key(domain, &value)?;
|
||||
opaque_affinity_key(
|
||||
"client-session",
|
||||
&format!(
|
||||
@@ -49,6 +23,33 @@ pub(crate) fn derive_codex_session_affinity_key(
|
||||
)
|
||||
}
|
||||
|
||||
/// 账号选择优先使用根会话事实;`local_conversation_id` 仍可按 child thread 隔离
|
||||
/// WebSocket 与 continuation 状态,但 child 身份不再拆分账号首选项。账号不可调度时
|
||||
/// selector 仍可按既有策略切换,并由 request scope 清理账号绑定状态。
|
||||
fn derive_account_affinity_anchor(
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Option<(&'static str, String)> {
|
||||
non_empty(request.client_session_id.as_deref())
|
||||
.map(|value| ("root-session", value.to_owned()))
|
||||
.or_else(|| {
|
||||
non_empty(request.client_conversation_id.as_deref())
|
||||
.map(|value| ("root-conversation", value.to_owned()))
|
||||
})
|
||||
.or_else(|| {
|
||||
request
|
||||
.explicit_prompt_cache_key
|
||||
.then(|| request.prompt_cache_key())
|
||||
.flatten()
|
||||
.and_then(|value| non_empty(Some(value)))
|
||||
.map(|value| ("root-prompt-cache", value.to_owned()))
|
||||
})
|
||||
.or_else(|| {
|
||||
non_empty(request.local_conversation_id.as_deref())
|
||||
.map(|value| ("local-conversation", value.to_owned()))
|
||||
})
|
||||
.or_else(|| derive_conversation_anchor(request))
|
||||
}
|
||||
|
||||
fn opaque_affinity_key(domain: &str, value: &str) -> Option<ProviderSessionAffinityKey> {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(b"codex-session-affinity-v1\0");
|
||||
|
||||
@@ -19,6 +19,7 @@ pub(crate) use affinity::{
|
||||
derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key,
|
||||
};
|
||||
pub(crate) use oauth::oauth_owner_ref;
|
||||
pub(crate) use types::parse_access_token_expiration;
|
||||
|
||||
pub use agent_identity::{
|
||||
CodexAgentIdentityError, CodexAgentIdentitySecret, CodexAgentIdentityTaskRegistrar,
|
||||
@@ -46,7 +47,10 @@ pub use oauth::{
|
||||
pub use quota::{
|
||||
CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService,
|
||||
CodexQuotaFact, CodexQuotaRefreshPolicy, CodexQuotaSyncSummary, CodexQuotaWindow,
|
||||
CodexQuotaWindowKind, CodexQuotaWindowRole, parse_codex_quota_usage,
|
||||
CodexQuotaWindowKind, CodexQuotaWindowRole, CodexResetCreditsError, CodexUsageStatistics,
|
||||
CodexUsageStatisticsCycle, CodexUsageStatisticsDay, CodexUsageStatisticsError,
|
||||
CodexUsageStatisticsMode, CodexUsageStatisticsModel, CodexUsageStatisticsServiceTier,
|
||||
CodexUsageStatisticsSummary, CodexUsageStatisticsTokens, parse_codex_quota_usage,
|
||||
};
|
||||
pub use refresh::{
|
||||
CodexCredentialRefreshError, CodexCredentialRefreshOutcome, CodexCredentialRefreshService,
|
||||
@@ -57,11 +61,11 @@ pub use security::{
|
||||
CodexCredentialCodec, CodexCredentialDataError, CodexRuntimeAuthentication,
|
||||
CodexRuntimeCredential,
|
||||
};
|
||||
pub(crate) use selector::CodexCyberPolicyScope;
|
||||
pub use selector::{
|
||||
CodexAccountFailure, CodexCredentialLease, CodexCredentialSelector, CredentialSelectionError,
|
||||
SelectCodexCredential,
|
||||
};
|
||||
pub(crate) use selector::{CodexCyberPolicyScope, SelectCodexProviderEndpointCredential};
|
||||
pub use types::{
|
||||
CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY, CODEX_AUTHENTICATION_KIND_OAUTH, CodexAccountProfile,
|
||||
CodexAgentIdentityAuthMode, CodexAgentIdentityCredentialData, CodexCookie,
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
//! Codex Authorization Code + PKCE/OIDC 管理流。
|
||||
|
||||
use std::{
|
||||
fmt,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
use std::{fmt, sync::Arc, time::Duration};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
@@ -35,7 +31,7 @@ use super::token_client::{
|
||||
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
|
||||
OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI,
|
||||
};
|
||||
use super::types::parse_chatgpt_jwt_claims;
|
||||
use super::types::{parse_access_token_expiration, parse_chatgpt_jwt_claims};
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
|
||||
const AUTHORIZATION_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
|
||||
@@ -527,13 +523,12 @@ impl CodexOAuthAdminService {
|
||||
.map(current_oauth_refresh_token)
|
||||
.transpose()?
|
||||
.flatten();
|
||||
let (mut secret, id_token, expires_in) = self
|
||||
let (mut secret, id_token) = self
|
||||
.exchange_pending(&pending, callback_url, fallback_refresh_token)
|
||||
.await?;
|
||||
let mutation = pending.mutation.clone();
|
||||
let access_token_expires_at = expires_in
|
||||
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
|
||||
.map(DateTime::<Utc>::from);
|
||||
let access_token_expires_at =
|
||||
parse_access_token_expiration(secret.access_token.expose_secret());
|
||||
let metadata = parse_chatgpt_jwt_claims(id_token.expose_secret())
|
||||
.map_err(|_| CodexOAuthAdminError::TokenRejected)?;
|
||||
secret.id_token = Some(id_token);
|
||||
@@ -671,14 +666,7 @@ impl CodexOAuthAdminService {
|
||||
pending: &CodexPendingAuthorization,
|
||||
callback_url: &str,
|
||||
fallback_refresh_token: Option<SecretString>,
|
||||
) -> Result<
|
||||
(
|
||||
super::types::CodexOAuthSecret,
|
||||
SecretString,
|
||||
Option<Duration>,
|
||||
),
|
||||
CodexOAuthAdminError,
|
||||
> {
|
||||
) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> {
|
||||
let (code, callback_state) = callback_parts(callback_url)?;
|
||||
if !constant_time_equal(
|
||||
pending.state.expose_secret().as_bytes(),
|
||||
@@ -709,7 +697,7 @@ impl CodexOAuthAdminService {
|
||||
.as_ref()
|
||||
.map(ExposeSecret::expose_secret),
|
||||
);
|
||||
Ok((secret, tokens.id_token, tokens.expires_in))
|
||||
Ok((secret, tokens.id_token))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,26 +1,34 @@
|
||||
//! Codex quota 服务编排:权威额度事实、展示投影、主动/被动同步与 429 冷却。
|
||||
//!
|
||||
//! - [`document`]:上游 `/usage` 输入到多桶 map 的单向规范化。
|
||||
//! - [`snapshot`]:快照/窗口解析、聚合、O(1) 滚动与调度信号。
|
||||
//! - [`evidence`]:额度接口错误到凭据/额度事实的分类。
|
||||
|
||||
mod document;
|
||||
pub(crate) mod evidence;
|
||||
pub(crate) mod snapshot;
|
||||
mod usage_statistics;
|
||||
|
||||
pub use snapshot::{
|
||||
CodexAccountQuotaSnapshot, CodexQuotaFact, CodexQuotaWindow, CodexQuotaWindowKind,
|
||||
CodexQuotaWindowRole, parse_codex_quota_usage,
|
||||
};
|
||||
pub use usage_statistics::{
|
||||
CodexUsageStatistics, CodexUsageStatisticsCycle, CodexUsageStatisticsDay,
|
||||
CodexUsageStatisticsMode, CodexUsageStatisticsModel, CodexUsageStatisticsServiceTier,
|
||||
CodexUsageStatisticsSummary, CodexUsageStatisticsTokens,
|
||||
};
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
|
||||
use chrono::Utc;
|
||||
use chrono::{DateTime, FixedOffset, Utc};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountQuotaSignals, CredentialRevision, CredentialState,
|
||||
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch,
|
||||
QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::provider_ports::{ProviderCooldown, ProviderCooldownPort};
|
||||
use gateway_protocol::openai::events::{
|
||||
@@ -34,10 +42,19 @@ use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext};
|
||||
use crate::transport::usage_statistics::CodexDailyUsageReport;
|
||||
use crate::transport::{
|
||||
CodexBackendClient, CodexClientError, CodexRateLimitResetCredits,
|
||||
CodexRateLimitResetCreditsConsumeResult, CodexRequestContext,
|
||||
};
|
||||
|
||||
use super::agent_identity::{CodexAgentIdentityTaskService, PreparedCodexRuntimeCredential};
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
use super::types::CODEX_AUTHENTICATION_KIND_OAUTH;
|
||||
use document::{
|
||||
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId,
|
||||
canonicalize_rate_limit_document,
|
||||
};
|
||||
use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure};
|
||||
use snapshot::{
|
||||
parse_account_quota_snapshot, quota_projection_ttl, quota_snapshot_from_observation,
|
||||
@@ -112,6 +129,105 @@ pub enum CodexCredentialQuotaError {
|
||||
Upstream { detail: String },
|
||||
}
|
||||
|
||||
/// 主动额度重置卡查询/消费失败。
|
||||
#[derive(Error)]
|
||||
pub enum CodexResetCreditsError {
|
||||
#[error("Codex reset-credit credential data is invalid")]
|
||||
InvalidCredentialData,
|
||||
#[error("Codex OAuth access token must be refreshed before using reset credits")]
|
||||
CredentialRefreshRequired { upstream_body: Option<String> },
|
||||
#[error("Codex reset-credit account was not found")]
|
||||
NotFound,
|
||||
#[error("provider account store is unavailable: {detail}")]
|
||||
Store { detail: String },
|
||||
#[error("Codex reset-credit upstream returned HTTP {status}")]
|
||||
Upstream {
|
||||
status: u16,
|
||||
body: String,
|
||||
retry_after_seconds: Option<u64>,
|
||||
},
|
||||
#[error("Codex reset-credit query transport is unavailable")]
|
||||
TransportUnavailable,
|
||||
#[error("Codex reset-credit consume result is unknown")]
|
||||
ConsumeResultUnknown,
|
||||
}
|
||||
|
||||
/// 官方每日用量统计查询失败。
|
||||
#[derive(Error)]
|
||||
pub enum CodexUsageStatisticsError {
|
||||
#[error("Codex usage-statistics request is invalid")]
|
||||
InvalidRequest,
|
||||
#[error("Codex usage-statistics credential data is invalid")]
|
||||
InvalidCredentialData,
|
||||
#[error("Codex OAuth access token must be refreshed before querying usage statistics")]
|
||||
CredentialRefreshRequired { upstream_body: Option<String> },
|
||||
#[error("Codex usage-statistics account was not found")]
|
||||
NotFound,
|
||||
#[error("provider account store is unavailable: {detail}")]
|
||||
Store { detail: String },
|
||||
#[error("Codex usage-statistics quota window is unavailable")]
|
||||
QuotaWindowUnavailable,
|
||||
#[error("Codex usage-statistics upstream returned HTTP {status}")]
|
||||
Upstream {
|
||||
status: u16,
|
||||
body: String,
|
||||
retry_after_seconds: Option<u64>,
|
||||
},
|
||||
#[error("Codex usage-statistics query transport is unavailable")]
|
||||
TransportUnavailable,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for CodexUsageStatisticsError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::InvalidRequest => formatter.write_str("InvalidRequest"),
|
||||
Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"),
|
||||
Self::CredentialRefreshRequired { .. } => {
|
||||
formatter.write_str("CredentialRefreshRequired { upstream_body: <redacted> }")
|
||||
}
|
||||
Self::NotFound => formatter.write_str("NotFound"),
|
||||
Self::Store { .. } => formatter.write_str("Store { detail: <redacted> }"),
|
||||
Self::QuotaWindowUnavailable => formatter.write_str("QuotaWindowUnavailable"),
|
||||
Self::Upstream {
|
||||
status,
|
||||
retry_after_seconds,
|
||||
..
|
||||
} => formatter
|
||||
.debug_struct("Upstream")
|
||||
.field("status", status)
|
||||
.field("body", &"<redacted>")
|
||||
.field("retry_after_seconds", retry_after_seconds)
|
||||
.finish(),
|
||||
Self::TransportUnavailable => formatter.write_str("TransportUnavailable"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for CodexResetCreditsError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::InvalidCredentialData => formatter.write_str("InvalidCredentialData"),
|
||||
Self::CredentialRefreshRequired { .. } => {
|
||||
formatter.write_str("CredentialRefreshRequired { upstream_body: <redacted> }")
|
||||
}
|
||||
Self::NotFound => formatter.write_str("NotFound"),
|
||||
Self::Store { .. } => formatter.write_str("Store { detail: <redacted> }"),
|
||||
Self::Upstream {
|
||||
status,
|
||||
retry_after_seconds,
|
||||
..
|
||||
} => formatter
|
||||
.debug_struct("Upstream")
|
||||
.field("status", status)
|
||||
.field("body", &"<redacted>")
|
||||
.field("retry_after_seconds", retry_after_seconds)
|
||||
.finish(),
|
||||
Self::TransportUnavailable => formatter.write_str("TransportUnavailable"),
|
||||
Self::ConsumeResultUnknown => formatter.write_str("ConsumeResultUnknown"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<gateway_core::error::StoreError> for CodexCredentialQuotaError {
|
||||
fn from(error: gateway_core::error::StoreError) -> Self {
|
||||
Self::Store {
|
||||
@@ -129,6 +245,7 @@ pub struct CodexCredentialQuotaService {
|
||||
agent_identity: Arc<CodexAgentIdentityTaskService>,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
scheduling: CodexQuotaSchedulingProjection,
|
||||
reset_consume_locks: Mutex<HashMap<ProviderAccountId, Arc<Mutex<()>>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -183,6 +300,31 @@ enum CodexQuotaFetchError {
|
||||
},
|
||||
}
|
||||
|
||||
enum ResetCreditAttemptError {
|
||||
InvalidCredential,
|
||||
Upstream(CodexClientError),
|
||||
}
|
||||
|
||||
enum DailyUsageAttemptError {
|
||||
InvalidCredential,
|
||||
Upstream(CodexClientError),
|
||||
}
|
||||
|
||||
enum DailyUsageFetchError {
|
||||
InvalidCredential,
|
||||
Recovery,
|
||||
Upstream {
|
||||
prepared: Box<PreparedCodexRuntimeCredential>,
|
||||
error: CodexClientError,
|
||||
},
|
||||
}
|
||||
|
||||
struct FetchedDailyUsage {
|
||||
mode: CodexUsageStatisticsMode,
|
||||
model_breakdown: Value,
|
||||
daily_totals: Option<Value>,
|
||||
}
|
||||
|
||||
impl CodexQuotaSchedulingProjection {
|
||||
fn invalidate(&self, account_ids: &[ProviderAccountId]) {
|
||||
let mut state = self
|
||||
@@ -407,9 +549,237 @@ impl CodexCredentialQuotaService {
|
||||
agent_identity,
|
||||
cooldowns,
|
||||
scheduling: CodexQuotaSchedulingProjection::default(),
|
||||
reset_consume_locks: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 查询当前账号由 Codex Desktop 暴露的主动额度重置卡。
|
||||
pub async fn list_reset_credits(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CodexRateLimitResetCredits, CodexResetCreditsError> {
|
||||
let account = self.reset_credit_account(account_id).await?;
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
let mut prepared = self
|
||||
.agent_identity
|
||||
.prepare(&account)
|
||||
.await
|
||||
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?;
|
||||
let request_id = format!("reset_credits_{}", Uuid::now_v7().simple());
|
||||
let mut result = list_reset_credits_once(&client, &prepared, &request_id).await;
|
||||
if let Err(ResetCreditAttemptError::Upstream(error)) = &result
|
||||
&& let Some(recovered) = self
|
||||
.agent_identity
|
||||
.recover_after_rejected_task(
|
||||
prepared.account.id(),
|
||||
&prepared.credential.authentication,
|
||||
error,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?
|
||||
{
|
||||
prepared = recovered;
|
||||
result = list_reset_credits_once(&client, &prepared, &request_id).await;
|
||||
}
|
||||
result.map_err(|error| map_reset_credit_attempt_error(error, &prepared, false))
|
||||
}
|
||||
|
||||
/// 查询当前账号的官方每日用量报表,并把报表范围锚定到最接近七天的额度窗口。
|
||||
pub async fn usage_statistics(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
cycle_offset: i8,
|
||||
utc_offset_minutes: i16,
|
||||
) -> Result<CodexUsageStatistics, CodexUsageStatisticsError> {
|
||||
const MAX_CYCLE_OFFSET: u8 = 8;
|
||||
const MAX_UTC_OFFSET_MINUTES: i16 = 14 * 60;
|
||||
const TARGET_WINDOW_SECONDS: u64 = 7 * 24 * 60 * 60;
|
||||
|
||||
if !(-(MAX_CYCLE_OFFSET as i8)..=0).contains(&cycle_offset)
|
||||
|| !(-MAX_UTC_OFFSET_MINUTES..=MAX_UTC_OFFSET_MINUTES).contains(&utc_offset_minutes)
|
||||
{
|
||||
return Err(CodexUsageStatisticsError::InvalidRequest);
|
||||
}
|
||||
let account = self
|
||||
.store
|
||||
.get_account(account_id)
|
||||
.await
|
||||
.map_err(|error| CodexUsageStatisticsError::Store {
|
||||
detail: error.to_string(),
|
||||
})?
|
||||
.filter(|account| account.provider().as_str() == "openai")
|
||||
.ok_or(CodexUsageStatisticsError::NotFound)?;
|
||||
let observed_at = SystemTime::now();
|
||||
if !access_token_is_current(&account, observed_at) {
|
||||
return Err(CodexUsageStatisticsError::CredentialRefreshRequired {
|
||||
upstream_body: None,
|
||||
});
|
||||
}
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
let FetchedCodexQuota { account, value } = self
|
||||
.fetch_usage_with_recovery(&client, &account)
|
||||
.await
|
||||
.map_err(map_statistics_quota_fetch_error)?;
|
||||
let snapshot = parse_account_quota_snapshot(
|
||||
account.id().clone(),
|
||||
account.revision(),
|
||||
observed_at,
|
||||
&value,
|
||||
)
|
||||
.map_err(|_| CodexUsageStatisticsError::InvalidCredentialData)?;
|
||||
let window = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.filter(|window| window.is_account_wide())
|
||||
.filter_map(|window| {
|
||||
Some((
|
||||
window,
|
||||
window.window_seconds()?.abs_diff(TARGET_WINDOW_SECONDS),
|
||||
))
|
||||
})
|
||||
.min_by_key(|(_, distance)| *distance)
|
||||
.map(|(window, _)| window)
|
||||
.ok_or(CodexUsageStatisticsError::QuotaWindowUnavailable)?;
|
||||
let window_seconds = window
|
||||
.window_seconds()
|
||||
.filter(|seconds| *seconds > 0)
|
||||
.ok_or(CodexUsageStatisticsError::QuotaWindowUnavailable)?;
|
||||
let end_at = window
|
||||
.reset_at()
|
||||
.ok_or(CodexUsageStatisticsError::QuotaWindowUnavailable)?;
|
||||
let window_delta = chrono::TimeDelta::seconds(
|
||||
i64::try_from(window_seconds)
|
||||
.map_err(|_| CodexUsageStatisticsError::QuotaWindowUnavailable)?,
|
||||
);
|
||||
let current_start_at = end_at - window_delta;
|
||||
let timezone = FixedOffset::east_opt(i32::from(utc_offset_minutes) * 60)
|
||||
.ok_or(CodexUsageStatisticsError::InvalidRequest)?;
|
||||
let now = DateTime::<Utc>::from(observed_at);
|
||||
let period = usage_statistics::statistics_period(
|
||||
current_start_at,
|
||||
end_at,
|
||||
window_seconds,
|
||||
cycle_offset,
|
||||
timezone,
|
||||
now,
|
||||
)
|
||||
.ok_or(CodexUsageStatisticsError::InvalidRequest)?;
|
||||
let fetched = self
|
||||
.fetch_daily_usage_with_recovery(
|
||||
&client,
|
||||
&account,
|
||||
period.query_start_date,
|
||||
period.query_end_date,
|
||||
)
|
||||
.await
|
||||
.map_err(map_daily_usage_fetch_error)?;
|
||||
|
||||
Ok(usage_statistics::build_usage_statistics(
|
||||
usage_statistics::BuildUsageStatistics {
|
||||
mode: fetched.mode,
|
||||
period,
|
||||
timezone,
|
||||
current_used_percent: window.used_percent(),
|
||||
now,
|
||||
model_breakdown: &fetched.model_breakdown,
|
||||
daily_totals: fetched.daily_totals.as_ref(),
|
||||
max_cycle_offset: MAX_CYCLE_OFFSET,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
/// 消费一张主动额度重置卡。相同账号在本进程内串行,且不做传输重试。
|
||||
pub async fn consume_reset_credit(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
credit_id: Option<&str>,
|
||||
redeem_request_id: Uuid,
|
||||
) -> Result<CodexRateLimitResetCreditsConsumeResult, CodexResetCreditsError> {
|
||||
let lock = {
|
||||
let mut locks = self.reset_consume_locks.lock().await;
|
||||
Arc::clone(
|
||||
locks
|
||||
.entry(account_id.clone())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(()))),
|
||||
)
|
||||
};
|
||||
let _guard = lock.lock().await;
|
||||
let account = self.reset_credit_account(account_id).await?;
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
let mut prepared = self
|
||||
.agent_identity
|
||||
.prepare(&account)
|
||||
.await
|
||||
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?;
|
||||
let request_id = format!("reset_credit_consume_{}", Uuid::now_v7().simple());
|
||||
let mut result = consume_reset_credit_once(
|
||||
&client,
|
||||
&prepared,
|
||||
&request_id,
|
||||
credit_id,
|
||||
redeem_request_id,
|
||||
)
|
||||
.await;
|
||||
// Agent Identity task recovery is the only automatic retry below. It keeps the exact
|
||||
// official credit selection and redeem_request_id; generic HTTP/transport failures are
|
||||
// never retried because the debit result may already be committed upstream.
|
||||
if let Err(ResetCreditAttemptError::Upstream(error)) = &result
|
||||
&& let Some(recovered) = self
|
||||
.agent_identity
|
||||
.recover_after_rejected_task(
|
||||
prepared.account.id(),
|
||||
&prepared.credential.authentication,
|
||||
error,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CodexResetCreditsError::InvalidCredentialData)?
|
||||
{
|
||||
prepared = recovered;
|
||||
result = consume_reset_credit_once(
|
||||
&client,
|
||||
&prepared,
|
||||
&request_id,
|
||||
credit_id,
|
||||
redeem_request_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
result.map_err(|error| map_reset_credit_attempt_error(error, &prepared, true))
|
||||
}
|
||||
|
||||
async fn reset_credit_account(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderAccount, CodexResetCreditsError> {
|
||||
let account = self
|
||||
.store
|
||||
.get_account(account_id)
|
||||
.await
|
||||
.map_err(|error| CodexResetCreditsError::Store {
|
||||
detail: error.to_string(),
|
||||
})?
|
||||
.filter(|account| account.provider().as_str() == "openai")
|
||||
.ok_or(CodexResetCreditsError::NotFound)?;
|
||||
if !access_token_is_current(&account, SystemTime::now()) {
|
||||
return Err(CodexResetCreditsError::CredentialRefreshRequired {
|
||||
upstream_body: None,
|
||||
});
|
||||
}
|
||||
Ok(account)
|
||||
}
|
||||
|
||||
/// 真实推理错误只更新额度访问事实,不伪造 Provider JSON 或展示百分比。
|
||||
pub(crate) async fn record_confirmed_exhaustion(
|
||||
&self,
|
||||
@@ -646,7 +1016,25 @@ impl CodexCredentialQuotaService {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !quota_refresh_can_update(account, &snapshot, previous.as_ref()) {
|
||||
if account.quota().is_exhausted()
|
||||
&& !quota_refresh_allows_recovery(account, &snapshot, previous.as_ref())
|
||||
{
|
||||
let outcome = self
|
||||
.store
|
||||
.touch_quota_observation(QuotaObservationTouch {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
observed_at,
|
||||
})
|
||||
.await?;
|
||||
if outcome == QuotaWriteOutcome::Conflict {
|
||||
summary.stale += 1;
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(previous) = previous {
|
||||
self.scheduling
|
||||
.observe(&previous.with_observed_at(observed_at));
|
||||
}
|
||||
summary.exhausted += 1;
|
||||
return Ok(());
|
||||
}
|
||||
@@ -688,10 +1076,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()))
|
||||
@@ -716,14 +1119,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(),
|
||||
@@ -930,8 +1333,29 @@ impl CodexCredentialQuotaService {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !quota_refresh_can_update(&account, &snapshot, previous.as_ref()) {
|
||||
return Ok(previous.unwrap_or_else(|| snapshot.with_quota_state(account.quota())));
|
||||
if !quota_refresh_allows_recovery(&account, &snapshot, previous.as_ref()) {
|
||||
if self
|
||||
.store
|
||||
.touch_quota_observation(QuotaObservationTouch {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
observed_at,
|
||||
})
|
||||
.await?
|
||||
== QuotaWriteOutcome::Conflict
|
||||
{
|
||||
return Err(CodexCredentialQuotaError::RevisionConflict);
|
||||
}
|
||||
return Ok(match previous {
|
||||
Some(previous) => {
|
||||
let previous = previous.with_observed_at(observed_at);
|
||||
self.scheduling.observe(&previous);
|
||||
previous
|
||||
}
|
||||
None => snapshot
|
||||
.with_quota_state(account.quota())
|
||||
.with_observed_at(observed_at),
|
||||
});
|
||||
}
|
||||
let state = if account.quota().is_exhausted() {
|
||||
QuotaState::allowed(observed_at)
|
||||
@@ -1025,9 +1449,286 @@ impl CodexCredentialQuotaService {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_daily_usage_with_recovery(
|
||||
&self,
|
||||
client: &CodexBackendClient,
|
||||
account: &ProviderAccount,
|
||||
start_date: chrono::NaiveDate,
|
||||
end_date: chrono::NaiveDate,
|
||||
) -> Result<FetchedDailyUsage, DailyUsageFetchError> {
|
||||
let mut prepared = self
|
||||
.agent_identity
|
||||
.prepare(account)
|
||||
.await
|
||||
.map_err(|_| DailyUsageFetchError::InvalidCredential)?;
|
||||
let mut result = fetch_daily_usage_once(client, &prepared, start_date, end_date).await;
|
||||
if let Err(DailyUsageAttemptError::Upstream(error)) = &result
|
||||
&& let Some(recovered) = self
|
||||
.agent_identity
|
||||
.recover_after_rejected_task(
|
||||
prepared.account.id(),
|
||||
&prepared.credential.authentication,
|
||||
error,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| DailyUsageFetchError::Recovery)?
|
||||
{
|
||||
prepared = recovered;
|
||||
result = fetch_daily_usage_once(client, &prepared, start_date, end_date).await;
|
||||
}
|
||||
match result {
|
||||
Ok(value) => Ok(value),
|
||||
Err(DailyUsageAttemptError::InvalidCredential) => {
|
||||
Err(DailyUsageFetchError::InvalidCredential)
|
||||
}
|
||||
Err(DailyUsageAttemptError::Upstream(error)) => Err(DailyUsageFetchError::Upstream {
|
||||
prepared: Box::new(prepared),
|
||||
error,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_refresh_can_update(
|
||||
async fn fetch_daily_usage_once(
|
||||
client: &CodexBackendClient,
|
||||
prepared: &PreparedCodexRuntimeCredential,
|
||||
start_date: chrono::NaiveDate,
|
||||
end_date: chrono::NaiveDate,
|
||||
) -> Result<FetchedDailyUsage, DailyUsageAttemptError> {
|
||||
let authorization = prepared
|
||||
.credential
|
||||
.authentication
|
||||
.authorization_header(Utc::now())
|
||||
.map_err(|_| DailyUsageAttemptError::InvalidCredential)?;
|
||||
let workspace_request_id = format!("usage_statistics_{}", Uuid::now_v7().simple());
|
||||
let workspace = client
|
||||
.fetch_daily_usage(
|
||||
CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
&workspace_request_id,
|
||||
None,
|
||||
),
|
||||
CodexDailyUsageReport::WorkspaceModelTokens,
|
||||
start_date,
|
||||
end_date,
|
||||
)
|
||||
.await;
|
||||
match workspace {
|
||||
Ok(model_breakdown) => {
|
||||
return Ok(FetchedDailyUsage {
|
||||
mode: CodexUsageStatisticsMode::Workspace,
|
||||
model_breakdown,
|
||||
daily_totals: None,
|
||||
});
|
||||
}
|
||||
Err(error) if workspace_is_unavailable_for_account(&error) => {}
|
||||
Err(error) => return Err(DailyUsageAttemptError::Upstream(error)),
|
||||
}
|
||||
|
||||
let model_request_id = format!("usage_statistics_{}", Uuid::now_v7().simple());
|
||||
let totals_request_id = format!("usage_statistics_{}", Uuid::now_v7().simple());
|
||||
let model_breakdown = client.fetch_daily_usage(
|
||||
CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
&model_request_id,
|
||||
None,
|
||||
),
|
||||
CodexDailyUsageReport::PersonalModelCredits,
|
||||
start_date,
|
||||
end_date,
|
||||
);
|
||||
let daily_totals = client.fetch_daily_usage(
|
||||
CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
&totals_request_id,
|
||||
None,
|
||||
),
|
||||
CodexDailyUsageReport::PersonalTokenTotals,
|
||||
start_date,
|
||||
end_date,
|
||||
);
|
||||
let (model_breakdown, daily_totals) = tokio::try_join!(model_breakdown, daily_totals)
|
||||
.map_err(DailyUsageAttemptError::Upstream)?;
|
||||
Ok(FetchedDailyUsage {
|
||||
mode: CodexUsageStatisticsMode::Personal,
|
||||
model_breakdown,
|
||||
daily_totals: Some(daily_totals),
|
||||
})
|
||||
}
|
||||
|
||||
fn workspace_is_unavailable_for_account(error: &CodexClientError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexClientError::Upstream { status, body, .. }
|
||||
if *status == reqwest::StatusCode::BAD_REQUEST
|
||||
&& body.to_ascii_lowercase().contains("no active workspace")
|
||||
)
|
||||
}
|
||||
|
||||
fn map_statistics_quota_fetch_error(error: CodexQuotaFetchError) -> CodexUsageStatisticsError {
|
||||
match error {
|
||||
CodexQuotaFetchError::InvalidCredential => CodexUsageStatisticsError::InvalidCredentialData,
|
||||
CodexQuotaFetchError::Recovery => CodexUsageStatisticsError::TransportUnavailable,
|
||||
CodexQuotaFetchError::Upstream { account, error } => map_statistics_client_error(
|
||||
error,
|
||||
account.authentication_kind() == CODEX_AUTHENTICATION_KIND_OAUTH,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_daily_usage_fetch_error(error: DailyUsageFetchError) -> CodexUsageStatisticsError {
|
||||
match error {
|
||||
DailyUsageFetchError::InvalidCredential => CodexUsageStatisticsError::InvalidCredentialData,
|
||||
DailyUsageFetchError::Recovery => CodexUsageStatisticsError::TransportUnavailable,
|
||||
DailyUsageFetchError::Upstream { prepared, error } => {
|
||||
map_statistics_client_error(error, prepared.credential.authentication.oauth().is_some())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_statistics_client_error(
|
||||
error: CodexClientError,
|
||||
credential_can_refresh: bool,
|
||||
) -> CodexUsageStatisticsError {
|
||||
match error {
|
||||
CodexClientError::Upstream { status, body, .. }
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED && credential_can_refresh =>
|
||||
{
|
||||
CodexUsageStatisticsError::CredentialRefreshRequired {
|
||||
upstream_body: Some(body),
|
||||
}
|
||||
}
|
||||
CodexClientError::Upstream {
|
||||
status,
|
||||
body,
|
||||
retry_after_seconds,
|
||||
..
|
||||
} => CodexUsageStatisticsError::Upstream {
|
||||
status: status.as_u16(),
|
||||
body,
|
||||
retry_after_seconds,
|
||||
},
|
||||
_ => CodexUsageStatisticsError::TransportUnavailable,
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_reset_credits_once(
|
||||
client: &CodexBackendClient,
|
||||
prepared: &PreparedCodexRuntimeCredential,
|
||||
request_id: &str,
|
||||
) -> Result<CodexRateLimitResetCredits, ResetCreditAttemptError> {
|
||||
let authorization = prepared
|
||||
.credential
|
||||
.authentication
|
||||
.authorization_header(Utc::now())
|
||||
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
|
||||
client
|
||||
.list_rate_limit_reset_credits(CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
request_id,
|
||||
None,
|
||||
))
|
||||
.await
|
||||
.map_err(ResetCreditAttemptError::Upstream)
|
||||
}
|
||||
|
||||
async fn consume_reset_credit_once(
|
||||
client: &CodexBackendClient,
|
||||
prepared: &PreparedCodexRuntimeCredential,
|
||||
request_id: &str,
|
||||
credit_id: Option<&str>,
|
||||
redeem_request_id: Uuid,
|
||||
) -> Result<CodexRateLimitResetCreditsConsumeResult, ResetCreditAttemptError> {
|
||||
let authorization = prepared
|
||||
.credential
|
||||
.authentication
|
||||
.authorization_header(Utc::now())
|
||||
.map_err(|_| ResetCreditAttemptError::InvalidCredential)?;
|
||||
client
|
||||
.consume_rate_limit_reset_credit(
|
||||
CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
prepared.account.upstream_account_id(),
|
||||
request_id,
|
||||
None,
|
||||
),
|
||||
credit_id,
|
||||
redeem_request_id,
|
||||
)
|
||||
.await
|
||||
.map_err(ResetCreditAttemptError::Upstream)
|
||||
}
|
||||
|
||||
fn map_reset_credit_attempt_error(
|
||||
error: ResetCreditAttemptError,
|
||||
prepared: &PreparedCodexRuntimeCredential,
|
||||
consume: bool,
|
||||
) -> CodexResetCreditsError {
|
||||
match error {
|
||||
ResetCreditAttemptError::InvalidCredential => CodexResetCreditsError::InvalidCredentialData,
|
||||
ResetCreditAttemptError::Upstream(error) => {
|
||||
map_reset_credit_client_error(error, prepared, consume)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_reset_credit_client_error(
|
||||
error: CodexClientError,
|
||||
prepared: &PreparedCodexRuntimeCredential,
|
||||
consume: bool,
|
||||
) -> CodexResetCreditsError {
|
||||
match error {
|
||||
CodexClientError::Upstream {
|
||||
status,
|
||||
body,
|
||||
diagnostics,
|
||||
..
|
||||
} if status == reqwest::StatusCode::UNAUTHORIZED
|
||||
&& reset_credit_response_was_explicit_rejection(status, &diagnostics)
|
||||
&& prepared.credential.authentication.oauth().is_some() =>
|
||||
{
|
||||
CodexResetCreditsError::CredentialRefreshRequired {
|
||||
upstream_body: Some(body),
|
||||
}
|
||||
}
|
||||
CodexClientError::Upstream {
|
||||
status,
|
||||
body,
|
||||
retry_after_seconds,
|
||||
diagnostics,
|
||||
..
|
||||
} => {
|
||||
// 2xx 后发生的解码/响应体上限错误使用 synthetic 502 表示,但额度卡
|
||||
// 可能已经消费。只有 transport 记录的真实非成功状态与错误状态一致时,
|
||||
// 才能把它当作确定的上游拒绝并允许前端清除 pending 幂等键。
|
||||
if consume && !reset_credit_response_was_explicit_rejection(status, &diagnostics) {
|
||||
CodexResetCreditsError::ConsumeResultUnknown
|
||||
} else {
|
||||
CodexResetCreditsError::Upstream {
|
||||
status: status.as_u16(),
|
||||
body,
|
||||
retry_after_seconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
_ if consume => CodexResetCreditsError::ConsumeResultUnknown,
|
||||
_ => CodexResetCreditsError::TransportUnavailable,
|
||||
}
|
||||
}
|
||||
|
||||
fn reset_credit_response_was_explicit_rejection(
|
||||
status: reqwest::StatusCode,
|
||||
diagnostics: &crate::transport::CodexUpstreamDiagnostics,
|
||||
) -> bool {
|
||||
!status.is_success() && diagnostics.status_code == Some(status.as_u16())
|
||||
}
|
||||
|
||||
fn quota_refresh_allows_recovery(
|
||||
account: &ProviderAccount,
|
||||
refreshed: &CodexAccountQuotaSnapshot,
|
||||
previous: Option<&CodexAccountQuotaSnapshot>,
|
||||
@@ -1138,18 +1839,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()));
|
||||
}
|
||||
@@ -1177,48 +1908,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` 占位。
|
||||
@@ -1230,7 +1932,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)]
|
||||
@@ -151,6 +154,13 @@ impl CodexAccountQuotaSnapshot {
|
||||
self.quota = quota;
|
||||
self
|
||||
}
|
||||
|
||||
/// 只更新展示与调度使用的最后成功查询时间,保留既有额度事实。
|
||||
#[must_use]
|
||||
pub(crate) const fn with_observed_at(mut self, observed_at: SystemTime) -> Self {
|
||||
self.observed_at = observed_at;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexQuotaFact {
|
||||
@@ -166,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)?;
|
||||
@@ -189,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> {
|
||||
@@ -270,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,
|
||||
@@ -309,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() {
|
||||
@@ -391,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)
|
||||
}
|
||||
@@ -457,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,
|
||||
|
||||
@@ -0,0 +1,970 @@
|
||||
use std::{
|
||||
cmp::{Ordering, Reverse},
|
||||
collections::BTreeMap,
|
||||
};
|
||||
|
||||
use chrono::{DateTime, FixedOffset, NaiveDate, NaiveTime, TimeDelta, Utc};
|
||||
use gateway_core::accounting::{Decimal, Money};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::transport::{OpenAiBillingUsage, openai_aggregate_billing_breakdown};
|
||||
|
||||
const EPSILON: f64 = 1e-9;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CodexUsageStatisticsMode {
|
||||
Workspace,
|
||||
Personal,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum CodexUsageStatisticsServiceTier {
|
||||
Standard,
|
||||
Fast,
|
||||
}
|
||||
|
||||
impl CodexUsageStatisticsServiceTier {
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Standard => "standard",
|
||||
Self::Fast => "fast",
|
||||
}
|
||||
}
|
||||
|
||||
const fn billing_service_tier(self) -> &'static str {
|
||||
match self {
|
||||
Self::Standard => "default",
|
||||
Self::Fast => "fast",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct CodexUsageStatisticsTokens {
|
||||
pub uncached_input: u64,
|
||||
pub cached_input: u64,
|
||||
pub output: u64,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
impl CodexUsageStatisticsTokens {
|
||||
fn checked_add(self, other: Self) -> Option<Self> {
|
||||
Some(Self {
|
||||
uncached_input: self.uncached_input.checked_add(other.uncached_input)?,
|
||||
cached_input: self.cached_input.checked_add(other.cached_input)?,
|
||||
output: self.output.checked_add(other.output)?,
|
||||
total: self.total.checked_add(other.total)?,
|
||||
})
|
||||
}
|
||||
|
||||
const fn is_empty(self) -> bool {
|
||||
self.total == 0 && self.uncached_input == 0 && self.cached_input == 0 && self.output == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexUsageStatisticsCycle {
|
||||
pub offset: i8,
|
||||
pub start_at: DateTime<Utc>,
|
||||
pub end_at: DateTime<Utc>,
|
||||
pub window_seconds: u64,
|
||||
pub used_percent: Option<f64>,
|
||||
pub is_current: bool,
|
||||
pub can_go_previous: bool,
|
||||
pub can_go_next: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexUsageStatisticsModel {
|
||||
pub key: String,
|
||||
pub model: String,
|
||||
pub service_tier: CodexUsageStatisticsServiceTier,
|
||||
pub credit_share: Option<f64>,
|
||||
pub quota_share: Option<f64>,
|
||||
pub tokens: CodexUsageStatisticsTokens,
|
||||
pub estimated_cost: Option<Money>,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_estimated_allocation: bool,
|
||||
pub has_rate_fallback: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexUsageStatisticsDay {
|
||||
pub date: NaiveDate,
|
||||
pub credit_share: Option<f64>,
|
||||
pub tokens: CodexUsageStatisticsTokens,
|
||||
pub estimated_cost: Option<Money>,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
pub is_boundary_day: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexUsageStatisticsSummary {
|
||||
pub tokens: CodexUsageStatisticsTokens,
|
||||
pub estimated_cost: Option<Money>,
|
||||
pub projected_tokens: Option<u64>,
|
||||
pub projected_cost: Option<Money>,
|
||||
pub day_count: u32,
|
||||
pub has_unknown_pricing: bool,
|
||||
pub has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexUsageStatistics {
|
||||
pub mode: CodexUsageStatisticsMode,
|
||||
pub cycle: CodexUsageStatisticsCycle,
|
||||
pub summary: CodexUsageStatisticsSummary,
|
||||
pub models: Vec<CodexUsageStatisticsModel>,
|
||||
pub daily: Vec<CodexUsageStatisticsDay>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct CodexUsageStatisticsPeriod {
|
||||
pub cycle_offset: i8,
|
||||
pub start_at: DateTime<Utc>,
|
||||
pub end_at: DateTime<Utc>,
|
||||
pub window_seconds: u64,
|
||||
pub query_start_date: NaiveDate,
|
||||
pub query_end_date: NaiveDate,
|
||||
}
|
||||
|
||||
pub fn statistics_period(
|
||||
current_start_at: DateTime<Utc>,
|
||||
current_end_at: DateTime<Utc>,
|
||||
window_seconds: u64,
|
||||
cycle_offset: i8,
|
||||
timezone: FixedOffset,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<CodexUsageStatisticsPeriod> {
|
||||
let window_delta = TimeDelta::seconds(i64::try_from(window_seconds).ok()?);
|
||||
let shift = window_delta.checked_mul(i32::from(cycle_offset))?;
|
||||
let start_at = current_start_at.checked_add_signed(shift)?;
|
||||
let end_at = current_end_at.checked_add_signed(shift)?;
|
||||
let query_start_date = start_at
|
||||
.checked_sub_signed(TimeDelta::days(1))?
|
||||
.with_timezone(&timezone)
|
||||
.date_naive();
|
||||
let effective_end = end_at.min(now);
|
||||
let query_end_date = effective_end
|
||||
.checked_add_signed(TimeDelta::days(1))?
|
||||
.with_timezone(&timezone)
|
||||
.date_naive();
|
||||
Some(CodexUsageStatisticsPeriod {
|
||||
cycle_offset,
|
||||
start_at,
|
||||
end_at,
|
||||
window_seconds,
|
||||
query_start_date,
|
||||
query_end_date,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) struct BuildUsageStatistics<'a> {
|
||||
pub mode: CodexUsageStatisticsMode,
|
||||
pub period: CodexUsageStatisticsPeriod,
|
||||
pub timezone: FixedOffset,
|
||||
pub current_used_percent: Option<f64>,
|
||||
pub now: DateTime<Utc>,
|
||||
pub model_breakdown: &'a Value,
|
||||
pub daily_totals: Option<&'a Value>,
|
||||
pub max_cycle_offset: u8,
|
||||
}
|
||||
|
||||
pub(super) fn build_usage_statistics(input: BuildUsageStatistics<'_>) -> CodexUsageStatistics {
|
||||
let BuildUsageStatistics {
|
||||
mode,
|
||||
period,
|
||||
timezone,
|
||||
current_used_percent,
|
||||
now,
|
||||
model_breakdown,
|
||||
daily_totals,
|
||||
max_cycle_offset,
|
||||
} = input;
|
||||
let all_days = match mode {
|
||||
CodexUsageStatisticsMode::Workspace => build_workspace_days(model_breakdown),
|
||||
CodexUsageStatisticsMode::Personal => {
|
||||
build_personal_days(daily_totals.unwrap_or(&Value::Null), model_breakdown)
|
||||
}
|
||||
};
|
||||
let selected_days = select_cycle_days(
|
||||
&all_days,
|
||||
period.start_at,
|
||||
period.end_at,
|
||||
timezone,
|
||||
now,
|
||||
true,
|
||||
);
|
||||
let is_current = period.cycle_offset == 0;
|
||||
let used_percent = is_current
|
||||
.then_some(current_used_percent)
|
||||
.flatten()
|
||||
.filter(|value| value.is_finite())
|
||||
.map(|value| value.clamp(0.0, 100.0));
|
||||
let used_ratio = used_percent
|
||||
.map(|value| value.clamp(0.0, 100.0) / 100.0)
|
||||
.unwrap_or_default();
|
||||
let summary = summarize(&selected_days, used_ratio, is_current);
|
||||
let models = aggregate_models(&selected_days, used_ratio, is_current);
|
||||
let total_credits = selected_days
|
||||
.iter()
|
||||
.map(|day| day.credit_weight)
|
||||
.sum::<f64>();
|
||||
let mut daily = selected_days
|
||||
.into_iter()
|
||||
.map(|day| CodexUsageStatisticsDay {
|
||||
date: day.date,
|
||||
credit_share: (total_credits > EPSILON).then_some(day.credit_weight / total_credits),
|
||||
tokens: day.tokens,
|
||||
estimated_cost: day.estimated_cost,
|
||||
has_unknown_pricing: day.has_unknown_pricing,
|
||||
has_missing_token_data: day.has_missing_token_data,
|
||||
is_boundary_day: day.is_boundary_day,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
daily.sort_by_key(|day| Reverse(day.date));
|
||||
|
||||
CodexUsageStatistics {
|
||||
mode,
|
||||
cycle: CodexUsageStatisticsCycle {
|
||||
offset: period.cycle_offset,
|
||||
start_at: period.start_at,
|
||||
end_at: period.end_at,
|
||||
window_seconds: period.window_seconds,
|
||||
used_percent,
|
||||
is_current,
|
||||
can_go_previous: period.cycle_offset > -(max_cycle_offset as i8),
|
||||
can_go_next: period.cycle_offset < 0,
|
||||
},
|
||||
summary,
|
||||
models,
|
||||
daily,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct DailyModelUsage {
|
||||
model: String,
|
||||
service_tier: CodexUsageStatisticsServiceTier,
|
||||
credits: f64,
|
||||
tokens: CodexUsageStatisticsTokens,
|
||||
estimated_cost: Option<Money>,
|
||||
has_unknown_pricing: bool,
|
||||
has_estimated_allocation: bool,
|
||||
has_rate_fallback: bool,
|
||||
has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct DailyUsage {
|
||||
date: NaiveDate,
|
||||
credit_weight: f64,
|
||||
tokens: CodexUsageStatisticsTokens,
|
||||
estimated_cost: Option<Money>,
|
||||
has_unknown_pricing: bool,
|
||||
has_missing_token_data: bool,
|
||||
is_boundary_day: bool,
|
||||
models: Vec<DailyModelUsage>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct MoneySum {
|
||||
value: Option<Money>,
|
||||
overflowed: bool,
|
||||
}
|
||||
|
||||
impl MoneySum {
|
||||
fn push(&mut self, value: Option<Money>) {
|
||||
let Some(value) = value else {
|
||||
return;
|
||||
};
|
||||
if self.overflowed {
|
||||
return;
|
||||
}
|
||||
self.value = match self.value {
|
||||
Some(current) => match current.checked_add(value) {
|
||||
Some(total) => Some(total),
|
||||
None => {
|
||||
self.overflowed = true;
|
||||
None
|
||||
}
|
||||
},
|
||||
None => Some(value),
|
||||
};
|
||||
}
|
||||
|
||||
fn finish(self) -> Option<Money> {
|
||||
(!self.overflowed).then_some(self.value).flatten()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_workspace_days(payload: &Value) -> Vec<DailyUsage> {
|
||||
extract_daily_list(payload)
|
||||
.into_iter()
|
||||
.filter_map(|row| {
|
||||
let date = date_value(row.get("date")?)?;
|
||||
let mut models = value_array(row.get("models"))
|
||||
.iter()
|
||||
.filter_map(direct_model_usage)
|
||||
.collect::<Vec<_>>();
|
||||
let model_tokens = sum_tokens(models.iter().map(|model| model.tokens));
|
||||
let fallback_tokens = token_parts(row.get("totals").unwrap_or(row));
|
||||
let tokens = if model_tokens.is_empty() {
|
||||
fallback_tokens
|
||||
} else {
|
||||
model_tokens
|
||||
};
|
||||
let credit_weight = day_credit_weight(row);
|
||||
if tokens.is_empty() && credit_weight <= EPSILON {
|
||||
return None;
|
||||
}
|
||||
let has_missing_token_data = tokens.is_empty() && credit_weight > EPSILON;
|
||||
if models.is_empty() && !tokens.is_empty() {
|
||||
models.push(unknown_model(tokens, false));
|
||||
}
|
||||
if has_missing_token_data {
|
||||
for model in &mut models {
|
||||
model.has_missing_token_data = true;
|
||||
}
|
||||
}
|
||||
Some(finalize_day(
|
||||
date,
|
||||
credit_weight,
|
||||
tokens,
|
||||
models,
|
||||
has_missing_token_data,
|
||||
))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn build_personal_days(totals_payload: &Value, breakdown_payload: &Value) -> Vec<DailyUsage> {
|
||||
let totals_by_date = rows_by_date(totals_payload);
|
||||
let breakdown_by_date = rows_by_date(breakdown_payload);
|
||||
let mut dates = totals_by_date
|
||||
.keys()
|
||||
.chain(breakdown_by_date.keys())
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
dates.sort_unstable();
|
||||
dates.dedup();
|
||||
|
||||
dates
|
||||
.into_iter()
|
||||
.filter_map(|date| {
|
||||
let totals_row = totals_by_date.get(&date).unwrap_or(&Value::Null);
|
||||
let breakdown_row = breakdown_by_date.get(&date).unwrap_or(&Value::Null);
|
||||
let totals = token_parts(totals_row.get("totals").unwrap_or(totals_row));
|
||||
let credit_weight = day_credit_weight(breakdown_row);
|
||||
let weighted = weighted_models(breakdown_row);
|
||||
if totals.is_empty() {
|
||||
if credit_weight <= EPSILON {
|
||||
return None;
|
||||
}
|
||||
let models = weighted
|
||||
.into_iter()
|
||||
.map(|model| DailyModelUsage {
|
||||
model: model.model,
|
||||
service_tier: model.service_tier,
|
||||
credits: model.weight,
|
||||
tokens: CodexUsageStatisticsTokens::default(),
|
||||
estimated_cost: None,
|
||||
has_unknown_pricing: false,
|
||||
has_estimated_allocation: false,
|
||||
has_rate_fallback: false,
|
||||
has_missing_token_data: true,
|
||||
})
|
||||
.collect();
|
||||
return Some(finalize_day(date, credit_weight, totals, models, true));
|
||||
}
|
||||
|
||||
let has_model_breakdown = !weighted.is_empty();
|
||||
let weighted = if has_model_breakdown {
|
||||
weighted
|
||||
} else {
|
||||
vec![WeightedModel {
|
||||
model: "unknown".to_owned(),
|
||||
service_tier: CodexUsageStatisticsServiceTier::Standard,
|
||||
weight: 1.0,
|
||||
}]
|
||||
};
|
||||
let allocation = rate_adjusted_shares(&weighted, totals);
|
||||
let uncached = allocate_integer_total(totals.uncached_input, &allocation.token_shares);
|
||||
let cached = allocate_integer_total(totals.cached_input, &allocation.token_shares);
|
||||
let output = allocate_integer_total(totals.output, &allocation.token_shares);
|
||||
let total = allocate_integer_total(totals.total, &allocation.token_shares);
|
||||
let estimated_allocation = weighted.len() > 1 || !has_model_breakdown;
|
||||
let rate_fallback = allocation.has_rate_fallback || !has_model_breakdown;
|
||||
let models = weighted
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, model)| {
|
||||
let tokens = CodexUsageStatisticsTokens {
|
||||
uncached_input: uncached.get(index).copied().unwrap_or_default(),
|
||||
cached_input: cached.get(index).copied().unwrap_or_default(),
|
||||
output: output.get(index).copied().unwrap_or_default(),
|
||||
total: total.get(index).copied().unwrap_or_default(),
|
||||
};
|
||||
let estimated_cost = aggregate_cost(&model.model, model.service_tier, tokens);
|
||||
DailyModelUsage {
|
||||
model: model.model,
|
||||
service_tier: model.service_tier,
|
||||
credits: if has_model_breakdown {
|
||||
model.weight
|
||||
} else {
|
||||
0.0
|
||||
},
|
||||
tokens,
|
||||
has_unknown_pricing: estimated_cost.is_none() && !tokens.is_empty(),
|
||||
estimated_cost,
|
||||
has_estimated_allocation: estimated_allocation,
|
||||
has_rate_fallback: rate_fallback,
|
||||
has_missing_token_data: false,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Some(finalize_day(date, credit_weight, totals, models, false))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn direct_model_usage(value: &Value) -> Option<DailyModelUsage> {
|
||||
let tokens = token_parts(value);
|
||||
let credits = non_negative_f64(value.get("credits"));
|
||||
if tokens.is_empty() && credits <= EPSILON {
|
||||
return None;
|
||||
}
|
||||
let model = model_name(value);
|
||||
let service_tier = model_service_tier(value.get("speed").or_else(|| value.get("mode")));
|
||||
let estimated_cost = (!tokens.is_empty())
|
||||
.then(|| aggregate_cost(&model, service_tier, tokens))
|
||||
.flatten();
|
||||
Some(DailyModelUsage {
|
||||
model,
|
||||
service_tier,
|
||||
credits,
|
||||
tokens,
|
||||
has_unknown_pricing: estimated_cost.is_none() && !tokens.is_empty(),
|
||||
estimated_cost,
|
||||
has_estimated_allocation: false,
|
||||
has_rate_fallback: false,
|
||||
has_missing_token_data: tokens.is_empty() && credits > EPSILON,
|
||||
})
|
||||
}
|
||||
|
||||
fn unknown_model(
|
||||
tokens: CodexUsageStatisticsTokens,
|
||||
estimated_allocation: bool,
|
||||
) -> DailyModelUsage {
|
||||
DailyModelUsage {
|
||||
model: "unknown".to_owned(),
|
||||
service_tier: CodexUsageStatisticsServiceTier::Standard,
|
||||
credits: 0.0,
|
||||
tokens,
|
||||
estimated_cost: None,
|
||||
has_unknown_pricing: !tokens.is_empty(),
|
||||
has_estimated_allocation: estimated_allocation,
|
||||
has_rate_fallback: estimated_allocation,
|
||||
has_missing_token_data: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn finalize_day(
|
||||
date: NaiveDate,
|
||||
credit_weight: f64,
|
||||
tokens: CodexUsageStatisticsTokens,
|
||||
models: Vec<DailyModelUsage>,
|
||||
has_missing_token_data: bool,
|
||||
) -> DailyUsage {
|
||||
let mut costs = MoneySum::default();
|
||||
for model in &models {
|
||||
costs.push(model.estimated_cost);
|
||||
}
|
||||
DailyUsage {
|
||||
date,
|
||||
credit_weight,
|
||||
tokens,
|
||||
estimated_cost: costs.finish(),
|
||||
has_unknown_pricing: models.iter().any(|model| model.has_unknown_pricing),
|
||||
has_missing_token_data,
|
||||
is_boundary_day: false,
|
||||
models,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct WeightedModel {
|
||||
model: String,
|
||||
service_tier: CodexUsageStatisticsServiceTier,
|
||||
weight: f64,
|
||||
}
|
||||
|
||||
struct Allocation {
|
||||
token_shares: Vec<f64>,
|
||||
has_rate_fallback: bool,
|
||||
}
|
||||
|
||||
fn weighted_models(row: &Value) -> Vec<WeightedModel> {
|
||||
value_array(row.get("models"))
|
||||
.iter()
|
||||
.filter_map(|value| {
|
||||
let weight = non_negative_f64(value.get("credits"));
|
||||
(weight > EPSILON).then(|| WeightedModel {
|
||||
model: model_name(value),
|
||||
service_tier: model_service_tier(value.get("speed").or_else(|| value.get("mode"))),
|
||||
weight,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn rate_adjusted_shares(
|
||||
models: &[WeightedModel],
|
||||
totals: CodexUsageStatisticsTokens,
|
||||
) -> Allocation {
|
||||
let total_weight = models.iter().map(|model| model.weight).sum::<f64>();
|
||||
let credit_shares = models
|
||||
.iter()
|
||||
.map(|model| model.weight / total_weight)
|
||||
.collect::<Vec<_>>();
|
||||
if models.len() == 1 {
|
||||
return Allocation {
|
||||
token_shares: vec![1.0],
|
||||
has_rate_fallback: false,
|
||||
};
|
||||
}
|
||||
let classified_tokens = totals
|
||||
.uncached_input
|
||||
.saturating_add(totals.cached_input)
|
||||
.saturating_add(totals.output);
|
||||
if classified_tokens == 0 {
|
||||
return Allocation {
|
||||
token_shares: credit_shares,
|
||||
has_rate_fallback: true,
|
||||
};
|
||||
}
|
||||
let rates = models
|
||||
.iter()
|
||||
.map(|model| blended_rate(&model.model, model.service_tier, totals, classified_tokens))
|
||||
.collect::<Vec<_>>();
|
||||
let known = rates
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, rate)| rate.map(|rate| (index, rate)))
|
||||
.collect::<Vec<_>>();
|
||||
if known.is_empty() {
|
||||
return Allocation {
|
||||
token_shares: credit_shares,
|
||||
has_rate_fallback: true,
|
||||
};
|
||||
}
|
||||
let known_weight = known
|
||||
.iter()
|
||||
.map(|(index, _)| models[*index].weight)
|
||||
.sum::<f64>();
|
||||
let fallback_rate = known
|
||||
.iter()
|
||||
.map(|(index, rate)| rate * models[*index].weight / known_weight)
|
||||
.sum::<f64>();
|
||||
let inverse = models
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, model)| model.weight / rates[index].unwrap_or(fallback_rate))
|
||||
.collect::<Vec<_>>();
|
||||
let inverse_total = inverse.iter().sum::<f64>();
|
||||
Allocation {
|
||||
token_shares: inverse
|
||||
.into_iter()
|
||||
.map(|value| value / inverse_total)
|
||||
.collect(),
|
||||
has_rate_fallback: known.len() != models.len(),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn blended_rate(
|
||||
model: &str,
|
||||
service_tier: CodexUsageStatisticsServiceTier,
|
||||
totals: CodexUsageStatisticsTokens,
|
||||
classified_tokens: u64,
|
||||
) -> Option<f64> {
|
||||
let cost = aggregate_cost(model, service_tier, totals)?;
|
||||
let rate = cost.amount().scaled() as f64 / classified_tokens as f64;
|
||||
(rate > EPSILON).then_some(rate)
|
||||
}
|
||||
|
||||
fn aggregate_cost(
|
||||
model: &str,
|
||||
service_tier: CodexUsageStatisticsServiceTier,
|
||||
tokens: CodexUsageStatisticsTokens,
|
||||
) -> Option<Money> {
|
||||
let input = tokens.uncached_input.checked_add(tokens.cached_input)?;
|
||||
openai_aggregate_billing_breakdown(
|
||||
model,
|
||||
OpenAiBillingUsage::new(input, tokens.output, tokens.cached_input, 0),
|
||||
Some(service_tier.billing_service_tier()),
|
||||
)
|
||||
.map(|breakdown| breakdown.total_amount())
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
|
||||
fn allocate_integer_total(total: u64, shares: &[f64]) -> Vec<u64> {
|
||||
if shares.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let share_total = shares.iter().copied().sum::<f64>();
|
||||
if !share_total.is_finite() || share_total <= EPSILON {
|
||||
return vec![0; shares.len()];
|
||||
}
|
||||
let raw = shares
|
||||
.iter()
|
||||
.map(|share| total as f64 * (share.max(0.0) / share_total))
|
||||
.collect::<Vec<_>>();
|
||||
let mut allocated = raw
|
||||
.iter()
|
||||
.map(|value| value.floor() as u64)
|
||||
.collect::<Vec<_>>();
|
||||
let allocated_total = allocated.iter().copied().sum::<u64>();
|
||||
let mut remainder = total.saturating_sub(allocated_total);
|
||||
let mut order = raw
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, value)| (index, value - value.floor()))
|
||||
.collect::<Vec<_>>();
|
||||
order.sort_by(|left, right| {
|
||||
right
|
||||
.1
|
||||
.total_cmp(&left.1)
|
||||
.then_with(|| left.0.cmp(&right.0))
|
||||
});
|
||||
let mut index = 0_usize;
|
||||
while remainder > 0 {
|
||||
let target = order[index % order.len()].0;
|
||||
allocated[target] = allocated[target].saturating_add(1);
|
||||
remainder -= 1;
|
||||
index += 1;
|
||||
}
|
||||
allocated
|
||||
}
|
||||
|
||||
fn select_cycle_days(
|
||||
days: &[DailyUsage],
|
||||
start_at: DateTime<Utc>,
|
||||
end_at: DateTime<Utc>,
|
||||
timezone: FixedOffset,
|
||||
now: DateTime<Utc>,
|
||||
mark_boundary: bool,
|
||||
) -> Vec<DailyUsage> {
|
||||
let local_start = start_at.with_timezone(&timezone);
|
||||
let start_date = local_start.date_naive();
|
||||
let effective_end = end_at.min(now);
|
||||
let end_date = effective_end.with_timezone(&timezone).date_naive();
|
||||
let exclude_end_date = end_at <= now && end_date > start_date;
|
||||
let is_boundary = local_start.time() != NaiveTime::MIN;
|
||||
days.iter()
|
||||
.filter(|day| {
|
||||
day.date >= start_date
|
||||
&& if exclude_end_date {
|
||||
day.date < end_date
|
||||
} else {
|
||||
day.date <= end_date
|
||||
}
|
||||
})
|
||||
.cloned()
|
||||
.map(|mut day| {
|
||||
day.is_boundary_day = mark_boundary && is_boundary && day.date == start_date;
|
||||
day
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn summarize(
|
||||
days: &[DailyUsage],
|
||||
used_ratio: f64,
|
||||
is_current: bool,
|
||||
) -> CodexUsageStatisticsSummary {
|
||||
let tokens = sum_tokens(days.iter().map(|day| day.tokens));
|
||||
let mut costs = MoneySum::default();
|
||||
for day in days {
|
||||
costs.push(day.estimated_cost);
|
||||
}
|
||||
let estimated_cost = costs.finish();
|
||||
let has_unknown_pricing = days.iter().any(|day| day.has_unknown_pricing);
|
||||
let projected_tokens = (is_current && used_ratio > EPSILON && tokens.total > 0)
|
||||
.then(|| project_u64(tokens.total, used_ratio))
|
||||
.flatten();
|
||||
let projected_cost = (is_current && !has_unknown_pricing && used_ratio > EPSILON)
|
||||
.then(|| estimated_cost.and_then(|cost| project_money(cost, used_ratio)))
|
||||
.flatten();
|
||||
CodexUsageStatisticsSummary {
|
||||
tokens,
|
||||
estimated_cost,
|
||||
projected_tokens,
|
||||
projected_cost,
|
||||
day_count: u32::try_from(days.len()).unwrap_or(u32::MAX),
|
||||
has_unknown_pricing,
|
||||
has_missing_token_data: days.iter().any(|day| day.has_missing_token_data),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ModelAggregate {
|
||||
credits: f64,
|
||||
tokens: CodexUsageStatisticsTokens,
|
||||
costs: MoneySum,
|
||||
has_unknown_pricing: bool,
|
||||
has_estimated_allocation: bool,
|
||||
has_rate_fallback: bool,
|
||||
has_missing_token_data: bool,
|
||||
}
|
||||
|
||||
fn aggregate_models(
|
||||
days: &[DailyUsage],
|
||||
used_ratio: f64,
|
||||
is_current: bool,
|
||||
) -> Vec<CodexUsageStatisticsModel> {
|
||||
let mut aggregates =
|
||||
BTreeMap::<(String, CodexUsageStatisticsServiceTier), ModelAggregate>::new();
|
||||
for model in days.iter().flat_map(|day| &day.models) {
|
||||
if model.tokens.is_empty() && model.credits <= EPSILON {
|
||||
continue;
|
||||
}
|
||||
let aggregate = aggregates
|
||||
.entry((model.model.clone(), model.service_tier))
|
||||
.or_default();
|
||||
aggregate.credits += model.credits;
|
||||
aggregate.tokens =
|
||||
aggregate
|
||||
.tokens
|
||||
.checked_add(model.tokens)
|
||||
.unwrap_or(CodexUsageStatisticsTokens {
|
||||
uncached_input: u64::MAX,
|
||||
cached_input: u64::MAX,
|
||||
output: u64::MAX,
|
||||
total: u64::MAX,
|
||||
});
|
||||
aggregate.costs.push(model.estimated_cost);
|
||||
aggregate.has_unknown_pricing |= model.has_unknown_pricing;
|
||||
aggregate.has_estimated_allocation |= model.has_estimated_allocation;
|
||||
aggregate.has_rate_fallback |= model.has_rate_fallback;
|
||||
aggregate.has_missing_token_data |= model.has_missing_token_data;
|
||||
}
|
||||
let total_credits = aggregates
|
||||
.values()
|
||||
.map(|aggregate| aggregate.credits)
|
||||
.sum::<f64>();
|
||||
let mut rows = aggregates
|
||||
.into_iter()
|
||||
.map(|((model, service_tier), aggregate)| {
|
||||
let credit_share =
|
||||
(total_credits > EPSILON).then_some(aggregate.credits / total_credits);
|
||||
let quota_share = (is_current && used_ratio > EPSILON)
|
||||
.then(|| credit_share.map(|share| share * used_ratio))
|
||||
.flatten();
|
||||
CodexUsageStatisticsModel {
|
||||
key: format!("{model}::{}", service_tier.as_str()),
|
||||
model,
|
||||
service_tier,
|
||||
credit_share,
|
||||
quota_share,
|
||||
tokens: aggregate.tokens,
|
||||
estimated_cost: aggregate.costs.finish(),
|
||||
has_unknown_pricing: aggregate.has_unknown_pricing,
|
||||
has_estimated_allocation: aggregate.has_estimated_allocation,
|
||||
has_rate_fallback: aggregate.has_rate_fallback,
|
||||
has_missing_token_data: aggregate.has_missing_token_data,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(compare_model_rows);
|
||||
rows
|
||||
}
|
||||
|
||||
fn compare_model_rows(
|
||||
left: &CodexUsageStatisticsModel,
|
||||
right: &CodexUsageStatisticsModel,
|
||||
) -> Ordering {
|
||||
right
|
||||
.credit_share
|
||||
.unwrap_or_default()
|
||||
.total_cmp(&left.credit_share.unwrap_or_default())
|
||||
.then_with(|| cost_ticks(right.estimated_cost).cmp(&cost_ticks(left.estimated_cost)))
|
||||
.then_with(|| right.tokens.total.cmp(&left.tokens.total))
|
||||
.then_with(|| left.model.cmp(&right.model))
|
||||
.then_with(|| left.service_tier.cmp(&right.service_tier))
|
||||
}
|
||||
|
||||
fn cost_ticks(cost: Option<Money>) -> u128 {
|
||||
cost.map(|value| value.amount().scaled())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
|
||||
fn project_u64(value: u64, used_ratio: f64) -> Option<u64> {
|
||||
let projected = value as f64 / used_ratio;
|
||||
(projected.is_finite() && projected >= 0.0 && projected <= u64::MAX as f64)
|
||||
.then(|| projected.round() as u64)
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
|
||||
fn project_money(cost: Money, used_ratio: f64) -> Option<Money> {
|
||||
let projected = cost.amount().scaled() as f64 / used_ratio;
|
||||
if !projected.is_finite() || projected < 0.0 || projected > u128::MAX as f64 {
|
||||
return None;
|
||||
}
|
||||
let amount = Decimal::from_scaled(projected.round() as u128).ok()?;
|
||||
Some(Money::new(amount, cost.currency()))
|
||||
}
|
||||
|
||||
fn sum_tokens(
|
||||
values: impl IntoIterator<Item = CodexUsageStatisticsTokens>,
|
||||
) -> CodexUsageStatisticsTokens {
|
||||
values
|
||||
.into_iter()
|
||||
.fold(CodexUsageStatisticsTokens::default(), |total, value| {
|
||||
total
|
||||
.checked_add(value)
|
||||
.unwrap_or(CodexUsageStatisticsTokens {
|
||||
uncached_input: u64::MAX,
|
||||
cached_input: u64::MAX,
|
||||
output: u64::MAX,
|
||||
total: u64::MAX,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn rows_by_date(payload: &Value) -> BTreeMap<NaiveDate, Value> {
|
||||
extract_daily_list(payload)
|
||||
.into_iter()
|
||||
.filter_map(|row| Some((date_value(row.get("date")?)?, row.clone())))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn extract_daily_list(payload: &Value) -> Vec<&Value> {
|
||||
if let Some(rows) = payload.as_array() {
|
||||
return rows.iter().collect();
|
||||
}
|
||||
[
|
||||
"data",
|
||||
"items",
|
||||
"results",
|
||||
"daily",
|
||||
"daily_usage",
|
||||
"dailyWorkspaceUsageCounts",
|
||||
"daily_workspace_usage_counts",
|
||||
"workspace_usage_counts",
|
||||
]
|
||||
.into_iter()
|
||||
.find_map(|key| payload.get(key).and_then(Value::as_array))
|
||||
.map(|rows| rows.iter().collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn value_array(value: Option<&Value>) -> &[Value] {
|
||||
value.and_then(Value::as_array).map_or(&[], Vec::as_slice)
|
||||
}
|
||||
|
||||
fn token_parts(value: &Value) -> CodexUsageStatisticsTokens {
|
||||
let uncached_input = non_negative_u64(value.get("uncached_text_input_tokens"));
|
||||
let cached_input = non_negative_u64(value.get("cached_text_input_tokens"));
|
||||
let output = non_negative_u64(value.get("text_output_tokens"));
|
||||
let derived = uncached_input
|
||||
.saturating_add(cached_input)
|
||||
.saturating_add(output);
|
||||
let reported = non_negative_u64(value.get("text_total_tokens"));
|
||||
CodexUsageStatisticsTokens {
|
||||
uncached_input,
|
||||
cached_input,
|
||||
output,
|
||||
total: if reported > 0 { reported } else { derived },
|
||||
}
|
||||
}
|
||||
|
||||
fn day_credit_weight(value: &Value) -> f64 {
|
||||
let model_credits = value_array(value.get("models"))
|
||||
.iter()
|
||||
.map(|model| non_negative_f64(model.get("credits")))
|
||||
.sum::<f64>();
|
||||
if model_credits > EPSILON {
|
||||
return model_credits;
|
||||
}
|
||||
value
|
||||
.get("product_surface_usage_values")
|
||||
.and_then(Value::as_object)
|
||||
.map(|values| {
|
||||
values
|
||||
.values()
|
||||
.map(|value| non_negative_f64(Some(value)))
|
||||
.sum()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn model_name(value: &Value) -> String {
|
||||
["model", "model_id", "model_name", "name", "id"]
|
||||
.into_iter()
|
||||
.find_map(|key| value.get(key).and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("unknown")
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn model_service_tier(value: Option<&Value>) -> CodexUsageStatisticsServiceTier {
|
||||
match value
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.map(str::to_ascii_lowercase)
|
||||
.as_deref()
|
||||
{
|
||||
Some("fast" | "priority") => CodexUsageStatisticsServiceTier::Fast,
|
||||
Some(_) | None => CodexUsageStatisticsServiceTier::Standard,
|
||||
}
|
||||
}
|
||||
|
||||
fn date_value(value: &Value) -> Option<NaiveDate> {
|
||||
NaiveDate::parse_from_str(value.as_str()?.get(..10)?, "%Y-%m-%d").ok()
|
||||
}
|
||||
|
||||
fn non_negative_u64(value: Option<&Value>) -> u64 {
|
||||
let Some(value) = value else {
|
||||
return 0;
|
||||
};
|
||||
if let Some(number) = value.as_u64() {
|
||||
return number;
|
||||
}
|
||||
let normalized = value
|
||||
.as_str()
|
||||
.map(|raw| raw.replace(',', ""))
|
||||
.unwrap_or_else(|| value.to_string());
|
||||
normalized
|
||||
.trim()
|
||||
.parse::<u64>()
|
||||
.ok()
|
||||
.or_else(|| {
|
||||
normalized
|
||||
.trim()
|
||||
.parse::<f64>()
|
||||
.ok()
|
||||
.filter(|number| number.is_finite() && *number >= 0.0)
|
||||
.map(|number| number.round() as u64)
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn non_negative_f64(value: Option<&Value>) -> f64 {
|
||||
let Some(value) = value else {
|
||||
return 0.0;
|
||||
};
|
||||
let parsed = value.as_f64().or_else(|| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|raw| raw.replace(',', ""))
|
||||
.and_then(|raw| raw.trim().parse::<f64>().ok())
|
||||
});
|
||||
parsed
|
||||
.filter(|number| number.is_finite() && *number >= 0.0)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
@@ -20,8 +20,8 @@ use thiserror::Error;
|
||||
|
||||
use super::recovery_log::{CodexOAuthRecoveryOperation, record_oauth_recovery};
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
use super::token_client::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
use super::types::CodexOAuthSecret;
|
||||
use super::token_client::{RefreshFailure, RefreshUpstreamFailure, TokenPair, TokenRefresher};
|
||||
use super::types::{CodexOAuthSecret, parse_access_token_expiration};
|
||||
|
||||
const PROVIDER_NAME: &str = "openai";
|
||||
const MAX_REFRESH_BATCH: u32 = 1_000;
|
||||
@@ -32,7 +32,7 @@ const REFRESH_BACKOFF_MAX_ATTEMPTS: u32 = 5;
|
||||
/// 耗尽指数退避后的 OAuth 恢复周期。
|
||||
const REFRESH_RECOVERY_DELAY: Duration = Duration::from_secs(10 * 60);
|
||||
/// 过期 AT 仍允许 RT 恢复的最长窗口。
|
||||
const REFRESH_RECOVERY_WINDOW: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
const REFRESH_RECOVERY_WINDOW: Duration = Duration::from_secs(2 * 60 * 60);
|
||||
|
||||
pub(crate) fn refresh_recovery_deadline(
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
@@ -91,12 +91,19 @@ fn log_refresh_deferred(
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
attempt: u32,
|
||||
reason: &'static str,
|
||||
upstream_message: Option<&str>,
|
||||
upstream: Option<&RefreshUpstreamFailure>,
|
||||
retry_at: SystemTime,
|
||||
) {
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
attempt,
|
||||
reason,
|
||||
upstream_message = ?upstream_message,
|
||||
upstream_status = ?upstream.map(RefreshUpstreamFailure::status),
|
||||
upstream_code = ?upstream.and_then(RefreshUpstreamFailure::code),
|
||||
upstream_type = ?upstream.and_then(RefreshUpstreamFailure::error_type),
|
||||
upstream_body = ?upstream.map(RefreshUpstreamFailure::body),
|
||||
retry_at = %DateTime::<Utc>::from(retry_at),
|
||||
access_token_expires_at = ?access_token_expires_at.map(DateTime::<Utc>::from),
|
||||
recovery_deadline = ?refresh_recovery_deadline(access_token_expires_at)
|
||||
@@ -105,6 +112,26 @@ fn log_refresh_deferred(
|
||||
);
|
||||
}
|
||||
|
||||
struct RefreshFailureContext<'a> {
|
||||
reason: &'static str,
|
||||
message: Option<String>,
|
||||
upstream: Option<&'a RefreshUpstreamFailure>,
|
||||
}
|
||||
|
||||
impl<'a> RefreshFailureContext<'a> {
|
||||
const fn new(
|
||||
reason: &'static str,
|
||||
message: Option<String>,
|
||||
upstream: Option<&'a RefreshUpstreamFailure>,
|
||||
) -> Self {
|
||||
Self {
|
||||
reason,
|
||||
message,
|
||||
upstream,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_due_at(
|
||||
account: &ProviderAccount,
|
||||
policy: ProviderRefreshPolicy,
|
||||
@@ -248,20 +275,10 @@ impl CodexCredentialRefreshService {
|
||||
policy: ProviderRefreshPolicy,
|
||||
) -> Result<CodexCredentialRefreshOutcome, CodexCredentialRefreshError> {
|
||||
let account_id = due.account.id().to_string();
|
||||
if refresh_recovery_window_exhausted(
|
||||
let recovery_window_exhausted = refresh_recovery_window_exhausted(
|
||||
due.account.access_token_expires_at(),
|
||||
SystemTime::now(),
|
||||
) {
|
||||
return self
|
||||
.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Expired,
|
||||
AccountErrorReason::CredentialExpired,
|
||||
"refresh_recovery_window_exhausted",
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await;
|
||||
}
|
||||
);
|
||||
let capacity = self
|
||||
.leases
|
||||
.try_acquire(ProviderLeaseRequest::RefreshCapacity(
|
||||
@@ -291,31 +308,72 @@ impl CodexCredentialRefreshService {
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?;
|
||||
match self.refresher.refresh(refresh_token.expose_secret()).await {
|
||||
let refresh_result = self.refresher.refresh(refresh_token.expose_secret()).await;
|
||||
if recovery_window_exhausted && let Err(failure) = &refresh_result {
|
||||
let message = failure.message().map(str::to_owned);
|
||||
let upstream = failure.upstream();
|
||||
return match failure {
|
||||
RefreshFailure::Banned { .. } => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Banned,
|
||||
AccountErrorReason::AccountBanned,
|
||||
RefreshFailureContext::new(
|
||||
"refresh_recovery_window_exhausted",
|
||||
message,
|
||||
upstream,
|
||||
),
|
||||
CodexCredentialRefreshOutcome::Banned { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
RefreshFailure::InvalidGrant { .. }
|
||||
| RefreshFailure::RetryableTransport { .. }
|
||||
| RefreshFailure::Transport { .. } => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Expired,
|
||||
AccountErrorReason::CredentialExpired,
|
||||
RefreshFailureContext::new(
|
||||
"refresh_recovery_window_exhausted",
|
||||
message,
|
||||
upstream,
|
||||
),
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
};
|
||||
}
|
||||
match refresh_result {
|
||||
Ok(tokens) => self.persist_success(due, tokens).await,
|
||||
Err(RefreshFailure::InvalidGrant) => {
|
||||
Err(RefreshFailure::InvalidGrant { message, upstream }) => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Expired,
|
||||
AccountErrorReason::CredentialExpired,
|
||||
"refresh_invalid_grant",
|
||||
RefreshFailureContext::new(
|
||||
"refresh_invalid_grant",
|
||||
message,
|
||||
upstream.as_deref(),
|
||||
),
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(RefreshFailure::Banned) => {
|
||||
Err(RefreshFailure::Banned { message, upstream }) => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Banned,
|
||||
AccountErrorReason::AccountBanned,
|
||||
"account_banned",
|
||||
RefreshFailureContext::new("account_banned", message, upstream.as_deref()),
|
||||
CodexCredentialRefreshOutcome::Banned { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(RefreshFailure::RetryableTransport) => {
|
||||
Err(RefreshFailure::RetryableTransport { message }) => {
|
||||
if self
|
||||
.defer_refresh(&due.account, "transport-not-sent")
|
||||
.defer_refresh(&due.account, "transport-not-sent", Some(&message), None)
|
||||
.await?
|
||||
{
|
||||
Ok(CodexCredentialRefreshOutcome::Transient { account_id })
|
||||
@@ -323,11 +381,16 @@ impl CodexCredentialRefreshService {
|
||||
Ok(CodexCredentialRefreshOutcome::Stale { account_id })
|
||||
}
|
||||
}
|
||||
Err(RefreshFailure::Transport) => {
|
||||
// 上游瞬态(429/5xx/超时/畸形响应等,发送状态未知)保留现有凭据、
|
||||
// 推进退避重试;仅 invalid_grant/banned 才终态失效账号。
|
||||
Err(RefreshFailure::Transport { message, upstream }) => {
|
||||
// 上游瞬态(401/429/5xx/超时/畸形响应等)保留现有凭据、
|
||||
// 记录最近一次失败并推进有界退避。
|
||||
if self
|
||||
.defer_refresh(&due.account, "transport-ambiguous")
|
||||
.defer_refresh(
|
||||
&due.account,
|
||||
"transport-ambiguous",
|
||||
message.as_deref(),
|
||||
upstream.as_deref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Ok(CodexCredentialRefreshOutcome::Transient { account_id })
|
||||
@@ -403,9 +466,11 @@ impl CodexCredentialRefreshService {
|
||||
access_token,
|
||||
refresh_token: rotated_refresh_token,
|
||||
id_token: rotated_id_token,
|
||||
expires_in,
|
||||
} = tokens;
|
||||
let access_token_rotated = access_token.is_some();
|
||||
let access_token_expires_at = match access_token.as_deref() {
|
||||
Some(access_token) => parse_access_token_expiration(access_token).map(SystemTime::from),
|
||||
None => due.account.access_token_expires_at(),
|
||||
};
|
||||
let CodexOAuthSecret {
|
||||
access_token: current_access_token,
|
||||
refresh_token: current_refresh_token,
|
||||
@@ -432,12 +497,6 @@ impl CodexCredentialRefreshService {
|
||||
.as_ref()
|
||||
.map(ExposeSecret::expose_secret),
|
||||
);
|
||||
let observed_at = SystemTime::now();
|
||||
let access_token_expires_at = if access_token_rotated {
|
||||
expires_in.and_then(|value| observed_at.checked_add(value))
|
||||
} else {
|
||||
due.account.access_token_expires_at()
|
||||
};
|
||||
// 成功轮换清掉仅用于失败退避的 retry-not-before;正常预刷新窗口由
|
||||
// worker 结合当前 runtime policy 动态判断,不写入账号时间字段。
|
||||
let next_refresh_at = None;
|
||||
@@ -483,9 +542,14 @@ impl CodexCredentialRefreshService {
|
||||
account: &ProviderAccount,
|
||||
credential_state: CredentialState,
|
||||
error_reason: AccountErrorReason,
|
||||
reason: &'static str,
|
||||
failure: RefreshFailureContext<'_>,
|
||||
outcome: CodexCredentialRefreshOutcome,
|
||||
) -> Result<CodexCredentialRefreshOutcome, CodexCredentialRefreshError> {
|
||||
let RefreshFailureContext {
|
||||
reason,
|
||||
message,
|
||||
upstream,
|
||||
} = failure;
|
||||
match self.repository.load_runtime_credential(account).await {
|
||||
Ok(_) => {}
|
||||
Err(CredentialRepositoryError::RevisionConflict) => {
|
||||
@@ -495,6 +559,12 @@ impl CodexCredentialRefreshService {
|
||||
}
|
||||
Err(error) => return Err(error.into()),
|
||||
}
|
||||
let persisted_message = message
|
||||
.as_ref()
|
||||
.filter(|message| !message.trim().is_empty())
|
||||
.map(String::as_str)
|
||||
.or_else(|| account.last_error_message())
|
||||
.map(str::to_owned);
|
||||
match self
|
||||
.repository
|
||||
.apply_state_with_reason(
|
||||
@@ -502,7 +572,7 @@ impl CodexCredentialRefreshService {
|
||||
credential_state,
|
||||
SystemTime::now(),
|
||||
Some(error_reason),
|
||||
Some(reason.to_owned()),
|
||||
persisted_message,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -511,6 +581,11 @@ impl CodexCredentialRefreshService {
|
||||
account_id = %account.id(),
|
||||
?credential_state,
|
||||
reason,
|
||||
upstream_message = ?message.as_deref(),
|
||||
upstream_status = ?upstream.map(RefreshUpstreamFailure::status),
|
||||
upstream_code = ?upstream.and_then(RefreshUpstreamFailure::code),
|
||||
upstream_type = ?upstream.and_then(RefreshUpstreamFailure::error_type),
|
||||
upstream_body = ?upstream.map(RefreshUpstreamFailure::body),
|
||||
access_token_expires_at = ?account.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from),
|
||||
recovery_deadline = ?refresh_recovery_deadline(account.access_token_expires_at())
|
||||
@@ -532,6 +607,8 @@ impl CodexCredentialRefreshService {
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
reason: &'static str,
|
||||
upstream_message: Option<&str>,
|
||||
upstream: Option<&RefreshUpstreamFailure>,
|
||||
) -> Result<bool, CodexCredentialRefreshError> {
|
||||
let attempt = self
|
||||
.credential_state
|
||||
@@ -545,13 +622,28 @@ impl CodexCredentialRefreshService {
|
||||
attempt,
|
||||
reason,
|
||||
)?;
|
||||
match self.repository.defer_refresh(account, retry_at).await {
|
||||
let message = upstream_message
|
||||
.filter(|message| !message.trim().is_empty())
|
||||
.map(str::to_owned);
|
||||
let error_reason = upstream.map(|_| {
|
||||
account
|
||||
.credential_state()
|
||||
.error_reason()
|
||||
.unwrap_or(AccountErrorReason::AccessTokenExpired)
|
||||
});
|
||||
match self
|
||||
.repository
|
||||
.defer_refresh(account, retry_at, error_reason, message)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
log_refresh_deferred(
|
||||
account.id(),
|
||||
account.access_token_expires_at(),
|
||||
attempt,
|
||||
reason,
|
||||
upstream_message,
|
||||
upstream,
|
||||
retry_at,
|
||||
);
|
||||
Ok(true)
|
||||
|
||||
@@ -120,15 +120,18 @@ impl CodexCredentialRepository {
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
)
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?
|
||||
.with_account_state(CredentialState::Ready, SystemTime::now(), None, None);
|
||||
cas_revision(self.store.compare_and_swap_credential(update).await?)
|
||||
}
|
||||
|
||||
/// token 签名边界不可用时,以相同 credential 推进持久刷新退避。
|
||||
/// 以相同 credential 原子推进刷新退避及本次上游错误事实。
|
||||
pub async fn defer_refresh(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
next_refresh_at: SystemTime,
|
||||
error_reason: Option<AccountErrorReason>,
|
||||
message: Option<String>,
|
||||
) -> Result<CredentialRevision, CredentialRepositoryError> {
|
||||
let current = self
|
||||
.store
|
||||
@@ -140,7 +143,7 @@ impl CodexCredentialRepository {
|
||||
let data = CodexCredentialCodec::decode_complete(¤t.credential)?;
|
||||
let has_refresh_token = data.has_refresh_token();
|
||||
let credential = CodexCredentialCodec::encode_complete(data)?;
|
||||
let update = CredentialCasUpdate::new(
|
||||
let mut update = CredentialCasUpdate::new(
|
||||
account.id().clone(),
|
||||
account.revision(),
|
||||
unchanged_profile(account),
|
||||
@@ -150,6 +153,14 @@ impl CodexCredentialRepository {
|
||||
Some(next_refresh_at),
|
||||
)
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
|
||||
if let Some(error_reason) = error_reason {
|
||||
update = update.with_account_state(
|
||||
account.credential_state(),
|
||||
SystemTime::now(),
|
||||
Some(error_reason),
|
||||
message,
|
||||
);
|
||||
}
|
||||
cas_revision(self.store.compare_and_swap_credential(update).await?)
|
||||
}
|
||||
|
||||
@@ -245,13 +256,11 @@ impl CodexCredentialRepository {
|
||||
if !account.enabled() {
|
||||
return Ok(());
|
||||
}
|
||||
let (error_reason, message) = if credential_state == CredentialState::Ready {
|
||||
(None, None)
|
||||
let message = message.filter(|value| !value.trim().is_empty());
|
||||
let error_reason = if credential_state == CredentialState::Ready {
|
||||
message.as_ref().and(error_reason)
|
||||
} else {
|
||||
(
|
||||
error_reason.or_else(|| credential_state.error_reason()),
|
||||
message.filter(|value| !value.trim().is_empty()),
|
||||
)
|
||||
error_reason.or_else(|| credential_state.error_reason())
|
||||
};
|
||||
self.store
|
||||
.apply_state_change(AccountStateChange {
|
||||
|
||||
@@ -95,6 +95,17 @@ pub struct SelectCodexCredential<'a> {
|
||||
pub session_affinity_key: Option<&'a ProviderSessionAffinityKey>,
|
||||
}
|
||||
|
||||
pub(crate) struct SelectCodexProviderEndpointCredential<'a> {
|
||||
pub request_url: &'a Url,
|
||||
pub attempt: &'a AttemptContext,
|
||||
}
|
||||
|
||||
struct CredentialSelectionInput<'a> {
|
||||
request_url: &'a Url,
|
||||
attempt: &'a AttemptContext,
|
||||
session_affinity_key: Option<&'a ProviderSessionAffinityKey>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CodexCyberPolicyScope {
|
||||
key: ProviderSessionAffinityKey,
|
||||
@@ -240,6 +251,12 @@ struct AffinityTelemetry {
|
||||
account_switch: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ModelCatalogEligibility<'a> {
|
||||
Required(&'a str),
|
||||
NotApplicable,
|
||||
}
|
||||
|
||||
impl CodexCredentialSelector {
|
||||
#[must_use]
|
||||
// 选择器显式持有各能力边界,避免把 Provider 私有服务重新包装成通用容器。
|
||||
@@ -275,13 +292,59 @@ impl CodexCredentialSelector {
|
||||
&self,
|
||||
request: &SelectCodexCredential<'_>,
|
||||
) -> Result<CodexCredentialLease, CredentialSelectionError> {
|
||||
self.select_with_cyber_policy(request, None).await
|
||||
let input = CredentialSelectionInput {
|
||||
request_url: request.request_url,
|
||||
attempt: request.attempt,
|
||||
session_affinity_key: request.session_affinity_key,
|
||||
};
|
||||
self.select_inner(
|
||||
&input,
|
||||
None,
|
||||
ModelCatalogEligibility::Required(request.upstream_model),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn select_with_cyber_policy(
|
||||
&self,
|
||||
request: &SelectCodexCredential<'_>,
|
||||
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
|
||||
) -> Result<CodexCredentialLease, CredentialSelectionError> {
|
||||
let input = CredentialSelectionInput {
|
||||
request_url: request.request_url,
|
||||
attempt: request.attempt,
|
||||
session_affinity_key: request.session_affinity_key,
|
||||
};
|
||||
self.select_inner(
|
||||
&input,
|
||||
cyber_policy_session_key,
|
||||
ModelCatalogEligibility::Required(request.upstream_model),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 为不属于 Responses 文本模型目录的 Provider 原生端点选择账号。
|
||||
///
|
||||
/// 账号范围、健康度、配额、并发租约、cookie 与认证准备仍走同一套选择链路;
|
||||
/// 唯一差异是不能用文本模型目录成员关系淘汰账号。
|
||||
pub(crate) async fn select_for_provider_endpoint(
|
||||
&self,
|
||||
request: &SelectCodexProviderEndpointCredential<'_>,
|
||||
) -> Result<CodexCredentialLease, CredentialSelectionError> {
|
||||
let input = CredentialSelectionInput {
|
||||
request_url: request.request_url,
|
||||
attempt: request.attempt,
|
||||
session_affinity_key: None,
|
||||
};
|
||||
self.select_inner(&input, None, ModelCatalogEligibility::NotApplicable)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn select_inner(
|
||||
&self,
|
||||
request: &CredentialSelectionInput<'_>,
|
||||
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
|
||||
model_catalog_eligibility: ModelCatalogEligibility<'_>,
|
||||
) -> Result<CodexCredentialLease, CredentialSelectionError> {
|
||||
let diagnostic = request.attempt.is_diagnostic_required_account();
|
||||
let accounts = self.repository.list_for_provider().await?;
|
||||
@@ -294,12 +357,15 @@ impl CodexCredentialSelector {
|
||||
.attempt
|
||||
.account_scope()
|
||||
.is_some_and(|scope| scope.allows(account.id())))
|
||||
&& (diagnostic || {
|
||||
let observed_support = self
|
||||
.catalog
|
||||
.observed_model_support(account, request.upstream_model);
|
||||
matches!(observed_support, Ok(None | Some(true)))
|
||||
})
|
||||
&& (diagnostic
|
||||
|| match model_catalog_eligibility {
|
||||
ModelCatalogEligibility::NotApplicable => true,
|
||||
ModelCatalogEligibility::Required(upstream_model) => {
|
||||
let observed_support =
|
||||
self.catalog.observed_model_support(account, upstream_model);
|
||||
matches!(observed_support, Ok(None | Some(true)))
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if !diagnostic {
|
||||
@@ -410,14 +476,16 @@ impl CodexCredentialSelector {
|
||||
{
|
||||
affinity.escape(AffinityEscapeReason::PinnedAccount);
|
||||
}
|
||||
let preferred = pinned_account
|
||||
.clone()
|
||||
.or_else(|| affinity.preferred_account().cloned());
|
||||
let mut observed_affinity_account = affinity.bound_account().cloned();
|
||||
let mut shortest_retry = None;
|
||||
let policy = request.attempt.account_selection_policy();
|
||||
|
||||
loop {
|
||||
let preferred = pinned_account
|
||||
.clone()
|
||||
.or_else(|| affinity.preferred_account().cloned());
|
||||
let context = AccountSelectionContext {
|
||||
policy: request.attempt.account_selection_policy(),
|
||||
policy,
|
||||
now: SystemTime::now(),
|
||||
excluded_accounts: excluded.clone(),
|
||||
preferred_account: preferred.clone(),
|
||||
@@ -445,7 +513,6 @@ impl CodexCredentialSelector {
|
||||
.map(|candidate| candidate.account.clone())
|
||||
.ok_or(CredentialSelectionError::InvalidCredential)?;
|
||||
let allows_account_state_mutation = !diagnostic || account.enabled();
|
||||
let policy = request.attempt.account_selection_policy();
|
||||
match self
|
||||
.leases
|
||||
.try_acquire(ProviderLeaseRequest::Scheduling(
|
||||
@@ -466,6 +533,40 @@ impl CodexCredentialSelector {
|
||||
excluded.insert(account.id().clone());
|
||||
}
|
||||
ProviderLeaseAcquisition::Acquired(guard) => {
|
||||
let initial_affinity_claim = if !diagnostic
|
||||
&& observed_affinity_account.is_none()
|
||||
&& let Some(key) = request.session_affinity_key
|
||||
{
|
||||
self.claim_initial_session_affinity(key, account.id()).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// 原生 continuation/required account 比根亲和绑定更严格;它可以
|
||||
// 使用 owner 账号,但不能把已经迁移的根会话拉回旧号。
|
||||
if pinned_account.is_none()
|
||||
&& let Some(effective_account) = initial_affinity_claim
|
||||
&& &effective_account != account.id()
|
||||
{
|
||||
drop(guard);
|
||||
observed_affinity_account = Some(effective_account.clone());
|
||||
affinity = affinity_selection_for_bound_account(
|
||||
effective_account,
|
||||
&candidates,
|
||||
SystemTime::now(),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let affinity_expected_account_id = if pinned_account
|
||||
.as_ref()
|
||||
.zip(observed_affinity_account.as_ref())
|
||||
.is_some_and(|(pinned, bound)| pinned != bound)
|
||||
{
|
||||
account.id().clone()
|
||||
} else {
|
||||
observed_affinity_account
|
||||
.clone()
|
||||
.unwrap_or_else(|| account.id().clone())
|
||||
};
|
||||
let affinity_telemetry = affinity.telemetry(account.id());
|
||||
tracing::info!(
|
||||
request_id = %request.attempt.request_id(),
|
||||
@@ -517,6 +618,7 @@ impl CodexCredentialSelector {
|
||||
cyber_policy_scope,
|
||||
allows_account_state_mutation,
|
||||
affinity_telemetry,
|
||||
affinity_expected_account_id,
|
||||
_guard: guard,
|
||||
});
|
||||
}
|
||||
@@ -537,24 +639,7 @@ impl CodexCredentialSelector {
|
||||
else {
|
||||
return AffinitySelection::default();
|
||||
};
|
||||
let Some(candidate) = candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.account.id() == &account_id)
|
||||
else {
|
||||
self.clear_session_affinity(key).await;
|
||||
return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable);
|
||||
};
|
||||
if candidate
|
||||
.account
|
||||
.status_projection(now, candidate.signals.rate_limited_until)
|
||||
.status
|
||||
== AccountStatus::Normal
|
||||
{
|
||||
return AffinitySelection::preferred(account_id);
|
||||
}
|
||||
let reason = affinity_unavailable_reason(candidate, now);
|
||||
self.clear_session_affinity(key).await;
|
||||
AffinitySelection::escaped(account_id, reason)
|
||||
affinity_selection_for_bound_account(account_id, candidates, now)
|
||||
}
|
||||
|
||||
async fn lookup_session_affinity(
|
||||
@@ -583,22 +668,38 @@ impl CodexCredentialSelector {
|
||||
}
|
||||
}
|
||||
|
||||
async fn clear_session_affinity(&self, key: &ProviderSessionAffinityKey) {
|
||||
async fn claim_initial_session_affinity(
|
||||
&self,
|
||||
key: &ProviderSessionAffinityKey,
|
||||
selected_account_id: &ProviderAccountId,
|
||||
) -> Option<ProviderAccountId> {
|
||||
match tokio::time::timeout(
|
||||
SESSION_AFFINITY_TIMEOUT,
|
||||
self.session_affinity.clear(&self.provider_kind, key),
|
||||
self.session_affinity.claim_or_load(
|
||||
&self.provider_kind,
|
||||
key,
|
||||
selected_account_id,
|
||||
SESSION_AFFINITY_TTL,
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(_)) => {}
|
||||
Ok(Ok(account_id)) => Some(account_id),
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(error = %error, "OpenAI stale session affinity clear failed open");
|
||||
tracing::warn!(
|
||||
account_id = %selected_account_id,
|
||||
error = %error,
|
||||
"OpenAI initial session affinity claim failed open"
|
||||
);
|
||||
None
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
account_id = %selected_account_id,
|
||||
timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(),
|
||||
"OpenAI stale session affinity clear timed out"
|
||||
"OpenAI initial session affinity claim timed out"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -720,7 +821,15 @@ impl CodexCredentialSelector {
|
||||
.map(chrono::DateTime::<chrono::Utc>::from),
|
||||
"OpenAI access token expired; retaining account for bounded OAuth refresh recovery"
|
||||
);
|
||||
return Ok(());
|
||||
return self
|
||||
.apply_credential_state(
|
||||
account,
|
||||
CredentialState::Ready,
|
||||
AccountErrorReason::AccessTokenExpired,
|
||||
now,
|
||||
message,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
self.apply_credential_state(
|
||||
account,
|
||||
@@ -827,6 +936,7 @@ impl CodexCredentialSelector {
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
session_affinity_key: Option<&ProviderSessionAffinityKey>,
|
||||
expected_affinity_account_id: &ProviderAccountId,
|
||||
) {
|
||||
self.restore_recoverable_account_state(account).await;
|
||||
self.risk_recovery
|
||||
@@ -836,24 +946,26 @@ impl CodexCredentialSelector {
|
||||
let Some(key) = session_affinity_key else {
|
||||
return;
|
||||
};
|
||||
if matches!(
|
||||
self.lookup_session_affinity(key).await,
|
||||
SessionAffinityLookup::Unavailable
|
||||
) {
|
||||
return;
|
||||
}
|
||||
match tokio::time::timeout(
|
||||
SESSION_AFFINITY_TIMEOUT,
|
||||
self.session_affinity.bind(
|
||||
self.session_affinity.compare_and_bind(
|
||||
&self.provider_kind,
|
||||
key,
|
||||
expected_affinity_account_id,
|
||||
account.id(),
|
||||
SESSION_AFFINITY_TTL,
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Ok(effective_account)) if &effective_account == account.id() => {}
|
||||
Ok(Ok(effective_account)) => {
|
||||
tracing::debug!(
|
||||
account_id = %account.id(),
|
||||
effective_account_id = %effective_account,
|
||||
"OpenAI late success left the newer session affinity binding unchanged"
|
||||
);
|
||||
}
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
@@ -1030,6 +1142,32 @@ impl CodexCredentialSelector {
|
||||
}
|
||||
}
|
||||
|
||||
fn affinity_selection_for_bound_account(
|
||||
account_id: ProviderAccountId,
|
||||
candidates: &[AccountCandidate],
|
||||
now: SystemTime,
|
||||
) -> AffinitySelection {
|
||||
let Some(candidate) = candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.account.id() == &account_id)
|
||||
else {
|
||||
return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable);
|
||||
};
|
||||
match candidate
|
||||
.account
|
||||
.status_projection(now, candidate.signals.rate_limited_until)
|
||||
.status
|
||||
{
|
||||
AccountStatus::Normal => AffinitySelection::preferred(account_id),
|
||||
AccountStatus::QuotaExhausted
|
||||
| AccountStatus::RateLimited
|
||||
| AccountStatus::Disabled
|
||||
| AccountStatus::Error => {
|
||||
AffinitySelection::escaped(account_id, affinity_unavailable_reason(candidate, now))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn affinity_unavailable_reason(
|
||||
candidate: &AccountCandidate,
|
||||
now: SystemTime,
|
||||
@@ -1085,6 +1223,7 @@ pub struct CodexCredentialLease {
|
||||
cyber_policy_scope: Option<CodexCyberPolicyScope>,
|
||||
allows_account_state_mutation: bool,
|
||||
affinity_telemetry: AffinityTelemetry,
|
||||
affinity_expected_account_id: ProviderAccountId,
|
||||
_guard: Box<dyn ProviderLeaseGuard>,
|
||||
}
|
||||
|
||||
@@ -1142,6 +1281,12 @@ impl CodexCredentialLease {
|
||||
pub const fn account_switch(&self) -> bool {
|
||||
self.affinity_telemetry.account_switch
|
||||
}
|
||||
|
||||
/// 成功反馈只能从选择时观察到的绑定迁移,避免迟到请求覆盖较新 winner。
|
||||
#[must_use]
|
||||
pub(crate) const fn affinity_expected_account_id(&self) -> &ProviderAccountId {
|
||||
&self.affinity_expected_account_id
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CodexCredentialLease {
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use async_trait::async_trait;
|
||||
use reqwest::{Client, StatusCode, redirect::Policy};
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fmt;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -27,9 +27,6 @@ pub struct TokenPair {
|
||||
pub refresh_token: Option<String>,
|
||||
/// 官方刷新响应省略时由持久化调用方保留当前 ID token。
|
||||
pub id_token: Option<String>,
|
||||
/// OAuth endpoint 可以省略或返回零 expiry;这不影响已经成功的 token
|
||||
/// exchange,只意味着调用方不能据此安排下一次刷新。
|
||||
pub expires_in: Option<Duration>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for TokenPair {
|
||||
@@ -45,22 +42,131 @@ impl fmt::Debug for TokenPair {
|
||||
&self.refresh_token.as_ref().map(|_| "[REDACTED]"),
|
||||
)
|
||||
.field("id_token", &self.id_token.as_ref().map(|_| "[REDACTED]"))
|
||||
.field("expires_in", &self.expires_in)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex token 刷新的稳定失败分类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
||||
#[derive(Clone, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum RefreshFailure {
|
||||
#[error("refresh token is invalid or expired")]
|
||||
InvalidGrant,
|
||||
InvalidGrant {
|
||||
message: Option<String>,
|
||||
upstream: Option<Box<RefreshUpstreamFailure>>,
|
||||
},
|
||||
#[error("account is banned")]
|
||||
Banned,
|
||||
Banned {
|
||||
message: Option<String>,
|
||||
upstream: Option<Box<RefreshUpstreamFailure>>,
|
||||
},
|
||||
#[error("refresh transport failed before server processing")]
|
||||
RetryableTransport,
|
||||
RetryableTransport { message: String },
|
||||
#[error("refresh transport failed after possible server processing")]
|
||||
Transport,
|
||||
Transport {
|
||||
message: Option<String>,
|
||||
upstream: Option<Box<RefreshUpstreamFailure>>,
|
||||
},
|
||||
}
|
||||
|
||||
impl fmt::Debug for RefreshFailure {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RefreshFailure")
|
||||
.field("classification", &self.classification())
|
||||
.field("message", &self.message().map(|_| "<redacted>"))
|
||||
.field("upstream", &self.upstream())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl RefreshFailure {
|
||||
#[must_use]
|
||||
pub fn message(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::InvalidGrant { message, .. }
|
||||
| Self::Banned { message, .. }
|
||||
| Self::Transport { message, .. } => message.as_deref(),
|
||||
Self::RetryableTransport { message } => Some(message),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn upstream(&self) -> Option<&RefreshUpstreamFailure> {
|
||||
match self {
|
||||
Self::InvalidGrant { upstream, .. }
|
||||
| Self::Banned { upstream, .. }
|
||||
| Self::Transport { upstream, .. } => upstream.as_deref(),
|
||||
Self::RetryableTransport { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn classification(&self) -> &'static str {
|
||||
match self {
|
||||
Self::InvalidGrant { .. } => "invalid-grant",
|
||||
Self::Banned { .. } => "account-banned",
|
||||
Self::RetryableTransport { .. } => "transport-not-sent",
|
||||
Self::Transport { .. } => "transport-ambiguous",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 当前 OAuth 刷新请求收到的完整非成功响应。
|
||||
///
|
||||
/// 该值不持久化;`Debug` 不输出正文。调用方仅在受控刷新失败日志中显式记录正文。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct RefreshUpstreamFailure {
|
||||
status: u16,
|
||||
code: Option<String>,
|
||||
error_type: Option<String>,
|
||||
body: String,
|
||||
}
|
||||
|
||||
impl RefreshUpstreamFailure {
|
||||
fn new(status: StatusCode, body: &[u8], error: Option<&RefreshErrorResponse>) -> Self {
|
||||
Self {
|
||||
status: status.as_u16(),
|
||||
code: error
|
||||
.and_then(RefreshErrorResponse::code)
|
||||
.map(str::to_owned),
|
||||
error_type: error
|
||||
.and_then(RefreshErrorResponse::error_type)
|
||||
.map(str::to_owned),
|
||||
body: String::from_utf8_lossy(body).into_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn status(&self) -> u16 {
|
||||
self.status
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn code(&self) -> Option<&str> {
|
||||
self.code.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn error_type(&self) -> Option<&str> {
|
||||
self.error_type.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn body(&self) -> &str {
|
||||
&self.body
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for RefreshUpstreamFailure {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RefreshUpstreamFailure")
|
||||
.field("status", &self.status)
|
||||
.field("code", &self.code)
|
||||
.field("error_type", &self.error_type)
|
||||
.field("body", &"<redacted>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex token 刷新端口。
|
||||
@@ -92,7 +198,6 @@ impl fmt::Debug for AuthorizationCodeGrant {
|
||||
pub struct AuthorizationTokenSet {
|
||||
pub secret: crate::credential::CodexOAuthSecret,
|
||||
pub id_token: SecretString,
|
||||
pub expires_in: Option<Duration>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for AuthorizationTokenSet {
|
||||
@@ -101,7 +206,6 @@ impl fmt::Debug for AuthorizationTokenSet {
|
||||
.debug_struct("AuthorizationTokenSet")
|
||||
.field("secret", &"[REDACTED]")
|
||||
.field("id_token", &"[REDACTED]")
|
||||
.field("expires_in", &self.expires_in)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -183,7 +287,60 @@ struct RefreshTokenResponse {
|
||||
access_token: Option<String>,
|
||||
refresh_token: Option<String>,
|
||||
id_token: Option<String>,
|
||||
expires_in: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct RefreshTokenRequest<'a> {
|
||||
client_id: &'a str,
|
||||
grant_type: &'static str,
|
||||
refresh_token: &'a str,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefreshErrorResponse {
|
||||
error: Option<RefreshError>,
|
||||
code: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum RefreshError {
|
||||
Details(RefreshErrorDetails),
|
||||
Code(String),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefreshErrorDetails {
|
||||
code: Option<String>,
|
||||
message: Option<String>,
|
||||
#[serde(rename = "type")]
|
||||
error_type: Option<String>,
|
||||
}
|
||||
|
||||
impl RefreshErrorResponse {
|
||||
fn code(&self) -> Option<&str> {
|
||||
match self.error.as_ref() {
|
||||
Some(RefreshError::Details(error)) => error.code.as_deref(),
|
||||
Some(RefreshError::Code(code)) => Some(code.as_str()),
|
||||
None => None,
|
||||
}
|
||||
.or(self.code.as_deref())
|
||||
}
|
||||
|
||||
fn message(&self) -> Option<String> {
|
||||
match self.error.as_ref() {
|
||||
Some(RefreshError::Details(error)) => error.message.as_deref(),
|
||||
Some(RefreshError::Code(_)) | None => None,
|
||||
}
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn error_type(&self) -> Option<&str> {
|
||||
match self.error.as_ref() {
|
||||
Some(RefreshError::Details(error)) => error.error_type.as_deref(),
|
||||
Some(RefreshError::Code(_)) | None => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -191,7 +348,6 @@ struct AuthorizationCodeResponse {
|
||||
access_token: String,
|
||||
refresh_token: String,
|
||||
id_token: String,
|
||||
expires_in: Option<u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -200,11 +356,11 @@ impl TokenRefresher for OpenAiTokenClient {
|
||||
let response = self
|
||||
.client
|
||||
.post(&self.config.token_endpoint)
|
||||
.form(&[
|
||||
("grant_type", "refresh_token"),
|
||||
("client_id", self.config.client_id.as_str()),
|
||||
("refresh_token", refresh_token),
|
||||
])
|
||||
.json(&RefreshTokenRequest {
|
||||
client_id: self.config.client_id.as_str(),
|
||||
grant_type: "refresh_token",
|
||||
refresh_token,
|
||||
})
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| refresh_transport_failure(&error))?;
|
||||
@@ -212,7 +368,10 @@ impl TokenRefresher for OpenAiTokenClient {
|
||||
if !status.is_success() {
|
||||
return Err(classify_refresh_failure(status, &body));
|
||||
}
|
||||
parse_token_pair(&body).map_err(|()| RefreshFailure::Transport)
|
||||
parse_token_pair(&body).map_err(|()| RefreshFailure::Transport {
|
||||
message: Some("OpenAI OAuth refresh returned an invalid success response".to_owned()),
|
||||
upstream: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -260,7 +419,6 @@ impl AuthorizationCodeExchanger for OpenAiTokenClient {
|
||||
id_token: None,
|
||||
},
|
||||
id_token,
|
||||
expires_in: optional_expiry(tokens.expires_in),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -273,13 +431,21 @@ async fn read_bounded_response(
|
||||
while let Some(chunk) = response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|_| RefreshFailure::Transport)?
|
||||
.map_err(|error| RefreshFailure::Transport {
|
||||
message: Some(error.to_string()),
|
||||
upstream: None,
|
||||
})?
|
||||
{
|
||||
let next_len = body
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.filter(|length| *length <= MAX_OAUTH_RESPONSE_BYTES)
|
||||
.ok_or(RefreshFailure::Transport)?;
|
||||
.ok_or_else(|| RefreshFailure::Transport {
|
||||
message: Some(format!(
|
||||
"OpenAI OAuth response exceeded {MAX_OAUTH_RESPONSE_BYTES} bytes"
|
||||
)),
|
||||
upstream: None,
|
||||
})?;
|
||||
body.reserve(next_len.saturating_sub(body.len()));
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
@@ -301,43 +467,68 @@ fn parse_token_pair(body: &[u8]) -> Result<TokenPair, ()> {
|
||||
refresh_token: tokens.refresh_token,
|
||||
// 官方刷新响应允许省略 ID token;缺失时由调用方保留当前值。
|
||||
id_token: tokens.id_token,
|
||||
expires_in: optional_expiry(tokens.expires_in),
|
||||
})
|
||||
}
|
||||
|
||||
fn optional_expiry(value: Option<u64>) -> Option<Duration> {
|
||||
value
|
||||
.filter(|seconds| *seconds > 0)
|
||||
.map(Duration::from_secs)
|
||||
}
|
||||
|
||||
fn classify_refresh_failure(status: StatusCode, body: &[u8]) -> RefreshFailure {
|
||||
// 5xx 与 429 按状态码判为瞬态:其正文(CDN/网关页等)不是权威 OAuth 错误响应,
|
||||
// 不据此终态失效账号。
|
||||
if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS {
|
||||
return RefreshFailure::Transport;
|
||||
// 官方刷新错误的消息与错误码分别位于 `error.message`、`error.code`;
|
||||
// `error` 字符串与顶层 `code` 仅用于兼容官方客户端自身的错误码提取契约。
|
||||
let error = serde_json::from_slice::<RefreshErrorResponse>(body).ok();
|
||||
let message = error.as_ref().and_then(RefreshErrorResponse::message);
|
||||
let upstream = || {
|
||||
Some(Box::new(RefreshUpstreamFailure::new(
|
||||
status,
|
||||
body,
|
||||
error.as_ref(),
|
||||
)))
|
||||
};
|
||||
let normalized_code = error
|
||||
.as_ref()
|
||||
.and_then(RefreshErrorResponse::code)
|
||||
.map(str::to_ascii_lowercase);
|
||||
// 生产策略故意与官方 Codex 的“任意 401 立即终态”不同:
|
||||
// 显式 401 先进入有界恢复退避,避免瞬时授权故障直接失效账号。
|
||||
if status == StatusCode::UNAUTHORIZED {
|
||||
return RefreshFailure::Transport {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
};
|
||||
}
|
||||
// 4xx 正文才是权威 OAuth 错误响应,据其区分真过期/封禁与其余瞬态。
|
||||
let lower = String::from_utf8_lossy(body).to_ascii_lowercase();
|
||||
if lower.contains("account has been deactivated") {
|
||||
return RefreshFailure::Banned;
|
||||
// 非 401 响应仍与官方一致:三个明确的 RT 原因是永久失败。
|
||||
if matches!(
|
||||
normalized_code.as_deref(),
|
||||
Some("refresh_token_expired" | "refresh_token_reused" | "refresh_token_invalidated")
|
||||
) {
|
||||
return RefreshFailure::InvalidGrant {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
};
|
||||
}
|
||||
if lower.contains("invalid_grant")
|
||||
|| lower.contains("invalid_token")
|
||||
|| lower.contains("access_denied")
|
||||
|| lower.contains("refresh_token_expired")
|
||||
|| lower.contains("refresh_token_reused")
|
||||
{
|
||||
return RefreshFailure::InvalidGrant;
|
||||
if message.as_deref().is_some_and(|message| {
|
||||
message
|
||||
.to_ascii_lowercase()
|
||||
.contains("account has been deactivated")
|
||||
}) {
|
||||
return RefreshFailure::Banned {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
};
|
||||
}
|
||||
RefreshFailure::Transport {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
}
|
||||
RefreshFailure::Transport
|
||||
}
|
||||
|
||||
fn refresh_transport_failure(error: &reqwest::Error) -> RefreshFailure {
|
||||
let message = error.to_string();
|
||||
if is_safe_to_retry_refresh_transport(error) {
|
||||
RefreshFailure::RetryableTransport
|
||||
RefreshFailure::RetryableTransport { message }
|
||||
} else {
|
||||
RefreshFailure::Transport
|
||||
RefreshFailure::Transport {
|
||||
message: Some(message),
|
||||
upstream: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
@@ -62,21 +71,26 @@ use crate::transport::profile::{
|
||||
use crate::transport::protocol::responses::{
|
||||
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
|
||||
};
|
||||
use crate::transport::protocol::websocket::{
|
||||
PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
|
||||
};
|
||||
use crate::transport::request::{
|
||||
CodexRequestEncodeError, 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 +125,8 @@ pub struct CodexProvider {
|
||||
account_feedback: Arc<AccountFeedbackStats>,
|
||||
client: CodexBackendClient,
|
||||
responses_url: Url,
|
||||
image_generations_url: Url,
|
||||
image_edits_url: Url,
|
||||
session_identity: Option<CodexSessionIdentity>,
|
||||
}
|
||||
|
||||
@@ -130,6 +146,11 @@ impl CodexProvider {
|
||||
) -> Result<Self, CodexProviderConfigError> {
|
||||
let responses_url = Url::parse(&endpoint_url(&base_url, CODEX_RESPONSES_PATH))
|
||||
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
|
||||
let image_generations_url =
|
||||
Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_GENERATIONS_PATH))
|
||||
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
|
||||
let image_edits_url = Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_EDITS_PATH))
|
||||
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
|
||||
let client =
|
||||
CodexBackendClient::new(http, base_url, profile).with_websocket_pool(websocket_pool);
|
||||
Ok(Self {
|
||||
@@ -140,6 +161,8 @@ impl CodexProvider {
|
||||
account_feedback,
|
||||
client,
|
||||
responses_url,
|
||||
image_generations_url,
|
||||
image_edits_url,
|
||||
session_identity: None,
|
||||
})
|
||||
}
|
||||
@@ -227,17 +250,25 @@ impl Provider for CodexProvider {
|
||||
UpstreamSendState::NotSent,
|
||||
));
|
||||
}
|
||||
if let Operation::GenerateImage(image) = request.operation() {
|
||||
return self.execute_image(image, candidate, context).await;
|
||||
}
|
||||
let Operation::Generate(generate) = request.operation() else {
|
||||
return Err(provider_error(
|
||||
ProviderErrorKind::Unsupported,
|
||||
UpstreamSendState::NotSent,
|
||||
));
|
||||
};
|
||||
let Some(upstream_model) = candidate.upstream_model() else {
|
||||
return Err(provider_error(
|
||||
ProviderErrorKind::Protocol,
|
||||
UpstreamSendState::NotSent,
|
||||
));
|
||||
};
|
||||
let previous_session = decode_openai_session_state(generate);
|
||||
let continuation_requested = generate.native_continuation_requested();
|
||||
let mut upstream_request =
|
||||
encode_generate_request(generate, candidate.upstream_model().as_str())
|
||||
.map_err(map_request_error)?;
|
||||
let mut upstream_request = encode_generate_request(generate, upstream_model.as_str())
|
||||
.map_err(map_request_error)?;
|
||||
if let Some(conversation_id) = previous_session
|
||||
.as_ref()
|
||||
.and_then(|state| state.conversation_id.as_ref())
|
||||
@@ -247,7 +278,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 +299,7 @@ impl Provider for CodexProvider {
|
||||
.selector
|
||||
.select_with_cyber_policy(
|
||||
&SelectCodexCredential {
|
||||
upstream_model: candidate.upstream_model().as_str(),
|
||||
upstream_model: upstream_model.as_str(),
|
||||
request_url: &self.responses_url,
|
||||
attempt: &context,
|
||||
session_affinity_key: session_affinity_key.as_ref(),
|
||||
@@ -277,57 +313,54 @@ 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 account_scope =
|
||||
if state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch()) {
|
||||
RequestAccountScope::Different
|
||||
} else if has_explicit_state_owner || lease.affinity_hit() {
|
||||
RequestAccountScope::Same
|
||||
} else {
|
||||
RequestAccountScope::Unknown
|
||||
};
|
||||
if context.continuation_attempt() != ContinuationAttempt::None
|
||||
&& let Some(continuation) = context.continuation()
|
||||
{
|
||||
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 +375,19 @@ impl Provider for CodexProvider {
|
||||
}
|
||||
}
|
||||
}
|
||||
if upstream_request.previous_response_id().is_some()
|
||||
&& !account_scope.can_reuse_account_state()
|
||||
{
|
||||
return Err(continuation_replay_required_error());
|
||||
}
|
||||
scope_request_to_account(
|
||||
&mut upstream_request,
|
||||
lease.installation_id(),
|
||||
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 +402,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 +411,7 @@ impl Provider for CodexProvider {
|
||||
client: self.client.clone(),
|
||||
response_origin: self.responses_url.clone(),
|
||||
request: upstream_request,
|
||||
upstream_model: candidate.upstream_model().clone(),
|
||||
upstream_model: upstream_model.clone(),
|
||||
transport_policy: transport,
|
||||
context,
|
||||
selector: Arc::clone(&self.selector),
|
||||
@@ -388,6 +432,76 @@ impl Provider for CodexProvider {
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexProvider {
|
||||
async fn execute_image(
|
||||
&self,
|
||||
image: &ImageRequest,
|
||||
candidate: &ProviderCandidate,
|
||||
context: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
if image.payload().protocol() != PROVIDER_NAME || candidate.upstream_model().is_some() {
|
||||
return Err(provider_error(
|
||||
ProviderErrorKind::InvalidRequest,
|
||||
UpstreamSendState::NotSent,
|
||||
));
|
||||
}
|
||||
let (endpoint_path, response_origin) = match image.kind() {
|
||||
ImageRequestKind::Generation => (
|
||||
CODEX_IMAGE_GENERATIONS_PATH,
|
||||
self.image_generations_url.clone(),
|
||||
),
|
||||
ImageRequestKind::Edit => (CODEX_IMAGE_EDITS_PATH, self.image_edits_url.clone()),
|
||||
};
|
||||
let lease = self
|
||||
.selector
|
||||
.select_for_provider_endpoint(&SelectCodexProviderEndpointCredential {
|
||||
request_url: &response_origin,
|
||||
attempt: &context,
|
||||
})
|
||||
.await
|
||||
.map_err(map_selection_error)?;
|
||||
let lease = Arc::new(lease);
|
||||
let allows_account_state_mutation = lease.allows_account_state_mutation();
|
||||
let provider_kind = ProviderKind::new(PROVIDER_NAME)
|
||||
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
|
||||
let metadata = ProviderCallMetadata::for_provider_endpoint(
|
||||
provider_kind,
|
||||
ProviderResource::Account {
|
||||
id: lease.account_id().clone(),
|
||||
revision: lease.account().revision(),
|
||||
},
|
||||
UpstreamTransport::new(HTTP_JSON_TRANSPORT).map_err(|_| {
|
||||
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
|
||||
})?,
|
||||
);
|
||||
let image_turn_id = image
|
||||
.payload()
|
||||
.context()
|
||||
.get("image_turn_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let events = cold_json_response_stream(ColdJsonResponse {
|
||||
client: self.client.clone(),
|
||||
response_origin,
|
||||
endpoint_path,
|
||||
body: image.payload().body().clone(),
|
||||
image_turn_id,
|
||||
context,
|
||||
selector: Arc::clone(&self.selector),
|
||||
quota: Arc::clone(&self.quota),
|
||||
agent_identity: Arc::clone(&self.agent_identity),
|
||||
lease: Arc::clone(&lease),
|
||||
output_started_at: Instant::now(),
|
||||
});
|
||||
let stream = ProviderStream::new(metadata, events, lease);
|
||||
Ok(if allows_account_state_mutation {
|
||||
stream.with_account_feedback(Arc::clone(&self.account_feedback))
|
||||
} else {
|
||||
stream
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ColdResponse {
|
||||
client: CodexBackendClient,
|
||||
response_origin: Url,
|
||||
@@ -405,12 +519,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 +556,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 +608,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 +684,178 @@ fn map_handshake_attempt_error(error: CodexHandshakeAttemptError) -> MappedProvi
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_json_attempt(
|
||||
request: &ColdJsonResponse,
|
||||
account: &ProviderAccount,
|
||||
installation_id: &str,
|
||||
authorization: &SecretString,
|
||||
cookie_header: Option<&SecretString>,
|
||||
account_selection: CodexAccountSelectionTelemetry<'_>,
|
||||
) -> Result<CodexBackendJsonResponse, CodexHandshakeAttemptError> {
|
||||
let Some(handshake_deadline) = remaining(request.context.deadline()) else {
|
||||
return Err(CodexHandshakeAttemptError::Timeout);
|
||||
};
|
||||
let request_id = request.context.request_id().as_str();
|
||||
let mut request_context = CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
account.upstream_account_id(),
|
||||
request_id,
|
||||
Some(installation_id),
|
||||
);
|
||||
request_context.cookie_header = cookie_header.map(ExposeSecret::expose_secret);
|
||||
request_context.account_selection = account_selection;
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = request.context.cancellation().cancelled() => Err(CodexHandshakeAttemptError::Cancelled),
|
||||
_ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout),
|
||||
response = request.client.post_raw_json(
|
||||
request.endpoint_path,
|
||||
request.body.clone(),
|
||||
request.image_turn_id.as_deref(),
|
||||
request_context,
|
||||
) => response.map_err(CodexHandshakeAttemptError::Client),
|
||||
}
|
||||
}
|
||||
|
||||
fn cold_json_response_stream(request: ColdJsonResponse) -> EventStream {
|
||||
Box::pin(async_stream::try_stream! {
|
||||
let allows_account_state_mutation = request.lease.allows_account_state_mutation();
|
||||
let failure_context = OpenAiFailureContext {
|
||||
client: &request.client,
|
||||
selector: &request.selector,
|
||||
quota: &request.quota,
|
||||
response_origin: &request.response_origin,
|
||||
cyber_policy_scope: None,
|
||||
allows_account_state_mutation,
|
||||
};
|
||||
let mut active_account = request.lease.account().clone();
|
||||
let cookie_header = build_cookie_header(request.lease.cookies())?;
|
||||
let mut authorization = request
|
||||
.lease
|
||||
.authentication()
|
||||
.authorization_header(chrono::Utc::now())
|
||||
.map_err(|_| {
|
||||
provider_error(
|
||||
ProviderErrorKind::Unauthorized,
|
||||
UpstreamSendState::NotSent,
|
||||
)
|
||||
})?;
|
||||
let account_selection = CodexAccountSelectionTelemetry::new(
|
||||
request.lease.affinity_hit(),
|
||||
request.lease.escape_reason(),
|
||||
request.lease.account_switch(),
|
||||
);
|
||||
let mut response = create_json_attempt(
|
||||
&request,
|
||||
&active_account,
|
||||
request.lease.installation_id(),
|
||||
&authorization,
|
||||
cookie_header.as_ref(),
|
||||
account_selection,
|
||||
)
|
||||
.await;
|
||||
if allows_account_state_mutation
|
||||
&& let Err(CodexHandshakeAttemptError::Client(error)) = &response
|
||||
{
|
||||
match request
|
||||
.agent_identity
|
||||
.recover_after_rejected_task(
|
||||
active_account.id(),
|
||||
request.lease.authentication(),
|
||||
error,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(recovered)) => {
|
||||
authorization = recovered
|
||||
.credential
|
||||
.authentication
|
||||
.authorization_header(chrono::Utc::now())
|
||||
.map_err(|_| {
|
||||
provider_error(
|
||||
ProviderErrorKind::Unauthorized,
|
||||
UpstreamSendState::NotSent,
|
||||
)
|
||||
})?;
|
||||
active_account = recovered.account;
|
||||
response = create_json_attempt(
|
||||
&request,
|
||||
&active_account,
|
||||
&recovered.credential.installation_id,
|
||||
&authorization,
|
||||
cookie_header.as_ref(),
|
||||
account_selection,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(_) => response = Err(CodexHandshakeAttemptError::AgentIdentityRecovery),
|
||||
}
|
||||
}
|
||||
let response = match response.map_err(map_handshake_attempt_error) {
|
||||
Ok(response) => response,
|
||||
Err(mut failure) => {
|
||||
if let Some(observation) = failure.observation.take() {
|
||||
yield ProviderEvent::observation(observation);
|
||||
}
|
||||
apply_failure(&failure_context, &active_account, &failure).await;
|
||||
Err(failure.error)?;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut metrics = response.transport_metrics.clone();
|
||||
metrics.first_event_ms = Some(
|
||||
i64::try_from(request.output_started_at.elapsed().as_millis()).unwrap_or(i64::MAX),
|
||||
);
|
||||
if let Some(observation) = codex_response_observation(
|
||||
CodexBackendTransport::HttpJson,
|
||||
&response.diagnostics,
|
||||
&response.response_metadata,
|
||||
&metrics,
|
||||
None,
|
||||
openai_response_timings(&metrics, &response.response_metadata),
|
||||
) {
|
||||
yield ProviderEvent::observation(observation);
|
||||
}
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota_headers(
|
||||
&request.quota,
|
||||
&active_account,
|
||||
&response.rate_limit_headers,
|
||||
)
|
||||
.await;
|
||||
if !response.set_cookie_headers.is_empty()
|
||||
&& let Err(error) = request
|
||||
.selector
|
||||
.capture_response_cookies(
|
||||
&active_account,
|
||||
&request.response_origin,
|
||||
&response.set_cookie_headers,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id = %active_account.id(),
|
||||
error = %error,
|
||||
"Failed to persist OpenAI image response cookies"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let response_meta =
|
||||
ResponseMeta::for_provider_endpoint(request.context.request_id().as_str());
|
||||
yield ProviderEvent::canonical(GatewayEvent::Started(response_meta.clone()));
|
||||
let wire = ProtocolWireEvent::raw_json(PROVIDER_NAME, response.body).map_err(|_| {
|
||||
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
|
||||
})?;
|
||||
yield ProviderEvent::wire(wire);
|
||||
yield ProviderEvent::canonical(GatewayEvent::Completed(
|
||||
response_meta.with_finish_reason(FinishReason::Stop),
|
||||
));
|
||||
})
|
||||
}
|
||||
|
||||
fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let ColdResponse {
|
||||
client,
|
||||
@@ -698,9 +1008,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("a, &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 +1035,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 +1050,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(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
|
||||
return;
|
||||
};
|
||||
@@ -773,15 +1090,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("a, &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(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
apply_failure(&failure_context, &active_account, &failure)
|
||||
.await;
|
||||
Err(failure.error)?;
|
||||
@@ -793,18 +1117,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("a, &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 +1168,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(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let Some((failure, _)) = terminal_failure.as_ref() {
|
||||
apply_failure(&failure_context, &active_account, failure)
|
||||
.await;
|
||||
@@ -852,7 +1185,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
// 完成事件一旦交给下游,Core 可以立刻停止轮询 Provider stream;
|
||||
// 在此之前持久化亲和关系,保证成功请求不会因流被提前 drop 而丢失绑定。
|
||||
selector
|
||||
.record_success(&active_account, session_affinity_key.as_ref())
|
||||
.record_success(
|
||||
&active_account,
|
||||
session_affinity_key.as_ref(),
|
||||
lease.affinity_expected_account_id(),
|
||||
)
|
||||
.await;
|
||||
selector
|
||||
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
|
||||
@@ -861,6 +1198,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 +1258,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("a, &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(
|
||||
"a,
|
||||
&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()
|
||||
@@ -945,13 +1293,21 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
if allows_account_state_mutation && completed && terminal_failure.is_none() {
|
||||
// 同上:尾部 finish() 也可能产出 completed,亲和记录必须先于任何下游 yield。
|
||||
selector
|
||||
.record_success(&active_account, session_affinity_key.as_ref())
|
||||
.record_success(
|
||||
&active_account,
|
||||
session_affinity_key.as_ref(),
|
||||
lease.affinity_expected_account_id(),
|
||||
)
|
||||
.await;
|
||||
selector
|
||||
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
|
||||
.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 +1357,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 +1496,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 +1828,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 +1867,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 +1884,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 +2289,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 { .. })
|
||||
@@ -1963,6 +2379,19 @@ fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure {
|
||||
map_client_error(error, UpstreamSendState::Ambiguous, true)
|
||||
}
|
||||
|
||||
fn continuation_replay_required_error() -> ProviderError {
|
||||
ProviderError::new(
|
||||
ProviderErrorKind::InvalidRequest,
|
||||
UpstreamSendState::NotSent,
|
||||
)
|
||||
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
|
||||
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
|
||||
PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
|
||||
Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
|
||||
Some("invalid_request_error".to_owned()),
|
||||
))
|
||||
}
|
||||
|
||||
fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
|
||||
let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error);
|
||||
let mut failure = map_client_error(error, UpstreamSendState::Sent, false);
|
||||
@@ -1974,7 +2403,9 @@ fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
|
||||
|
||||
fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool {
|
||||
match error {
|
||||
CodexClientError::Http(_) | CodexClientError::StreamIdleTimeout { .. } => true,
|
||||
CodexClientError::Http(_)
|
||||
| CodexClientError::HttpJson(_)
|
||||
| CodexClientError::StreamIdleTimeout { .. } => true,
|
||||
CodexClientError::WebSocket(error) => matches!(
|
||||
error,
|
||||
CodexWebSocketExchangeError::Transport(_)
|
||||
@@ -2055,7 +2486,7 @@ fn map_client_error(
|
||||
ProviderErrorKind::Protocol,
|
||||
UpstreamSendState::Sent,
|
||||
)),
|
||||
CodexClientError::Http(error) => {
|
||||
CodexClientError::Http(error) | CodexClientError::HttpJson(error) => {
|
||||
let send_state = if error.is_connect() {
|
||||
UpstreamSendState::NotSent
|
||||
} else {
|
||||
@@ -2152,7 +2583,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,
|
||||
})
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user