mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
21a436b76a | ||
|
|
849aa73a28 | ||
|
|
fb89bda177 | ||
|
|
9e6c8c9734 | ||
|
|
3f69bc861c | ||
|
|
cea4e1d8e9 | ||
|
|
539764f7eb | ||
|
|
c9be17778d | ||
|
|
d8bc91efce | ||
|
|
e8e2f18927 | ||
|
|
0977de3713 | ||
|
|
eaa37d9faa | ||
|
|
4e588ffa1d | ||
|
|
b9e919f510 | ||
|
|
32f9f3675d | ||
|
|
221d557d5e | ||
|
|
4105c1ece6 | ||
|
|
4d59687cbc | ||
|
|
8621e24e7c | ||
|
|
3bebccac5d | ||
|
|
98287195ca | ||
|
|
fd1b159f71 | ||
|
|
7b03fdc8ab | ||
|
|
851a50962d | ||
|
|
b6e17db0c6 | ||
|
|
b98c2f01a5 | ||
|
|
14e7967e1c | ||
|
|
a4efb77223 | ||
|
|
5041d6ba26 | ||
|
|
322b602154 | ||
|
|
de15722ee1 | ||
|
|
380f04bace | ||
|
|
194cc0e0c9 | ||
|
|
d3acb42d37 | ||
|
|
bfca754714 | ||
|
|
881e78975e | ||
|
|
eb3baef3aa |
@@ -10,12 +10,14 @@ description: 当前仓库开发指南。修改或排查本仓库的 Rust 网关
|
||||
|
||||
## 按任务读取
|
||||
|
||||
| 任务 | 文档入口 |
|
||||
| 任务 | 技能与文档入口 |
|
||||
| --- | --- |
|
||||
| 用户能力、快速开始 | [README](../../../README.md) |
|
||||
| 模块职责、调用链、状态归属 | [系统架构](../../../docs/architecture.md) |
|
||||
| Rust 开发 | 加载并使用 `$rust-best-practices`,结合 [系统架构](../../../docs/architecture.md) 与 [项目约定](../../../CONTRIBUTING.md#项目约定) |
|
||||
| HTTP 路由、请求与响应、导入导出 | [API 文档](../../../docs/api.md) |
|
||||
| 管理端组件、主题 Token、视觉验证 | [管理端主题](../../../docs/theme.md) |
|
||||
| 前端开发、管理端组件、主题 Token、视觉验证 | 加载并使用 `$frontend-design`,结合 [管理端主题](../../../docs/theme.md) 与 [界面验证](../../../CONTRIBUTING.md#界面验证) |
|
||||
| 常规文档更新 | [文档职责与更新条件](../../../CONTRIBUTING.md#文档职责与更新条件):维护当前说明,在原章节就地修订,不记录变更历史 |
|
||||
| 部署、客户端配置、备份与恢复 | [部署说明](../../../deploy/README.md) |
|
||||
| 数据迁移、冻结清单与本地测试库 | [迁移说明](../../../backend/migrations/README.md) |
|
||||
| 验证命令与提交规范 | [贡献与审查](../../../CONTRIBUTING.md) |
|
||||
@@ -27,8 +29,10 @@ description: 当前仓库开发指南。修改或排查本仓库的 Rust 网关
|
||||
|
||||
1. 从仓库根目录确认工作区、当前任务范围和相关运行实例;保留已有改动。
|
||||
2. 按 `AGENTS.md` 的 CodeGraph 约定定位入口、调用链和职责归属,再阅读相关实现。
|
||||
3. 结合上表读取本次需要的文档,先找可复用的同类实现,再修改所属模块。
|
||||
4. 文档和运行结果不一致时,追查当前源码、配置与版本,区分实现缺陷和文档过期;同步受影响的说明。
|
||||
3. 按 [问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据) 核实触发条件、问题发生版本与目标分支现状,
|
||||
优先依据适用的官方源码确定方案。结合上表读取所需文档,先找可复用的同类实现,再修改所属模块。
|
||||
4. 文档和运行结果不一致时,追查当前源码、配置与版本,区分实现缺陷和文档过期;
|
||||
按 [文档职责与更新条件](../../../CONTRIBUTING.md#文档职责与更新条件) 判断是否修订,不为每次改动追加文档或发布说明。
|
||||
5. 按 `CONTRIBUTING.md` 验证并报告结果。后端 Cargo 命令从 `backend/` 执行,或显式指定该目录的 manifest。
|
||||
|
||||
发布任务按 `release` 技能读取当前版本、发布入口和工作流,发布说明语言遵循 `CONTRIBUTING.md`。
|
||||
|
||||
@@ -23,10 +23,13 @@ description: 整理与提交 GitHub Issue(问题反馈 / 功能建议)。用
|
||||
1. 读取目标仓库实际 Issue 模板和配置,区分缺陷、功能建议与协作规则提案。
|
||||
本仓库后两类允许普通 Issue;其他仓库按其配置选择入口。
|
||||
2. 先搜索 open 和 closed Issues,结合关键词、错误信息与触发条件检查重复,而不只比较标题。
|
||||
同时核对相关 PR 与发布记录,关联已有修复及其发布版本;仅主分支已修复时明确尚未发布。
|
||||
同一问题已有记录且没有新增证据时返回已有链接;用户仍明确要求新建时按其要求处理。
|
||||
不因发现重复而擅自评论、重开或关闭已有 Issue。
|
||||
3. 从用户描述、当前版本与已有诊断中整理事实:区分实际结果、预期结果、稳定复现条件和未经证实的推测。
|
||||
不把本地开发版本当作发生问题的部署版本;无法复现时如实说明已知条件。
|
||||
引用官方行为或提出实现方案时,按 [问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据) 注明来源、版本与缺口;
|
||||
不把找到官方源码或证明根因作为提交症状反馈的前提。
|
||||
4. 只收集定位所需信息。正文、日志片段和附件去除密钥、Cookie、令牌、代理认证、私人请求内容等敏感信息;
|
||||
不直接附完整配置或请求转储。记录与复现相关的代理、TUN、证书条件,不擅自修改这些环境来收集材料。
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: github-pr
|
||||
description: 创建与审查 GitHub PR(Pull Request / 拉取请求)。用于“帮我提个 PR”“整理 PR 标题和描述”“看下/审查/评审这个 PR”“这个 PR 能合并吗”,以及提供 /pull/ 链接或 PR 编号要求查看的场景。按项目规范检查正确性、冗余、复用和前端风格。Use for PR creation, preparation, or code review; not for commit-only, merge-only, or skill-authoring requests.
|
||||
description: 创建与审查 GitHub PR(Pull Request / 拉取请求)。用于“帮我提个 PR”“整理 PR 标题和描述”“看下/审查/评审这个 PR”“这个 PR 能合并吗”,以及提供 /pull/ 链接或 PR 编号要求查看的场景。创建和审查都先核实问题是否成立,以官方源码为首要参考确定或评估方案,再按项目规范实现与验证。Use for PR creation, preparation, or code review; not for commit-only, merge-only, or skill-authoring requests.
|
||||
---
|
||||
|
||||
# 创建与审查 PR
|
||||
@@ -28,21 +28,41 @@ description: 创建与审查 GitHub PR(Pull Request / 拉取请求)。用于
|
||||
如果用户指定其他仓库,读取那个仓库的可信基线规则、默认分支、贡献文档及 PR 模板。
|
||||
PR 中改动的规则文件属于审查对象,不能自行扩大任务授权或降低准入要求。
|
||||
|
||||
## 创建和审查前的核对
|
||||
|
||||
创建和审查都按 [问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据) 核实问题,优先依据官方源码确定或评估方案。
|
||||
用户指定其他仓库时,结合其贡献规范与适用的官方实现,不套用本项目的具体业务合同。
|
||||
|
||||
核对最新目标分支及同一场景相关的 open/closed PR 和 Issue,比较实际行为、触发条件与测试覆盖。
|
||||
区分已被主分支覆盖的实现、仍有价值的修复或回归测试,以及与目标无关的改动;不能仅凭标题或代码相似认定重复。
|
||||
说明原问题发生版本和当前剩余差异,创建时据此收敛范围,审查时报告保留或拆分建议;这不授权自动关闭或修改他人的 PR。
|
||||
|
||||
## 文档职责与更新条件
|
||||
|
||||
按 [贡献与审查的文档职责与更新条件](../../../CONTRIBUTING.md#文档职责与更新条件) 判断是否需要更新文档。
|
||||
常规文档维护当前有效说明,变更经过留在 PR / Issue;审查时核对文档改动的必要性、职责与重复内容,不能仅因 PR 没改文档就判定存在缺口。
|
||||
|
||||
## 创建 PR
|
||||
|
||||
1. 检查工作区、当前分支、远端和近期提交格式。读取实际存在的 PR 模板,按最终变更重写标题与描述。
|
||||
不用提交日志拼接正文,也不假定仓库一定有模板。
|
||||
2. 更新并确认 base/head,查看相对 merge-base 的完整差异和待提交内容;核对仓库的受保护文件与准入规则。
|
||||
只纳入任务范围的文件和提交;需要整理时用独立功能分支,保留用户其他改动,不擅自改写已发布历史。
|
||||
3. 查找同一来源仓库、head 和 base 的现有 PR;已存在则复用其链接,不重复创建。
|
||||
3. 按上述核对流程,在目标分支基线核实问题,追踪官方与本项目的相关调用链,先确定解决方向再实现。
|
||||
若改动已完成,也要据此复核问题与方案;需要修正时只修改已授权的范围。仅整理标题或正文时报告差异与缺口,
|
||||
不自行修改实现。原问题不成立时不将它包装成已修复的缺陷;证据不足时明确保留结论。
|
||||
实现应职责清晰、表达直接、复用合理;提交前主动清理本次变更引入的重复逻辑、冗余状态、死代码、
|
||||
多余抽象及无依据的兼容分支。功能可用和检查通过不能替代可维护性自查。
|
||||
4. 查找同一来源仓库、head 和 base 的现有 PR;已存在则复用其链接,不重复创建。
|
||||
更新现有标题或正文须在本次整理/更新请求范围内。
|
||||
4. 按贡献文档执行适用检查。记录实际验证的提交或工作区状态,区分通过、失败、跳过和未执行。
|
||||
5. 按贡献文档执行适用检查。记录实际验证的提交或工作区状态,区分通过、失败、跳过和未执行。
|
||||
前端改动提供实际界面证据或预览步骤;缺少证据时写明,不用构建成功替代视觉验收。
|
||||
5. 将任务范围内的提交推送到正确远端,核对远端 head。标题遵循项目约定,正文先说明问题和最终行为,
|
||||
再给必要的设计取舍、验证与遗留事项;按规模压缩,不强填无关章节。
|
||||
6. 明确要求草稿或仍有阻塞项时创建 Draft PR,并解释原因;具备审查条件且未要求草稿时创建普通 PR。
|
||||
6. 将任务范围内的提交推送到正确远端,核对远端 head。标题遵循项目约定,正文先说明问题成立的依据和最终行为,
|
||||
再给必要的设计取舍、验证与遗留事项。涉及官方行为时,写明官方源码位置与版本或提交、官方处理方式、
|
||||
本次实现与它的关系及必要的偏离原因;区分源码推断和实际复现,不用检查通过替代方案论证。按规模压缩,不强填无关章节。
|
||||
7. 明确要求草稿或仍有阻塞项时创建 Draft PR,并解释原因;具备审查条件且未要求草稿时创建普通 PR。
|
||||
GitHub Draft 状态与“只写本地草稿正文”不同,后者不创建远程对象。
|
||||
7. 创建后核对链接、base、head、草稿状态和最新检查结果。网络超时或结果不确定时先查询远端,再决定是否重试。
|
||||
8. 创建后核对链接、base、head、草稿状态和最新检查结果。网络超时或结果不确定时先查询远端,再决定是否重试。
|
||||
|
||||
CLI 示例中的变量必须先根据上述检查赋值;`pr_body_file` 是精确写入 UTF-8 正文的临时文件:
|
||||
|
||||
@@ -59,16 +79,24 @@ gh pr create --repo "$pr_repo" --base "$pr_base" --head "$pr_head" \
|
||||
|
||||
1. 获取 PR 的 base/head SHA、描述、文件差异、已有讨论及 CI 状态。记录实际审查的 head,使用 merge-base 差异。
|
||||
需要源码上下文时先遵守目标仓库的定位规则;需检出时使用隔离工作区,保留用户当前分支和运行环境。
|
||||
2. 按贡献文档的审查维度追踪完整调用链,找到共享业务事实、同类实现及已有验证。
|
||||
2. 先核实 PR 要解决的问题是否成立:提取触发场景、预期行为和声称的原因,在目标分支基线检查实际行为。
|
||||
按上述核对流程收集一手证据,区分原有问题、环境或配置差异与 PR 新引入的问题,说明主分支已覆盖的部分。
|
||||
3. 从官方方案评估解决方向:说明官方如何处理同一场景,检查 PR 是否遵循其协议语义和职责归属、
|
||||
是否复用已有能力或正确传递官方合同;核对偏离官方方案的理由与运行证据。
|
||||
4. 按贡献文档的审查维度追踪完整调用链,找到共享业务事实、同类实现及已有验证。
|
||||
对既有特殊分支查清服务的场景,再判断本次改动是否引入或加重问题。
|
||||
3. UI 改动同时对照主题、共用组件、数据展示逻辑和同类页面。浏览器可用时检查相关交互及状态;
|
||||
主动检查本次变更新引入或加重的冗余与代码异味,说明具体位置、维护成本或行为影响及最小修复方向;
|
||||
不能仅因功能可用、CI 通过而略过,也不把个人写法偏好判为缺陷。
|
||||
5. UI 改动同时对照主题、共用组件、数据展示逻辑和同类页面。浏览器可用时检查相关交互及状态;
|
||||
未运行的页面检查列为验证缺口,不从截图或源码推断所有状态都正确。
|
||||
4. 核对 CI 与本地证据是否覆盖当前 head;根据风险补充必要验证,不触发无关部署或真实账号业务请求。
|
||||
6. 核对 CI 与本地证据是否覆盖当前 head;根据风险补充必要验证,不触发无关部署或真实账号业务请求。
|
||||
工具不可用、集成测试跳过或环境不足时具体说明影响范围。
|
||||
5. 依照贡献文档用中文输出必须修复项、可选建议和验证缺口;没有的类别可省略。
|
||||
7. 按 [审查标准](../../../CONTRIBUTING.md#审查标准) 用中文先回答原问题是否成立、解决方向是否合理,
|
||||
再列出实现问题与验证缺口,结论附具体依据。
|
||||
实现问题区分必须修复项、可选建议和验证缺口;没有的类别可省略。
|
||||
每个问题给变更中的文件位置、触发条件或规范依据、影响和最小修复方向。
|
||||
复用建议指出具体已有实现,不把个人偏好写成缺陷;未发现问题时直接说明并交代覆盖范围。
|
||||
6. 结束前核对 head 是否变化。发生变化时检查新增差异并更新结论,或明确结论只适用于原 SHA。
|
||||
8. 结束前核对 head 是否变化。发生变化时检查新增差异并更新结论,或明确结论只适用于原 SHA。
|
||||
|
||||
审查本身不修改源码、不发表评论、不提交 GitHub 的批准或请求修改状态。
|
||||
用户明确要求发布时,先复核当前 head 和重复评论,再通过结构化字段或 `--body-file` 发布对应内容;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
interface:
|
||||
display_name: "github-pr"
|
||||
short_description: "按项目规范创建拉取请求,或基于明确提交审查代码与界面"
|
||||
default_prompt: "使用 $github-pr 审查当前 PR,按仓库规范报告问题与验证缺口。"
|
||||
short_description: "创建与审查拉取请求,先核实问题,以官方源码为首要依据确定或评估实现方案"
|
||||
default_prompt: "使用 $github-pr 按当前请求创建或审查 PR,先核实问题是否成立,以官方源码为首要参考确定或评估解决方案,说明源码依据及必要的偏离原因,并按仓库规范实现与验证或报告审查结论。"
|
||||
|
||||
@@ -94,8 +94,17 @@ body:
|
||||
attributes:
|
||||
label: 请求与响应
|
||||
description: |
|
||||
接口问题请提供脱敏后的 cURL 或最小请求体、HTTP 状态码、响应片段,以及请求 ID / 发生时间。
|
||||
如需详细报文,可按部署文档开启 `host.logging.request_dump`,仅摘取定位所需片段并脱敏,不要直接上传完整转储。
|
||||
接口问题请提供脱敏后的最小请求、网关实际状态/响应片段和客户端终端显示,区分两者。
|
||||
优先附管理端默认诊断包并检查其缺失说明;原始错误与日志片段需另行审阅脱敏。
|
||||
先按部署文档的“请求错误排查”检查关联 ID、错误详情与 application 日志,不默认开启完整 dump。
|
||||
placeholder: |
|
||||
发生时间与时区:
|
||||
传输(HTTP JSON / SSE / WebSocket):
|
||||
网关关联 ID(x-gateway-request-id / 配置的入口头):
|
||||
上游关联 ID(x-request-id / x-oai-request-id,注明来源):
|
||||
网关实际状态与响应片段:
|
||||
客户端终端显示:
|
||||
诊断包或无法取得的信息:
|
||||
render: text
|
||||
|
||||
- type: textarea
|
||||
|
||||
+17
-1
@@ -21,11 +21,27 @@ updates:
|
||||
day: monday
|
||||
time: '09:10'
|
||||
timezone: Asia/Shanghai
|
||||
# 固定官方 Codex Desktop 的 TLS 指纹;升级前必须抓取握手并比对扩展顺序。
|
||||
# 按官方 Codex 锁文件、协议指纹或已审计 fork 固定;完成对应核对与实测后再手动升级。
|
||||
ignore:
|
||||
- dependency-name: aws-lc-rs
|
||||
- dependency-name: hyper
|
||||
- dependency-name: hyper-rustls
|
||||
- dependency-name: hyper-tls
|
||||
- dependency-name: hyper-util
|
||||
- dependency-name: native-tls
|
||||
- dependency-name: openssl
|
||||
- dependency-name: openssl-src
|
||||
- dependency-name: openssl-sys
|
||||
- dependency-name: reqwest
|
||||
- dependency-name: rustls
|
||||
- dependency-name: rustls-native-certs
|
||||
- dependency-name: rustls-pki-types
|
||||
- dependency-name: tokio
|
||||
- dependency-name: tokio-native-tls
|
||||
- dependency-name: tokio-rustls
|
||||
- dependency-name: tokio-tungstenite
|
||||
- dependency-name: tungstenite
|
||||
- dependency-name: zstd
|
||||
groups:
|
||||
minor-and-patch:
|
||||
update-types:
|
||||
|
||||
@@ -159,7 +159,7 @@ jobs:
|
||||
ref: ${{ inputs.ref }}
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@0977fd99725f1db4007ccb2928dbb4e90d06cc86 # v6.0.10
|
||||
uses: pnpm/action-setup@ea17c68df8912ef543352723c149a84f56e3d413 # v6.1.0
|
||||
with:
|
||||
version: 11.7.0
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ jobs:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@0977fd99725f1db4007ccb2928dbb4e90d06cc86 # v6.0.10
|
||||
uses: pnpm/action-setup@ea17c68df8912ef543352723c149a84f56e3d413 # v6.1.0
|
||||
with:
|
||||
version: 11.7.0
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ If there is no `.codegraph/` directory, skip CodeGraph entirely — indexing is
|
||||
|
||||
## 工作方式
|
||||
|
||||
- 开发和审查按 [问题与方案依据](CONTRIBUTING.md#问题与方案依据) 先核实问题,以官方源码为首要参考确定或评估方案。
|
||||
- 先定位变更所属模块、同类实现和上下游调用关系,再选择最小合理改动。复用已有能力,也避免为尚不存在的需求增加抽象。
|
||||
- 代码注释使用中文,解释原因与边界;提交信息使用英文,沿用历史中的 Conventional Commits 格式。
|
||||
- 按变更范围执行验证,记录命令、结果和缺口。跳过的测试不算通过;界面与集成行为需要对应运行证据,构建通过不能替代实际验收。
|
||||
@@ -40,6 +41,7 @@ If there is no `.codegraph/` directory, skip CodeGraph entirely — indexing is
|
||||
|
||||
### 审查输出
|
||||
|
||||
- 用中文报告本次变更新引入或加重的、可定位的问题,包含文件位置、触发场景或规范依据、影响和最小修复方向。
|
||||
- 用中文先回答原问题是否成立、解决方向是否合理,依据与判断标准见 [贡献与审查](CONTRIBUTING.md#审查标准),再报告实现问题和验证缺口。
|
||||
- 实现问题聚焦本次变更新引入或加重的、可定位的问题,包含文件位置、触发场景或规范依据、影响和最小修复方向。
|
||||
- 区分必须修复的问题、可选建议和验证缺口。个人偏好标为建议,不作为必须修复项;没有问题时不要凑数。
|
||||
- 严重程度按实际影响判断,不为绕过工具的优先级过滤而升级普通规范问题。说明审查和验证的实际覆盖范围,AI 未报告问题不等于维护者批准。
|
||||
|
||||
+36
-5
@@ -10,10 +10,20 @@
|
||||
|
||||
日志和配置只提供定位所需片段,不上传真实凭据或完整请求转储。功能建议与协作规则调整可直接提交 Issue,说明场景、原因和建议方案。
|
||||
|
||||
## 问题与方案依据
|
||||
|
||||
开发、创建 PR 和审查都先核实待解决的问题、触发条件与预期行为,再确定或评估实现。区分问题发生版本、目标分支现状和本次变更后的行为,不把需求描述、PR 描述或第三方项目的做法直接当成事实。目标分支已修复的问题,应说明当前仍需补齐的差异,不能据此否定旧版本中的问题。
|
||||
|
||||
涉及上游协议、客户端行为或兼容性时,以对应版本的官方源码为首要资料。优先读取用户提供或本机已有的官方源码,核对仓库来源、版本或提交;本地资料不足时查官方仓库的对应版本源码,再结合官方测试、文档、发布说明和运行证据。第三方资料只作线索或补充。引用具体文件、符号或行号,对外引用优先使用固定提交的源码链接;源码不可得或版本不匹配时说明证据缺口,不把最新实现当成当前部署行为。无上游依赖的内部变更,以本项目代码、合同和运行证据为依据。
|
||||
|
||||
官方已有同一场景的处理方式时,优先沿用其协议语义、职责归属和解决路径,检查是否只需复用已有能力或正确传递官方合同。新增覆盖配置、补偿逻辑或采用第三方兼容方案,应说明官方路径为何不足;偏离官方方案需有本项目的实际约束和验证依据,不机械照搬源码,也不默认接受已有提案。
|
||||
|
||||
反馈 Issue 时区分观察到的症状、根因假设和方案建议,关联已有 Issue、PR 及包含修复的发布版本。能说明问题的反馈可以先提交,不要求报告者先找到官方源码或证明根因;引用官方行为或提出实现方案时,按上述原则注明依据与缺口。
|
||||
|
||||
## PR 流程
|
||||
|
||||
1. **准备变更。** 找到问题所属模块、相关调用链和最接近的现有实现,在功能分支完成修改与自查。涉及架构、公开合同或较大交互变化时,先说明方案与取舍。
|
||||
2. **发起 PR。** 目标分支使用 `main`。说明触发场景、改动后的行为和实际验证结果;界面改动提供截图、录屏或可复现的预览方式。一个 PR 围绕一个明确目标,必要重构说明关联原因。
|
||||
1. **准备变更。** 按上述依据核实问题与方案,找到所属模块、相关调用链和最接近的现有实现,在功能分支完成修改与自查。涉及架构、公开合同或较大交互变化时,先说明方案与取舍。
|
||||
2. **发起 PR。** 目标分支使用 `main`。说明问题成立的依据、触发场景、改动后的行为和实际验证结果;涉及官方行为时注明源码位置、版本和方案关系。界面改动提供截图、录屏或可复现的预览方式。一个 PR 围绕一个明确目标,必要重构说明关联原因。
|
||||
3. **自动检查与初审。** CI 执行适用的格式、静态分析、测试和构建检查;启用 AI 审查后,由它按本规范辅助检查语义、复用和一致性。提交者先处理明确问题,维护者重点审阅方案、风险与用户体验。
|
||||
4. **修改与合并。** 更新代码后重新执行受影响的检查,确认审查意见已处理,验证证据对应当前变更。维护者审阅批准后再合并;需要升级说明的变更同步相关文档。
|
||||
|
||||
@@ -42,6 +52,8 @@ PR 不接受根目录或子目录中 `AGENTS.md` 的新增、修改、删除或
|
||||
|
||||
按改动涉及的范围选择检查深度,无需为每个 PR 填写所有维度的报告。
|
||||
|
||||
结论先回答原问题是否成立(成立、部分成立、不成立或证据不足),说明触发条件和依据;再判断解决方向是否合理,交代适用的官方方案、本次实现与它的关系及必要调整,最后报告实现问题与验证缺口。原问题是否成立与 PR 实现是否正确分别判断。源码推断、实际复现和检查结果分别说明,CI 通过不能代替问题与方案的论证。
|
||||
|
||||
| 维度 | 检查重点 |
|
||||
| --- | --- |
|
||||
| 设计与架构 | 职责归属、依赖方向、公开接口是否合理;是否融入已有体系,有无绕过边界的旁路 |
|
||||
@@ -64,21 +76,40 @@ PR 不接受根目录或子目录中 `AGENTS.md` 的新增、修改、删除或
|
||||
|
||||
## 项目约定
|
||||
|
||||
- Rust 模块、可见性、依赖方向和测试目录按架构文档执行,使用现有 Rustfmt 与 Clippy 配置。
|
||||
- AI 进行 Rust 开发时,加载并使用 `$rust-best-practices` 技能。模块、可见性、依赖方向和测试目录按架构文档执行,使用现有 Rustfmt 与 Clippy 配置。
|
||||
- Vue 使用 TypeScript、`<script setup>` 和现有 ESLint 配置;API 请求通过已有模块及 `request` 封装,沿用逐行请求配置格式。
|
||||
- 业务数据的解析、校验和展示优先复用已有逻辑。新增抽象放在拥有该职责的模块,避免重复事实源或无实际需求的通用层。
|
||||
- 新增代码注释使用中文。提交信息使用英文 Conventional Commits,参考近期历史,按完整功能归并;确由 Codex 参与的提交保留 `Co-authored-by: Codex <noreply@openai.com>`。
|
||||
- 发布说明写入 `release/notes.md`,标题首行为目标版本的 `# vX.Y.Z`,章节标题和正文使用中文;命令、配置键、协议名等保留原文。只说明本次版本的实际变化,按影响交代升级事项。
|
||||
- PR 聚焦目标;无关格式化、重构和依赖升级单独处理。行为或合同变更同步到架构文档指定的文档位置,发布说明按维护者的发布安排更新。
|
||||
- PR 聚焦目标;无关格式化、重构和依赖升级单独处理。
|
||||
|
||||
## 文档职责与更新条件
|
||||
|
||||
常规文档说明当前有效的用法、接口合同、模块职责和维护约束,不承担变更历史。
|
||||
只有改动使对应说明失真、缺少必要信息,或改变读者需要执行的操作与判断时,才更新文档。
|
||||
内部修复、重构和界面微调不默认附带文档改动;审查时不能仅因 PR 没改文档就判定存在缺口。
|
||||
|
||||
- 按 [开发指南的文档入口](.agents/skills/dev-guide/SKILL.md#按任务读取) 找到内容所属文档,在原有章节就地修订,
|
||||
不按每次变更追加说明。同一事实只保留一处完整说明,其他文档按读者任务提供必要摘要和链接。
|
||||
- API 文档保留调用所需的字段、状态、语义和限制;架构文档保留职责与不变量;部署文档保留配置、操作和排障。
|
||||
不把界面布局、内部锁或索引实现写进接口说明,也不把前端库的实现选择写进部署说明。
|
||||
- 改动原因、前后对比、排查经过和本次验证结果放在 PR / Issue 或任务报告。版本变更与升级影响按维护者的发版安排维护
|
||||
`release/notes.md`,不为普通开发任务或 PR 自动追加发布说明。
|
||||
- 仍影响实际升级、恢复或兼容使用的限制与操作应保留,并注明适用条件;不能仅因涉及旧版就删除。
|
||||
审查文档改动时核对必要性、所属职责与重复内容,区分当前使用说明和仅解释此次修复的内容。
|
||||
|
||||
## 界面验证
|
||||
|
||||
AI 进行前端开发时,加载并使用 `$frontend-design` 技能。设计方案以本项目的用户任务、[管理端主题](docs/theme.md) 和已有组件为依据,在任务范围内改进布局、视觉与交互。
|
||||
|
||||
开发前查找同类页面及 `components/base` 中的基础组件,优先使用已有主题 Token、交互和数据展示逻辑。需要扩展时在对应职责层修改;局部样式的适用范围遵循主题文档,不一律禁止,也不靠页面覆盖掩盖共用组件的问题。
|
||||
|
||||
列表列宽、序号、按钮、分页和弹窗尺寸由实际任务与数据规模决定,以同类页面为参照。新增或调整交互时说明用户需要完成什么操作,不把某个页面的取舍推广成全局禁令。
|
||||
|
||||
提交者在浏览器中验证受影响路径,提供前后截图、录屏或预览步骤,并标明实际覆盖的状态。按变更选择浅色/深色、窄窗口、空/加载/错误状态、大量数据及键盘操作;共用主题或基础组件变更按主题文档扩大验证范围。模拟数据与预览材料不包含真实凭据。
|
||||
|
||||
涉及浏览器能力或访问环境差异时,按项目支持的部署方式核对相关 HTTP/HTTPS、localhost/实际部署地址与浏览器版本,记录实际覆盖的条件;本地预览正常不能替代部署环境中的交互验收。
|
||||
|
||||
## 验证
|
||||
|
||||
选择与变更相关的检查,纯文档修改无需重建前后端。行为修复补充能覆盖触发条件和结果的测试,不增加仅复述实现的断言。
|
||||
@@ -99,7 +130,7 @@ pnpm --dir frontend build
|
||||
git diff --check
|
||||
```
|
||||
|
||||
`build` 已包含类型检查;按当前架构约定,不维护独立前端测试代码。后端的工具链、命令与线程栈设置统一见 [修改与验收](docs/architecture.md#12-修改与验收)。
|
||||
`build` 已包含类型检查;不维护独立前端测试代码。后端的工具链、命令与线程栈设置统一见 [修改与验收](docs/architecture.md#12-修改与验收)。
|
||||
|
||||
依赖 PostgreSQL / Redis 的测试使用 [专用测试环境](backend/migrations/README.md#本地测试库),未配置导致跳过时明确记录。CI 中已有对应服务和检查,见 [质量工作流](.github/workflows/_quality.yml) 与 [安全扫描](.github/workflows/security-scan.yml)。
|
||||
|
||||
|
||||
@@ -5,30 +5,17 @@
|
||||
|
||||
# Codex Proxy RS
|
||||
|
||||
面向 Codex 的自托管多账号 AI 网关。
|
||||
面向 Codex 的自托管多账号 AI 网关
|
||||
|
||||
[](https://github.com/zyycn/codex-proxy-rs/actions/workflows/ci.yml)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/releases)
|
||||
[](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs)
|
||||
[](LICENSE)
|
||||
|
||||
[功能](#功能) · [快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档) · [社区](#社区) · [许可证](#许可证)
|
||||
[快速开始](#快速开始) · [客户端接入](#客户端接入) · [文档](#文档) · [社区](#社区) · [许可证](#许可证)
|
||||
|
||||
</div>
|
||||
|
||||
将 OpenAI、xAI 账号接入同一个服务,通过网页管理账号、分配客户端密钥和查看使用情况。
|
||||
客户端使用统一的访问地址与代理密钥,上游账号凭据由服务端管理。
|
||||
|
||||
## 功能
|
||||
|
||||
- **账号管理**:OAuth 授权、账号导入、凭据刷新与额度查看;代理管理支持独立保存、连接测试和账号绑定,以及 sub2api OpenAI 账号代理导入。
|
||||
- **访问控制**:通过账号分组分配密钥权限,设置日/周 USD 限额、并发和请求频率。
|
||||
- **客户端接入**:支持 Codex CLI、桌面端与 Responses API 客户端,提供 HTTP 流式响应和 WebSocket。
|
||||
- **请求诊断**:查看请求记录、用量、费用、延迟和错误详情。
|
||||
- **部署运维**:Docker Compose 部署、S3/R2 数据库备份、同大版本在线更新与回滚。
|
||||
|
||||
图片生成、编辑和 Codex 独立搜索由 OpenAI 账号提供;模型与功能的可用性取决于上游账号权限。
|
||||
|
||||
> [!NOTE]
|
||||
> 本项目提供 Responses API,不支持 `/v1/chat/completions`。接入前请确认客户端支持 Responses 协议。
|
||||
|
||||
@@ -41,21 +28,16 @@
|
||||
### 1. 下载部署文件并配置
|
||||
|
||||
```bash
|
||||
mkdir -p codex-proxy-rs/deploy
|
||||
cd codex-proxy-rs
|
||||
mkdir -p codex-proxy-rs/deploy && cd codex-proxy-rs
|
||||
|
||||
curl -fsSL https://raw.githubusercontent.com/zyycn/codex-proxy-rs/main/deploy/compose.yaml \
|
||||
-o deploy/compose.yaml
|
||||
curl -fsSL https://raw.githubusercontent.com/zyycn/codex-proxy-rs/main/deploy/config.example.yaml \
|
||||
-o deploy/config.example.yaml
|
||||
|
||||
mkdir -p .runtime/data .runtime/logs
|
||||
install -d -m 0750 .runtime/postgres .runtime/redis
|
||||
cp deploy/config.example.yaml deploy/config.yaml
|
||||
sudo chown "$(id -u):10001" deploy/config.yaml
|
||||
chmod 0640 deploy/config.yaml
|
||||
sudo chown -R "$(id -u):10001" .runtime/data .runtime/logs
|
||||
chmod 0770 .runtime/data .runtime/logs
|
||||
sudo install -d -m 0770 -o "$(id -u)" -g 10001 .runtime/data .runtime/logs
|
||||
sudo install -m 0640 -o "$(id -u)" -g 10001 deploy/config.example.yaml deploy/config.yaml
|
||||
```
|
||||
|
||||
分别生成数据库和 Redis 密码:
|
||||
@@ -73,8 +55,6 @@ openssl rand -hex 24
|
||||
| `store.redis.password` | 第二个生成的 48 位十六进制密码 |
|
||||
| `admin.default_password` | 管理员初始密码,至少 12 位,不能包含 `$` |
|
||||
|
||||
不需要创建 `.env`。管理员初始密码只在首次创建账号时使用。
|
||||
|
||||
### 2. 启动服务
|
||||
|
||||
```bash
|
||||
@@ -92,7 +72,7 @@ curl -i http://127.0.0.1:8080/healthz
|
||||
|
||||
### 3. 添加账号与客户端密钥
|
||||
|
||||
1. 在「账号」中添加 OpenAI 或 xAI 账号,完成授权或导入。
|
||||
1. 在「账号」中添加账号,完成授权或导入。
|
||||
2. 按需建立账号分组,再创建客户端密钥并选择可用分组。**不选分组表示可使用全部账号**。
|
||||
3. 打开密钥的「使用密钥」,复制客户端配置。
|
||||
|
||||
@@ -106,7 +86,7 @@ curl -i http://127.0.0.1:8080/healthz
|
||||
| 配置 | 值 |
|
||||
| --- | --- |
|
||||
| Base URL | `http://127.0.0.1:8080/v1`;远程接入使用服务器的 HTTPS 地址 |
|
||||
| API Key | 管理端创建的 `sk_...` 客户端密钥 |
|
||||
| API Key | 管理端创建的客户端密钥 |
|
||||
|
||||
可用模型以该密钥查询到的模型列表为准:
|
||||
|
||||
|
||||
Generated
+214
-37
@@ -116,7 +116,7 @@ checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.4",
|
||||
"syn 3.0.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -154,9 +154,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-rs"
|
||||
version = "1.18.1"
|
||||
version = "1.16.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b281d307588d634de920874890732659e2e7672f72b5e10e81badc1a8a83621e"
|
||||
checksum = "a054912289d18629dc78375ba2c3726a3afe3ff71b4edba9dedfca0e3446d1fc"
|
||||
dependencies = [
|
||||
"aws-lc-sys",
|
||||
"untrusted 0.7.1",
|
||||
@@ -165,22 +165,21 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-sys"
|
||||
version = "0.45.0"
|
||||
version = "0.39.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9bff6c3b54fad79a2e60b8102caf565819711497c1f5f092f49508e2f5c31b27"
|
||||
checksum = "1fa7e52a4c5c547c741610a2c6f123f3881e409b714cd27e6798ef020c514f0a"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"cmake",
|
||||
"dunce",
|
||||
"fs_extra",
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-runtime"
|
||||
version = "1.9.1"
|
||||
version = "1.9.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c9007227e10b5fed2f3e0a2beff489211e2b5604c400b7a9d5d81ca9d64c24bb"
|
||||
checksum = "ef47857a1d4488b528f4a5d5715fa7c3300820897824152234d3fa22b1426657"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-sigv4",
|
||||
@@ -206,9 +205,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-sdk-s3"
|
||||
version = "1.144.0"
|
||||
version = "1.146.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "30dc8bf6baaf7d46336a0ca2c69f223d9b90d7a801fb3e28f7ea17b00dc6b1de"
|
||||
checksum = "37b5ffaae346b9bd486ebdc3eb7053ec8ab8a1ad0f19a332eefeff036ed559e6"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"aws-credential-types",
|
||||
@@ -479,9 +478,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-types"
|
||||
version = "1.5.0"
|
||||
version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eec1cd5469f328c782dc3e33d4153cf118a54e33cbb3356d60d16f89883e1f94"
|
||||
checksum = "209f3a6d82a6e9e5f94abbed94c7a26e1c052341002bf57a5fb5481f625896fc"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-smithy-async",
|
||||
@@ -740,7 +739,7 @@ dependencies = [
|
||||
"provider-xai",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"syn 3.0.4",
|
||||
"syn 3.0.5",
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
@@ -749,9 +748,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "combine"
|
||||
version = "4.6.7"
|
||||
version = "4.6.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd"
|
||||
checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
@@ -808,6 +807,16 @@ dependencies = [
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.9.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f"
|
||||
dependencies = [
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.10.1"
|
||||
@@ -1192,7 +1201,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1291,6 +1300,21 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
|
||||
|
||||
[[package]]
|
||||
name = "foreign-types"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
|
||||
dependencies = [
|
||||
"foreign-types-shared",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "foreign-types-shared"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.2.2"
|
||||
@@ -1383,7 +1407,7 @@ checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.4",
|
||||
"syn 3.0.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1879,6 +1903,22 @@ dependencies = [
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper-tls"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"native-tls",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper-util"
|
||||
version = "0.1.20"
|
||||
@@ -1896,9 +1936,11 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"socket2",
|
||||
"system-configuration",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
"windows-registry",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2236,6 +2278,23 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "native-tls"
|
||||
version = "0.2.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "87de3442987e9dbec73158d5c715e7ad9072fda936bb03d19d7fa10e00520f0e"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"log",
|
||||
"openssl",
|
||||
"openssl-probe 0.1.6",
|
||||
"openssl-sys",
|
||||
"schannel",
|
||||
"security-framework 2.11.1",
|
||||
"security-framework-sys",
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "new_debug_unreachable"
|
||||
version = "1.0.6"
|
||||
@@ -2291,12 +2350,66 @@ version = "1.21.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
|
||||
|
||||
[[package]]
|
||||
name = "openssl"
|
||||
version = "0.10.75"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08838db121398ad17ab8531ce9de97b244589089e290a384c900cb9ff7434328"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"cfg-if",
|
||||
"foreign-types",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"openssl-macros",
|
||||
"openssl-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-macros"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-probe"
|
||||
version = "0.1.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e"
|
||||
|
||||
[[package]]
|
||||
name = "openssl-probe"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe"
|
||||
|
||||
[[package]]
|
||||
name = "openssl-src"
|
||||
version = "300.6.1+3.6.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "46eb8fb9fb3b61ce1c0f8a026c4c1a0714d3a9e138e7fbde78753ce2babc3846"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-sys"
|
||||
version = "0.9.111"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "82cab2d520aa75e3c58898289429321eb788c3106963d0dc886ec7a5f4adc321"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"openssl-src",
|
||||
"pkg-config",
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "outref"
|
||||
version = "0.5.2"
|
||||
@@ -2604,6 +2717,7 @@ dependencies = [
|
||||
"base64 0.23.1",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"chrono-tz",
|
||||
"cookie",
|
||||
"flate2",
|
||||
"futures",
|
||||
@@ -2615,6 +2729,9 @@ dependencies = [
|
||||
"hyper",
|
||||
"hyper-rustls",
|
||||
"hyper-util",
|
||||
"native-tls",
|
||||
"openssl",
|
||||
"openssl-sys",
|
||||
"reqwest",
|
||||
"roxmltree",
|
||||
"rustls",
|
||||
@@ -2741,7 +2858,7 @@ dependencies = [
|
||||
"once_cell",
|
||||
"socket2",
|
||||
"tracing",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2840,9 +2957,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.6.0"
|
||||
version = "1.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e37a4ca5c6ca42aa3e6df2fd32b987a65d32a4c2159a6f3fe0fd1df306a2658f"
|
||||
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
@@ -2919,6 +3036,7 @@ dependencies = [
|
||||
"bytes",
|
||||
"cookie",
|
||||
"cookie_store",
|
||||
"encoding_rs",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
@@ -2927,9 +3045,12 @@ dependencies = [
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-rustls",
|
||||
"hyper-tls",
|
||||
"hyper-util",
|
||||
"js-sys",
|
||||
"log",
|
||||
"mime",
|
||||
"native-tls",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"quinn",
|
||||
@@ -2941,6 +3062,7 @@ dependencies = [
|
||||
"serde_urlencoded",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
@@ -3011,7 +3133,7 @@ dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3032,21 +3154,21 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rustls-native-certs"
|
||||
version = "0.8.4"
|
||||
version = "0.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d"
|
||||
checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63"
|
||||
dependencies = [
|
||||
"openssl-probe",
|
||||
"openssl-probe 0.2.1",
|
||||
"rustls-pki-types",
|
||||
"schannel",
|
||||
"security-framework",
|
||||
"security-framework 3.7.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-pki-types"
|
||||
version = "1.15.1"
|
||||
version = "1.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96"
|
||||
checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd"
|
||||
dependencies = [
|
||||
"web-time",
|
||||
"zeroize",
|
||||
@@ -3130,6 +3252,19 @@ dependencies = [
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "security-framework"
|
||||
version = "2.11.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"core-foundation 0.9.4",
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
"security-framework-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "security-framework"
|
||||
version = "3.7.0"
|
||||
@@ -3137,7 +3272,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"core-foundation",
|
||||
"core-foundation 0.10.1",
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
"security-framework-sys",
|
||||
@@ -3205,7 +3340,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.4",
|
||||
"syn 3.0.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3663,9 +3798,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.4"
|
||||
version = "3.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f"
|
||||
checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -3692,6 +3827,27 @@ dependencies = [
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "system-configuration"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"core-foundation 0.9.4",
|
||||
"system-configuration-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "system-configuration-sys"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4"
|
||||
dependencies = [
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tar"
|
||||
version = "0.4.46"
|
||||
@@ -3713,7 +3869,7 @@ dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3742,7 +3898,7 @@ checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.4",
|
||||
"syn 3.0.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3811,9 +3967,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.53.1"
|
||||
version = "1.52.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
|
||||
checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"libc",
|
||||
@@ -3836,6 +3992,16 @@ dependencies = [
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-native-tls"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2"
|
||||
dependencies = [
|
||||
"native-tls",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-rustls"
|
||||
version = "0.26.4"
|
||||
@@ -4184,9 +4350,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.26.0"
|
||||
version = "1.26.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812"
|
||||
checksum = "2ef6dac1e96601b4fb3acccccff2139741fcb757cb9a36089bf5be91cfb285ce"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"js-sys",
|
||||
@@ -4419,6 +4585,17 @@ version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-registry"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
"windows-result",
|
||||
"windows-strings",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.4.1"
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
mod architecture;
|
||||
mod bootstrap;
|
||||
|
||||
// 组合根是否混入业务策略由代码审查判断,行数和标识符/注释关键词不能证明职责越界。
|
||||
// 依赖 DAG、公开模块面和模块镜像继续由 architecture 检查,源码纪律由下方测试维护。
|
||||
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
fs,
|
||||
@@ -75,28 +78,6 @@ fn workspace_production_files_have_no_hidden_modules_or_test_hooks() {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bootstrap_owns_only_bundle_wiring() {
|
||||
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src");
|
||||
let bootstrap = fs::read_to_string(root.join("bootstrap.rs")).expect("read bootstrap");
|
||||
for forbidden in [
|
||||
"sqlx::",
|
||||
"redis::",
|
||||
"Repository",
|
||||
"impl Provider",
|
||||
"impl ExecutionService",
|
||||
"tokio::spawn",
|
||||
"access_token",
|
||||
"refresh_token",
|
||||
] {
|
||||
assert!(
|
||||
!bootstrap.contains(forbidden),
|
||||
"bootstrap owns `{forbidden}`"
|
||||
);
|
||||
}
|
||||
assert!(bootstrap.lines().count() <= 300);
|
||||
}
|
||||
|
||||
fn rust_files(root: &Path) -> BTreeSet<PathBuf> {
|
||||
let mut result = BTreeSet::new();
|
||||
collect_rust_files(root, root, &mut result);
|
||||
|
||||
@@ -6,7 +6,7 @@ use chrono::{DateTime, Utc};
|
||||
|
||||
use gateway_core::{
|
||||
engine::budget::{ClientBudgetLimits, ClientBudgetStatus},
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
};
|
||||
|
||||
@@ -144,6 +144,7 @@ impl fmt::Debug for ClientKeySecret {
|
||||
/// API 提交的 Client Key 创建命令。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateClientKey {
|
||||
pub custom_key: Option<PlaintextClientApiKey>,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
|
||||
@@ -11,6 +11,8 @@ pub mod client_keys;
|
||||
pub mod observability;
|
||||
pub mod provider_credentials;
|
||||
pub mod proxies;
|
||||
pub mod quota_forecast;
|
||||
pub mod quota_forecast_sampling;
|
||||
pub mod settings;
|
||||
pub mod system;
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ use serde_json::{Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{
|
||||
AdminError, MutationActor, MutationContext, PageSize, Revision,
|
||||
AdminError, MutationActor, MutationContext, Revision,
|
||||
accounts::{
|
||||
AccountImportSettings, AccountRecord, AccountSummary, AccountUsage, CredentialState,
|
||||
},
|
||||
@@ -49,56 +49,6 @@ impl fmt::Debug for ProviderDocument {
|
||||
}
|
||||
}
|
||||
|
||||
/// Credential 列表稳定游标。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialCursor {
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub account_id: ProviderAccountId,
|
||||
}
|
||||
|
||||
/// Provider credential 列表查询。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialListQuery {
|
||||
pub credential_state: Option<CredentialStateFilter>,
|
||||
pub enabled: Option<bool>,
|
||||
pub window: CredentialListWindow,
|
||||
}
|
||||
|
||||
/// Credential 状态筛选。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum CredentialStateFilter {
|
||||
Exact(CredentialState),
|
||||
AnyOf(Vec<CredentialState>),
|
||||
}
|
||||
|
||||
impl CredentialStateFilter {
|
||||
#[must_use]
|
||||
pub fn matches(&self, credential_state: CredentialState) -> bool {
|
||||
match self {
|
||||
Self::Exact(expected) => *expected == credential_state,
|
||||
Self::AnyOf(expected) => expected.contains(&credential_state),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Credential 目录的集合窗口;完整列表与游标分页互斥。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum CredentialListWindow {
|
||||
All,
|
||||
Page {
|
||||
cursor: Option<CredentialCursor>,
|
||||
page_size: PageSize,
|
||||
},
|
||||
}
|
||||
|
||||
/// Provider credential 列表。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialPage {
|
||||
pub config_revision: Revision,
|
||||
pub items: Vec<AccountRecord>,
|
||||
pub next_cursor: Option<CredentialCursor>,
|
||||
}
|
||||
|
||||
/// Provider credential 详情。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialDetails {
|
||||
@@ -807,6 +757,24 @@ pub(crate) fn explicit_plan_type(value: Option<&str>) -> Option<&str> {
|
||||
})
|
||||
}
|
||||
|
||||
/// 按需查询的订阅周期,不参与额度或账号可用性判断。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderSubscription {
|
||||
pub starts_at: Option<DateTime<Utc>>,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub will_renew: Option<bool>,
|
||||
pub billing_period: Option<String>,
|
||||
pub billing_currency: Option<String>,
|
||||
pub observed_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 按需汇聚的个人信息;资料查询失败不丢弃可用的订阅结果。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct AccountPersonalInfo {
|
||||
pub profile: Result<ProviderProfileStatistics, AdminError>,
|
||||
pub subscription: Option<ProviderSubscription>,
|
||||
}
|
||||
|
||||
/// Provider 官方个人资料中的累计摘要。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderProfileStatisticsSummary {
|
||||
@@ -934,28 +902,31 @@ impl ProviderQuota {
|
||||
/// Provider 投影顺序。没有符合条件的窗口时不回退到短期或历史累计。
|
||||
#[must_use]
|
||||
pub fn usage_window(&self) -> Option<(&ProviderQuotaWindow, AccountUsagePeriod)> {
|
||||
self.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, window)| {
|
||||
if window.local_usage_attribution != QuotaLocalUsageAttribution::AccountWide
|
||||
|| window.reset_at.is_none()
|
||||
|| window.local_usage.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let seconds = window.window_seconds?;
|
||||
let period = if is_week_window(seconds) {
|
||||
AccountUsagePeriod::Weekly
|
||||
} else if window.group == "monthly" && seconds >= 7 * 24 * 60 * 60 {
|
||||
AccountUsagePeriod::Monthly
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
Some((index, window, period))
|
||||
})
|
||||
.min_by_key(|(index, _, period)| (*period, *index))
|
||||
.map(|(_, window, period)| (window, period))
|
||||
self.usage_windows()
|
||||
.filter(|(window, _)| window.local_usage.is_some())
|
||||
.min_by_key(|(_, period)| *period)
|
||||
}
|
||||
|
||||
/// 统计与预测共用窗口归属规则,避免把模型专属桶或短期限流当作账号容量。
|
||||
pub(crate) fn usage_windows(
|
||||
&self,
|
||||
) -> impl Iterator<Item = (&ProviderQuotaWindow, AccountUsagePeriod)> {
|
||||
self.windows.iter().filter_map(|window| {
|
||||
if window.local_usage_attribution != QuotaLocalUsageAttribution::AccountWide
|
||||
|| window.reset_at.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let seconds = window.window_seconds?;
|
||||
let period = if is_week_window(seconds) {
|
||||
AccountUsagePeriod::Weekly
|
||||
} else if window.group == "monthly" && seconds >= 7 * 24 * 60 * 60 {
|
||||
AccountUsagePeriod::Monthly
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
Some((window, period))
|
||||
})
|
||||
}
|
||||
|
||||
/// 返回 Dashboard 使用的代表性额度比例。
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
//! 账号额度的只读容量估算,不参与金额结算或账号调度。
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
|
||||
use super::provider_credentials::{AccountUsagePeriod, ProviderQuota, ProviderQuotaWindow};
|
||||
use super::quota_forecast_sampling::{QuotaForecastMethod, QuotaForecastSample};
|
||||
|
||||
const DAY_SECONDS: u64 = 86_400;
|
||||
const MIN_USED_PERCENT: f64 = 5.0;
|
||||
const LOW_SAMPLE_PERCENT: f64 = 10.0;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct AccountQuotaForecastReport {
|
||||
pub account_id: String,
|
||||
pub generated_at: DateTime<Utc>,
|
||||
pub forecasts: [AccountQuotaForecast; 2],
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct AccountQuotaForecast {
|
||||
pub period: AccountUsagePeriod,
|
||||
pub target_seconds: u64,
|
||||
pub extrapolated: bool,
|
||||
pub source: Option<QuotaForecastSource>,
|
||||
pub unavailable_reason: Option<&'static str>,
|
||||
pub low_sample: bool,
|
||||
pub incomplete_cost: bool,
|
||||
pub incomplete_tokens: bool,
|
||||
pub estimated_tokens: Option<u64>,
|
||||
pub estimated_usd: Option<f64>,
|
||||
/// 剩余估算始终属于源窗口,不随目标周期折算。
|
||||
pub remaining_tokens: Option<u64>,
|
||||
pub remaining_usd: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct QuotaForecastSource {
|
||||
pub label: String,
|
||||
pub used_percent: Option<f64>,
|
||||
pub observed_at: Option<DateTime<Utc>>,
|
||||
pub reset_at: DateTime<Utc>,
|
||||
pub tokens: Option<u64>,
|
||||
pub usd: Option<f64>,
|
||||
}
|
||||
|
||||
/// 优先预测真实的对应窗口;缺少对应周期时只给出明确标识的 7/30 天容量折算。
|
||||
/// 本地日志不能证明站外消耗或完整留存,因此即使样本充足也不声称官方额度。
|
||||
#[must_use]
|
||||
pub fn account_quota_forecasts(
|
||||
quota: &ProviderQuota,
|
||||
account_added_at: DateTime<Utc>,
|
||||
now: DateTime<Utc>,
|
||||
samples: &[QuotaForecastSample],
|
||||
) -> [AccountQuotaForecast; 2] {
|
||||
[AccountUsagePeriod::Weekly, AccountUsagePeriod::Monthly].map(|period| {
|
||||
let selected = quota
|
||||
.usage_windows()
|
||||
.find(|(_, source_period)| *source_period == period)
|
||||
.or_else(|| quota.usage_windows().min_by_key(|(_, period)| *period));
|
||||
let mut forecast = AccountQuotaForecast {
|
||||
period,
|
||||
target_seconds: match period {
|
||||
AccountUsagePeriod::Weekly => 7 * DAY_SECONDS,
|
||||
AccountUsagePeriod::Monthly => 30 * DAY_SECONDS,
|
||||
},
|
||||
extrapolated: false,
|
||||
source: None,
|
||||
unavailable_reason: Some("没有可统计的周/月额度窗口,请先刷新账号额度。"),
|
||||
low_sample: false,
|
||||
incomplete_cost: false,
|
||||
incomplete_tokens: false,
|
||||
estimated_tokens: None,
|
||||
estimated_usd: None,
|
||||
remaining_tokens: None,
|
||||
remaining_usd: None,
|
||||
};
|
||||
if let Some((window, source_period)) = selected {
|
||||
forecast.project(
|
||||
window,
|
||||
source_period,
|
||||
quota.observed_at,
|
||||
account_added_at,
|
||||
now,
|
||||
samples.iter().find(|sample| sample.key == window.key),
|
||||
);
|
||||
}
|
||||
forecast
|
||||
})
|
||||
}
|
||||
|
||||
impl AccountQuotaForecast {
|
||||
fn project(
|
||||
&mut self,
|
||||
window: &ProviderQuotaWindow,
|
||||
source_period: AccountUsagePeriod,
|
||||
observed_at: Option<DateTime<Utc>>,
|
||||
account_added_at: DateTime<Utc>,
|
||||
now: DateTime<Utc>,
|
||||
sample: Option<&QuotaForecastSample>,
|
||||
) {
|
||||
let (Some(seconds), Some(reset_at)) = (window.window_seconds, window.reset_at) else {
|
||||
return;
|
||||
};
|
||||
self.extrapolated = source_period != self.period;
|
||||
if !self.extrapolated {
|
||||
self.target_seconds = seconds;
|
||||
}
|
||||
let percent = window
|
||||
.used_percent
|
||||
.filter(|p| p.is_finite() && (0.0..=100.0).contains(p));
|
||||
let usage = sample.map(|sample| &sample.usage);
|
||||
let usd = usage
|
||||
.filter(|usage| usage.known_cost_count > 0)
|
||||
.map(|usage| usage.usd)
|
||||
.filter(|value| value.is_finite() && *value >= 0.0);
|
||||
let start = i64::try_from(seconds)
|
||||
.ok()
|
||||
.and_then(Duration::try_seconds)
|
||||
.and_then(|duration| reset_at.checked_sub_signed(duration));
|
||||
self.source = Some(QuotaForecastSource {
|
||||
label: window.label.clone(),
|
||||
used_percent: percent,
|
||||
observed_at,
|
||||
reset_at,
|
||||
tokens: usage.map(|usage| usage.tokens),
|
||||
usd,
|
||||
});
|
||||
self.incomplete_cost = usage.is_none_or(|usage| {
|
||||
usage.unavailable_cost_count > 0
|
||||
|| usage.known_cost_count == 0
|
||||
|| usage.known_cost_count != usage.request_count
|
||||
});
|
||||
self.incomplete_tokens = usage.is_some_and(|usage| usage.missing_token_count > 0);
|
||||
let method = sample.map_or(QuotaForecastMethod::Cumulative, |sample| sample.method);
|
||||
let Some(start) = start.filter(|start| *start <= now && now < reset_at) else {
|
||||
self.unavailable_reason = Some("额度窗口已过期或边界无效,请刷新账号额度后重试。");
|
||||
return;
|
||||
};
|
||||
if !observed_at.is_some_and(|observed| start <= observed && observed <= now) {
|
||||
self.unavailable_reason = Some("缺少本周期的额度快照,请先刷新账号额度。");
|
||||
return;
|
||||
}
|
||||
if sample.is_some_and(|sample| sample.discontinuous) {
|
||||
self.unavailable_reason =
|
||||
Some("额度观测出现回落或累计记录不连续,正在重新积累配对样本。");
|
||||
return;
|
||||
}
|
||||
if account_added_at > start && method == QuotaForecastMethod::Cumulative {
|
||||
self.unavailable_reason = Some(
|
||||
"本周期开始时的记录不完整,正在积累至少 5 个百分点的配对观测,无需等待下次重置。",
|
||||
);
|
||||
return;
|
||||
}
|
||||
let Some(percent) = percent else {
|
||||
self.unavailable_reason = Some("已用比例未知,请刷新额度后查看预测。");
|
||||
return;
|
||||
};
|
||||
let Some(usage) = usage.filter(|usage| usage.request_count > 0) else {
|
||||
self.unavailable_reason = Some("本周期没有网关用量记录,暂时无法预测额度。");
|
||||
return;
|
||||
};
|
||||
let Some(sample) = sample.filter(|sample| {
|
||||
sample.end_at == observed_at.unwrap_or(now)
|
||||
&& start <= sample.start_at
|
||||
&& sample.start_at <= sample.end_at
|
||||
&& sample.sampled_percent.is_finite()
|
||||
&& sample.sampled_percent >= MIN_USED_PERCENT
|
||||
&& sample.sampled_percent <= percent
|
||||
}) else {
|
||||
self.unavailable_reason =
|
||||
Some("有效额度进度不足 5 个百分点或采样边界无效,请继续积累用量。");
|
||||
return;
|
||||
};
|
||||
self.low_sample = sample.sampled_percent < LOW_SAMPLE_PERCENT
|
||||
|| (method == QuotaForecastMethod::Incremental && sample.block_count < 2);
|
||||
// 漏记和个别缺失只影响精度,仍按已记录数值估算,不按请求数补齐未知消耗。
|
||||
// 预测是近似展示值;不复用为账单金额,也不把月折算当成自然月或额外余额。
|
||||
let capacity_factor = 100.0 / sample.sampled_percent;
|
||||
let factor = capacity_factor * self.target_seconds as f64 / seconds as f64;
|
||||
let tokens = Some(usage.tokens).filter(|tokens| *tokens > 0);
|
||||
self.estimated_tokens = tokens.and_then(|value| estimate_tokens(value, factor));
|
||||
self.estimated_usd = usd.and_then(|value| estimate(value, factor));
|
||||
let remaining_factor = (100.0 - percent) / sample.sampled_percent;
|
||||
self.remaining_tokens = tokens.and_then(|value| estimate_tokens(value, remaining_factor));
|
||||
self.remaining_usd = usd.and_then(|value| estimate(value, remaining_factor));
|
||||
self.unavailable_reason = if self.estimated_tokens.is_none() && self.estimated_usd.is_none()
|
||||
{
|
||||
Some("本周期暂无可用于估算的 Token 或费用数据,请积累用量后重试。")
|
||||
} else {
|
||||
None
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
fn estimate(value: f64, factor: f64) -> Option<f64> {
|
||||
let estimate = value * factor;
|
||||
(estimate.is_finite() && estimate >= 0.0).then_some(estimate)
|
||||
}
|
||||
|
||||
fn estimate_tokens(value: u64, factor: f64) -> Option<u64> {
|
||||
estimate(value as f64, factor)
|
||||
.filter(|value| value.round() < u64::MAX as f64)
|
||||
.map(|value| value.round() as u64)
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
//! 预测专用的配对采样;不改变完整交付用量或费用账本的统计口径。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use super::provider_credentials::ProviderDocument;
|
||||
|
||||
/// 限制带 Provider 文档的历史点数量;累计用量仍覆盖整个查询区间。
|
||||
pub const MAX_FORECAST_HISTORY_POINTS: usize = 128;
|
||||
const MIN_BLOCK_PERCENT: f64 = 5.0;
|
||||
const RECENT_BLOCKS: usize = 3;
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct QuotaForecastUsage {
|
||||
pub request_count: u64,
|
||||
pub tokens: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub missing_token_count: u64,
|
||||
pub known_cost_count: u64,
|
||||
pub unavailable_cost_count: u64,
|
||||
pub usd: f64,
|
||||
pub excluded_request_count: u64,
|
||||
}
|
||||
|
||||
impl QuotaForecastUsage {
|
||||
/// 累计计数倒退表示不能配对,不用饱和减法掩盖数据断点。
|
||||
fn difference(&self, baseline: &Self) -> Option<Self> {
|
||||
let usd = self.usd - baseline.usd;
|
||||
if !usd.is_finite() || usd < -1e-9 {
|
||||
return None;
|
||||
}
|
||||
Some(Self {
|
||||
request_count: self.request_count.checked_sub(baseline.request_count)?,
|
||||
tokens: self.tokens.checked_sub(baseline.tokens)?,
|
||||
input_tokens: self.input_tokens.checked_sub(baseline.input_tokens)?,
|
||||
output_tokens: self.output_tokens.checked_sub(baseline.output_tokens)?,
|
||||
cached_tokens: self.cached_tokens.checked_sub(baseline.cached_tokens)?,
|
||||
missing_token_count: self
|
||||
.missing_token_count
|
||||
.checked_sub(baseline.missing_token_count)?,
|
||||
known_cost_count: self
|
||||
.known_cost_count
|
||||
.checked_sub(baseline.known_cost_count)?,
|
||||
unavailable_cost_count: self
|
||||
.unavailable_cost_count
|
||||
.checked_sub(baseline.unavailable_cost_count)?,
|
||||
usd: usd.max(0.0),
|
||||
excluded_request_count: self
|
||||
.excluded_request_count
|
||||
.checked_sub(baseline.excluded_request_count)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct QuotaForecastHistoryPoint {
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub completed_at: DateTime<Utc>,
|
||||
pub usage: QuotaForecastUsage,
|
||||
pub provider_observation: ProviderDocument,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct QuotaForecastHistory {
|
||||
pub points: Vec<QuotaForecastHistoryPoint>,
|
||||
pub usage: QuotaForecastUsage,
|
||||
pub pending_request_count: u64,
|
||||
}
|
||||
|
||||
/// 具体 Provider 从自己的历史文档解释出的额度事实。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct QuotaForecastObservation {
|
||||
pub used_percent: f64,
|
||||
pub reset_at: DateTime<Utc>,
|
||||
pub plan_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct QuotaForecastPoint {
|
||||
pub observed_at: DateTime<Utc>,
|
||||
pub used_percent: f64,
|
||||
pub usage: QuotaForecastUsage,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum QuotaForecastMethod {
|
||||
Cumulative,
|
||||
Incremental,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct QuotaForecastSample {
|
||||
pub key: String,
|
||||
pub method: QuotaForecastMethod,
|
||||
pub start_at: DateTime<Utc>,
|
||||
pub end_at: DateTime<Utc>,
|
||||
pub baseline_percent: f64,
|
||||
pub sampled_percent: f64,
|
||||
pub block_count: usize,
|
||||
pub observation_count: usize,
|
||||
pub pending_request_count: u64,
|
||||
pub discontinuous: bool,
|
||||
pub usage: QuotaForecastUsage,
|
||||
}
|
||||
|
||||
/// 对同一额度段的累计点形成至少 5 个百分点的连续块,使用最近三个完整块及尾部。
|
||||
/// 不平均逐请求比值;重复百分比不增加块数,5/10 个百分点仅是量化噪声门槛。
|
||||
#[must_use]
|
||||
pub fn select_forecast_sample(
|
||||
key: String,
|
||||
start_at: DateTime<Utc>,
|
||||
current: QuotaForecastPoint,
|
||||
mut history: Vec<QuotaForecastPoint>,
|
||||
pending_request_count: u64,
|
||||
) -> QuotaForecastSample {
|
||||
history.retain(|point| {
|
||||
start_at <= point.observed_at
|
||||
&& point.observed_at < current.observed_at
|
||||
&& point.used_percent.is_finite()
|
||||
&& (0.0..100.0).contains(&point.used_percent)
|
||||
});
|
||||
history.sort_by_key(|point| point.observed_at);
|
||||
let mut sample = QuotaForecastSample {
|
||||
key,
|
||||
method: QuotaForecastMethod::Cumulative,
|
||||
start_at,
|
||||
end_at: current.observed_at,
|
||||
baseline_percent: 0.0,
|
||||
sampled_percent: current.used_percent,
|
||||
block_count: 0,
|
||||
observation_count: history.len(),
|
||||
pending_request_count,
|
||||
discontinuous: false,
|
||||
usage: current.usage.clone(),
|
||||
};
|
||||
history.push(current);
|
||||
let mut anchors: Vec<&QuotaForecastPoint> = Vec::new();
|
||||
let mut high_water: Option<&QuotaForecastPoint> = None;
|
||||
for point in &history {
|
||||
if let Some(previous) = high_water {
|
||||
if point.observed_at <= previous.observed_at {
|
||||
continue;
|
||||
}
|
||||
if point.used_percent < previous.used_percent {
|
||||
// 一个百分点内的回落可能是量化/并发乱序;更大的回落不能跨越。
|
||||
if previous.used_percent - point.used_percent <= 1.0 {
|
||||
continue;
|
||||
}
|
||||
sample.discontinuous = true;
|
||||
anchors.clear();
|
||||
}
|
||||
}
|
||||
high_water = Some(point);
|
||||
if anchors
|
||||
.last()
|
||||
.is_none_or(|anchor| point.used_percent - anchor.used_percent >= MIN_BLOCK_PERCENT)
|
||||
{
|
||||
anchors.push(point);
|
||||
}
|
||||
}
|
||||
let Some(current) = history.last() else {
|
||||
return sample;
|
||||
};
|
||||
if anchors.len() < 2 {
|
||||
return sample;
|
||||
}
|
||||
let baseline_index = anchors.len().saturating_sub(RECENT_BLOCKS + 1);
|
||||
let baseline = anchors[baseline_index];
|
||||
let progress = current.used_percent - baseline.used_percent;
|
||||
if !progress.is_finite() || progress < MIN_BLOCK_PERCENT {
|
||||
return sample;
|
||||
}
|
||||
let Some(usage) = current.usage.difference(&baseline.usage) else {
|
||||
sample.discontinuous = true;
|
||||
return sample;
|
||||
};
|
||||
if usage.request_count == 0 {
|
||||
return sample;
|
||||
}
|
||||
sample.method = QuotaForecastMethod::Incremental;
|
||||
sample.start_at = baseline.observed_at;
|
||||
sample.baseline_percent = baseline.used_percent;
|
||||
sample.sampled_percent = progress;
|
||||
sample.block_count = anchors.len() - baseline_index - 1;
|
||||
sample.usage = usage;
|
||||
// 断点后的新基线已有足够进度时,只消费新段;旧段从不参与本次估算。
|
||||
sample.discontinuous = false;
|
||||
sample
|
||||
}
|
||||
@@ -19,7 +19,11 @@ use crate::model::provider_credentials::{
|
||||
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
|
||||
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, ProviderModels,
|
||||
ProviderProfileAvatar, ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderResetCreditResult, ProviderResetCredits, explicit_plan_type,
|
||||
ProviderResetCreditResult, ProviderResetCredits, ProviderSubscription, explicit_plan_type,
|
||||
};
|
||||
use crate::model::{
|
||||
provider_credentials::{ProviderDocument, ProviderQuotaWindow},
|
||||
quota_forecast_sampling::QuotaForecastObservation,
|
||||
};
|
||||
|
||||
/// Provider 管理失败的稳定分类。
|
||||
@@ -171,7 +175,24 @@ pub trait ProviderAdmin: Send + Sync {
|
||||
request: ProviderQuotaRequest,
|
||||
) -> Result<ProviderQuota, ProviderAdminError>;
|
||||
|
||||
/// 历史观测的协议字段仅由具体 Provider 解释;不支持时保留累计估算。
|
||||
fn quota_forecast_observation(
|
||||
&self,
|
||||
_document: &ProviderDocument,
|
||||
_window: &ProviderQuotaWindow,
|
||||
) -> Option<QuotaForecastObservation> {
|
||||
None
|
||||
}
|
||||
|
||||
/// 查询 Provider 官方个人资料统计;不支持该能力的 Provider 使用默认拒绝。
|
||||
/// 只由个人信息显式查询,不在额度刷新或后台任务中预取。
|
||||
async fn subscription(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
) -> Result<Option<ProviderSubscription>, ProviderAdminError> {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
async fn profile_statistics(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
|
||||
@@ -32,9 +32,9 @@ use crate::model::{
|
||||
},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, CredentialDetails, CredentialImportCommit, CredentialImportResult,
|
||||
CredentialListQuery, CredentialMutationResult, CredentialPage, CredentialRotationCommit,
|
||||
ProviderExportCredentialInput,
|
||||
CredentialMutationResult, CredentialRotationCommit, ProviderExportCredentialInput,
|
||||
},
|
||||
quota_forecast_sampling::QuotaForecastHistory,
|
||||
settings::{AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RuntimeSettings},
|
||||
};
|
||||
|
||||
@@ -44,6 +44,7 @@ pub enum AdminStoreErrorKind {
|
||||
Invalid,
|
||||
NotFound,
|
||||
StaleRevision,
|
||||
DuplicateName,
|
||||
Conflict,
|
||||
Unavailable,
|
||||
}
|
||||
@@ -110,11 +111,11 @@ pub trait AccountStore: Send + Sync {
|
||||
windows: &[AccountUsageWindowQuery],
|
||||
) -> AdminStoreResult<Vec<AccountUsageWindowResult>>;
|
||||
|
||||
async fn list_credentials(
|
||||
/// 从同一数据库语句取得截止快照的累计用量和有界历史观测。
|
||||
async fn load_quota_forecast_history(
|
||||
&self,
|
||||
provider_kind: &gateway_core::routing::ProviderKind,
|
||||
query: CredentialListQuery,
|
||||
) -> AdminStoreResult<CredentialPage>;
|
||||
window: &AccountUsageWindowQuery,
|
||||
) -> AdminStoreResult<QuotaForecastHistory>;
|
||||
|
||||
async fn credential_details(
|
||||
&self,
|
||||
|
||||
@@ -22,12 +22,14 @@ use crate::{
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
ConsumeProviderResetCredit, PrepareCredentialRefresh, ProviderModels,
|
||||
ProviderProfileAvatar, ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest,
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountPersonalInfo,
|
||||
AccountRefreshResult, ConsumeProviderResetCredit, PrepareCredentialRefresh,
|
||||
ProviderModels, ProviderProfileAvatar, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderQuotaWindow, ProviderResetCreditResult, ProviderResetCredits,
|
||||
QuotaLocalUsageAttribution,
|
||||
},
|
||||
quota_forecast::{AccountQuotaForecastReport, account_quota_forecasts},
|
||||
quota_forecast_sampling::{QuotaForecastPoint, select_forecast_sample},
|
||||
},
|
||||
ports::{
|
||||
provider::ProviderAdminRegistry,
|
||||
@@ -83,11 +85,16 @@ pub trait AccountsService: Send + Sync {
|
||||
refresh: bool,
|
||||
) -> Result<AccountDirectoryItem, AdminError>;
|
||||
|
||||
async fn profile_statistics(
|
||||
async fn quota_forecast(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<AccountQuotaForecastReport, AdminError>;
|
||||
|
||||
async fn personal_info(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileStatistics, AdminError> {
|
||||
Err(AdminError::invalid("当前 Provider 不支持账号统计"))
|
||||
) -> Result<AccountPersonalInfo, AdminError> {
|
||||
Err(AdminError::invalid("当前 Provider 不支持个人信息"))
|
||||
}
|
||||
|
||||
async fn profile_avatar(
|
||||
@@ -213,6 +220,7 @@ impl DefaultAccountsService {
|
||||
quota
|
||||
.windows
|
||||
.iter()
|
||||
.filter(|window| window.local_usage.is_none())
|
||||
.filter_map(|window| quota_usage_window(&item.account.id, window))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -581,15 +589,127 @@ impl AccountsService for DefaultAccountsService {
|
||||
self.load_directory_item(account_id, refresh).await
|
||||
}
|
||||
|
||||
async fn profile_statistics(
|
||||
async fn quota_forecast(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileStatistics, AdminError> {
|
||||
let (_, provider) = self.provider_for_account(account_id).await?;
|
||||
provider
|
||||
.profile_statistics(account_id)
|
||||
) -> Result<AccountQuotaForecastReport, AdminError> {
|
||||
let (stored, provider) = self.provider_for_account(account_id).await?;
|
||||
let quota = provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: account_id.clone(),
|
||||
refresh: false,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|error| map_provider_error(error, "provider profile statistics"))
|
||||
.map_err(|error| map_provider_error(error, "forecast quota snapshot"))?;
|
||||
let now = Utc::now();
|
||||
let mut samples = Vec::new();
|
||||
for (window, _) in quota.usage_windows() {
|
||||
let (Some(mut query), Some(observed), Some(percent)) = (
|
||||
quota_usage_window(account_id.as_str(), window),
|
||||
quota.observed_at,
|
||||
window.used_percent,
|
||||
) else {
|
||||
continue;
|
||||
};
|
||||
query.range.start = query.range.start.max(stored.account.created_at);
|
||||
if query.range.start >= observed
|
||||
|| observed > now
|
||||
|| now >= query.range.end
|
||||
|| !percent.is_finite()
|
||||
|| !(0.0..=100.0).contains(&percent)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let reset_at = query.range.end;
|
||||
query.range.end = observed;
|
||||
let history = self
|
||||
.accounts
|
||||
.load_quota_forecast_history(&query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "forecast paired usage"))?;
|
||||
let mut points = Vec::new();
|
||||
let mut interrupted = false;
|
||||
for point in history.points {
|
||||
let Some(fact) =
|
||||
provider.quota_forecast_observation(&point.provider_observation, window)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let same_plan = quota
|
||||
.plan_type
|
||||
.as_deref()
|
||||
.zip(fact.plan_type.as_deref())
|
||||
.is_some_and(|(current, previous)| current.eq_ignore_ascii_case(previous));
|
||||
// 仅容纳已观测到的秒级量化抖动,不用宽时间容差合并实际重置。
|
||||
// 不匹配的段截断基线;之后的有效观测可以重新积累。
|
||||
if !same_plan || (fact.reset_at - reset_at).abs() > Duration::seconds(2) {
|
||||
points.clear();
|
||||
interrupted = true;
|
||||
continue;
|
||||
}
|
||||
points.push(QuotaForecastPoint {
|
||||
observed_at: point.completed_at,
|
||||
used_percent: fact.used_percent,
|
||||
usage: point.usage,
|
||||
});
|
||||
}
|
||||
let mut sample = select_forecast_sample(
|
||||
window.key.clone(),
|
||||
query.range.start,
|
||||
QuotaForecastPoint {
|
||||
observed_at: observed,
|
||||
used_percent: percent,
|
||||
usage: history.usage,
|
||||
},
|
||||
points,
|
||||
history.pending_request_count,
|
||||
);
|
||||
if interrupted
|
||||
&& sample.method
|
||||
== crate::model::quota_forecast_sampling::QuotaForecastMethod::Cumulative
|
||||
{
|
||||
sample.discontinuous = true;
|
||||
}
|
||||
samples.push(sample);
|
||||
}
|
||||
Ok(AccountQuotaForecastReport {
|
||||
account_id: account_id.to_string(),
|
||||
generated_at: now,
|
||||
forecasts: account_quota_forecasts("a, stored.account.created_at, now, &samples),
|
||||
})
|
||||
}
|
||||
|
||||
async fn personal_info(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<AccountPersonalInfo, AdminError> {
|
||||
let (initial, provider) = self.provider_for_account(account_id).await?;
|
||||
// 两项读取相互独立;不因其中一项失败而取消另一项,也不触发凭据或额度刷新。
|
||||
let (profile, subscription) = futures::join!(
|
||||
provider.profile_statistics(account_id),
|
||||
provider.subscription(account_id),
|
||||
);
|
||||
let profile =
|
||||
profile.map_err(|error| map_provider_error(error, "provider profile statistics"));
|
||||
let subscription = subscription
|
||||
.map_err(|error| map_provider_error(error, "provider subscription"))
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
// 汇聚等待期间发生重新授权、换绑或删除时,不返回混合身份的数据。
|
||||
let current = self.load_account(account_id).await?;
|
||||
if current.account.provider_kind != initial.account.provider_kind
|
||||
|| current.account.credential_revision != initial.account.credential_revision
|
||||
|| current.account.upstream_user_id != initial.account.upstream_user_id
|
||||
|| current.account.upstream_account_id != initial.account.upstream_account_id
|
||||
{
|
||||
return Err(AdminError::conflict("账号身份已变化,请刷新信息后重试"));
|
||||
}
|
||||
Ok(AccountPersonalInfo {
|
||||
profile,
|
||||
subscription,
|
||||
})
|
||||
}
|
||||
|
||||
async fn profile_avatar(
|
||||
@@ -762,14 +882,12 @@ fn quota_usage_window(
|
||||
account_id: &str,
|
||||
window: &ProviderQuotaWindow,
|
||||
) -> Option<AccountUsageWindowQuery> {
|
||||
if window.local_usage.is_some()
|
||||
|| window.local_usage_attribution != QuotaLocalUsageAttribution::AccountWide
|
||||
{
|
||||
if window.local_usage_attribution != QuotaLocalUsageAttribution::AccountWide {
|
||||
return None;
|
||||
}
|
||||
let reset_at = window.reset_at?;
|
||||
let seconds = i64::try_from(window.window_seconds?).ok()?;
|
||||
let start = reset_at.checked_sub_signed(Duration::seconds(seconds))?;
|
||||
let start = reset_at.checked_sub_signed(Duration::try_seconds(seconds)?)?;
|
||||
let range = TimeRange::new(start, reset_at).ok()?;
|
||||
// 上游百分比以该 reset 边界定义;以当前时间回推会让本地 Token 属于另一窗口。
|
||||
Some(AccountUsageWindowQuery {
|
||||
|
||||
@@ -18,7 +18,7 @@ use crate::{
|
||||
DeleteClientKey, NewClientKey, SetClientKeyEnabled, UpdateClientKey,
|
||||
},
|
||||
},
|
||||
ports::store::ClientKeyStore,
|
||||
ports::store::{AdminStoreError, AdminStoreErrorKind, ClientKeyStore},
|
||||
};
|
||||
|
||||
use super::{map_store_error, publish_committed};
|
||||
@@ -87,9 +87,13 @@ impl ClientKeyService for DefaultClientKeyService {
|
||||
) -> Result<CreatedClientKey, AdminError> {
|
||||
let id = ClientApiKeyId::new(format!("key_{}", Uuid::now_v7().simple()))
|
||||
.map_err(|_| AdminError::internal("创建 Client API Key ID 失败"))?;
|
||||
let mut bytes = [0_u8; 32];
|
||||
OsRng.fill_bytes(&mut bytes);
|
||||
let plaintext = format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes));
|
||||
let plaintext = if let Some(key) = command.custom_key {
|
||||
key.expose_for_auth().to_owned()
|
||||
} else {
|
||||
let mut bytes = [0_u8; 32];
|
||||
OsRng.fill_bytes(&mut bytes);
|
||||
format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes))
|
||||
};
|
||||
let (config_revision, record) = self
|
||||
.store
|
||||
.create_client_key(
|
||||
@@ -105,7 +109,7 @@ impl ClientKeyService for DefaultClientKeyService {
|
||||
context,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "client API key"))?;
|
||||
.map_err(map_client_key_write_error)?;
|
||||
publish_committed(self.snapshot.as_ref(), config_revision).await?;
|
||||
Ok(CreatedClientKey {
|
||||
config_revision,
|
||||
@@ -123,7 +127,7 @@ impl ClientKeyService for DefaultClientKeyService {
|
||||
.store
|
||||
.update_client_key(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "client API key"))?;
|
||||
.map_err(map_client_key_write_error)?;
|
||||
publish_committed(self.snapshot.as_ref(), config_revision).await?;
|
||||
Ok(ClientKeyMutation {
|
||||
config_revision,
|
||||
@@ -171,6 +175,14 @@ impl ClientKeyService for DefaultClientKeyService {
|
||||
}
|
||||
}
|
||||
|
||||
fn map_client_key_write_error(error: AdminStoreError) -> AdminError {
|
||||
match error.kind() {
|
||||
AdminStoreErrorKind::DuplicateName => AdminError::conflict("名称已存在"),
|
||||
AdminStoreErrorKind::Conflict => AdminError::conflict("API Key 已存在,请使用其他密钥"),
|
||||
_ => map_store_error(error, "client API key"),
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_cursor(query: &ClientKeyListQuery) -> Result<(), AdminError> {
|
||||
let Some(cursor) = &query.cursor else {
|
||||
return Ok(());
|
||||
|
||||
@@ -13,6 +13,8 @@ pub mod settings;
|
||||
pub mod system;
|
||||
pub mod xai;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind, MutationContext,
|
||||
@@ -22,7 +24,7 @@ use crate::{
|
||||
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
|
||||
PendingAuthorizationMutation, PreparedAuthorizationCommit,
|
||||
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
|
||||
StartAuthorization,
|
||||
ProviderQuotaRequest, StartAuthorization,
|
||||
},
|
||||
},
|
||||
ports::{
|
||||
@@ -40,9 +42,9 @@ fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError
|
||||
let kind = match error.kind() {
|
||||
AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid,
|
||||
AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound,
|
||||
AdminStoreErrorKind::StaleRevision | AdminStoreErrorKind::Conflict => {
|
||||
AdminErrorKind::Conflict
|
||||
}
|
||||
AdminStoreErrorKind::StaleRevision
|
||||
| AdminStoreErrorKind::DuplicateName
|
||||
| AdminStoreErrorKind::Conflict => AdminErrorKind::Conflict,
|
||||
AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable,
|
||||
};
|
||||
tracing::warn!(resource, error_kind = ?error.kind(), "admin store operation failed");
|
||||
@@ -97,6 +99,43 @@ async fn publish_committed(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn publish_credentials_and_observe_quota(
|
||||
provider: &Arc<dyn ProviderAdmin>,
|
||||
snapshot: &dyn SnapshotControl,
|
||||
revision: crate::model::Revision,
|
||||
account_ids: &[ProviderAccountId],
|
||||
request_id: &str,
|
||||
) -> Result<(), AdminError> {
|
||||
provider.account_facts_changed(account_ids).await;
|
||||
publish_committed(snapshot, revision).await?;
|
||||
|
||||
// 凭据已经提交;额度只是可重建的观察,不能拖住管理请求或回滚提交结果。
|
||||
let provider = Arc::clone(provider);
|
||||
let account_ids = account_ids.to_vec();
|
||||
let request_id = request_id.to_owned();
|
||||
tokio::spawn(async move {
|
||||
for account_id in account_ids {
|
||||
if let Err(error) = provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: account_id.clone(),
|
||||
refresh: true,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
request_id,
|
||||
provider = %provider.provider_kind().as_str(),
|
||||
account_id = %account_id.as_str(),
|
||||
quota_error = ?error.kind(),
|
||||
"initial quota observation failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn required_credential(
|
||||
accounts: &dyn AccountStore,
|
||||
provider_kind: &ProviderKind,
|
||||
|
||||
@@ -3,17 +3,16 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::{account::ProviderAccountId, runtime::SnapshotControl};
|
||||
use gateway_core::runtime::SnapshotControl;
|
||||
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError,
|
||||
provider_credentials::{
|
||||
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
|
||||
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
|
||||
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
|
||||
ImportCredentials, PrepareCredentialImport, PrepareCredentialRotation,
|
||||
ProviderQuotaRequest, RotateCredential, StartAuthorization,
|
||||
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
|
||||
CredentialMutationResult, ImportCredentials, PrepareCredentialImport,
|
||||
PrepareCredentialRotation, RotateCredential, StartAuthorization,
|
||||
},
|
||||
},
|
||||
ports::{provider::ProviderAdmin, store::AccountStore},
|
||||
@@ -21,18 +20,14 @@ use crate::{
|
||||
|
||||
use super::{
|
||||
commit_authorization, commit_credential_rotation, delete_credentials, map_provider_error,
|
||||
map_store_error, pending_authorization, publish_committed, required_credential,
|
||||
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
|
||||
map_store_error, pending_authorization, publish_committed,
|
||||
publish_credentials_and_observe_quota, required_credential, validate_authorization_commit,
|
||||
validate_prepared_import, validate_prepared_rotation,
|
||||
};
|
||||
|
||||
/// OpenAI 固定管理路由消费的服务。
|
||||
#[async_trait]
|
||||
pub trait OpenAiService: Send + Sync {
|
||||
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError>;
|
||||
async fn details(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CredentialDetails, AdminError>;
|
||||
async fn import_document(
|
||||
&self,
|
||||
command: ImportCredentials,
|
||||
@@ -77,55 +72,10 @@ impl DefaultOpenAiService {
|
||||
snapshot,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_initial_quotas(&self, account_ids: &[ProviderAccountId], request_id: &str) {
|
||||
let provider = Arc::clone(&self.provider);
|
||||
let account_ids = account_ids.to_vec();
|
||||
let request_id = request_id.to_owned();
|
||||
tokio::spawn(async move {
|
||||
for account_id in account_ids {
|
||||
if let Err(error) = provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: account_id.clone(),
|
||||
refresh: true,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
request_id,
|
||||
account_id = %account_id.as_str(),
|
||||
quota_error = ?error.kind(),
|
||||
"OpenAI initial quota observation failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl OpenAiService for DefaultOpenAiService {
|
||||
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError> {
|
||||
self.accounts
|
||||
.list_credentials(self.provider.provider_kind(), query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "OpenAI credential"))
|
||||
}
|
||||
|
||||
async fn details(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CredentialDetails, AdminError> {
|
||||
required_credential(
|
||||
self.accounts.as_ref(),
|
||||
self.provider.provider_kind(),
|
||||
account_id,
|
||||
"OpenAI credential",
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn import_document(
|
||||
&self,
|
||||
command: ImportCredentials,
|
||||
@@ -167,12 +117,14 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
|
||||
drop(proxy_reservation);
|
||||
self.provider
|
||||
.account_facts_changed(&result.credential_ids)
|
||||
.await;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
// 导入完成后立即做一次观察;失败不回滚已经提交的 credential。
|
||||
self.observe_initial_quotas(&result.credential_ids, &context.request_id);
|
||||
publish_credentials_and_observe_quota(
|
||||
&self.provider,
|
||||
self.snapshot.as_ref(),
|
||||
result.config_revision,
|
||||
&result.credential_ids,
|
||||
&context.request_id,
|
||||
)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
@@ -220,14 +172,14 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
"OpenAI authorization",
|
||||
)
|
||||
.await?;
|
||||
self.provider
|
||||
.account_facts_changed(std::slice::from_ref(&result.account_id))
|
||||
.await;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
self.observe_initial_quotas(
|
||||
publish_credentials_and_observe_quota(
|
||||
&self.provider,
|
||||
self.snapshot.as_ref(),
|
||||
result.config_revision,
|
||||
std::slice::from_ref(&result.account_id),
|
||||
&context.request_id,
|
||||
);
|
||||
)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
|
||||
@@ -11,8 +11,8 @@ use crate::{
|
||||
provider_credentials::{
|
||||
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
|
||||
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
|
||||
CredentialListQuery, CredentialMutationResult, CredentialPage, ImportCredentials,
|
||||
PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
|
||||
CredentialMutationResult, ImportCredentials, PrepareCredentialImport,
|
||||
StartAuthorization,
|
||||
},
|
||||
},
|
||||
ports::{provider::ProviderAdmin, store::AccountStore},
|
||||
@@ -20,14 +20,13 @@ use crate::{
|
||||
|
||||
use super::{
|
||||
commit_authorization, delete_credentials, map_provider_error, map_store_error,
|
||||
pending_authorization, publish_committed, validate_authorization_commit,
|
||||
validate_prepared_import,
|
||||
pending_authorization, publish_committed, publish_credentials_and_observe_quota,
|
||||
validate_authorization_commit, validate_prepared_import,
|
||||
};
|
||||
|
||||
/// xAI 固定管理路由消费的服务。
|
||||
#[async_trait]
|
||||
pub trait XaiService: Send + Sync {
|
||||
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError>;
|
||||
async fn import_document(
|
||||
&self,
|
||||
command: ImportCredentials,
|
||||
@@ -72,13 +71,6 @@ impl DefaultXaiService {
|
||||
|
||||
#[async_trait]
|
||||
impl XaiService for DefaultXaiService {
|
||||
async fn list(&self, query: CredentialListQuery) -> Result<CredentialPage, AdminError> {
|
||||
self.accounts
|
||||
.list_credentials(self.provider.provider_kind(), query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "xAI credential"))
|
||||
}
|
||||
|
||||
async fn import_document(
|
||||
&self,
|
||||
command: ImportCredentials,
|
||||
@@ -120,18 +112,14 @@ impl XaiService for DefaultXaiService {
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "xAI credential import"))?;
|
||||
drop(proxy_reservation);
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
// 导入已经提交,额度属于可重建观察事实;单个账号查询失败由周期任务重试。
|
||||
for account_id in &result.credential_ids {
|
||||
let _ = self
|
||||
.provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: account_id.clone(),
|
||||
refresh: true,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
publish_credentials_and_observe_quota(
|
||||
&self.provider,
|
||||
self.snapshot.as_ref(),
|
||||
result.config_revision,
|
||||
&result.credential_ids,
|
||||
&context.request_id,
|
||||
)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
@@ -179,7 +167,14 @@ impl XaiService for DefaultXaiService {
|
||||
"xAI authorization",
|
||||
)
|
||||
.await?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
publish_credentials_and_observe_quota(
|
||||
&self.provider,
|
||||
self.snapshot.as_ref(),
|
||||
result.config_revision,
|
||||
std::slice::from_ref(&result.account_id),
|
||||
&context.request_id,
|
||||
)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
|
||||
@@ -2,3 +2,5 @@
|
||||
|
||||
mod backup;
|
||||
mod provider_credentials;
|
||||
mod quota_forecast;
|
||||
mod quota_forecast_sampling;
|
||||
|
||||
@@ -0,0 +1,382 @@
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use gateway_admin::model::{
|
||||
accounts::{AccountCost, AccountUsage},
|
||||
observability::CostCoverage,
|
||||
provider_credentials::{ProviderQuota, ProviderQuotaWindow, QuotaLocalUsageAttribution},
|
||||
quota_forecast::{AccountQuotaForecast, account_quota_forecasts},
|
||||
quota_forecast_sampling::{QuotaForecastMethod, QuotaForecastSample, QuotaForecastUsage},
|
||||
};
|
||||
|
||||
fn now() -> DateTime<Utc> {
|
||||
"2026-09-12T00:00:00Z".parse().unwrap()
|
||||
}
|
||||
|
||||
fn window(key: &str, days: u64) -> ProviderQuotaWindow {
|
||||
ProviderQuotaWindow {
|
||||
key: key.to_owned(),
|
||||
group: if days > 7 { "monthly" } else { "shortTerm" }.to_owned(),
|
||||
label: key.to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(days * 86_400),
|
||||
used_percent: Some(20.0),
|
||||
reset_at: Some(now() + Duration::days(1)),
|
||||
limit_reached: false,
|
||||
local_usage: Some(AccountUsage {
|
||||
account_id: "acct_forecast".to_owned(),
|
||||
request_count: 10,
|
||||
success_count: 10,
|
||||
input_tokens: Some(900),
|
||||
output_tokens: Some(100),
|
||||
cached_tokens: Some(800),
|
||||
cache_write_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
image_input_tokens: None,
|
||||
image_output_tokens: None,
|
||||
image_request_count: 0,
|
||||
image_request_failed_count: 0,
|
||||
total_tokens: Some(1_000),
|
||||
cost_coverage: CostCoverage {
|
||||
calculated_count: 10,
|
||||
..Default::default()
|
||||
},
|
||||
costs: vec![AccountCost {
|
||||
currency: "USD".to_owned(),
|
||||
amount: "2".parse().unwrap(),
|
||||
}],
|
||||
last_used_at: Some(now()),
|
||||
request_buckets: Vec::new(),
|
||||
models: Vec::new(),
|
||||
}),
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota(windows: Vec<ProviderQuotaWindow>) -> ProviderQuota {
|
||||
ProviderQuota {
|
||||
plan_type: None,
|
||||
observed_at: Some(now()),
|
||||
refresh_token_expires_at: None,
|
||||
windows,
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn forecast(quota: &ProviderQuota) -> [AccountQuotaForecast; 2] {
|
||||
account_quota_forecasts(quota, now() - Duration::days(60), now(), &samples(quota))
|
||||
}
|
||||
|
||||
fn samples(quota: &ProviderQuota) -> Vec<QuotaForecastSample> {
|
||||
quota
|
||||
.windows
|
||||
.iter()
|
||||
.filter_map(|window| {
|
||||
let usage = window.local_usage.as_ref()?;
|
||||
let usd = usage.costs.iter().find(|cost| cost.currency == "USD");
|
||||
Some(QuotaForecastSample {
|
||||
key: window.key.clone(),
|
||||
method: QuotaForecastMethod::Cumulative,
|
||||
start_at: window.reset_at?.checked_sub_signed(Duration::try_seconds(
|
||||
i64::try_from(window.window_seconds?).ok()?,
|
||||
)?)?,
|
||||
end_at: quota.observed_at.unwrap_or(now()),
|
||||
baseline_percent: 0.0,
|
||||
sampled_percent: window.used_percent.unwrap_or(0.0),
|
||||
block_count: 0,
|
||||
observation_count: 0,
|
||||
pending_request_count: 0,
|
||||
discontinuous: false,
|
||||
usage: QuotaForecastUsage {
|
||||
request_count: usage.request_count,
|
||||
tokens: usage.total_tokens.unwrap_or(0),
|
||||
input_tokens: usage.input_tokens.unwrap_or(0),
|
||||
output_tokens: usage.output_tokens.unwrap_or(0),
|
||||
cached_tokens: usage.cached_tokens.unwrap_or(0),
|
||||
missing_token_count: u64::from(usage.total_tokens.is_none()),
|
||||
known_cost_count: if usd.is_some() {
|
||||
usage.cost_coverage.known_count()
|
||||
} else {
|
||||
0
|
||||
},
|
||||
unavailable_cost_count: usage.cost_coverage.partial_count
|
||||
+ usage.cost_coverage.unavailable_count,
|
||||
usd: usd.map_or(0.0, |cost| cost.amount.as_str().parse().unwrap()),
|
||||
excluded_request_count: 0,
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forecasts_use_actual_periods_and_do_not_extrapolate_remaining_capacity() {
|
||||
let [week, month] = forecast("a(vec![window("week", 7)]));
|
||||
assert_eq!(week.estimated_tokens, Some(5_000));
|
||||
assert_eq!(week.estimated_usd, Some(10.0));
|
||||
assert_eq!(week.remaining_tokens, Some(4_000));
|
||||
assert_eq!(week.remaining_usd, Some(8.0));
|
||||
assert!(!week.extrapolated);
|
||||
assert!(month.extrapolated);
|
||||
assert_eq!(month.estimated_tokens, Some(21_429));
|
||||
assert!((month.estimated_usd.unwrap() - 10.0 * 30.0 / 7.0).abs() < 1e-10);
|
||||
assert_eq!(month.remaining_tokens, week.remaining_tokens);
|
||||
assert_eq!(month.remaining_usd, week.remaining_usd);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corresponding_actual_window_takes_priority_and_retains_its_duration() {
|
||||
let [week, month] = forecast("a(vec![window("month", 31), window("week", 7)]));
|
||||
assert_eq!(week.source.unwrap().label, "week");
|
||||
assert_eq!(month.source.unwrap().label, "month");
|
||||
assert!(!month.extrapolated);
|
||||
assert_eq!(month.target_seconds, 31 * 86_400);
|
||||
assert_eq!(month.estimated_tokens, Some(5_000));
|
||||
let [week, _] = forecast("a(vec![window("month", 31)]));
|
||||
assert!(week.extrapolated);
|
||||
assert_eq!(week.estimated_tokens, Some(1_129));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn short_or_model_specific_windows_are_not_account_capacity() {
|
||||
let mut model = window("model", 7);
|
||||
model.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
|
||||
for item in forecast("a(vec![window("day", 1), model])) {
|
||||
assert!(item.source.is_none());
|
||||
assert!(item.unavailable_reason.is_some());
|
||||
assert!(item.estimated_tokens.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_or_tiny_percentages_cannot_produce_estimates() {
|
||||
for percent in [
|
||||
None,
|
||||
Some(f64::NAN),
|
||||
Some(f64::INFINITY),
|
||||
Some(-1.0),
|
||||
Some(101.0),
|
||||
Some(0.0),
|
||||
Some(0.5),
|
||||
Some(1.0),
|
||||
Some(4.9),
|
||||
] {
|
||||
let mut sample = window("week", 7);
|
||||
sample.used_percent = percent;
|
||||
let [week, _] = forecast("a(vec![sample]));
|
||||
assert!(week.unavailable_reason.is_some(), "{percent:?}");
|
||||
assert!(week.estimated_tokens.is_none());
|
||||
assert!(week.estimated_usd.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn low_samples_warn_and_exhaustion_uses_raw_not_display_percentages() {
|
||||
let mut sample = window("week", 7);
|
||||
sample.used_percent = Some(5.0);
|
||||
sample.limit_reached = true;
|
||||
let mut source = quota(vec![sample]);
|
||||
source.limit_reached = true;
|
||||
let [week, _] = forecast(&source);
|
||||
assert!(week.low_sample);
|
||||
assert_eq!(week.estimated_tokens, Some(20_000));
|
||||
source.windows[0].used_percent = Some(100.0);
|
||||
let [week, _] = forecast(&source);
|
||||
assert!(!week.low_sample);
|
||||
assert_eq!(week.remaining_tokens, Some(0));
|
||||
assert_eq!(week.remaining_usd, Some(0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_future_or_missing_snapshots_do_not_predict() {
|
||||
for observed in [
|
||||
None,
|
||||
Some(now() - Duration::days(8)),
|
||||
Some(now() + Duration::seconds(1)),
|
||||
] {
|
||||
let mut source = quota(vec![window("week", 7)]);
|
||||
source.observed_at = observed;
|
||||
assert!(forecast(&source)[0].unavailable_reason.is_some());
|
||||
}
|
||||
let mut sample = window("week", 7);
|
||||
sample.reset_at = Some(now());
|
||||
assert!(
|
||||
forecast("a(vec![sample]))[0]
|
||||
.unavailable_reason
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_account_history_or_missing_usage_keep_diagnostics_without_estimates() {
|
||||
let source = quota(vec![window("week", 7)]);
|
||||
let [week, _] =
|
||||
account_quota_forecasts(&source, now() - Duration::days(1), now(), &samples(&source));
|
||||
assert!(week.unavailable_reason.unwrap().contains("记录不完整"));
|
||||
assert!(week.estimated_tokens.is_none());
|
||||
let mut sample = window("week", 7);
|
||||
sample.local_usage = None;
|
||||
let [week, _] = forecast("a(vec![sample]));
|
||||
assert!(week.source.is_some());
|
||||
assert!(
|
||||
week.unavailable_reason
|
||||
.unwrap()
|
||||
.contains("没有网关用量记录")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_sample_supports_mid_cycle_accounts_and_remaining_uses_current_percent() {
|
||||
let source = quota(vec![window("week", 7)]);
|
||||
let mut sample = samples(&source).remove(0);
|
||||
sample.method = QuotaForecastMethod::Incremental;
|
||||
sample.start_at = now() - Duration::hours(6);
|
||||
sample.baseline_percent = 10.0;
|
||||
sample.sampled_percent = 10.0;
|
||||
sample.block_count = 2;
|
||||
let [week, month] =
|
||||
account_quota_forecasts(&source, now() - Duration::days(1), now(), &[sample]);
|
||||
assert_eq!(week.estimated_tokens, Some(10_000));
|
||||
assert_eq!(week.remaining_tokens, Some(8_000));
|
||||
assert_eq!(month.remaining_tokens, Some(8_000));
|
||||
assert!(week.unavailable_reason.is_none());
|
||||
assert!(!week.low_sample);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_tokens_keep_estimates_from_recorded_usage() {
|
||||
let source = quota(vec![window("week", 7)]);
|
||||
let mut sample = samples(&source).remove(0);
|
||||
sample.usage.missing_token_count = 1;
|
||||
let [week, _] = account_quota_forecasts(&source, now() - Duration::days(60), now(), &[sample]);
|
||||
assert!(week.incomplete_tokens);
|
||||
assert!(week.unavailable_reason.is_none());
|
||||
assert_eq!(week.estimated_tokens, Some(5_000));
|
||||
assert_eq!(week.remaining_tokens, Some(4_000));
|
||||
assert_eq!(week.estimated_usd, Some(10.0));
|
||||
assert_eq!(week.source.unwrap().tokens, Some(1_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discontinuous_samples_never_fall_back_to_cumulative_predictions() {
|
||||
let source = quota(vec![window("week", 7)]);
|
||||
let mut sample = samples(&source).remove(0);
|
||||
sample.discontinuous = true;
|
||||
let [week, _] = account_quota_forecasts(&source, now() - Duration::days(60), now(), &[sample]);
|
||||
assert!(week.unavailable_reason.unwrap().contains("不连续"));
|
||||
assert!(week.estimated_tokens.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_costs_keep_estimates_from_known_amounts() {
|
||||
for coverage in [
|
||||
CostCoverage {
|
||||
calculated_count: 9,
|
||||
partial_count: 1,
|
||||
..Default::default()
|
||||
},
|
||||
CostCoverage {
|
||||
calculated_count: 9,
|
||||
unavailable_count: 1,
|
||||
..Default::default()
|
||||
},
|
||||
] {
|
||||
let mut sample = window("week", 7);
|
||||
sample.local_usage.as_mut().unwrap().cost_coverage = coverage;
|
||||
let [week, _] = forecast("a(vec![sample]));
|
||||
assert!(week.incomplete_cost);
|
||||
assert!(week.unavailable_reason.is_none());
|
||||
assert_eq!(week.estimated_usd, Some(10.0));
|
||||
assert_eq!(week.remaining_usd, Some(8.0));
|
||||
assert_eq!(week.estimated_tokens, Some(5_000));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_tokens_and_costs_do_not_block_recorded_capacity_or_remaining_estimates() {
|
||||
let source = quota(vec![window("week", 7)]);
|
||||
let mut sample = samples(&source).remove(0);
|
||||
// 完整请求之外混有缺少计量的成功请求,不把已有用量和费用一并作废。
|
||||
sample.usage.request_count += 2;
|
||||
sample.usage.missing_token_count = 2;
|
||||
sample.usage.unavailable_cost_count = 2;
|
||||
sample.usage.excluded_request_count = 1;
|
||||
sample.pending_request_count = 1;
|
||||
let [week, month] =
|
||||
account_quota_forecasts(&source, now() - Duration::days(60), now(), &[sample]);
|
||||
assert!(week.incomplete_tokens);
|
||||
assert!(week.incomplete_cost);
|
||||
assert!(week.unavailable_reason.is_none());
|
||||
assert_eq!(week.estimated_tokens, Some(5_000));
|
||||
assert_eq!(week.estimated_usd, Some(10.0));
|
||||
assert_eq!(week.remaining_tokens, Some(4_000));
|
||||
assert_eq!(week.remaining_usd, Some(8.0));
|
||||
assert_eq!(month.estimated_tokens, Some(21_429));
|
||||
assert_eq!(month.remaining_tokens, week.remaining_tokens);
|
||||
assert_eq!(month.remaining_usd, week.remaining_usd);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entirely_unknown_costs_leave_only_money_estimates_unavailable() {
|
||||
let mut sample = window("week", 7);
|
||||
let usage = sample.local_usage.as_mut().unwrap();
|
||||
usage.cost_coverage = CostCoverage {
|
||||
unavailable_count: usage.request_count,
|
||||
..Default::default()
|
||||
};
|
||||
usage.costs.clear();
|
||||
let [week, _] = forecast("a(vec![sample]));
|
||||
assert!(week.incomplete_cost);
|
||||
assert!(week.unavailable_reason.is_none());
|
||||
assert_eq!(week.estimated_tokens, Some(5_000));
|
||||
assert_eq!(week.estimated_usd, None);
|
||||
assert_eq!(week.remaining_usd, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entirely_unknown_tokens_and_costs_do_not_invent_zero_estimates() {
|
||||
let source = quota(vec![window("week", 7)]);
|
||||
let mut sample = samples(&source).remove(0);
|
||||
sample.usage = QuotaForecastUsage {
|
||||
request_count: 2,
|
||||
missing_token_count: 2,
|
||||
unavailable_cost_count: 2,
|
||||
..Default::default()
|
||||
};
|
||||
let [week, _] = account_quota_forecasts(&source, now() - Duration::days(60), now(), &[sample]);
|
||||
assert!(week.unavailable_reason.is_some());
|
||||
assert_eq!(week.estimated_tokens, None);
|
||||
assert_eq!(week.estimated_usd, None);
|
||||
assert_eq!(week.remaining_tokens, None);
|
||||
assert_eq!(week.remaining_usd, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_known_cost_is_valid_but_missing_usd_is_not_invented() {
|
||||
let mut sample = window("week", 7);
|
||||
sample.local_usage.as_mut().unwrap().costs[0].amount = "0".parse().unwrap();
|
||||
assert_eq!(
|
||||
forecast("a(vec![sample.clone()]))[0].estimated_usd,
|
||||
Some(0.0)
|
||||
);
|
||||
sample.local_usage.as_mut().unwrap().costs[0].currency = "EUR".to_owned();
|
||||
assert_eq!(forecast("a(vec![sample]))[0].estimated_usd, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oversized_window_or_estimate_is_unavailable_without_panics_or_saturation() {
|
||||
let mut sample = window("month", 30);
|
||||
sample.window_seconds = Some(i64::MAX as u64);
|
||||
assert!(
|
||||
forecast("a(vec![sample]))[1]
|
||||
.unavailable_reason
|
||||
.is_some()
|
||||
);
|
||||
let mut sample = window("week", 7);
|
||||
sample.local_usage.as_mut().unwrap().total_tokens = Some(u64::MAX);
|
||||
let [week, _] = forecast("a(vec![sample]));
|
||||
assert!(week.estimated_tokens.is_none());
|
||||
assert_eq!(week.estimated_usd, Some(10.0));
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use gateway_admin::model::quota_forecast_sampling::{
|
||||
QuotaForecastMethod, QuotaForecastPoint, QuotaForecastSample, QuotaForecastUsage,
|
||||
select_forecast_sample,
|
||||
};
|
||||
|
||||
fn time(minutes: i64) -> DateTime<Utc> {
|
||||
"2026-09-12T00:00:00Z".parse::<DateTime<Utc>>().unwrap() + Duration::minutes(minutes)
|
||||
}
|
||||
|
||||
fn point(minute: i64, percent: f64, tokens: u64) -> QuotaForecastPoint {
|
||||
QuotaForecastPoint {
|
||||
observed_at: time(minute),
|
||||
used_percent: percent,
|
||||
usage: QuotaForecastUsage {
|
||||
request_count: minute as u64,
|
||||
tokens,
|
||||
input_tokens: tokens,
|
||||
known_cost_count: minute as u64,
|
||||
usd: tokens as f64 / 1_000.0,
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn select(points: Vec<QuotaForecastPoint>, current: QuotaForecastPoint) -> QuotaForecastSample {
|
||||
select_forecast_sample("week".to_owned(), time(0), current, points, 2)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repeated_readings_do_not_make_independent_blocks_or_shift_the_baseline() {
|
||||
let points = (1..100)
|
||||
.map(|minute| point(minute, 20.0, minute as u64 * 100))
|
||||
.collect();
|
||||
let sample = select(points, point(100, 25.0, 10_000));
|
||||
assert_eq!(sample.method, QuotaForecastMethod::Incremental);
|
||||
assert_eq!(sample.block_count, 1);
|
||||
assert_eq!(sample.start_at, time(1));
|
||||
assert_eq!(sample.sampled_percent, 5.0);
|
||||
assert_eq!(sample.usage.tokens, 9_900);
|
||||
assert_eq!(sample.pending_request_count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_blocks_and_tail_replace_old_workload_without_averaging_individual_ratios() {
|
||||
let points = vec![
|
||||
point(1, 0.0, 0),
|
||||
point(2, 5.0, 100),
|
||||
point(3, 10.0, 200),
|
||||
point(4, 15.0, 400),
|
||||
point(5, 20.0, 600),
|
||||
point(6, 25.0, 800),
|
||||
];
|
||||
let sample = select(points, point(7, 27.0, 900));
|
||||
assert_eq!(sample.baseline_percent, 10.0);
|
||||
assert_eq!(sample.sampled_percent, 17.0);
|
||||
assert_eq!(sample.usage.tokens, 700);
|
||||
assert_eq!(sample.block_count, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_at_a_block_boundary_keeps_only_three_recent_blocks() {
|
||||
let points = (0..6)
|
||||
.map(|i| point(i + 1, i as f64 * 5.0, i as u64 * 100))
|
||||
.collect();
|
||||
let sample = select(points, point(7, 30.0, 600));
|
||||
assert_eq!(sample.baseline_percent, 15.0);
|
||||
assert_eq!(sample.sampled_percent, 15.0);
|
||||
assert_eq!(sample.usage.tokens, 300);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_percentage_reversal_restarts_sampling_without_crossing_old_consumption() {
|
||||
let sample = select(
|
||||
vec![point(1, 40.0, 100), point(2, 50.0, 200), point(3, 1.0, 210)],
|
||||
point(4, 8.0, 280),
|
||||
);
|
||||
assert_eq!(sample.baseline_percent, 1.0);
|
||||
assert_eq!(sample.usage.tokens, 70);
|
||||
assert_eq!(sample.sampled_percent, 7.0);
|
||||
assert!(!sample.discontinuous);
|
||||
let waiting = select(
|
||||
vec![point(1, 40.0, 100), point(2, 50.0, 200)],
|
||||
point(3, 2.0, 210),
|
||||
);
|
||||
assert!(waiting.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn one_point_reversal_is_skipped_without_counting_an_extra_block() {
|
||||
let sample = select(
|
||||
vec![
|
||||
point(1, 10.0, 100),
|
||||
point(2, 15.0, 200),
|
||||
point(3, 14.0, 210),
|
||||
],
|
||||
point(4, 20.0, 300),
|
||||
);
|
||||
assert_eq!(sample.block_count, 2);
|
||||
assert_eq!(sample.usage.tokens, 200);
|
||||
assert!(!sample.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_token_and_cost_counts_are_local_to_the_selected_interval() {
|
||||
let mut baseline = point(1, 30.0, 1_000);
|
||||
baseline.usage.missing_token_count = 2;
|
||||
baseline.usage.unavailable_cost_count = 3;
|
||||
let mut end = point(3, 40.0, 2_000);
|
||||
end.usage.missing_token_count = 2;
|
||||
end.usage.unavailable_cost_count = 4;
|
||||
let sample = select(vec![baseline], end);
|
||||
assert_eq!(sample.usage.missing_token_count, 0);
|
||||
assert_eq!(sample.usage.unavailable_cost_count, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn counter_regressions_fail_closed_and_future_points_are_not_used() {
|
||||
let invalid = select(vec![point(1, 10.0, 1_000)], point(2, 20.0, 500));
|
||||
assert!(invalid.discontinuous);
|
||||
let future = select(vec![point(10, 1.0, 1)], point(2, 20.0, 500));
|
||||
assert_eq!(future.method, QuotaForecastMethod::Cumulative);
|
||||
assert_eq!(future.observation_count, 0);
|
||||
}
|
||||
@@ -33,16 +33,18 @@ use gateway_admin::{
|
||||
AuthorizationCommit, AuthorizationCommitGuard, AuthorizationCredentialCommit,
|
||||
AuthorizationMutationTarget, AuthorizationStarted, CompleteAuthorization,
|
||||
ConsumeProviderResetCredit, CredentialCommitGuard, CredentialDetails,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
|
||||
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
|
||||
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
|
||||
PrepareCredentialRotation, PreparedAuthorizationCommit,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
|
||||
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
|
||||
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
|
||||
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
|
||||
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
|
||||
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
|
||||
ProviderQuotaRequest, ProviderQuotaWindow, ProviderResetCreditResult,
|
||||
ProviderExport, ProviderExportCredentialInput, ProviderModels,
|
||||
ProviderProfileActivityInsights, ProviderProfileStatistics,
|
||||
ProviderProfileStatisticsSummary, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderQuotaWindow, ProviderResetCreditResult, ProviderSubscription,
|
||||
QuotaLocalUsageAttribution,
|
||||
},
|
||||
quota_forecast_sampling::{QuotaForecastHistory, QuotaForecastUsage},
|
||||
settings::{
|
||||
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
|
||||
RuntimeSettings,
|
||||
@@ -71,10 +73,15 @@ pub(super) struct FakeProviderAdmin {
|
||||
export_inputs: Mutex<Vec<ProviderExportCredentialInput>>,
|
||||
import_account_ids: Mutex<Vec<String>>,
|
||||
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
|
||||
quota_started: tokio::sync::Notify,
|
||||
quota_gate: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
|
||||
quota: Mutex<ProviderQuota>,
|
||||
quota_refresh_account: Mutex<Option<(Arc<FakeAccountStore>, AccountRecord)>>,
|
||||
current_credential_revision: Mutex<Revision>,
|
||||
reset_credit_commands: Mutex<Vec<ConsumeProviderResetCredit>>,
|
||||
profile_result: Mutex<Result<ProviderProfileStatistics, ProviderAdminErrorKind>>,
|
||||
subscription_result: Mutex<Result<Option<ProviderSubscription>, ProviderAdminErrorKind>>,
|
||||
personal_info_barrier: Mutex<Option<Arc<tokio::sync::Barrier>>>,
|
||||
}
|
||||
|
||||
impl FakeProviderAdmin {
|
||||
@@ -89,10 +96,15 @@ impl FakeProviderAdmin {
|
||||
export_inputs: Mutex::new(Vec::new()),
|
||||
import_account_ids: Mutex::new(vec!["acct_prepared".to_owned()]),
|
||||
quota_requests: Mutex::new(Vec::new()),
|
||||
quota_started: tokio::sync::Notify::new(),
|
||||
quota_gate: Mutex::new(None),
|
||||
quota: Mutex::new(empty_quota()),
|
||||
quota_refresh_account: Mutex::new(None),
|
||||
current_credential_revision: Mutex::new(revision(1)),
|
||||
reset_credit_commands: Mutex::new(Vec::new()),
|
||||
profile_result: Mutex::new(Ok(empty_profile_statistics())),
|
||||
subscription_result: Mutex::new(Ok(None)),
|
||||
personal_info_barrier: Mutex::new(None),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -100,6 +112,14 @@ impl FakeProviderAdmin {
|
||||
*self.failure.lock().expect("provider failure") = Some(ProviderAdminError::new(kind));
|
||||
}
|
||||
|
||||
async fn wait_personal_info_queries(&self) {
|
||||
let barrier = self.personal_info_barrier.lock().unwrap().clone();
|
||||
if let Some(barrier) = barrier {
|
||||
barrier.wait().await;
|
||||
barrier.wait().await;
|
||||
}
|
||||
}
|
||||
|
||||
fn fail_next_with_message(&self, kind: ProviderAdminErrorKind, message: &str) {
|
||||
*self.failure.lock().expect("provider failure") =
|
||||
Some(ProviderAdminError::new(kind).with_message(message));
|
||||
@@ -138,6 +158,26 @@ impl FakeProviderAdmin {
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub(super) fn pause_next_quota(&self) -> tokio::sync::oneshot::Sender<()> {
|
||||
let (sender, receiver) = tokio::sync::oneshot::channel();
|
||||
*self.quota_gate.lock().expect("quota gate") = Some(receiver);
|
||||
sender
|
||||
}
|
||||
|
||||
pub(super) async fn wait_for_quota_requests(&self, expected: usize) {
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
loop {
|
||||
let notified = self.quota_started.notified();
|
||||
if self.quota_requests().len() >= expected {
|
||||
return;
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("quota observation started");
|
||||
}
|
||||
|
||||
fn reset_credit_commands(&self) -> Vec<ConsumeProviderResetCredit> {
|
||||
self.reset_credit_commands
|
||||
.lock()
|
||||
@@ -213,6 +253,32 @@ impl ProviderAdmin for FakeProviderAdmin {
|
||||
&self.kind
|
||||
}
|
||||
|
||||
async fn profile_statistics(
|
||||
&self,
|
||||
_: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileStatistics, ProviderAdminError> {
|
||||
self.record("provider.profile_statistics");
|
||||
self.wait_personal_info_queries().await;
|
||||
self.profile_result
|
||||
.lock()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.map_err(ProviderAdminError::new)
|
||||
}
|
||||
|
||||
async fn subscription(
|
||||
&self,
|
||||
_: &ProviderAccountId,
|
||||
) -> Result<Option<ProviderSubscription>, ProviderAdminError> {
|
||||
self.record("provider.subscription");
|
||||
self.wait_personal_info_queries().await;
|
||||
self.subscription_result
|
||||
.lock()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.map_err(ProviderAdminError::new)
|
||||
}
|
||||
|
||||
fn plan_type_display(&self, plan_type: &str) -> String {
|
||||
format!("{} display: {plan_type}", self.kind)
|
||||
}
|
||||
@@ -377,6 +443,11 @@ impl ProviderAdmin for FakeProviderAdmin {
|
||||
.lock()
|
||||
.expect("provider quota requests")
|
||||
.push(request);
|
||||
self.quota_started.notify_one();
|
||||
let gate = self.quota_gate.lock().expect("quota gate").take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("release paused quota");
|
||||
}
|
||||
if let Some(kind) = self
|
||||
.quota_failure
|
||||
.lock()
|
||||
@@ -404,6 +475,21 @@ impl ProviderAdmin for FakeProviderAdmin {
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_forecast_observation(
|
||||
&self,
|
||||
document: &ProviderDocument,
|
||||
_: &ProviderQuotaWindow,
|
||||
) -> Option<gateway_admin::model::quota_forecast_sampling::QuotaForecastObservation> {
|
||||
let fields = document.expose_to_provider().expose_to_provider();
|
||||
Some(
|
||||
gateway_admin::model::quota_forecast_sampling::QuotaForecastObservation {
|
||||
used_percent: fields.get("percent")?.as_f64()?,
|
||||
reset_at: fields.get("reset")?.as_str()?.parse().ok()?,
|
||||
plan_type: fields.get("plan")?.as_str().map(str::to_owned),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
async fn models(
|
||||
&self,
|
||||
_: &ProviderAccountId,
|
||||
@@ -444,6 +530,7 @@ pub(super) struct FakeAccountStore {
|
||||
import_settings: Mutex<Vec<Option<gateway_admin::model::accounts::AccountImportSettings>>>,
|
||||
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
|
||||
quota_window_queries: Mutex<Vec<AccountUsageWindowQuery>>,
|
||||
quota_forecast_history: Mutex<QuotaForecastHistory>,
|
||||
}
|
||||
|
||||
impl FakeAccountStore {
|
||||
@@ -461,6 +548,7 @@ impl FakeAccountStore {
|
||||
import_settings: Mutex::new(Vec::new()),
|
||||
quota_window_usage: Mutex::new(Vec::new()),
|
||||
quota_window_queries: Mutex::new(Vec::new()),
|
||||
quota_forecast_history: Mutex::new(QuotaForecastHistory::default()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -612,21 +700,15 @@ impl AccountStore for FakeAccountStore {
|
||||
.clone())
|
||||
}
|
||||
|
||||
async fn list_credentials(
|
||||
async fn load_quota_forecast_history(
|
||||
&self,
|
||||
provider_kind: &ProviderKind,
|
||||
_: CredentialListQuery,
|
||||
) -> AdminStoreResult<CredentialPage> {
|
||||
self.record("store.list_credentials");
|
||||
let accounts = self.accounts.lock().expect("accounts").clone();
|
||||
Ok(CredentialPage {
|
||||
config_revision: revision(1),
|
||||
items: accounts
|
||||
.into_iter()
|
||||
.filter(|account| &account.provider_kind == provider_kind)
|
||||
.collect(),
|
||||
next_cursor: None,
|
||||
})
|
||||
window: &AccountUsageWindowQuery,
|
||||
) -> AdminStoreResult<QuotaForecastHistory> {
|
||||
self.quota_window_queries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(window.clone());
|
||||
Ok(self.quota_forecast_history.lock().unwrap().clone())
|
||||
}
|
||||
|
||||
async fn credential_details(
|
||||
@@ -1700,6 +1782,172 @@ async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_forecast_reads_raw_snapshot_and_limits_usage_to_observation_time() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let now = Utc::now();
|
||||
let observed = now - TimeDelta::hours(1);
|
||||
let reset = now + TimeDelta::days(1);
|
||||
provider.set_quota(ProviderQuota {
|
||||
observed_at: Some(observed),
|
||||
limit_reached: true,
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "week".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "周额度".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(7 * 86_400),
|
||||
used_percent: Some(20.0),
|
||||
reset_at: Some(reset),
|
||||
limit_reached: true,
|
||||
// Provider 自带的统计不保证与快照同一时间,预测必须重新采样。
|
||||
local_usage: Some(quota_local_usage("acct_test", 999_999)),
|
||||
provider_data: None,
|
||||
}],
|
||||
..empty_quota()
|
||||
});
|
||||
let mut account = account_record("openai");
|
||||
account.created_at = now - TimeDelta::days(60);
|
||||
let store = FakeAccountStore::with_account(account, events());
|
||||
store.set_quota_window_usage(vec![AccountUsageWindowResult {
|
||||
account_id: "acct_test".to_owned(),
|
||||
key: "week".to_owned(),
|
||||
usage: quota_local_usage("acct_test", 1_000),
|
||||
}]);
|
||||
store.quota_forecast_history.lock().unwrap().usage = QuotaForecastUsage {
|
||||
request_count: 10,
|
||||
tokens: 1_000,
|
||||
known_cost_count: 10,
|
||||
..Default::default()
|
||||
};
|
||||
let report = accounts_service(provider.clone(), store.clone())
|
||||
.await
|
||||
.accounts()
|
||||
.quota_forecast(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.expect("quota forecast");
|
||||
assert_eq!(report.account_id, "acct_test");
|
||||
assert_eq!(report.forecasts[0].estimated_tokens, Some(5_000));
|
||||
assert_eq!(
|
||||
report.forecasts[0].source.as_ref().unwrap().observed_at,
|
||||
Some(observed)
|
||||
);
|
||||
let queries = store.quota_window_queries();
|
||||
assert_eq!(queries.len(), 1);
|
||||
assert_eq!(queries[0].range.start, reset - TimeDelta::days(7));
|
||||
assert_eq!(queries[0].range.end, observed);
|
||||
let requests = provider.quota_requests.lock().unwrap();
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert!(!requests[0].refresh);
|
||||
assert!(requests[0].rolling_usage.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_forecast_does_not_query_usage_without_a_current_snapshot() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
let report = accounts_service(provider.clone(), store.clone())
|
||||
.await
|
||||
.accounts()
|
||||
.quota_forecast(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.expect("empty forecast");
|
||||
assert!(
|
||||
report
|
||||
.forecasts
|
||||
.iter()
|
||||
.all(|item| item.estimated_tokens.is_none())
|
||||
);
|
||||
assert!(store.quota_window_queries().is_empty());
|
||||
assert!(!provider.quota_requests.lock().unwrap()[0].refresh);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_forecast_mid_cycle_sampling_accepts_small_reset_jitter_but_not_a_changed_bucket() {
|
||||
use gateway_admin::model::quota_forecast_sampling::QuotaForecastHistoryPoint;
|
||||
let now = Utc::now();
|
||||
let observed = now - TimeDelta::minutes(1);
|
||||
let reset = now + TimeDelta::days(1);
|
||||
let added = now - TimeDelta::hours(5);
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
plan_type: Some("pro".to_owned()),
|
||||
observed_at: Some(observed),
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "week".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "周额度".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(604_800),
|
||||
used_percent: Some(40.0),
|
||||
reset_at: Some(reset),
|
||||
limit_reached: false,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}],
|
||||
..empty_quota()
|
||||
});
|
||||
let mut account = account_record("openai");
|
||||
account.created_at = added;
|
||||
let store = FakeAccountStore::with_account(account, events());
|
||||
let make_point = |hours, percent, tokens, delta| {
|
||||
QuotaForecastHistoryPoint {
|
||||
started_at: now - TimeDelta::hours(hours) - TimeDelta::seconds(10),
|
||||
completed_at: now - TimeDelta::hours(hours),
|
||||
usage: QuotaForecastUsage { request_count: tokens / 100, tokens, ..Default::default() },
|
||||
provider_observation: ProviderDocument::new(OpaqueProviderData::new(serde_json::json!({
|
||||
"percent": percent, "reset": (reset + TimeDelta::seconds(delta)).to_rfc3339(), "plan": "pro",
|
||||
}).as_object().unwrap().clone())),
|
||||
}
|
||||
};
|
||||
*store.quota_forecast_history.lock().unwrap() = QuotaForecastHistory {
|
||||
points: vec![
|
||||
make_point(3, 20.0, 1_000, 1),
|
||||
make_point(2, 30.0, 1_500, -1),
|
||||
],
|
||||
usage: QuotaForecastUsage {
|
||||
request_count: 20,
|
||||
tokens: 2_000,
|
||||
..Default::default()
|
||||
},
|
||||
pending_request_count: 1,
|
||||
};
|
||||
let services = accounts_service(provider, store.clone()).await;
|
||||
let result = services
|
||||
.accounts()
|
||||
.quota_forecast(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.forecasts[0].unavailable_reason.is_none());
|
||||
assert_eq!(result.forecasts[0].estimated_tokens, Some(5_000));
|
||||
assert_eq!(result.forecasts[0].remaining_tokens, Some(3_000));
|
||||
assert_eq!(store.quota_window_queries()[0].range.start, added);
|
||||
store
|
||||
.quota_forecast_history
|
||||
.lock()
|
||||
.unwrap()
|
||||
.points
|
||||
.push(make_point(1, 35.0, 1_700, 3_600));
|
||||
let result = services
|
||||
.accounts()
|
||||
.quota_forecast(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.forecasts[0].estimated_tokens.is_none());
|
||||
assert!(
|
||||
result.forecasts[0]
|
||||
.unavailable_reason
|
||||
.unwrap()
|
||||
.contains("不连续")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
@@ -2223,6 +2471,167 @@ fn rotation_result(command: CredentialRotationCommit) -> CredentialMutationResul
|
||||
}
|
||||
}
|
||||
|
||||
fn empty_profile_statistics() -> ProviderProfileStatistics {
|
||||
ProviderProfileStatistics {
|
||||
display_name: Some("Preview user".to_owned()),
|
||||
username: None,
|
||||
image_url: None,
|
||||
has_stats_error: false,
|
||||
summary: ProviderProfileStatisticsSummary {
|
||||
total_text_tokens: Some(123),
|
||||
peak_tokens: None,
|
||||
longest_task_duration_ms: None,
|
||||
current_streak_days: None,
|
||||
longest_streak_days: None,
|
||||
},
|
||||
daily_usage: None,
|
||||
activity_insights: ProviderProfileActivityInsights {
|
||||
fast_mode_percent: None,
|
||||
reasoning_effort: None,
|
||||
reasoning_effort_percent: None,
|
||||
skills_explored: None,
|
||||
total_skills_used: None,
|
||||
total_threads: None,
|
||||
invocations: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn personal_info_subscription() -> ProviderSubscription {
|
||||
ProviderSubscription {
|
||||
starts_at: None,
|
||||
expires_at: Utc::now(),
|
||||
will_renew: Some(true),
|
||||
billing_period: Some("monthly".to_owned()),
|
||||
billing_currency: Some("USD".to_owned()),
|
||||
observed_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_info_should_query_both_parts_concurrently_once_on_every_request() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let subscription = personal_info_subscription();
|
||||
*provider.subscription_result.lock().unwrap() = Ok(Some(subscription.clone()));
|
||||
*provider.personal_info_barrier.lock().unwrap() = Some(Arc::new(tokio::sync::Barrier::new(2)));
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let services = accounts_service(provider, store).await;
|
||||
let account_id = ProviderAccountId::new("acct_test").unwrap();
|
||||
|
||||
for _ in 0..2 {
|
||||
let info = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(1),
|
||||
services.accounts().personal_info(&account_id),
|
||||
)
|
||||
.await
|
||||
.expect("两项查询应并发执行")
|
||||
.unwrap();
|
||||
assert_eq!(info.profile.unwrap(), empty_profile_statistics());
|
||||
assert_eq!(info.subscription, Some(subscription.clone()));
|
||||
}
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"store.load_account",
|
||||
"provider.profile_statistics",
|
||||
"provider.subscription",
|
||||
"store.load_account",
|
||||
"store.load_account",
|
||||
"provider.profile_statistics",
|
||||
"provider.subscription",
|
||||
"store.load_account",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_info_should_keep_profile_when_subscription_is_missing_or_fails() {
|
||||
for subscription in [Ok(None), Err(ProviderAdminErrorKind::Unavailable)] {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
*provider.subscription_result.lock().unwrap() = subscription;
|
||||
let services = accounts_service(provider, FakeAccountStore::new("openai", events())).await;
|
||||
let info = services
|
||||
.accounts()
|
||||
.personal_info(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(info.profile.unwrap(), empty_profile_statistics());
|
||||
assert_eq!(info.subscription, None);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_info_should_keep_subscription_when_profile_fails() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
*provider.profile_result.lock().unwrap() = Err(ProviderAdminErrorKind::BadGateway);
|
||||
let subscription = personal_info_subscription();
|
||||
*provider.subscription_result.lock().unwrap() = Ok(Some(subscription.clone()));
|
||||
let services = accounts_service(provider, FakeAccountStore::new("openai", events())).await;
|
||||
let info = services
|
||||
.accounts()
|
||||
.personal_info(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
info.profile.unwrap_err().kind(),
|
||||
gateway_admin::model::AdminErrorKind::BadGateway
|
||||
);
|
||||
assert_eq!(info.subscription, Some(subscription));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_info_should_reject_missing_account_before_querying_provider() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
store.set_accounts(Vec::new());
|
||||
let services = accounts_service(provider, store).await;
|
||||
let error = services
|
||||
.accounts()
|
||||
.personal_info(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind(), gateway_admin::model::AdminErrorKind::NotFound);
|
||||
assert_eq!(recorded(&events), ["store.load_account"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_info_should_discard_results_when_account_changes_during_query() {
|
||||
for change in 0..5 {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let barrier = Arc::new(tokio::sync::Barrier::new(3));
|
||||
*provider.personal_info_barrier.lock().unwrap() = Some(barrier.clone());
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
let services = accounts_service(provider, store.clone()).await;
|
||||
let account_id = ProviderAccountId::new("acct_test").unwrap();
|
||||
let mutate = async {
|
||||
barrier.wait().await;
|
||||
let mut accounts = store.accounts.lock().unwrap().clone();
|
||||
match change {
|
||||
0 => accounts[0].credential_revision = revision(99),
|
||||
1 => accounts[0].upstream_account_id = Some("changed-account".to_owned()),
|
||||
2 => accounts[0].upstream_user_id = Some("changed-user".to_owned()),
|
||||
3 => accounts[0].provider_kind = ProviderKind::new("xai").unwrap(),
|
||||
_ => accounts.clear(),
|
||||
}
|
||||
store.set_accounts(accounts);
|
||||
barrier.wait().await;
|
||||
};
|
||||
let (result, ()) = tokio::time::timeout(std::time::Duration::from_secs(1), async {
|
||||
tokio::join!(services.accounts().personal_info(&account_id), mutate)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let expected = if change == 4 {
|
||||
gateway_admin::model::AdminErrorKind::NotFound
|
||||
} else {
|
||||
gateway_admin::model::AdminErrorKind::Conflict
|
||||
};
|
||||
assert_eq!(result.unwrap_err().kind(), expected);
|
||||
}
|
||||
}
|
||||
|
||||
async fn accounts_service(
|
||||
provider: Arc<FakeProviderAdmin>,
|
||||
store: Arc<FakeAccountStore>,
|
||||
|
||||
@@ -1,22 +1,31 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use chrono::Utc;
|
||||
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
|
||||
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
AdminErrorKind, MutationContext, Revision,
|
||||
AdminErrorKind, MutationActor, MutationContext, Revision,
|
||||
client_keys::{
|
||||
ClientKeyCursor, ClientKeyCursorValue, ClientKeyListQuery, ClientKeyPage,
|
||||
ClientKeyPageSize, ClientKeyRecord, ClientKeySecret, ClientKeySort, ClientKeySortField,
|
||||
DeleteClientKey, NewClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey,
|
||||
CreateClientKey, DeleteClientKey, NewClientKey, SetClientKeyEnabled, SortDirection,
|
||||
UpdateClientKey,
|
||||
},
|
||||
},
|
||||
ports::store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult, ClientKeyStore},
|
||||
};
|
||||
|
||||
struct UnusedClientKeyStore;
|
||||
#[derive(Default)]
|
||||
struct TestClientKeyStore {
|
||||
plaintexts: Mutex<Vec<String>>,
|
||||
create_error: Option<AdminStoreErrorKind>,
|
||||
update_error: Option<AdminStoreErrorKind>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyStore for UnusedClientKeyStore {
|
||||
impl ClientKeyStore for TestClientKeyStore {
|
||||
async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage> {
|
||||
assert_eq!(query.page_size.get(), u16::MAX);
|
||||
Ok(ClientKeyPage {
|
||||
@@ -36,10 +45,32 @@ impl ClientKeyStore for UnusedClientKeyStore {
|
||||
|
||||
async fn create_client_key(
|
||||
&self,
|
||||
_: NewClientKey,
|
||||
key: NewClientKey,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<(Revision, ClientKeyRecord)> {
|
||||
Err(unused())
|
||||
if let Some(kind) = self.create_error {
|
||||
return Err(AdminStoreError::new(
|
||||
kind,
|
||||
"client key",
|
||||
"test create failure",
|
||||
));
|
||||
}
|
||||
self.plaintexts.lock().unwrap().push(key.plaintext);
|
||||
let record = ClientKeyRecord {
|
||||
id: key.id,
|
||||
name: key.name,
|
||||
label: key.label,
|
||||
groups: Vec::new(),
|
||||
provider_kinds: Vec::new(),
|
||||
prefix: String::new(),
|
||||
enabled: true,
|
||||
limits: key.limits,
|
||||
budget: Default::default(),
|
||||
last_used_at: None,
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
};
|
||||
Ok((Revision::new(2).unwrap(), record))
|
||||
}
|
||||
|
||||
async fn update_client_key(
|
||||
@@ -47,6 +78,13 @@ impl ClientKeyStore for UnusedClientKeyStore {
|
||||
_: UpdateClientKey,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<(Revision, ClientKeyRecord)> {
|
||||
if let Some(kind) = self.update_error {
|
||||
return Err(AdminStoreError::new(
|
||||
kind,
|
||||
"client key",
|
||||
"test update failure",
|
||||
));
|
||||
}
|
||||
Err(unused())
|
||||
}
|
||||
|
||||
@@ -70,7 +108,7 @@ impl ClientKeyStore for UnusedClientKeyStore {
|
||||
#[tokio::test]
|
||||
async fn client_key_cursor_should_reject_value_that_does_not_match_sort() {
|
||||
let services = super::AdminHarness::new()
|
||||
.client_keys(std::sync::Arc::new(UnusedClientKeyStore))
|
||||
.client_keys(Arc::new(TestClientKeyStore::default()))
|
||||
.build()
|
||||
.await;
|
||||
let sort = ClientKeySort {
|
||||
@@ -98,7 +136,7 @@ async fn client_key_cursor_should_reject_value_that_does_not_match_sort() {
|
||||
#[tokio::test]
|
||||
async fn client_key_list_should_forward_the_full_nonzero_u16_page_size() {
|
||||
let services = super::AdminHarness::new()
|
||||
.client_keys(std::sync::Arc::new(UnusedClientKeyStore))
|
||||
.client_keys(Arc::new(TestClientKeyStore::default()))
|
||||
.build()
|
||||
.await;
|
||||
let page = services
|
||||
@@ -118,6 +156,105 @@ async fn client_key_list_should_forward_the_full_nonzero_u16_page_size() {
|
||||
assert_eq!(page.total, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_preserves_migrated_keys_and_keeps_default_generation() {
|
||||
let store = Arc::new(TestClientKeyStore::default());
|
||||
let services = super::AdminHarness::new()
|
||||
.client_keys(store.clone())
|
||||
.build()
|
||||
.await;
|
||||
for value in [Some("q".to_owned()), Some("legacy+/=:!".repeat(1024)), None] {
|
||||
let created = services
|
||||
.client_keys()
|
||||
.create(&mutation_context(), create_command(value.as_deref()))
|
||||
.await
|
||||
.unwrap();
|
||||
let plaintext = created.secret.expose_for_response();
|
||||
assert_eq!(store.plaintexts.lock().unwrap().last().unwrap(), plaintext);
|
||||
if let Some(value) = value {
|
||||
assert_eq!(plaintext, value);
|
||||
} else {
|
||||
assert!(plaintext.starts_with("sk_"));
|
||||
assert_eq!(plaintext.len(), 46);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_keys_return_actionable_conflicts_without_disclosing_the_key() {
|
||||
let services = super::AdminHarness::new()
|
||||
.client_keys(Arc::new(TestClientKeyStore {
|
||||
create_error: Some(AdminStoreErrorKind::Conflict),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let key = "legacy-duplicate-must-stay-private";
|
||||
let error = services
|
||||
.client_keys()
|
||||
.create(&mutation_context(), create_command(Some(key)))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind(), AdminErrorKind::Conflict);
|
||||
assert_eq!(error.message(), "API Key 已存在,请使用其他密钥");
|
||||
assert!(!format!("{error:?}").contains(key));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_names_report_the_same_actionable_conflict_on_create_and_update() {
|
||||
let services = super::AdminHarness::new()
|
||||
.client_keys(Arc::new(TestClientKeyStore {
|
||||
create_error: Some(AdminStoreErrorKind::DuplicateName),
|
||||
update_error: Some(AdminStoreErrorKind::DuplicateName),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let created = services
|
||||
.client_keys()
|
||||
.create(&mutation_context(), create_command(None))
|
||||
.await
|
||||
.unwrap_err();
|
||||
let updated = services
|
||||
.client_keys()
|
||||
.update(
|
||||
&mutation_context(),
|
||||
UpdateClientKey {
|
||||
id: ClientApiKeyId::new("key_existing").unwrap(),
|
||||
name: "Migration".to_owned(),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits::unlimited(),
|
||||
daily_limit_usd: None,
|
||||
weekly_limit_usd: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
for error in [created, updated] {
|
||||
assert_eq!(error.kind(), AdminErrorKind::Conflict);
|
||||
assert_eq!(error.message(), "名称已存在");
|
||||
}
|
||||
}
|
||||
|
||||
fn create_command(key: Option<&str>) -> CreateClientKey {
|
||||
CreateClientKey {
|
||||
custom_key: key.map(|value| PlaintextClientApiKey::new(value).unwrap()),
|
||||
name: "Migration".to_owned(),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits::unlimited(),
|
||||
budget: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn mutation_context() -> MutationContext {
|
||||
MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "custom-key-test".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn unused() -> AdminStoreError {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
|
||||
@@ -45,12 +45,11 @@ use gateway_admin::{
|
||||
},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
|
||||
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
|
||||
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
|
||||
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
|
||||
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput,
|
||||
ProviderModels, ProviderQuota,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
|
||||
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
|
||||
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
|
||||
PreparedCredentialImport, PreparedCredentialRotation, ProviderExport,
|
||||
ProviderExportCredentialInput, ProviderModels, ProviderQuota,
|
||||
},
|
||||
settings::{AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RuntimeSettings},
|
||||
system::{SystemOperationAccepted, SystemUpdateDetail, SystemUpdateStatus, SystemVersion},
|
||||
@@ -372,12 +371,11 @@ impl AccountStore for UnavailableStore {
|
||||
Err(unavailable("account quota window usage"))
|
||||
}
|
||||
|
||||
async fn list_credentials(
|
||||
async fn load_quota_forecast_history(
|
||||
&self,
|
||||
_: &ProviderKind,
|
||||
_: CredentialListQuery,
|
||||
) -> AdminStoreResult<CredentialPage> {
|
||||
Err(unavailable("credentials"))
|
||||
_: &AccountUsageWindowQuery,
|
||||
) -> AdminStoreResult<gateway_admin::model::quota_forecast_sampling::QuotaForecastHistory> {
|
||||
Err(unavailable("quota forecast history"))
|
||||
}
|
||||
|
||||
async fn credential_details(
|
||||
|
||||
@@ -58,7 +58,7 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
|
||||
.await
|
||||
.expect("import credential");
|
||||
|
||||
await_quota_requests(&provider, 1).await;
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
@@ -142,7 +142,7 @@ async fn openai_import_should_refresh_quota_for_every_imported_account() {
|
||||
.await
|
||||
.expect("import credentials");
|
||||
|
||||
await_quota_requests(&provider, 2).await;
|
||||
provider.wait_for_quota_requests(2).await;
|
||||
|
||||
assert_eq!(
|
||||
result.credential_ids,
|
||||
@@ -187,7 +187,7 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
|
||||
.await
|
||||
.expect("committed import remains successful");
|
||||
|
||||
await_quota_requests(&provider, 1).await;
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
|
||||
assert_eq!(
|
||||
result.credential_ids,
|
||||
@@ -224,7 +224,7 @@ async fn openai_authorization_create_should_observe_initial_quota() {
|
||||
.await
|
||||
.expect("complete authorization");
|
||||
|
||||
await_quota_requests(&provider, 1).await;
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
@@ -280,7 +280,7 @@ async fn openai_authorization_create_should_remain_successful_when_initial_quota
|
||||
.await
|
||||
.expect("quota observation must not fail authorization");
|
||||
|
||||
await_quota_requests(&provider, 1).await;
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
|
||||
assert_eq!(result.account_id.as_str(), "acct_prepared");
|
||||
assert_eq!(provider.quota_requests().len(), 1);
|
||||
@@ -322,7 +322,7 @@ async fn openai_authorization_store_failure_should_release_claim_for_retry() {
|
||||
.await
|
||||
.expect("released claim must permit retry");
|
||||
|
||||
await_quota_requests(&provider, 1).await;
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
@@ -406,7 +406,7 @@ async fn openai_reauthorization_should_commit_after_credential_revision_advances
|
||||
.await
|
||||
.expect("complete reauthorization");
|
||||
|
||||
await_quota_requests(&provider, 1).await;
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
|
||||
assert_eq!(result.credential_revision, Some(revision(3)));
|
||||
|
||||
@@ -434,16 +434,6 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn await_quota_requests(provider: &FakeProviderAdmin, expected: usize) {
|
||||
for _ in 0..100 {
|
||||
if provider.quota_requests().len() == expected {
|
||||
return;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
assert_eq!(provider.quota_requests().len(), expected);
|
||||
}
|
||||
|
||||
fn deletion(account_id: &str) -> CredentialDeletion {
|
||||
CredentialDeletion {
|
||||
context: context("request-openai"),
|
||||
|
||||
@@ -4,7 +4,9 @@ use gateway_core::account::ProviderAccountId;
|
||||
|
||||
use gateway_admin::{
|
||||
AdminServices,
|
||||
model::provider_credentials::{ImportCredentials, ProviderQuotaRequest},
|
||||
model::provider_credentials::{
|
||||
CompleteAuthorization, ImportCredentials, ProviderQuotaRequest, StartAuthorization,
|
||||
},
|
||||
ports::provider::ProviderAdminErrorKind,
|
||||
};
|
||||
|
||||
@@ -28,11 +30,13 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
|
||||
.await
|
||||
.expect("import xAI credential");
|
||||
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"provider.prepare_import",
|
||||
"store.commit_import",
|
||||
"provider.account_facts_changed",
|
||||
"provider.quota"
|
||||
]
|
||||
);
|
||||
@@ -70,6 +74,7 @@ async fn xai_import_should_refresh_quota_for_every_imported_account() {
|
||||
.await
|
||||
.expect("import credentials");
|
||||
|
||||
provider.wait_for_quota_requests(2).await;
|
||||
assert_eq!(
|
||||
provider.quota_requests(),
|
||||
[
|
||||
@@ -106,10 +111,120 @@ async fn xai_import_should_remain_successful_when_quota_refresh_fails() {
|
||||
.await
|
||||
.expect("committed import remains successful");
|
||||
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
assert_eq!(result.credential_ids.len(), 1);
|
||||
assert_eq!(provider.quota_requests().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn xai_authorization_should_observe_quota_without_affecting_commit() {
|
||||
for quota_failure in [None, Some(ProviderAdminErrorKind::Unavailable)] {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("xai", events.clone());
|
||||
if let Some(kind) = quota_failure {
|
||||
provider.fail_next_quota(kind);
|
||||
}
|
||||
let store = FakeAccountStore::new("xai", events.clone());
|
||||
let services = service(provider.clone(), store.clone()).await;
|
||||
services
|
||||
.xai()
|
||||
.start_authorization(StartAuthorization {
|
||||
outbound_proxy: None,
|
||||
context: context("oauth-xai-create"),
|
||||
name: "new xAI credential".to_owned(),
|
||||
reauthorization: None,
|
||||
})
|
||||
.await
|
||||
.expect("start authorization");
|
||||
let result = services
|
||||
.xai()
|
||||
.complete_authorization(CompleteAuthorization {
|
||||
settings: None,
|
||||
context: context("oauth-xai-create"),
|
||||
flow_id: "flow-test".to_owned(),
|
||||
callback_url: "http://localhost/callback?code=test&state=test".to_owned(),
|
||||
})
|
||||
.await
|
||||
.expect("quota observation must not fail authorization");
|
||||
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"provider.start_authorization",
|
||||
"provider.complete_authorization",
|
||||
"store.commit_authorization",
|
||||
"provider.account_facts_changed",
|
||||
"provider.quota",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
provider.quota_requests(),
|
||||
[ProviderQuotaRequest {
|
||||
account_id: result.account_id,
|
||||
refresh: true,
|
||||
rolling_usage: None,
|
||||
}]
|
||||
);
|
||||
assert_eq!(store.audit_requests(), ["oauth-xai-create"]);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_import_should_return_while_initial_quota_is_pending() {
|
||||
for kind in ["openai", "xai"] {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new(kind, events.clone());
|
||||
let quota_release = provider.pause_next_quota();
|
||||
let store = FakeAccountStore::new(kind, events);
|
||||
let services = service(provider.clone(), store).await;
|
||||
let command = ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-slow-quota"),
|
||||
document: document(),
|
||||
};
|
||||
let result = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
match kind {
|
||||
"openai" => services.openai().import_document(command).await,
|
||||
_ => services.xai().import_document(command).await,
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("committed import must not wait for quota")
|
||||
.expect("import credentials");
|
||||
provider.wait_for_quota_requests(1).await;
|
||||
assert_eq!(result.credential_ids.len(), 1);
|
||||
quota_release.send(()).expect("finish quota observation");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn xai_import_store_failure_should_not_observe_quota_or_invalidate_provider_facts() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("xai", events.clone());
|
||||
let store = FakeAccountStore::new("xai", events.clone());
|
||||
store.fail_next_commit();
|
||||
let services = service(provider.clone(), store).await;
|
||||
|
||||
services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-xai-store-failure"),
|
||||
document: document(),
|
||||
})
|
||||
.await
|
||||
.expect_err("store commit fails");
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
["provider.prepare_import", "store.commit_import"]
|
||||
);
|
||||
assert!(provider.quota_requests().is_empty());
|
||||
}
|
||||
|
||||
async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>) -> AdminServices {
|
||||
super::AdminHarness::new()
|
||||
.provider(provider)
|
||||
|
||||
@@ -23,8 +23,12 @@ where
|
||||
)
|
||||
.route("/api/admin/accounts/quota", get(account_quota::<S>))
|
||||
.route(
|
||||
"/api/admin/accounts/profile-statistics",
|
||||
get(account_profile_statistics::<S>),
|
||||
"/api/admin/accounts/personal-info",
|
||||
get(account_personal_info::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/quota-forecast",
|
||||
get(account_quota_forecast::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/profile-avatar",
|
||||
@@ -364,7 +368,7 @@ where
|
||||
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
|
||||
}
|
||||
|
||||
async fn account_profile_statistics<S>(
|
||||
async fn account_quota_forecast<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<AccountIdQuery>,
|
||||
@@ -376,12 +380,33 @@ where
|
||||
let result = state
|
||||
.admin_services()
|
||||
.accounts()
|
||||
.profile_statistics(&account_id)
|
||||
.quota_forecast(&account_id)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountProfileStatisticsData::from(result)),
|
||||
AdminEnvelope::ok(AccountQuotaForecastData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn account_personal_info<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<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()
|
||||
.personal_info(&account_id)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountPersonalInfoData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
@@ -25,14 +25,15 @@ use gateway_admin::model::{
|
||||
AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount,
|
||||
},
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
AccountUsagePeriod, AuthorizationStarted, CompleteAuthorization,
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountPersonalInfo,
|
||||
AccountRefreshResult, AccountUsagePeriod, AuthorizationStarted, CompleteAuthorization,
|
||||
ConsumeProviderResetCredit, CredentialDeletion, CredentialDeletionResult,
|
||||
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
|
||||
ProviderDocument, ProviderModels, ProviderProfileActivityInsights, ProviderProfileAvatar,
|
||||
ProviderProfileDailyUsage, ProviderProfileInvocation, ProviderProfileStatistics,
|
||||
ProviderProfileStatisticsSummary, ProviderQuota, ProviderQuotaWindow, ProviderResetCredit,
|
||||
ProviderResetCreditResult, ProviderResetCredits, RotateCredential, StartAuthorization,
|
||||
ProviderResetCreditResult, ProviderResetCredits, ProviderSubscription, RotateCredential,
|
||||
StartAuthorization,
|
||||
},
|
||||
};
|
||||
use gateway_core::{
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! Admin 领域结果到安全 HTTP wire 的展示投影。
|
||||
|
||||
use super::*;
|
||||
use gateway_admin::model::quota_forecast::{AccountQuotaForecast, AccountQuotaForecastReport};
|
||||
|
||||
pub(super) fn account_page_data(
|
||||
result: AccountDirectoryPage,
|
||||
@@ -122,6 +123,64 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> Ac
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AccountQuotaForecastReport> for AccountQuotaForecastData {
|
||||
fn from(report: AccountQuotaForecastReport) -> Self {
|
||||
Self {
|
||||
account_id: report.account_id,
|
||||
generated_at: china_rfc3339(&report.generated_at),
|
||||
forecasts: report
|
||||
.forecasts
|
||||
.into_iter()
|
||||
.map(quota_forecast_view)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_forecast_view(forecast: AccountQuotaForecast) -> AccountQuotaForecastView {
|
||||
AccountQuotaForecastView {
|
||||
period: match forecast.period {
|
||||
AccountUsagePeriod::Weekly => "weekly",
|
||||
AccountUsagePeriod::Monthly => "monthly",
|
||||
},
|
||||
target_days: forecast.target_seconds as f64 / 86_400.0,
|
||||
extrapolated: forecast.extrapolated,
|
||||
source: forecast.source.map(|source| QuotaForecastSourceView {
|
||||
label: source.label,
|
||||
used_percent: source.used_percent,
|
||||
used_percent_display: source
|
||||
.used_percent
|
||||
.map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")),
|
||||
observed_at: source.observed_at.map(|value| china_rfc3339(&value)),
|
||||
observed_at_display: source
|
||||
.observed_at
|
||||
.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)),
|
||||
reset_at: china_rfc3339(&source.reset_at),
|
||||
tokens_display: display_optional_tokens(source.tokens),
|
||||
usd_display: forecast_usd_display(source.usd),
|
||||
}),
|
||||
unavailable_reason: forecast.unavailable_reason,
|
||||
low_sample: forecast.low_sample,
|
||||
incomplete_cost: forecast.incomplete_cost,
|
||||
incomplete_tokens: forecast.incomplete_tokens,
|
||||
estimated_tokens: forecast.estimated_tokens,
|
||||
estimated_tokens_display: display_optional_tokens(forecast.estimated_tokens),
|
||||
estimated_usd: forecast.estimated_usd,
|
||||
estimated_usd_display: forecast_usd_display(forecast.estimated_usd),
|
||||
remaining_tokens: forecast.remaining_tokens,
|
||||
remaining_tokens_display: display_optional_tokens(forecast.remaining_tokens),
|
||||
remaining_usd: forecast.remaining_usd,
|
||||
remaining_usd_display: forecast_usd_display(forecast.remaining_usd),
|
||||
}
|
||||
}
|
||||
|
||||
fn forecast_usd_display(value: Option<f64>) -> String {
|
||||
value.map_or_else(
|
||||
|| "—".to_owned(),
|
||||
|value| format_decimal_currency(&format!("{value:.10}"), "USD"),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn account_quota_view(
|
||||
mut quota: ProviderQuota,
|
||||
rate_limited_until: Option<DateTime<Utc>>,
|
||||
|
||||
@@ -272,6 +272,49 @@ pub struct AccountView {
|
||||
pub usage: AccountUsageView,
|
||||
}
|
||||
|
||||
/// 容量估算仅供管理端展示;金额不是订阅账单或可消费余额。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountQuotaForecastData {
|
||||
pub account_id: String,
|
||||
pub generated_at: String,
|
||||
pub forecasts: Vec<AccountQuotaForecastView>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountQuotaForecastView {
|
||||
pub period: &'static str,
|
||||
pub target_days: f64,
|
||||
pub extrapolated: bool,
|
||||
pub source: Option<QuotaForecastSourceView>,
|
||||
pub unavailable_reason: Option<&'static str>,
|
||||
pub low_sample: bool,
|
||||
pub incomplete_cost: bool,
|
||||
pub incomplete_tokens: bool,
|
||||
pub estimated_tokens: Option<u64>,
|
||||
pub estimated_tokens_display: String,
|
||||
pub estimated_usd: Option<f64>,
|
||||
pub estimated_usd_display: String,
|
||||
pub remaining_tokens: Option<u64>,
|
||||
pub remaining_tokens_display: String,
|
||||
pub remaining_usd: Option<f64>,
|
||||
pub remaining_usd_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct QuotaForecastSourceView {
|
||||
pub label: String,
|
||||
pub used_percent: Option<f64>,
|
||||
pub used_percent_display: String,
|
||||
pub observed_at: Option<String>,
|
||||
pub observed_at_display: String,
|
||||
pub reset_at: String,
|
||||
pub tokens_display: String,
|
||||
pub usd_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountGroupRefView {
|
||||
@@ -610,6 +653,54 @@ pub struct AccountQuotaData {
|
||||
pub account: AccountView,
|
||||
}
|
||||
|
||||
/// 个人资料、累计统计与订阅的统一响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountPersonalInfoData {
|
||||
pub profile: Option<AccountProfileStatisticsData>,
|
||||
pub profile_error: Option<String>,
|
||||
pub subscription: Option<AccountSubscriptionData>,
|
||||
}
|
||||
|
||||
impl From<AccountPersonalInfo> for AccountPersonalInfoData {
|
||||
fn from(info: AccountPersonalInfo) -> Self {
|
||||
let (profile, profile_error) = match info.profile {
|
||||
Ok(profile) => (Some(AccountProfileStatisticsData::from(profile)), None),
|
||||
Err(error) => (None, Some(error.message().to_owned())),
|
||||
};
|
||||
Self {
|
||||
profile,
|
||||
profile_error,
|
||||
subscription: info.subscription.map(AccountSubscriptionData::from),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 按需读取的订阅安全字段,不暴露原始上游响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountSubscriptionData {
|
||||
pub starts_at: Option<String>,
|
||||
pub expires_at: String,
|
||||
pub will_renew: Option<bool>,
|
||||
pub billing_period: Option<String>,
|
||||
pub billing_currency: Option<String>,
|
||||
pub observed_at: String,
|
||||
}
|
||||
|
||||
impl From<ProviderSubscription> for AccountSubscriptionData {
|
||||
fn from(subscription: ProviderSubscription) -> Self {
|
||||
Self {
|
||||
starts_at: subscription.starts_at.map(|value| value.to_rfc3339()),
|
||||
expires_at: subscription.expires_at.to_rfc3339(),
|
||||
will_renew: subscription.will_renew,
|
||||
billing_period: subscription.billing_period,
|
||||
billing_currency: subscription.billing_currency,
|
||||
observed_at: subscription.observed_at.to_rfc3339(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 官方个人资料统计响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
|
||||
@@ -5,7 +5,11 @@ use std::fmt;
|
||||
use axum::{
|
||||
Router,
|
||||
extract::{FromRequestParts, State},
|
||||
http::{HeaderMap, HeaderValue, StatusCode, header::SET_COOKIE, request::Parts},
|
||||
http::{
|
||||
HeaderMap, HeaderValue, StatusCode,
|
||||
header::{ORIGIN, SET_COOKIE},
|
||||
request::Parts,
|
||||
},
|
||||
response::{IntoResponse, Response},
|
||||
routing::{get, post},
|
||||
};
|
||||
@@ -15,18 +19,37 @@ use gateway_admin::{
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tower_http::request_id::RequestId;
|
||||
use url::Url;
|
||||
|
||||
use super::{AdminEnvelope, AdminError, AdminJson, AdminResponse};
|
||||
|
||||
const REQUEST_ID_HEADER: &str = "x-request-id";
|
||||
const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session";
|
||||
const ADMIN_SESSION_COOKIE_ATTRS: &str = "Path=/; Secure; HttpOnly; SameSite=Lax";
|
||||
|
||||
/// 所有管理 HTTP 模块从 state 消费同一个认证用例端口。
|
||||
pub trait AdminSessionState {
|
||||
fn admin_services(&self) -> &AdminServices;
|
||||
}
|
||||
|
||||
fn admin_session_cookie_attrs(headers: &HeaderMap) -> &'static str {
|
||||
// 同源管理端的 POST 由浏览器携带 Origin,HTTPS 反代回源 HTTP 不改变它。
|
||||
// 只对明确的 HTTP 来源放宽;缺失、opaque 或非法来源继续使用 Secure。
|
||||
let http_origin = headers.get_all(ORIGIN).iter().count() == 1
|
||||
&& headers
|
||||
.get(ORIGIN)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.is_some_and(|origin| {
|
||||
Url::parse(origin).is_ok_and(|url| {
|
||||
url.scheme() == "http" && url.origin().ascii_serialization() == origin
|
||||
})
|
||||
});
|
||||
if http_origin {
|
||||
"Path=/; HttpOnly; SameSite=Lax"
|
||||
} else {
|
||||
"Path=/; Secure; HttpOnly; SameSite=Lax"
|
||||
}
|
||||
}
|
||||
|
||||
/// 已通过管理员会话或部署级管理 API Key 鉴权的请求。
|
||||
pub struct AdminAuth {
|
||||
context: AdminRequestContext,
|
||||
@@ -203,6 +226,7 @@ where
|
||||
|
||||
async fn login<S>(
|
||||
State(state): State<S>,
|
||||
headers: HeaderMap,
|
||||
AdminJson(payload): AdminJson<AdminLoginRequest>,
|
||||
) -> Result<Response, AdminError>
|
||||
where
|
||||
@@ -221,8 +245,9 @@ where
|
||||
)
|
||||
.into_response();
|
||||
let cookie = format!(
|
||||
"{ADMIN_SESSION_COOKIE}={}; {ADMIN_SESSION_COOKIE_ATTRS}",
|
||||
session.session_id
|
||||
"{ADMIN_SESSION_COOKIE}={}; {}",
|
||||
session.session_id,
|
||||
admin_session_cookie_attrs(&headers)
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
SET_COOKIE,
|
||||
@@ -260,7 +285,10 @@ where
|
||||
let mut response =
|
||||
AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(AdminLogoutData::new()))
|
||||
.into_response();
|
||||
let cookie = format!("{ADMIN_SESSION_COOKIE}=; {ADMIN_SESSION_COOKIE_ATTRS}; Max-Age=0");
|
||||
let cookie = format!(
|
||||
"{ADMIN_SESSION_COOKIE}=; {}; Max-Age=0",
|
||||
admin_session_cookie_attrs(&headers)
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
SET_COOKIE,
|
||||
HeaderValue::from_str(&cookie).map_err(|_| AdminError::internal())?,
|
||||
|
||||
@@ -13,7 +13,7 @@ use gateway_admin::model::client_keys::{
|
||||
use gateway_core::{
|
||||
engine::budget::ClientBudgetLimits,
|
||||
metering::Decimal,
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
|
||||
routing::AccountGroupId,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -70,9 +70,7 @@ impl ListClientKeysQuery {
|
||||
}
|
||||
let search = self.search.map(|search| search.trim().to_owned());
|
||||
if search.as_deref().is_some_and(|search| {
|
||||
search.len() > MAX_SEARCH_BYTES
|
||||
|| search.chars().any(char::is_control)
|
||||
|| contains_client_key_material(search)
|
||||
search.len() > MAX_SEARCH_BYTES || search.chars().any(char::is_control)
|
||||
}) {
|
||||
return Err(WireValidationError::new("search"));
|
||||
}
|
||||
@@ -143,9 +141,10 @@ impl ClientKeySort {
|
||||
}
|
||||
|
||||
/// 创建 Client Key 请求。
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize)]
|
||||
#[derive(Clone, PartialEq, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct CreateClientKeyRequest {
|
||||
custom_key: Option<String>,
|
||||
name: String,
|
||||
label: Option<String>,
|
||||
group_ids: Vec<String>,
|
||||
@@ -155,6 +154,16 @@ pub struct CreateClientKeyRequest {
|
||||
weekly_limit_usd: Option<String>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for CreateClientKeyRequest {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CreateClientKeyRequest")
|
||||
.field("name", &self.name)
|
||||
.field("custom_key", &"[REDACTED]")
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl CreateClientKeyRequest {
|
||||
/// 校验 wire 边界并直接构造管理用例命令。
|
||||
pub fn into_command(self) -> Result<CreateClientKey, WireValidationError> {
|
||||
@@ -164,6 +173,12 @@ impl CreateClientKeyRequest {
|
||||
validate_limit(self.max_concurrency, "maxConcurrency")?;
|
||||
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
|
||||
Ok(CreateClientKey {
|
||||
custom_key: self
|
||||
.custom_key
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(PlaintextClientApiKey::new)
|
||||
.transpose()
|
||||
.map_err(|_| WireValidationError::new("customKey"))?,
|
||||
name: self.name,
|
||||
label: self.label,
|
||||
group_ids,
|
||||
@@ -661,15 +676,6 @@ fn validate_optional_text(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn contains_client_key_material(value: &str) -> bool {
|
||||
value.as_bytes().windows(46).any(|window| {
|
||||
&window[..3] == b"sk_"
|
||||
&& window[3..]
|
||||
.iter()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
|
||||
})
|
||||
}
|
||||
|
||||
/// 构造固定 GET/POST 且 ID 仅位于 query/body 的 Client API Key 路由。
|
||||
pub fn router<S>() -> Router<S>
|
||||
where
|
||||
|
||||
@@ -331,7 +331,7 @@ impl AdminError {
|
||||
pub(crate) fn map_admin_service_error(error: gateway_admin::model::AdminError) -> AdminError {
|
||||
use gateway_admin::model::AdminErrorKind;
|
||||
|
||||
match error.kind() {
|
||||
let mut response = match error.kind() {
|
||||
AdminErrorKind::Invalid => AdminError::bad_request(error.message()),
|
||||
AdminErrorKind::Unauthorized => AdminError::admin_session_required(),
|
||||
AdminErrorKind::NotFound => AdminError::not_found(error.message()),
|
||||
@@ -341,7 +341,19 @@ pub(crate) fn map_admin_service_error(error: gateway_admin::model::AdminError) -
|
||||
AdminErrorKind::UpstreamResultUnknown => AdminError::upstream_result_unknown(),
|
||||
AdminErrorKind::Unavailable => AdminError::service_unavailable(),
|
||||
AdminErrorKind::Internal => AdminError::internal(),
|
||||
};
|
||||
// Admin owner 已负责公开文案,不能在 HTTP 边界再次丢掉安全的业务原因。
|
||||
// 认证与未知内部异常仍使用固定提示,不放行底层诊断或任意 500 消息。
|
||||
if matches!(
|
||||
error.kind(),
|
||||
AdminErrorKind::BadGateway
|
||||
| AdminErrorKind::UpstreamResultUnknown
|
||||
| AdminErrorKind::Unavailable
|
||||
) && !error.message().trim().is_empty()
|
||||
{
|
||||
response.body.message = error.message().to_owned();
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
impl IntoResponse for AdminError {
|
||||
|
||||
@@ -86,7 +86,7 @@ impl ClientApiKeyAuthError {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Header 缺失、Bearer 语法错误或不是网关 Client Key 前缀时返回稳定错误。
|
||||
/// Header 缺失、Bearer 语法错误或 Key 不能作为 HTTP Bearer 值时返回稳定错误。
|
||||
pub fn bearer_client_api_key(headers: &HeaderMap) -> Result<&str, ClientApiKeyAuthError> {
|
||||
let raw = headers
|
||||
.get(AUTHORIZATION)
|
||||
@@ -100,7 +100,7 @@ pub fn bearer_client_api_key(headers: &HeaderMap) -> Result<&str, ClientApiKeyAu
|
||||
if token.is_empty() {
|
||||
return Err(ClientApiKeyAuthError::MalformedAuthorization);
|
||||
}
|
||||
if !token.starts_with("sk_") {
|
||||
if gateway_core::policy::PlaintextClientApiKey::validate(token).is_err() {
|
||||
return Err(ClientApiKeyAuthError::InvalidKeyFormat);
|
||||
}
|
||||
Ok(token)
|
||||
|
||||
@@ -5,7 +5,7 @@ use axum::{
|
||||
http::{HeaderName, HeaderValue, StatusCode, header::CONTENT_TYPE},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use gateway_core::engine::execution::StartedExecution;
|
||||
use gateway_core::engine::execution::{ExecutionSession, StartedExecution};
|
||||
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
|
||||
|
||||
use super::{
|
||||
@@ -15,7 +15,12 @@ use super::{
|
||||
|
||||
/// 收集一个 Provider 原生 JSON 响应,并在 Core commit 后按原始 bytes 交付。
|
||||
pub(super) async fn collect_raw_json_response(started: StartedExecution) -> Response {
|
||||
let mut execution = PendingExecution::new(started.session);
|
||||
let response = collect_raw_json_session(started.session).await;
|
||||
super::with_model_request_id(response, &started.request_id)
|
||||
}
|
||||
|
||||
async fn collect_raw_json_session(session: Box<dyn ExecutionSession>) -> Response {
|
||||
let mut execution = PendingExecution::new(session);
|
||||
let Some(session) = execution.session_mut() else {
|
||||
return invalid_upstream_response();
|
||||
};
|
||||
|
||||
@@ -190,7 +190,15 @@ pub fn engine_error_response(error: &EngineError) -> Response {
|
||||
}
|
||||
return response;
|
||||
}
|
||||
gateway_error_response(&gateway_error_from_engine(error))
|
||||
let mut response = gateway_error_response(&gateway_error_from_engine(error));
|
||||
// 正文未完整取得时不能透传残缺响应,但已确认的上游关联 ID 仍应交付。
|
||||
if let EngineError::Provider(error) = error
|
||||
&& let Some(request_id) = error.upstream_request_id()
|
||||
&& let Ok(value) = HeaderValue::from_str(request_id.as_str())
|
||||
{
|
||||
response.headers_mut().insert("x-request-id", value);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
/// Gateway 错误稳定映射,供 HTTP、SSE 和 WebSocket 共用。
|
||||
|
||||
@@ -9,3 +9,28 @@ pub mod responses;
|
||||
pub mod router;
|
||||
pub mod search;
|
||||
pub(crate) mod service;
|
||||
|
||||
/// 客户端可读的 ID 必须能按现有请求记录检索;不以无关入口 ID 补位。
|
||||
fn with_model_request_id(
|
||||
mut response: axum::response::Response,
|
||||
request_id: &gateway_core::engine::ModelRequestId,
|
||||
) -> axum::response::Response {
|
||||
use axum::http::HeaderValue;
|
||||
|
||||
let Ok(gateway_id) = HeaderValue::from_str(request_id.as_str()) else {
|
||||
return response;
|
||||
};
|
||||
let headers = response.headers_mut();
|
||||
let usable = |value: &HeaderValue| value.to_str().is_ok_and(|id| !id.trim().is_empty());
|
||||
if !headers.get("x-request-id").is_some_and(usable) {
|
||||
// 只有别名时沿用同一个上游值,避免外层 middleware 再填入入口 ID。
|
||||
let client_id = headers
|
||||
.get("x-oai-request-id")
|
||||
.filter(|value| usable(value))
|
||||
.unwrap_or(&gateway_id)
|
||||
.clone();
|
||||
headers.insert("x-request-id", client_id);
|
||||
}
|
||||
headers.insert("x-gateway-request-id", gateway_id);
|
||||
response
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ use axum::{
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use gateway_core::routing::{PublicModelId, PublicModelProfile};
|
||||
use gateway_core::routing::{PublicModelDescriptor, PublicModelId, PublicModelProfile};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
@@ -14,7 +14,7 @@ use crate::ApiState;
|
||||
|
||||
use super::{
|
||||
auth::{authenticate_client, client_access_error_response},
|
||||
error::model_not_found_response,
|
||||
error::{model_not_found_response, openai_error_response},
|
||||
};
|
||||
|
||||
const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000;
|
||||
@@ -37,19 +37,53 @@ pub(crate) async fn models(
|
||||
Err(error) => return client_access_error_response(error),
|
||||
};
|
||||
|
||||
if query
|
||||
if let Some(version) = query
|
||||
.client_version
|
||||
.as_deref()
|
||||
.is_some_and(|version| !version.trim().is_empty())
|
||||
.filter(|version| !version.trim().is_empty())
|
||||
{
|
||||
let models = service
|
||||
.public_model_profiles(&client)
|
||||
if version.len() > 64
|
||||
|| !version
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || b".-+".contains(&byte))
|
||||
{
|
||||
return openai_error_response(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Invalid client_version",
|
||||
"invalid_request_error",
|
||||
"invalid_client_version",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let catalog = match service.client_model_catalog(&client, version).await {
|
||||
Ok(catalog) => catalog,
|
||||
Err(_) => return catalog_unavailable_response(),
|
||||
};
|
||||
let models = catalog
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, profile)| codex_model_json(profile, index))
|
||||
.collect::<Vec<_>>();
|
||||
.map(|(index, entry)| match entry {
|
||||
PublicModelDescriptor::Adapted(profile) => Ok(codex_model_json(profile, index)),
|
||||
PublicModelDescriptor::Native { model, payload } => {
|
||||
if payload.protocol() != "codex" {
|
||||
return Err(());
|
||||
}
|
||||
let mut value: Value =
|
||||
serde_json::from_slice(payload.body()).map_err(|_| ())?;
|
||||
let object = value.as_object_mut().ok_or(())?;
|
||||
// 别名只改变请求标识,能力、提示词、顺序等仍由上游目录决定。
|
||||
object.insert("slug".to_owned(), Value::String(model.as_str().to_owned()));
|
||||
Ok(value)
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>, ()>>();
|
||||
let Ok(models) = models else {
|
||||
return catalog_unavailable_response();
|
||||
};
|
||||
return (
|
||||
StatusCode::OK,
|
||||
// 目录经过账号选择/别名/多 Provider 聚合,不能借用上游 ETag 表示改写后的正文。
|
||||
[("cache-control", "private, no-store")],
|
||||
Json(json!({
|
||||
"models": models,
|
||||
})),
|
||||
@@ -72,6 +106,16 @@ pub(crate) async fn models(
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn catalog_unavailable_response() -> Response {
|
||||
openai_error_response(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Model catalog is temporarily unavailable",
|
||||
"server_error",
|
||||
"model_catalog_unavailable",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
/// `GET /v1/models/{model_id}`。
|
||||
pub(crate) async fn model_detail(
|
||||
State(state): State<ApiState>,
|
||||
@@ -124,6 +168,7 @@ fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
|
||||
.iter()
|
||||
.any(|effort| effort != "none");
|
||||
let context_window = presentation.context_window_tokens();
|
||||
let max_context_window = presentation.max_context_window_tokens();
|
||||
let input_modalities = if presentation.image_input() {
|
||||
vec!["text", "image"]
|
||||
} else {
|
||||
@@ -173,7 +218,7 @@ fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
|
||||
"supports_parallel_tool_calls": presentation.parallel_tool_calls(),
|
||||
"supports_image_detail_original": presentation.image_detail_original(),
|
||||
"context_window": context_window,
|
||||
"max_context_window": context_window,
|
||||
"max_context_window": max_context_window,
|
||||
"effective_context_window_percent": 95,
|
||||
"experimental_supported_tools": [],
|
||||
"input_modalities": input_modalities,
|
||||
|
||||
@@ -93,20 +93,17 @@ pub(crate) async fn responses(
|
||||
"ingressRequestId": ingress_id.as_ref().and_then(|Extension(id)| id.header_value().to_str().ok()),
|
||||
}), headers.iter().map(|(name, value)| (name.as_str(), value.as_bytes())));
|
||||
trace.capture("client.request.body", &body);
|
||||
let request_id_header = HeaderValue::from_str(started.request_id.as_str()).ok();
|
||||
let request_id = started.request_id;
|
||||
let StartedExecution {
|
||||
stream, session, ..
|
||||
} = started;
|
||||
let mut response = if stream {
|
||||
let response = if stream {
|
||||
stream_execution_response(session, connection_guard).await
|
||||
} else {
|
||||
drop(connection_guard);
|
||||
collect_execution_response(session).await
|
||||
};
|
||||
if let Some(id) = request_id_header {
|
||||
response.headers_mut().insert("x-gateway-request-id", id);
|
||||
}
|
||||
response
|
||||
crate::openai::with_model_request_id(response, &request_id)
|
||||
}
|
||||
|
||||
/// 从 socket 与标准转发头提取旧 Usage 页面使用的诊断事实。
|
||||
@@ -271,7 +268,26 @@ fn engine_error_response_with_headers(
|
||||
error: &EngineError,
|
||||
response_headers: &[ProviderResponseHeader],
|
||||
) -> Response {
|
||||
apply_response_headers(engine_error_response(error), response_headers)
|
||||
let response = engine_error_response(error);
|
||||
let failure_request_ids: Vec<_> = ["x-request-id", "x-oai-request-id"]
|
||||
.into_iter()
|
||||
.flat_map(|name| {
|
||||
response
|
||||
.headers()
|
||||
.get_all(name)
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(move |value| (name, value))
|
||||
})
|
||||
.collect();
|
||||
// 失败后仍须交付已采集的 turn state;只隔离 opening 身份,不丢弃会话状态。
|
||||
let mut response = apply_response_headers(response, response_headers);
|
||||
response.headers_mut().remove("x-request-id");
|
||||
response.headers_mut().remove("x-oai-request-id");
|
||||
for (name, value) in failure_request_ids {
|
||||
response.headers_mut().append(name, value);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
/// 编码首个 SSE frame 后提交下游,再持续驱动同一执行会话。
|
||||
|
||||
@@ -14,7 +14,7 @@ use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
|
||||
use super::{
|
||||
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, PendingExecution, ProtocolErrorBody},
|
||||
connection::{FramePhase, ResponsesWebSocketConnection, WriteContext},
|
||||
protocol::{error_event, response_metadata_event},
|
||||
protocol::{error_event, initial_engine_error_event, response_metadata_event},
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
@@ -273,8 +273,12 @@ async fn send_initial_engine_error(
|
||||
{
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
let error = gateway_error_from_engine(error);
|
||||
send_gateway_error(connection, &error, request_id).await
|
||||
send_error_message(
|
||||
connection,
|
||||
initial_engine_error_event(error, request_id, &response_headers),
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn commit_connection_replay(
|
||||
@@ -393,10 +397,31 @@ async fn send_error(
|
||||
message: &str,
|
||||
param: Option<&str>,
|
||||
request_id: &Arc<str>,
|
||||
) -> ForwardOutcome {
|
||||
send_error_message(
|
||||
connection,
|
||||
error_event(
|
||||
status,
|
||||
error_type,
|
||||
code,
|
||||
message,
|
||||
param,
|
||||
Some(request_id),
|
||||
serde_json::Map::new(),
|
||||
),
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_error_message(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
message: String,
|
||||
request_id: &Arc<str>,
|
||||
) -> ForwardOutcome {
|
||||
if connection
|
||||
.send_text(
|
||||
error_event(status, error_type, code, message, param, Some(request_id)),
|
||||
message,
|
||||
WriteContext::request(request_id, FramePhase::Error),
|
||||
)
|
||||
.await
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
//! Responses WebSocket 入站请求与下行事件的纯协议映射。
|
||||
|
||||
use axum::http::{HeaderName, StatusCode};
|
||||
use axum::http::{HeaderName, HeaderValue, StatusCode};
|
||||
use gateway_core::engine::EngineError;
|
||||
use gateway_core::event::ProviderResponseHeader;
|
||||
use serde_json::{Map, Value, json};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
|
||||
|
||||
use super::super::{
|
||||
DecodedResponsesRequest, ProtocolErrorBody, RequestDecodeError,
|
||||
request::{OpenAiRequestHeaders, RequestDecodeSource},
|
||||
@@ -91,21 +94,7 @@ pub(super) fn response_metadata_event(
|
||||
request_id: &str,
|
||||
response_headers: &[ProviderResponseHeader],
|
||||
) -> String {
|
||||
let connection_options = super::super::response_connection_options(response_headers);
|
||||
let mut headers = Map::new();
|
||||
for header in response_headers {
|
||||
if !super::super::response_header_is_forwardable(header.name(), &connection_options) {
|
||||
continue;
|
||||
}
|
||||
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(value) = std::str::from_utf8(header.value()) else {
|
||||
continue;
|
||||
};
|
||||
// 官方 response.metadata 使用 string map,无法表达同名多值;后值保持既有覆盖语义。
|
||||
headers.insert(name.as_str().to_owned(), Value::String(value.to_owned()));
|
||||
}
|
||||
let mut headers = response_event_headers(response_headers);
|
||||
headers
|
||||
.entry("x-request-id".to_owned())
|
||||
.or_insert_with(|| Value::String(request_id.to_owned()));
|
||||
@@ -116,6 +105,84 @@ pub(super) fn response_metadata_event(
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn response_event_headers(response_headers: &[ProviderResponseHeader]) -> Map<String, Value> {
|
||||
let connection_options = super::super::response_connection_options(response_headers);
|
||||
let mut headers = Map::new();
|
||||
for header in response_headers {
|
||||
if !super::super::response_header_is_forwardable(header.name(), &connection_options) {
|
||||
continue;
|
||||
}
|
||||
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(http_value) = HeaderValue::from_bytes(header.value()) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(value) = std::str::from_utf8(header.value()) else {
|
||||
continue;
|
||||
};
|
||||
if matches!(name.as_str(), "x-request-id" | "x-oai-request-id")
|
||||
&& (value.trim().is_empty() || http_value.to_str().is_err())
|
||||
{
|
||||
continue;
|
||||
}
|
||||
// 官方 WS 事件使用 string map,无法表达同名多值;后值保持既有覆盖语义。
|
||||
headers.insert(name.as_str().to_owned(), Value::String(value.to_owned()));
|
||||
}
|
||||
headers
|
||||
}
|
||||
|
||||
pub(super) fn initial_engine_error_event(
|
||||
error: &EngineError,
|
||||
request_id: &str,
|
||||
response_headers: &[ProviderResponseHeader],
|
||||
) -> String {
|
||||
let gateway = gateway_error_from_engine(error);
|
||||
let (default_status, default_type, default_code) = gateway_error_contract(gateway.kind());
|
||||
let provider = match error {
|
||||
EngineError::Provider(provider) => Some(provider),
|
||||
_ => None,
|
||||
};
|
||||
let upstream = provider.and_then(|error| error.client_visible_upstream_response());
|
||||
// 最终失败响应优先于会话快照,避免混入先前 attempt 或 opening 的关联头。
|
||||
let mut headers =
|
||||
response_event_headers(upstream.map_or(response_headers, |response| response.headers()));
|
||||
if upstream.is_none() {
|
||||
// 会话 opening 身份不证明当前失败属于该请求;仅保留其他允许的会话头。
|
||||
headers.remove("x-request-id");
|
||||
headers.remove("x-oai-request-id");
|
||||
}
|
||||
let status = upstream
|
||||
.map(|response| response.status())
|
||||
.or_else(|| provider.and_then(|error| error.upstream_status()))
|
||||
.and_then(|value| StatusCode::from_u16(value).ok())
|
||||
.filter(|value| !value.is_success() && !value.is_informational())
|
||||
.unwrap_or(default_status);
|
||||
// 无原响应关联头时使用 Provider 的失败事实;不能让旧会话 ID 遮住当前失败。
|
||||
if !has_request_id(&headers)
|
||||
&& let Some(request_id) = provider.and_then(|error| error.upstream_request_id())
|
||||
&& let Ok(value) = HeaderValue::from_str(request_id.as_str())
|
||||
&& let Ok(value) = value.to_str()
|
||||
&& !value.trim().is_empty()
|
||||
{
|
||||
headers.insert("x-request-id".to_owned(), Value::String(value.to_owned()));
|
||||
}
|
||||
// 只桥接 Provider 已提取的结构化错误;原始 HTML/文本正文不成为安全 fallback 的 message。
|
||||
error_event(
|
||||
status,
|
||||
gateway.client_error_type().unwrap_or(default_type),
|
||||
gateway.client_error_code().unwrap_or(default_code),
|
||||
gateway.client_message(),
|
||||
None,
|
||||
Some(request_id),
|
||||
headers,
|
||||
)
|
||||
}
|
||||
|
||||
fn has_request_id(headers: &Map<String, Value>) -> bool {
|
||||
headers.contains_key("x-request-id") || headers.contains_key("x-oai-request-id")
|
||||
}
|
||||
|
||||
pub(super) fn error_event(
|
||||
status: StatusCode,
|
||||
error_type: &str,
|
||||
@@ -123,6 +190,7 @@ pub(super) fn error_event(
|
||||
message: &str,
|
||||
param: Option<&str>,
|
||||
request_id: Option<&str>,
|
||||
mut headers: Map<String, Value>,
|
||||
) -> String {
|
||||
let mut error = Map::new();
|
||||
error.insert("type".to_owned(), Value::String(error_type.to_owned()));
|
||||
@@ -136,11 +204,26 @@ pub(super) fn error_event(
|
||||
event.insert("status".to_owned(), Value::Number(status.as_u16().into()));
|
||||
event.insert("error".to_owned(), Value::Object(error));
|
||||
if let Some(request_id) = request_id {
|
||||
// Codex 从本 error 事件的 headers 读 ID,独立 metadata 和顶层 request_id 不能替代。
|
||||
// 网关身份独立保存;有上游关联头时不冒充或覆盖上游身份。
|
||||
if !has_request_id(&headers) {
|
||||
headers.insert(
|
||||
"x-request-id".to_owned(),
|
||||
Value::String(request_id.to_owned()),
|
||||
);
|
||||
}
|
||||
headers.insert(
|
||||
"x-gateway-request-id".to_owned(),
|
||||
Value::String(request_id.to_owned()),
|
||||
);
|
||||
event.insert(
|
||||
"request_id".to_owned(),
|
||||
Value::String(request_id.to_owned()),
|
||||
);
|
||||
}
|
||||
if !headers.is_empty() {
|
||||
event.insert("headers".to_owned(), Value::Object(headers));
|
||||
}
|
||||
Value::Object(event).to_string()
|
||||
}
|
||||
|
||||
@@ -152,5 +235,6 @@ pub(super) fn connection_limit_event() -> String {
|
||||
"Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.",
|
||||
None,
|
||||
None,
|
||||
Map::new(),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use gateway_core::engine::execution::{
|
||||
};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle};
|
||||
use gateway_core::routing::{PublicModelId, PublicModelProfile};
|
||||
use gateway_core::routing::{ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::auth::ClientApiKeyAuthError;
|
||||
@@ -52,11 +52,14 @@ impl OpenAiService {
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn public_model_profiles(
|
||||
pub(crate) async fn client_model_catalog(
|
||||
&self,
|
||||
client: &AuthenticatedClient,
|
||||
) -> Vec<PublicModelProfile> {
|
||||
self.execution.public_model_profiles(client)
|
||||
client_version: &str,
|
||||
) -> Result<Vec<PublicModelDescriptor>, ProviderCatalogUnavailable> {
|
||||
self.execution
|
||||
.client_model_catalog(client, "codex", client_version)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn contains_public_model(
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header},
|
||||
};
|
||||
use gateway_api::admin;
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::super::{AdminTestFixture, AdminTestState};
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_info_requires_admin_and_a_valid_account_query() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (query, authenticated, expected) in [
|
||||
("?accountId=acct_test", false, StatusCode::UNAUTHORIZED),
|
||||
("", true, StatusCode::BAD_REQUEST),
|
||||
("?accountId=bad", true, StatusCode::BAD_REQUEST),
|
||||
(
|
||||
"?accountId=acct_test&refresh=true",
|
||||
true,
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
// 此夹具未提供账号 Store,合法查询应透传服务不可用,而非绕过查询。
|
||||
(
|
||||
"?accountId=acct_test",
|
||||
true,
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
),
|
||||
] {
|
||||
let mut request = Request::builder()
|
||||
.uri(format!("/api/admin/accounts/personal-info{query}"))
|
||||
.header("x-request-id", "req_personal_info");
|
||||
if authenticated {
|
||||
request = request.header(header::COOKIE, "cpr_admin_session=valid-session");
|
||||
}
|
||||
let response = admin::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(request.body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), expected, "{query}");
|
||||
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_forecast_requires_admin_and_valid_account_query() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (uri, authenticated, expected) in [
|
||||
(
|
||||
"/api/admin/accounts/quota-forecast?accountId=acct_test",
|
||||
false,
|
||||
StatusCode::UNAUTHORIZED,
|
||||
),
|
||||
(
|
||||
"/api/admin/accounts/quota-forecast",
|
||||
true,
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
"/api/admin/accounts/quota-forecast?accountId=bad",
|
||||
true,
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
"/api/admin/accounts/quota-forecast?accountId=acct_test&refresh=true",
|
||||
true,
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
] {
|
||||
let mut request = Request::builder()
|
||||
.uri(uri)
|
||||
.header("x-request-id", "req_forecast");
|
||||
if authenticated {
|
||||
request = request.header(header::COOKIE, "cpr_admin_session=valid-session");
|
||||
}
|
||||
let response = admin::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(request.body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), expected, "{uri}");
|
||||
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
|
||||
let body = to_bytes(response.into_body(), 8192).await.unwrap();
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
assert!(value["data"].is_null());
|
||||
assert!(value["message"].is_string());
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,77 @@
|
||||
mod handlers;
|
||||
mod presenter;
|
||||
|
||||
mod personal_info {
|
||||
use chrono::{TimeZone as _, Utc};
|
||||
use gateway_admin::model::{
|
||||
AdminError,
|
||||
provider_credentials::{AccountPersonalInfo, ProviderSubscription},
|
||||
};
|
||||
use gateway_api::admin::accounts::{AccountPersonalInfoData, AccountSubscriptionData};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn subscription_response_exposes_only_display_fields_and_preserves_unknowns() {
|
||||
let observed = Utc.with_ymd_and_hms(2026, 9, 14, 0, 0, 0).unwrap();
|
||||
let response = AccountSubscriptionData::from(ProviderSubscription {
|
||||
starts_at: None,
|
||||
expires_at: observed,
|
||||
will_renew: None,
|
||||
billing_period: None,
|
||||
billing_currency: Some("USD".to_owned()),
|
||||
observed_at: observed,
|
||||
});
|
||||
assert_eq!(
|
||||
serde_json::to_value(response).unwrap(),
|
||||
json!({
|
||||
"startsAt": null, "expiresAt": "2026-09-14T00:00:00+00:00",
|
||||
"willRenew": null, "billingPeriod": null, "billingCurrency": "USD",
|
||||
"observedAt": "2026-09-14T00:00:00+00:00"
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn personal_info_response_preserves_subscription_when_profile_fails() {
|
||||
let observed = Utc.with_ymd_and_hms(2026, 9, 14, 0, 0, 0).unwrap();
|
||||
let response = AccountPersonalInfoData::from(AccountPersonalInfo {
|
||||
profile: Err(AdminError::bad_gateway("上游服务请求失败")),
|
||||
subscription: Some(ProviderSubscription {
|
||||
starts_at: None,
|
||||
expires_at: observed,
|
||||
will_renew: None,
|
||||
billing_period: None,
|
||||
billing_currency: None,
|
||||
observed_at: observed,
|
||||
}),
|
||||
});
|
||||
let value = serde_json::to_value(response).unwrap();
|
||||
assert!(value["profile"].is_null());
|
||||
assert_eq!(value["profileError"], "上游服务请求失败");
|
||||
assert_eq!(
|
||||
value["subscription"]["expiresAt"],
|
||||
"2026-09-14T00:00:00+00:00"
|
||||
);
|
||||
assert_eq!(value.as_object().unwrap().len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn personal_info_response_preserves_unknown_subscription_and_profile_error() {
|
||||
let response = AccountPersonalInfoData::from(AccountPersonalInfo {
|
||||
profile: Err(AdminError::unavailable("Provider 服务暂不可用")),
|
||||
subscription: None,
|
||||
});
|
||||
assert_eq!(
|
||||
serde_json::to_value(response).unwrap(),
|
||||
json!({
|
||||
"profile": null,
|
||||
"profileError": "Provider 服务暂不可用",
|
||||
"subscription": null,
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
mod query {
|
||||
use gateway_admin::model::accounts::{AccountSortField, AccountStatus, SortDirection};
|
||||
use gateway_api::admin::accounts::ListQuery;
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
use gateway_admin::model::{
|
||||
provider_credentials::AccountUsagePeriod,
|
||||
quota_forecast::{AccountQuotaForecast, AccountQuotaForecastReport, QuotaForecastSource},
|
||||
};
|
||||
use gateway_api::admin::accounts::AccountQuotaForecastData;
|
||||
|
||||
#[test]
|
||||
fn quota_forecast_projection_only_exposes_capacity_and_preserves_null_zero() {
|
||||
let now = "2026-09-12T00:00:00Z".parse().unwrap();
|
||||
let forecast = AccountQuotaForecast {
|
||||
period: AccountUsagePeriod::Weekly,
|
||||
target_seconds: 7 * 86_400,
|
||||
extrapolated: false,
|
||||
source: Some(QuotaForecastSource {
|
||||
label: "周额度".to_owned(),
|
||||
used_percent: Some(100.0),
|
||||
observed_at: Some(now),
|
||||
reset_at: now,
|
||||
tokens: Some(1_000_000),
|
||||
usd: Some(0.1234),
|
||||
}),
|
||||
unavailable_reason: None,
|
||||
low_sample: false,
|
||||
incomplete_cost: true,
|
||||
incomplete_tokens: false,
|
||||
estimated_tokens: Some(1_000_000),
|
||||
estimated_usd: None,
|
||||
remaining_tokens: Some(0),
|
||||
remaining_usd: None,
|
||||
};
|
||||
let mut monthly = forecast.clone();
|
||||
monthly.period = AccountUsagePeriod::Monthly;
|
||||
monthly.extrapolated = true;
|
||||
monthly.target_seconds = 30 * 86_400;
|
||||
let view = AccountQuotaForecastData::from(AccountQuotaForecastReport {
|
||||
account_id: "acct_forecast".to_owned(),
|
||||
generated_at: now,
|
||||
forecasts: [forecast, monthly],
|
||||
});
|
||||
let value = serde_json::to_value(view).unwrap();
|
||||
assert_eq!(value["accountId"], "acct_forecast");
|
||||
assert_eq!(value["generatedAt"], "2026-09-12T08:00:00+08:00");
|
||||
let week = &value["forecasts"][0];
|
||||
assert_eq!(week["estimatedTokens"], 1_000_000);
|
||||
assert_eq!(week["estimatedTokensDisplay"], "1M");
|
||||
assert!(week["estimatedUsd"].is_null());
|
||||
assert_eq!(week["estimatedUsdDisplay"], "—");
|
||||
assert_eq!(week["remainingTokensDisplay"], "0");
|
||||
assert_eq!(
|
||||
week["source"],
|
||||
serde_json::json!({
|
||||
"label": "周额度",
|
||||
"usedPercent": 100.0,
|
||||
"usedPercentDisplay": "100.0%",
|
||||
"observedAt": "2026-09-12T08:00:00+08:00",
|
||||
"observedAtDisplay": "2026-09-12 08:00:00",
|
||||
"resetAt": "2026-09-12T08:00:00+08:00",
|
||||
"tokensDisplay": "1M",
|
||||
"usdDisplay": "$0.1234"
|
||||
})
|
||||
);
|
||||
assert!(week.get("method").is_none());
|
||||
assert!(week.get("methodDisplay").is_none());
|
||||
assert!(value.get("generatedAtDisplay").is_none());
|
||||
assert_eq!(value["forecasts"][1]["period"], "monthly");
|
||||
assert_eq!(value["forecasts"][1]["targetDays"], 30.0);
|
||||
assert_eq!(value["forecasts"][1]["extrapolated"], true);
|
||||
assert!(value.get("account").is_none());
|
||||
}
|
||||
@@ -10,6 +10,71 @@ use serde_json::json;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState};
|
||||
|
||||
#[test]
|
||||
fn custom_client_keys_preserve_migrated_values_and_redact_debug() {
|
||||
for value in [
|
||||
"q".to_owned(),
|
||||
"sk-old!@/key+=value".to_owned(),
|
||||
"x".repeat(8192),
|
||||
] {
|
||||
let request: CreateClientKeyRequest = serde_json::from_value(json!({
|
||||
"name": "migration", "customKey": value, "groupIds": [],
|
||||
"maxConcurrency": 2, "requestsPerMinute": 0
|
||||
}))
|
||||
.unwrap();
|
||||
let debug = format!("{request:?}");
|
||||
assert!(debug.contains("[REDACTED]"));
|
||||
let command = request.into_command().unwrap();
|
||||
assert_eq!(command.custom_key.unwrap().expose_for_auth(), value);
|
||||
}
|
||||
let secret = "legacy-key-should-not-leak";
|
||||
let request: CreateClientKeyRequest = serde_json::from_value(json!({
|
||||
"name": "migration", "customKey": secret, "groupIds": [],
|
||||
"maxConcurrency": 0, "requestsPerMinute": 0
|
||||
}))
|
||||
.unwrap();
|
||||
assert!(!format!("{request:?}").contains(secret));
|
||||
assert!(!format!("{:?}", request.into_command().unwrap()).contains(secret));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_client_key_is_optional_and_only_rejects_untransportable_values() {
|
||||
let payload = json!({"name": "migration", "groupIds": [], "maxConcurrency": 0,
|
||||
"requestsPerMinute": 0});
|
||||
for value in [None, Some(json!(null)), Some(json!(""))] {
|
||||
let mut payload = payload.clone();
|
||||
if let Some(value) = value {
|
||||
payload["customKey"] = value;
|
||||
}
|
||||
let command = serde_json::from_value::<CreateClientKeyRequest>(payload)
|
||||
.unwrap()
|
||||
.into_command()
|
||||
.unwrap();
|
||||
assert!(command.custom_key.is_none());
|
||||
}
|
||||
for value in [
|
||||
" key",
|
||||
"key ",
|
||||
"key with spaces",
|
||||
"key\n",
|
||||
"key\tvalue",
|
||||
"密钥",
|
||||
] {
|
||||
let mut payload = payload.clone();
|
||||
payload["customKey"] = json!(value);
|
||||
let error = serde_json::from_value::<CreateClientKeyRequest>(payload)
|
||||
.unwrap()
|
||||
.into_command()
|
||||
.unwrap_err();
|
||||
assert_eq!(error.field(), "customKey");
|
||||
assert!(!format!("{error:?}").contains(value));
|
||||
}
|
||||
let mut update = payload;
|
||||
update["id"] = json!("key_existing");
|
||||
update["customKey"] = json!("replacement-credential");
|
||||
assert!(serde_json::from_value::<UpdateClientKeyRequest>(update).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn budget_inputs_preserve_decimal_precision_and_omitted_updates() {
|
||||
let payload = json!({"name": "budget", "groupIds": [], "maxConcurrency": 3,
|
||||
@@ -63,10 +128,6 @@ fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
|
||||
"search": "a".repeat(257)
|
||||
}))
|
||||
.expect("deserialize oversized search");
|
||||
let secret_search = serde_json::from_value::<ListClientKeysQuery>(json!({
|
||||
"search": format!("sk_{}", "a".repeat(43))
|
||||
}))
|
||||
.expect("deserialize secret search");
|
||||
let invalid_sort = serde_json::from_value::<ListClientKeysQuery>(json!({
|
||||
"sortBy": "plaintextKey",
|
||||
"sortDirection": "asc"
|
||||
@@ -92,13 +153,6 @@ fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
|
||||
.field(),
|
||||
"search"
|
||||
);
|
||||
assert_eq!(
|
||||
secret_search
|
||||
.into_command()
|
||||
.expect_err("reject full key in search")
|
||||
.field(),
|
||||
"search"
|
||||
);
|
||||
assert_eq!(
|
||||
invalid_sort
|
||||
.into_command()
|
||||
@@ -140,6 +194,16 @@ fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_key_name_search_does_not_interpret_names_as_credentials() {
|
||||
let name = format!("sk_{}", "a".repeat(43));
|
||||
let command = serde_json::from_value::<ListClientKeysQuery>(json!({ "search": name }))
|
||||
.unwrap()
|
||||
.into_command()
|
||||
.unwrap();
|
||||
assert_eq!(command.search.as_deref(), Some(name.as_str()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
|
||||
let valid = serde_json::from_value::<CreateClientKeyRequest>(json!({
|
||||
|
||||
@@ -318,3 +318,187 @@ async fn admin_auth_failures_should_use_stable_chinese_contracts() {
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
mod provider {
|
||||
use std::time::SystemTime;
|
||||
|
||||
use chrono::Utc;
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
Revision,
|
||||
accounts::{AccountPageItem, AccountRecord},
|
||||
},
|
||||
ports::provider::{ProviderAdminError, ProviderAdminErrorKind as Kind},
|
||||
};
|
||||
use gateway_core::{
|
||||
account::{
|
||||
AccountStatusFacts, AccountWeight, CredentialState, QuotaState, resolve_account_status,
|
||||
},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_public_errors_survive_import_and_refresh_http_handlers() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for provider in ["openai", "xai"] {
|
||||
*fixture.account.lock().unwrap() = Some(account(provider));
|
||||
for (kind, status, code, message) in [
|
||||
(Kind::Invalid, 400, 40001, "刷新令牌已被使用,请重新授权"),
|
||||
(
|
||||
Kind::Conflict,
|
||||
409,
|
||||
40901,
|
||||
"令牌刷新繁忙,请等待当前刷新完成后重试",
|
||||
),
|
||||
(
|
||||
Kind::BadGateway,
|
||||
502,
|
||||
50201,
|
||||
"OpenAI 返回的 Codex PAT 身份资料不完整或格式无效,请稍后重试",
|
||||
),
|
||||
(
|
||||
Kind::Unavailable,
|
||||
503,
|
||||
50301,
|
||||
"暂时无法向 OpenAI 验证 Codex PAT,请稍后重试",
|
||||
),
|
||||
(
|
||||
Kind::Ambiguous,
|
||||
502,
|
||||
50202,
|
||||
"令牌刷新结果未知,请先核对账号状态,不要立即重复刷新",
|
||||
),
|
||||
] {
|
||||
*fixture.provider_error.lock().unwrap() = Some(
|
||||
ProviderAdminError::new(kind)
|
||||
.with_public_message(message)
|
||||
.with_message("raw-upstream-secret-marker"),
|
||||
);
|
||||
for (path, body) in [
|
||||
(
|
||||
"/api/admin/accounts/refresh",
|
||||
json!({"accountId": "acct_error_test"}),
|
||||
),
|
||||
(
|
||||
"/api/admin/accounts/import",
|
||||
json!({"provider": provider, "data": {"accessToken": "at-synthetic-test"}}),
|
||||
),
|
||||
] {
|
||||
let response = send(&fixture, path, body).await;
|
||||
assert_eq!(response.headers()["x-request-id"], "req_admin_errors");
|
||||
let actual = response_json(response).await;
|
||||
assert_eq!(
|
||||
actual,
|
||||
(
|
||||
StatusCode::from_u16(status).unwrap(),
|
||||
"application/json".to_owned(),
|
||||
json!({"code": code, "message": message, "data": null}),
|
||||
)
|
||||
);
|
||||
assert!(!actual.2.to_string().contains("raw-upstream-secret-marker"));
|
||||
assert_eq!(fixture.auth.audit_count(), 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_diagnostics_without_public_messages_still_use_safe_fallbacks() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (kind, status, code, message) in [
|
||||
(Kind::BadGateway, 502, 50201, "上游服务请求失败"),
|
||||
(Kind::Unavailable, 503, 50301, "Provider 服务暂不可用"),
|
||||
(
|
||||
Kind::Ambiguous,
|
||||
502,
|
||||
50202,
|
||||
"上游执行结果未知,请刷新状态后再决定是否重试",
|
||||
),
|
||||
(Kind::Internal, 500, 50001, "服务内部错误"),
|
||||
] {
|
||||
let error = ProviderAdminError::new(kind).with_message("private-provider-diagnostics");
|
||||
// 即使 Provider 错标了公开文案,未知内部异常仍不得通过 500 信封下发。
|
||||
*fixture.provider_error.lock().unwrap() = Some(if kind == Kind::Internal {
|
||||
error.with_public_message("internal-detail-must-stay-hidden")
|
||||
} else {
|
||||
error
|
||||
});
|
||||
let response = send(
|
||||
&fixture,
|
||||
"/api/admin/accounts/import",
|
||||
json!({"provider": "openai", "data": {}}),
|
||||
)
|
||||
.await;
|
||||
let (actual_status, _, body) = response_json(response).await;
|
||||
assert_eq!(actual_status.as_u16(), status);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({"code": code, "message": message, "data": null})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(fixture: &AdminTestFixture, path: &str, body: Value) -> axum::response::Response {
|
||||
let mut request = request(Method::POST, path, Body::from(body.to_string()));
|
||||
request
|
||||
.headers_mut()
|
||||
.insert(header::COOKIE, SESSION_COOKIE.parse().unwrap());
|
||||
request
|
||||
.headers_mut()
|
||||
.insert(header::CONTENT_TYPE, "application/json".parse().unwrap());
|
||||
let name = header::HeaderName::from_static("x-request-id");
|
||||
app(fixture.state())
|
||||
.layer(PropagateRequestIdLayer::new(name.clone()))
|
||||
.layer(SetRequestIdLayer::new(name, MakeRequestUuid))
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn account(provider: &str) -> AccountPageItem {
|
||||
let now = Utc::now();
|
||||
let facts = AccountStatusFacts {
|
||||
enabled: true,
|
||||
credential_state: CredentialState::Ready,
|
||||
access_token_expires_at: None,
|
||||
quota: QuotaState::unknown(),
|
||||
rate_limited_until: None,
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
};
|
||||
AccountPageItem {
|
||||
account: AccountRecord {
|
||||
id: "acct_error_test".to_owned(),
|
||||
provider_kind: ProviderKind::new(provider).unwrap(),
|
||||
groups: Vec::new(),
|
||||
name: "synthetic account".to_owned(),
|
||||
email: None,
|
||||
upstream_user_id: None,
|
||||
upstream_account_id: None,
|
||||
plan_type: None,
|
||||
authentication_kind: "oauth".to_owned(),
|
||||
credential_revision: Revision::new(1).unwrap(),
|
||||
has_refresh_token: true,
|
||||
access_token_expires_at: None,
|
||||
next_refresh_at: None,
|
||||
enabled: facts.enabled,
|
||||
concurrency_limit: None,
|
||||
weight: AccountWeight::default(),
|
||||
outbound_proxy: None,
|
||||
credential_state: facts.credential_state,
|
||||
credential_observed_at: now,
|
||||
quota: facts.quota,
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
},
|
||||
projection: resolve_account_status(&facts, SystemTime::now()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,12 +40,11 @@ use gateway_admin::{
|
||||
},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialListQuery,
|
||||
CredentialMutationResult, CredentialPage, CredentialRotationCommit,
|
||||
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
|
||||
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
|
||||
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput,
|
||||
ProviderModels, ProviderQuota,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
|
||||
CredentialRotationCommit, PendingAuthorizationMutation, PrepareCredentialImport,
|
||||
PrepareCredentialRefresh, PrepareCredentialRotation, PreparedAuthorizationCommit,
|
||||
PreparedCredentialImport, PreparedCredentialRotation, ProviderExport,
|
||||
ProviderExportCredentialInput, ProviderModels, ProviderQuota,
|
||||
},
|
||||
settings::{
|
||||
AdminApiKey, AdminApiKeyMutation, ModelMappings, ReplaceRuntimeSettings,
|
||||
@@ -97,6 +96,8 @@ pub(super) struct AdminTestFixture {
|
||||
pub ops_errors: Arc<Mutex<Vec<OpsError>>>,
|
||||
pub dashboard_observation: Arc<Mutex<Option<DashboardObservation>>>,
|
||||
pub dashboard_summary_range: Arc<Mutex<Option<TimeRange>>>,
|
||||
pub provider_error: Arc<Mutex<Option<ProviderAdminError>>>,
|
||||
pub account: Arc<Mutex<Option<AccountPageItem>>>,
|
||||
}
|
||||
|
||||
impl AdminTestFixture {
|
||||
@@ -116,6 +117,8 @@ impl AdminTestFixture {
|
||||
let ops_errors = Arc::new(Mutex::new(Vec::new()));
|
||||
let dashboard_observation = Arc::new(Mutex::new(None));
|
||||
let dashboard_summary_range = Arc::new(Mutex::new(None));
|
||||
let provider_error = Arc::new(Mutex::new(None));
|
||||
let account = Arc::new(Mutex::new(None));
|
||||
let unused = Arc::new(UnusedStore {
|
||||
usage_records: Arc::clone(&usage_records),
|
||||
usage_detail: Arc::clone(&usage_detail),
|
||||
@@ -123,6 +126,7 @@ impl AdminTestFixture {
|
||||
ops_errors: Arc::clone(&ops_errors),
|
||||
dashboard_observation: Arc::clone(&dashboard_observation),
|
||||
dashboard_summary_range: Arc::clone(&dashboard_summary_range),
|
||||
account: Arc::clone(&account),
|
||||
});
|
||||
let stores = AdminStorePorts::new(
|
||||
AdminAccountStorePorts::new(
|
||||
@@ -138,8 +142,8 @@ impl AdminTestFixture {
|
||||
gateway_admin::ports::backup::BackupStorePorts::disabled(),
|
||||
);
|
||||
let providers: Vec<Arc<dyn ProviderAdmin>> = vec![
|
||||
Arc::new(UnusedProvider::new("openai")),
|
||||
Arc::new(UnusedProvider::new("xai")),
|
||||
Arc::new(UnusedProvider::new("openai", Arc::clone(&provider_error))),
|
||||
Arc::new(UnusedProvider::new("xai", Arc::clone(&provider_error))),
|
||||
];
|
||||
let bundle = gateway_admin::initialize(
|
||||
AdminConfig {
|
||||
@@ -166,6 +170,8 @@ impl AdminTestFixture {
|
||||
ops_errors,
|
||||
dashboard_observation,
|
||||
dashboard_summary_range,
|
||||
provider_error,
|
||||
account,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -703,6 +709,7 @@ struct UnusedStore {
|
||||
ops_errors: Arc<Mutex<Vec<OpsError>>>,
|
||||
dashboard_observation: Arc<Mutex<Option<DashboardObservation>>>,
|
||||
dashboard_summary_range: Arc<Mutex<Option<TimeRange>>>,
|
||||
account: Arc<Mutex<Option<AccountPageItem>>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -717,9 +724,12 @@ impl AccountStore for UnusedStore {
|
||||
|
||||
async fn load_account(
|
||||
&self,
|
||||
_: &str,
|
||||
id: &str,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
if let Some(account) = self.account.lock().expect("account").as_ref() {
|
||||
return Ok((account.account.id == id).then(|| account.clone()));
|
||||
}
|
||||
Err(unavailable("account"))
|
||||
}
|
||||
|
||||
@@ -738,12 +748,11 @@ impl AccountStore for UnusedStore {
|
||||
Err(unavailable("account quota window usage"))
|
||||
}
|
||||
|
||||
async fn list_credentials(
|
||||
async fn load_quota_forecast_history(
|
||||
&self,
|
||||
_: &ProviderKind,
|
||||
_: CredentialListQuery,
|
||||
) -> AdminStoreResult<CredentialPage> {
|
||||
Err(unavailable("credential list"))
|
||||
_: &AccountUsageWindowQuery,
|
||||
) -> AdminStoreResult<gateway_admin::model::quota_forecast_sampling::QuotaForecastHistory> {
|
||||
Err(unavailable("quota forecast history"))
|
||||
}
|
||||
|
||||
async fn credential_details(
|
||||
@@ -936,12 +945,14 @@ impl ObservabilityStore for UnusedStore {
|
||||
|
||||
struct UnusedProvider {
|
||||
kind: ProviderKind,
|
||||
error: Arc<Mutex<Option<ProviderAdminError>>>,
|
||||
}
|
||||
|
||||
impl UnusedProvider {
|
||||
fn new(kind: &str) -> Self {
|
||||
fn new(kind: &str, error: Arc<Mutex<Option<ProviderAdminError>>>) -> Self {
|
||||
Self {
|
||||
kind: ProviderKind::new(kind).expect("provider kind"),
|
||||
error,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -982,7 +993,12 @@ impl ProviderAdmin for UnusedProvider {
|
||||
&self,
|
||||
_: PrepareCredentialImport,
|
||||
) -> Result<PreparedCredentialImport, ProviderAdminError> {
|
||||
Err(unsupported_provider())
|
||||
Err(self
|
||||
.error
|
||||
.lock()
|
||||
.expect("provider error")
|
||||
.clone()
|
||||
.unwrap_or_else(unsupported_provider))
|
||||
}
|
||||
|
||||
async fn start_authorization(
|
||||
@@ -1010,7 +1026,12 @@ impl ProviderAdmin for UnusedProvider {
|
||||
&self,
|
||||
_: PrepareCredentialRefresh,
|
||||
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
|
||||
Err(unsupported_provider())
|
||||
Err(self
|
||||
.error
|
||||
.lock()
|
||||
.expect("provider error")
|
||||
.clone()
|
||||
.unwrap_or_else(unsupported_provider))
|
||||
}
|
||||
|
||||
async fn quota(
|
||||
|
||||
@@ -518,6 +518,98 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_failure_detail_keeps_missing_upstream_facts_and_preparation_trace() {
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header},
|
||||
};
|
||||
use gateway_admin::model::observability::{RequestOutcome, UsageDetail};
|
||||
use gateway_api::admin::observability;
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use crate::admin::{AdminTestFixture, AdminTestState};
|
||||
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let mut record = usage_record_with_account("req_early", "", "", "", "", Utc::now());
|
||||
record.provider_kind = None;
|
||||
record.provider_account_ref = None;
|
||||
record.provider_account_name = None;
|
||||
record.provider_account_email = None;
|
||||
record.provider_account_authentication_kind = None;
|
||||
record.upstream_model_id = None;
|
||||
record.upstream_transport = None;
|
||||
record.http_version = None;
|
||||
record.attempt_count = 0;
|
||||
record.upstream_send_state = "not_sent".to_owned();
|
||||
record.downstream_committed_at = None;
|
||||
record.outcome = RequestOutcome::Failed;
|
||||
record.client_status_code = Some(503);
|
||||
record.upstream_status_code = None;
|
||||
record.error_kind = Some("no_available_provider".to_owned());
|
||||
record.error_message = Some("no upstream account is available".to_owned());
|
||||
record.input_tokens = None;
|
||||
record.output_tokens = None;
|
||||
record.cached_tokens = None;
|
||||
record.cache_write_tokens = None;
|
||||
record.reasoning_tokens = None;
|
||||
record.image_input_tokens = None;
|
||||
record.image_output_tokens = None;
|
||||
record.total_tokens = None;
|
||||
let trace = json!({
|
||||
"events": [
|
||||
{"stage": "attempt.started", "attemptIndex": 1},
|
||||
{"stage": "attempt.failed", "attemptIndex": 1},
|
||||
{"stage": "request.finished", "data": {"attempts": 0}},
|
||||
]
|
||||
});
|
||||
fixture.usage_detail.lock().unwrap().replace(UsageDetail {
|
||||
request: record,
|
||||
attempts: Vec::new(),
|
||||
trace: Some(trace.clone()),
|
||||
related_requests: Vec::new(),
|
||||
});
|
||||
let response = observability::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/api/admin/usage/records/detail?id=req_early")
|
||||
.header(header::COOKIE, "cpr_admin_session=valid-session")
|
||||
.header("x-request-id", "req_zero_attempt_detail")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap();
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
let data = &value["data"];
|
||||
assert_eq!(data["requestId"], "req_early");
|
||||
assert_eq!(data["attemptCount"], 0);
|
||||
assert_eq!(data["attempts"], json!([]));
|
||||
// 明细仍沿用全局 best-effort 合同,不从预备 trace 推断完整尝试列表。
|
||||
assert_eq!(data["attemptsComplete"], false);
|
||||
assert_eq!(data["trace"], trace);
|
||||
assert_eq!(data["logicalOutcome"], "failed");
|
||||
assert_eq!(data["clientStatusCode"], 503);
|
||||
assert_eq!(data["message"], "no upstream account is available");
|
||||
assert_eq!(data["model"], "coding");
|
||||
assert!(data.get("upstreamStatusCode").is_none());
|
||||
assert!(data.get("httpVersion").is_none());
|
||||
for field in [
|
||||
"provider",
|
||||
"accountId",
|
||||
"upstreamModel",
|
||||
"upstreamTransport",
|
||||
"inputTokens",
|
||||
"outputTokens",
|
||||
] {
|
||||
assert_eq!(data.get(field), Some(&serde_json::Value::Null), "{field}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ops_errors_should_keep_account_label_and_authentication_contract() {
|
||||
use axum::{
|
||||
|
||||
@@ -89,7 +89,9 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
actual.sort();
|
||||
let mut expected = vec![
|
||||
"tests/admin/account_groups.rs",
|
||||
"tests/admin/accounts/handlers.rs",
|
||||
"tests/admin/accounts/mod.rs",
|
||||
"tests/admin/accounts/presenter.rs",
|
||||
"tests/admin/auth.rs",
|
||||
"tests/admin/client_keys.rs",
|
||||
"tests/admin/errors.rs",
|
||||
@@ -114,6 +116,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/openai/responses/mod.rs",
|
||||
"tests/openai/responses/request.rs",
|
||||
"tests/openai/responses/websocket/connection.rs",
|
||||
"tests/openai/responses/websocket/forward.rs",
|
||||
"tests/openai/responses/websocket/mod.rs",
|
||||
"tests/openai/responses/websocket/protocol.rs",
|
||||
"tests/openai/router.rs",
|
||||
|
||||
@@ -76,10 +76,28 @@ fn bearer_client_api_key_should_reject_empty_bearer_token() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bearer_client_api_key_should_reject_non_gateway_key_prefix() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer xai-secret"));
|
||||
fn bearer_client_api_key_should_accept_migrated_formats_without_rewriting() {
|
||||
for key in [
|
||||
"q".to_owned(),
|
||||
"sk-old/key+value=:!@".to_owned(),
|
||||
"x".repeat(8192),
|
||||
] {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {key}")).unwrap(),
|
||||
);
|
||||
assert_eq!(bearer_client_api_key(&headers), Ok(key.as_str()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bearer_client_api_key_should_reject_whitespace_inside_tokens() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer key with spaces"),
|
||||
);
|
||||
assert_eq!(
|
||||
bearer_client_api_key(&headers),
|
||||
Err(ClientApiKeyAuthError::InvalidKeyFormat)
|
||||
|
||||
@@ -258,6 +258,7 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio
|
||||
.oneshot(
|
||||
Request::post(*path)
|
||||
.header(AUTHORIZATION, "Bearer sk_images_test")
|
||||
.header("x-request-id", "caller-image-request")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-codex-image-turn-id", "turn_image_route")
|
||||
.header("session-id", "root-image-session")
|
||||
@@ -268,6 +269,11 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio
|
||||
.await
|
||||
.expect("image response");
|
||||
assert_eq!(response.status(), StatusCode::CREATED);
|
||||
assert_eq!(response.headers()["x-request-id"], "req_images_test");
|
||||
assert_eq!(
|
||||
response.headers()["x-gateway-request-id"],
|
||||
"req_images_test"
|
||||
);
|
||||
assert_eq!(
|
||||
response.headers().get("x-image-rate-limit"),
|
||||
Some(&"42".parse().expect("header value"))
|
||||
@@ -315,6 +321,11 @@ async fn image_route_should_return_the_exact_upstream_error_response() {
|
||||
.expect("image response");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
|
||||
assert_eq!(response.headers()["x-request-id"], "req_images_test");
|
||||
assert_eq!(
|
||||
response.headers()["x-gateway-request-id"],
|
||||
"req_images_test"
|
||||
);
|
||||
assert_eq!(
|
||||
response.headers().get("content-type"),
|
||||
Some(&"application/problem+json".parse().expect("content type"))
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
@@ -11,7 +11,8 @@ use gateway_core::engine::execution::{
|
||||
};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::routing::{
|
||||
ModelPresentation, ModelServiceTier, PublicModelId, PublicModelProfile,
|
||||
ModelPresentation, ModelServiceTier, ProviderCatalogUnavailable, PublicModelDescriptor,
|
||||
PublicModelId, PublicModelProfile,
|
||||
};
|
||||
use tower::ServiceExt;
|
||||
|
||||
@@ -19,28 +20,63 @@ use super::{api_router, authenticated_client, authenticated_client_with_min_vers
|
||||
|
||||
pub(super) struct ModelsExecution {
|
||||
client: AuthenticatedClient,
|
||||
profiles: bool,
|
||||
presentation: Option<ModelPresentation>,
|
||||
native: Option<Result<Vec<PublicModelDescriptor>, ProviderCatalogUnavailable>>,
|
||||
requested_versions: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl ModelsExecution {
|
||||
pub(super) fn new() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
profiles: false,
|
||||
presentation: None,
|
||||
native: None,
|
||||
requested_versions: Mutex::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn with_profiles() -> Arc<Self> {
|
||||
Self::with_presentation(
|
||||
ModelPresentation::new(
|
||||
Some("Grok 4.5".to_owned()),
|
||||
Some("xAI Grok 4.5 frontier model.".to_owned()),
|
||||
)
|
||||
.with_reasoning(
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
)
|
||||
.with_context_window_tokens(Some(500_000))
|
||||
.with_max_context_window_tokens(Some(900_000))
|
||||
.with_image_input(true)
|
||||
.with_agent_tools(true, true)
|
||||
.with_service_tiers(vec![
|
||||
ModelServiceTier::new(
|
||||
"priority",
|
||||
"Fast",
|
||||
"Route the request through Codex fast mode.",
|
||||
)
|
||||
.with_speed_tier("fast"),
|
||||
]),
|
||||
)
|
||||
}
|
||||
|
||||
fn with_presentation(presentation: ModelPresentation) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
profiles: true,
|
||||
presentation: Some(presentation),
|
||||
native: None,
|
||||
requested_versions: Mutex::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn with_cli_min() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
client: authenticated_client_with_min_versions("sk_models_test", None, Some("0.40.0")),
|
||||
profiles: false,
|
||||
presentation: None,
|
||||
native: None,
|
||||
requested_versions: Mutex::default(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -51,12 +87,49 @@ impl ModelsExecution {
|
||||
Some("26.825.51511"),
|
||||
None,
|
||||
),
|
||||
profiles: false,
|
||||
presentation: None,
|
||||
native: None,
|
||||
requested_versions: Mutex::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn profiles(&self) -> Vec<PublicModelProfile> {
|
||||
self.presentation
|
||||
.as_ref()
|
||||
.map(|presentation| {
|
||||
PublicModelProfile::new(
|
||||
PublicModelId::new("grok-4.5").expect("model"),
|
||||
presentation.clone(),
|
||||
)
|
||||
})
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionService for ModelsExecution {
|
||||
fn client_model_catalog<'a>(
|
||||
&'a self,
|
||||
_client: &'a AuthenticatedClient,
|
||||
protocol: &'a str,
|
||||
version: &'a str,
|
||||
) -> BoxFuture<'a, Result<Vec<PublicModelDescriptor>, ProviderCatalogUnavailable>> {
|
||||
Box::pin(async move {
|
||||
assert_eq!(protocol, "codex");
|
||||
self.requested_versions
|
||||
.lock()
|
||||
.expect("versions")
|
||||
.push(version.to_owned());
|
||||
self.native.clone().unwrap_or_else(|| {
|
||||
Ok(self
|
||||
.profiles()
|
||||
.into_iter()
|
||||
.map(PublicModelDescriptor::Adapted)
|
||||
.collect())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn authenticate(
|
||||
&self,
|
||||
plaintext: &str,
|
||||
@@ -75,36 +148,6 @@ impl ExecutionService for ModelsExecution {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn public_model_profiles(&self, _: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
if !self.profiles {
|
||||
return Vec::new();
|
||||
}
|
||||
vec![PublicModelProfile::new(
|
||||
PublicModelId::new("grok-4.5").expect("model"),
|
||||
ModelPresentation::new(
|
||||
Some("Grok 4.5".to_owned()),
|
||||
Some("xAI Grok 4.5 frontier model.".to_owned()),
|
||||
)
|
||||
.with_reasoning(
|
||||
Some("medium".to_owned()),
|
||||
["low", "medium", "high", "xhigh"]
|
||||
.map(str::to_owned)
|
||||
.to_vec(),
|
||||
)
|
||||
.with_context_window_tokens(Some(500_000))
|
||||
.with_image_input(true)
|
||||
.with_agent_tools(true, true)
|
||||
.with_service_tiers(vec![
|
||||
ModelServiceTier::new(
|
||||
"priority",
|
||||
"Fast",
|
||||
"Route the request through Codex fast mode.",
|
||||
)
|
||||
.with_speed_tier("fast"),
|
||||
]),
|
||||
)]
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
|
||||
matches!(model.as_str(), "model-a" | "model-b")
|
||||
}
|
||||
@@ -138,6 +181,102 @@ fn authorized_request(path: &str) -> Request<Body> {
|
||||
.expect("build models request")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_catalog_preserves_complete_objects_and_only_rewrites_alias_slug() {
|
||||
let original = serde_json::json!({
|
||||
"slug": "gpt-native", "display_name": "Official name", "priority": 47,
|
||||
"base_instructions": "Official instructions\nDo not replace.",
|
||||
"model_messages": {"token_budget": "upstream template", "future_section": [null, {"a": true}]},
|
||||
"service_tiers": [{"id": "priority", "name": "Fast", "description": "Official speed tier"}],
|
||||
"tool_mode": "code_mode_only", "use_responses_lite": true,
|
||||
"future_field": {"nested": [1, null, "value"]}, "explicit_null": null
|
||||
});
|
||||
let entry = |model: &str| PublicModelDescriptor::Native {
|
||||
model: PublicModelId::new(model).expect("model"),
|
||||
payload: gateway_core::operation::RawJsonPayload::new(
|
||||
"codex",
|
||||
serde_json::to_vec(&original).expect("JSON").into(),
|
||||
)
|
||||
.expect("payload"),
|
||||
};
|
||||
let execution = Arc::new(ModelsExecution {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
presentation: None,
|
||||
native: Some(Ok(vec![entry("gpt-native"), entry("my-alias")])),
|
||||
requested_versions: Mutex::default(),
|
||||
});
|
||||
let response = api_router(execution.clone())
|
||||
.await
|
||||
.oneshot(authorized_request(
|
||||
"/v1/models?client_version=0.154.0-alpha.3",
|
||||
))
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()["cache-control"], "private, no-store");
|
||||
assert!(response.headers().get("etag").is_none());
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("JSON");
|
||||
assert_eq!(value["models"][0], original);
|
||||
let mut alias = original;
|
||||
alias["slug"] = serde_json::json!("my-alias");
|
||||
assert_eq!(value["models"][1], alias);
|
||||
assert_eq!(
|
||||
*execution.requested_versions.lock().expect("versions"),
|
||||
["0.154.0-alpha.3"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_catalog_failure_does_not_publish_a_synthetic_or_empty_success() {
|
||||
let execution = Arc::new(ModelsExecution {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
presentation: None,
|
||||
native: Some(Err(ProviderCatalogUnavailable)),
|
||||
requested_versions: Mutex::default(),
|
||||
});
|
||||
let response = api_router(execution)
|
||||
.await
|
||||
.oneshot(authorized_request("/v1/models?client_version=0.154.0"))
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("JSON");
|
||||
assert_eq!(value["error"]["code"], "model_catalog_unavailable");
|
||||
assert!(value.get("models").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn catalog_rejects_unbounded_or_control_character_versions_before_fetching() {
|
||||
for version in [
|
||||
"x".repeat(65),
|
||||
"0.154.0%0A".to_owned(),
|
||||
"0.154.0%20".to_owned(),
|
||||
] {
|
||||
let execution = ModelsExecution::new();
|
||||
let response = api_router(execution.clone())
|
||||
.await
|
||||
.oneshot(authorized_request(&format!(
|
||||
"/v1/models?client_version={version}"
|
||||
)))
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert!(
|
||||
execution
|
||||
.requested_versions
|
||||
.lock()
|
||||
.expect("versions")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_should_encode_the_service_visible_catalog() {
|
||||
let response = api_router(ModelsExecution::new())
|
||||
@@ -289,6 +428,7 @@ async fn models_should_encode_provider_profiles_for_current_codex_clients() {
|
||||
assert_eq!(model["slug"], "grok-4.5");
|
||||
assert_eq!(model["default_reasoning_level"], "medium");
|
||||
assert_eq!(model["context_window"], 500_000);
|
||||
assert_eq!(model["max_context_window"], 900_000);
|
||||
assert_eq!(model["apply_patch_tool_type"], "freeform");
|
||||
assert_eq!(model["additional_speed_tiers"], serde_json::json!(["fast"]));
|
||||
assert_eq!(
|
||||
@@ -322,6 +462,34 @@ async fn models_should_encode_provider_profiles_for_current_codex_clients() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_should_preserve_unknown_context_windows() {
|
||||
for (context_window, max_context_window) in
|
||||
[(None, None), (Some(272_000), None), (None, Some(872_000))]
|
||||
{
|
||||
let presentation = ModelPresentation::default()
|
||||
.with_context_window_tokens(context_window)
|
||||
.with_max_context_window_tokens(max_context_window);
|
||||
let response = api_router(ModelsExecution::with_presentation(presentation))
|
||||
.await
|
||||
.oneshot(authorized_request("/v1/models?client_version=0.145.0"))
|
||||
.await
|
||||
.expect("list Codex models response");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("read Codex models body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("models JSON");
|
||||
let model = &value["models"][0];
|
||||
|
||||
assert_eq!(model["context_window"], serde_json::json!(context_window));
|
||||
assert_eq!(
|
||||
model["max_context_window"],
|
||||
serde_json::json!(max_context_window)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_should_keep_the_codex_contract_when_profiles_are_empty() {
|
||||
let response = api_router(ModelsExecution::new())
|
||||
|
||||
@@ -22,7 +22,7 @@ use gateway_core::engine::execution::{
|
||||
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
|
||||
StartExecution, StartProviderExecution, StartedExecution,
|
||||
};
|
||||
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
|
||||
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, ModelRequestId};
|
||||
use gateway_core::error::{
|
||||
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, GatewayError, GatewayErrorKind,
|
||||
ProviderError, ProviderErrorKind,
|
||||
@@ -33,7 +33,7 @@ use gateway_core::event::{
|
||||
};
|
||||
use gateway_core::operation::{Operation, OperationKind};
|
||||
use gateway_core::routing::PublicModelId;
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_core::upstream::{OpaqueUpstreamValue, UpstreamSendState};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use gateway_api::openai::responses::{collect_execution_response, stream_execution_response};
|
||||
@@ -1819,3 +1819,223 @@ fn response_status_from_body(body: &[u8]) -> Option<String> {
|
||||
let value: serde_json::Value = serde_json::from_slice(body).ok()?;
|
||||
value.pointer("/error/code")?.as_str().map(str::to_owned)
|
||||
}
|
||||
|
||||
const MODEL_REQUEST_ID: &str = "req_model_correlation";
|
||||
|
||||
struct SessionExecution {
|
||||
client: AuthenticatedClient,
|
||||
session: Mutex<Option<Box<dyn ExecutionSession>>>,
|
||||
}
|
||||
|
||||
impl ExecutionService for SessionExecution {
|
||||
fn authenticate(&self, _: &str) -> Result<AuthenticatedClient, ClientAuthenticationError> {
|
||||
Ok(self.client.clone())
|
||||
}
|
||||
|
||||
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, _: &AuthenticatedClient, _: &PublicModelId) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn start(
|
||||
&self,
|
||||
request: StartExecution,
|
||||
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
|
||||
Box::pin(async move {
|
||||
Ok(StartedExecution {
|
||||
request_id: ModelRequestId::new(MODEL_REQUEST_ID).unwrap(),
|
||||
created_at: std::time::SystemTime::now(),
|
||||
stream: request.metadata.stream,
|
||||
session: self.session.lock().unwrap().take().expect("one execution"),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn start_provider_endpoint(
|
||||
&self,
|
||||
_: StartProviderExecution,
|
||||
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
|
||||
Box::pin(async { panic!("Responses must use model execution") })
|
||||
}
|
||||
}
|
||||
|
||||
async fn response(
|
||||
provider: &str,
|
||||
session: FakeSession,
|
||||
streaming: bool,
|
||||
) -> axum::response::Response {
|
||||
let execution = Arc::new(SessionExecution {
|
||||
client: authenticated_client_for_provider("sk_correlation_test", provider),
|
||||
session: Mutex::new(Some(Box::new(session))),
|
||||
});
|
||||
api_router(execution)
|
||||
.await
|
||||
.oneshot(
|
||||
Request::post("/v1/responses")
|
||||
.header(AUTHORIZATION, "Bearer sk_correlation_test")
|
||||
.header("x-request-id", "caller-chosen-not-a-model-id")
|
||||
.header("user-agent", "AnotherTerminal/1.0")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "model-a", "input": "synthetic", "stream": streaming
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn headers(values: &[(&str, &'static str)]) -> Vec<ProviderResponseHeader> {
|
||||
values
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
ProviderResponseHeader::new(*name, Bytes::from_static(value.as_bytes()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn failed_session(streaming: bool, error: EngineError) -> FakeSession {
|
||||
let trace = Arc::new(Trace::default());
|
||||
if streaming {
|
||||
FakeSession::streaming(trace, vec![NextStep::Error(error)])
|
||||
} else {
|
||||
FakeSession::buffered(trace, Vec::new()).with_collect_error(error)
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_ids_are_provider_and_client_independent() {
|
||||
let cases: &[(&[(&str, &str)], &str)] = &[
|
||||
(&[], MODEL_REQUEST_ID),
|
||||
(&[("x-request-id", "upstream-primary")], "upstream-primary"),
|
||||
(&[("x-oai-request-id", "upstream-alias")], "upstream-alias"),
|
||||
(
|
||||
&[
|
||||
("x-request-id", "upstream-primary"),
|
||||
("x-oai-request-id", "upstream-alias"),
|
||||
],
|
||||
"upstream-primary",
|
||||
),
|
||||
(
|
||||
&[
|
||||
("x-request-id", " "),
|
||||
("x-oai-request-id", "upstream-alias"),
|
||||
],
|
||||
"upstream-alias",
|
||||
),
|
||||
(&[("x-request-id", "")], MODEL_REQUEST_ID),
|
||||
];
|
||||
for provider in ["openai", "xai"] {
|
||||
for streaming in [false, true] {
|
||||
for (upstream_headers, expected) in cases {
|
||||
let trace = Arc::new(Trace::default());
|
||||
let session = if streaming {
|
||||
FakeSession::streaming(
|
||||
trace,
|
||||
vec![
|
||||
NextStep::Event(delivery(
|
||||
started(),
|
||||
CommitRequirement::CommitBeforeDelivery,
|
||||
)),
|
||||
NextStep::Event(delivery(
|
||||
completed(),
|
||||
CommitRequirement::AlreadyCommitted,
|
||||
)),
|
||||
NextStep::FinalizeSuccess,
|
||||
],
|
||||
)
|
||||
} else {
|
||||
FakeSession::buffered(trace, vec![started(), completed()])
|
||||
}
|
||||
.with_response_headers(headers(upstream_headers));
|
||||
let response = response(provider, session, streaming).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()["x-request-id"], *expected);
|
||||
assert_eq!(response.headers()["x-gateway-request-id"], MODEL_REQUEST_ID);
|
||||
assert!(
|
||||
!to_bytes(response.into_body(), 64 * 1024)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failure_ids_prefer_the_actual_error_over_session_headers() {
|
||||
for provider in ["openai", "xai"] {
|
||||
for streaming in [false, true] {
|
||||
for upstream_headers in [vec![], headers(&[("x-oai-request-id", "actual-error")])] {
|
||||
let expected = if upstream_headers.is_empty() {
|
||||
MODEL_REQUEST_ID
|
||||
} else {
|
||||
"actual-error"
|
||||
};
|
||||
let raw =
|
||||
Bytes::from_static(br#"{"error":{"message":"synthetic upstream detail"}}"#);
|
||||
let error =
|
||||
ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
|
||||
.with_client_visible_upstream_response(
|
||||
ClientVisibleUpstreamResponse::new(
|
||||
422,
|
||||
Some(b"application/json".to_vec()),
|
||||
raw.clone(),
|
||||
)
|
||||
.with_headers(upstream_headers),
|
||||
);
|
||||
let session = failed_session(streaming, EngineError::Provider(error))
|
||||
.with_response_headers(headers(&[("x-request-id", "earlier-session")]));
|
||||
let response = response(provider, session, streaming).await;
|
||||
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
|
||||
assert_eq!(response.headers()["x-request-id"], expected);
|
||||
assert_eq!(response.headers()["x-gateway-request-id"], MODEL_REQUEST_ID);
|
||||
assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), raw);
|
||||
}
|
||||
|
||||
let error =
|
||||
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
|
||||
.with_status(503)
|
||||
.with_upstream_request_id(OpaqueUpstreamValue::new("body-read-error"));
|
||||
let session = failed_session(streaming, EngineError::Provider(error))
|
||||
.with_response_headers(headers(&[
|
||||
("x-request-id", "earlier-session"),
|
||||
("x-oai-request-id", "earlier-alias"),
|
||||
]));
|
||||
let response = response(provider, session, streaming).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
|
||||
assert_eq!(response.headers()["x-request-id"], "body-read-error");
|
||||
assert_eq!(response.headers().get_all("x-request-id").iter().count(), 1);
|
||||
assert!(response.headers().get("x-oai-request-id").is_none());
|
||||
assert_eq!(response.headers()["x-gateway-request-id"], MODEL_REQUEST_ID);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_failure_uses_the_existing_model_id() {
|
||||
for provider in ["openai", "xai"] {
|
||||
for streaming in [false, true] {
|
||||
let session =
|
||||
failed_session(streaming, EngineError::Deadline).with_response_headers(headers(&[
|
||||
("x-request-id", "opening-primary"),
|
||||
("x-oai-request-id", "opening-alias"),
|
||||
("x-codex-turn-state", "retained-turn-state"),
|
||||
]));
|
||||
let response = response(provider, session, streaming).await;
|
||||
assert_eq!(response.headers()["x-request-id"], MODEL_REQUEST_ID);
|
||||
assert_eq!(response.headers()["x-gateway-request-id"], MODEL_REQUEST_ID);
|
||||
assert!(response.headers().get("x-oai-request-id").is_none());
|
||||
assert_eq!(
|
||||
response.headers()["x-codex-turn-state"],
|
||||
"retained-turn-state"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,449 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex, atomic::Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::http::{StatusCode, header::AUTHORIZATION};
|
||||
use bytes::Bytes;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
|
||||
use gateway_core::error::{
|
||||
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ProviderError, ProviderErrorKind,
|
||||
};
|
||||
use gateway_core::event::{ProtocolWireEvent, ProviderEvent, ProviderResponseHeader};
|
||||
use gateway_core::upstream::{OpaqueUpstreamValue, UpstreamSendState};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_tungstenite::{
|
||||
MaybeTlsStream, WebSocketStream, connect_async,
|
||||
tungstenite::{Message, client::IntoClientRequest},
|
||||
};
|
||||
|
||||
use super::{AtomicFailureExecution, AtomicFailureTrace};
|
||||
use crate::openai::{api_router, authenticated_client};
|
||||
|
||||
type TestSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
|
||||
// 全部使用合成响应事实和真实 API 路由,不读取账号、日志或请求转储。
|
||||
struct TestServer(tokio::task::JoinHandle<()>);
|
||||
|
||||
impl Drop for TestServer {
|
||||
fn drop(&mut self) {
|
||||
self.0.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect(
|
||||
trace: Arc<AtomicFailureTrace>,
|
||||
response_headers: Vec<ProviderResponseHeader>,
|
||||
) -> (TestSocket, TestServer) {
|
||||
let execution = Arc::new(AtomicFailureExecution {
|
||||
client: authenticated_client("sk_ws_atomic"),
|
||||
trace,
|
||||
response_headers,
|
||||
fail_before_first_event: false,
|
||||
});
|
||||
let app = api_router(execution).await;
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind local API");
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = TestServer(tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.expect("serve local API");
|
||||
}));
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.unwrap();
|
||||
request
|
||||
.headers_mut()
|
||||
.insert(AUTHORIZATION, "Bearer sk_ws_atomic".parse().unwrap());
|
||||
let (socket, response) = connect_async(request).await.expect("upgrade WebSocket");
|
||||
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
|
||||
(socket, server)
|
||||
}
|
||||
|
||||
async fn send_request(socket: &mut TestSocket) {
|
||||
socket
|
||||
.send(Message::Text(
|
||||
json!({"type": "response.create", "model": "model-a", "input": "synthetic"})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn next_event(socket: &mut TestSocket) -> Value {
|
||||
let frame = tokio::time::timeout(Duration::from_secs(2), socket.next())
|
||||
.await
|
||||
.expect("event timeout")
|
||||
.expect("connection remains open")
|
||||
.expect("valid frame");
|
||||
let Message::Text(text) = frame else {
|
||||
panic!("expected a JSON text frame, got {frame:?}");
|
||||
};
|
||||
serde_json::from_str(&text).expect("JSON event")
|
||||
}
|
||||
|
||||
async fn next_error(socket: &mut TestSocket) -> Value {
|
||||
let first = next_event(socket).await;
|
||||
let error = if first["type"] == "response.metadata" {
|
||||
next_event(socket).await
|
||||
} else {
|
||||
first
|
||||
};
|
||||
assert_eq!(error["type"], "error");
|
||||
error
|
||||
}
|
||||
|
||||
async fn assert_no_duplicate_failure(socket: &mut TestSocket) {
|
||||
// 后续入站帧只会在本轮收敛后处理,比固定等待窗口更可靠地检查没有重复失败终态。
|
||||
socket.send(Message::Text("{".into())).await.unwrap();
|
||||
let error = next_event(socket).await;
|
||||
assert_eq!(error["error"]["code"], "invalid_json");
|
||||
assert_eq!(error["status"], 400);
|
||||
assert_eq!(error["headers"]["x-request-id"], error["request_id"]);
|
||||
assert_eq!(
|
||||
error["headers"]["x-gateway-request-id"],
|
||||
error["request_id"]
|
||||
);
|
||||
}
|
||||
|
||||
async fn initial_error(error: EngineError, response_headers: Vec<ProviderResponseHeader>) -> Value {
|
||||
let trace = Arc::new(AtomicFailureTrace {
|
||||
initial_errors: Mutex::new(VecDeque::from([error])),
|
||||
..AtomicFailureTrace::default()
|
||||
});
|
||||
let (mut socket, _server) = connect(Arc::clone(&trace), response_headers).await;
|
||||
send_request(&mut socket).await;
|
||||
let error = next_error(&mut socket).await;
|
||||
assert!(!trace.committed.load(Ordering::Acquire));
|
||||
assert!(trace.finalized.load(Ordering::Acquire));
|
||||
assert_eq!(error["request_id"], "req_ws_atomic");
|
||||
assert_eq!(error["headers"]["x-gateway-request-id"], "req_ws_atomic");
|
||||
assert_no_duplicate_failure(&mut socket).await;
|
||||
socket.close(None).await.unwrap();
|
||||
error
|
||||
}
|
||||
|
||||
fn header(name: &str, value: &'static [u8]) -> ProviderResponseHeader {
|
||||
ProviderResponseHeader::new(name, Bytes::from_static(value))
|
||||
}
|
||||
|
||||
fn upstream_failure(
|
||||
status: u16,
|
||||
body: &'static [u8],
|
||||
headers: Vec<ProviderResponseHeader>,
|
||||
) -> ProviderError {
|
||||
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent)
|
||||
.with_status(status)
|
||||
.with_client_visible_upstream_response(
|
||||
ClientVisibleUpstreamResponse::new(status, None, Bytes::from_static(body))
|
||||
.with_headers(headers),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_http_or_opening_failure_preserves_status_and_structured_error() {
|
||||
// 分类故意都用 Unavailable:不能把收到的 401/429/500 等按网关分类改成 502。
|
||||
for status in [302, 400, 401, 403, 429, 500, 502, 503] {
|
||||
let provider = upstream_failure(
|
||||
status,
|
||||
br#"{"error":{"message":"synthetic upstream marker","code":"custom_code","type":"custom_type"}}"#,
|
||||
vec![header("x-request-id", b"req_current"), header("retry-after", b"7")],
|
||||
)
|
||||
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
|
||||
"synthetic upstream marker",
|
||||
Some("custom_code".to_owned()),
|
||||
Some("custom_type".to_owned()),
|
||||
));
|
||||
let error = initial_error(EngineError::Provider(provider), Vec::new()).await;
|
||||
assert_eq!(error["status"], status);
|
||||
assert_eq!(
|
||||
error["error"],
|
||||
json!({
|
||||
"message": "synthetic upstream marker", "code": "custom_code", "type": "custom_type",
|
||||
})
|
||||
);
|
||||
assert_eq!(error["headers"]["x-request-id"], "req_current");
|
||||
assert_eq!(error["headers"]["retry-after"], "7");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_structured_bodies_use_safe_fallback_without_losing_status_or_ids() {
|
||||
for body in [
|
||||
b"<html>synthetic-private-body-marker</html>".as_slice(),
|
||||
b"synthetic-private-body-marker",
|
||||
b"",
|
||||
br#""synthetic-private-body-marker""#,
|
||||
br#"{"error":"synthetic-private-body-marker"}"#,
|
||||
br#"{"detail":"synthetic-private-body-marker"}"#,
|
||||
br#"{"error":{"code":"synthetic-private-body-marker"}}"#,
|
||||
] {
|
||||
let provider = upstream_failure(
|
||||
503,
|
||||
body,
|
||||
vec![header("x-request-id", b"req_body_fallback")],
|
||||
);
|
||||
let error = initial_error(EngineError::Provider(provider), Vec::new()).await;
|
||||
assert_eq!(error["status"], 503);
|
||||
assert_eq!(error["headers"]["x-request-id"], "req_body_fallback");
|
||||
assert_eq!(
|
||||
error["error"],
|
||||
json!({
|
||||
"message": "upstream service is unavailable",
|
||||
"code": "upstream_unavailable",
|
||||
"type": "server_error",
|
||||
})
|
||||
);
|
||||
assert!(!error.to_string().contains("synthetic-private-body-marker"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn final_failure_headers_take_precedence_over_observed_opening_headers() {
|
||||
let provider = upstream_failure(403, b"", vec![header("x-oai-request-id", b"req_final")])
|
||||
.with_status(502)
|
||||
.with_upstream_request_id(OpaqueUpstreamValue::new("req_stale_fact"));
|
||||
let error = initial_error(
|
||||
EngineError::Provider(provider),
|
||||
vec![
|
||||
header("x-request-id", b"req_old_opening"),
|
||||
header("x-codex-turn-state", b"old-turn-state"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(error["status"], 403);
|
||||
assert_eq!(error["headers"]["x-oai-request-id"], "req_final");
|
||||
assert!(error["headers"].get("x-request-id").is_none());
|
||||
assert!(!error.to_string().contains("req_old_opening"));
|
||||
assert!(!error.to_string().contains("req_stale_fact"));
|
||||
assert!(!error.to_string().contains("old-turn-state"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_failure_facts_survive_without_a_complete_http_response() {
|
||||
let provider = ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
|
||||
.with_status(429)
|
||||
.with_upstream_request_id(OpaqueUpstreamValue::new("req_body_read_failed"));
|
||||
let error = initial_error(
|
||||
EngineError::Provider(provider),
|
||||
vec![
|
||||
header("x-request-id", b"req_old"),
|
||||
header("x-oai-request-id", b"req_old_alias"),
|
||||
header("x-codex-turn-state", b"observed-turn-state"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(error["status"], 429);
|
||||
assert_eq!(error["headers"]["x-request-id"], "req_body_read_failed");
|
||||
assert_eq!(
|
||||
error["headers"]["x-codex-turn-state"],
|
||||
"observed-turn-state"
|
||||
);
|
||||
assert!(error["headers"].get("x-oai-request-id").is_none());
|
||||
assert_eq!(error["error"]["message"], "upstream service is unavailable");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_error_headers_apply_the_response_security_boundary() {
|
||||
let provider = upstream_failure(
|
||||
502,
|
||||
b"",
|
||||
vec![
|
||||
header("X-Request-Id", b"req_safe"),
|
||||
header("x-gateway-request-id", b"upstream-cannot-claim-gateway-id"),
|
||||
header("x-future-header", b"first"),
|
||||
header("x-future-header", b"second"),
|
||||
header("cf-ray", b"synthetic-ray"),
|
||||
header("retry-after", b"11"),
|
||||
header("authorization", b"synthetic-private-header-marker"),
|
||||
header("set-cookie", b"synthetic-private-header-marker"),
|
||||
header("chatgpt-account-id", b"synthetic-private-header-marker"),
|
||||
header(
|
||||
"x-codex-installation-id",
|
||||
b"synthetic-private-header-marker",
|
||||
),
|
||||
header("connection", b"x-private-hop"),
|
||||
header("x-private-hop", b"synthetic-private-header-marker"),
|
||||
header("sec-websocket-accept", b"synthetic-private-header-marker"),
|
||||
header("content-length", b"123"),
|
||||
header("content-type", b"text/html"),
|
||||
header("bad\0name", b"synthetic-private-header-marker"),
|
||||
header("x-bad-value", b"synthetic-private-header-marker\r\n"),
|
||||
header("x-non-utf8", b"\xff"),
|
||||
],
|
||||
);
|
||||
let error = initial_error(EngineError::Provider(provider), Vec::new()).await;
|
||||
assert_eq!(
|
||||
error["headers"],
|
||||
json!({
|
||||
"x-request-id": "req_safe",
|
||||
"x-gateway-request-id": "req_ws_atomic",
|
||||
"x-future-header": "second",
|
||||
"cf-ray": "synthetic-ray",
|
||||
"retry-after": "11",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_error_without_upstream_facts_keeps_the_local_contract() {
|
||||
for (error, status, code) in [
|
||||
(EngineError::Deadline, 504, "request_timeout"),
|
||||
(EngineError::EmptyRoutingPlan, 503, "no_available_provider"),
|
||||
(
|
||||
EngineError::Provider(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
UpstreamSendState::NotSent,
|
||||
)),
|
||||
502,
|
||||
"upstream_unavailable",
|
||||
),
|
||||
] {
|
||||
let error = initial_error(error, Vec::new()).await;
|
||||
assert_eq!(error["status"], status);
|
||||
assert_eq!(error["error"]["code"], code);
|
||||
assert_eq!(error["headers"]["x-request-id"], "req_ws_atomic");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_error_never_uses_a_success_or_upgrade_status_as_an_error_status() {
|
||||
for status in [0, 99, 101, 200, 204] {
|
||||
let error = initial_error(
|
||||
EngineError::Provider(upstream_failure(status, b"", Vec::new())),
|
||||
Vec::new(),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(error["status"], 502);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn opening_ids_stay_in_metadata_but_do_not_identify_initial_failures() {
|
||||
for (failure, status) in [
|
||||
(EngineError::Deadline, 504),
|
||||
(
|
||||
EngineError::Provider(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
UpstreamSendState::Ambiguous,
|
||||
)),
|
||||
502,
|
||||
),
|
||||
] {
|
||||
let trace = Arc::new(AtomicFailureTrace {
|
||||
initial_errors: Mutex::new(VecDeque::from([failure])),
|
||||
..AtomicFailureTrace::default()
|
||||
});
|
||||
let (mut socket, _server) = connect(
|
||||
trace,
|
||||
vec![
|
||||
header("x-request-id", b"req_old_opening"),
|
||||
header("x-oai-request-id", b"req_observed"),
|
||||
header("x-codex-turn-state", b"observed-turn-state"),
|
||||
header("x-future-header", b"future-value"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
send_request(&mut socket).await;
|
||||
// metadata 仍交付原会话观测;opening ID 不能因此成为当前失败的请求身份。
|
||||
assert_eq!(
|
||||
next_event(&mut socket).await,
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": {
|
||||
"x-request-id": "req_old_opening",
|
||||
"x-oai-request-id": "req_observed",
|
||||
"x-codex-turn-state": "observed-turn-state",
|
||||
"x-future-header": "future-value",
|
||||
},
|
||||
})
|
||||
);
|
||||
let error = next_event(&mut socket).await;
|
||||
assert_eq!(error["type"], "error");
|
||||
assert_eq!(error["status"], status);
|
||||
assert_eq!(error["request_id"], "req_ws_atomic");
|
||||
assert_eq!(
|
||||
error["headers"],
|
||||
json!({
|
||||
"x-request-id": "req_ws_atomic",
|
||||
"x-gateway-request-id": "req_ws_atomic",
|
||||
"x-codex-turn-state": "observed-turn-state",
|
||||
"x-future-header": "future-value",
|
||||
})
|
||||
);
|
||||
assert_no_duplicate_failure(&mut socket).await;
|
||||
socket.close(None).await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn errors_on_a_reused_connection_keep_their_own_upstream_ids() {
|
||||
let errors = ["req_first", "req_second"].map(|request_id| {
|
||||
EngineError::Provider(
|
||||
upstream_failure(502, b"", Vec::new())
|
||||
.with_upstream_request_id(OpaqueUpstreamValue::new(request_id)),
|
||||
)
|
||||
});
|
||||
let trace = Arc::new(AtomicFailureTrace {
|
||||
initial_errors: Mutex::new(VecDeque::from(errors)),
|
||||
..AtomicFailureTrace::default()
|
||||
});
|
||||
let (mut socket, _server) = connect(trace, Vec::new()).await;
|
||||
for request_id in ["req_first", "req_second"] {
|
||||
send_request(&mut socket).await;
|
||||
let error = next_error(&mut socket).await;
|
||||
assert_eq!(error["headers"]["x-request-id"], request_id);
|
||||
assert_eq!(error["headers"]["x-gateway-request-id"], "req_ws_atomic");
|
||||
}
|
||||
assert_no_duplicate_failure(&mut socket).await;
|
||||
socket.close(None).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deliverable_business_failures_are_not_rewritten_or_followed_by_a_second_error() {
|
||||
for raw in [
|
||||
json!({
|
||||
"type": "error",
|
||||
"status": 429,
|
||||
"headers": {"x-request-id": "req_raw"},
|
||||
"error": {"message": "raw marker", "type": "custom_type", "code": "custom_code"},
|
||||
"future": {"keep": true},
|
||||
}),
|
||||
json!({
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
"id": "resp_raw",
|
||||
"status": "failed",
|
||||
"error": {"message": "raw marker", "code": "custom_code"},
|
||||
},
|
||||
"future": {"keep": true},
|
||||
}),
|
||||
] {
|
||||
let event = ProviderEvent::wire(
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some(raw["type"].as_str().unwrap().to_owned()),
|
||||
raw.clone(),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
let trace = Arc::new(AtomicFailureTrace {
|
||||
first_batch: Mutex::new(Some(
|
||||
CoordinatedEvent::try_batch(vec![event], CommitRequirement::CommitBeforeDelivery)
|
||||
.unwrap(),
|
||||
)),
|
||||
..AtomicFailureTrace::default()
|
||||
});
|
||||
let (mut socket, _server) = connect(Arc::clone(&trace), Vec::new()).await;
|
||||
send_request(&mut socket).await;
|
||||
assert_eq!(next_event(&mut socket).await["type"], "response.metadata");
|
||||
assert_eq!(next_event(&mut socket).await, raw);
|
||||
assert_no_duplicate_failure(&mut socket).await;
|
||||
assert!(trace.committed.load(Ordering::Acquire));
|
||||
assert!(trace.finalized.load(Ordering::Acquire));
|
||||
assert_eq!(trace.next_calls.load(Ordering::Acquire), 2);
|
||||
socket.close(None).await.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
mod connection;
|
||||
mod forward;
|
||||
mod protocol;
|
||||
|
||||
use std::sync::{
|
||||
@@ -7,6 +8,7 @@ use std::sync::{
|
||||
};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{HeaderMap, HeaderValue, Request, StatusCode, header::AUTHORIZATION},
|
||||
@@ -17,18 +19,40 @@ use gateway_api::openai::responses::{
|
||||
DecodedResponsesRequest, OpenAiRequestHeaders, ResponseCreateFrameError,
|
||||
decode_response_create_with_context,
|
||||
};
|
||||
use gateway_core::account::ProviderAccountId;
|
||||
use gateway_core::engine::admission::{
|
||||
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
|
||||
ClientAdmissionRequest, ClientAdmissionRestoreResult,
|
||||
};
|
||||
use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
|
||||
use gateway_core::engine::execution::{
|
||||
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
|
||||
AuthenticatedClient, ClientAuthenticationError, DefaultExecutionService, ExecutionService,
|
||||
ExecutionSession, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
|
||||
StartExecution, StartProviderExecution, StartedExecution,
|
||||
};
|
||||
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
|
||||
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind};
|
||||
use gateway_core::engine::provider::{
|
||||
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AttemptContext, AttemptRecord, CommitRequirement, CoordinatedEvent, EngineError,
|
||||
ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
|
||||
RecoveryReport,
|
||||
};
|
||||
use gateway_core::error::{GatewayError, ProviderError, ProviderErrorKind, StoreError};
|
||||
use gateway_core::event::{
|
||||
GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta,
|
||||
};
|
||||
use gateway_core::lifecycle::{
|
||||
CancellationToken, ConnectionDraining, ConnectionGuard, ConnectionLifecycle,
|
||||
};
|
||||
use gateway_core::metering::{Decimal, ProviderReportedCost};
|
||||
use gateway_core::operation::Operation;
|
||||
use gateway_core::routing::PublicModelId;
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{
|
||||
ProviderCatalogGeneration, ProviderKind, ProviderModelCapabilities, PublicModelId,
|
||||
};
|
||||
use gateway_core::runtime::RuntimeSnapshotHandle;
|
||||
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
|
||||
use gateway_protocol::openai::codex_responses_request_semantics;
|
||||
use serde_json::{Value, json};
|
||||
use tokio_tungstenite::connect_async;
|
||||
@@ -37,7 +61,10 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use super::decode_response_create;
|
||||
use crate::openai::{api_router, authenticated_client, models::ModelsExecution};
|
||||
use crate::openai::{
|
||||
EmptyWorkerHealth, IgnoredClientApiKeyUsage, UnusedContinuation, api_router,
|
||||
authenticated_client, models::ModelsExecution, snapshot,
|
||||
};
|
||||
|
||||
fn decode_response_create_with_turn_header(
|
||||
payload: Value,
|
||||
@@ -281,6 +308,8 @@ struct AtomicFailureTrace {
|
||||
committed: AtomicBool,
|
||||
finalized: AtomicBool,
|
||||
downstream_websocket_connection_ids: Mutex<Vec<String>>,
|
||||
initial_errors: Mutex<std::collections::VecDeque<EngineError>>,
|
||||
first_batch: Mutex<Option<CoordinatedEvent>>,
|
||||
}
|
||||
|
||||
struct AtomicFailureSession {
|
||||
@@ -293,6 +322,10 @@ impl ExecutionSession for AtomicFailureSession {
|
||||
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
|
||||
Box::pin(async move {
|
||||
let next_call = self.trace.next_calls.fetch_add(1, Ordering::AcqRel);
|
||||
if let Some(error) = self.trace.initial_errors.lock().unwrap().pop_front() {
|
||||
self.trace.finalized.store(true, Ordering::Release);
|
||||
return Err(error);
|
||||
}
|
||||
if self.fail_before_first_event && next_call == 0 {
|
||||
return Err(EngineError::Provider(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
@@ -300,7 +333,14 @@ impl ExecutionSession for AtomicFailureSession {
|
||||
)));
|
||||
}
|
||||
match next_call {
|
||||
0 => Ok(Some(atomic_failure_batch())),
|
||||
0 => Ok(Some(
|
||||
self.trace
|
||||
.first_batch
|
||||
.lock()
|
||||
.unwrap()
|
||||
.take()
|
||||
.unwrap_or_else(atomic_failure_batch),
|
||||
)),
|
||||
1 => {
|
||||
self.trace.finalized.store(true, Ordering::Release);
|
||||
Err(EngineError::Provider(ProviderError::new(
|
||||
@@ -811,6 +851,352 @@ async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit(
|
||||
|
||||
const TEST_WEBSOCKET_KEY: &str = "AAAAAAAAAAAAAAAAAAAAAA==";
|
||||
|
||||
#[derive(Default)]
|
||||
struct SettlementPorts {
|
||||
request_id: Mutex<Option<ModelRequestId>>,
|
||||
finalizations: Mutex<Vec<ModelRequestFinalization>>,
|
||||
active: AtomicBool,
|
||||
settlement_calls: AtomicUsize,
|
||||
settlement_started: tokio::sync::Notify,
|
||||
settlement_gate: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
|
||||
charges: Mutex<Vec<ClientBudgetCharge>>,
|
||||
releases: AtomicUsize,
|
||||
released: tokio::sync::Notify,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ExecutionStore for SettlementPorts {
|
||||
async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> {
|
||||
*self.request_id.lock().unwrap() = Some(request.id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mark_send_state(
|
||||
&self,
|
||||
_: &ModelRequestId,
|
||||
_: UpstreamSendState,
|
||||
) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mark_downstream_committed(
|
||||
&self,
|
||||
_: &ModelRequestId,
|
||||
_: SystemTime,
|
||||
_: Option<u16>,
|
||||
) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_model_request(
|
||||
&self,
|
||||
finalization: ModelRequestFinalization,
|
||||
) -> Result<(), StoreError> {
|
||||
self.finalizations.lock().unwrap().push(finalization);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn recover_expired(&self, _: SystemTime) -> Result<RecoveryReport, StoreError> {
|
||||
unreachable!("the WebSocket test does not recover requests")
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientAdmissionPort for SettlementPorts {
|
||||
fn admit(
|
||||
&self,
|
||||
_: ClientAdmissionRequest,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
|
||||
Box::pin(async {
|
||||
assert!(!self.active.swap(true, Ordering::SeqCst));
|
||||
Ok(ClientAdmissionDecision::Granted)
|
||||
})
|
||||
}
|
||||
|
||||
fn release<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
request_id: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
|
||||
Box::pin(async move {
|
||||
assert_eq!(
|
||||
self.charges.lock().unwrap().len(),
|
||||
1,
|
||||
"settlement must finish first"
|
||||
);
|
||||
assert_eq!(self.request_id.lock().unwrap().as_ref(), Some(request_id));
|
||||
self.releases.fetch_add(1, Ordering::SeqCst);
|
||||
let active = self.active.swap(false, Ordering::SeqCst);
|
||||
self.released.notify_one();
|
||||
Ok(active)
|
||||
})
|
||||
}
|
||||
|
||||
fn restore(
|
||||
&self,
|
||||
_: ClientAdmissionRecovery,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
|
||||
Box::pin(async { unreachable!("the WebSocket test does not restore admission") })
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for SettlementPorts {
|
||||
fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
|
||||
Box::pin(async move {
|
||||
self.settlement_calls.fetch_add(1, Ordering::SeqCst);
|
||||
assert!(self.active.load(Ordering::SeqCst));
|
||||
let gate = self
|
||||
.settlement_gate
|
||||
.lock()
|
||||
.unwrap()
|
||||
.take()
|
||||
.expect("one settlement");
|
||||
self.settlement_started.notify_one();
|
||||
gate.await
|
||||
.expect("release settlement after client disconnect");
|
||||
assert!(self.active.load(Ordering::SeqCst));
|
||||
self.charges.lock().unwrap().push(charge);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCircuitPort for SettlementPorts {
|
||||
fn decision<'a>(
|
||||
&'a self,
|
||||
_: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
|
||||
Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
|
||||
}
|
||||
|
||||
fn observe_failure<'a>(
|
||||
&'a self,
|
||||
_: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn observe_success<'a>(
|
||||
&'a self,
|
||||
_: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
struct ChargedWebSocketProvider;
|
||||
|
||||
fn websocket_charge() -> Decimal {
|
||||
"1.25".parse().unwrap()
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for ChargedWebSocketProvider {
|
||||
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 candidate = request.candidate();
|
||||
let metadata = ProviderCallMetadata::new(
|
||||
candidate.provider().clone(),
|
||||
candidate.upstream_model().unwrap().clone(),
|
||||
ProviderAccountId::new("acct_api_test").unwrap(),
|
||||
UpstreamTransport::new("websocket").unwrap(),
|
||||
);
|
||||
let response = ResponseMeta::new("resp_ws_settlement", "model-a");
|
||||
let events = [
|
||||
ProviderEvent::canonical_with_wire(
|
||||
vec![GatewayEvent::Started(response.clone())],
|
||||
ProtocolWireEvent::json("openai", Some("response.created".to_owned()), json!({
|
||||
"type": "response.created",
|
||||
"response": {"id": "resp_ws_settlement", "model": "model-a", "status": "in_progress"}
|
||||
})).unwrap(),
|
||||
),
|
||||
ProviderEvent::canonical_with_wire(
|
||||
vec![
|
||||
GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(websocket_charge().scaled()).unwrap(),
|
||||
),
|
||||
GatewayEvent::Completed(response),
|
||||
],
|
||||
ProtocolWireEvent::json("openai", Some("response.completed".to_owned()), json!({
|
||||
"type": "response.completed",
|
||||
"response": {"id": "resp_ws_settlement", "model": "model-a", "status": "completed"}
|
||||
})).unwrap(),
|
||||
),
|
||||
];
|
||||
Ok(ProviderStream::new(
|
||||
metadata,
|
||||
futures::stream::iter(events.map(Ok)),
|
||||
(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SettlementConnectionLifecycle {
|
||||
closed: Arc<tokio::sync::Notify>,
|
||||
}
|
||||
|
||||
struct SettlementConnectionGuard(Arc<tokio::sync::Notify>);
|
||||
|
||||
impl ConnectionGuard for SettlementConnectionGuard {}
|
||||
|
||||
impl Drop for SettlementConnectionGuard {
|
||||
fn drop(&mut self) {
|
||||
self.0.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionLifecycle for SettlementConnectionLifecycle {
|
||||
fn try_register(&self) -> Result<Box<dyn ConnectionGuard>, ConnectionDraining> {
|
||||
Ok(Box::new(SettlementConnectionGuard(self.closed.clone())))
|
||||
}
|
||||
|
||||
fn cancellation(&self) -> CancellationToken {
|
||||
CancellationToken::new()
|
||||
}
|
||||
|
||||
fn is_draining(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_disconnect_during_core_settlement_finishes_charge_before_releasing_admission() {
|
||||
let (release_settlement, gate) = tokio::sync::oneshot::channel();
|
||||
let ports = Arc::new(SettlementPorts {
|
||||
settlement_gate: Mutex::new(Some(gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let execution = Arc::new(
|
||||
DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot("sk_ws_settlement", "openai")),
|
||||
ports.clone(),
|
||||
ProviderRegistry::new([Arc::new(ChargedWebSocketProvider) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
ports.clone(),
|
||||
ports.clone(),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(IgnoredClientApiKeyUsage),
|
||||
)
|
||||
.with_budget(ports.clone()),
|
||||
);
|
||||
let lifecycle = Arc::new(SettlementConnectionLifecycle::default());
|
||||
let admin = crate::admin::AdminTestFixture::new().await;
|
||||
let app = gateway_api::initialize(
|
||||
gateway_api::ApiConfig {
|
||||
asset_directory: std::env::temp_dir(),
|
||||
cors_allowed_origins: Vec::new(),
|
||||
request_timeout_seconds: None,
|
||||
request_id_header: "x-request-id".to_owned(),
|
||||
},
|
||||
execution,
|
||||
admin.services,
|
||||
Vec::new(),
|
||||
Arc::new(EmptyWorkerHealth),
|
||||
lifecycle.clone(),
|
||||
)
|
||||
.unwrap()
|
||||
.router();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.unwrap();
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer sk_ws_settlement"),
|
||||
);
|
||||
let (mut socket, response) = connect_async(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
|
||||
socket
|
||||
.send(ClientMessage::Text(
|
||||
json!({
|
||||
"type": "response.create", "model": "model-a", "input": "hello"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), ports.settlement_started.notified())
|
||||
.await
|
||||
.expect("production Core reached settlement");
|
||||
{
|
||||
let finalizations = ports.finalizations.lock().unwrap();
|
||||
assert_eq!(finalizations.len(), 1);
|
||||
assert_eq!(
|
||||
finalizations[0].outcome,
|
||||
gateway_core::engine::ExecutionOutcome::Succeeded,
|
||||
"unexpected execution error: {:?}",
|
||||
finalizations[0].error,
|
||||
);
|
||||
}
|
||||
assert!(ports.charges.lock().unwrap().is_empty());
|
||||
assert_eq!(ports.releases.load(Ordering::SeqCst), 0);
|
||||
|
||||
// 等到生产连接 guard 释放,保证 forward 的 select 已因断连取消 next_event;
|
||||
// 在此之前不能打开结算屏障,否则只会验证正常完成而错过取消窗口。
|
||||
socket.close(None).await.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(5), lifecycle.closed.notified())
|
||||
.await
|
||||
.expect("production WebSocket handler exited");
|
||||
assert!(ports.active.load(Ordering::SeqCst));
|
||||
release_settlement
|
||||
.send(())
|
||||
.expect("the original settlement future survived disconnect");
|
||||
tokio::time::timeout(Duration::from_secs(5), ports.released.notified())
|
||||
.await
|
||||
.expect("detached Core cleanup released admission");
|
||||
assert!(!ports.active.load(Ordering::SeqCst));
|
||||
assert_eq!(ports.releases.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(ports.settlement_calls.load(Ordering::SeqCst), 1);
|
||||
{
|
||||
let charges = ports.charges.lock().unwrap();
|
||||
assert_eq!(charges.len(), 1);
|
||||
assert_eq!(charges[0].amount_usd, websocket_charge());
|
||||
assert_eq!(charges[0].key_id.as_str(), "key_api_test");
|
||||
assert_eq!(
|
||||
Some(&charges[0].request_id),
|
||||
ports.request_id.lock().unwrap().as_ref()
|
||||
);
|
||||
}
|
||||
server.abort();
|
||||
let _ = server.await;
|
||||
}
|
||||
|
||||
fn upgrade_request(authorization: &str) -> Request<Body> {
|
||||
Request::get("/v1/responses")
|
||||
.header(AUTHORIZATION, authorization)
|
||||
|
||||
@@ -9,6 +9,106 @@ use super::models::ModelsExecution;
|
||||
|
||||
const REMOVED_BODY_LIMIT_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
#[tokio::test]
|
||||
async fn browser_origin_controls_http_admin_sessions_without_configuration() {
|
||||
use axum::body::to_bytes;
|
||||
use serde_json::{Value, json};
|
||||
for (origins, secure) in [
|
||||
(vec!["http://admin.example.test"], false),
|
||||
(vec!["http://192.0.2.1:8080"], false),
|
||||
(vec!["http://[::1]:8080"], false),
|
||||
(vec!["https://admin.example.test"], true),
|
||||
(vec![], true),
|
||||
(vec!["null"], true),
|
||||
(vec!["invalid-origin"], true),
|
||||
(vec!["http://admin.example.test/path"], true),
|
||||
(vec!["http://admin.example.test?https"], true),
|
||||
(vec!["http://user@admin.example.test"], true),
|
||||
(
|
||||
vec!["http://admin.example.test", "https://admin.example.test"],
|
||||
true,
|
||||
),
|
||||
] {
|
||||
let app = api_router_with_origins(ModelsExecution::new(), Vec::new()).await;
|
||||
let request = |path: &str| {
|
||||
let mut builder = Request::post(path)
|
||||
.header("content-type", "application/json")
|
||||
.header("x-forwarded-proto", "http")
|
||||
.header("forwarded", "proto=http");
|
||||
for origin in &origins {
|
||||
builder = builder.header("origin", *origin);
|
||||
}
|
||||
builder
|
||||
};
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
request("/api/admin/auth/login")
|
||||
.body(Body::from(
|
||||
json!({"username": "admin_1", "password": "strong-admin-password"})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let cookie = response.headers()["set-cookie"].to_str().unwrap();
|
||||
let session = cookie.split(';').next().unwrap().to_owned();
|
||||
let attrs: Vec<_> = cookie.split(';').map(str::trim).collect();
|
||||
assert_eq!(attrs.contains(&"Secure"), secure);
|
||||
for attribute in ["Path=/", "HttpOnly", "SameSite=Lax"] {
|
||||
assert!(attrs.contains(&attribute));
|
||||
}
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::get("/api/admin/auth/status")
|
||||
.header("cookie", &session)
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let body: Value =
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap();
|
||||
assert_eq!(body["data"]["authenticated"], true);
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
request("/api/admin/auth/logout")
|
||||
.header("cookie", &session)
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let cookie = response.headers()["set-cookie"].to_str().unwrap();
|
||||
assert!(cookie.starts_with("cpr_admin_session=;"));
|
||||
let attrs: Vec<_> = cookie.split(';').map(str::trim).collect();
|
||||
assert_eq!(attrs.contains(&"Secure"), secure);
|
||||
for attribute in ["Path=/", "HttpOnly", "SameSite=Lax", "Max-Age=0"] {
|
||||
assert!(attrs.contains(&attribute));
|
||||
}
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::get("/api/admin/auth/status")
|
||||
.header("cookie", &session)
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let body: Value =
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap();
|
||||
assert_eq!(body["data"]["authenticated"], false);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn removed_responses_review_route_should_not_reach_the_responses_handler() {
|
||||
let response = api_router_with_origins(ModelsExecution::new(), Vec::new())
|
||||
|
||||
@@ -171,6 +171,11 @@ async fn search_route_should_preserve_request_and_success_response_bytes() {
|
||||
.expect("search response");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::CREATED);
|
||||
assert_eq!(response.headers()["x-request-id"], "req_search_test");
|
||||
assert_eq!(
|
||||
response.headers()["x-gateway-request-id"],
|
||||
"req_search_test"
|
||||
);
|
||||
assert_eq!(
|
||||
response.headers().get("x-search-rate-limit"),
|
||||
Some(&"42".parse().expect("header value"))
|
||||
@@ -216,6 +221,11 @@ async fn search_route_should_return_the_exact_upstream_error_response() {
|
||||
.expect("search response");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
|
||||
assert_eq!(response.headers()["x-request-id"], "req_search_test");
|
||||
assert_eq!(
|
||||
response.headers()["x-gateway-request-id"],
|
||||
"req_search_test"
|
||||
);
|
||||
assert_eq!(
|
||||
response.headers().get("content-type"),
|
||||
Some(&"application/problem+json".parse().expect("content type"))
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 诊断内容的统一筛选边界。未知字符串只保留长度和摘要。
|
||||
//! 诊断内容的统一筛选边界。未知键名与值只保留有界结构和摘要。
|
||||
|
||||
use serde_json::{Map, Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -6,8 +6,12 @@ use sha2::{Digest, Sha256};
|
||||
/// 原始字节的长度与 SHA-256,不复制或保留正文。
|
||||
#[must_use]
|
||||
pub fn body_fingerprint(bytes: &[u8]) -> Value {
|
||||
json!({"bytes": bytes.len(), "sha256": fingerprint_hex(bytes)})
|
||||
}
|
||||
|
||||
fn fingerprint_hex(bytes: &[u8]) -> String {
|
||||
let digest = Sha256::digest(bytes);
|
||||
let hex: String = digest
|
||||
digest
|
||||
.iter()
|
||||
.flat_map(|byte| {
|
||||
let alphabet = b"0123456789abcdef";
|
||||
@@ -16,17 +20,40 @@ pub fn body_fingerprint(bytes: &[u8]) -> Value {
|
||||
char::from(alphabet[usize::from(byte & 15)]),
|
||||
]
|
||||
})
|
||||
.collect();
|
||||
json!({"bytes": bytes.len(), "sha256": hex})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 保留可诊断的协议结构;凭据永不进入普通日志或数据库。
|
||||
/// 仅保留有限协议骨架;任意 JSON(包括 headers 和 metadata)不能声明自己可信。
|
||||
#[must_use]
|
||||
pub fn diagnostic_json(value: &Value) -> Value {
|
||||
capture(value, "", 0, &mut 96)
|
||||
capture(value, "", Context::Protocol, 0, &mut 96)
|
||||
}
|
||||
|
||||
fn capture(value: &Value, key: &str, depth: usize, budget: &mut usize) -> Value {
|
||||
/// 只有 transport 的上游 metadata / error 帧顶层才是受控头部边界。
|
||||
pub(super) fn diagnostic_event_json(value: &Value, stage: &str) -> Value {
|
||||
let context = if stage == "upstream.event"
|
||||
&& value.get("headers").is_some_and(Value::is_object)
|
||||
&& matches!(
|
||||
value.get("type").and_then(Value::as_str),
|
||||
Some("response.metadata" | "codex.response.metadata" | "error")
|
||||
) {
|
||||
Context::UpstreamEvent
|
||||
} else {
|
||||
Context::Protocol
|
||||
};
|
||||
capture(value, "", context, 0, &mut 96)
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Context {
|
||||
Protocol,
|
||||
UpstreamEvent,
|
||||
Opaque,
|
||||
Headers,
|
||||
HeaderValue,
|
||||
}
|
||||
|
||||
fn capture(value: &Value, key: &str, context: Context, depth: usize, budget: &mut usize) -> Value {
|
||||
if *budget == 0 || depth > 6 {
|
||||
return json!({"omitted": true});
|
||||
}
|
||||
@@ -42,7 +69,11 @@ fn capture(value: &Value, key: &str, depth: usize, budget: &mut usize) -> Value
|
||||
if *budget == 0 {
|
||||
break;
|
||||
}
|
||||
result.insert(bounded(name, 96), capture(value, name, depth + 1, budget));
|
||||
let (captured_name, child_context) = field(context, name);
|
||||
result.insert(
|
||||
captured_name,
|
||||
capture(value, name, child_context, depth + 1, budget),
|
||||
);
|
||||
}
|
||||
if result.len() < object.len() {
|
||||
result.insert(
|
||||
@@ -55,16 +86,86 @@ fn capture(value: &Value, key: &str, depth: usize, budget: &mut usize) -> Value
|
||||
Value::Array(items) => json!({
|
||||
"length": items.len(),
|
||||
"sample": items.iter().take(4)
|
||||
.map(|item| capture(item, key.as_str(), depth + 1, budget))
|
||||
.map(|item| {
|
||||
// 受控头值只接受一层字符串数组;畸形嵌套不能继承头部放行规则。
|
||||
let context = if matches!(context, Context::HeaderValue) && !item.is_string() {
|
||||
Context::Opaque
|
||||
} else {
|
||||
context
|
||||
};
|
||||
capture(item, key.as_str(), context, depth + 1, budget)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
}),
|
||||
Value::String(text)
|
||||
if safe_header(&key) || safe_value_key(&key) && protocol_label(text) =>
|
||||
if matches!(context, Context::HeaderValue) && safe_header(&key)
|
||||
|| matches!(context, Context::Protocol | Context::UpstreamEvent)
|
||||
&& protocol_value(&key, text) =>
|
||||
{
|
||||
Value::String(bounded(text, 256))
|
||||
}
|
||||
Value::String(text) => body_fingerprint(text.as_bytes()),
|
||||
other => other.clone(),
|
||||
Value::Number(number)
|
||||
if matches!(context, Context::Protocol)
|
||||
&& key == "status"
|
||||
&& number
|
||||
.as_u64()
|
||||
.is_some_and(|status| (100..=599).contains(&status)) =>
|
||||
{
|
||||
value.clone()
|
||||
}
|
||||
other => body_fingerprint(other.to_string().as_bytes()),
|
||||
}
|
||||
}
|
||||
|
||||
fn field(context: Context, name: &str) -> (String, Context) {
|
||||
let child = match (context, name) {
|
||||
(Context::UpstreamEvent, "headers") => Some(Context::Headers),
|
||||
(Context::Headers, _) if safe_header(&name.to_ascii_lowercase()) => {
|
||||
Some(Context::HeaderValue)
|
||||
}
|
||||
(
|
||||
Context::Protocol | Context::UpstreamEvent,
|
||||
"response" | "error" | "input" | "output" | "item" | "part" | "content" | "type"
|
||||
| "role" | "status",
|
||||
) => Some(Context::Protocol),
|
||||
(
|
||||
Context::Protocol | Context::UpstreamEvent,
|
||||
"id" | "model" | "message" | "text" | "delta" | "code" | "metadata" | "headers",
|
||||
) => Some(Context::Opaque),
|
||||
_ => None,
|
||||
};
|
||||
match child {
|
||||
Some(context) => (name.to_owned(), context),
|
||||
None => (
|
||||
format!("field_{}", fingerprint_hex(name.as_bytes())),
|
||||
Context::Opaque,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn protocol_value(key: &str, value: &str) -> bool {
|
||||
match key {
|
||||
"type" => {
|
||||
known_event_type(value)
|
||||
|| matches!(
|
||||
value,
|
||||
"message"
|
||||
| "function_call"
|
||||
| "function_call_output"
|
||||
| "reasoning"
|
||||
| "input_text"
|
||||
| "output_text"
|
||||
| "input_image"
|
||||
| "refusal"
|
||||
)
|
||||
}
|
||||
"role" => matches!(
|
||||
value,
|
||||
"system" | "developer" | "user" | "assistant" | "tool"
|
||||
),
|
||||
"status" => matches!(value, "in_progress" | "completed" | "incomplete" | "failed"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -131,14 +232,6 @@ fn safe_header(key: &str) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
fn safe_value_key(key: &str) -> bool {
|
||||
safe_header(key)
|
||||
|| matches!(
|
||||
key,
|
||||
"type" | "event" | "status" | "code" | "model" | "role" | "object" | "service_tier"
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn bounded(value: &str, limit: usize) -> String {
|
||||
value
|
||||
.chars()
|
||||
@@ -147,11 +240,39 @@ pub(super) fn bounded(value: &str, limit: usize) -> String {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Event names/codes are protocol labels, never arbitrary upstream text.
|
||||
pub(super) fn protocol_label(value: &str) -> bool {
|
||||
!value.is_empty()
|
||||
&& value.len() <= 128
|
||||
&& value
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-' | b':'))
|
||||
pub(super) fn diagnostic_event_type(value: &str) -> Value {
|
||||
if known_event_type(value) {
|
||||
Value::String(value.to_owned())
|
||||
} else {
|
||||
body_fingerprint(value.as_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
/// 这里只决定诊断可明文展示的阶段,不校验协议;未知事件仍采集摘要,不能按字符形状放行。
|
||||
fn known_event_type(value: &str) -> bool {
|
||||
matches!(
|
||||
value,
|
||||
"error"
|
||||
| "response.create"
|
||||
| "response.created"
|
||||
| "response.in_progress"
|
||||
| "response.completed"
|
||||
| "response.failed"
|
||||
| "response.incomplete"
|
||||
| "response.metadata"
|
||||
| "codex.response.metadata"
|
||||
| "codex.rate_limits"
|
||||
| "response.output_item.added"
|
||||
| "response.output_item.done"
|
||||
| "response.content_part.added"
|
||||
| "response.content_part.done"
|
||||
| "response.output_text.delta"
|
||||
| "response.output_text.done"
|
||||
| "response.function_call_arguments.delta"
|
||||
| "response.function_call_arguments.done"
|
||||
| "response.reasoning_summary_text.delta"
|
||||
| "response.reasoning_summary_text.done"
|
||||
| "response.reasoning_text.delta"
|
||||
| "response.reasoning_text.done"
|
||||
)
|
||||
}
|
||||
|
||||
@@ -10,8 +10,8 @@ use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use serde::Serialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::capture::bounded;
|
||||
use super::{body_fingerprint, diagnostic_headers, diagnostic_json};
|
||||
use super::capture::{bounded, diagnostic_event_json, diagnostic_event_type};
|
||||
use super::{body_fingerprint, diagnostic_headers};
|
||||
|
||||
const MAX_EVENTS: usize = 128;
|
||||
const HEAD_EVENTS: usize = 8;
|
||||
@@ -124,7 +124,7 @@ impl TraceContext {
|
||||
}
|
||||
|
||||
/// 同一个头部边界只采集一次:默认脱敏,显式开启 dump 时另存完整字节。
|
||||
/// facts 必须为调用方构造的安全对象;原始头值只进入独立 dump。
|
||||
/// facts 必须为调用方构造的安全对象;仅受控诊断头可明文保留,完整头部另存 dump。
|
||||
pub fn headers<'a>(
|
||||
&self,
|
||||
stage: &'static str,
|
||||
@@ -182,27 +182,27 @@ impl TraceContext {
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.or(name)
|
||||
.filter(|kind| super::capture::protocol_label(kind));
|
||||
.or(name);
|
||||
// 后缀只用于观测合并,不授予明文权限;未知事件用摘要区分,避免合并成同一类。
|
||||
let delta = event_type.is_some_and(|kind| kind.ends_with(".delta"));
|
||||
let data = json!({
|
||||
"body": fingerprint,
|
||||
"eventType": event_type.map(|kind| bounded(kind, 128)),
|
||||
"eventType": event_type.map(diagnostic_event_type),
|
||||
"jsonValid": parsed.is_some(),
|
||||
"metadata": parsed.as_ref().filter(|_| !delta).map(diagnostic_json),
|
||||
"metadata": parsed.as_ref().filter(|_| !delta)
|
||||
.map(|value| diagnostic_event_json(value, stage)),
|
||||
});
|
||||
self.push(stage, data, delta);
|
||||
}
|
||||
|
||||
/// 只记录业务解析边界;metadata 已由 transport 在解析前捕获。
|
||||
pub fn wire_event(&self, protocol: &str, event_type: Option<&str>) {
|
||||
let event_type = event_type.filter(|kind| super::capture::protocol_label(kind));
|
||||
let delta = event_type.is_some_and(|kind| kind.ends_with(".delta"));
|
||||
self.push(
|
||||
"provider.event",
|
||||
json!({
|
||||
"protocol": protocol,
|
||||
"eventType": event_type.map(|kind| bounded(kind, 128)),
|
||||
"eventType": event_type.map(diagnostic_event_type),
|
||||
}),
|
||||
delta,
|
||||
);
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use crate::diagnostics::TraceContext;
|
||||
use serde_json::json;
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::future::Future;
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
@@ -33,7 +33,7 @@ use crate::lifecycle::CancellationToken;
|
||||
use crate::metering::Decimal;
|
||||
use crate::operation::{Operation, ProviderSessionState};
|
||||
use crate::routing::RoutingPlan;
|
||||
use futures::future::Fuse;
|
||||
use futures::future::{BoxFuture, Fuse};
|
||||
use futures::{FutureExt, StreamExt, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
|
||||
@@ -63,7 +63,7 @@ impl AccountSelection {
|
||||
|
||||
impl<S: ?Sized> AttemptCoordinator<S>
|
||||
where
|
||||
S: ExecutionStore,
|
||||
S: ExecutionStore + 'static,
|
||||
{
|
||||
#[must_use]
|
||||
pub fn new(engine: GatewayEngine<S>) -> Self {
|
||||
@@ -72,11 +72,11 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号选择成功后创建唯一请求行,再返回由 Core 完整拥有 retry/commit 的流式会话。
|
||||
/// 接纳请求并返回由 Core 完整拥有 retry/commit 的流式会话。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 账号选择后的持久化、取消或已过 deadline 时返回稳定错误。
|
||||
/// 取消、已过 deadline 或 continuation 与路由不匹配时返回稳定错误。
|
||||
pub async fn start(
|
||||
&self,
|
||||
request: NewModelRequest,
|
||||
@@ -192,14 +192,15 @@ where
|
||||
excluded_accounts: BTreeSet::new(),
|
||||
credential_recovery_attempted_accounts: BTreeSet::new(),
|
||||
recovery_account: None,
|
||||
transport_recovery: None,
|
||||
pending_retry: None,
|
||||
transient_retry_counts: BTreeMap::new(),
|
||||
current: None,
|
||||
send_state_watermark: UpstreamSendState::NotSent,
|
||||
downstream_committed_at: None,
|
||||
client_status_code: None,
|
||||
delivery_pending: false,
|
||||
upstream_complete: false,
|
||||
finalized: false,
|
||||
finalization: None,
|
||||
finalized_at: None,
|
||||
image_generation_requested,
|
||||
last_retryable_failure: None,
|
||||
@@ -209,11 +210,11 @@ where
|
||||
};
|
||||
|
||||
if session.cancellation.is_cancelled() {
|
||||
session.finish_interruption(EngineError::Cancelled).await?;
|
||||
session.finish_interruption(&EngineError::Cancelled).await?;
|
||||
return Err(EngineError::Cancelled);
|
||||
}
|
||||
if SystemTime::now() >= deadline {
|
||||
session.finish_interruption(EngineError::Deadline).await?;
|
||||
session.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
Ok(session)
|
||||
@@ -224,17 +225,29 @@ struct CurrentAttempt {
|
||||
stream: ProviderStream,
|
||||
metadata: ProviderCallMetadata,
|
||||
trigger: AttemptTrigger,
|
||||
transport: AttemptTransport,
|
||||
index: NonZeroU32,
|
||||
started_at: SystemTime,
|
||||
send_observed: bool,
|
||||
response_observation: Option<ProviderResponseObservation>,
|
||||
}
|
||||
|
||||
impl CurrentAttempt {
|
||||
fn upstream_request_id(&self) -> Option<&str> {
|
||||
self.response_observation
|
||||
.as_ref()
|
||||
.and_then(ProviderResponseObservation::request_id)
|
||||
.or_else(|| self.metadata.upstream_request_id())
|
||||
.map(|id| id.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
struct FailureFinalization {
|
||||
outcome: ExecutionOutcome,
|
||||
send_state: UpstreamSendState,
|
||||
error: GatewayError,
|
||||
upstream_status_code: Option<u16>,
|
||||
upstream_request_id: Option<String>,
|
||||
provider_error_code: Option<String>,
|
||||
raw_upstream_error: Option<String>,
|
||||
retry_after_ms: Option<u64>,
|
||||
@@ -242,10 +255,11 @@ struct FailureFinalization {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct PendingTransportRecovery {
|
||||
struct PendingAttemptRetry {
|
||||
account: crate::account::ProviderAccountId,
|
||||
transport: AttemptTransport,
|
||||
delay: Duration,
|
||||
transport_recovery: bool,
|
||||
}
|
||||
|
||||
/// API 可逐事件消费的 Core 执行会话。
|
||||
@@ -283,9 +297,10 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
/// attempt 建立时即被消费,后续可重试错误仍可换号消耗剩余重试预算。
|
||||
/// 与 `required_account`(外部指定、贯穿整个请求)语义不同,不可合并。
|
||||
recovery_account: Option<crate::account::ProviderAccountId>,
|
||||
/// 上游传输在提交边界前失败后的一次性同账号重试或备用传输钉选。
|
||||
/// 与凭据恢复分开保存,避免把 transport recovery 误记为 OAuth 恢复。
|
||||
transport_recovery: Option<PendingTransportRecovery>,
|
||||
/// 提交前同账号退避/传输恢复共用的等待与钉选状态;业务重试消耗路由预算。
|
||||
pending_retry: Option<PendingAttemptRetry>,
|
||||
/// 请求内按账号累计的瞬时拒绝重试次数,不能跨请求污染账号健康状态。
|
||||
transient_retry_counts: BTreeMap<crate::account::ProviderAccountId, u32>,
|
||||
current: Option<CurrentAttempt>,
|
||||
/// 请求级发送状态水位;跨 attempt 单调不降,终态写回不得低于此档。
|
||||
send_state_watermark: UpstreamSendState,
|
||||
@@ -293,7 +308,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
client_status_code: Option<u16>,
|
||||
delivery_pending: bool,
|
||||
upstream_complete: bool,
|
||||
finalized: bool,
|
||||
finalization: Option<RequestFinalization>,
|
||||
finalized_at: Option<SystemTime>,
|
||||
image_generation_requested: bool,
|
||||
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
|
||||
@@ -310,9 +325,14 @@ struct PendingTerminalFailure {
|
||||
send_state: UpstreamSendState,
|
||||
}
|
||||
|
||||
enum RequestFinalization {
|
||||
Pending(BoxFuture<'static, bool>),
|
||||
Complete,
|
||||
}
|
||||
|
||||
impl<S: ?Sized> ResponseExecutionSession<S>
|
||||
where
|
||||
S: ExecutionStore,
|
||||
S: ExecutionStore + 'static,
|
||||
{
|
||||
/// 当前请求共享的诊断上下文。
|
||||
pub fn trace(&self) -> TraceContext {
|
||||
@@ -323,8 +343,9 @@ where
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 未提交上一条首事件、Provider/Store 失败、取消或超时时返回错误。
|
||||
/// 未提交上一条首事件、Provider 失败、取消或超时时返回错误;观测写入失败只记录告警。
|
||||
pub async fn next_event(&mut self) -> Result<Option<CoordinatedEvent>, EngineError> {
|
||||
self.resume_finalization().await;
|
||||
if self.delivery_pending {
|
||||
return Err(EngineError::DownstreamCommitRequired);
|
||||
}
|
||||
@@ -333,7 +354,7 @@ where
|
||||
.await?;
|
||||
return Err(provider_engine_error(pending.error));
|
||||
}
|
||||
if self.finalized {
|
||||
if self.is_finalized() {
|
||||
return Ok(None);
|
||||
}
|
||||
loop {
|
||||
@@ -383,10 +404,11 @@ where
|
||||
///
|
||||
/// 会话已提交、已有待提交结果或执行失败时返回错误。
|
||||
pub async fn collect_uncommitted(&mut self) -> Result<Vec<ProviderEvent>, EngineError> {
|
||||
self.resume_finalization().await;
|
||||
if self.downstream_committed_at.is_some() || self.delivery_pending {
|
||||
return Err(EngineError::InvalidDeliveryState);
|
||||
}
|
||||
if self.finalized {
|
||||
if self.is_finalized() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
@@ -422,7 +444,7 @@ where
|
||||
&mut self,
|
||||
client_status_code: Option<u16>,
|
||||
) -> Result<(), EngineError> {
|
||||
if !self.delivery_pending || self.downstream_committed_at.is_some() || self.finalized {
|
||||
if !self.delivery_pending || self.downstream_committed_at.is_some() || self.is_finalized() {
|
||||
return Err(EngineError::InvalidDeliveryState);
|
||||
}
|
||||
let committed_at = SystemTime::now();
|
||||
@@ -479,7 +501,7 @@ where
|
||||
|
||||
#[must_use]
|
||||
pub const fn is_finalized(&self) -> bool {
|
||||
self.finalized
|
||||
matches!(self.finalization, Some(RequestFinalization::Complete))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
@@ -573,14 +595,13 @@ where
|
||||
self.cancellation.cancel();
|
||||
}
|
||||
|
||||
/// 丢弃尚未提交的 delivery,并立即把请求收敛为取消终态。
|
||||
/// 丢弃尚未提交的 delivery,并收敛为取消终态;已有终态写入继续保留原结果。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Store 无法写回发送状态或最终请求行时返回错误。
|
||||
/// 观测写入失败只记录告警,不改变取消结果。
|
||||
pub async fn cancel_and_finalize(&mut self) -> Result<(), EngineError> {
|
||||
self.cancellation.cancel();
|
||||
if self.finalized {
|
||||
self.resume_finalization().await;
|
||||
if self.is_finalized() {
|
||||
return Ok(());
|
||||
}
|
||||
self.delivery_pending = false;
|
||||
@@ -589,15 +610,22 @@ where
|
||||
.finish_provider_error_with_send_state(&pending.error, pending.send_state)
|
||||
.await;
|
||||
}
|
||||
self.finish_interruption(EngineError::Cancelled).await
|
||||
self.finish_interruption(&EngineError::Cancelled).await
|
||||
}
|
||||
|
||||
async fn pull(&mut self) -> Result<PullOutcome, EngineError> {
|
||||
loop {
|
||||
if self.current.is_none()
|
||||
&& let Some(outcome) = self.prepare_attempt().await?
|
||||
{
|
||||
return Ok(outcome);
|
||||
if self.current.is_none() {
|
||||
match self.prepare_attempt().await {
|
||||
Ok(Some(outcome)) => return Ok(outcome),
|
||||
Ok(None) => {}
|
||||
Err(error) => {
|
||||
if !self.is_finalized() {
|
||||
self.finish_interruption(&error).await?;
|
||||
}
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let boundary = {
|
||||
@@ -613,14 +641,14 @@ where
|
||||
|
||||
match boundary {
|
||||
PollBoundary::Cancelled => {
|
||||
self.finish_interruption(EngineError::Cancelled).await?;
|
||||
self.finish_interruption(&EngineError::Cancelled).await?;
|
||||
return Err(EngineError::Cancelled);
|
||||
}
|
||||
PollBoundary::Deadline => {
|
||||
// 会话 deadline 是网关自身的请求预算,不是上游超时;
|
||||
// 真正的上游超时会作为流错误进入 `handle_stream_error` 记账。
|
||||
// 这里不写 provider 失败,避免长流集中到期误触 provider 熔断。
|
||||
self.finish_interruption(EngineError::Deadline).await?;
|
||||
self.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
PollBoundary::Item(Some(Ok(mut event))) => {
|
||||
@@ -669,21 +697,24 @@ where
|
||||
}
|
||||
|
||||
async fn prepare_attempt(&mut self) -> Result<Option<PullOutcome>, EngineError> {
|
||||
let transport_recovery = self.transport_recovery.take();
|
||||
if transport_recovery.is_none() && self.routing_attempts >= self.plan.max_attempts().get() {
|
||||
let pending_retry = self.pending_retry.take();
|
||||
let is_transport_recovery = pending_retry
|
||||
.as_ref()
|
||||
.is_some_and(|retry| retry.transport_recovery);
|
||||
if !is_transport_recovery && self.routing_attempts >= self.plan.max_attempts().get() {
|
||||
return Err(EngineError::EmptyRoutingPlan);
|
||||
}
|
||||
if let Some(recovery) = transport_recovery.as_ref()
|
||||
if let Some(recovery) = pending_retry.as_ref()
|
||||
&& !recovery.delay.is_zero()
|
||||
{
|
||||
match poll_retry_delay(recovery.delay, self.cancellation.clone(), self.deadline).await {
|
||||
RetryDelayBoundary::Elapsed => {}
|
||||
RetryDelayBoundary::Cancelled => {
|
||||
self.finish_interruption(EngineError::Cancelled).await?;
|
||||
self.finish_interruption(&EngineError::Cancelled).await?;
|
||||
return Err(EngineError::Cancelled);
|
||||
}
|
||||
RetryDelayBoundary::Deadline => {
|
||||
self.finish_interruption(EngineError::Deadline).await?;
|
||||
self.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
}
|
||||
@@ -698,6 +729,7 @@ where
|
||||
send_state: self.current_send_state(),
|
||||
error,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
retry_after_ms: None,
|
||||
@@ -713,8 +745,7 @@ where
|
||||
.ok_or(EngineError::EmptyRoutingPlan)?;
|
||||
// 请求局部恢复钉选在此被一次性消费,只绑定本次 replay attempt;
|
||||
// 外部 required_account 每次 attempt 都重新生效。
|
||||
let is_transport_recovery = transport_recovery.is_some();
|
||||
let (pinned_account, attempt_transport) = if let Some(recovery) = transport_recovery {
|
||||
let (pinned_account, attempt_transport) = if let Some(recovery) = pending_retry {
|
||||
(Some(recovery.account), recovery.transport)
|
||||
} else {
|
||||
match &self.account_selection {
|
||||
@@ -795,14 +826,14 @@ where
|
||||
.await
|
||||
{
|
||||
ProviderBoundary::Cancelled => {
|
||||
self.finish_interruption(EngineError::Cancelled).await?;
|
||||
self.finish_interruption(&EngineError::Cancelled).await?;
|
||||
return Err(EngineError::Cancelled);
|
||||
}
|
||||
ProviderBoundary::Deadline => {
|
||||
// 网关预算到期同样不是候选 Provider 的上游超时,不计入熔断;
|
||||
// Provider 自身的握手/传输超时会以 `ProviderErrorKind::Timeout`
|
||||
// 错误返回并在下方 `Result` 分支记账。
|
||||
self.finish_interruption(EngineError::Deadline).await?;
|
||||
self.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
ProviderBoundary::Result(result) => match *result {
|
||||
@@ -861,6 +892,7 @@ where
|
||||
},
|
||||
};
|
||||
if !stream.metadata().confirms(&candidate) {
|
||||
drop(stream);
|
||||
self.record_provider_failure(candidate.provider().clone(), ProviderErrorKind::Protocol);
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
@@ -871,6 +903,7 @@ where
|
||||
send_state: self.current_send_state(),
|
||||
error,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
retry_after_ms: None,
|
||||
@@ -893,6 +926,7 @@ where
|
||||
.account_scope()
|
||||
.allows(metadata.provider_account_id())
|
||||
{
|
||||
drop(stream);
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"provider selected an account outside the frozen client scope",
|
||||
@@ -902,6 +936,7 @@ where
|
||||
send_state: self.current_send_state(),
|
||||
error,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
retry_after_ms: None,
|
||||
@@ -914,6 +949,7 @@ where
|
||||
.as_ref()
|
||||
.is_some_and(|required| metadata.provider_account_id() != required)
|
||||
{
|
||||
drop(stream);
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"provider did not use the required account",
|
||||
@@ -923,6 +959,7 @@ where
|
||||
send_state: self.current_send_state(),
|
||||
error,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
retry_after_ms: None,
|
||||
@@ -938,6 +975,7 @@ where
|
||||
&& self.continuation_attempt == ContinuationAttempt::Native
|
||||
&& !pin.matches(metadata.provider(), metadata.provider_account_id())
|
||||
{
|
||||
drop(stream);
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"native continuation binding did not match selected account",
|
||||
@@ -947,6 +985,7 @@ where
|
||||
send_state: self.current_send_state(),
|
||||
error,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
retry_after_ms: None,
|
||||
@@ -976,6 +1015,21 @@ where
|
||||
capacity_used_slots: capacity.map(|snapshot| snapshot.used_slots()),
|
||||
capacity_total_slots: capacity.map(|snapshot| snapshot.total_slots()),
|
||||
};
|
||||
// 合法冷流在可取消的观测写入前归会话所有,首写等待中断也不能退回零次或重新选号。
|
||||
self.attempts = next_attempt.get();
|
||||
if !is_transport_recovery {
|
||||
self.routing_attempts = self.routing_attempts.saturating_add(1);
|
||||
}
|
||||
self.current = Some(CurrentAttempt {
|
||||
stream,
|
||||
metadata,
|
||||
trigger,
|
||||
transport: attempt_transport,
|
||||
index: next_attempt,
|
||||
started_at: SystemTime::now(),
|
||||
send_observed: false,
|
||||
response_observation: None,
|
||||
});
|
||||
if self.request_persisted {
|
||||
best_effort_store_write(
|
||||
"record_attempt",
|
||||
@@ -1002,19 +1056,6 @@ where
|
||||
self.request_persisted = true;
|
||||
}
|
||||
}
|
||||
self.attempts = next_attempt.get();
|
||||
if !is_transport_recovery {
|
||||
self.routing_attempts = self.routing_attempts.saturating_add(1);
|
||||
}
|
||||
self.current = Some(CurrentAttempt {
|
||||
stream,
|
||||
metadata,
|
||||
trigger,
|
||||
index: next_attempt,
|
||||
started_at: SystemTime::now(),
|
||||
send_observed: false,
|
||||
response_observation: None,
|
||||
});
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
@@ -1156,6 +1197,27 @@ where
|
||||
&& attempt_send_state != UpstreamSendState::Ambiguous
|
||||
&& provider_proved_replay_safe
|
||||
&& self.routing_attempts < self.plan.max_attempts().get();
|
||||
let transient_retry = match error.pre_delivery_retry() {
|
||||
Some(crate::error::PreDeliveryRetry::SameAccountTransientRetry {
|
||||
max_retries,
|
||||
initial_delay,
|
||||
max_delay,
|
||||
}) if ordinary_retry => {
|
||||
let retries = self
|
||||
.transient_retry_counts
|
||||
.entry(current.metadata.provider_account_id().clone())
|
||||
.or_default();
|
||||
if *retries < max_retries.get() {
|
||||
let multiplier = 1_u32.checked_shl(*retries).unwrap_or(u32::MAX);
|
||||
let delay = initial_delay.saturating_mul(multiplier).min(max_delay);
|
||||
*retries = retries.saturating_add(1);
|
||||
Some(delay)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let same_account_retry = error.retries_same_account()
|
||||
&& provider_proved_replay_safe
|
||||
&& self.downstream_committed_at.is_none()
|
||||
@@ -1175,7 +1237,8 @@ where
|
||||
"retryable": retryable, "continuationRetry": continuation_retry,
|
||||
"sameAccountRetry": same_account_retry, "accountRotationRetry": account_rotation_retry,
|
||||
"ordinaryRetry": ordinary_retry, "transportRecovery": transport_recovery.is_some(),
|
||||
"delayMs": transport_recovery.map(|(_, delay)| duration_ms(delay)),
|
||||
"transientRetry": transient_retry.is_some(),
|
||||
"delayMs": transient_retry.or(transport_recovery.map(|(_, delay)| delay)).map(duration_ms),
|
||||
"downstreamCommitted": self.downstream_committed_at.is_some(),
|
||||
"sendState": format!("{attempt_send_state:?}"),
|
||||
}));
|
||||
@@ -1195,11 +1258,19 @@ where
|
||||
// 只钉住紧随其后的 replay attempt;replay 再遇可重试错误时,
|
||||
// ordinary/continuation 重试门不受影响,仍可换号。
|
||||
self.recovery_account = Some(account);
|
||||
} else if let Some(delay) = transient_retry {
|
||||
self.pending_retry = Some(PendingAttemptRetry {
|
||||
account: current.metadata.provider_account_id().clone(),
|
||||
transport: current.transport,
|
||||
delay,
|
||||
transport_recovery: false,
|
||||
});
|
||||
} else if let Some((transport, delay)) = transport_recovery {
|
||||
self.transport_recovery = Some(PendingTransportRecovery {
|
||||
self.pending_retry = Some(PendingAttemptRetry {
|
||||
account: current.metadata.provider_account_id().clone(),
|
||||
transport,
|
||||
delay,
|
||||
transport_recovery: true,
|
||||
});
|
||||
} else if !continuation_retry {
|
||||
self.excluded_accounts
|
||||
@@ -1220,15 +1291,17 @@ where
|
||||
provider_kind: current.metadata.provider().clone(),
|
||||
account_id: Some(current.metadata.provider_account_id().clone()),
|
||||
upstream_model_id: current.metadata.upstream_model().cloned(),
|
||||
upstream_status_code: current
|
||||
.response_observation
|
||||
.as_ref()
|
||||
.and_then(ProviderResponseObservation::status_code),
|
||||
upstream_request_id: current
|
||||
.response_observation
|
||||
.as_ref()
|
||||
.and_then(ProviderResponseObservation::request_id)
|
||||
.map(|value| value.as_str().to_owned()),
|
||||
upstream_status_code: error.upstream_status().or_else(|| {
|
||||
current
|
||||
.response_observation
|
||||
.as_ref()
|
||||
.and_then(ProviderResponseObservation::status_code)
|
||||
}),
|
||||
upstream_request_id: error
|
||||
.upstream_request_id()
|
||||
.map(|id| id.as_str())
|
||||
.or_else(|| current.upstream_request_id())
|
||||
.map(str::to_owned),
|
||||
latency: current.started_at.elapsed().unwrap_or_default(),
|
||||
error,
|
||||
}),
|
||||
@@ -1370,28 +1443,16 @@ where
|
||||
}
|
||||
|
||||
async fn finish_success(&mut self) -> Result<(), EngineError> {
|
||||
if self.finalized {
|
||||
if self.is_finalized() {
|
||||
return Ok(());
|
||||
}
|
||||
let completed_at = SystemTime::now();
|
||||
self.finalized_at = Some(completed_at);
|
||||
self.observation.finish();
|
||||
let upstream_request_id = self.current.as_ref().and_then(|current| {
|
||||
current
|
||||
.response_observation
|
||||
.as_ref()
|
||||
.and_then(|observation| {
|
||||
observation
|
||||
.request_id()
|
||||
.map(|value| value.as_str().to_owned())
|
||||
})
|
||||
.or_else(|| {
|
||||
current
|
||||
.metadata
|
||||
.upstream_request_id()
|
||||
.map(|value| value.as_str().to_owned())
|
||||
})
|
||||
});
|
||||
let upstream_request_id = self
|
||||
.current
|
||||
.as_ref()
|
||||
.and_then(CurrentAttempt::upstream_request_id)
|
||||
.map(str::to_owned);
|
||||
let upstream_status_code = self
|
||||
.current
|
||||
.as_ref()
|
||||
@@ -1405,50 +1466,44 @@ where
|
||||
"request.finished",
|
||||
json!({"outcome": "succeeded", "attempts": self.attempts}),
|
||||
);
|
||||
if self.request_persisted {
|
||||
best_effort_store_write(
|
||||
"finalize_model_request",
|
||||
&self.request_id,
|
||||
self.engine
|
||||
.store()
|
||||
.finalize_model_request(ModelRequestFinalization {
|
||||
request_id: self.request_id.clone(),
|
||||
outcome: ExecutionOutcome::Succeeded,
|
||||
send_state: UpstreamSendState::Sent,
|
||||
attempt_count: self.attempts,
|
||||
downstream_committed_at: self.downstream_committed_at,
|
||||
client_status_code: self.client_status_code,
|
||||
upstream_status_code,
|
||||
client_response_id: self.observation.client_response_id.clone(),
|
||||
upstream_request_id,
|
||||
upstream_response_id: self.observation.upstream_response_id.clone(),
|
||||
upstream_transport,
|
||||
http_version,
|
||||
websocket_pool,
|
||||
service_tier,
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
|
||||
error: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
failure_observation: ModelRequestFailureObservation::default(),
|
||||
retry_after_ms: None,
|
||||
usage: self.observation.usage.clone(),
|
||||
image_generation_succeeded: self.image_generation_succeeded(),
|
||||
cost: self.observation.cost.clone(),
|
||||
timings: self.observation.timings.clone(),
|
||||
completed_at,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
self.finalized = true;
|
||||
self.persist_finalization(ModelRequestFinalization {
|
||||
request_id: self.request_id.clone(),
|
||||
outcome: ExecutionOutcome::Succeeded,
|
||||
send_state: UpstreamSendState::Sent,
|
||||
attempt_count: self.attempts,
|
||||
downstream_committed_at: self.downstream_committed_at,
|
||||
client_status_code: self.client_status_code,
|
||||
upstream_status_code,
|
||||
client_response_id: self.observation.client_response_id.clone(),
|
||||
upstream_request_id,
|
||||
upstream_response_id: self.observation.upstream_response_id.clone(),
|
||||
upstream_transport,
|
||||
http_version,
|
||||
websocket_pool,
|
||||
service_tier,
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
|
||||
error: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
failure_observation: ModelRequestFailureObservation::default(),
|
||||
retry_after_ms: None,
|
||||
usage: self.observation.usage.clone(),
|
||||
image_generation_succeeded: self.image_generation_succeeded(),
|
||||
cost: self.observation.cost.clone(),
|
||||
timings: self.observation.timings.clone(),
|
||||
completed_at,
|
||||
})
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finish_provider_error(&mut self, error: &ProviderError) -> Result<(), EngineError> {
|
||||
self.finish_provider_error_with_send_state(error, self.current_send_state())
|
||||
.await
|
||||
self.finish_provider_error_with_send_state(
|
||||
error,
|
||||
escalate_send_state(self.current_send_state(), error.send_state()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn finish_provider_error_with_send_state(
|
||||
@@ -1469,6 +1524,7 @@ where
|
||||
send_state,
|
||||
error: GatewayError::from_provider(error),
|
||||
upstream_status_code: error.upstream_status(),
|
||||
upstream_request_id: error.upstream_request_id().map(|id| id.as_str().to_owned()),
|
||||
provider_error_code: error.upstream_code().map(|code| code.as_str().to_owned()),
|
||||
raw_upstream_error: error
|
||||
.raw_upstream_error()
|
||||
@@ -1484,7 +1540,7 @@ where
|
||||
.await
|
||||
}
|
||||
|
||||
async fn finish_interruption(&mut self, error: EngineError) -> Result<(), EngineError> {
|
||||
async fn finish_interruption(&mut self, error: &EngineError) -> Result<(), EngineError> {
|
||||
let (outcome, gateway_error) = match error {
|
||||
EngineError::Cancelled => (
|
||||
ExecutionOutcome::Cancelled,
|
||||
@@ -1531,6 +1587,7 @@ where
|
||||
send_state,
|
||||
error: gateway_error,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
retry_after_ms: None,
|
||||
@@ -1543,7 +1600,7 @@ where
|
||||
&mut self,
|
||||
finalization: FailureFinalization,
|
||||
) -> Result<(), EngineError> {
|
||||
if self.finalized {
|
||||
if self.is_finalized() {
|
||||
return Ok(());
|
||||
}
|
||||
self.trace.record(
|
||||
@@ -1555,28 +1612,12 @@ where
|
||||
}),
|
||||
);
|
||||
let completed_at = SystemTime::now();
|
||||
self.finalized_at = Some(completed_at);
|
||||
if !self.request_persisted {
|
||||
self.finalized = true;
|
||||
return Ok(());
|
||||
}
|
||||
self.observation.finish();
|
||||
let upstream_request_id = self.current.as_ref().and_then(|current| {
|
||||
current
|
||||
.response_observation
|
||||
.as_ref()
|
||||
.and_then(|observation| {
|
||||
observation
|
||||
.request_id()
|
||||
.map(|value| value.as_str().to_owned())
|
||||
})
|
||||
.or_else(|| {
|
||||
current
|
||||
.metadata
|
||||
.upstream_request_id()
|
||||
.map(|value| value.as_str().to_owned())
|
||||
})
|
||||
});
|
||||
let upstream_request_id = self
|
||||
.current
|
||||
.as_ref()
|
||||
.and_then(CurrentAttempt::upstream_request_id)
|
||||
.map(str::to_owned);
|
||||
let observed_status_code = self
|
||||
.current
|
||||
.as_ref()
|
||||
@@ -1586,47 +1627,83 @@ where
|
||||
self.current_transport_observation();
|
||||
let service_tier = self.current_service_tier();
|
||||
let provider_metadata_json = self.current_provider_metadata_json();
|
||||
best_effort_store_write(
|
||||
"finalize_model_request",
|
||||
&self.request_id,
|
||||
self.engine
|
||||
.store()
|
||||
.finalize_model_request(ModelRequestFinalization {
|
||||
request_id: self.request_id.clone(),
|
||||
outcome: finalization.outcome,
|
||||
send_state: finalization.send_state,
|
||||
attempt_count: self.attempts,
|
||||
downstream_committed_at: self.downstream_committed_at,
|
||||
client_status_code: self.client_status_code,
|
||||
upstream_status_code: finalization
|
||||
.upstream_status_code
|
||||
.or(observed_status_code),
|
||||
client_response_id: self.observation.client_response_id.clone(),
|
||||
upstream_request_id,
|
||||
upstream_response_id: self.observation.upstream_response_id.clone(),
|
||||
upstream_transport,
|
||||
http_version,
|
||||
websocket_pool,
|
||||
service_tier,
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
|
||||
error: Some(finalization.error),
|
||||
provider_error_code: finalization.provider_error_code,
|
||||
raw_upstream_error: finalization.raw_upstream_error,
|
||||
failure_observation: finalization.observation,
|
||||
retry_after_ms: finalization.retry_after_ms,
|
||||
usage: self.observation.usage.clone(),
|
||||
image_generation_succeeded: self.image_generation_succeeded(),
|
||||
cost: self.observation.cost.clone(),
|
||||
timings: self.observation.timings.clone(),
|
||||
completed_at,
|
||||
}),
|
||||
)
|
||||
self.persist_finalization(ModelRequestFinalization {
|
||||
request_id: self.request_id.clone(),
|
||||
outcome: finalization.outcome,
|
||||
send_state: finalization.send_state,
|
||||
attempt_count: self.attempts,
|
||||
downstream_committed_at: self.downstream_committed_at,
|
||||
client_status_code: self.client_status_code,
|
||||
upstream_status_code: finalization.upstream_status_code.or(observed_status_code),
|
||||
client_response_id: self.observation.client_response_id.clone(),
|
||||
upstream_request_id: finalization.upstream_request_id.or(upstream_request_id),
|
||||
upstream_response_id: self.observation.upstream_response_id.clone(),
|
||||
upstream_transport,
|
||||
http_version,
|
||||
websocket_pool,
|
||||
service_tier,
|
||||
provider_metadata_json,
|
||||
diagnostic_trace_json: self.trace.snapshot().map(|value| value.to_string()),
|
||||
error: Some(finalization.error),
|
||||
provider_error_code: finalization.provider_error_code,
|
||||
raw_upstream_error: finalization.raw_upstream_error,
|
||||
failure_observation: finalization.observation,
|
||||
retry_after_ms: finalization.retry_after_ms,
|
||||
usage: self.observation.usage.clone(),
|
||||
image_generation_succeeded: self.image_generation_succeeded(),
|
||||
cost: self.observation.cost.clone(),
|
||||
timings: self.observation.timings.clone(),
|
||||
completed_at,
|
||||
})
|
||||
.await;
|
||||
self.finalized = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn persist_finalization(&mut self, finalization: ModelRequestFinalization) {
|
||||
self.finalized_at = Some(finalization.completed_at);
|
||||
let mut persisted = self.request_persisted;
|
||||
// 首次合并写失败后不能以零 attempt 补行;建流前也只有确定未发送的请求可按零次收敛。
|
||||
let request = if !persisted
|
||||
&& finalization.attempt_count == 0
|
||||
&& finalization.send_state == UpstreamSendState::NotSent
|
||||
{
|
||||
self.pending_request.take()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let store = Arc::clone(self.engine.store());
|
||||
let request_id = self.request_id.clone();
|
||||
// 终态写入由会话持有;取消事件等待只暂停同一个 future,不重建请求或覆盖原失败。
|
||||
self.finalization = Some(RequestFinalization::Pending(Box::pin(async move {
|
||||
if let Some(request) = request {
|
||||
persisted = best_effort_store_write(
|
||||
"create_model_request",
|
||||
&request_id,
|
||||
store.create_model_request(request),
|
||||
)
|
||||
.await
|
||||
.is_some();
|
||||
}
|
||||
if persisted {
|
||||
best_effort_store_write(
|
||||
"finalize_model_request",
|
||||
&request_id,
|
||||
store.finalize_model_request(finalization),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
persisted
|
||||
})));
|
||||
self.resume_finalization().await;
|
||||
}
|
||||
|
||||
async fn resume_finalization(&mut self) {
|
||||
if let Some(RequestFinalization::Pending(write)) = self.finalization.as_mut() {
|
||||
self.request_persisted = write.await;
|
||||
self.finalization = Some(RequestFinalization::Complete);
|
||||
}
|
||||
}
|
||||
|
||||
fn current_transport_observation(&self) -> (Option<String>, Option<String>, Option<String>) {
|
||||
let Some(current) = self.current.as_ref() else {
|
||||
return (None, None, None);
|
||||
@@ -1908,6 +1985,7 @@ fn record_trace_error(trace: &TraceContext, error: &ProviderError) {
|
||||
"stage": diagnostic.stage(), "code": diagnostic.code(), "message": diagnostic.as_str(),
|
||||
})),
|
||||
"upstreamStatus": error.upstream_status(),
|
||||
"upstreamRequestId": error.upstream_request_id().map(|id| id.as_str()),
|
||||
"upstreamCode": error.upstream_code().map(|code| code.as_str()),
|
||||
"rawError": error.raw_upstream_error().map(|raw| {
|
||||
let value = serde_json::from_str(raw.as_str()).unwrap_or_else(|_| json!(raw.as_str()));
|
||||
|
||||
@@ -37,7 +37,8 @@ use crate::lifecycle::CancellationToken;
|
||||
use crate::operation::{Operation, ProviderSessionState};
|
||||
use crate::policy::{ClientApiKeyId, ClientPolicy};
|
||||
use crate::routing::{
|
||||
PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot, UpstreamModelId,
|
||||
ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId, RoutingContext,
|
||||
RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
use crate::runtime::RuntimeSnapshotHandle;
|
||||
|
||||
@@ -173,8 +174,11 @@ pub trait ExecutionSession: Send {
|
||||
&mut self,
|
||||
client_status_code: u16,
|
||||
) -> BoxFuture<'_, Result<(), EngineError>>;
|
||||
/// 执行已终结且结算、准入释放均已返回,协议层才可以放弃清理责任。
|
||||
/// 结算失败仍由 Store 保留费用重试,释放失败仍按租约 TTL 收敛。
|
||||
fn is_finalized(&self) -> bool;
|
||||
fn cancel(&self);
|
||||
/// 将会话交给宿主持续驱动清理;取消请求事件的等待不会丢弃已开始的结算。
|
||||
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()>;
|
||||
}
|
||||
|
||||
@@ -184,8 +188,13 @@ pub trait ExecutionService: Send + Sync {
|
||||
plaintext: &str,
|
||||
) -> Result<AuthenticatedClient, ClientAuthenticationError>;
|
||||
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId>;
|
||||
fn public_model_profiles(&self, _client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
Vec::new()
|
||||
fn client_model_catalog<'a>(
|
||||
&'a self,
|
||||
_client: &'a AuthenticatedClient,
|
||||
_protocol: &'a str,
|
||||
_client_version: &'a str,
|
||||
) -> BoxFuture<'a, Result<Vec<PublicModelDescriptor>, ProviderCatalogUnavailable>> {
|
||||
Box::pin(async { Err(ProviderCatalogUnavailable) })
|
||||
}
|
||||
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool;
|
||||
fn start(
|
||||
@@ -917,10 +926,77 @@ impl ExecutionService for DefaultExecutionService {
|
||||
.public_models_for_scope(client.policy.account_scope())
|
||||
}
|
||||
|
||||
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
client
|
||||
.snapshot
|
||||
.public_model_profiles_for_scope(client.policy.account_scope())
|
||||
fn client_model_catalog<'a>(
|
||||
&'a self,
|
||||
client: &'a AuthenticatedClient,
|
||||
protocol: &'a str,
|
||||
client_version: &'a str,
|
||||
) -> BoxFuture<'a, Result<Vec<PublicModelDescriptor>, ProviderCatalogUnavailable>> {
|
||||
Box::pin(async move {
|
||||
let scope = client.policy.account_scope();
|
||||
let mut result = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
for kind in scope.provider_kinds() {
|
||||
let provider = self.providers.get(kind).ok_or(ProviderCatalogUnavailable)?;
|
||||
let Some(models) = provider
|
||||
.query_client_model_catalog(scope, protocol, client_version)
|
||||
.await?
|
||||
else {
|
||||
for profile in client.snapshot.public_model_profiles_for_provider(kind) {
|
||||
if seen.insert(profile.model().clone()) {
|
||||
result.push(PublicModelDescriptor::Adapted(profile));
|
||||
}
|
||||
}
|
||||
continue;
|
||||
};
|
||||
// 保留上游顺序;映射只选择完整的目标对象,不能跨账号/模型混拼字段。
|
||||
let by_id = models
|
||||
.iter()
|
||||
.map(|entry| (entry.model.as_str(), entry))
|
||||
.collect::<std::collections::BTreeMap<_, _>>();
|
||||
for entry in &models {
|
||||
let target = client.snapshot.mapped_model(entry.model.as_str());
|
||||
let Some(source) = by_id.get(target.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
let model = PublicModelId::new(entry.model.as_str().to_owned())
|
||||
.map_err(|_| ProviderCatalogUnavailable)?;
|
||||
// 原生目录可能比路由快照更新,不能向客户端公布当前已知不可路由的模型。
|
||||
if !client
|
||||
.snapshot
|
||||
.contains_public_model_for_provider(&model, kind)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if seen.insert(model.clone()) {
|
||||
result.push(PublicModelDescriptor::Native {
|
||||
model,
|
||||
payload: source.payload.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
for model in client.snapshot.public_models_for_provider(kind) {
|
||||
let target = client.snapshot.mapped_model(model.as_str());
|
||||
if target == model.as_str() || seen.contains(&model) {
|
||||
continue;
|
||||
}
|
||||
if !client
|
||||
.snapshot
|
||||
.contains_public_model_for_provider(&model, kind)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if let Some(source) = by_id.get(target.as_str()) {
|
||||
seen.insert(model.clone());
|
||||
result.push(PublicModelDescriptor::Native {
|
||||
model,
|
||||
payload: source.payload.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
})
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
|
||||
@@ -978,8 +1054,9 @@ impl AdmissionLease {
|
||||
}
|
||||
|
||||
struct DefaultExecutionSession {
|
||||
core: Option<ResponseExecutionSession<dyn ExecutionStore>>,
|
||||
core: ResponseExecutionSession<dyn ExecutionStore>,
|
||||
admission: Option<AdmissionLease>,
|
||||
cleanup: Option<BoxFuture<'static, ()>>,
|
||||
circuits: Arc<dyn ProviderCircuitPort>,
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
observed_provider_outcomes: usize,
|
||||
@@ -996,8 +1073,9 @@ impl DefaultExecutionSession {
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
core: Some(core),
|
||||
core,
|
||||
admission: Some(admission),
|
||||
cleanup: None,
|
||||
circuits,
|
||||
continuation,
|
||||
observed_provider_outcomes: 0,
|
||||
@@ -1006,33 +1084,29 @@ impl DefaultExecutionSession {
|
||||
}
|
||||
}
|
||||
|
||||
fn core_mut(
|
||||
&mut self,
|
||||
) -> Result<&mut ResponseExecutionSession<dyn ExecutionStore>, EngineError> {
|
||||
self.core.as_mut().ok_or(EngineError::InvalidDeliveryState)
|
||||
}
|
||||
|
||||
async fn settle_if_finalized(&mut self) {
|
||||
if self
|
||||
.core
|
||||
.as_ref()
|
||||
.is_some_and(ResponseExecutionSession::is_finalized)
|
||||
if self.core.is_finalized()
|
||||
&& let Some(admission) = self.admission.take()
|
||||
{
|
||||
if let Some(budget) = self.budget.take()
|
||||
&& let Some(core) = &self.core
|
||||
{
|
||||
settle_budget(budget.as_ref(), core.budget_charge()).await;
|
||||
}
|
||||
admission.release().await;
|
||||
let budget = self.budget.take();
|
||||
let charge = self.core.budget_charge();
|
||||
// 在首次 await 前把完整清理责任留在会话内。事件等待被取消后,后续 poll
|
||||
// 或 detach 继续同一个 future,既不丢失费用,也不重启已完成的结算。
|
||||
self.cleanup = Some(Box::pin(async move {
|
||||
if let Some(budget) = budget {
|
||||
settle_budget(budget.as_ref(), charge).await;
|
||||
}
|
||||
admission.release().await;
|
||||
}));
|
||||
}
|
||||
if let Some(cleanup) = self.cleanup.as_mut() {
|
||||
cleanup.await;
|
||||
self.cleanup = None;
|
||||
}
|
||||
}
|
||||
|
||||
async fn observe_provider_outcomes(&mut self) {
|
||||
let Some(core) = self.core.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let outcomes = core.provider_attempt_outcomes();
|
||||
let outcomes = self.core.provider_attempt_outcomes();
|
||||
let new_outcomes = outcomes
|
||||
.get(self.observed_provider_outcomes..)
|
||||
.unwrap_or_default()
|
||||
@@ -1048,11 +1122,7 @@ impl DefaultExecutionSession {
|
||||
let Some(state) = state else {
|
||||
return;
|
||||
};
|
||||
let Some(pin) = self
|
||||
.core
|
||||
.as_ref()
|
||||
.and_then(|core| core.native_continuation_pin(state))
|
||||
else {
|
||||
let Some(pin) = self.core.native_continuation_pin(state) else {
|
||||
return;
|
||||
};
|
||||
self.continuation_recorded = true;
|
||||
@@ -1060,47 +1130,27 @@ impl DefaultExecutionSession {
|
||||
}
|
||||
|
||||
async fn finalize_detached(&mut self) {
|
||||
let Some(mut core) = self.core.take() else {
|
||||
return;
|
||||
};
|
||||
core.cancel();
|
||||
if !core.is_finalized()
|
||||
&& let Err(error) = core.cancel_and_finalize().await
|
||||
{
|
||||
if let Err(error) = self.core.cancel_and_finalize().await {
|
||||
tracing::warn!(%error, "Detached execution 终态收敛失败");
|
||||
}
|
||||
let pending = core
|
||||
.provider_attempt_outcomes()
|
||||
.get(self.observed_provider_outcomes..)
|
||||
.unwrap_or_default();
|
||||
publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await;
|
||||
if let Some(budget) = self.budget.take() {
|
||||
settle_budget(budget.as_ref(), core.budget_charge()).await;
|
||||
}
|
||||
if let Some(admission) = self.admission.take() {
|
||||
admission.release().await;
|
||||
}
|
||||
self.observe_provider_outcomes().await;
|
||||
self.settle_if_finalized().await;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DefaultExecutionSession {
|
||||
fn drop(&mut self) {
|
||||
if let Some(core) = &self.core {
|
||||
core.cancel();
|
||||
}
|
||||
self.core.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionSession for DefaultExecutionSession {
|
||||
fn trace(&self) -> crate::diagnostics::TraceContext {
|
||||
self.core
|
||||
.as_ref()
|
||||
.map(ResponseExecutionSession::trace)
|
||||
.unwrap_or_default()
|
||||
self.core.trace()
|
||||
}
|
||||
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
|
||||
Box::pin(async move {
|
||||
let result = self.core_mut()?.next_event().await;
|
||||
let result = self.core.next_event().await;
|
||||
if let Ok(Some(event)) = result.as_ref() {
|
||||
self.record_continuation(event.session_update()).await;
|
||||
}
|
||||
@@ -1112,7 +1162,7 @@ impl ExecutionSession for DefaultExecutionSession {
|
||||
|
||||
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
|
||||
Box::pin(async move {
|
||||
let result = self.core_mut()?.collect_uncommitted().await;
|
||||
let result = self.core.collect_uncommitted().await;
|
||||
if let Ok(events) = result.as_ref() {
|
||||
let state = events.iter().find_map(ProviderEvent::session_update);
|
||||
self.record_continuation(state).await;
|
||||
@@ -1124,16 +1174,11 @@ impl ExecutionSession for DefaultExecutionSession {
|
||||
}
|
||||
|
||||
fn response_headers(&self) -> &[ProviderResponseHeader] {
|
||||
self.core
|
||||
.as_ref()
|
||||
.map(ResponseExecutionSession::response_headers)
|
||||
.unwrap_or_default()
|
||||
self.core.response_headers()
|
||||
}
|
||||
|
||||
fn response_status_code(&self) -> Option<u16> {
|
||||
self.core
|
||||
.as_ref()
|
||||
.and_then(ResponseExecutionSession::response_status_code)
|
||||
self.core.response_status_code()
|
||||
}
|
||||
|
||||
fn commit_downstream(
|
||||
@@ -1141,7 +1186,7 @@ impl ExecutionSession for DefaultExecutionSession {
|
||||
client_status_code: Option<u16>,
|
||||
) -> BoxFuture<'_, Result<(), EngineError>> {
|
||||
Box::pin(async move {
|
||||
let result = self.core_mut()?.commit_downstream(client_status_code).await;
|
||||
let result = self.core.commit_downstream(client_status_code).await;
|
||||
self.observe_provider_outcomes().await;
|
||||
self.settle_if_finalized().await;
|
||||
result
|
||||
@@ -1153,10 +1198,7 @@ impl ExecutionSession for DefaultExecutionSession {
|
||||
client_status_code: u16,
|
||||
) -> BoxFuture<'_, Result<(), EngineError>> {
|
||||
Box::pin(async move {
|
||||
let result = self
|
||||
.core_mut()?
|
||||
.record_client_status(client_status_code)
|
||||
.await;
|
||||
let result = self.core.record_client_status(client_status_code).await;
|
||||
self.observe_provider_outcomes().await;
|
||||
self.settle_if_finalized().await;
|
||||
result
|
||||
@@ -1164,15 +1206,11 @@ impl ExecutionSession for DefaultExecutionSession {
|
||||
}
|
||||
|
||||
fn is_finalized(&self) -> bool {
|
||||
self.core
|
||||
.as_ref()
|
||||
.is_none_or(ResponseExecutionSession::is_finalized)
|
||||
self.core.is_finalized() && self.admission.is_none() && self.cleanup.is_none()
|
||||
}
|
||||
|
||||
fn cancel(&self) {
|
||||
if let Some(core) = &self.core {
|
||||
core.cancel();
|
||||
}
|
||||
self.core.cancel();
|
||||
}
|
||||
|
||||
fn detach_finalize(mut self: Box<Self>) -> BoxFuture<'static, ()> {
|
||||
|
||||
@@ -8,7 +8,7 @@ use std::task::{Context, Poll};
|
||||
use std::time::Instant;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::Stream;
|
||||
use futures::{Stream, future::BoxFuture};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::account::{
|
||||
@@ -21,7 +21,8 @@ use crate::identity::ProviderKind;
|
||||
use crate::operation::Operation;
|
||||
use crate::policy::ClientApiKeyId;
|
||||
use crate::routing::{
|
||||
ModelCapabilities, ModelPresentation, ProviderCandidate, PublicModelId, UpstreamModelId,
|
||||
ProviderCandidate, ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
|
||||
ProviderModelCapabilities, PublicModelId, UpstreamModelId,
|
||||
};
|
||||
use crate::upstream::OpaqueUpstreamValue;
|
||||
use crate::upstream::{UpstreamSendState, UpstreamTransport};
|
||||
@@ -437,64 +438,6 @@ pub struct ContinuationRequestObservation {
|
||||
pub requested: bool,
|
||||
}
|
||||
|
||||
/// Provider 实时目录编译后的单模型能力。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderModelCapabilities {
|
||||
upstream_model: UpstreamModelId,
|
||||
capabilities: ModelCapabilities,
|
||||
presentation: Option<ModelPresentation>,
|
||||
}
|
||||
|
||||
/// Provider 实时目录成功发布后的进程内单调代次。
|
||||
///
|
||||
/// 代次只表达“目录内容已经变化”,不承载模型、ETag 或 Provider 私有数据。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct ProviderCatalogGeneration(u64);
|
||||
|
||||
impl ProviderCatalogGeneration {
|
||||
#[must_use]
|
||||
pub const fn new(value: u64) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u64 {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderModelCapabilities {
|
||||
#[must_use]
|
||||
pub const fn new(upstream_model: UpstreamModelId, capabilities: ModelCapabilities) -> Self {
|
||||
Self {
|
||||
upstream_model,
|
||||
capabilities,
|
||||
presentation: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
|
||||
self.presentation = Some(presentation);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn upstream_model(&self) -> &UpstreamModelId {
|
||||
&self.upstream_model
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn capabilities(&self) -> &ModelCapabilities {
|
||||
&self.capabilities
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn presentation(&self) -> Option<&ModelPresentation> {
|
||||
self.presentation.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderRequest {
|
||||
/// 绑定 operation 与请求计划中冻结的 Provider 候选。
|
||||
#[must_use]
|
||||
@@ -560,11 +503,30 @@ pub trait Provider: Send + Sync {
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError>;
|
||||
|
||||
/// 读取当前客户端协议的原生目录,必须限定到认证时冻结的账号范围。
|
||||
/// `None` 表示不提供该协议的原生目录;读取失败不能伪装成不支持。
|
||||
async fn query_client_model_catalog(
|
||||
&self,
|
||||
_scope: &crate::account::scope::FrozenAccountScope,
|
||||
_protocol: &str,
|
||||
_client_version: &str,
|
||||
) -> Result<
|
||||
Option<Vec<crate::routing::ProviderModelDescriptor>>,
|
||||
crate::routing::ProviderCatalogUnavailable,
|
||||
> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// 选择一个未被排除的资源并返回 cold [`ProviderStream`]。
|
||||
///
|
||||
/// 返回成功、返回错误或准备 future 被取消前,均不得发送本次请求的上游握手或业务
|
||||
/// 载荷,也不得启动可独立完成这些发送的后台任务。发送只在返回的流被 poll 后发生,
|
||||
/// 保证 Core 能在真实出站前校验账号范围并登记 attempt。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 没有可用资源、请求无效或在产生 stream 前失败时返回 Provider 错误。
|
||||
/// 没有可用资源、请求无效或准备失败时返回 `NotSent` 错误;
|
||||
/// 可能已发送的失败必须通过 stream 返回,不得降级发送事实。
|
||||
async fn execute(
|
||||
&self,
|
||||
request: ProviderRequest,
|
||||
@@ -586,15 +548,6 @@ pub enum RegistryError {
|
||||
},
|
||||
}
|
||||
|
||||
/// Registry 实时目录查询错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ProviderCatalogError {
|
||||
#[error("provider `{provider}` is not registered")]
|
||||
NotRegistered { provider: String },
|
||||
#[error("provider model catalog query failed")]
|
||||
Query(#[source] ProviderError),
|
||||
}
|
||||
|
||||
/// 唯一保存 `Arc<dyn Provider>` 的异构注册表。
|
||||
#[derive(Default)]
|
||||
pub struct ProviderRegistryBuilder {
|
||||
@@ -665,11 +618,6 @@ impl ProviderRegistry {
|
||||
self.providers.get(provider)
|
||||
}
|
||||
|
||||
/// 返回全部已注册 Provider。
|
||||
pub fn provider_kinds(&self) -> impl Iterator<Item = &ProviderKind> {
|
||||
self.providers.keys()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn request_observation(
|
||||
&self,
|
||||
@@ -684,35 +632,6 @@ impl ProviderRegistry {
|
||||
})
|
||||
}
|
||||
|
||||
/// 读取全部 Provider 的目录代次,用于 Core 快照稳定性校验与对账。
|
||||
#[must_use]
|
||||
pub fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
self.providers
|
||||
.iter()
|
||||
.map(|(kind, provider)| (kind.clone(), provider.catalog_generation()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 通过编译期 Provider adapter 查询并编译实时能力目录。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Provider 未注册或目录查询失败时返回错误。
|
||||
pub async fn query_model_capabilities(
|
||||
&self,
|
||||
provider_kind: &ProviderKind,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderCatalogError> {
|
||||
let provider = self.providers.get(provider_kind).ok_or_else(|| {
|
||||
ProviderCatalogError::NotRegistered {
|
||||
provider: provider_kind.as_str().to_owned(),
|
||||
}
|
||||
})?;
|
||||
provider
|
||||
.query_model_capabilities()
|
||||
.await
|
||||
.map_err(ProviderCatalogError::Query)
|
||||
}
|
||||
|
||||
/// 判断 Provider 是否已注册。
|
||||
#[must_use]
|
||||
pub fn contains(&self, provider: &ProviderKind) -> bool {
|
||||
@@ -731,3 +650,26 @@ impl ProviderRegistry {
|
||||
self.providers.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCatalogPort for ProviderRegistry {
|
||||
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
self.providers
|
||||
.iter()
|
||||
.map(|(kind, provider)| (kind.clone(), provider.catalog_generation()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn query_model_capabilities(
|
||||
&self,
|
||||
provider_kind: &ProviderKind,
|
||||
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
|
||||
let provider = self.providers.get(provider_kind);
|
||||
Box::pin(async move {
|
||||
provider
|
||||
.ok_or(ProviderCatalogUnavailable)?
|
||||
.query_model_capabilities()
|
||||
.await
|
||||
.map_err(|_| ProviderCatalogUnavailable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,6 +46,8 @@ pub enum ProviderErrorKind {
|
||||
Protocol,
|
||||
/// Provider 暂不可用。
|
||||
Unavailable,
|
||||
/// 上游明确拒绝当前请求的模型容量;不证明整个 Provider 的基础设施故障。
|
||||
UpstreamCapacityUnavailable,
|
||||
/// 请求被取消。
|
||||
Cancelled,
|
||||
/// 进程终止后由恢复流程收敛。
|
||||
@@ -71,6 +73,7 @@ impl ProviderErrorKind {
|
||||
Self::Transport => "transport",
|
||||
Self::Protocol => "protocol",
|
||||
Self::Unavailable => "unavailable",
|
||||
Self::UpstreamCapacityUnavailable => "upstream_capacity_unavailable",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::ProcessTerminated => "process_terminated",
|
||||
}
|
||||
@@ -105,6 +108,15 @@ pub enum ContinuationRecoveryDisposition {
|
||||
pub enum PreDeliveryRetry {
|
||||
/// 排除本次账号,按普通调度策略重新选号。
|
||||
AccountRotation,
|
||||
/// 上游明确拒绝后的同账号退避;耗尽后仍由普通安全重放规则决定是否换号。
|
||||
SameAccountTransientRetry {
|
||||
/// 每个账号在本请求中的重试上限,同时受总路由预算约束。
|
||||
max_retries: NonZeroU32,
|
||||
/// 指数退避的起始间隔。
|
||||
initial_delay: Duration,
|
||||
/// 单次等待上限,避免上游提示把请求拖入长时间等待。
|
||||
max_delay: Duration,
|
||||
},
|
||||
/// 固定本次账号,并按 Provider 给出的序号重试当前传输。
|
||||
SameAccountTransportRetry {
|
||||
/// Provider-owned 传输重试序号,从 1 开始。
|
||||
@@ -409,7 +421,7 @@ pub struct ProviderError {
|
||||
continuation_recovery_disposition: Option<ContinuationRecoveryDisposition>,
|
||||
failure_observation: Option<Box<ProviderErrorFailureObservation>>,
|
||||
replay_safe: bool,
|
||||
pre_delivery_retry: Option<PreDeliveryRetry>,
|
||||
pre_delivery_retry: Option<Box<PreDeliveryRetry>>,
|
||||
credential_recovery_required: bool,
|
||||
retry_same_account: bool,
|
||||
sensitive_context_redacted: bool,
|
||||
@@ -557,8 +569,8 @@ impl ProviderError {
|
||||
/// Provider 选择了“客户端无感恢复优先”的传输策略。Core 在 continuation、
|
||||
/// 指定账号或下游已经进入提交状态时必须忽略它。
|
||||
#[must_use]
|
||||
pub const fn with_pre_delivery_retry(mut self) -> Self {
|
||||
self.pre_delivery_retry = Some(PreDeliveryRetry::AccountRotation);
|
||||
pub fn with_pre_delivery_retry(mut self) -> Self {
|
||||
self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::AccountRotation));
|
||||
self
|
||||
}
|
||||
|
||||
@@ -567,14 +579,14 @@ impl ProviderError {
|
||||
/// 与普通换号恢复相同,Core 仍会拒绝 continuation、指定账号、预算耗尽或
|
||||
/// 下游已提交后的隐藏重放。
|
||||
#[must_use]
|
||||
pub const fn with_pre_delivery_transport_fallback(mut self) -> Self {
|
||||
pub fn with_pre_delivery_transport_fallback(mut self) -> Self {
|
||||
self.set_pre_delivery_transport_fallback();
|
||||
self
|
||||
}
|
||||
|
||||
/// 原地设置同账号备用传输恢复,保留当前错误携带的原客户端响应。
|
||||
pub const fn set_pre_delivery_transport_fallback(&mut self) {
|
||||
self.pre_delivery_retry = Some(PreDeliveryRetry::SameAccountTransportFallback);
|
||||
pub fn set_pre_delivery_transport_fallback(&mut self) {
|
||||
self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountTransportFallback));
|
||||
}
|
||||
|
||||
/// 允许 Core 在首个客户端事件前固定原账号并重试当前 Provider 传输。
|
||||
@@ -582,7 +594,7 @@ impl ProviderError {
|
||||
/// 重试预算和退避策略由 Provider 所属协议定义;Core 只负责同账号钉选、
|
||||
/// deadline/取消边界以及下游 commit barrier。
|
||||
#[must_use]
|
||||
pub const fn with_pre_delivery_transport_retry(
|
||||
pub fn with_pre_delivery_transport_retry(
|
||||
mut self,
|
||||
retry_index: NonZeroU32,
|
||||
delay: Duration,
|
||||
@@ -592,13 +604,11 @@ impl ProviderError {
|
||||
}
|
||||
|
||||
/// 原地设置同账号当前传输重试,保留当前错误携带的原客户端响应。
|
||||
pub const fn set_pre_delivery_transport_retry(
|
||||
&mut self,
|
||||
retry_index: NonZeroU32,
|
||||
delay: Duration,
|
||||
) {
|
||||
self.pre_delivery_retry =
|
||||
Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, delay });
|
||||
pub fn set_pre_delivery_transport_retry(&mut self, retry_index: NonZeroU32, delay: Duration) {
|
||||
self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountTransportRetry {
|
||||
retry_index,
|
||||
delay,
|
||||
}));
|
||||
}
|
||||
|
||||
/// 要求 Core 在账号凭据已恢复后仅对原账号重放一次。
|
||||
@@ -608,6 +618,22 @@ impl ProviderError {
|
||||
self
|
||||
}
|
||||
|
||||
/// 请求有界同账号退避;该意图本身不提供安全重放证据。
|
||||
#[must_use]
|
||||
pub fn with_transient_retry(
|
||||
mut self,
|
||||
max_retries: NonZeroU32,
|
||||
initial_delay: Duration,
|
||||
max_delay: Duration,
|
||||
) -> Self {
|
||||
self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountTransientRetry {
|
||||
max_retries,
|
||||
initial_delay,
|
||||
max_delay,
|
||||
}));
|
||||
self
|
||||
}
|
||||
|
||||
/// 标记上游明确拒绝应先触发 OAuth 恢复,而不是直接废弃账号。
|
||||
#[must_use]
|
||||
pub const fn with_credential_recovery(mut self) -> Self {
|
||||
@@ -748,7 +774,10 @@ impl ProviderError {
|
||||
/// 返回下游提交前由 Provider 选择的隐藏恢复方式。
|
||||
#[must_use]
|
||||
pub const fn pre_delivery_retry(&self) -> Option<PreDeliveryRetry> {
|
||||
self.pre_delivery_retry
|
||||
match self.pre_delivery_retry.as_ref() {
|
||||
Some(retry) => Some(**retry),
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回 Provider 是否已完成凭据恢复并要求原账号重放。
|
||||
@@ -822,7 +851,7 @@ impl Clone for ProviderError {
|
||||
continuation_recovery_disposition: self.continuation_recovery_disposition,
|
||||
failure_observation: self.failure_observation.clone(),
|
||||
replay_safe: self.replay_safe,
|
||||
pre_delivery_retry: self.pre_delivery_retry,
|
||||
pre_delivery_retry: self.pre_delivery_retry.clone(),
|
||||
credential_recovery_required: self.credential_recovery_required,
|
||||
retry_same_account: self.retry_same_account,
|
||||
sensitive_context_redacted: self.sensitive_context_redacted,
|
||||
@@ -1064,6 +1093,7 @@ impl GatewayError {
|
||||
ProviderErrorKind::Transport
|
||||
| ProviderErrorKind::Protocol
|
||||
| ProviderErrorKind::Unavailable
|
||||
| ProviderErrorKind::UpstreamCapacityUnavailable
|
||||
| ProviderErrorKind::ProcessTerminated => Self::new(
|
||||
GatewayErrorKind::UpstreamUnavailable,
|
||||
"upstream service is unavailable",
|
||||
|
||||
@@ -521,7 +521,7 @@ impl fmt::Debug for ProviderResponseHeader {
|
||||
}
|
||||
}
|
||||
|
||||
/// Core 消费的实际上游响应事实。
|
||||
/// Core 消费的 Provider 执行观测,不会原样进入客户端响应。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderResponseObservation {
|
||||
transport: UpstreamTransport,
|
||||
@@ -575,7 +575,7 @@ impl ProviderResponseObservation {
|
||||
self
|
||||
}
|
||||
|
||||
/// 附加 Provider 在响应中确认的服务档位。
|
||||
/// 附加 Provider 确定的统计档位,Core 不重新解释其来源。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@@ -590,7 +590,7 @@ impl ProviderResponseObservation {
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// 附加合法服务档位;不适合持久化的上游观测值会被忽略。
|
||||
/// 附加 Provider 确定的统计档位;不适合持久化的值会被忽略。
|
||||
#[must_use]
|
||||
pub fn with_service_tier_if_valid(mut self, service_tier: impl Into<String>) -> Self {
|
||||
let service_tier = service_tier.into();
|
||||
|
||||
@@ -140,7 +140,7 @@ pub async fn initialize(
|
||||
.map_err(|_| CoreError::AdmissionRecoveryUnavailable)?;
|
||||
let compiler = Arc::new(RuntimeSnapshotCompiler::new(
|
||||
Arc::clone(&ports.snapshots),
|
||||
providers.clone(),
|
||||
Arc::new(providers.clone()),
|
||||
));
|
||||
let initial = compiler
|
||||
.compile()
|
||||
|
||||
@@ -225,10 +225,10 @@ impl fmt::Debug for ProtocolPayload {
|
||||
}
|
||||
}
|
||||
|
||||
/// 客户端协议交给 Provider 的未经重编码 JSON 正文。
|
||||
/// 协议 adapter 与 Provider 之间不透明传递的 JSON 正文。
|
||||
///
|
||||
/// 协议 adapter 只确定路由端点;Core 保留原始字节与非 wire 上下文,
|
||||
/// 不验证、解析或改写正文。
|
||||
/// Core 保留字节与非 wire 上下文,不验证、解析或改写正文;
|
||||
/// 请求载荷与原生模型目录均由协议 adapter 和对应 Provider 共同拥有语义。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct RawJsonPayload {
|
||||
protocol: String,
|
||||
@@ -237,7 +237,7 @@ pub struct RawJsonPayload {
|
||||
}
|
||||
|
||||
impl RawJsonPayload {
|
||||
/// 创建未经重编码的协议 JSON 正文。
|
||||
/// 创建不透明的协议 JSON 正文。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@@ -260,7 +260,7 @@ impl RawJsonPayload {
|
||||
&self.protocol
|
||||
}
|
||||
|
||||
/// 返回 adapter 收到的原始 JSON 字节。
|
||||
/// 返回协议正文的 JSON 字节。
|
||||
#[must_use]
|
||||
pub const fn body(&self) -> &Bytes {
|
||||
&self.body
|
||||
|
||||
@@ -53,13 +53,28 @@ impl PlaintextClientApiKey {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Key 为空、过长或包含控制字符时返回错误。
|
||||
/// Key 为空或无法作为 HTTP Bearer 值发送时返回错误。
|
||||
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
|
||||
let value = value.into();
|
||||
validate_text(&value, 512, false, None)?;
|
||||
Self::validate(&value)?;
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
/// 迁入的 Key 不限定前缀或长度;保持原值,仅校验 HTTP 可传输的非空可见 ASCII。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Key 为空或包含空白、控制字符、非 ASCII 字符时返回错误。
|
||||
pub fn validate(value: &str) -> Result<(), IdentifierError> {
|
||||
if value.is_empty() {
|
||||
return Err(IdentifierError::Empty);
|
||||
}
|
||||
if !value.bytes().all(|byte| byte.is_ascii_graphic()) {
|
||||
return Err(IdentifierError::InvalidFormat);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 仅借给同步认证器做常量时间比较。
|
||||
#[must_use]
|
||||
pub fn expose_for_auth(&self) -> &str {
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
//! 快照编译与客户端读取需要的 Provider 目录合同,不暴露执行或账号选择能力。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::identity::ProviderKind;
|
||||
use crate::operation::RawJsonPayload;
|
||||
|
||||
use super::{
|
||||
ModelCapabilities, ModelPresentation, PublicModelId, PublicModelProfile, UpstreamModelId,
|
||||
};
|
||||
|
||||
/// Provider 原生客户端目录条目;Core 只解释模型标识,不解释协议正文。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderModelDescriptor {
|
||||
pub model: UpstreamModelId,
|
||||
pub payload: RawJsonPayload,
|
||||
}
|
||||
|
||||
/// 对外目录保留原生正文;只有没有原生目录的 Provider 才使用通用画像适配。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PublicModelDescriptor {
|
||||
Native {
|
||||
model: PublicModelId,
|
||||
payload: RawJsonPayload,
|
||||
},
|
||||
Adapted(PublicModelProfile),
|
||||
}
|
||||
|
||||
/// Provider 实时目录编译后的单模型能力。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderModelCapabilities {
|
||||
upstream_model: UpstreamModelId,
|
||||
capabilities: ModelCapabilities,
|
||||
presentation: Option<ModelPresentation>,
|
||||
}
|
||||
|
||||
impl ProviderModelCapabilities {
|
||||
#[must_use]
|
||||
pub const fn new(upstream_model: UpstreamModelId, capabilities: ModelCapabilities) -> Self {
|
||||
Self {
|
||||
upstream_model,
|
||||
capabilities,
|
||||
presentation: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
|
||||
self.presentation = Some(presentation);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn upstream_model(&self) -> &UpstreamModelId {
|
||||
&self.upstream_model
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn capabilities(&self) -> &ModelCapabilities {
|
||||
&self.capabilities
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn presentation(&self) -> Option<&ModelPresentation> {
|
||||
self.presentation.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 实时目录成功发布后的进程内单调代次。
|
||||
///
|
||||
/// 代次只表达“目录内容已经变化”,不承载模型、ETag 或 Provider 私有数据。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct ProviderCatalogGeneration(u64);
|
||||
|
||||
impl ProviderCatalogGeneration {
|
||||
#[must_use]
|
||||
pub const fn new(value: u64) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u64 {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// 目录未知,不能据此判定某个模型不存在;不携带执行错误或上游协议细节。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
||||
#[error("provider model catalog is unavailable")]
|
||||
pub struct ProviderCatalogUnavailable;
|
||||
|
||||
/// 快照编译与对账使用的对象安全目录端口。
|
||||
pub trait ProviderCatalogPort: Send + Sync {
|
||||
/// 返回全部已注册 Provider 的目录代次;注册集合在初始化后保持不变。
|
||||
///
|
||||
/// 即使某个目录暂时不可读,也必须保留它的 Provider 与最近成功发布的代次。
|
||||
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration>;
|
||||
|
||||
/// 查询指定 Provider 的模型事实;成功的空列表表示已知无模型,失败表示未知。
|
||||
fn query_model_capabilities(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>>;
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
//! Provider、模型目录、精确模型映射与请求级候选计划。
|
||||
|
||||
mod catalog;
|
||||
pub mod snapshot;
|
||||
|
||||
pub use crate::account::scope::{
|
||||
@@ -7,6 +8,10 @@ pub use crate::account::scope::{
|
||||
FrozenAccountScope, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory,
|
||||
};
|
||||
pub use crate::identity::ProviderKind;
|
||||
pub use catalog::{
|
||||
ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
|
||||
ProviderModelCapabilities, ProviderModelDescriptor, PublicModelDescriptor,
|
||||
};
|
||||
pub use snapshot::RuntimeSnapshot;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
@@ -139,6 +144,7 @@ pub struct ModelPresentation {
|
||||
default_reasoning_effort: Option<String>,
|
||||
supported_reasoning_efforts: Vec<String>,
|
||||
context_window_tokens: Option<u64>,
|
||||
max_context_window_tokens: Option<u64>,
|
||||
image_input: bool,
|
||||
agent_tools: bool,
|
||||
parallel_tool_calls: bool,
|
||||
@@ -227,6 +233,16 @@ impl ModelPresentation {
|
||||
self
|
||||
}
|
||||
|
||||
/// 目录声明的可覆盖上限;客户端本地 context window 覆盖会被它钳制。
|
||||
#[must_use]
|
||||
pub const fn with_max_context_window_tokens(
|
||||
mut self,
|
||||
max_context_window_tokens: Option<u64>,
|
||||
) -> Self {
|
||||
self.max_context_window_tokens = max_context_window_tokens;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_image_input(mut self, image_input: bool) -> Self {
|
||||
self.image_input = image_input;
|
||||
@@ -295,6 +311,11 @@ impl ModelPresentation {
|
||||
self.context_window_tokens
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn max_context_window_tokens(&self) -> Option<u64> {
|
||||
self.max_context_window_tokens
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn image_input(&self) -> bool {
|
||||
self.image_input
|
||||
|
||||
@@ -8,7 +8,6 @@ use std::time::Duration;
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
|
||||
use crate::operation::Operation;
|
||||
use crate::policy::{
|
||||
ClientApiKeyId, ClientPolicy, CodexClientMinVersions, CodexClientVersion,
|
||||
@@ -18,8 +17,9 @@ use crate::validation::RoutingError;
|
||||
|
||||
use super::{
|
||||
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
|
||||
ProviderCandidate, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
|
||||
RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, UpstreamModelId,
|
||||
ProviderCandidate, ProviderCatalogGeneration, ProviderCatalogPort, ProviderKind, ProviderModel,
|
||||
PublicModelId, RoutingContext, RoutingGroupSnapshot, RoutingPlan, RuntimeAccount,
|
||||
RuntimeAccountDirectory, UpstreamModelId,
|
||||
};
|
||||
|
||||
const MAXIMUM_CATALOG_STABILITY_ATTEMPTS: usize = 4;
|
||||
@@ -212,13 +212,16 @@ pub enum RuntimeSnapshotCompileError {
|
||||
#[derive(Clone)]
|
||||
pub struct RuntimeSnapshotCompiler {
|
||||
store: Arc<dyn SnapshotStorePort>,
|
||||
providers: ProviderRegistry,
|
||||
catalogs: Arc<dyn ProviderCatalogPort>,
|
||||
}
|
||||
|
||||
impl RuntimeSnapshotCompiler {
|
||||
#[must_use]
|
||||
pub const fn new(store: Arc<dyn SnapshotStorePort>, providers: ProviderRegistry) -> Self {
|
||||
Self { store, providers }
|
||||
pub const fn new(
|
||||
store: Arc<dyn SnapshotStorePort>,
|
||||
catalogs: Arc<dyn ProviderCatalogPort>,
|
||||
) -> Self {
|
||||
Self { store, catalogs }
|
||||
}
|
||||
|
||||
pub(crate) fn store(&self) -> Arc<dyn SnapshotStorePort> {
|
||||
@@ -228,13 +231,13 @@ impl RuntimeSnapshotCompiler {
|
||||
pub(crate) fn provider_catalog_generations(
|
||||
&self,
|
||||
) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
self.providers.catalog_generations()
|
||||
self.catalogs.catalog_generations()
|
||||
}
|
||||
|
||||
/// 读取一个 revision,并为已注册 Provider 查询实时模型目录。
|
||||
pub async fn compile(&self) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
|
||||
for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS {
|
||||
let catalog_generations = self.providers.catalog_generations();
|
||||
let catalog_generations = self.catalogs.catalog_generations();
|
||||
let facts = self
|
||||
.store
|
||||
.load_snapshot_facts()
|
||||
@@ -243,8 +246,10 @@ impl RuntimeSnapshotCompiler {
|
||||
if facts.config_revision != facts.observed_current_revision {
|
||||
return Err(RuntimeSnapshotCompileError::RevisionChanged);
|
||||
}
|
||||
let snapshot = compile_runtime_snapshot(facts, &self.providers).await?;
|
||||
let observed_generations = self.providers.catalog_generations();
|
||||
let provider_kinds = catalog_generations.keys().cloned().collect();
|
||||
let snapshot =
|
||||
compile_runtime_snapshot(facts, self.catalogs.as_ref(), provider_kinds).await?;
|
||||
let observed_generations = self.catalogs.catalog_generations();
|
||||
if catalog_generations == observed_generations {
|
||||
return Ok(snapshot.with_provider_catalog_generations(observed_generations));
|
||||
}
|
||||
@@ -255,16 +260,16 @@ impl RuntimeSnapshotCompiler {
|
||||
|
||||
async fn compile_runtime_snapshot(
|
||||
facts: SnapshotFacts,
|
||||
providers: &ProviderRegistry,
|
||||
catalogs: &dyn ProviderCatalogPort,
|
||||
provider_kinds: Vec<ProviderKind>,
|
||||
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
|
||||
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
|
||||
let registered_providers = provider_kinds.iter().cloned().collect::<BTreeSet<_>>();
|
||||
|
||||
// 目录查询失败表示未知;查询成功后,即使为空,也必须与“已知缺少模型”区分。
|
||||
let mut provider_models = Vec::new();
|
||||
let mut known_provider_catalogs = BTreeSet::new();
|
||||
for provider in &provider_kinds {
|
||||
let Ok(models) = providers.query_model_capabilities(provider).await else {
|
||||
let Ok(models) = catalogs.query_model_capabilities(provider).await else {
|
||||
continue;
|
||||
};
|
||||
known_provider_catalogs.insert(provider.clone());
|
||||
|
||||
@@ -6,17 +6,17 @@ use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use futures::lock::Mutex;
|
||||
use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
|
||||
use crate::engine::provider::ProviderCatalogGeneration;
|
||||
use crate::health::{HealthProbe, HealthState};
|
||||
use crate::identity::ProviderKind;
|
||||
use crate::lifecycle::CancellationToken;
|
||||
use crate::routing::ConfigRevision;
|
||||
use crate::routing::snapshot::{
|
||||
RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotStorePort,
|
||||
RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler,
|
||||
};
|
||||
use crate::routing::{ConfigRevision, ProviderCatalogGeneration};
|
||||
use crate::task::{
|
||||
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
|
||||
WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable,
|
||||
@@ -135,11 +135,17 @@ pub struct RuntimeSnapshotPublisher {
|
||||
compiler: Arc<RuntimeSnapshotCompiler>,
|
||||
snapshots: RuntimeSnapshotHandle,
|
||||
subscriptions: Arc<dyn SnapshotSubscriptionPort>,
|
||||
refresh_lock: Arc<Mutex<()>>,
|
||||
}
|
||||
|
||||
enum RefreshMode {
|
||||
Required,
|
||||
Reconcile,
|
||||
}
|
||||
|
||||
impl RuntimeSnapshotPublisher {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
pub fn new(
|
||||
compiler: Arc<RuntimeSnapshotCompiler>,
|
||||
snapshots: RuntimeSnapshotHandle,
|
||||
subscriptions: Arc<dyn SnapshotSubscriptionPort>,
|
||||
@@ -148,21 +154,55 @@ impl RuntimeSnapshotPublisher {
|
||||
compiler,
|
||||
snapshots,
|
||||
subscriptions,
|
||||
refresh_lock: Arc::new(Mutex::new(())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 重新编译并原子替换本进程快照。
|
||||
/// 串行重编译并替换本进程快照;无法确认配置时暂停新请求。
|
||||
pub async fn refresh(&self) -> Result<ConfigRevision, RuntimeSnapshotCompileError> {
|
||||
let snapshot = self.compiler.compile().await?;
|
||||
self.refresh_with_mode(RefreshMode::Required).await
|
||||
}
|
||||
|
||||
async fn refresh_with_mode(
|
||||
&self,
|
||||
mode: RefreshMode,
|
||||
) -> Result<ConfigRevision, RuntimeSnapshotCompileError> {
|
||||
// 所有入口共享从读取事实到发布或暂停的完整临界区;只锁最后的替换会让
|
||||
// 旧编译覆盖新授权,或让晚到的失败暂停新快照。请求读取不等待此锁。
|
||||
let _refresh = self.refresh_lock.lock().await;
|
||||
let configuration_changed = if matches!(mode, RefreshMode::Reconcile) {
|
||||
let persisted_revision = self
|
||||
.compiler
|
||||
.store()
|
||||
.current_config_revision()
|
||||
.await
|
||||
.map_err(|_| {
|
||||
self.snapshots.suspend();
|
||||
RuntimeSnapshotCompileError::StoreUnavailable
|
||||
})?;
|
||||
let changed = runtime_revision_needs_refresh(
|
||||
self.published_revision().map(ConfigRevision::get),
|
||||
persisted_revision.get(),
|
||||
);
|
||||
if !changed && !self.provider_catalogs_need_refresh() {
|
||||
return Ok(persisted_revision);
|
||||
}
|
||||
changed
|
||||
} else {
|
||||
true
|
||||
};
|
||||
let snapshot = self.compiler.compile().await.inspect_err(|_| {
|
||||
// 仅目录代次变化时保留旧快照,下一周期继续对账;配置缺失、变化
|
||||
// 或持久 revision 回退均须 fail closed,不按 revision 大小丢弃刷新。
|
||||
if configuration_changed {
|
||||
self.snapshots.suspend();
|
||||
}
|
||||
})?;
|
||||
let revision = snapshot.revision();
|
||||
self.snapshots.publish(snapshot);
|
||||
Ok(revision)
|
||||
}
|
||||
|
||||
pub fn suspend(&self) {
|
||||
self.snapshots.suspend();
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn published_revision(&self) -> Option<ConfigRevision> {
|
||||
self.snapshots.revision()
|
||||
@@ -176,9 +216,7 @@ impl RuntimeSnapshotPublisher {
|
||||
|
||||
/// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚。
|
||||
async fn publish_committed_inner(&self, committed_revision: ConfigRevision) {
|
||||
if self.refresh().await.is_err() {
|
||||
self.suspend();
|
||||
}
|
||||
let _ = self.refresh().await;
|
||||
let _ = self
|
||||
.subscriptions
|
||||
.publish_snapshot_revision(committed_revision)
|
||||
@@ -204,7 +242,6 @@ impl RuntimeSnapshotPublisher {
|
||||
schedule,
|
||||
lease: None,
|
||||
task: Box::new(RuntimeSnapshotReconciliationTask {
|
||||
store: self.compiler.store(),
|
||||
publisher: self.clone(),
|
||||
}),
|
||||
},
|
||||
@@ -238,7 +275,6 @@ impl SnapshotControl for RuntimeSnapshotPublisher {
|
||||
}
|
||||
|
||||
struct RuntimeSnapshotReconciliationTask {
|
||||
store: Arc<dyn SnapshotStorePort>,
|
||||
publisher: RuntimeSnapshotPublisher,
|
||||
}
|
||||
|
||||
@@ -248,33 +284,11 @@ impl ScheduledTask for RuntimeSnapshotReconciliationTask {
|
||||
_context: WorkerCycleContext,
|
||||
) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let persisted_revision = match self.store.current_config_revision().await {
|
||||
Ok(revision) => revision,
|
||||
Err(_) => {
|
||||
self.publisher.suspend();
|
||||
return Err(WorkerTaskError::safe(
|
||||
"runtime snapshot revision is unavailable",
|
||||
));
|
||||
}
|
||||
};
|
||||
let configuration_changed = runtime_revision_needs_refresh(
|
||||
self.publisher.published_revision().map(ConfigRevision::get),
|
||||
persisted_revision.get(),
|
||||
);
|
||||
if !configuration_changed && !self.publisher.provider_catalogs_need_refresh() {
|
||||
return Ok(());
|
||||
}
|
||||
if self.publisher.refresh().await.is_err() {
|
||||
// 已提交的配置无法确认时必须 fail closed;仅目录重编译失败时继续
|
||||
// 服务旧的不可变快照,单调代次会让下一周期再次尝试。
|
||||
if configuration_changed {
|
||||
self.publisher.suspend();
|
||||
}
|
||||
return Err(WorkerTaskError::safe(
|
||||
"runtime snapshot reconciliation failed",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
self.publisher
|
||||
.refresh_with_mode(RefreshMode::Reconcile)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(|_| WorkerTaskError::safe("runtime snapshot reconciliation failed"))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -314,9 +328,7 @@ impl DaemonTask for RuntimeSnapshotSubscriptionTask {
|
||||
};
|
||||
match notified {
|
||||
Some(Ok(_)) => {
|
||||
if self.publisher.refresh().await.is_err() {
|
||||
self.publisher.suspend();
|
||||
}
|
||||
let _ = self.publisher.refresh().await;
|
||||
}
|
||||
Some(Err(_)) | None => break,
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use gateway_core::diagnostics::{diagnostic_headers, diagnostic_json};
|
||||
use serde_json::json;
|
||||
use gateway_core::diagnostics::{body_fingerprint, diagnostic_headers, diagnostic_json};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
#[test]
|
||||
fn credentials_and_content_never_enter_diagnostics_but_trace_ids_survive() {
|
||||
@@ -42,8 +42,175 @@ fn arbitrary_event_type_text_is_fingerprinted_instead_of_logged() {
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
assert_eq!(
|
||||
snapshot["events"][0]["data"]["eventType"],
|
||||
serde_json::Value::Null
|
||||
body_fingerprint(b"secret@example.com")
|
||||
);
|
||||
assert!(!snapshot.to_string().contains("secret@example.com"));
|
||||
assert!(!snapshot.to_string().contains("private-content"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_map_keys_and_header_named_values_are_not_diagnostic_facts() {
|
||||
let content = diagnostic_json(&json!({
|
||||
"type": "response.create",
|
||||
"metadata": {
|
||||
"PRIVATE_OWNER": "PRIVATE_VALUE",
|
||||
"x-request-id": "PRIVATE_NESTED_REQUEST",
|
||||
"headers": {"x-oai-request-id": "PRIVATE_NESTED_HEADER"},
|
||||
"type": "PRIVATE_TYPE",
|
||||
"model": "PRIVATE_MODEL",
|
||||
"account": 1234567890123_u64,
|
||||
},
|
||||
"input": [{
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "PRIVATE_PROMPT"}],
|
||||
"PRIVATE_ITEM_KEY": [{"cf-ray": "PRIVATE_ARRAY_HEADER"}],
|
||||
}],
|
||||
}));
|
||||
assert!(
|
||||
!content.to_string().contains("PRIVATE_"),
|
||||
"user keys and values must be summarized: {content}"
|
||||
);
|
||||
assert!(!content.to_string().contains("1234567890123"));
|
||||
assert_eq!(content["type"], "response.create");
|
||||
assert_eq!(content["input"]["length"], 1);
|
||||
assert_eq!(content["input"]["sample"][0]["role"], "user");
|
||||
assert_eq!(
|
||||
content["input"]["sample"][0]["content"]["sample"][0]["type"],
|
||||
"input_text"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_headers_never_gain_trust_from_their_shape_or_event_type() {
|
||||
for value in [
|
||||
json!({"headers": {"x-request-id": "PRIVATE_ROOT_HEADER"}}),
|
||||
json!({
|
||||
"type": "codex.response.metadata",
|
||||
"headers": {"x-request-id": "PRIVATE_ROOT_HEADER"},
|
||||
}),
|
||||
json!([{"x-request-id": "PRIVATE_ARRAY_HEADER"}]),
|
||||
] {
|
||||
let content = diagnostic_json(&value);
|
||||
assert!(!content.to_string().contains("PRIVATE_"), "{content}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protocol_shaped_strings_are_not_a_safety_guarantee() {
|
||||
let content = diagnostic_json(&json!({
|
||||
"type": "PRIVATE_TYPE",
|
||||
"code": "PRIVATE_CODE",
|
||||
"model": "PRIVATE_MODEL",
|
||||
"role": "PRIVATE_ROLE",
|
||||
"status": "PRIVATE_STATUS",
|
||||
"object": "PRIVATE_OBJECT",
|
||||
"service_tier": "PRIVATE_TIER",
|
||||
}));
|
||||
assert!(!content.to_string().contains("PRIVATE_"), "{content}");
|
||||
assert_eq!(content["type"], body_fingerprint(b"PRIVATE_TYPE"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn opaque_maps_do_not_inherit_even_known_protocol_keys_or_values() {
|
||||
let content = diagnostic_json(&json!({
|
||||
"status": 429,
|
||||
"metadata": {
|
||||
"type": "response.completed",
|
||||
"role": "user",
|
||||
"status": 429,
|
||||
"Authorization_PRIVATE_OWNER": "PRIVATE_CREDENTIAL",
|
||||
"私人标记@example.test": "PRIVATE_VALUE",
|
||||
},
|
||||
}));
|
||||
assert_eq!(content["status"], 429);
|
||||
let metadata = content["metadata"].as_object().unwrap();
|
||||
for value in ["response.completed", "user", "429"] {
|
||||
assert!(
|
||||
metadata
|
||||
.values()
|
||||
.any(|v| *v == body_fingerprint(value.as_bytes()))
|
||||
);
|
||||
}
|
||||
let text = content.to_string();
|
||||
for original in ["PRIVATE_", "私人标记", "response.completed"] {
|
||||
assert!(!text.contains(original), "{text}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnostic_capture_keeps_bounded_samples_of_unknown_keys_without_collisions() {
|
||||
let prefix = "PRIVATE_".repeat(16);
|
||||
let value = json!({
|
||||
format!("{prefix}one"): "PRIVATE_ONE",
|
||||
format!("{prefix}two"): "PRIVATE_TWO",
|
||||
});
|
||||
let content = diagnostic_json(&value);
|
||||
let entries = content.as_object().unwrap();
|
||||
assert_eq!(entries.len(), 2);
|
||||
assert!(
|
||||
entries
|
||||
.values()
|
||||
.any(|v| *v == body_fingerprint(b"PRIVATE_ONE"))
|
||||
);
|
||||
assert!(
|
||||
entries
|
||||
.values()
|
||||
.any(|v| *v == body_fingerprint(b"PRIVATE_TWO"))
|
||||
);
|
||||
assert!(!content.to_string().contains("PRIVATE_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnostic_capture_bounds_fields_arrays_depth_and_total_nodes() {
|
||||
let fields: serde_json::Map<_, _> = (0..100)
|
||||
.map(|i| (format!("PRIVATE_FIELD_{i}"), json!("PRIVATE_VALUE")))
|
||||
.collect();
|
||||
let content = diagnostic_json(&Value::Object(fields));
|
||||
assert_eq!(content.as_object().unwrap().len(), 49);
|
||||
assert_eq!(content["_omittedFields"], 52);
|
||||
|
||||
let array = diagnostic_json(&json!(vec!["PRIVATE_VALUE"; 8]));
|
||||
assert_eq!(array["length"], 8);
|
||||
assert_eq!(array["sample"].as_array().unwrap().len(), 4);
|
||||
|
||||
let mut deep = json!("PRIVATE_DEEP");
|
||||
for _ in 0..12 {
|
||||
deep = json!({"PRIVATE_DEPTH": deep});
|
||||
}
|
||||
let content = diagnostic_json(&deep);
|
||||
assert!(content.to_string().contains("\"omitted\":true"));
|
||||
assert!(!content.to_string().contains("PRIVATE_"));
|
||||
|
||||
let wide: serde_json::Map<_, _> = (0..48)
|
||||
.map(|i| (format!("PRIVATE_NODE_{i}"), json!(vec!["PRIVATE_VALUE"; 4])))
|
||||
.collect();
|
||||
let content = diagnostic_json(&Value::Object(wide));
|
||||
assert!(content["_omittedFields"].as_u64().unwrap() > 0);
|
||||
assert!(content.to_string().len() < 32 * 1024);
|
||||
assert!(!content.to_string().contains("PRIVATE_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn credentials_remain_redacted_at_nested_json_and_real_header_boundaries() {
|
||||
for name in [
|
||||
"Authorization",
|
||||
"Cookie",
|
||||
"Set-Cookie",
|
||||
"api_key",
|
||||
"api-key",
|
||||
"access_token",
|
||||
"refresh_token",
|
||||
"password",
|
||||
"client_secret",
|
||||
"attestation",
|
||||
"x-oai-is",
|
||||
"x-oai-is-update",
|
||||
] {
|
||||
let content = diagnostic_json(&json!({"metadata": {name: "PRIVATE_CREDENTIAL"}}));
|
||||
assert!(!content.to_string().contains("PRIVATE_"));
|
||||
assert!(content.to_string().contains("<redacted>"));
|
||||
let headers = diagnostic_headers([(name, "PRIVATE_CREDENTIAL")]);
|
||||
assert!(!headers.to_string().contains("PRIVATE_"));
|
||||
assert_eq!(headers[name.to_ascii_lowercase()][0], "<redacted>");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use gateway_core::diagnostics::{StreamCapture, StreamFormat, TraceContext};
|
||||
use gateway_core::diagnostics::{StreamCapture, StreamFormat, TraceContext, body_fingerprint};
|
||||
|
||||
#[test]
|
||||
fn sse_capture_reassembles_every_chunk_boundary_and_preserves_unknown_events() {
|
||||
@@ -12,11 +12,13 @@ fn sse_capture_reassembles_every_chunk_boundary_and_preserves_unknown_events() {
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
let events = snapshot["events"].as_array().unwrap();
|
||||
assert_eq!(events.len(), 3, "boundary {boundary}");
|
||||
assert_eq!(events[0]["data"]["eventType"], "future.metadata");
|
||||
// 未知 SSE 名称和伪装成头部的 JSON 仍有摘要,但不能获得协议字段的明文权限。
|
||||
assert_eq!(
|
||||
events[0]["data"]["metadata"]["headers"]["x-request-id"],
|
||||
"up-1"
|
||||
events[0]["data"]["eventType"],
|
||||
body_fingerprint(b"future.metadata")
|
||||
);
|
||||
assert!(!snapshot.to_string().contains("up-1"));
|
||||
assert!(!snapshot.to_string().contains("中文"));
|
||||
assert_eq!(events[1]["data"]["eventType"], "response.completed");
|
||||
assert_eq!(events[2]["data"]["partialFrame"], false);
|
||||
}
|
||||
|
||||
@@ -1,5 +1,15 @@
|
||||
use gateway_core::diagnostics::TraceContext;
|
||||
use std::{
|
||||
fmt::Write as _,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use gateway_core::diagnostics::{TraceContext, body_fingerprint};
|
||||
use serde_json::json;
|
||||
use tracing::{
|
||||
Event, Metadata, Subscriber,
|
||||
field::{Field, Visit},
|
||||
span::{Attributes, Id, Record},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn bounded_history_preserves_start_failure_and_final_result() {
|
||||
@@ -140,3 +150,232 @@ fn a_long_successful_retry_keeps_the_earlier_failure_ahead_of_routine_events() {
|
||||
assert_eq!(events.last().unwrap()["stage"], "request.finished");
|
||||
assert!(events.len() <= 128);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_top_level_headers_in_controlled_upstream_frames_keep_correlation_values() {
|
||||
for event_type in ["response.metadata", "codex.response.metadata", "error"] {
|
||||
let trace = TraceContext::new("req_headers");
|
||||
let body = serde_json::to_vec(&json!({
|
||||
"type": event_type,
|
||||
"status": 429,
|
||||
"headers": {
|
||||
"x-request-id": "upstream-controlled",
|
||||
"x-oai-request-id": ["upstream-one", "upstream-two"],
|
||||
"authorization": "PRIVATE_CREDENTIAL",
|
||||
"PRIVATE_HEADER_KEY": "PRIVATE_HEADER_VALUE",
|
||||
"cf-ray": [{"x-request-id": "PRIVATE_MALFORMED_HEADER"}],
|
||||
"request-id": [["PRIVATE_NESTED_ARRAY"]],
|
||||
},
|
||||
"metadata": {
|
||||
"x-request-id": "PRIVATE_METADATA_VALUE",
|
||||
"headers": {"x-request-id": "PRIVATE_NESTED_VALUE"},
|
||||
},
|
||||
"response": {
|
||||
"type": "codex.response.metadata",
|
||||
"headers": {"x-request-id": "PRIVATE_RESPONSE_HEADER"},
|
||||
},
|
||||
}))
|
||||
.unwrap();
|
||||
trace.capture("upstream.event", &body);
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
let metadata = &snapshot["events"][0]["data"]["metadata"];
|
||||
assert_eq!(metadata["headers"]["x-request-id"], "upstream-controlled");
|
||||
assert_eq!(
|
||||
metadata["headers"]["x-oai-request-id"]["sample"],
|
||||
json!(["upstream-one", "upstream-two"])
|
||||
);
|
||||
assert_eq!(metadata["status"], 429);
|
||||
assert!(!snapshot.to_string().contains("PRIVATE_"), "{snapshot}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_capture_never_infers_controlled_headers_from_client_supplied_type() {
|
||||
for stage in [
|
||||
"client.request.body",
|
||||
"upstream.request.body",
|
||||
"upstream.response.body",
|
||||
"upstream.error.body",
|
||||
] {
|
||||
let trace = TraceContext::new("req_untrusted");
|
||||
trace.capture(
|
||||
stage,
|
||||
br#"{"type":"codex.response.metadata","headers":{"x-request-id":"PRIVATE_HEADER"}}"#,
|
||||
);
|
||||
assert!(
|
||||
!trace.snapshot().unwrap().to_string().contains("PRIVATE_"),
|
||||
"{stage}"
|
||||
);
|
||||
}
|
||||
let trace = TraceContext::new("req_malformed");
|
||||
trace.capture(
|
||||
"upstream.event",
|
||||
br#"{"type":"codex.response.metadata","headers":[{"x-request-id":"PRIVATE_HEADER"}]}"#,
|
||||
);
|
||||
assert!(!trace.snapshot().unwrap().to_string().contains("PRIVATE_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_event_names_are_fingerprinted_at_all_trace_entry_points() {
|
||||
let trace = TraceContext::new("req_names");
|
||||
trace.capture("upstream.event", br#"{"type":"PRIVATE_EVENT.delta"}"#);
|
||||
trace.capture_named_event("upstream.event", b"{}", Some("PRIVATE_SSE_EVENT"));
|
||||
trace.wire_event("openai", Some("PRIVATE_WIRE_EVENT"));
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
for (event, name) in snapshot["events"].as_array().unwrap().iter().zip([
|
||||
"PRIVATE_EVENT.delta",
|
||||
"PRIVATE_SSE_EVENT",
|
||||
"PRIVATE_WIRE_EVENT",
|
||||
]) {
|
||||
assert_eq!(
|
||||
event["data"]["eventType"],
|
||||
body_fingerprint(name.as_bytes())
|
||||
);
|
||||
}
|
||||
assert!(!snapshot.to_string().contains("PRIVATE_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_deltas_coalesce_by_summary_without_merging_distinct_event_types() {
|
||||
let trace = TraceContext::new("req_unknown_delta");
|
||||
for name in ["PRIVATE_FIRST.delta", "PRIVATE_SECOND.delta"] {
|
||||
let bytes = serde_json::to_vec(&json!({"type": name, "delta": "PRIVATE_TEXT"})).unwrap();
|
||||
for _ in 0..100 {
|
||||
trace.capture("upstream.event", &bytes);
|
||||
trace.wire_event("openai", Some(name));
|
||||
}
|
||||
}
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
let events = snapshot["events"].as_array().unwrap();
|
||||
assert_eq!(events.len(), 4);
|
||||
assert!(events.iter().all(|event| event["count"] == 100));
|
||||
assert_eq!(snapshot["totalEvents"], 400);
|
||||
assert_eq!(snapshot["droppedEvents"], 0);
|
||||
assert_ne!(
|
||||
events[0]["data"]["eventType"],
|
||||
events[2]["data"]["eventType"]
|
||||
);
|
||||
assert!(!snapshot.to_string().contains("PRIVATE_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn size_gaps_and_truncated_events_do_not_reintroduce_user_content() {
|
||||
let trace = TraceContext::new("req_limits");
|
||||
let private = "PRIVATE_OVERSIZE".repeat(10_000);
|
||||
trace.capture("client.request.body", private.as_bytes());
|
||||
let fields: serde_json::Map<_, _> = (0..48)
|
||||
.map(|i| (format!("PRIVATE_FIELD_{i}"), json!("PRIVATE_VALUE")))
|
||||
.collect();
|
||||
let bytes = serde_json::to_vec(&fields).unwrap();
|
||||
trace.capture_named_event("upstream.event", &bytes, Some("PRIVATE_EVENT"));
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
assert_eq!(snapshot["events"][0]["stage"], "capture.gap");
|
||||
assert_eq!(
|
||||
snapshot["events"][0]["data"]["body"],
|
||||
body_fingerprint(private.as_bytes())
|
||||
);
|
||||
assert_eq!(snapshot["events"][1]["data"]["truncated"], true);
|
||||
assert_eq!(
|
||||
snapshot["events"][1]["data"]["eventType"],
|
||||
body_fingerprint(b"PRIVATE_EVENT")
|
||||
);
|
||||
assert!(!snapshot.to_string().contains("PRIVATE_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_trace_logs_and_serialized_snapshots_share_the_sanitized_capture_boundary() {
|
||||
// tracing 的 callsite 缓存是进程级的;隔离并行测试的注册,不修改生产日志或全局 subscriber。
|
||||
const CHILD: &str = "GATEWAY_CORE_DIAGNOSTICS_LOG_TEST";
|
||||
if std::env::var_os(CHILD).is_none() {
|
||||
let thread = std::thread::current();
|
||||
let output = std::process::Command::new(std::env::current_exe().unwrap())
|
||||
.args(["--exact", thread.name().unwrap()])
|
||||
.env(CHILD, "1")
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success() && String::from_utf8_lossy(&output.stdout).contains("1 passed"),
|
||||
"isolated log regression failed:\n{}\n{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
return;
|
||||
}
|
||||
// 仅捕获本线程的普通 trace 事件,不打开 request_dump,也不读写真实日志文件。
|
||||
let logs = Arc::new(Mutex::new(Vec::new()));
|
||||
let snapshot = tracing::subscriber::with_default(TraceLog(Arc::clone(&logs)), || {
|
||||
let trace = TraceContext::new("req_feedback");
|
||||
let exchange = trace.attempt(2).exchange("websocket");
|
||||
exchange.capture(
|
||||
"client.request.body",
|
||||
br#"{"type":"PRIVATE_EVENT","metadata":{"PRIVATE_KEY":"PRIVATE_VALUE","x-request-id":"PRIVATE_HEADER"},"input":[{"role":"user","content":[{"type":"input_text","text":"PRIVATE_PROMPT"}]}]}"#,
|
||||
);
|
||||
exchange.capture_named_event("upstream.event", b"{}", Some("PRIVATE_SSE_EVENT"));
|
||||
exchange.wire_event("openai", Some("PRIVATE_WIRE_EVENT"));
|
||||
exchange.headers(
|
||||
"upstream.response.headers",
|
||||
json!({"status": 429}),
|
||||
[
|
||||
("x-request-id", b"upstream-controlled".as_slice()),
|
||||
("authorization", b"PRIVATE_CREDENTIAL".as_slice()),
|
||||
("set-cookie", b"PRIVATE_COOKIE".as_slice()),
|
||||
],
|
||||
);
|
||||
let snapshot = trace.snapshot().unwrap();
|
||||
assert_eq!(snapshot["wireDumpEnabled"], false);
|
||||
assert_eq!(snapshot["wireFrames"], 0);
|
||||
snapshot
|
||||
});
|
||||
// API / 反馈导出使用的是这个 Value 的序列化结果,而不只是 Debug 表示。
|
||||
let exported = String::from_utf8(serde_json::to_vec(&snapshot).unwrap()).unwrap();
|
||||
assert!(!exported.contains("PRIVATE_"), "{exported}");
|
||||
let logs = logs.lock().unwrap();
|
||||
let output = logs.join("\n");
|
||||
assert!(!logs.is_empty(), "必须真实捕获 tracing 事件,不能空集通过");
|
||||
assert!(!output.contains("PRIVATE_"), "{output}");
|
||||
for fact in [
|
||||
"req_feedback",
|
||||
"upstream-controlled",
|
||||
"upstream.response.headers",
|
||||
] {
|
||||
assert!(exported.contains(fact), "{exported}");
|
||||
assert!(output.contains(fact), "{output}");
|
||||
}
|
||||
assert!(output.contains("attempt_index=2"));
|
||||
assert!(output.contains("exchange_id=1"));
|
||||
assert!(output.contains("request.trace.closed"));
|
||||
}
|
||||
|
||||
struct TraceLog(Arc<Mutex<Vec<String>>>);
|
||||
|
||||
impl Subscriber for TraceLog {
|
||||
fn enabled(&self, metadata: &Metadata<'_>) -> bool {
|
||||
metadata.target() == "request_trace"
|
||||
}
|
||||
|
||||
fn new_span(&self, _span: &Attributes<'_>) -> Id {
|
||||
Id::from_u64(1)
|
||||
}
|
||||
|
||||
fn record(&self, _span: &Id, _values: &Record<'_>) {}
|
||||
|
||||
fn record_follows_from(&self, _span: &Id, _follows: &Id) {}
|
||||
|
||||
fn event(&self, event: &Event<'_>) {
|
||||
let mut text = String::new();
|
||||
event.record(&mut LogFields(&mut text));
|
||||
self.0.lock().unwrap().push(text);
|
||||
}
|
||||
|
||||
fn enter(&self, _span: &Id) {}
|
||||
|
||||
fn exit(&self, _span: &Id) {}
|
||||
}
|
||||
|
||||
struct LogFields<'a>(&'a mut String);
|
||||
|
||||
impl Visit for LogFields<'_> {
|
||||
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
|
||||
write!(self.0, "{}={value:?} ", field.name()).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,7 +2,7 @@ use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
},
|
||||
time::{Duration, Instant, SystemTime},
|
||||
};
|
||||
@@ -14,6 +14,8 @@ use gateway_core::error::GatewayError;
|
||||
struct Admissions {
|
||||
active: Arc<AtomicBool>,
|
||||
limits: Mutex<Vec<RateLimits>>,
|
||||
release_gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
releases: AtomicUsize,
|
||||
}
|
||||
|
||||
impl ClientAdmissionPort for Admissions {
|
||||
@@ -32,7 +34,14 @@ impl ClientAdmissionPort for Admissions {
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
|
||||
Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) })
|
||||
Box::pin(async {
|
||||
self.releases.fetch_add(1, Ordering::SeqCst);
|
||||
let gate = self.release_gate.lock().unwrap().take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("release gate");
|
||||
}
|
||||
Ok(self.active.swap(false, Ordering::SeqCst))
|
||||
})
|
||||
}
|
||||
fn restore(
|
||||
&self,
|
||||
@@ -42,10 +51,14 @@ impl ClientAdmissionPort for Admissions {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Budget {
|
||||
reject: bool,
|
||||
active: Arc<AtomicBool>,
|
||||
charges: Mutex<Vec<ClientBudgetCharge>>,
|
||||
settlement_gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
settlements: AtomicUsize,
|
||||
fail_settlement: bool,
|
||||
}
|
||||
|
||||
impl ClientBudgetPort for Budget {
|
||||
@@ -64,12 +77,22 @@ impl ClientBudgetPort for Budget {
|
||||
}
|
||||
fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
|
||||
Box::pin(async move {
|
||||
self.settlements.fetch_add(1, Ordering::SeqCst);
|
||||
assert!(
|
||||
self.active.load(Ordering::SeqCst),
|
||||
"settle before releasing the concurrent request slot"
|
||||
);
|
||||
let gate = self.settlement_gate.lock().unwrap().take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("settlement gate");
|
||||
}
|
||||
assert!(self.active.load(Ordering::SeqCst));
|
||||
self.charges.lock().unwrap().push(charge);
|
||||
Ok(())
|
||||
if self.fail_settlement {
|
||||
Err(ClientBudgetError)
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -111,6 +134,7 @@ fn budget_rejection_releases_client_concurrency_without_creating_a_charge() {
|
||||
reject: true,
|
||||
active: admissions.active.clone(),
|
||||
charges: Mutex::default(),
|
||||
..Default::default()
|
||||
});
|
||||
let service = service(admissions.clone(), budget.clone());
|
||||
for transport in [
|
||||
@@ -200,6 +224,7 @@ fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all
|
||||
reject: false,
|
||||
active: admissions.active.clone(),
|
||||
charges: Mutex::default(),
|
||||
..Default::default()
|
||||
});
|
||||
let service = service(admissions.clone(), budget.clone());
|
||||
for transport in [
|
||||
@@ -232,9 +257,554 @@ fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all
|
||||
);
|
||||
}
|
||||
|
||||
fn early_failure_service(
|
||||
store: Arc<TrackingExecutionStore>,
|
||||
admissions: Arc<Admissions>,
|
||||
budget: Arc<Budget>,
|
||||
) -> DefaultExecutionService {
|
||||
DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
store,
|
||||
ProviderRegistry::new([Arc::new(LocalFailingProvider) as Arc<dyn Provider>]).unwrap(),
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
.with_budget(budget)
|
||||
}
|
||||
|
||||
fn assert_early_failure_recorded(
|
||||
store: &TrackingExecutionStore,
|
||||
request_id: &ModelRequestId,
|
||||
transport: ClientTransport,
|
||||
) {
|
||||
assert_eq!(store.creates.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(store.finalizes.load(Ordering::SeqCst), 1);
|
||||
assert!(store.attempts.lock().unwrap().is_empty());
|
||||
let requests = store.requests.lock().unwrap();
|
||||
assert_eq!(requests.len(), 1);
|
||||
let request = &requests[0];
|
||||
assert_eq!(request.id, *request_id);
|
||||
assert_eq!(request.client_api_key_ref.as_str(), "key_start_test");
|
||||
assert_eq!(
|
||||
request.client_api_key_id.as_ref(),
|
||||
Some(&request.client_api_key_ref)
|
||||
);
|
||||
assert_eq!(request.protocol, "openai");
|
||||
assert_eq!(request.endpoint, "/v1/responses");
|
||||
assert_eq!(request.client_transport, transport.as_str());
|
||||
assert_eq!(request.operation, OperationKind::Generate);
|
||||
assert_eq!(
|
||||
request.requested_model.as_ref().map(PublicModelId::as_str),
|
||||
Some("gpt-start")
|
||||
);
|
||||
let finalizations = store.finalizations.lock().unwrap();
|
||||
assert_eq!(finalizations.len(), 1);
|
||||
let finalization = &finalizations[0];
|
||||
assert_eq!(finalization.request_id, *request_id);
|
||||
assert_eq!(finalization.outcome, ExecutionOutcome::Failed);
|
||||
assert_eq!(finalization.send_state, UpstreamSendState::NotSent);
|
||||
assert_eq!(finalization.attempt_count, 0);
|
||||
assert!(finalization.downstream_committed_at.is_none());
|
||||
assert!(finalization.upstream_request_id.is_none());
|
||||
assert!(finalization.upstream_response_id.is_none());
|
||||
assert!(finalization.upstream_status_code.is_none());
|
||||
assert!(finalization.upstream_transport.is_none());
|
||||
assert!(finalization.http_version.is_none());
|
||||
assert!(finalization.websocket_pool.is_none());
|
||||
assert!(finalization.provider_metadata_json.is_none());
|
||||
assert_eq!(
|
||||
finalization
|
||||
.error
|
||||
.as_ref()
|
||||
.expect("original failure")
|
||||
.kind(),
|
||||
GatewayErrorKind::Unsupported,
|
||||
"detached cancellation must not replace the original provider failure"
|
||||
);
|
||||
assert_eq!(finalization.usage, Default::default());
|
||||
assert!(finalization.cost.total().is_none());
|
||||
}
|
||||
|
||||
fn assert_zero_cleanup_completed(
|
||||
admissions: &Admissions,
|
||||
budget: &Budget,
|
||||
store: &TrackingExecutionStore,
|
||||
request_id: &ModelRequestId,
|
||||
) {
|
||||
assert!(!admissions.active.load(Ordering::SeqCst));
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
|
||||
let charges = budget.charges.lock().unwrap();
|
||||
assert_eq!(charges.len(), 1);
|
||||
assert_eq!(charges[0].request_id, *request_id);
|
||||
assert_eq!(charges[0].key_id.as_str(), "key_start_test");
|
||||
assert_eq!(charges[0].amount_usd, Decimal::ZERO);
|
||||
assert_eq!(
|
||||
charges[0].completed_at,
|
||||
store.finalizations.lock().unwrap()[0].completed_at,
|
||||
"settlement must retain the original failure completion time"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn early_provider_failure_records_request_and_settles_zero_for_all_transports() {
|
||||
block_on(async {
|
||||
for transport in [
|
||||
ClientTransport::HttpJson,
|
||||
ClientTransport::HttpSse,
|
||||
ClientTransport::WebSocket,
|
||||
] {
|
||||
let store = Arc::new(TrackingExecutionStore::default());
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
let service = early_failure_service(store.clone(), admissions.clone(), budget.clone());
|
||||
let mut started = service.start(request(&service, transport)).await.unwrap();
|
||||
assert!(admissions.active.load(Ordering::SeqCst));
|
||||
assert_eq!(store.creates.load(Ordering::SeqCst), 0);
|
||||
let error = if transport == ClientTransport::HttpJson {
|
||||
started.session.collect_uncommitted().await.unwrap_err()
|
||||
} else {
|
||||
started.session.next_event().await.unwrap_err()
|
||||
};
|
||||
assert!(matches!(
|
||||
error,
|
||||
EngineError::Provider(error)
|
||||
if error.kind() == ProviderErrorKind::Unsupported
|
||||
&& error.send_state() == UpstreamSendState::NotSent
|
||||
));
|
||||
assert!(started.session.is_finalized());
|
||||
assert_early_failure_recorded(&store, &started.request_id, transport);
|
||||
assert_zero_cleanup_completed(&admissions, &budget, &store, &started.request_id);
|
||||
assert!(started.session.next_event().await.unwrap().is_none());
|
||||
started.session.detach_finalize().await;
|
||||
assert_early_failure_recorded(&store, &started.request_id, transport);
|
||||
assert_zero_cleanup_completed(&admissions, &budget, &store, &started.request_id);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detached_early_failure_resumes_cancelled_store_write_and_settles_once_for_all_transports() {
|
||||
block_on(async {
|
||||
for transport in [
|
||||
ClientTransport::HttpJson,
|
||||
ClientTransport::HttpSse,
|
||||
ClientTransport::WebSocket,
|
||||
] {
|
||||
for suspend_create in [true, false] {
|
||||
let (complete_write, write_gate) = oneshot::channel();
|
||||
let store = Arc::new(TrackingExecutionStore::default());
|
||||
if suspend_create {
|
||||
*store.create_gate.lock().unwrap() = Some(write_gate);
|
||||
} else {
|
||||
*store.finalize_gate.lock().unwrap() = Some(write_gate);
|
||||
}
|
||||
let (complete_release, release_gate) = oneshot::channel();
|
||||
let admissions = Arc::new(Admissions {
|
||||
release_gate: Mutex::new(Some(release_gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let (complete_settlement, settlement_gate) = oneshot::channel();
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
settlement_gate: Mutex::new(Some(settlement_gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let service =
|
||||
early_failure_service(store.clone(), admissions.clone(), budget.clone());
|
||||
let mut started = service.start(request(&service, transport)).await.unwrap();
|
||||
let mut next = started.session.next_event();
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
drop(next);
|
||||
assert!(!started.session.is_finalized());
|
||||
assert_eq!(store.creates.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(
|
||||
store.finalizes.load(Ordering::SeqCst),
|
||||
usize::from(!suspend_create)
|
||||
);
|
||||
assert_eq!(
|
||||
store.requests.lock().unwrap().len(),
|
||||
usize::from(!suspend_create)
|
||||
);
|
||||
assert!(store.finalizations.lock().unwrap().is_empty());
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
|
||||
assert!(admissions.active.load(Ordering::SeqCst));
|
||||
|
||||
let mut detached = started.session.detach_finalize();
|
||||
assert!(futures::poll!(detached.as_mut()).is_pending());
|
||||
assert_eq!(store.creates.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(
|
||||
store.finalizes.load(Ordering::SeqCst),
|
||||
usize::from(!suspend_create)
|
||||
);
|
||||
// receiver 已从 Store 替身取走;只有延续原 future 才能继续接收此信号。
|
||||
complete_write
|
||||
.send(())
|
||||
.expect("detached cleanup retains the original store write");
|
||||
assert!(futures::poll!(detached.as_mut()).is_pending());
|
||||
assert_early_failure_recorded(&store, &started.request_id, transport);
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
|
||||
assert!(budget.charges.lock().unwrap().is_empty());
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
|
||||
assert!(admissions.active.load(Ordering::SeqCst));
|
||||
|
||||
complete_settlement.send(()).expect("original settlement");
|
||||
assert!(futures::poll!(detached.as_mut()).is_pending());
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
|
||||
{
|
||||
let charges = budget.charges.lock().unwrap();
|
||||
assert_eq!(charges.len(), 1);
|
||||
assert_eq!(charges[0].amount_usd, Decimal::ZERO);
|
||||
}
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
|
||||
assert!(admissions.active.load(Ordering::SeqCst));
|
||||
complete_release
|
||||
.send(())
|
||||
.expect("original admission release");
|
||||
detached.await;
|
||||
assert_early_failure_recorded(&store, &started.request_id, transport);
|
||||
assert_zero_cleanup_completed(&admissions, &budget, &store, &started.request_id);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
struct ChargedProvider {
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
fn known_charge() -> Decimal {
|
||||
"1.25".parse().unwrap()
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for ChargedProvider {
|
||||
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 candidate = request.candidate();
|
||||
let metadata = ProviderCallMetadata::new(
|
||||
candidate.provider().clone(),
|
||||
candidate.upstream_model().unwrap().clone(),
|
||||
ProviderAccountId::new("acct_start").unwrap(),
|
||||
UpstreamTransport::new("websocket").unwrap(),
|
||||
);
|
||||
let response = ResponseMeta::new("resp_cleanup", "gpt-start");
|
||||
let terminal = if self.fail {
|
||||
Err(ProviderError::new(
|
||||
ProviderErrorKind::InvalidRequest,
|
||||
UpstreamSendState::Sent,
|
||||
))
|
||||
} else {
|
||||
Ok(GatewayEvent::Completed(response.clone()).into())
|
||||
};
|
||||
let events: Vec<Result<ProviderEvent, ProviderError>> = vec![
|
||||
Ok(GatewayEvent::Started(response).into()),
|
||||
Ok(GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(known_charge().scaled()).unwrap(),
|
||||
)
|
||||
.into()),
|
||||
terminal,
|
||||
];
|
||||
Ok(ProviderStream::new(
|
||||
metadata,
|
||||
futures::stream::iter(events),
|
||||
(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn charged_service(
|
||||
admissions: Arc<Admissions>,
|
||||
budget: Option<Arc<Budget>>,
|
||||
fail: bool,
|
||||
) -> DefaultExecutionService {
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::new([Arc::new(ChargedProvider { fail }) as Arc<dyn Provider>]).unwrap(),
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
match budget {
|
||||
Some(budget) => service.with_budget(budget),
|
||||
None => service,
|
||||
}
|
||||
}
|
||||
|
||||
async fn consume_charged_prefix(session: &mut dyn ExecutionSession) {
|
||||
let first = session.next_event().await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
first.commit_requirement(),
|
||||
CommitRequirement::CommitBeforeDelivery
|
||||
);
|
||||
session.commit_downstream(None).await.unwrap();
|
||||
let cost = session.next_event().await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
cost.commit_requirement(),
|
||||
CommitRequirement::AlreadyCommitted
|
||||
);
|
||||
}
|
||||
|
||||
fn assert_cleanup_completed(admissions: &Admissions, budget: &Budget, request_id: &ModelRequestId) {
|
||||
assert!(!admissions.active.load(Ordering::SeqCst));
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
|
||||
let charges = budget.charges.lock().unwrap();
|
||||
assert_eq!(charges.len(), 1);
|
||||
assert_eq!(charges[0].request_id, *request_id);
|
||||
assert_eq!(charges[0].key_id.as_str(), "key_start_test");
|
||||
assert_eq!(charges[0].amount_usd, known_charge());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_completion_settles_known_charge_once_for_all_transports() {
|
||||
block_on(async {
|
||||
for transport in [
|
||||
ClientTransport::HttpJson,
|
||||
ClientTransport::HttpSse,
|
||||
ClientTransport::WebSocket,
|
||||
] {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
let service = charged_service(admissions.clone(), Some(budget.clone()), false);
|
||||
let mut started = service.start(request(&service, transport)).await.unwrap();
|
||||
if transport == ClientTransport::HttpJson {
|
||||
assert_eq!(
|
||||
started.session.collect_uncommitted().await.unwrap().len(),
|
||||
3
|
||||
);
|
||||
assert!(!started.session.is_finalized());
|
||||
started.session.commit_downstream(Some(200)).await.unwrap();
|
||||
} else {
|
||||
consume_charged_prefix(started.session.as_mut()).await;
|
||||
assert!(started.session.next_event().await.unwrap().is_some());
|
||||
}
|
||||
assert!(started.session.next_event().await.unwrap().is_none());
|
||||
assert!(started.session.is_finalized());
|
||||
started.session.detach_finalize().await;
|
||||
assert_cleanup_completed(&admissions, &budget, &started.request_id);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_survives_cancelled_event_wait_and_resumes_without_restart() {
|
||||
block_on(async {
|
||||
for fail in [false, true] {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let (release, gate) = oneshot::channel();
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
settlement_gate: Mutex::new(Some(gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let service = charged_service(admissions.clone(), Some(budget.clone()), fail);
|
||||
let mut started = service
|
||||
.start(request(&service, ClientTransport::WebSocket))
|
||||
.await
|
||||
.unwrap();
|
||||
consume_charged_prefix(started.session.as_mut()).await;
|
||||
if !fail {
|
||||
assert!(started.session.next_event().await.unwrap().is_some());
|
||||
}
|
||||
for _ in 0..2 {
|
||||
let mut terminal = started.session.next_event();
|
||||
assert!(futures::poll!(terminal.as_mut()).is_pending());
|
||||
drop(terminal);
|
||||
assert!(!started.session.is_finalized());
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
|
||||
assert!(budget.charges.lock().unwrap().is_empty());
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
release
|
||||
.send(())
|
||||
.expect("the original settlement is still alive");
|
||||
assert!(started.session.next_event().await.unwrap().is_none());
|
||||
assert!(started.session.is_finalized());
|
||||
started.session.detach_finalize().await;
|
||||
assert_cleanup_completed(&admissions, &budget, &started.request_id);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detached_cleanup_settles_cancelled_execution_and_resumes_existing_settlement() {
|
||||
block_on(async {
|
||||
for cancel_event_wait in [false, true] {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let (release, gate) = oneshot::channel();
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
settlement_gate: Mutex::new(Some(gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let service = charged_service(admissions.clone(), Some(budget.clone()), false);
|
||||
let mut started = service
|
||||
.start(request(&service, ClientTransport::WebSocket))
|
||||
.await
|
||||
.unwrap();
|
||||
consume_charged_prefix(started.session.as_mut()).await;
|
||||
if cancel_event_wait {
|
||||
started.session.cancel();
|
||||
let mut terminal = started.session.next_event();
|
||||
assert!(futures::poll!(terminal.as_mut()).is_pending());
|
||||
drop(terminal);
|
||||
assert!(!started.session.is_finalized());
|
||||
}
|
||||
let mut detached = started.session.detach_finalize();
|
||||
assert!(futures::poll!(detached.as_mut()).is_pending());
|
||||
assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
|
||||
release
|
||||
.send(())
|
||||
.expect("detached cleanup retains settlement");
|
||||
detached.await;
|
||||
assert_cleanup_completed(&admissions, &budget, &started.request_id);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancelled_release_resumes_without_repeating_settlement_with_or_without_budget() {
|
||||
block_on(async {
|
||||
for with_budget in [false, true] {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let admissions = Arc::new(Admissions {
|
||||
release_gate: Mutex::new(Some(gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
let service = charged_service(
|
||||
admissions.clone(),
|
||||
with_budget.then(|| budget.clone()),
|
||||
false,
|
||||
);
|
||||
let mut started = service
|
||||
.start(request(&service, ClientTransport::HttpSse))
|
||||
.await
|
||||
.unwrap();
|
||||
consume_charged_prefix(started.session.as_mut()).await;
|
||||
assert!(started.session.next_event().await.unwrap().is_some());
|
||||
let mut terminal = started.session.next_event();
|
||||
assert!(futures::poll!(terminal.as_mut()).is_pending());
|
||||
drop(terminal);
|
||||
assert!(!started.session.is_finalized());
|
||||
assert!(admissions.active.load(Ordering::SeqCst));
|
||||
assert_eq!(
|
||||
budget.settlements.load(Ordering::SeqCst),
|
||||
usize::from(with_budget)
|
||||
);
|
||||
let mut detached = started.session.detach_finalize();
|
||||
assert!(futures::poll!(detached.as_mut()).is_pending());
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
|
||||
release
|
||||
.send(())
|
||||
.expect("the original release is still alive");
|
||||
detached.await;
|
||||
if with_budget {
|
||||
assert_cleanup_completed(&admissions, &budget, &started.request_id);
|
||||
} else {
|
||||
assert!(!admissions.active.load(Ordering::SeqCst));
|
||||
assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
|
||||
assert!(budget.charges.lock().unwrap().is_empty());
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancelled_buffered_commit_keeps_settlement_for_detached_cleanup() {
|
||||
block_on(async {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let (release, gate) = oneshot::channel();
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
settlement_gate: Mutex::new(Some(gate)),
|
||||
..Default::default()
|
||||
});
|
||||
let service = charged_service(admissions.clone(), Some(budget.clone()), false);
|
||||
let mut started = service
|
||||
.start(request(&service, ClientTransport::HttpJson))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
started.session.collect_uncommitted().await.unwrap().len(),
|
||||
3
|
||||
);
|
||||
let mut commit = started.session.commit_downstream(Some(200));
|
||||
assert!(futures::poll!(commit.as_mut()).is_pending());
|
||||
drop(commit);
|
||||
assert!(!started.session.is_finalized());
|
||||
release
|
||||
.send(())
|
||||
.expect("buffered commit retains settlement");
|
||||
started.session.detach_finalize().await;
|
||||
assert_cleanup_completed(&admissions, &budget, &started.request_id);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_failure_keeps_provider_error_and_releases_concurrency_once() {
|
||||
block_on(async {
|
||||
for fail_settlement in [false, true] {
|
||||
let admissions = Arc::new(Admissions::default());
|
||||
let budget = Arc::new(Budget {
|
||||
active: admissions.active.clone(),
|
||||
fail_settlement,
|
||||
..Default::default()
|
||||
});
|
||||
let service = charged_service(admissions.clone(), Some(budget.clone()), true);
|
||||
let mut started = service
|
||||
.start(request(&service, ClientTransport::WebSocket))
|
||||
.await
|
||||
.unwrap();
|
||||
consume_charged_prefix(started.session.as_mut()).await;
|
||||
assert!(matches!(
|
||||
started.session.next_event().await,
|
||||
Err(EngineError::Provider(error)) if error.kind() == ProviderErrorKind::InvalidRequest
|
||||
));
|
||||
assert!(started.session.is_finalized());
|
||||
started.session.detach_finalize().await;
|
||||
// Store 端口已接管精确费用后,结算错误不能改写 Provider 错误或触发第二次结算。
|
||||
assert_cleanup_completed(&admissions, &budget, &started.request_id);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use futures::{executor::block_on, future::BoxFuture};
|
||||
use futures::{channel::oneshot, executor::block_on, future::BoxFuture};
|
||||
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
use gateway_core::engine::admission::{
|
||||
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
|
||||
@@ -245,31 +815,34 @@ use gateway_core::engine::continuation::{
|
||||
};
|
||||
use gateway_core::engine::execution::{
|
||||
ClientApiKeyUsageSink, ClientTransport, DefaultExecutionService, ExecutionRequestMetadata,
|
||||
ExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
|
||||
StartExecution, StartProviderExecution, provider_failure_affects_circuit,
|
||||
ExecutionService, ExecutionSession, ProviderCircuitDecision, ProviderCircuitError,
|
||||
ProviderCircuitPort, StartExecution, StartProviderExecution, provider_failure_affects_circuit,
|
||||
};
|
||||
use gateway_core::engine::probe::{AccountProbe, AccountProbeErrorSource, AccountProbeRequest};
|
||||
use gateway_core::engine::provider::{
|
||||
Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities,
|
||||
ProviderRegistry, ProviderRequest, ProviderRequestObservation, ProviderStream,
|
||||
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderRequestObservation,
|
||||
ProviderStream,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AttemptContext, AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
|
||||
ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport,
|
||||
AttemptContext, AttemptRecord, CommitRequirement, EngineError, ExecutionOutcome,
|
||||
ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
|
||||
ProbeFailure, RecoveryReport,
|
||||
};
|
||||
use gateway_core::error::{
|
||||
ClientVisibleUpstreamResponse, GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError,
|
||||
StoreErrorKind,
|
||||
};
|
||||
use gateway_core::event::{GatewayEvent, ProviderEvent, ResponseMeta};
|
||||
use gateway_core::metering::{Decimal, ProviderReportedCost};
|
||||
use gateway_core::operation::{
|
||||
GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProtocolPayload,
|
||||
RawJsonPayload,
|
||||
};
|
||||
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::{
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
|
||||
ProviderModel, PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
|
||||
ProviderCatalogGeneration, ProviderKind, ProviderModel, ProviderModelCapabilities,
|
||||
PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
use gateway_core::runtime::RuntimeSnapshotHandle;
|
||||
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
|
||||
@@ -290,6 +863,9 @@ fn only_provider_attributable_failures_should_affect_circuit() {
|
||||
assert!(!provider_failure_affects_circuit(
|
||||
ProviderErrorKind::ContinuationRecoveryRequired
|
||||
));
|
||||
assert!(!provider_failure_affects_circuit(
|
||||
ProviderErrorKind::UpstreamCapacityUnavailable
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -423,6 +999,108 @@ fn probe_observation_store_failure_preserves_the_provider_error() {
|
||||
|
||||
struct FailingProvider;
|
||||
|
||||
struct NativeCatalogProvider {
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for NativeCatalogProvider {
|
||||
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 query_client_model_catalog(
|
||||
&self,
|
||||
scope: &FrozenAccountScope,
|
||||
protocol: &str,
|
||||
version: &str,
|
||||
) -> Result<
|
||||
Option<Vec<gateway_core::routing::ProviderModelDescriptor>>,
|
||||
gateway_core::routing::ProviderCatalogUnavailable,
|
||||
> {
|
||||
assert!(scope.allows(&ProviderAccountId::new("acct_start").expect("account")));
|
||||
assert!(!scope.allows(&ProviderAccountId::new("acct_other").expect("account")));
|
||||
assert_eq!((protocol, version), ("codex", "0.154.0"));
|
||||
if self.fail {
|
||||
return Err(gateway_core::routing::ProviderCatalogUnavailable);
|
||||
}
|
||||
Ok(Some(
|
||||
["gpt-new", "gpt-start"]
|
||||
.into_iter()
|
||||
.map(|id| gateway_core::routing::ProviderModelDescriptor {
|
||||
model: UpstreamModelId::new(id).expect("id"),
|
||||
payload: RawJsonPayload::new(
|
||||
"codex",
|
||||
serde_json::to_vec(&json!({"slug":id,"future":null}))
|
||||
.expect("JSON")
|
||||
.into(),
|
||||
)
|
||||
.expect("payload"),
|
||||
})
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
_: ProviderRequest,
|
||||
_: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
panic!("catalog must not use metered execution")
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models() {
|
||||
for fail in [false, true] {
|
||||
let snapshot = start_snapshot().with_model_mappings(BTreeMap::from([
|
||||
("alias".to_owned(), "gpt-start".to_owned()),
|
||||
("missing-alias".to_owned(), "unavailable-model".to_owned()),
|
||||
]));
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::new([Arc::new(NativeCatalogProvider { fail }) as Arc<dyn Provider>])
|
||||
.expect("registry"),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
let client = service.authenticate("sk_start_test").expect("authenticate");
|
||||
let result = block_on(service.client_model_catalog(&client, "codex", "0.154.0"));
|
||||
if fail {
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"native failure cannot use the global synthesized catalog"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let models = result.expect("models");
|
||||
let pairs = models
|
||||
.iter()
|
||||
.map(|entry| {
|
||||
let gateway_core::routing::PublicModelDescriptor::Native { model, payload } = entry
|
||||
else {
|
||||
panic!("native required")
|
||||
};
|
||||
(model.as_str(), payload.body())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
pairs.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
|
||||
["gpt-start", "alias"]
|
||||
);
|
||||
assert_eq!(pairs[0].1, pairs[1].1, "Core preserves alias payload bytes");
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for FailingProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
@@ -821,9 +1499,13 @@ 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>>>,
|
||||
requests: Mutex<Vec<NewModelRequest>>,
|
||||
attempts: Mutex<Vec<AttemptRecord>>,
|
||||
finalizations: Mutex<Vec<ModelRequestFinalization>>,
|
||||
create_gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
finalize_gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
creates: AtomicUsize,
|
||||
finalizes: AtomicUsize,
|
||||
fail_probe_observation: AtomicBool,
|
||||
}
|
||||
|
||||
@@ -840,24 +1522,40 @@ impl TrackingExecutionStore {
|
||||
}
|
||||
|
||||
fn model_request_endpoints(&self) -> Vec<String> {
|
||||
self.model_request_endpoints
|
||||
self.requests
|
||||
.lock()
|
||||
.expect("model request endpoints lock")
|
||||
.clone()
|
||||
.expect("model requests lock")
|
||||
.iter()
|
||||
.map(|request| request.endpoint.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn requested_models(&self) -> Vec<Option<String>> {
|
||||
self.requested_models
|
||||
self.requests
|
||||
.lock()
|
||||
.expect("requested models lock")
|
||||
.clone()
|
||||
.expect("model requests lock")
|
||||
.iter()
|
||||
.map(|request| {
|
||||
request
|
||||
.requested_model
|
||||
.as_ref()
|
||||
.map(|model| model.as_str().to_owned())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn upstream_models(&self) -> Vec<Option<String>> {
|
||||
self.upstream_models
|
||||
self.attempts
|
||||
.lock()
|
||||
.expect("upstream models lock")
|
||||
.clone()
|
||||
.expect("attempts lock")
|
||||
.iter()
|
||||
.map(|attempt| {
|
||||
attempt
|
||||
.upstream_model_id
|
||||
.as_ref()
|
||||
.map(|model| model.as_str().to_owned())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -865,33 +1563,21 @@ impl TrackingExecutionStore {
|
||||
impl ExecutionStore for TrackingExecutionStore {
|
||||
async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> {
|
||||
self.touch();
|
||||
self.requested_models
|
||||
self.creates.fetch_add(1, Ordering::SeqCst);
|
||||
let gate = self.create_gate.lock().unwrap().take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("create model request gate");
|
||||
}
|
||||
self.requests
|
||||
.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);
|
||||
.expect("model requests lock")
|
||||
.push(request);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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()),
|
||||
);
|
||||
self.attempts.lock().expect("attempts lock").push(attempt);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -935,8 +1621,20 @@ impl ExecutionStore for TrackingExecutionStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> {
|
||||
async fn finalize_model_request(
|
||||
&self,
|
||||
finalization: ModelRequestFinalization,
|
||||
) -> Result<(), StoreError> {
|
||||
self.touch();
|
||||
self.finalizes.fetch_add(1, Ordering::SeqCst);
|
||||
let gate = self.finalize_gate.lock().unwrap().take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("finalize model request gate");
|
||||
}
|
||||
self.finalizations
|
||||
.lock()
|
||||
.expect("finalizations lock")
|
||||
.push(finalization);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -11,8 +11,8 @@ use futures_timer::Delay;
|
||||
use gateway_core::account::{AccountFeedbackStats, ProviderAccountId};
|
||||
use gateway_core::engine::AttemptContext;
|
||||
use gateway_core::engine::provider::{
|
||||
EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration,
|
||||
ProviderModelCapabilities, ProviderRegistry, ProviderRequest, ProviderStream, RegistryError,
|
||||
EventStream, Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
|
||||
RegistryError,
|
||||
};
|
||||
use gateway_core::error::{ProviderError, ProviderErrorKind};
|
||||
use gateway_core::event::{
|
||||
@@ -20,7 +20,10 @@ use gateway_core::event::{
|
||||
TextDelta,
|
||||
};
|
||||
use gateway_core::operation::OperationKind;
|
||||
use gateway_core::routing::{ModelCapabilities, ProviderKind, UpstreamModelId};
|
||||
use gateway_core::routing::{
|
||||
ModelCapabilities, ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
|
||||
ProviderKind, ProviderModelCapabilities, UpstreamModelId,
|
||||
};
|
||||
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -84,9 +87,14 @@ fn registry_should_query_provider_compiled_model_capabilities() {
|
||||
.expect("provider");
|
||||
let registry = builder.build();
|
||||
let provider = ProviderKind::new("openai").expect("provider kind");
|
||||
let models = futures::executor::block_on(registry.query_model_capabilities(&provider))
|
||||
let catalog: &dyn ProviderCatalogPort = ®istry;
|
||||
let models = futures::executor::block_on(catalog.query_model_capabilities(&provider))
|
||||
.expect("live catalog");
|
||||
|
||||
assert_eq!(
|
||||
catalog.catalog_generations(),
|
||||
std::collections::BTreeMap::from([(provider, ProviderCatalogGeneration::default())]),
|
||||
);
|
||||
assert_eq!(models[0].upstream_model().as_str(), "live-model");
|
||||
assert!(
|
||||
models[0]
|
||||
@@ -98,6 +106,75 @@ fn registry_should_query_provider_compiled_model_capabilities() {
|
||||
);
|
||||
}
|
||||
|
||||
struct CatalogProvider(Result<Vec<ProviderModelCapabilities>, ProviderErrorKind>);
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for CatalogProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"alpha"
|
||||
}
|
||||
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::new(7)
|
||||
}
|
||||
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
self.0
|
||||
.clone()
|
||||
.map_err(|kind| ProviderError::new(kind, UpstreamSendState::NotSent))
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_: ProviderRequest,
|
||||
_: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
panic!("目录端口不得触发执行")
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registry_catalog_port_should_preserve_unknown_and_known_empty_catalogs() {
|
||||
let provider = ProviderKind::new("alpha").expect("provider");
|
||||
for (response, expected) in [
|
||||
(
|
||||
Err(ProviderErrorKind::Unavailable),
|
||||
Err(ProviderCatalogUnavailable),
|
||||
),
|
||||
(Ok(Vec::new()), Ok(Vec::new())),
|
||||
] {
|
||||
let registry =
|
||||
ProviderRegistry::new([Arc::new(CatalogProvider(response)) as Arc<dyn Provider>])
|
||||
.expect("registry");
|
||||
let catalog: &dyn ProviderCatalogPort = ®istry;
|
||||
|
||||
assert_eq!(
|
||||
futures::executor::block_on(catalog.query_model_capabilities(&provider)),
|
||||
expected,
|
||||
);
|
||||
assert_eq!(
|
||||
catalog.catalog_generations().get(&provider),
|
||||
Some(&ProviderCatalogGeneration::new(7)),
|
||||
"目录不可读也不能丢失注册身份和最近成功发布的代次",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registry_catalog_port_should_report_unregistered_provider_as_unknown() {
|
||||
let registry = ProviderRegistry::default();
|
||||
let catalog: &dyn ProviderCatalogPort = ®istry;
|
||||
let provider = ProviderKind::new("missing").expect("provider");
|
||||
|
||||
assert_eq!(
|
||||
futures::executor::block_on(catalog.query_model_capabilities(&provider)),
|
||||
Err(ProviderCatalogUnavailable),
|
||||
);
|
||||
assert!(catalog.catalog_generations().is_empty());
|
||||
}
|
||||
|
||||
struct DropLease(Arc<AtomicBool>);
|
||||
|
||||
impl Drop for DropLease {
|
||||
|
||||
@@ -57,6 +57,29 @@ fn plaintext_client_key_debug_should_be_redacted() {
|
||||
assert_eq!(key.expose_for_auth(), "sk_must_not_appear");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migrated_client_keys_keep_their_exact_value_without_prefix_or_length_limits() {
|
||||
for value in [
|
||||
"q".to_owned(),
|
||||
"legacy-!@#$%^&*()[]{}:;,.?/+=~|\\".to_owned(),
|
||||
"x".repeat(8192),
|
||||
] {
|
||||
assert_eq!(plaintext(&value).expose_for_auth(), value);
|
||||
}
|
||||
for value in [
|
||||
"",
|
||||
" ",
|
||||
" key",
|
||||
"key ",
|
||||
"key\n",
|
||||
"key\tvalue",
|
||||
"key\u{7f}",
|
||||
"密钥",
|
||||
] {
|
||||
assert!(PlaintextClientApiKey::new(value).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_version_should_require_strict_semver() {
|
||||
assert!(CodexClientVersion::parse("0.152.0").is_ok());
|
||||
|
||||
@@ -2,16 +2,9 @@ use std::collections::BTreeMap;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::executor::block_on;
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use gateway_core::engine::AttemptContext;
|
||||
use gateway_core::engine::provider::{
|
||||
Provider, ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRegistry,
|
||||
ProviderRequest, ProviderStream,
|
||||
};
|
||||
use gateway_core::error::{ProviderError, ProviderErrorKind};
|
||||
use gateway_core::operation::OperationKind;
|
||||
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::snapshot::{
|
||||
@@ -20,10 +13,10 @@ use gateway_core::routing::snapshot::{
|
||||
SnapshotProviderAccountFacts, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
|
||||
};
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ModelPresentation, ProviderKind, PublicModelId,
|
||||
UpstreamModelId,
|
||||
ConfigRevision, ModelCapabilities, ModelPresentation, ProviderCatalogGeneration,
|
||||
ProviderCatalogPort, ProviderCatalogUnavailable, ProviderKind, ProviderModelCapabilities,
|
||||
PublicModelId, UpstreamModelId,
|
||||
};
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestSnapshotStore {
|
||||
@@ -51,90 +44,72 @@ impl SnapshotStorePort for TestSnapshotStore {
|
||||
}
|
||||
}
|
||||
|
||||
struct PublishingCatalogProvider {
|
||||
struct PublishingCatalog {
|
||||
generation: AtomicU64,
|
||||
queries: AtomicUsize,
|
||||
}
|
||||
|
||||
enum TestCatalogProvider {
|
||||
enum TestCatalog {
|
||||
NoProviders,
|
||||
Unavailable,
|
||||
Empty,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for TestCatalogProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"alpha"
|
||||
}
|
||||
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::new(0)
|
||||
}
|
||||
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
impl ProviderCatalogPort for TestCatalog {
|
||||
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
match self {
|
||||
Self::Unavailable => Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
)),
|
||||
Self::Empty => Ok(Vec::new()),
|
||||
Self::NoProviders => BTreeMap::new(),
|
||||
Self::Unavailable | Self::Empty => catalog_generations(0),
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
fn query_model_capabilities(
|
||||
&self,
|
||||
_: ProviderRequest,
|
||||
_: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
))
|
||||
_: &ProviderKind,
|
||||
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
|
||||
Box::pin(async move {
|
||||
match self {
|
||||
Self::NoProviders | Self::Unavailable => Err(ProviderCatalogUnavailable),
|
||||
Self::Empty => Ok(Vec::new()),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for PublishingCatalogProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"alpha"
|
||||
fn catalog_generations(generation: u64) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
BTreeMap::from([(
|
||||
ProviderKind::new("alpha").expect("provider"),
|
||||
ProviderCatalogGeneration::new(generation),
|
||||
)])
|
||||
}
|
||||
|
||||
impl ProviderCatalogPort for PublishingCatalog {
|
||||
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
catalog_generations(self.generation.load(Ordering::SeqCst))
|
||||
}
|
||||
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::new(self.generation.load(Ordering::SeqCst))
|
||||
}
|
||||
|
||||
async fn query_model_capabilities(
|
||||
fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
self.generation.store(1, Ordering::SeqCst);
|
||||
}
|
||||
Ok(vec![
|
||||
ProviderModelCapabilities::new(
|
||||
UpstreamModelId::new("upstream-model").expect("model"),
|
||||
ModelCapabilities::new(
|
||||
std::collections::BTreeSet::from([OperationKind::Generate]),
|
||||
_: &ProviderKind,
|
||||
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
|
||||
Box::pin(async move {
|
||||
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
self.generation.store(1, Ordering::SeqCst);
|
||||
}
|
||||
Ok(vec![
|
||||
ProviderModelCapabilities::new(
|
||||
UpstreamModelId::new("upstream-model").expect("model"),
|
||||
ModelCapabilities::new(
|
||||
std::collections::BTreeSet::from([OperationKind::Generate]),
|
||||
None,
|
||||
),
|
||||
)
|
||||
.with_presentation(ModelPresentation::new(
|
||||
Some("Upstream Model".to_owned()),
|
||||
None,
|
||||
),
|
||||
)
|
||||
.with_presentation(ModelPresentation::new(
|
||||
Some("Upstream Model".to_owned()),
|
||||
None,
|
||||
)),
|
||||
])
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_: ProviderRequest,
|
||||
_: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
))
|
||||
)),
|
||||
])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -150,11 +125,10 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
|
||||
|
||||
#[test]
|
||||
fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
|
||||
let providers =
|
||||
ProviderRegistry::new([Arc::new(TestCatalogProvider::Unavailable) as Arc<dyn Provider>])
|
||||
.expect("provider registry");
|
||||
let compiler =
|
||||
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
|
||||
let compiler = RuntimeSnapshotCompiler::new(
|
||||
Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))),
|
||||
Arc::new(TestCatalog::Unavailable),
|
||||
);
|
||||
|
||||
let snapshot = block_on(compiler.compile()).expect("compile snapshot");
|
||||
let provider = ProviderKind::new("alpha").expect("provider");
|
||||
@@ -170,11 +144,10 @@ fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
|
||||
|
||||
#[test]
|
||||
fn known_empty_catalog_should_report_model_not_found() {
|
||||
let providers =
|
||||
ProviderRegistry::new([Arc::new(TestCatalogProvider::Empty) as Arc<dyn Provider>])
|
||||
.expect("provider registry");
|
||||
let compiler =
|
||||
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
|
||||
let compiler = RuntimeSnapshotCompiler::new(
|
||||
Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))),
|
||||
Arc::new(TestCatalog::Empty),
|
||||
);
|
||||
let snapshot = block_on(compiler.compile()).expect("compile empty catalog");
|
||||
let error = snapshot
|
||||
.plan(
|
||||
@@ -199,18 +172,18 @@ fn known_empty_catalog_should_report_model_not_found() {
|
||||
|
||||
#[test]
|
||||
fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
|
||||
let provider = Arc::new(PublishingCatalogProvider {
|
||||
let catalog = Arc::new(PublishingCatalog {
|
||||
generation: AtomicU64::new(0),
|
||||
queries: AtomicUsize::new(0),
|
||||
});
|
||||
let providers =
|
||||
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).expect("provider registry");
|
||||
let compiler =
|
||||
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
|
||||
let compiler = RuntimeSnapshotCompiler::new(
|
||||
Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))),
|
||||
catalog.clone(),
|
||||
);
|
||||
|
||||
let snapshot = block_on(compiler.compile()).expect("stable catalog snapshot");
|
||||
|
||||
assert_eq!(provider.queries.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(catalog.queries.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(
|
||||
snapshot
|
||||
.provider_catalog_generations()
|
||||
@@ -311,7 +284,7 @@ fn facts_revision(facts: &SnapshotFacts) -> ConfigRevision {
|
||||
}
|
||||
|
||||
fn compiler(store: Arc<dyn SnapshotStorePort>) -> RuntimeSnapshotCompiler {
|
||||
RuntimeSnapshotCompiler::new(store, ProviderRegistry::default())
|
||||
RuntimeSnapshotCompiler::new(store, Arc::new(TestCatalog::NoProviders))
|
||||
}
|
||||
|
||||
fn revision(value: u64) -> ConfigRevision {
|
||||
|
||||
@@ -1,30 +1,42 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::channel::{mpsc, oneshot};
|
||||
use futures::executor::block_on;
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use gateway_core::account::{AccountSelectionPolicy, RotationStrategy};
|
||||
use gateway_core::engine::provider::ProviderRegistry;
|
||||
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::snapshot::{
|
||||
RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
|
||||
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotProviderAccountFacts, SnapshotSettingsFacts,
|
||||
SnapshotStoreError, SnapshotStorePort,
|
||||
};
|
||||
use gateway_core::routing::{ConfigRevision, RuntimeSnapshot};
|
||||
use gateway_core::routing::{
|
||||
AccountGroupId, ConfigRevision, ProviderCatalogGeneration, ProviderCatalogPort,
|
||||
ProviderCatalogUnavailable, ProviderKind, ProviderModelCapabilities, RuntimeSnapshot,
|
||||
};
|
||||
use gateway_core::runtime::{
|
||||
RuntimeSnapshotHandle, RuntimeSnapshotPublisher, SnapshotControl, SnapshotRevisionStream,
|
||||
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
|
||||
};
|
||||
use gateway_core::task::WorkerKind;
|
||||
use gateway_core::task::{
|
||||
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerKind, WorkerRunnable,
|
||||
};
|
||||
|
||||
type ReadGate<T> = Arc<Mutex<Option<oneshot::Receiver<Result<T, SnapshotStoreError>>>>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestSnapshotStore {
|
||||
facts: Arc<Mutex<Result<SnapshotFacts, SnapshotStoreError>>>,
|
||||
current_revision: Arc<Mutex<Result<ConfigRevision, SnapshotStoreError>>>,
|
||||
next_load: ReadGate<SnapshotFacts>,
|
||||
next_revision: ReadGate<ConfigRevision>,
|
||||
loads: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl TestSnapshotStore {
|
||||
@@ -36,23 +48,45 @@ impl TestSnapshotStore {
|
||||
Self {
|
||||
facts: Arc::new(Mutex::new(facts)),
|
||||
current_revision: Arc::new(Mutex::new(current_revision)),
|
||||
next_load: Arc::default(),
|
||||
next_revision: Arc::default(),
|
||||
loads: Arc::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SnapshotStorePort for TestSnapshotStore {
|
||||
fn load_snapshot_facts(&self) -> BoxFuture<'_, Result<SnapshotFacts, SnapshotStoreError>> {
|
||||
Box::pin(async move { self.facts.lock().expect("facts lock").clone() })
|
||||
Box::pin(async move {
|
||||
self.loads.fetch_add(1, Ordering::SeqCst);
|
||||
let gate = self.next_load.lock().expect("load gate lock").take();
|
||||
if let Some(gate) = gate {
|
||||
return gate.await.expect("release load");
|
||||
}
|
||||
self.facts.lock().expect("facts lock").clone()
|
||||
})
|
||||
}
|
||||
|
||||
fn current_config_revision(&self) -> BoxFuture<'_, Result<ConfigRevision, SnapshotStoreError>> {
|
||||
Box::pin(async move { self.current_revision.lock().expect("revision lock").clone() })
|
||||
Box::pin(async move {
|
||||
let gate = self
|
||||
.next_revision
|
||||
.lock()
|
||||
.expect("revision gate lock")
|
||||
.take();
|
||||
if let Some(gate) = gate {
|
||||
return gate.await.expect("release revision read");
|
||||
}
|
||||
self.current_revision.lock().expect("revision lock").clone()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestSnapshotSubscriptions {
|
||||
published: Mutex<Vec<ConfigRevision>>,
|
||||
stream: Mutex<Option<SnapshotRevisionStream>>,
|
||||
next_publish: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
}
|
||||
|
||||
impl SnapshotSubscriptionPort for TestSnapshotSubscriptions {
|
||||
@@ -61,6 +95,10 @@ impl SnapshotSubscriptionPort for TestSnapshotSubscriptions {
|
||||
revision: ConfigRevision,
|
||||
) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>> {
|
||||
Box::pin(async move {
|
||||
let gate = self.next_publish.lock().expect("publish gate lock").take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("release notification");
|
||||
}
|
||||
self.published
|
||||
.lock()
|
||||
.expect("published lock")
|
||||
@@ -72,7 +110,14 @@ impl SnapshotSubscriptionPort for TestSnapshotSubscriptions {
|
||||
fn subscribe_snapshot_revisions(
|
||||
&self,
|
||||
) -> BoxFuture<'_, Result<SnapshotRevisionStream, SnapshotSubscriptionError>> {
|
||||
Box::pin(async move { Ok(Box::pin(futures::stream::empty()) as SnapshotRevisionStream) })
|
||||
Box::pin(async move {
|
||||
Ok(self
|
||||
.stream
|
||||
.lock()
|
||||
.expect("subscription stream lock")
|
||||
.take()
|
||||
.unwrap_or_else(|| Box::pin(futures::stream::pending())))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -187,6 +232,398 @@ fn publisher_should_contribute_reconciliation_and_subscription_workers() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overlapping_refreshes_should_not_restore_revoked_account_scope() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let catalog = Arc::new(TestCatalog::default());
|
||||
*catalog.next_query.lock().expect("catalog gate lock") = Some(gate);
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(catalog_compiler(store.clone(), catalog.clone())),
|
||||
handle.clone(),
|
||||
Arc::new(TestSnapshotSubscriptions::default()),
|
||||
);
|
||||
let other_publisher = publisher.clone();
|
||||
|
||||
// 旧刷新已读完 revision 10 的完整事实,停在目录查询;不用 sleep 碰调度概率。
|
||||
let mut old = Box::pin(publisher.refresh());
|
||||
assert!(futures::poll!(old.as_mut()).is_pending());
|
||||
assert_eq!(catalog.queries.load(Ordering::SeqCst), 1);
|
||||
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
|
||||
*store.current_revision.lock().expect("revision lock") = Ok(revision(11));
|
||||
let mut new = Box::pin(other_publisher.refresh());
|
||||
assert!(futures::poll!(new.as_mut()).is_pending());
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(handle.revision(), Some(revision(9)));
|
||||
|
||||
release.send(()).expect("release old catalog query");
|
||||
assert_eq!(old.await.expect("old refresh"), revision(10));
|
||||
let frozen_old = handle.acquire().expect("old request snapshot");
|
||||
assert!(removed_account_allowed(&frozen_old));
|
||||
assert_eq!(new.await.expect("new refresh"), revision(11));
|
||||
let current = handle.acquire().expect("new request snapshot");
|
||||
assert!(!removed_account_allowed(¤t));
|
||||
assert_eq!(current.revision(), revision(11));
|
||||
assert!(removed_account_allowed(&frozen_old));
|
||||
assert_eq!(frozen_old.revision(), revision(10));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_committed_refresh_should_suspend_before_next_publication() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(11, 11))));
|
||||
*store.next_load.lock().expect("load gate lock") = Some(gate);
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
|
||||
let subscriptions = Arc::new(TestSnapshotSubscriptions::default());
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store.clone())),
|
||||
handle.clone(),
|
||||
subscriptions.clone(),
|
||||
);
|
||||
let other_publisher = publisher.clone();
|
||||
let mut old = publisher.publish_committed(revision(10));
|
||||
assert!(futures::poll!(old.as_mut()).is_pending());
|
||||
let mut new = other_publisher.publish_committed(revision(11));
|
||||
assert!(futures::poll!(new.as_mut()).is_pending());
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
|
||||
|
||||
release
|
||||
.send(Err(SnapshotStoreError::unavailable()))
|
||||
.expect("release failed facts read");
|
||||
old.await;
|
||||
assert!(handle.acquire().is_err());
|
||||
new.await;
|
||||
assert_eq!(handle.revision(), Some(revision(11)));
|
||||
assert_eq!(
|
||||
*subscriptions.published.lock().expect("published lock"),
|
||||
vec![revision(10), revision(11)],
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconciliation_revision_failure_should_not_suspend_later_committed_refresh() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(11, 11))));
|
||||
*store.next_revision.lock().expect("revision gate lock") = Some(gate);
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store.clone())),
|
||||
handle.clone(),
|
||||
Arc::new(TestSnapshotSubscriptions::default()),
|
||||
);
|
||||
let (task, context) = reconciliation_task(&publisher);
|
||||
let mut reconcile = task.run_cycle(context);
|
||||
assert!(futures::poll!(reconcile.as_mut()).is_pending());
|
||||
let mut committed = publisher.publish_committed(revision(11));
|
||||
assert!(futures::poll!(committed.as_mut()).is_pending());
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 0);
|
||||
|
||||
release
|
||||
.send(Err(SnapshotStoreError::unavailable()))
|
||||
.expect("release failed revision read");
|
||||
assert!(reconcile.await.is_err());
|
||||
assert!(handle.acquire().is_err());
|
||||
committed.await;
|
||||
assert_eq!(handle.revision(), Some(revision(11)));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subscription_refresh_should_wait_for_inflight_compile_and_reload_authoritative_facts() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let catalog = Arc::new(TestCatalog::default());
|
||||
*catalog.next_query.lock().expect("catalog gate lock") = Some(gate);
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
|
||||
let (notifications, stream) = mpsc::unbounded();
|
||||
let subscriptions = Arc::new(TestSnapshotSubscriptions {
|
||||
stream: Mutex::new(Some(Box::pin(stream))),
|
||||
..TestSnapshotSubscriptions::default()
|
||||
});
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(catalog_compiler(store.clone(), catalog)),
|
||||
handle.clone(),
|
||||
subscriptions,
|
||||
);
|
||||
let task = publisher
|
||||
.worker_contributions()
|
||||
.expect("workers")
|
||||
.into_iter()
|
||||
.find_map(|contribution| match contribution {
|
||||
WorkerContribution::Registration(registration) => match registration.runnable {
|
||||
WorkerRunnable::Daemon { task, .. } => Some(task),
|
||||
WorkerRunnable::Scheduled { .. } => None,
|
||||
},
|
||||
WorkerContribution::Disabled { .. } => None,
|
||||
})
|
||||
.expect("subscription task");
|
||||
let cancellation = CancellationToken::new();
|
||||
let mut daemon = task.run(cancellation.clone());
|
||||
let mut old = Box::pin(publisher.refresh());
|
||||
assert!(futures::poll!(old.as_mut()).is_pending());
|
||||
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
|
||||
// 通知只是提示,即使版本过期也要在取得发布权后重新读取 Store。
|
||||
notifications
|
||||
.unbounded_send(Ok(revision(3)))
|
||||
.expect("notify");
|
||||
assert!(futures::poll!(daemon.as_mut()).is_pending());
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(handle.revision(), Some(revision(9)));
|
||||
|
||||
release.send(()).expect("release old catalog query");
|
||||
old.await.expect("old refresh");
|
||||
assert!(futures::poll!(daemon.as_mut()).is_pending());
|
||||
let current = handle.acquire().expect("subscription snapshot");
|
||||
assert_eq!(current.revision(), revision(11));
|
||||
assert!(!removed_account_allowed(¤t));
|
||||
cancellation.cancel();
|
||||
daemon.await.expect("stop subscription");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconciliation_should_fail_closed_on_persisted_revision_rollback_and_recover() {
|
||||
block_on(async {
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(7, 7))));
|
||||
*store.facts.lock().expect("facts lock") = Err(SnapshotStoreError::unavailable());
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(8));
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store.clone())),
|
||||
handle.clone(),
|
||||
Arc::new(TestSnapshotSubscriptions::default()),
|
||||
);
|
||||
let (task, context) = reconciliation_task(&publisher);
|
||||
|
||||
assert!(task.run_cycle(context.clone()).await.is_err());
|
||||
assert!(handle.acquire().is_err());
|
||||
*store.facts.lock().expect("facts lock") = Ok(facts(7, 7));
|
||||
task.run_cycle(context.clone()).await.expect("recover");
|
||||
assert_eq!(handle.revision(), Some(revision(7)));
|
||||
|
||||
// 有效的持久回退也必须发布,不能用 revision 数值单调性掩盖竞争。
|
||||
*store.facts.lock().expect("facts lock") = Ok(facts(6, 6));
|
||||
*store.current_revision.lock().expect("revision lock") = Ok(revision(6));
|
||||
task.run_cycle(context).await.expect("publish rollback");
|
||||
assert_eq!(handle.revision(), Some(revision(6)));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn catalog_only_reconciliation_should_keep_snapshot_on_failure_and_retry_generation() {
|
||||
block_on(async {
|
||||
let catalog = Arc::new(TestCatalog::default());
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
|
||||
let handle = RuntimeSnapshotHandle::default();
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(catalog_compiler(store.clone(), catalog.clone())),
|
||||
handle.clone(),
|
||||
Arc::new(TestSnapshotSubscriptions::default()),
|
||||
);
|
||||
publisher.refresh().await.expect("initial snapshot");
|
||||
let frozen = handle.acquire().expect("frozen snapshot");
|
||||
let (task, context) = reconciliation_task(&publisher);
|
||||
task.run_cycle(context.clone()).await.expect("unchanged");
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 1);
|
||||
assert!(Arc::ptr_eq(
|
||||
&frozen,
|
||||
&handle.acquire().expect("unchanged snapshot")
|
||||
));
|
||||
|
||||
catalog.generation.store(1, Ordering::SeqCst);
|
||||
*store.facts.lock().expect("facts lock") = Err(SnapshotStoreError::unavailable());
|
||||
assert!(task.run_cycle(context.clone()).await.is_err());
|
||||
assert!(Arc::ptr_eq(
|
||||
&frozen,
|
||||
&handle.acquire().expect("retained snapshot")
|
||||
));
|
||||
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(10, true));
|
||||
task.run_cycle(context)
|
||||
.await
|
||||
.expect("retry catalog refresh");
|
||||
let current = handle.acquire().expect("refreshed catalog");
|
||||
assert_eq!(current.revision(), revision(10));
|
||||
assert_eq!(
|
||||
current
|
||||
.provider_catalog_generations()
|
||||
.values()
|
||||
.copied()
|
||||
.collect::<Vec<_>>(),
|
||||
vec![ProviderCatalogGeneration::new(1)],
|
||||
);
|
||||
assert_eq!(
|
||||
frozen
|
||||
.provider_catalog_generations()
|
||||
.values()
|
||||
.copied()
|
||||
.collect::<Vec<_>>(),
|
||||
vec![ProviderCatalogGeneration::new(0)],
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dropping_inflight_refresh_should_release_publication_for_next_refresh() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(11, 11))));
|
||||
*store.next_load.lock().expect("load gate lock") = Some(gate);
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store)),
|
||||
handle.clone(),
|
||||
Arc::new(TestSnapshotSubscriptions::default()),
|
||||
);
|
||||
let mut cancelled = Box::pin(publisher.refresh());
|
||||
assert!(futures::poll!(cancelled.as_mut()).is_pending());
|
||||
let mut next = Box::pin(publisher.refresh());
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
drop(cancelled);
|
||||
assert!(release.send(Ok(facts(10, 10))).is_err());
|
||||
assert_eq!(next.await.expect("next refresh"), revision(11));
|
||||
assert_eq!(handle.revision(), Some(revision(11)));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_notification_should_not_block_next_local_refresh() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(10, 10))));
|
||||
let handle = RuntimeSnapshotHandle::new(empty_snapshot(9));
|
||||
let subscriptions = Arc::new(TestSnapshotSubscriptions {
|
||||
next_publish: Mutex::new(Some(gate)),
|
||||
..TestSnapshotSubscriptions::default()
|
||||
});
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store.clone())),
|
||||
handle.clone(),
|
||||
subscriptions,
|
||||
);
|
||||
let mut committed = publisher.publish_committed(revision(10));
|
||||
assert!(futures::poll!(committed.as_mut()).is_pending());
|
||||
assert_eq!(handle.revision(), Some(revision(10)));
|
||||
*store.facts.lock().expect("facts lock") = Ok(facts(11, 11));
|
||||
let mut next = Box::pin(publisher.refresh());
|
||||
assert_eq!(
|
||||
futures::poll!(next.as_mut()),
|
||||
std::task::Poll::Ready(Ok(revision(11))),
|
||||
);
|
||||
release.send(()).expect("release notification");
|
||||
committed.await;
|
||||
assert_eq!(handle.revision(), Some(revision(11)));
|
||||
});
|
||||
}
|
||||
|
||||
fn reconciliation_task(
|
||||
publisher: &RuntimeSnapshotPublisher,
|
||||
) -> (Box<dyn ScheduledTask>, WorkerCycleContext) {
|
||||
publisher
|
||||
.worker_contributions()
|
||||
.expect("workers")
|
||||
.into_iter()
|
||||
.find_map(|contribution| match contribution {
|
||||
WorkerContribution::Registration(registration) => match registration.runnable {
|
||||
WorkerRunnable::Scheduled { task, .. } => Some((
|
||||
task,
|
||||
WorkerCycleContext::new(registration.id, None, CancellationToken::new()),
|
||||
)),
|
||||
WorkerRunnable::Daemon { .. } => None,
|
||||
},
|
||||
WorkerContribution::Disabled { .. } => None,
|
||||
})
|
||||
.expect("reconciliation task")
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestCatalog {
|
||||
generation: AtomicU64,
|
||||
queries: AtomicUsize,
|
||||
next_query: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
}
|
||||
|
||||
impl ProviderCatalogPort for TestCatalog {
|
||||
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
BTreeMap::from([(
|
||||
ProviderKind::new("audit").expect("provider"),
|
||||
ProviderCatalogGeneration::new(self.generation.load(Ordering::SeqCst)),
|
||||
)])
|
||||
}
|
||||
|
||||
fn query_model_capabilities(
|
||||
&self,
|
||||
_: &ProviderKind,
|
||||
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
|
||||
Box::pin(async move {
|
||||
self.queries.fetch_add(1, Ordering::SeqCst);
|
||||
let gate = self.next_query.lock().expect("catalog gate lock").take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("release catalog query");
|
||||
}
|
||||
Ok(Vec::new())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn catalog_compiler(
|
||||
store: Arc<dyn SnapshotStorePort>,
|
||||
catalog: Arc<TestCatalog>,
|
||||
) -> RuntimeSnapshotCompiler {
|
||||
RuntimeSnapshotCompiler::new(store, catalog)
|
||||
}
|
||||
|
||||
fn scoped_facts(value: u64, allow_removed: bool) -> SnapshotFacts {
|
||||
let group = AccountGroupId::new(format!("grp_{:032x}", 1)).expect("group");
|
||||
let keep = ProviderAccountId::new("acct_audit_keep").expect("kept account");
|
||||
let removed = ProviderAccountId::new("acct_audit_removed").expect("removed account");
|
||||
let mut memberships = vec![SnapshotAccountGroupMemberFacts::new(
|
||||
group.clone(),
|
||||
keep.clone(),
|
||||
)];
|
||||
if allow_removed {
|
||||
memberships.push(SnapshotAccountGroupMemberFacts::new(
|
||||
group.clone(),
|
||||
removed.clone(),
|
||||
));
|
||||
}
|
||||
SnapshotFacts::new(
|
||||
revision(value),
|
||||
revision(value),
|
||||
SnapshotSettingsFacts::new(3, 0, "smart", BTreeMap::new(), None, None),
|
||||
vec![SnapshotClientPolicyFacts::new(
|
||||
ClientApiKeyId::new("key_audit_synthetic").expect("key ID"),
|
||||
PlaintextClientApiKey::new("sk_audit_synthetic_not_a_real_key").expect("synthetic key"),
|
||||
vec![group.clone()],
|
||||
RateLimits::unlimited(),
|
||||
)],
|
||||
vec![SnapshotAccountGroupFacts::new(
|
||||
group,
|
||||
"Synthetic group".to_owned(),
|
||||
true,
|
||||
)],
|
||||
vec![
|
||||
SnapshotProviderAccountFacts::new(keep, "audit"),
|
||||
SnapshotProviderAccountFacts::new(removed, "audit"),
|
||||
],
|
||||
memberships,
|
||||
)
|
||||
}
|
||||
|
||||
fn removed_account_allowed(snapshot: &RuntimeSnapshot) -> bool {
|
||||
snapshot
|
||||
.client_policies()
|
||||
.next()
|
||||
.expect("client policy")
|
||||
.account_scope()
|
||||
.allows(&ProviderAccountId::new("acct_audit_removed").expect("removed account"))
|
||||
}
|
||||
|
||||
fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts {
|
||||
SnapshotFacts::new(
|
||||
revision(config_revision),
|
||||
@@ -212,7 +649,7 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
|
||||
}
|
||||
|
||||
fn compiler(store: Arc<dyn SnapshotStorePort>) -> RuntimeSnapshotCompiler {
|
||||
RuntimeSnapshotCompiler::new(store, ProviderRegistry::default())
|
||||
catalog_compiler(store, Arc::new(TestCatalog::default()))
|
||||
}
|
||||
|
||||
fn empty_snapshot(value: u64) -> RuntimeSnapshot {
|
||||
|
||||
@@ -47,7 +47,6 @@ use crate::{
|
||||
use super::{ControlPlaneRepository, PgControlPlaneRepository};
|
||||
|
||||
const ENTITY: &str = "client API key";
|
||||
const KEY_LENGTH: usize = 46;
|
||||
const CLIENT_API_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1);
|
||||
const CLIENT_API_KEY_LAST_USED_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
|
||||
@@ -355,7 +354,8 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
query.validate()?;
|
||||
let total = count_client_api_keys(&self.pool, query.search.as_deref()).await?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new(
|
||||
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
|
||||
"select k.id, k.name, k.label,
|
||||
left(k.key, least(10, length(k.key) / 2)) as prefix, k.enabled,
|
||||
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
|
||||
k.updated_at, '[]'::jsonb as groups, '{}'::text[] as provider_kinds
|
||||
from client_api_keys k
|
||||
@@ -414,7 +414,8 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
|
||||
require_nonempty(ENTITY, "id", id)?;
|
||||
let record = sqlx::query(
|
||||
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
|
||||
"select k.id, k.name, k.label,
|
||||
left(k.key, least(10, length(k.key) / 2)) as prefix, k.enabled,
|
||||
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
|
||||
k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups,
|
||||
case
|
||||
@@ -977,6 +978,7 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
key: &NewClientApiKey,
|
||||
) -> StoreResult<()> {
|
||||
key.validate()?;
|
||||
ensure_client_key_name_available(transaction, &key.id, key.name.trim()).await?;
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, label, key, enabled, max_concurrency, requests_per_minute,
|
||||
@@ -984,7 +986,7 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now(), $7::text::numeric, $8::text::numeric)",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.bind(&key.name)
|
||||
.bind(key.name.trim())
|
||||
.bind(&key.label)
|
||||
.bind(&key.key)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
@@ -993,7 +995,20 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
.bind(key.budget.weekly_usd.canonical())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
|
||||
.map_err(|error| {
|
||||
if error.as_database_error().is_some_and(|database| {
|
||||
database.is_unique_violation()
|
||||
&& database.constraint() == Some("client_api_keys_key_sha256_idx")
|
||||
}) {
|
||||
StoreError::Conflict {
|
||||
entity: ENTITY,
|
||||
id: key.id.clone(),
|
||||
kind: crate::ConflictKind::InvalidTransition,
|
||||
}
|
||||
} else {
|
||||
postgres_unavailable("insert client API key in transaction")
|
||||
}
|
||||
})?;
|
||||
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1003,6 +1018,7 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
key: &UpdateClientApiKeyDetails,
|
||||
) -> StoreResult<()> {
|
||||
key.validate()?;
|
||||
ensure_client_key_name_available(transaction, &key.id, key.name.trim()).await?;
|
||||
let result = sqlx::query(
|
||||
"update client_api_keys
|
||||
set name = $2, label = $3, max_concurrency = $4,
|
||||
@@ -1012,7 +1028,7 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
where id = $1",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.bind(&key.name)
|
||||
.bind(key.name.trim())
|
||||
.bind(&key.label)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
@@ -1025,6 +1041,32 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await
|
||||
}
|
||||
|
||||
async fn ensure_client_key_name_available(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
id: &str,
|
||||
name: &str,
|
||||
) -> StoreResult<()> {
|
||||
// 调用方已通过递增配置版本持有控制面行锁,查重与写入在同一事务内串行执行。
|
||||
// 不回填历史重名数据;创建和保存时统一校验,更新排除当前记录。
|
||||
let duplicate: bool = sqlx::query_scalar(
|
||||
"select exists(select 1 from client_api_keys
|
||||
where lower(btrim(name)) = lower($1) and id <> $2)",
|
||||
)
|
||||
.bind(name)
|
||||
.bind(id)
|
||||
.fetch_one(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("check client API key name"))?;
|
||||
if duplicate {
|
||||
return Err(StoreError::Conflict {
|
||||
entity: ENTITY,
|
||||
id: id.to_owned(),
|
||||
kind: crate::ConflictKind::DuplicateName,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn set_client_api_key_enabled_in_transaction(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
id: &str,
|
||||
@@ -1119,18 +1161,8 @@ fn validate_group_ids(group_ids: &[String]) -> StoreResult<()> {
|
||||
}
|
||||
|
||||
fn validate_key(key: &str) -> StoreResult<()> {
|
||||
let valid = key.len() == KEY_LENGTH
|
||||
&& key.starts_with("sk_")
|
||||
&& key[3..]
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'));
|
||||
if valid {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(invalid(
|
||||
"key must be sk_ followed by 43 URL-safe characters",
|
||||
))
|
||||
}
|
||||
PlaintextClientApiKey::validate(key)
|
||||
.map_err(|_| invalid("key must contain nonempty visible ASCII characters"))
|
||||
}
|
||||
|
||||
fn client_secret_from_row(
|
||||
@@ -1204,9 +1236,6 @@ fn push_client_key_search(statement: &mut QueryBuilder<Postgres>, search: Option
|
||||
statement.push_bind(prefix.clone());
|
||||
statement.push(" escape '\\'");
|
||||
statement.push(" or lower(coalesce(label, '')) like ");
|
||||
statement.push_bind(prefix.clone());
|
||||
statement.push(" escape '\\'");
|
||||
statement.push(" or lower(left(key, 10)) like ");
|
||||
statement.push_bind(prefix);
|
||||
statement.push(" escape '\\'");
|
||||
statement.push(")");
|
||||
|
||||
@@ -254,5 +254,6 @@ fn push_prefix_search(statement: &mut QueryBuilder<Postgres>, columns: &[&str],
|
||||
statement.push_bind(pattern.clone());
|
||||
statement.push(" escape '\\'");
|
||||
}
|
||||
push_client_key_name_search(statement, "mr.client_api_key_ref", pattern);
|
||||
statement.push(")");
|
||||
}
|
||||
|
||||
@@ -85,6 +85,8 @@ pub(crate) fn push_usage_filter(
|
||||
for column in [
|
||||
"client_api_key_ref",
|
||||
"provider_account_ref",
|
||||
"provider_account_email_snapshot",
|
||||
"provider_account_name_snapshot",
|
||||
"requested_model_id",
|
||||
"upstream_model_id",
|
||||
"upstream_request_id",
|
||||
@@ -92,22 +94,25 @@ pub(crate) fn push_usage_filter(
|
||||
query.push(format!(" escape '\\' or {alias}.{column} like "));
|
||||
query.push_bind(pattern.clone());
|
||||
}
|
||||
if value.starts_with("sk_") {
|
||||
query.push(format!(
|
||||
" escape '\\' or exists (
|
||||
select 1 from client_api_keys searched_client_key
|
||||
where searched_client_key.id = {alias}.client_api_key_ref
|
||||
and searched_client_key.key like "
|
||||
));
|
||||
query.push_bind(pattern);
|
||||
query.push(" escape '\\')");
|
||||
} else {
|
||||
query.push(" escape '\\'");
|
||||
}
|
||||
query.push(" escape '\\'");
|
||||
push_client_key_name_search(query, &format!("{alias}.client_api_key_ref"), pattern);
|
||||
query.push(")");
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn push_client_key_name_search(
|
||||
query: &mut QueryBuilder<Postgres>,
|
||||
key_ref: &str,
|
||||
pattern: String,
|
||||
) {
|
||||
query.push(format!(
|
||||
" or exists (select 1 from client_api_keys searched_client_key
|
||||
where searched_client_key.id = {key_ref} and searched_client_key.name ilike "
|
||||
));
|
||||
query.push_bind(pattern);
|
||||
query.push(" escape '\\')");
|
||||
}
|
||||
|
||||
pub(crate) fn literal_prefix_pattern(value: &str) -> String {
|
||||
format!(
|
||||
"{}%",
|
||||
|
||||
@@ -476,85 +476,11 @@ impl AccountStore for PgAdminAccountStore {
|
||||
self.usage_by_windows(windows).await
|
||||
}
|
||||
|
||||
async fn list_credentials(
|
||||
async fn load_quota_forecast_history(
|
||||
&self,
|
||||
provider_kind: &ProviderKind,
|
||||
query: CredentialListQuery,
|
||||
) -> AdminStoreResult<CredentialPage> {
|
||||
let (control_plane, accounts) = futures::try_join!(
|
||||
self.control_plane.load_control_plane(),
|
||||
self.accounts
|
||||
.list_provider_accounts(Some(provider_kind.as_str()), true),
|
||||
)
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let mut accounts = accounts
|
||||
.into_iter()
|
||||
.filter(|account| account.provider_kind == provider_kind.as_str())
|
||||
.filter(|account| {
|
||||
query
|
||||
.credential_state
|
||||
.as_ref()
|
||||
.is_none_or(|expected| expected.matches(account.credential_state))
|
||||
})
|
||||
.filter(|account| {
|
||||
query
|
||||
.enabled
|
||||
.is_none_or(|enabled| account.enabled == enabled)
|
||||
})
|
||||
.filter(|account| {
|
||||
let CredentialListWindow::Page {
|
||||
cursor: Some(cursor),
|
||||
..
|
||||
} = &query.window
|
||||
else {
|
||||
return true;
|
||||
};
|
||||
account.created_at > cursor.created_at
|
||||
|| (account.created_at == cursor.created_at
|
||||
&& account.id.as_str() > cursor.account_id.as_str())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if matches!(&query.window, CredentialListWindow::Page { .. }) {
|
||||
accounts.sort_by(|left, right| {
|
||||
left.created_at
|
||||
.cmp(&right.created_at)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
}
|
||||
let next_cursor = match query.window {
|
||||
CredentialListWindow::All => None,
|
||||
CredentialListWindow::Page { page_size, .. } => {
|
||||
let page_size = usize::from(page_size.get());
|
||||
let has_more = accounts.len() > page_size;
|
||||
accounts.truncate(page_size);
|
||||
has_more
|
||||
.then(|| accounts.last())
|
||||
.flatten()
|
||||
.map(|account| {
|
||||
Ok(CredentialCursor {
|
||||
created_at: account.created_at,
|
||||
account_id: CoreProviderAccountId::new(account.id.clone()).map_err(
|
||||
|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"persisted Provider account ID is invalid",
|
||||
)
|
||||
},
|
||||
)?,
|
||||
})
|
||||
})
|
||||
.transpose()?
|
||||
}
|
||||
};
|
||||
Ok(CredentialPage {
|
||||
config_revision: admin_revision(control_plane.settings.config_revision)?,
|
||||
items: accounts
|
||||
.into_iter()
|
||||
.map(admin_account_record)
|
||||
.collect::<AdminStoreResult<Vec<_>>>()?,
|
||||
next_cursor,
|
||||
})
|
||||
window: &AccountUsageWindowQuery,
|
||||
) -> AdminStoreResult<gateway_admin::model::quota_forecast_sampling::QuotaForecastHistory> {
|
||||
super::quota_forecast::load_history(&self.pool, &self.query_budget, window).await
|
||||
}
|
||||
|
||||
async fn credential_details(
|
||||
|
||||
@@ -21,9 +21,8 @@ use gateway_admin::{
|
||||
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
|
||||
},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, AuthorizationCredentialCommit, CredentialCursor,
|
||||
CredentialDetails, CredentialImportCommit, CredentialImportResult, CredentialListQuery,
|
||||
CredentialListWindow, CredentialMutationResult, CredentialPage,
|
||||
AuthorizationCommit, AuthorizationCredentialCommit, CredentialDetails,
|
||||
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
|
||||
CredentialRotationCommit, PreparedCredentialCreate, PreparedCredentialImport,
|
||||
PreparedCredentialRotationFacts, ProviderDocument, ProviderExportCredentialInput,
|
||||
},
|
||||
@@ -62,6 +61,7 @@ mod admin_adapter;
|
||||
mod admin_queries;
|
||||
mod core_adapter;
|
||||
mod mapping;
|
||||
mod quota_forecast;
|
||||
mod repository;
|
||||
mod rows;
|
||||
mod runtime;
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
//! 预测只读采样:先聚合数值、再读取有界的 Provider 文档,不扫描模型排行维度。
|
||||
|
||||
use gateway_admin::model::quota_forecast_sampling::{
|
||||
MAX_FORECAST_HISTORY_POINTS, QuotaForecastHistory, QuotaForecastHistoryPoint,
|
||||
QuotaForecastUsage,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
pub(super) async fn load_history(
|
||||
pool: &PgPool,
|
||||
budget: &super::super::ObservabilityQueryBudget,
|
||||
window: &AccountUsageWindowQuery,
|
||||
) -> AdminStoreResult<QuotaForecastHistory> {
|
||||
validate_admin_account_ids(std::slice::from_ref(&window.account_id))
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let range = ObservabilityRange::new(window.range.start, window.range.end)
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
if range.end - range.start > TimeDelta::days(32) {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"forecast range exceeds supported quota windows",
|
||||
));
|
||||
}
|
||||
let rows = budget
|
||||
.run("load quota forecast history", async {
|
||||
sqlx::query(sqlx::AssertSqlSafe(history_sql()))
|
||||
.bind(&window.account_id)
|
||||
.bind(range.start)
|
||||
.bind(range.end)
|
||||
.bind(MAX_FORECAST_HISTORY_POINTS as i32)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load quota forecast history"))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let mut history = QuotaForecastHistory::default();
|
||||
for row in &rows {
|
||||
let usage = QuotaForecastUsage {
|
||||
request_count: window_usage_count(row, "request_count")?,
|
||||
tokens: window_usage_count(row, "tokens")?,
|
||||
input_tokens: window_usage_count(row, "input_tokens")?,
|
||||
output_tokens: window_usage_count(row, "output_tokens")?,
|
||||
cached_tokens: window_usage_count(row, "cached_tokens")?,
|
||||
missing_token_count: window_usage_count(row, "missing_token_count")?,
|
||||
known_cost_count: window_usage_count(row, "known_cost_count")?,
|
||||
unavailable_cost_count: window_usage_count(row, "unavailable_cost_count")?,
|
||||
usd: window_usage_value(row, "usd")?,
|
||||
excluded_request_count: window_usage_count(row, "excluded_request_count")?,
|
||||
};
|
||||
if window_usage_value::<bool>(row, "is_total")? {
|
||||
history.usage = usage;
|
||||
history.pending_request_count = window_usage_count(row, "pending_count")?;
|
||||
} else if let Some(document) = window_usage_value::<
|
||||
Option<sqlx::types::Json<serde_json::Map<String, serde_json::Value>>>,
|
||||
>(row, "document")?
|
||||
{
|
||||
history.points.push(QuotaForecastHistoryPoint {
|
||||
started_at: window_usage_value(row, "started_at")?,
|
||||
completed_at: window_usage_value(row, "completed_at")?,
|
||||
usage,
|
||||
provider_observation: ProviderDocument::new(OpaqueProviderData::new(document.0)),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(history)
|
||||
}
|
||||
|
||||
fn history_sql() -> String {
|
||||
let completed_usage = completed_usage_fact_predicate("mr");
|
||||
// 一个语句共享 MVCC 快照。用 RANGE 帧让相同完成时间的点拥有相同累计值,
|
||||
// 避免并发请求的任意行顺序制造不同分子;未完成请求只进入待决计数。
|
||||
format!(
|
||||
"with scoped as materialized (
|
||||
select mr.id, mr.started_at, mr.completed_at,
|
||||
mr.completed_at <= $3 and mr.outcome <> 'running' as settled,
|
||||
coalesce(({completed_usage}), false) as included,
|
||||
mr.provider_observation_json is not null as has_document,
|
||||
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.total_tokens,
|
||||
mr.cost_amount, mr.cost_currency,
|
||||
coalesce(mr.cost_source in ('calculated', 'provider_reported')
|
||||
and mr.cost_amount is not null and mr.cost_currency = 'USD', false) as known_cost
|
||||
from model_requests mr
|
||||
where mr.provider_account_ref = $1
|
||||
and mr.started_at >= $2 and mr.started_at < $3
|
||||
), facts as materialized (
|
||||
select *,
|
||||
(included and settled)::integer as request_count,
|
||||
case when included and settled then coalesce(total_tokens,
|
||||
coalesce(input_tokens, 0) + coalesce(output_tokens, 0)) else 0 end as tokens,
|
||||
case when included and settled then coalesce(input_tokens, 0) else 0 end as inputs,
|
||||
case when included and settled then coalesce(output_tokens, 0) else 0 end as outputs,
|
||||
case when included and settled then coalesce(cached_tokens, 0) else 0 end as cached,
|
||||
(included and settled and total_tokens is null
|
||||
and (input_tokens is null or output_tokens is null))::integer as missing_tokens,
|
||||
(included and settled and known_cost)::integer as known_costs,
|
||||
(included and settled and not known_cost)::integer as missing_costs,
|
||||
case when included and settled and known_cost then cost_amount else 0 end as usd,
|
||||
(not included and settled)::integer as excluded
|
||||
from scoped
|
||||
), cumulative as (
|
||||
select id, started_at, completed_at, has_document,
|
||||
least($4 - 1, floor(extract(epoch from (completed_at - $2))
|
||||
/ greatest(1, extract(epoch from ($3::timestamptz - $2::timestamptz)) / $4))) as bucket,
|
||||
sum(request_count) over w as request_count,
|
||||
sum(tokens) over w as tokens,
|
||||
sum(inputs) over w as input_tokens,
|
||||
sum(outputs) over w as output_tokens,
|
||||
sum(cached) over w as cached_tokens,
|
||||
sum(missing_tokens) over w as missing_token_count,
|
||||
sum(known_costs) over w as known_cost_count,
|
||||
sum(missing_costs) over w as unavailable_cost_count,
|
||||
sum(usd) over w as usd,
|
||||
sum(excluded) over w as excluded_request_count
|
||||
from facts where settled
|
||||
window w as (order by completed_at range between unbounded preceding and current row)
|
||||
), selected as (
|
||||
select distinct on (bucket) * from cumulative
|
||||
where has_document order by bucket, completed_at desc, id desc
|
||||
)
|
||||
select false as is_total, s.started_at, s.completed_at,
|
||||
s.request_count::bigint, s.tokens::bigint, s.input_tokens::bigint,
|
||||
s.output_tokens::bigint, s.cached_tokens::bigint, s.missing_token_count::bigint,
|
||||
s.known_cost_count::bigint, s.unavailable_cost_count::bigint,
|
||||
s.usd::double precision, s.excluded_request_count::bigint,
|
||||
0::bigint as pending_count, mr.provider_observation_json as document
|
||||
from selected s join model_requests mr on mr.id = s.id
|
||||
union all
|
||||
select true, $3, $3,
|
||||
coalesce(sum(request_count), 0)::bigint, coalesce(sum(tokens), 0)::bigint,
|
||||
coalesce(sum(inputs), 0)::bigint, coalesce(sum(outputs), 0)::bigint,
|
||||
coalesce(sum(cached), 0)::bigint, coalesce(sum(missing_tokens), 0)::bigint,
|
||||
coalesce(sum(known_costs), 0)::bigint, coalesce(sum(missing_costs), 0)::bigint,
|
||||
coalesce(sum(usd), 0)::double precision, coalesce(sum(excluded), 0)::bigint,
|
||||
count(*) filter (where settled is not true), null::jsonb
|
||||
from facts
|
||||
order by completed_at, is_total"
|
||||
)
|
||||
}
|
||||
@@ -11,11 +11,15 @@ use sqlx::{Postgres, QueryBuilder};
|
||||
/// 因此不属于用量记录。图片端点由请求意图证明其用量语义,不依赖文本模型字段。
|
||||
/// 客户端 WebSocket 的每个 `response.create` 没有独立 HTTP 状态码;成功终态与下游
|
||||
/// 提交边界已足以证明交付,连接握手的 101 不能冒充单次请求状态。
|
||||
/// OpenAI Provider 将 `generate=false` 归类为 `prewarm`:它只准备连接与上下文,
|
||||
/// 即使响应带有模型和输入 Token,也不是推理用量;原始审计和费用事实仍保留。
|
||||
pub(crate) fn completed_usage_fact_predicate(alias: &str) -> String {
|
||||
let evidence = usage_evidence_predicate(alias);
|
||||
format!(
|
||||
"{alias}.outcome = 'succeeded'
|
||||
and {alias}.downstream_committed_at is not null
|
||||
and ({alias}.provider_kind is distinct from 'openai'
|
||||
or {alias}.request_kind is distinct from 'prewarm')
|
||||
and (({alias}.client_transport = 'websocket' and {alias}.client_status_code is null)
|
||||
or {alias}.client_status_code between 200 and 399)
|
||||
and ({evidence})"
|
||||
|
||||
@@ -13,6 +13,7 @@ pub enum StoreBackend {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ConflictKind {
|
||||
StaleRevision,
|
||||
DuplicateName,
|
||||
AlreadyFinalized,
|
||||
DownstreamAlreadyCommitted,
|
||||
RequestNotRunning,
|
||||
@@ -105,6 +106,10 @@ pub(crate) fn admin_store_error(resource: &'static str, error: StoreError) -> Ad
|
||||
kind: ConflictKind::StaleRevision,
|
||||
..
|
||||
} => AdminStoreErrorKind::StaleRevision,
|
||||
StoreError::Conflict {
|
||||
kind: ConflictKind::DuplicateName,
|
||||
..
|
||||
} => AdminStoreErrorKind::DuplicateName,
|
||||
StoreError::Conflict { .. } => AdminStoreErrorKind::Conflict,
|
||||
StoreError::InvalidData { .. } => AdminStoreErrorKind::Invalid,
|
||||
StoreError::Unavailable { .. } => AdminStoreErrorKind::Unavailable,
|
||||
|
||||
@@ -21,8 +21,122 @@ use gateway_store::postgres::{
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
#[tokio::test]
|
||||
async fn migrated_keys_persist_exactly_and_keep_short_keys_masked() {
|
||||
use gateway_admin::model::{MutationActor, MutationContext, client_keys::NewClientKey};
|
||||
use gateway_core::policy::RateLimits;
|
||||
use gateway_store::postgres::{PgRuntimeSnapshotRepository, RuntimeSnapshotRepository};
|
||||
|
||||
let Some(database) = TestDatabase::create("migrated_keys").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "migration-test".to_owned(),
|
||||
};
|
||||
let long_key: String = (0..160)
|
||||
.map(|_| uuid::Uuid::new_v4().simple().to_string())
|
||||
.collect();
|
||||
let values = [
|
||||
"q".to_owned(),
|
||||
"abc".to_owned(),
|
||||
"sk-legacy/key+value=:!@\\".to_owned(),
|
||||
long_key,
|
||||
];
|
||||
for (index, key) in values.iter().enumerate() {
|
||||
let id = ClientApiKeyId::new(format!("key_migrated_{index}")).unwrap();
|
||||
let (_, record) = store
|
||||
.create_client_key(
|
||||
NewClientKey {
|
||||
id: id.clone(),
|
||||
name: format!("Migrated {index}"),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits {
|
||||
max_concurrency: 3,
|
||||
requests_per_minute: 25,
|
||||
},
|
||||
budget: Default::default(),
|
||||
plaintext: key.clone(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(record.prefix.len(), 10.min(key.len() / 2));
|
||||
assert_ne!(record.prefix, *key);
|
||||
let revealed = store.reveal_client_key(&id).await.unwrap().unwrap();
|
||||
assert_eq!(revealed.expose_for_response(), key);
|
||||
}
|
||||
let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone())
|
||||
.load_runtime_snapshot()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(snapshot.client_api_keys.len(), values.len());
|
||||
for key in &values {
|
||||
let record = snapshot
|
||||
.client_api_keys
|
||||
.iter()
|
||||
.find(|item| item.plaintext_key.expose_for_auth() == key)
|
||||
.unwrap();
|
||||
assert_eq!(record.limits.max_concurrency, 3);
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_migrated_keys_conflict_atomically_without_extra_audits() {
|
||||
use gateway_admin::{
|
||||
model::{MutationActor, MutationContext, client_keys::NewClientKey},
|
||||
ports::store::AdminStoreErrorKind,
|
||||
};
|
||||
use gateway_core::policy::RateLimits;
|
||||
|
||||
let Some(database) = TestDatabase::create("duplicate_keys").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "duplicate-test".to_owned(),
|
||||
};
|
||||
let key = "legacy-key-case-sensitive!";
|
||||
let command = |id: &str| NewClientKey {
|
||||
id: ClientApiKeyId::new(id).unwrap(),
|
||||
name: id.to_owned(),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits::unlimited(),
|
||||
budget: Default::default(),
|
||||
plaintext: key.to_owned(),
|
||||
};
|
||||
let (first, second) = tokio::join!(
|
||||
store.create_client_key(command("key_first"), &context),
|
||||
store.create_client_key(command("key_second"), &context)
|
||||
);
|
||||
assert_ne!(first.is_ok(), second.is_ok());
|
||||
let error = first.err().or_else(|| second.err()).unwrap();
|
||||
assert_eq!(error.kind(), AdminStoreErrorKind::Conflict);
|
||||
assert!(!format!("{error:?}").contains(key));
|
||||
let counts: (i64, i64) = sqlx::query_as(
|
||||
"select (select count(*) from client_api_keys), (select count(*) from admin_audit_events)",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts, (1, 1));
|
||||
let mut different_case = command("key_case_sensitive");
|
||||
different_case.plaintext = key.to_uppercase();
|
||||
store
|
||||
.create_client_key(different_case, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_key_requires_the_frozen_plaintext_format() {
|
||||
fn generated_client_key_format_remains_valid() {
|
||||
let key = NewClientApiKey {
|
||||
budget: Default::default(),
|
||||
id: "key-1".to_owned(),
|
||||
@@ -36,6 +150,142 @@ fn client_key_requires_the_frozen_plaintext_format() {
|
||||
assert!(key.validate().is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_key_names_are_checked_atomically_on_create_and_rename() {
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationActor, MutationContext,
|
||||
client_keys::{NewClientKey, UpdateClientKey},
|
||||
},
|
||||
ports::store::AdminStoreErrorKind,
|
||||
};
|
||||
use gateway_core::policy::RateLimits;
|
||||
|
||||
let Some(database) = TestDatabase::create("duplicate_key_names").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "duplicate-name-test".to_owned(),
|
||||
};
|
||||
let create = |id: &str, name: &str| NewClientKey {
|
||||
id: ClientApiKeyId::new(id).unwrap(),
|
||||
name: name.to_owned(),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits::unlimited(),
|
||||
budget: Default::default(),
|
||||
plaintext: format!("synthetic-credential-{id}"),
|
||||
};
|
||||
let update = |id: ClientApiKeyId, name: &str| UpdateClientKey {
|
||||
id,
|
||||
name: name.to_owned(),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits::unlimited(),
|
||||
daily_limit_usd: None,
|
||||
weekly_limit_usd: None,
|
||||
};
|
||||
let (first, second) = tokio::join!(
|
||||
store.create_client_key(create("key_name_first", "Production"), &context),
|
||||
store.create_client_key(create("key_name_second", " production "), &context),
|
||||
);
|
||||
let (winner, conflict) = match (first, second) {
|
||||
(Ok((_, record)), Err(error)) | (Err(error), Ok((_, record))) => (record, error),
|
||||
results => panic!("expected one committed create: {results:?}"),
|
||||
};
|
||||
assert_eq!(conflict.kind(), AdminStoreErrorKind::DuplicateName);
|
||||
assert_eq!(winner.name.trim(), winner.name);
|
||||
let counts: (i64, i64) = sqlx::query_as(
|
||||
"select (select count(*) from client_api_keys), (select count(*) from admin_audit_events)",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts, (1, 1));
|
||||
|
||||
let (revision, _) = store
|
||||
.update_client_key(update(winner.id.clone(), "PRODUCTION"), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let (_, other) = store
|
||||
.create_client_key(create("key_name_other", "Staging"), &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let error = store
|
||||
.update_client_key(update(other.id.clone(), " production "), &context)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind(), AdminStoreErrorKind::DuplicateName);
|
||||
let unchanged = store.reveal_client_key(&other.id).await.unwrap().unwrap();
|
||||
assert_eq!(unchanged.record.name, "Staging");
|
||||
let persisted_revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(persisted_revision as u64, revision.get() + 1);
|
||||
|
||||
let (first, second) = tokio::join!(
|
||||
store.update_client_key(update(winner.id, "Shared"), &context),
|
||||
store.update_client_key(update(other.id, "SHARED"), &context),
|
||||
);
|
||||
assert_ne!(first.is_ok(), second.is_ok());
|
||||
assert_eq!(
|
||||
first.err().or_else(|| second.err()).unwrap().kind(),
|
||||
AdminStoreErrorKind::DuplicateName
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_key_search_matches_names_and_labels_but_never_credential_values() {
|
||||
let Some(database) = TestDatabase::create("key_name_search").await else {
|
||||
return;
|
||||
};
|
||||
let generated = format!("sk_{}", "K".repeat(43));
|
||||
for (id, name, key) in [
|
||||
("key_generated", "Production", generated.as_str()),
|
||||
("key_custom", "生产_%专用", "legacy-token+/123"),
|
||||
] {
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (id, name, label, key, created_at, updated_at)
|
||||
values ($1, $2, 'backend label', $3, now(), now())",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(name)
|
||||
.bind(key)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let repository = PgClientApiKeyRepository::new(database.pool.clone());
|
||||
for (search, expected) in [
|
||||
("PROD", 1),
|
||||
("生产_%", 1),
|
||||
("生产X", 0),
|
||||
("backend", 2),
|
||||
(generated.as_str(), 0),
|
||||
(&generated[..10], 0),
|
||||
("legacy-token+/123", 0),
|
||||
("legacy-", 0),
|
||||
] {
|
||||
let page = repository
|
||||
.list_client_api_keys(ClientApiKeyListQuery {
|
||||
cursor: None,
|
||||
page_size: 10,
|
||||
search: Some(search.to_owned()),
|
||||
sort: ClientApiKeySort::default(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.total, expected, "search: {search}");
|
||||
assert_eq!(page.items.len() as u64, expected);
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_key_list_uses_safe_keyset_search_and_filtered_total() {
|
||||
let Some(database) = TestDatabase::create("client_key_page").await else {
|
||||
|
||||
@@ -1,18 +1,34 @@
|
||||
use chrono::{Duration, Utc};
|
||||
use std::time::{Duration as StdDuration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use futures::TryStreamExt as _;
|
||||
use gateway_admin::{
|
||||
model::{PageSize, observability as admin_observability},
|
||||
ports::store::ObservabilityStore as _,
|
||||
};
|
||||
use gateway_core::diagnostics::TraceContext;
|
||||
use gateway_core::engine::{
|
||||
ExecutionOutcome, ExecutionStore, ModelRequestFailureObservation,
|
||||
ModelRequestFinalization as CoreModelRequestFinalization, ModelRequestId,
|
||||
ModelRequestTimings as CoreModelRequestTimings,
|
||||
ModelRequestTimings as CoreModelRequestTimings, NewModelRequest as CoreNewModelRequest,
|
||||
};
|
||||
use gateway_core::error::{
|
||||
GatewayError, GatewayErrorKind, ProviderConnectionObservation, StoreErrorKind,
|
||||
};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind, ProviderConnectionObservation};
|
||||
use gateway_core::metering::{CalculatedCost, CostEstimate, Usage};
|
||||
use gateway_core::operation::OperationKind;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{AccountRoutingSnapshot, ConfigRevision, PublicModelId};
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_store::postgres::{
|
||||
ModelRequestAttemptStart, ModelRequestRepository, NewModelRequest, ObservabilityRepository,
|
||||
PgExecutionStore,
|
||||
AttemptMetrics, ModelRequestAttemptStart, ModelRequestRepository, NewModelRequest,
|
||||
ObservabilityPageSize, ObservabilityRange, ObservabilityRepository, OpsErrorFilter,
|
||||
OpsErrorQuery, PgExecutionStore, UsageRecordFilter, UsageRecordQuery,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::PgPool;
|
||||
|
||||
use super::TestDatabase;
|
||||
use super::{TestDatabase, admin_observability_store, observability_repository};
|
||||
|
||||
#[test]
|
||||
fn postgres_execution_adapter_implements_core_port() {
|
||||
@@ -1150,3 +1166,596 @@ async fn diagnostic_trace_is_finalized_atomically_and_available_for_failed_reque
|
||||
assert!(detail.related_requests.is_empty());
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
// 只构造已接纳的入口事实;未选择账号、未出站,也没有上游用量。
|
||||
pub(super) fn accepted_request(id: &str) -> CoreNewModelRequest {
|
||||
let started_at = SystemTime::from(
|
||||
DateTime::from_timestamp_micros(Utc::now().timestamp_micros())
|
||||
.expect("PostgreSQL timestamp precision"),
|
||||
);
|
||||
CoreNewModelRequest {
|
||||
id: ModelRequestId::new(id).expect("request id"),
|
||||
client_api_key_id: None,
|
||||
client_api_key_ref: ClientApiKeyId::new("key_zero_attempt").expect("client key ref"),
|
||||
config_revision: ConfigRevision::new(1).expect("revision"),
|
||||
routing: AccountRoutingSnapshot::all(),
|
||||
protocol: "openai".to_owned(),
|
||||
operation: OperationKind::Generate,
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model: Some(PublicModelId::new("coding").expect("public model")),
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
reasoning_preset: None,
|
||||
request_kind: None,
|
||||
subagent_kind: None,
|
||||
compact: false,
|
||||
continuation: Default::default(),
|
||||
image_generation_requested: false,
|
||||
admission_decision_ms: Some(2),
|
||||
started_at,
|
||||
deadline_at: started_at + StdDuration::from_secs(30),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn early_failure(request: &CoreNewModelRequest) -> CoreModelRequestFinalization {
|
||||
let trace = TraceContext::new(request.id.as_str());
|
||||
trace.record("request.started", json!({"operation": "generate"}));
|
||||
// trace 的 index 是预备阶段关联,不证明已经拿到账号、建流或实际发送。
|
||||
let preparation = trace.attempt(1);
|
||||
preparation.record(
|
||||
"attempt.started",
|
||||
json!({"provider": "openai", "model": "coding", "pinnedAccount": null}),
|
||||
);
|
||||
preparation.record(
|
||||
"attempt.failed",
|
||||
json!({"kind": "no_eligible_account", "sendState": "NotSent"}),
|
||||
);
|
||||
trace.record(
|
||||
"request.finished",
|
||||
json!({"outcome": "failed", "errorKind": "no_available_provider", "attemptCount": 0}),
|
||||
);
|
||||
CoreModelRequestFinalization {
|
||||
request_id: request.id.clone(),
|
||||
outcome: ExecutionOutcome::Failed,
|
||||
send_state: UpstreamSendState::NotSent,
|
||||
attempt_count: 0,
|
||||
downstream_committed_at: None,
|
||||
client_status_code: Some(503),
|
||||
client_response_id: None,
|
||||
upstream_status_code: None,
|
||||
upstream_request_id: None,
|
||||
upstream_response_id: None,
|
||||
upstream_transport: None,
|
||||
http_version: None,
|
||||
websocket_pool: None,
|
||||
service_tier: None,
|
||||
provider_metadata_json: None,
|
||||
diagnostic_trace_json: trace.snapshot().map(|value| value.to_string()),
|
||||
error: Some(GatewayError::new(
|
||||
GatewayErrorKind::NoAvailableProvider,
|
||||
"no available provider",
|
||||
)),
|
||||
provider_error_code: None,
|
||||
raw_upstream_error: None,
|
||||
failure_observation: Default::default(),
|
||||
retry_after_ms: None,
|
||||
usage: Usage::new(),
|
||||
image_generation_succeeded: None,
|
||||
cost: CostEstimate::unavailable(),
|
||||
timings: CoreModelRequestTimings {
|
||||
latency_ms: Some(1_000),
|
||||
..Default::default()
|
||||
},
|
||||
completed_at: request.started_at + StdDuration::from_secs(1),
|
||||
}
|
||||
}
|
||||
|
||||
async fn stored_row(pool: &PgPool, request_id: &str) -> Value {
|
||||
sqlx::query_scalar("select to_jsonb(mr) from model_requests mr where id = $1")
|
||||
.bind(request_id)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.expect("load complete synthetic request row")
|
||||
}
|
||||
|
||||
fn zero_attempt_range(request: &CoreNewModelRequest) -> ObservabilityRange {
|
||||
ObservabilityRange::new(
|
||||
DateTime::from(request.started_at - StdDuration::from_secs(60)),
|
||||
DateTime::from(request.deadline_at + StdDuration::from_secs(60)),
|
||||
)
|
||||
.expect("observation range")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_failure_is_queryable_without_fabricating_upstream_facts() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_queries").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
let request = accepted_request("req_zero_attempt_queries");
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("create accepted request without attempt");
|
||||
let finalization = early_failure(&request);
|
||||
let trace: Value = serde_json::from_str(
|
||||
finalization
|
||||
.diagnostic_trace_json
|
||||
.as_deref()
|
||||
.expect("trace snapshot"),
|
||||
)
|
||||
.expect("trace JSON");
|
||||
ExecutionStore::finalize_model_request(&store, finalization)
|
||||
.await
|
||||
.expect("finalize zero-attempt failure with trace");
|
||||
|
||||
let repository = observability_repository(&database.pool);
|
||||
let detail = repository
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("load failed request detail");
|
||||
assert_eq!(detail.request.outcome, "failed");
|
||||
assert_eq!(detail.request.attempt_count, 0);
|
||||
assert_eq!(detail.request.upstream_send_state, "not_sent");
|
||||
assert_eq!(detail.request.client_api_key_ref, "key_zero_attempt");
|
||||
assert_eq!(detail.request.requested_model_id.as_deref(), Some("coding"));
|
||||
assert_eq!(detail.request.admission_decision_ms, Some(2));
|
||||
assert_eq!(detail.request.latency_ms, Some(1_000));
|
||||
assert_eq!(detail.request.client_status_code, Some(503));
|
||||
assert_eq!(
|
||||
detail.request.error_kind.as_deref(),
|
||||
Some("no_available_provider")
|
||||
);
|
||||
assert_eq!(
|
||||
detail.request.error_message.as_deref(),
|
||||
Some("no available provider")
|
||||
);
|
||||
assert!(detail.attempts.is_empty());
|
||||
assert!(detail.related_requests.is_empty());
|
||||
assert_eq!(detail.trace, Some(trace.clone()));
|
||||
let row = stored_row(&database.pool, request.id.as_str()).await;
|
||||
for field in [
|
||||
"provider_kind",
|
||||
"provider_account_id",
|
||||
"provider_account_ref",
|
||||
"provider_account_name_snapshot",
|
||||
"provider_account_email_snapshot",
|
||||
"provider_account_authentication_kind_snapshot",
|
||||
"upstream_model_id",
|
||||
"upstream_transport",
|
||||
"http_version",
|
||||
"upstream_status_code",
|
||||
"upstream_request_id",
|
||||
"upstream_response_id",
|
||||
"client_response_id",
|
||||
"downstream_committed_at",
|
||||
"service_tier",
|
||||
"provider_observation_json",
|
||||
"raw_upstream_error",
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"total_tokens",
|
||||
"cost_amount",
|
||||
"cost_currency",
|
||||
"account_selection_wait_ms",
|
||||
"capacity_used_slots",
|
||||
"capacity_total_slots",
|
||||
] {
|
||||
assert_eq!(
|
||||
row.get(field),
|
||||
Some(&Value::Null),
|
||||
"{field} must remain absent"
|
||||
);
|
||||
}
|
||||
assert_eq!(row["cost_source"], "unavailable");
|
||||
let events = trace["events"].as_array().expect("events");
|
||||
assert!(events.iter().all(|event| event["exchangeId"].is_null()));
|
||||
let prepared_attempts: Vec<_> = events
|
||||
.iter()
|
||||
.filter(|event| event["attemptIndex"] == 1)
|
||||
.map(|event| event["stage"].as_str().expect("stage"))
|
||||
.collect();
|
||||
assert_eq!(prepared_attempts, ["attempt.started", "attempt.failed"]);
|
||||
|
||||
let query = OpsErrorQuery {
|
||||
range: zero_attempt_range(&request),
|
||||
filter: OpsErrorFilter {
|
||||
request_id: Some(request.id.as_str().to_owned()),
|
||||
..Default::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
};
|
||||
let errors = repository
|
||||
.list_ops_errors(query.clone())
|
||||
.await
|
||||
.expect("ops errors");
|
||||
assert_eq!(errors.total, 1);
|
||||
assert_eq!(errors.items.len(), 1);
|
||||
let error = &errors.items[0];
|
||||
assert_eq!(error.source, "model_request");
|
||||
assert_eq!(error.request_id.as_deref(), Some(request.id.as_str()));
|
||||
assert_eq!(error.attempt_index, None);
|
||||
assert_eq!(error.failure_kind, "no_available_provider");
|
||||
assert_eq!(error.message, "no available provider");
|
||||
assert_eq!(error.upstream_send_state.as_deref(), Some("not_sent"));
|
||||
assert_eq!(error.provider_kind, None);
|
||||
assert_eq!(error.provider_account_ref, None);
|
||||
assert_eq!(error.upstream_transport, None);
|
||||
assert_eq!(error.upstream_status_code, None);
|
||||
assert_eq!(error.upstream_request_id, None);
|
||||
assert_eq!(error.raw_upstream_error, None);
|
||||
|
||||
// 路由入口不等于已选择 Provider;没有平台事实时不能被平台/attempt 筛选命中。
|
||||
for filter in [
|
||||
OpsErrorFilter {
|
||||
provider_kind: Some("openai".to_owned()),
|
||||
..query.filter.clone()
|
||||
},
|
||||
OpsErrorFilter {
|
||||
attempt_index: Some(1),
|
||||
..query.filter.clone()
|
||||
},
|
||||
] {
|
||||
let page = repository
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
filter,
|
||||
..query.clone()
|
||||
})
|
||||
.await
|
||||
.expect("filter absent upstream facts");
|
||||
assert_eq!(page.total, 0);
|
||||
assert!(page.items.is_empty());
|
||||
}
|
||||
let admin = admin_observability_store(&database.pool);
|
||||
let admin_detail = admin
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("admin detail");
|
||||
assert_eq!(admin_detail.trace, Some(trace));
|
||||
assert_eq!(admin_detail.request.attempt_count, 0);
|
||||
assert!(admin_detail.attempts.is_empty());
|
||||
let admin_errors = admin
|
||||
.list_ops_errors(admin_observability::OpsErrorQuery {
|
||||
range: admin_observability::TimeRange::new(query.range.start, query.range.end)
|
||||
.expect("admin time range"),
|
||||
filter: admin_observability::OpsErrorFilter {
|
||||
request_id: Some(request.id.as_str().to_owned()),
|
||||
..Default::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("admin ops errors without provider facts");
|
||||
assert_eq!(admin_errors.total, 1);
|
||||
assert_eq!(admin_errors.items.len(), 1);
|
||||
assert_eq!(admin_errors.items[0].attempt_index, None);
|
||||
assert_eq!(admin_errors.items[0].provider_kind, None);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_finalization_rejects_sent_or_missing_attempt_facts_atomically() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_constraints").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
let request = accepted_request("req_zero_attempt_constraints");
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("create request");
|
||||
let before = stored_row(&database.pool, request.id.as_str()).await;
|
||||
|
||||
for (count, send_state) in [
|
||||
(0, UpstreamSendState::Sent),
|
||||
(0, UpstreamSendState::Ambiguous),
|
||||
(1, UpstreamSendState::NotSent),
|
||||
] {
|
||||
let mut finalization = early_failure(&request);
|
||||
finalization.attempt_count = count;
|
||||
finalization.send_state = send_state;
|
||||
assert!(
|
||||
ExecutionStore::finalize_model_request(&store, finalization)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
stored_row(&database.pool, request.id.as_str()).await,
|
||||
before
|
||||
);
|
||||
}
|
||||
for invalid_trace in [
|
||||
"not JSON".to_owned(),
|
||||
"[]".to_owned(),
|
||||
json!({"data": "x".repeat(65_536)}).to_string(),
|
||||
] {
|
||||
let mut finalization = early_failure(&request);
|
||||
finalization.diagnostic_trace_json = Some(invalid_trace);
|
||||
assert_eq!(
|
||||
ExecutionStore::finalize_model_request(&store, finalization)
|
||||
.await
|
||||
.expect_err("invalid trace")
|
||||
.kind(),
|
||||
StoreErrorKind::InvalidData,
|
||||
);
|
||||
assert_eq!(
|
||||
stored_row(&database.pool, request.id.as_str()).await,
|
||||
before
|
||||
);
|
||||
}
|
||||
ExecutionStore::finalize_model_request(&store, early_failure(&request))
|
||||
.await
|
||||
.expect("valid zero attempt");
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_duplicate_create_and_finalization_never_overwrite_terminal_facts() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_idempotency").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
let request = accepted_request("req_zero_attempt_idempotency");
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("create request");
|
||||
let (first, second) = tokio::join!(
|
||||
ExecutionStore::finalize_model_request(&store, early_failure(&request)),
|
||||
ExecutionStore::finalize_model_request(&store, early_failure(&request)),
|
||||
);
|
||||
assert_eq!(usize::from(first.is_ok()) + usize::from(second.is_ok()), 1);
|
||||
assert_eq!(
|
||||
first
|
||||
.err()
|
||||
.or_else(|| second.err())
|
||||
.expect("CAS loser")
|
||||
.kind(),
|
||||
StoreErrorKind::InvalidState
|
||||
);
|
||||
let terminal = stored_row(&database.pool, request.id.as_str()).await;
|
||||
assert!(store.create_model_request(request.clone()).await.is_err());
|
||||
let mut late = early_failure(&request);
|
||||
late.outcome = ExecutionOutcome::Cancelled;
|
||||
late.error = Some(GatewayError::new(
|
||||
GatewayErrorKind::Cancelled,
|
||||
"request was cancelled",
|
||||
));
|
||||
late.diagnostic_trace_json = None;
|
||||
late.completed_at += StdDuration::from_secs(1);
|
||||
assert_eq!(
|
||||
ExecutionStore::finalize_model_request(&store, late)
|
||||
.await
|
||||
.expect_err("already finalized")
|
||||
.kind(),
|
||||
StoreErrorKind::InvalidState
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("recover")
|
||||
.requests,
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
stored_row(&database.pool, request.id.as_str()).await,
|
||||
terminal
|
||||
);
|
||||
let count: i64 = sqlx::query_scalar("select count(*) from model_requests")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("one request row");
|
||||
assert_eq!(count, 1);
|
||||
let ops: i64 = sqlx::query_scalar("select count(*) from ops_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("no extra error ledger");
|
||||
assert_eq!(ops, 0);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_recovery_respects_deadline_and_does_not_invent_a_trace() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_recovery").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
let request = accepted_request("req_zero_attempt_recovery");
|
||||
let mut active = request.clone();
|
||||
active.id = ModelRequestId::new("req_zero_attempt_still_running").expect("request id");
|
||||
active.deadline_at += StdDuration::from_secs(30);
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("create interrupted request");
|
||||
store
|
||||
.create_model_request(active.clone())
|
||||
.await
|
||||
.expect("create active request");
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at - StdDuration::from_micros(1))
|
||||
.await
|
||||
.expect("before deadline")
|
||||
.requests,
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("at deadline")
|
||||
.requests,
|
||||
1
|
||||
);
|
||||
let recovered = stored_row(&database.pool, request.id.as_str()).await;
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("repeat recovery")
|
||||
.requests,
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
ExecutionStore::finalize_model_request(&store, early_failure(&request))
|
||||
.await
|
||||
.expect_err("late finalize after recovery")
|
||||
.kind(),
|
||||
StoreErrorKind::InvalidState
|
||||
);
|
||||
assert_eq!(
|
||||
stored_row(&database.pool, request.id.as_str()).await,
|
||||
recovered
|
||||
);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let detail = repository
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("recovered detail");
|
||||
assert_eq!(detail.request.outcome, "incomplete");
|
||||
assert_eq!(
|
||||
detail.request.error_kind.as_deref(),
|
||||
Some("process_interrupted")
|
||||
);
|
||||
assert_eq!(
|
||||
detail.request.completed_at,
|
||||
Some(DateTime::from(request.deadline_at))
|
||||
);
|
||||
assert_eq!(detail.request.attempt_count, 0);
|
||||
assert_eq!(detail.request.upstream_send_state, "not_sent");
|
||||
assert_eq!(detail.trace, None);
|
||||
assert!(detail.attempts.is_empty());
|
||||
let errors = repository
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
range: zero_attempt_range(&request),
|
||||
filter: OpsErrorFilter::default(),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("recovered ops error");
|
||||
assert_eq!(errors.total, 1);
|
||||
assert_eq!(errors.items[0].failure_kind, "process_interrupted");
|
||||
assert_eq!(errors.items[0].attempt_index, None);
|
||||
assert_eq!(
|
||||
repository
|
||||
.usage_record_detail(active.id.as_str())
|
||||
.await
|
||||
.expect("active detail")
|
||||
.request
|
||||
.outcome,
|
||||
"running"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_failure_does_not_enter_successful_usage_or_cost_aggregates() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_metrics").await else {
|
||||
return;
|
||||
};
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
let request = accepted_request("req_zero_attempt_metrics");
|
||||
seed_running_request(&database.pool, "req_zero_attempt_success_control")
|
||||
.await
|
||||
.expect("success control");
|
||||
let mut success = successful_core_finalization("req_zero_attempt_success_control");
|
||||
success.downstream_committed_at = Some(success.completed_at);
|
||||
success.usage = Usage {
|
||||
input_tokens: Some(100),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(120),
|
||||
..Usage::new()
|
||||
};
|
||||
success.timings.latency_ms = Some(200);
|
||||
success.cost = gateway_core::metering::CalculatedCost::from_usd_ticks(12_500_000_000)
|
||||
.expect("calculated cost")
|
||||
.into_estimate();
|
||||
ExecutionStore::finalize_model_request(&store, success)
|
||||
.await
|
||||
.expect("finalize success control");
|
||||
let repository = observability_repository(&database.pool);
|
||||
let before = repository
|
||||
.usage_summary(zero_attempt_range(&request), UsageRecordFilter::default())
|
||||
.await
|
||||
.expect("baseline usage");
|
||||
assert_eq!(before.requests.total_tokens, 120);
|
||||
assert!(!before.attempts.costs.is_empty());
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("create zero attempt");
|
||||
ExecutionStore::finalize_model_request(&store, early_failure(&request))
|
||||
.await
|
||||
.expect("finalize failure");
|
||||
let after = repository
|
||||
.usage_summary(zero_attempt_range(&request), UsageRecordFilter::default())
|
||||
.await
|
||||
.expect("usage after failure");
|
||||
assert_eq!(
|
||||
after.requests.request_count,
|
||||
before.requests.request_count + 1
|
||||
);
|
||||
assert_eq!(
|
||||
after.requests.failure_count,
|
||||
before.requests.failure_count + 1
|
||||
);
|
||||
assert_eq!(after.requests.success_count, before.requests.success_count);
|
||||
assert_eq!(after.requests.total_tokens, before.requests.total_tokens);
|
||||
assert_eq!(after.requests.latency_sum, before.requests.latency_sum);
|
||||
assert_eq!(after.requests.latency_count, before.requests.latency_count);
|
||||
assert_eq!(after.attempts, before.attempts);
|
||||
let successes = repository
|
||||
.usage_summary(
|
||||
zero_attempt_range(&request),
|
||||
UsageRecordFilter {
|
||||
outcome: Some("succeeded".to_owned()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("successful aggregation unchanged");
|
||||
assert_eq!(successes, before);
|
||||
let filter = UsageRecordFilter {
|
||||
request_id: Some(request.id.as_str().to_owned()),
|
||||
..Default::default()
|
||||
};
|
||||
let failed = repository
|
||||
.usage_summary(zero_attempt_range(&request), filter.clone())
|
||||
.await
|
||||
.expect("failed request metrics");
|
||||
assert_eq!(failed.requests.request_count, 1);
|
||||
assert_eq!(failed.requests.failure_count, 1);
|
||||
assert_eq!(failed.requests.total_tokens, 0);
|
||||
assert_eq!(failed.requests.latency_count, 0);
|
||||
assert_eq!(failed.attempts, AttemptMetrics::default());
|
||||
let page = repository
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range: zero_attempt_range(&request),
|
||||
filter: filter.clone(),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("successful usage list excludes failure");
|
||||
assert_eq!(page.total, 0);
|
||||
assert!(page.items.is_empty());
|
||||
let billing: Vec<_> = repository
|
||||
.usage_calculated_billing_facts(zero_attempt_range(&request), filter)
|
||||
.try_collect()
|
||||
.await
|
||||
.expect("billing facts");
|
||||
assert!(billing.is_empty());
|
||||
let charges: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("observations do not settle budgets");
|
||||
assert_eq!(charges, 0);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -15,7 +15,15 @@ use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::routing::{ProviderKind, UpstreamModelId};
|
||||
use gateway_core::task::DaemonTask as _;
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_store::postgres::BufferedExecutionStore;
|
||||
use gateway_store::postgres::{
|
||||
BufferedExecutionStore, ObservabilityRepository as _, PgExecutionStore,
|
||||
};
|
||||
|
||||
use super::{
|
||||
TestDatabase,
|
||||
execution::{accepted_request, early_failure},
|
||||
observability_repository,
|
||||
};
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingStore {
|
||||
@@ -266,3 +274,293 @@ async fn observation_byte_budget_drops_payload_without_waiting_for_the_database(
|
||||
assert_eq!(store.stats().queued_items, 0);
|
||||
assert_eq!(store.stats().dropped_total, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_create_and_finalize_drain_in_order_with_trace() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_queue_order").await else {
|
||||
return;
|
||||
};
|
||||
let (store, writer) = BufferedExecutionStore::with_capacity(
|
||||
Arc::new(PgExecutionStore::new(database.pool.clone())),
|
||||
NonZeroUsize::new(2).expect("capacity"),
|
||||
);
|
||||
let request = accepted_request("req_zero_attempt_queue_order");
|
||||
let finalization = early_failure(&request);
|
||||
let expected_trace: serde_json::Value = serde_json::from_str(
|
||||
finalization
|
||||
.diagnostic_trace_json
|
||||
.as_deref()
|
||||
.expect("trace"),
|
||||
)
|
||||
.expect("trace JSON");
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("enqueue create");
|
||||
store
|
||||
.finalize_model_request(finalization)
|
||||
.await
|
||||
.expect("enqueue finalization");
|
||||
let before: i64 = sqlx::query_scalar("select count(*) from model_requests")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("not yet persisted");
|
||||
assert_eq!(before, 0);
|
||||
assert_eq!(store.stats().queued_items, 2);
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.await
|
||||
.expect("shutdown drains accepted commands");
|
||||
let stats = store.stats();
|
||||
assert_eq!(stats.persisted_total, 2);
|
||||
assert_eq!(stats.write_failure_total, 0);
|
||||
assert_eq!(stats.dropped_total, 0);
|
||||
assert_eq!(stats.queued_items, 0);
|
||||
assert_eq!(stats.queued_bytes, 0);
|
||||
let detail = observability_repository(&database.pool)
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("queued failure detail");
|
||||
assert_eq!(detail.request.outcome, "failed");
|
||||
assert_eq!(detail.request.attempt_count, 0);
|
||||
assert_eq!(detail.request.upstream_send_state, "not_sent");
|
||||
assert_eq!(detail.trace, Some(expected_trace));
|
||||
assert!(detail.attempts.is_empty());
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_dropped_finalize_is_fail_open_and_recoverable_at_deadline() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_queue_full").await else {
|
||||
return;
|
||||
};
|
||||
let (store, writer) = BufferedExecutionStore::with_capacity(
|
||||
Arc::new(PgExecutionStore::new(database.pool.clone())),
|
||||
NonZeroUsize::new(1).expect("capacity"),
|
||||
);
|
||||
let request = accepted_request("req_zero_attempt_queue_full");
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("enqueue create");
|
||||
store
|
||||
.finalize_model_request(early_failure(&request))
|
||||
.await
|
||||
.expect("full queue is fail open");
|
||||
})
|
||||
.await
|
||||
.expect("no database wait");
|
||||
assert_eq!(store.stats().enqueued_total, 1);
|
||||
assert_eq!(store.stats().dropped_total, 1);
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer.run(cancellation).await.expect("drain create only");
|
||||
let repository = observability_repository(&database.pool);
|
||||
let pending = repository
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("pending record");
|
||||
assert_eq!(pending.request.outcome, "running");
|
||||
assert_eq!(pending.trace, None);
|
||||
assert!(pending.attempts.is_empty());
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("recover without queue")
|
||||
.requests,
|
||||
1
|
||||
);
|
||||
let recovered = repository
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("recovered record");
|
||||
assert_eq!(recovered.request.outcome, "incomplete");
|
||||
assert_eq!(
|
||||
recovered.request.error_kind.as_deref(),
|
||||
Some("process_interrupted")
|
||||
);
|
||||
assert_eq!(recovered.request.attempt_count, 0);
|
||||
assert_eq!(recovered.trace, None);
|
||||
assert!(recovered.attempts.is_empty());
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("repeat recovery")
|
||||
.requests,
|
||||
0
|
||||
);
|
||||
|
||||
// 已关闭队列同样不把观测失败返回客户端,也不能覆盖已经恢复的终态。
|
||||
store
|
||||
.finalize_model_request(early_failure(&request))
|
||||
.await
|
||||
.expect("closed queue is fail open");
|
||||
assert_eq!(store.stats().dropped_total, 2);
|
||||
assert_eq!(store.stats().persisted_total, 1);
|
||||
assert_eq!(store.stats().write_failure_total, 0);
|
||||
assert_eq!(
|
||||
repository
|
||||
.usage_record_detail(request.id.as_str())
|
||||
.await
|
||||
.expect("unchanged recovery"),
|
||||
recovered
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_dropped_create_does_not_make_finalize_an_upsert() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_create_dropped").await else {
|
||||
return;
|
||||
};
|
||||
let (store, writer) = BufferedExecutionStore::with_limits(
|
||||
Arc::new(PgExecutionStore::new(database.pool.clone())),
|
||||
NonZeroUsize::new(2).expect("capacity"),
|
||||
NonZeroUsize::new(8_192).expect("byte budget"),
|
||||
);
|
||||
let mut request = accepted_request("req_zero_attempt_create_dropped");
|
||||
// 合成大 header 只触发 Create 字节预算;Finalize 不携带 user_agent,仍可入队。
|
||||
request.user_agent = Some("x".repeat(16_384));
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("oversized create is fail open");
|
||||
assert_eq!(store.stats().dropped_total, 1);
|
||||
store
|
||||
.finalize_model_request(early_failure(&request))
|
||||
.await
|
||||
.expect("enqueue finalization without create");
|
||||
assert_eq!(store.stats().enqueued_total, 1);
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.await
|
||||
.expect("drain orphan finalization");
|
||||
assert_eq!(store.stats().persisted_total, 0);
|
||||
assert_eq!(store.stats().write_failure_total, 1);
|
||||
assert_eq!(store.stats().queued_bytes, 0);
|
||||
let count: i64 = sqlx::query_scalar("select count(*) from model_requests")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("no fabricated row");
|
||||
assert_eq!(count, 0);
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("nothing to recover")
|
||||
.requests,
|
||||
0
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_attempt_postgres_write_failures_are_not_retried_and_do_not_stop_the_queue() {
|
||||
let Some(database) = TestDatabase::create("zero_attempt_write_failure").await else {
|
||||
return;
|
||||
};
|
||||
// 只在随机测试 schema 注入语句失败;序列不随语句回滚,用于核对实际写入次数。
|
||||
sqlx::raw_sql(
|
||||
"create sequence zero_attempt_rejected_writes;
|
||||
create function reject_zero_attempt_observation() returns trigger language plpgsql as $$
|
||||
begin
|
||||
if (TG_OP = 'INSERT' and NEW.id = 'req_zero_attempt_create_rejected')
|
||||
or (TG_OP = 'UPDATE' and NEW.id = 'req_zero_attempt_finalize_rejected'
|
||||
and NEW.outcome = 'failed') then
|
||||
perform nextval('zero_attempt_rejected_writes');
|
||||
raise exception 'synthetic observation write failure';
|
||||
end if;
|
||||
return NEW;
|
||||
end $$;
|
||||
create trigger reject_zero_attempt_observation
|
||||
before insert or update on model_requests
|
||||
for each row execute function reject_zero_attempt_observation();",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("install isolated write failure");
|
||||
let (store, writer) = BufferedExecutionStore::with_capacity(
|
||||
Arc::new(PgExecutionStore::new(database.pool.clone())),
|
||||
NonZeroUsize::new(6).expect("capacity"),
|
||||
);
|
||||
let request = accepted_request("req_zero_attempt_create_rejected");
|
||||
for id in [
|
||||
"req_zero_attempt_create_rejected",
|
||||
"req_zero_attempt_finalize_rejected",
|
||||
"req_zero_attempt_after_rejection",
|
||||
] {
|
||||
let mut request = request.clone();
|
||||
request.id = ModelRequestId::new(id).expect("request id");
|
||||
store
|
||||
.create_model_request(request.clone())
|
||||
.await
|
||||
.expect("enqueue create");
|
||||
store
|
||||
.finalize_model_request(early_failure(&request))
|
||||
.await
|
||||
.expect("enqueue finalization");
|
||||
}
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.await
|
||||
.expect("continue draining after PG errors");
|
||||
let stats = store.stats();
|
||||
assert_eq!(stats.enqueued_total, 6);
|
||||
assert_eq!(stats.persisted_total, 3);
|
||||
assert_eq!(stats.write_failure_total, 3);
|
||||
assert_eq!(stats.dropped_total, 0);
|
||||
assert_eq!(stats.queued_items, 0);
|
||||
assert_eq!(stats.queued_bytes, 0);
|
||||
let rejected_writes: i64 =
|
||||
sqlx::query_scalar("select last_value from zero_attempt_rejected_writes")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("count failed PostgreSQL writes without rollback");
|
||||
assert_eq!(rejected_writes, 2);
|
||||
let rows: Vec<(String, String, Option<serde_json::Value>)> =
|
||||
sqlx::query_as("select id, outcome, diagnostic_trace_json from model_requests order by id")
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.expect("load successfully persisted rows");
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].0, "req_zero_attempt_after_rejection");
|
||||
assert_eq!(rows[0].1, "failed");
|
||||
assert!(rows[0].2.is_some());
|
||||
assert_eq!(
|
||||
rows[1],
|
||||
(
|
||||
"req_zero_attempt_finalize_rejected".to_owned(),
|
||||
"running".to_owned(),
|
||||
None
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.recover_expired(request.deadline_at)
|
||||
.await
|
||||
.expect("recover only inserted unfinished row")
|
||||
.requests,
|
||||
1
|
||||
);
|
||||
let recovered = observability_repository(&database.pool)
|
||||
.usage_record_detail("req_zero_attempt_finalize_rejected")
|
||||
.await
|
||||
.expect("recovered detail");
|
||||
assert_eq!(
|
||||
recovered.request.error_kind.as_deref(),
|
||||
Some("process_interrupted")
|
||||
);
|
||||
assert_eq!(recovered.request.attempt_count, 0);
|
||||
assert!(recovered.attempts.is_empty());
|
||||
assert_eq!(recovered.trace, None);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -188,7 +188,150 @@ async fn usage_search_should_match_literal_prefix_instead_of_substring() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_search_should_match_client_api_key_prefix() {
|
||||
async fn usage_search_should_match_account_email_and_name_prefixes() {
|
||||
let Some(database) = TestDatabase::create("usage_account_prefix").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
seed_observability_facts(&database.pool, now)
|
||||
.await
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
for search in ["account@example.invalid", "account@", "primary", "pri"] {
|
||||
assert_usage_search_ids(&database.pool, range, search, &["req_observe_success"]).await;
|
||||
}
|
||||
for search in ["count@example.invalid", "rimary", "missing@example.invalid"] {
|
||||
assert_usage_search_ids(&database.pool, range, search, &[]).await;
|
||||
}
|
||||
|
||||
sqlx::query("update model_requests set provider_account_email_snapshot = null")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("clear account email snapshots");
|
||||
assert_usage_search_ids(&database.pool, range, "primary", &["req_observe_success"]).await;
|
||||
assert_usage_search_ids(&database.pool, range, "account@", &[]).await;
|
||||
|
||||
sqlx::query("update model_requests set provider_account_name_snapshot = null")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("clear account name snapshots");
|
||||
assert_usage_search_ids(
|
||||
&database.pool,
|
||||
range,
|
||||
"acct_observe",
|
||||
&["req_observe_success"],
|
||||
)
|
||||
.await;
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_search_should_preserve_account_snapshots_after_account_changes() {
|
||||
let Some(database) = TestDatabase::create("usage_account_history").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
seed_observability_facts(&database.pool, now)
|
||||
.await
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
// 搜索与列表都使用请求发生时的快照,不能随当前账号资料变化或删除而改变。
|
||||
for mutation in [
|
||||
"update provider_accounts set name = 'renamed', email = 'renamed@example.invalid'
|
||||
where id = 'acct_observe'",
|
||||
"delete from provider_accounts where id = 'acct_observe'",
|
||||
] {
|
||||
sqlx::query(mutation)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("change current account");
|
||||
for search in ["account@example.invalid", "primary"] {
|
||||
assert_usage_search_ids(&database.pool, range, search, &["req_observe_success"]).await;
|
||||
}
|
||||
for search in ["renamed@example.invalid", "renamed"] {
|
||||
assert_usage_search_ids(&database.pool, range, search, &[]).await;
|
||||
}
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_search_should_treat_account_snapshot_wildcards_as_literals() {
|
||||
let Some(database) = TestDatabase::create("usage_account_literals").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
seed_observability_facts(&database.pool, now)
|
||||
.await
|
||||
.expect("seed observability facts");
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set provider_account_email_snapshot = $1, provider_account_name_snapshot = $2",
|
||||
)
|
||||
.bind("account_team%tag@example.invalid")
|
||||
.bind(r"primary\ops_100%")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("set account snapshots with literal wildcard characters");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
for search in [
|
||||
"account_team%tag@example.invalid",
|
||||
"account_",
|
||||
"account_team%",
|
||||
r"primary\",
|
||||
r"primary\ops_",
|
||||
r"primary\ops_100%",
|
||||
] {
|
||||
assert_usage_search_ids(&database.pool, range, search, &["req_observe_success"]).await;
|
||||
}
|
||||
for search in ["account%", "account_team_", "primary_", r"primary\ops%"] {
|
||||
assert_usage_search_ids(&database.pool, range, search, &[]).await;
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
async fn assert_usage_search_ids(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
search: &str,
|
||||
expected_ids: &[&str],
|
||||
) {
|
||||
let page = observability_repository(pool)
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter {
|
||||
search: Some(search.to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("search usage by account snapshot");
|
||||
|
||||
assert_eq!(
|
||||
page.total,
|
||||
u64::try_from(expected_ids.len()).expect("expected total"),
|
||||
"search: {search}"
|
||||
);
|
||||
assert_eq!(
|
||||
page.items
|
||||
.iter()
|
||||
.map(|item| item.id.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
expected_ids,
|
||||
"search: {search}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_and_error_search_match_key_names_instead_of_credentials() {
|
||||
let Some(database) = TestDatabase::create("usage_client_api_key_search").await else {
|
||||
return;
|
||||
};
|
||||
@@ -199,7 +342,7 @@ async fn usage_search_should_match_client_api_key_prefix() {
|
||||
let plaintext_key = format!("sk_{}", "K".repeat(43));
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (id, name, key, enabled, created_at, updated_at)
|
||||
values ('key_observe', 'usage search', $1, true, $2, $2)",
|
||||
values ('key_observe', 'Production_%专用', $1, true, $2, $2)",
|
||||
)
|
||||
.bind(&plaintext_key)
|
||||
.bind(now)
|
||||
@@ -209,22 +352,53 @@ async fn usage_search_should_match_client_api_key_prefix() {
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
for search in [plaintext_key.as_str(), &plaintext_key[..10]] {
|
||||
let page = observability_repository(&database.pool)
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter {
|
||||
search: Some(search.to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
let repository = observability_repository(&database.pool);
|
||||
for key in [plaintext_key.as_str(), "legacy-key+/123"] {
|
||||
sqlx::query("update client_api_keys set key = $1 where id = 'key_observe'")
|
||||
.bind(key)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("search usage by client API key");
|
||||
.unwrap();
|
||||
for (search, expected) in [
|
||||
("production", 1),
|
||||
("PRODUCTION_%", 1),
|
||||
("ProductionX", 0),
|
||||
(key, 0),
|
||||
(&key[..10], 0),
|
||||
] {
|
||||
let page = repository
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter {
|
||||
search: Some(search.to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("search usage by client API key name");
|
||||
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].id, "req_observe_success");
|
||||
assert_eq!(page.total, expected, "search: {search}");
|
||||
assert_eq!(page.items.len() as u64, expected);
|
||||
if expected == 1 {
|
||||
assert_eq!(page.items[0].id, "req_observe_success");
|
||||
}
|
||||
let errors = repository
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
range,
|
||||
filter: OpsErrorFilter {
|
||||
search: Some(search.to_owned()),
|
||||
..Default::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).unwrap(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(errors.total, expected * 2, "error search: {search}");
|
||||
assert_eq!(errors.items.len() as u64, expected * 2);
|
||||
}
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user