mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
14e7967e1c | ||
|
|
a4efb77223 | ||
|
|
5041d6ba26 | ||
|
|
322b602154 | ||
|
|
de15722ee1 | ||
|
|
380f04bace | ||
|
|
194cc0e0c9 | ||
|
|
d3acb42d37 | ||
|
|
bfca754714 | ||
|
|
881e78975e | ||
|
|
eb3baef3aa | ||
|
|
5205c07d60 | ||
|
|
60eb08cb96 | ||
|
|
fd4a49e64f | ||
|
|
be2ee7f3b8 | ||
|
|
757abdb4d1 | ||
|
|
cb5e0ba4c0 | ||
|
|
d2b418f8cc | ||
|
|
2390783900 | ||
|
|
c75f0e7fd9 | ||
|
|
d5e9360ed7 | ||
|
|
3fcb8e5f2d | ||
|
|
8cdaf8b2e6 | ||
|
|
68f05a4481 | ||
|
|
bb4d178b5d | ||
|
|
73487fd98b | ||
|
|
66fd09eba6 | ||
|
|
bd152dec00 | ||
|
|
7090dca502 | ||
|
|
029f3124bf | ||
|
|
840707e5a1 | ||
|
|
704b652f86 | ||
|
|
2b98839c29 |
@@ -0,0 +1,35 @@
|
||||
---
|
||||
name: dev-guide
|
||||
description: 当前仓库开发指南。修改或排查本仓库的 Rust 网关、Vue 管理端、账号与出站代理、接口协议、部署、CI、发布或文档时使用,定位权威规范和验证入口。Use for development in the repository containing this skill; not for generic Rust/Vue questions in unrelated repositories.
|
||||
---
|
||||
|
||||
# 开发指南
|
||||
|
||||
先读根目录的 [AGENTS.md](../../../AGENTS.md) 和 [CONTRIBUTING.md](../../../CONTRIBUTING.md)。
|
||||
本技能只组织工作入口,架构、协议和视觉规则从下列权威文档读取,不维护第二份副本。
|
||||
|
||||
## 按任务读取
|
||||
|
||||
| 任务 | 文档入口 |
|
||||
| --- | --- |
|
||||
| 用户能力、快速开始 | [README](../../../README.md) |
|
||||
| 模块职责、调用链、状态归属 | [系统架构](../../../docs/architecture.md) |
|
||||
| HTTP 路由、请求与响应、导入导出 | [API 文档](../../../docs/api.md) |
|
||||
| 管理端组件、主题 Token、视觉验证 | [管理端主题](../../../docs/theme.md) |
|
||||
| 部署、客户端配置、备份与恢复 | [部署说明](../../../deploy/README.md) |
|
||||
| 数据迁移、冻结清单与本地测试库 | [迁移说明](../../../backend/migrations/README.md) |
|
||||
| 验证命令与提交规范 | [贡献与审查](../../../CONTRIBUTING.md) |
|
||||
| 创建或审查 PR | [github-pr](../github-pr/SKILL.md) |
|
||||
| 整理或提交 Issue | [github-issue](../github-issue/SKILL.md) |
|
||||
| 发布说明、版本规划与发版 | [release](../release/SKILL.md) |
|
||||
|
||||
## 执行顺序
|
||||
|
||||
1. 从仓库根目录确认工作区、当前任务范围和相关运行实例;保留已有改动。
|
||||
2. 按 `AGENTS.md` 的 CodeGraph 约定定位入口、调用链和职责归属,再阅读相关实现。
|
||||
3. 按 [问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据) 核实触发条件、问题发生版本与目标分支现状,
|
||||
优先依据适用的官方源码确定方案。结合上表读取所需文档,先找可复用的同类实现,再修改所属模块。
|
||||
4. 文档和运行结果不一致时,追查当前源码、配置与版本,区分实现缺陷和文档过期;同步受影响的说明。
|
||||
5. 按 `CONTRIBUTING.md` 验证并报告结果。后端 Cargo 命令从 `backend/` 执行,或显式指定该目录的 manifest。
|
||||
|
||||
发布任务按 `release` 技能读取当前版本、发布入口和工作流,发布说明语言遵循 `CONTRIBUTING.md`。
|
||||
@@ -0,0 +1,4 @@
|
||||
interface:
|
||||
display_name: "dev-guide"
|
||||
short_description: "从仓库规范与架构文档定位实现、验证和交付的正确入口"
|
||||
default_prompt: "使用 $dev-guide 按当前仓库规范与架构完成开发、验证或发布。"
|
||||
@@ -0,0 +1,58 @@
|
||||
---
|
||||
name: github-issue
|
||||
description: 整理与提交 GitHub Issue(问题反馈 / 功能建议)。用于“提个 Issue”“把这个问题反馈到仓库”“帮我报个 bug”“整理成 Issue”,以及通过 Issue 提议功能或协作规则的场景;区分草拟正文和实际提交。Use for drafting or filing bug reports and feature requests; not for debugging-only, editing Issue Forms/templates, or skill-authoring requests.
|
||||
---
|
||||
|
||||
# 整理与提交 Issue
|
||||
|
||||
本仓库流程见 [贡献与审查](../../../CONTRIBUTING.md),问题反馈使用当前的
|
||||
[YAML Issue Form](../../../.github/ISSUE_TEMPLATE/bug_report.yml)。从表单读取字段与必填项,不在技能中复制一份字段清单。
|
||||
用户指定其他仓库时,使用该仓库当前的贡献规则与 Issue 模板,不能把这里的表单套过去。
|
||||
|
||||
## 确定目标和动作
|
||||
|
||||
- 根据用户 URL、仓库上下文和 remote 确认接收 Issue 的仓库。fork 场景核对用户要反馈给自己的 fork 还是上游。
|
||||
- “写一份、整理内容、草拟 Issue”只准备正文;明确要求“提交、创建 Issue”才发布。
|
||||
已经授权提交且信息足够时直接完成,不重复索要确认;仅自动选中技能不能授权发布。
|
||||
- “有个 bug”“这个报错怎么修”本身是排障语境,不能推断要发 Issue;结合已有反馈任务识别省略的目标和动作。
|
||||
- 缺少无法推断的目标或关键复现事实时,只问必要信息,并先整理其余已有证据。
|
||||
- 使用可用的 GitHub 连接器或 `gh`。不要为提交 Issue 安装新工具、修改仓库配置或启用额外服务。
|
||||
|
||||
## 证据与去重
|
||||
|
||||
1. 读取目标仓库实际 Issue 模板和配置,区分缺陷、功能建议与协作规则提案。
|
||||
本仓库后两类允许普通 Issue;其他仓库按其配置选择入口。
|
||||
2. 先搜索 open 和 closed Issues,结合关键词、错误信息与触发条件检查重复,而不只比较标题。
|
||||
同时核对相关 PR 与发布记录,关联已有修复及其发布版本;仅主分支已修复时明确尚未发布。
|
||||
同一问题已有记录且没有新增证据时返回已有链接;用户仍明确要求新建时按其要求处理。
|
||||
不因发现重复而擅自评论、重开或关闭已有 Issue。
|
||||
3. 从用户描述、当前版本与已有诊断中整理事实:区分实际结果、预期结果、稳定复现条件和未经证实的推测。
|
||||
不把本地开发版本当作发生问题的部署版本;无法复现时如实说明已知条件。
|
||||
引用官方行为或提出实现方案时,按 [问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据) 注明来源、版本与缺口;
|
||||
不把找到官方源码或证明根因作为提交症状反馈的前提。
|
||||
4. 只收集定位所需信息。正文、日志片段和附件去除密钥、Cookie、令牌、代理认证、私人请求内容等敏感信息;
|
||||
不直接附完整配置或请求转储。记录与复现相关的代理、TUN、证书条件,不擅自修改这些环境来收集材料。
|
||||
|
||||
## 编写正文
|
||||
|
||||
- 标题写具体症状或所需行为,不预先断言尚未证明的根因。语言与目标仓库及用户要求一致。
|
||||
- YAML Issue Form 按字段顺序使用 `attributes.label` 作为 Markdown 小标题,填写相应答案。
|
||||
`gh issue create --body-file` 直接提交正文,不会替你完成网页表单校验;手动核对必填字段、选项和必选确认项。
|
||||
不把原始 YAML、占位文字或尚未确认的勾选项当成答案提交。
|
||||
- 无法获知的信息明确写未提供或不适用,并说明理由;如果该事实关系到能否理解或复现问题,发布前补齐。
|
||||
功能或规则提案说明当前场景、缺少的能力、建议结果和必要取舍,不强行填写缺陷复现步骤。
|
||||
- 附件只使用用户授权发布且实际可访问的地址;本机 `/tmp/...` 路径不能作为 GitHub 附件。
|
||||
没有可用附件时保留操作描述并说明缺口,不声称截图已上传。
|
||||
- 使用模板规定或用户要求的 labels、assignees;先确认它们在目标仓库存在,不自行猜测。
|
||||
|
||||
## 提交与收尾
|
||||
|
||||
把核对后的 UTF-8 正文写入临时文件,或直接传给连接器的结构化字段。CLI 变量先从已确认的目标与正文赋值:
|
||||
|
||||
```bash
|
||||
gh issue create --repo "$issue_repo" --title "$issue_title" --body-file "$issue_body_file"
|
||||
```
|
||||
|
||||
发布前快速复核目标、必填内容、脱敏及重复记录。网络超时或返回不明确时先查询远端是否已经创建,避免重复提交。
|
||||
发布后读取新 Issue 确认仓库、标题和正文,返回链接及必要的信息缺口。
|
||||
只准备草稿或发布失败时提供可直接使用的标题与正文,明确尚未提交。
|
||||
@@ -0,0 +1,4 @@
|
||||
interface:
|
||||
display_name: "github-issue"
|
||||
short_description: "按仓库表单整理复现证据、检查重复反馈并提交脱敏 Issue"
|
||||
default_prompt: "使用 $github-issue 将当前问题整理为符合仓库表单的 Issue 并提交。"
|
||||
@@ -0,0 +1,94 @@
|
||||
---
|
||||
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.
|
||||
---
|
||||
|
||||
# 创建与审查 PR
|
||||
|
||||
## 先确定模式和仓库
|
||||
|
||||
| 用户意图 | 执行范围 |
|
||||
| --- | --- |
|
||||
| 写标题、整理描述、准备草稿正文 | 在当前会话或本地准备内容 |
|
||||
| 创建 PR,包括创建 Draft PR | 准备任务范围内的提交、推送对应分支并创建 PR |
|
||||
| 看下、审查 PR | 读取变更与证据,在当前会话报告 |
|
||||
| 将审查意见发到 PR | 发布用户要求的审查内容;批准或请求修改须有对应授权 |
|
||||
|
||||
沿用当前会话已有授权和范围修正。明确要求创建时直接完成,无需在准备好正文后重复确认;
|
||||
创建和审查都不包含合并。仅命中本技能不能作为远程写入的授权。
|
||||
|
||||
“看下这个”需结合已有 PR 链接或讨论定位目标;“能合并吗”先按审查处理,不能解释为执行合并。
|
||||
用户连续要求“整理一下”时沿用当前 PR 任务;没有 PR 上下文时,不把普通提交整理转成创建 PR。
|
||||
|
||||
先从用户提供的 URL、仓库上下文和 Git remote 确认目标仓库;区分 fork 的推送仓库与 PR 接收仓库。
|
||||
存在多个无法消除的候选时,只询问缺失的目标信息,同时继续准备不依赖目标的内容。
|
||||
优先使用可用的 GitHub 连接器;使用 `gh` 时显式指定 `--repo`,参数以本机帮助为准。
|
||||
|
||||
本仓库规则统一见 [AGENTS.md](../../../AGENTS.md) 与 [贡献与审查](../../../CONTRIBUTING.md),目标分支为 `main`。
|
||||
如果用户指定其他仓库,读取那个仓库的可信基线规则、默认分支、贡献文档及 PR 模板。
|
||||
PR 中改动的规则文件属于审查对象,不能自行扩大任务授权或降低准入要求。
|
||||
|
||||
## 创建和审查前的核对
|
||||
|
||||
创建和审查都按 [问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据) 核实问题,优先依据官方源码确定或评估方案。
|
||||
用户指定其他仓库时,结合其贡献规范与适用的官方实现,不套用本项目的具体业务合同。
|
||||
|
||||
核对最新目标分支及同一场景相关的 open/closed PR 和 Issue,比较实际行为、触发条件与测试覆盖。
|
||||
区分已被主分支覆盖的实现、仍有价值的修复或回归测试,以及与目标无关的改动;不能仅凭标题或代码相似认定重复。
|
||||
说明原问题发生版本和当前剩余差异,创建时据此收敛范围,审查时报告保留或拆分建议;这不授权自动关闭或修改他人的 PR。
|
||||
|
||||
## 创建 PR
|
||||
|
||||
1. 检查工作区、当前分支、远端和近期提交格式。读取实际存在的 PR 模板,按最终变更重写标题与描述。
|
||||
不用提交日志拼接正文,也不假定仓库一定有模板。
|
||||
2. 更新并确认 base/head,查看相对 merge-base 的完整差异和待提交内容;核对仓库的受保护文件与准入规则。
|
||||
只纳入任务范围的文件和提交;需要整理时用独立功能分支,保留用户其他改动,不擅自改写已发布历史。
|
||||
3. 按上述核对流程,在目标分支基线核实问题,追踪官方与本项目的相关调用链,先确定解决方向再实现。
|
||||
若改动已完成,也要据此复核问题与方案;需要修正时只修改已授权的范围。仅整理标题或正文时报告差异与缺口,
|
||||
不自行修改实现。原问题不成立时不将它包装成已修复的缺陷;证据不足时明确保留结论。
|
||||
4. 查找同一来源仓库、head 和 base 的现有 PR;已存在则复用其链接,不重复创建。
|
||||
更新现有标题或正文须在本次整理/更新请求范围内。
|
||||
5. 按贡献文档执行适用检查。记录实际验证的提交或工作区状态,区分通过、失败、跳过和未执行。
|
||||
前端改动提供实际界面证据或预览步骤;缺少证据时写明,不用构建成功替代视觉验收。
|
||||
6. 将任务范围内的提交推送到正确远端,核对远端 head。标题遵循项目约定,正文先说明问题成立的依据和最终行为,
|
||||
再给必要的设计取舍、验证与遗留事项。涉及官方行为时,写明官方源码位置与版本或提交、官方处理方式、
|
||||
本次实现与它的关系及必要的偏离原因;区分源码推断和实际复现,不用检查通过替代方案论证。按规模压缩,不强填无关章节。
|
||||
7. 明确要求草稿或仍有阻塞项时创建 Draft PR,并解释原因;具备审查条件且未要求草稿时创建普通 PR。
|
||||
GitHub Draft 状态与“只写本地草稿正文”不同,后者不创建远程对象。
|
||||
8. 创建后核对链接、base、head、草稿状态和最新检查结果。网络超时或结果不确定时先查询远端,再决定是否重试。
|
||||
|
||||
CLI 示例中的变量必须先根据上述检查赋值;`pr_body_file` 是精确写入 UTF-8 正文的临时文件:
|
||||
|
||||
```bash
|
||||
gh pr create --repo "$pr_repo" --base "$pr_base" --head "$pr_head" \
|
||||
--title "$pr_title" --body-file "$pr_body_file"
|
||||
```
|
||||
|
||||
按需添加 `--draft`。不要用 `gh pr create --dry-run` 做只读预览,它仍可能推送 Git 变更。
|
||||
多行正文通过文件或连接器的结构化字段传递,不把未转义正文拼进 shell 命令。
|
||||
创建完成返回链接、目标分支和验证摘要;不要称为已合并或已获维护者批准。
|
||||
|
||||
## 审查 PR
|
||||
|
||||
1. 获取 PR 的 base/head SHA、描述、文件差异、已有讨论及 CI 状态。记录实际审查的 head,使用 merge-base 差异。
|
||||
需要源码上下文时先遵守目标仓库的定位规则;需检出时使用隔离工作区,保留用户当前分支和运行环境。
|
||||
2. 先核实 PR 要解决的问题是否成立:提取触发场景、预期行为和声称的原因,在目标分支基线检查实际行为。
|
||||
按上述核对流程收集一手证据,区分原有问题、环境或配置差异与 PR 新引入的问题,说明主分支已覆盖的部分。
|
||||
3. 从官方方案评估解决方向:说明官方如何处理同一场景,检查 PR 是否遵循其协议语义和职责归属、
|
||||
是否复用已有能力或正确传递官方合同;核对偏离官方方案的理由与运行证据。
|
||||
4. 按贡献文档的审查维度追踪完整调用链,找到共享业务事实、同类实现及已有验证。
|
||||
对既有特殊分支查清服务的场景,再判断本次改动是否引入或加重问题。
|
||||
5. UI 改动同时对照主题、共用组件、数据展示逻辑和同类页面。浏览器可用时检查相关交互及状态;
|
||||
未运行的页面检查列为验证缺口,不从截图或源码推断所有状态都正确。
|
||||
6. 核对 CI 与本地证据是否覆盖当前 head;根据风险补充必要验证,不触发无关部署或真实账号业务请求。
|
||||
工具不可用、集成测试跳过或环境不足时具体说明影响范围。
|
||||
7. 按 [审查标准](../../../CONTRIBUTING.md#审查标准) 用中文先回答原问题是否成立、解决方向是否合理,
|
||||
再列出实现问题与验证缺口,结论附具体依据。
|
||||
实现问题区分必须修复项、可选建议和验证缺口;没有的类别可省略。
|
||||
每个问题给变更中的文件位置、触发条件或规范依据、影响和最小修复方向。
|
||||
复用建议指出具体已有实现,不把个人偏好写成缺陷;未发现问题时直接说明并交代覆盖范围。
|
||||
8. 结束前核对 head 是否变化。发生变化时检查新增差异并更新结论,或明确结论只适用于原 SHA。
|
||||
|
||||
审查本身不修改源码、不发表评论、不提交 GitHub 的批准或请求修改状态。
|
||||
用户明确要求发布时,先复核当前 head 和重复评论,再通过结构化字段或 `--body-file` 发布对应内容;
|
||||
需要绑定具体提交时使用支持 commit ID 的审查接口。不要在旧提交结论上批准新的 head。
|
||||
@@ -0,0 +1,4 @@
|
||||
interface:
|
||||
display_name: "github-pr"
|
||||
short_description: "创建与审查拉取请求,先核实问题,以官方源码为首要依据确定或评估实现方案"
|
||||
default_prompt: "使用 $github-pr 按当前请求创建或审查 PR,先核实问题是否成立,以官方源码为首要参考确定或评估解决方案,说明源码依据及必要的偏离原因,并按仓库规范实现与验证或报告审查结论。"
|
||||
@@ -0,0 +1,79 @@
|
||||
---
|
||||
name: release
|
||||
description: 准备或执行当前仓库发版,维护中文发布说明。用于“发版”“发个新版本”“准备 vX.Y.Z”“推荐版本号”“整理更新日志/发布说明”“更新 release/notes.md”“打版本 tag”等发行任务。Use for version planning, release notes, and publishing a repository release; not for deploying/restarting an instance, publishing application data, or authoring a release skill.
|
||||
---
|
||||
|
||||
# 发版
|
||||
|
||||
## 识别任务范围
|
||||
|
||||
| 用户意图 | 执行范围 |
|
||||
| --- | --- |
|
||||
| 推荐版本号 | 根据本次变更与已有版本给出建议及理由 |
|
||||
| 整理更新日志、写 note、准备发版 | 核对变更、编写说明、检查发布条件;未要求提交时保留本地改动 |
|
||||
| 正式发布、发布指定版本、打版本 tag | 在已授权范围内执行对应步骤;仅打 tag 的请求不自动扩展为完整发版 |
|
||||
| 重试失败的发版 | 核对上次执行状态,从未完成的步骤恢复 |
|
||||
|
||||
沿用当前会话明确的版本、范围与授权;用户已要求正式发版且条件满足时直接执行,不重复确认。
|
||||
“准备好了”“可以发版吗”属于准备或检查,不能解释成正式发布;创建本技能也不触发发版。
|
||||
发行版发布与运行实例升级分开判断,“上线”需结合上下文确定用户指哪一种。
|
||||
|
||||
## 读取当前流程
|
||||
|
||||
先读 [AGENTS.md](../../../AGENTS.md) 与 [贡献与审查](../../../CONTRIBUTING.md),再按任务检查:
|
||||
|
||||
| 事实 | 权威入口 |
|
||||
| --- | --- |
|
||||
| 当前版本 | [release/version.yaml](../../../release/version.yaml) |
|
||||
| 本次发布说明及语言约定 | [release/notes.md](../../../release/notes.md)、贡献文档的项目约定 |
|
||||
| 版本提交、tag、推送与工作流触发 | [release/publish](../../../release/publish) |
|
||||
| 构建、检查、GitHub Release 和镜像发布 | [发布工作流](../../../.github/workflows/release.yml) |
|
||||
| 平台与产物清单 | [release/platforms.yaml](../../../release/platforms.yaml) |
|
||||
|
||||
命令从仓库根目录执行;每次以当前脚本和工作流为准,不另写发布脚本或绕过已有质量检查。
|
||||
用户明确指定其他仓库时先读取其发布流程,不能直接套用这里的路径和命令。
|
||||
|
||||
## 准备版本与说明
|
||||
|
||||
1. 确认发布分支、upstream、目标仓库、工作区和完整待推送范围。本仓库沿用 `main`;
|
||||
未提交或未跟踪内容分别核对,不因脚本只检查已跟踪文件就认定所有内容都已纳入发布。
|
||||
2. 读取版本文件、已有 tag 和实际发行记录,确认上一发布基线与本次目标提交。
|
||||
根据基线到目标的提交和最终差异归纳变更,剔除回退、重复和未进入发布范围的内容。
|
||||
不能只用当前版本文件或最近几条提交判断“本次更新”。
|
||||
3. 使用用户指定或当前会话已确定的目标版本;需要建议时按兼容性、能力新增和修复范围说明理由。
|
||||
目标 tag 必须未被占用;遇到已有同名版本先判断是重试还是新发版,不自动递增版本掩盖冲突。
|
||||
4. 按贡献文档维护 `release/notes.md`。正文写中文,保留版本号和必要技术标识;
|
||||
按用户可感知的最终行为归并,不逐条翻译 commit,也不把排查过程或未验证承诺写入说明。
|
||||
可按新增功能、问题修复、体验优化、升级说明组织,省略没有内容的章节。
|
||||
5. 替换为本次版本说明,首行与目标 tag 一致且正文非空。升级事项应能追溯到实际迁移、配置或兼容变化;
|
||||
敏感信息脱敏,未合入功能不宣称已发布。已有准确说明仅需翻译时保持其事实和范围。
|
||||
6. 按贡献规范检查变更与现有 CI 证据,明确通过、失败、跳过和未执行。
|
||||
只修改说明时无需重建前后端;正式发版仍要跟踪发布工作流针对目标提交的检查与产物。
|
||||
|
||||
## 执行正式发版
|
||||
|
||||
调用入口前重新核对当前脚本的副作用、目标版本和待推送提交,保留无关工作区内容,不自动清理或暂存用户文件。
|
||||
当前脚本要求先提交 `release/notes.md`,再由脚本更新 `release/version.yaml` 并创建版本提交,
|
||||
不要提前把版本文件改成目标版本,否则脚本会判定该版本已设置。
|
||||
本任务需要的提交按贡献文档保留英文 Conventional Commits 与相应署名。
|
||||
|
||||
在准备完成且用户已授权正式发布后,将 `release_version` 设为已确定的版本,执行:
|
||||
|
||||
```bash
|
||||
release/publish "$release_version"
|
||||
```
|
||||
|
||||
该命令会创建版本提交与 annotated tag,原子推送分支和 tag,再从长期分支手动触发 `release.yml`。
|
||||
它没有只读预演模式;不能为了检查或演练而执行。发布工作流读取 tag 中的说明,修改 tag 之后的本地文件不会改变该次构建。
|
||||
|
||||
## 恢复与验收
|
||||
|
||||
- 脚本中断或网络超时时,先核对本地提交与 tag、远端分支与 tag、对应工作流 run 及 Release,确定已完成到哪一步。
|
||||
若本地已创建版本提交或 tag,不直接重跑入口,也不通过删 tag、回退版本或强推来消除报错。
|
||||
- 提交和 tag 已推送但工作流未触发时,核对目标后使用脚本提示的 dispatch 命令补触发;
|
||||
工作流已存在则跟踪或恢复对应 run,避免重复发布。发现 tag 指向不同提交时报告冲突,不移动已发布 tag。
|
||||
- 跟踪与本次 tag 对应的 run ID,核对工作流解析的目标提交。触发工作流成功只表示已开始,不能当作发布完成。
|
||||
- 发布完成后核对 GitHub Release 的版本、正式版/预发行状态和中文变更正文;对照平台清单与工作流确认归档、校验和及镜像产物。
|
||||
工作流附加的安装说明和自动提交列表按其当前配置处理,不用它们替代手写的中文变更说明。
|
||||
- 返回版本、tag 对应提交、Release/工作流链接及实际结果。失败、仍在构建或缺失产物时明确说明,不能称为发布成功;
|
||||
Release 和镜像已生成也不表示用户的运行实例已经升级。
|
||||
@@ -0,0 +1,4 @@
|
||||
interface:
|
||||
display_name: "release"
|
||||
short_description: "按仓库流程整理中文发布说明、核对版本并完成发行与验证"
|
||||
default_prompt: "使用 $release 准备下一个版本的中文发布说明,并核对发版条件。"
|
||||
@@ -0,0 +1,120 @@
|
||||
name: 问题反馈
|
||||
description: 提交可复现的问题,帮助定位和修复
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
请提供最小复现信息。请求、日志和截图发布前请移除密钥、令牌、Cookie、代理认证及私人内容。
|
||||
|
||||
- type: dropdown
|
||||
id: issue-type
|
||||
attributes:
|
||||
label: 问题类型
|
||||
options:
|
||||
- 接口请求或响应异常
|
||||
- 页面或交互异常
|
||||
- 配置或部署异常
|
||||
- 其他问题
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: description
|
||||
attributes:
|
||||
label: 问题描述
|
||||
description: 说明发生了什么、实际结果以及出现的频率。
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: input
|
||||
id: version
|
||||
attributes:
|
||||
label: 服务版本
|
||||
description: 填写具体版本号;从源码运行时填写提交号。
|
||||
placeholder: 版本号或提交号
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: dropdown
|
||||
id: deployment
|
||||
attributes:
|
||||
label: 部署方式
|
||||
options:
|
||||
- Docker Compose
|
||||
- 其他容器部署
|
||||
- 二进制运行
|
||||
- 源码运行
|
||||
- 其他方式(请在环境信息中说明)
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: environment
|
||||
attributes:
|
||||
label: 环境信息
|
||||
description: 填写操作系统与版本;页面问题同时填写浏览器与版本。
|
||||
placeholder: |
|
||||
操作系统与版本:
|
||||
浏览器与版本(页面问题):
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: account-client
|
||||
attributes:
|
||||
label: 账号与客户端
|
||||
description: 接口问题请填写平台、认证类型、模型和客户端信息,不要提供账号凭据。
|
||||
placeholder: |
|
||||
平台与认证类型:
|
||||
模型名称:
|
||||
客户端名称与版本:
|
||||
|
||||
- type: textarea
|
||||
id: reproduction
|
||||
attributes:
|
||||
label: 复现步骤
|
||||
description: 描述触发问题的操作;无法稳定复现时,说明已知触发条件。
|
||||
placeholder: |
|
||||
1. 使用什么配置或输入
|
||||
2. 执行什么操作
|
||||
3. 出现什么结果
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: expected-behavior
|
||||
attributes:
|
||||
label: 预期行为
|
||||
description: 说明原本期望得到的结果。
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: request-response
|
||||
attributes:
|
||||
label: 请求与响应
|
||||
description: |
|
||||
接口问题请提供脱敏后的最小请求、网关实际状态/响应片段和客户端终端显示,区分两者。
|
||||
优先附管理端默认诊断包并检查其缺失说明;原始错误与日志片段需另行审阅脱敏。
|
||||
先按部署文档的“请求错误排查”检查关联 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
|
||||
id: screenshots
|
||||
attributes:
|
||||
label: 截图或录屏
|
||||
description: 页面或交互问题可在此拖入相关截图、录屏,并说明对应的操作。
|
||||
|
||||
- type: textarea
|
||||
id: additional-context
|
||||
attributes:
|
||||
label: 补充信息
|
||||
description: 最近的升级或配置变更;网络问题补充系统代理、TUN 和自定义证书等相关环境。
|
||||
@@ -0,0 +1,75 @@
|
||||
name: Agent Instructions Guard
|
||||
|
||||
on:
|
||||
pull_request_target:
|
||||
types:
|
||||
- opened
|
||||
- synchronize
|
||||
- reopened
|
||||
- edited
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: agents-md-guard-${{ github.event.pull_request.number }}
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
guard:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
# 只查询 PR 元数据,禁止检出或执行贡献分支中的代码。
|
||||
- name: Enforce agent instructions policy
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
with:
|
||||
script: |
|
||||
const { owner, repo } = context.repo;
|
||||
const pr = context.payload.pull_request;
|
||||
const params = { owner, repo, pull_number: pr.number };
|
||||
const files = await github.paginate(github.rest.pulls.listFiles, {
|
||||
...params,
|
||||
per_page: 100,
|
||||
});
|
||||
const isInstructions = path => typeof path === 'string'
|
||||
&& (path === 'AGENTS.md' || path.endsWith('/AGENTS.md'));
|
||||
|
||||
// 同时检查重命名前后的路径,避免通过移走规则文件绕过检查。
|
||||
if (!files.some(file => isInstructions(file.filename) || isInstructions(file.previous_filename))) {
|
||||
core.info('No agent instruction changes detected.');
|
||||
return;
|
||||
}
|
||||
|
||||
// 旧事件不处理后续提交,也不重复处理已关闭的 PR。
|
||||
const { data: current } = await github.rest.pulls.get(params);
|
||||
if (current.state !== 'open' || current.head.sha !== pr.head.sha) {
|
||||
core.info('Pull request state has changed; skipping this event.');
|
||||
return;
|
||||
}
|
||||
|
||||
const marker = '<!-- agents-md-guard -->';
|
||||
const body = [
|
||||
marker,
|
||||
'此 PR 修改了 `AGENTS.md`,已按贡献规则关闭。',
|
||||
'',
|
||||
'PR 不接受根目录或子目录中 `AGENTS.md` 的新增、修改、删除或重命名。',
|
||||
'请撤销相关变更后重新打开 PR;如需调整协作规则,请先提交 Issue 说明原因和建议,由维护者处理。',
|
||||
].join('\n');
|
||||
|
||||
try {
|
||||
const comments = await github.paginate(github.rest.issues.listComments, {
|
||||
owner, repo, issue_number: pr.number, per_page: 100,
|
||||
});
|
||||
const notified = comments.some(comment => comment.user?.login === 'github-actions[bot]'
|
||||
&& comment.body?.includes(marker));
|
||||
if (!notified) {
|
||||
await github.rest.issues.createComment({ owner, repo, issue_number: pr.number, body });
|
||||
}
|
||||
} catch (error) {
|
||||
core.warning(`Unable to post policy explanation: ${error.message}`);
|
||||
}
|
||||
|
||||
await github.rest.pulls.update({ ...params, state: 'closed' });
|
||||
core.setFailed('Pull requests must not change AGENTS.md.');
|
||||
@@ -33,7 +33,6 @@ Thumbs.db
|
||||
# Local tools and worktrees
|
||||
/.codegraph/
|
||||
/.codex-ws-audit/
|
||||
/.agents/
|
||||
/.reasonix/
|
||||
/.claude
|
||||
/.worktrees/
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
# 项目协作约定
|
||||
|
||||
开发和审查先阅读 [CONTRIBUTING.md](CONTRIBUTING.md),使用其中的协作流程和审查标准。
|
||||
模块职责以 [系统架构](docs/architecture.md) 为准,界面约定以 [管理端主题](docs/theme.md) 为准;
|
||||
没有明确约定时保持与相邻实现一致,不把单次页面反馈扩展成全仓规则。
|
||||
|
||||
<!-- CODEGRAPH_START -->
|
||||
## CodeGraph
|
||||
|
||||
In repositories indexed by CodeGraph (a `.codegraph/` directory exists at the repo root), reach for it BEFORE grep/find or reading files when you need to understand or locate code:
|
||||
|
||||
- **MCP tools** (when available): `codegraph_explore` answers most code questions in one call — the relevant symbols' verbatim source plus the call paths between them. `codegraph_node` returns one symbol's source + callers, or reads a whole file with line numbers. If the tools are listed but deferred, load them by name via tool search.
|
||||
- **Shell** (always works): `codegraph explore "<symbol names or question>"` and `codegraph node <symbol-or-file>` print the same output.
|
||||
|
||||
If there is no `.codegraph/` directory, skip CodeGraph entirely — indexing is the user's decision.
|
||||
<!-- CODEGRAPH_END -->
|
||||
|
||||
## 工作方式
|
||||
|
||||
- 开发和审查按 [问题与方案依据](CONTRIBUTING.md#问题与方案依据) 先核实问题,以官方源码为首要参考确定或评估方案。
|
||||
- 先定位变更所属模块、同类实现和上下游调用关系,再选择最小合理改动。复用已有能力,也避免为尚不存在的需求增加抽象。
|
||||
- 代码注释使用中文,解释原因与边界;提交信息使用英文,沿用历史中的 Conventional Commits 格式。
|
||||
- 按变更范围执行验证,记录命令、结果和缺口。跳过的测试不算通过;界面与集成行为需要对应运行证据,构建通过不能替代实际验收。
|
||||
- 审查请求默认只读。修复、提交、推送和合并按用户当前授权执行;“本地验证通过”不等于用户已经审阅批准。
|
||||
- 不读取或输出无关凭据,不把真实密钥、代理认证、账号令牌或请求转储放入提交、截图及审查报告。
|
||||
|
||||
## Code Review Rules
|
||||
|
||||
### 审查范围
|
||||
|
||||
- 确认目标分支与当前变更范围,结合调用方、被调用方及已有测试审查;结论对应实际审查的提交或工作区状态。
|
||||
- 按 CONTRIBUTING.md 的维度检查设计、正确性、安全、复杂度、性能、可读性、验证、界面和文档,深度与变更风险相称。
|
||||
- 重点核对现有架构和数据合同:共享业务事实是否重复解释,异步与并发操作是否破坏状态,权限和敏感数据是否越界,协议适配是否影响兼容性。
|
||||
- 修改既有特殊分支时先查场景和历史;放宽检查、修改测试或新增例外需要业务依据,不能仅以让检查通过为理由。
|
||||
|
||||
### 代码与界面一致性
|
||||
|
||||
- 对重复逻辑、冗余状态和多余抽象给出具体位置及影响;建议复用时指出已有实现,不能仅凭行数或语法相似判定质量。
|
||||
- UI 变更对照项目基础组件、主题和同类页面,检查交互、状态、响应式与可访问性。局部问题是否应在共用组件修复,按职责判断,不一律禁止局部样式。
|
||||
- 核对截图、预览和实际交互证据;未检查的状态明确列为验证缺口,不推断视觉验收已经通过。
|
||||
|
||||
### 审查输出
|
||||
|
||||
- 用中文先回答原问题是否成立、解决方向是否合理,依据与判断标准见 [贡献与审查](CONTRIBUTING.md#审查标准),再报告实现问题和验证缺口。
|
||||
- 实现问题聚焦本次变更新引入或加重的、可定位的问题,包含文件位置、触发场景或规范依据、影响和最小修复方向。
|
||||
- 区分必须修复的问题、可选建议和验证缺口。个人偏好标为建议,不作为必须修复项;没有问题时不要凑数。
|
||||
- 严重程度按实际影响判断,不为绕过工具的优先级过滤而升级普通规范问题。说明审查和验证的实际覆盖范围,AI 未报告问题不等于维护者批准。
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
# 贡献与审查
|
||||
|
||||
贡献应解决明确的问题,并保持或改善项目的可维护性。提交者、自动检查和审查者使用同一套标准,让重复问题在提交和初审阶段得到处理。
|
||||
|
||||
本文定义协作流程与审查标准;模块职责以 [系统架构](docs/architecture.md) 为准,视觉规则以 [管理端主题](docs/theme.md) 为准。没有明确约定时遵循相邻实现;具体页面的布局决定不自动成为全项目限制。
|
||||
|
||||
## 提交 Issue
|
||||
|
||||
问题反馈使用仓库的 [Issue 表单](.github/ISSUE_TEMPLATE/bug_report.yml),填写问题类型、描述、版本与环境、复现步骤和预期结果。接口问题补充平台与认证类型、模型、客户端,以及脱敏后的请求、响应和请求 ID;页面问题附相关截图或录屏。
|
||||
|
||||
日志和配置只提供定位所需片段,不上传真实凭据或完整请求转储。功能建议与协作规则调整可直接提交 Issue,说明场景、原因和建议方案。
|
||||
|
||||
## 问题与方案依据
|
||||
|
||||
开发、创建 PR 和审查都先核实待解决的问题、触发条件与预期行为,再确定或评估实现。区分问题发生版本、目标分支现状和本次变更后的行为,不把需求描述、PR 描述或第三方项目的做法直接当成事实。目标分支已修复的问题,应说明当前仍需补齐的差异,不能据此否定旧版本中的问题。
|
||||
|
||||
涉及上游协议、客户端行为或兼容性时,以对应版本的官方源码为首要资料。优先读取用户提供或本机已有的官方源码,核对仓库来源、版本或提交;本地资料不足时查官方仓库的对应版本源码,再结合官方测试、文档、发布说明和运行证据。第三方资料只作线索或补充。引用具体文件、符号或行号,对外引用优先使用固定提交的源码链接;源码不可得或版本不匹配时说明证据缺口,不把最新实现当成当前部署行为。无上游依赖的内部变更,以本项目代码、合同和运行证据为依据。
|
||||
|
||||
官方已有同一场景的处理方式时,优先沿用其协议语义、职责归属和解决路径,检查是否只需复用已有能力或正确传递官方合同。新增覆盖配置、补偿逻辑或采用第三方兼容方案,应说明官方路径为何不足;偏离官方方案需有本项目的实际约束和验证依据,不机械照搬源码,也不默认接受已有提案。
|
||||
|
||||
反馈 Issue 时区分观察到的症状、根因假设和方案建议,关联已有 Issue、PR 及包含修复的发布版本。能说明问题的反馈可以先提交,不要求报告者先找到官方源码或证明根因;引用官方行为或提出实现方案时,按上述原则注明依据与缺口。
|
||||
|
||||
## PR 流程
|
||||
|
||||
1. **准备变更。** 按上述依据核实问题与方案,找到所属模块、相关调用链和最接近的现有实现,在功能分支完成修改与自查。涉及架构、公开合同或较大交互变化时,先说明方案与取舍。
|
||||
2. **发起 PR。** 目标分支使用 `main`。说明问题成立的依据、触发场景、改动后的行为和实际验证结果;涉及官方行为时注明源码位置、版本和方案关系。界面改动提供截图、录屏或可复现的预览方式。一个 PR 围绕一个明确目标,必要重构说明关联原因。
|
||||
3. **自动检查与初审。** CI 执行适用的格式、静态分析、测试和构建检查;启用 AI 审查后,由它按本规范辅助检查语义、复用和一致性。提交者先处理明确问题,维护者重点审阅方案、风险与用户体验。
|
||||
4. **修改与合并。** 更新代码后重新执行受影响的检查,确认审查意见已处理,验证证据对应当前变更。维护者审阅批准后再合并;需要升级说明的变更同步相关文档。
|
||||
|
||||
AI 开发和审查入口是根目录的 [AGENTS.md](AGENTS.md)。在仓库上下文中要求“审查 PR #编号”时,应按下列标准执行,无需再次列举要求。
|
||||
|
||||
### 项目技能
|
||||
|
||||
仓库技能位于 `.agents/skills/`,随代码共享。技能提供执行流程,协作与审查标准仍以本文和 `AGENTS.md` 为准。
|
||||
|
||||
| 技能 | 用途 | 调用示例 |
|
||||
| --- | --- | --- |
|
||||
| [dev-guide](.agents/skills/dev-guide/SKILL.md) | 项目开发与文档导航 | `$dev-guide 修复当前问题` |
|
||||
| [github-pr](.agents/skills/github-pr/SKILL.md) | 创建或审查 PR | `$github-pr 创建当前分支的 PR` / `$github-pr 审查 PR #编号` |
|
||||
| [github-issue](.agents/skills/github-issue/SKILL.md) | 整理或提交 Issue | `$github-issue 提交当前问题` |
|
||||
| [release](.agents/skills/release/SKILL.md) | 准备发布说明或正式发版 | `$release 整理本次发布说明` / `$release 发布 vX.Y.Z` |
|
||||
|
||||
审查请求默认在当前会话报告;明确要求创建 PR、提交 Issue 或正式发版时才执行对应远程操作。技能安装不等于启用 GitHub 自动审查。
|
||||
|
||||
### 规则文件保护
|
||||
|
||||
PR 不接受根目录或子目录中 `AGENTS.md` 的新增、修改、删除或重命名。调整这些文件时先提交 Issue 说明原因和建议,由维护者处理;[规则保护工作流](.github/workflows/agents-md-guard.yml) 会提示并关闭包含此类变更的 PR。撤销相关变更后可以重新打开 PR,继续正常检查。
|
||||
|
||||
源码变更遵守现有架构边界与迁移冻结规则,由对应 CI 检查执行;不额外设立整块业务目录的 PR 禁改清单。构建、测试和安全扫描复用现有工作流。
|
||||
|
||||
## 审查标准
|
||||
|
||||
按改动涉及的范围选择检查深度,无需为每个 PR 填写所有维度的报告。
|
||||
|
||||
结论先回答原问题是否成立(成立、部分成立、不成立或证据不足),说明触发条件和依据;再判断解决方向是否合理,交代适用的官方方案、本次实现与它的关系及必要调整,最后报告实现问题与验证缺口。原问题是否成立与 PR 实现是否正确分别判断。源码推断、实际复现和检查结果分别说明,CI 通过不能代替问题与方案的论证。
|
||||
|
||||
| 维度 | 检查重点 |
|
||||
| --- | --- |
|
||||
| 设计与架构 | 职责归属、依赖方向、公开接口是否合理;是否融入已有体系,有无绕过边界的旁路 |
|
||||
| 正确性与兼容性 | 正常、异常和边界输入;并发、取消、重试与副作用;是否改变既有 API、协议或数据合同 |
|
||||
| 安全与数据 | 权限、输入校验、敏感信息、状态隔离和持久化一致性;日志与诊断是否符合脱敏边界 |
|
||||
| 简洁性与复用 | 重复业务规则、冗余状态、无效代码和过早抽象;能否用更直接的实现解决当前问题 |
|
||||
| 性能与资源 | 实际数据规模下的查询、请求、内存与连接成本;资源释放、背压和异常时的生命周期 |
|
||||
| 可读性与风格 | 命名是否准确,控制流是否易懂,格式是否一致;注释解释约束与原因而非复述代码 |
|
||||
| 测试与验证 | 是否覆盖真实触发条件和结果,能否发现回归;证据是否对应本次变更,有无跳过或缺失项 |
|
||||
| 界面与交互 | 共用组件和主题、信息层级、状态反馈、响应式、键盘操作及可访问性;是否有实际页面证据 |
|
||||
| 文档与升级 | 行为、配置、API 和部署说明是否同步;数据迁移和兼容影响是否交代清楚 |
|
||||
|
||||
审查意见区分三类:
|
||||
|
||||
- **必须修复:** 有具体触发条件的缺陷、违反现有合同或明确项目规范的问题。说明位置、依据、影响和最小修复方向。
|
||||
- **建议:** 不影响合同或规范的替代实现、命名偏好和局部优化。个人偏好不能单独阻止合并,也不应强求与本次目标无关的重构。
|
||||
- **验证缺口:** 尚未检查的行为、状态或运行环境。说明缺少什么证据,由维护者评估是否需要补齐,不能把未验证当作已通过或已发现缺陷。
|
||||
|
||||
复杂度不按行数判断,复用不只看语法相似。建议合并重复逻辑时应指出共享的业务含义;建议保留特殊分支时应能说明它服务的场景。例外应有依据,不能仅为绕过检查而放宽规则。
|
||||
|
||||
## 项目约定
|
||||
|
||||
- Rust 模块、可见性、依赖方向和测试目录按架构文档执行,使用现有 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 聚焦目标;无关格式化、重构和依赖升级单独处理。行为或合同变更同步到架构文档指定的文档位置,发布说明按维护者的发布安排更新。
|
||||
|
||||
## 界面验证
|
||||
|
||||
开发前查找同类页面及 `components/base` 中的基础组件,优先使用已有主题 Token、交互和数据展示逻辑。需要扩展时在对应职责层修改;局部样式的适用范围遵循主题文档,不一律禁止,也不靠页面覆盖掩盖共用组件的问题。
|
||||
|
||||
列表列宽、序号、按钮、分页和弹窗尺寸由实际任务与数据规模决定,以同类页面为参照。新增或调整交互时说明用户需要完成什么操作,不把某个页面的取舍推广成全局禁令。
|
||||
|
||||
提交者在浏览器中验证受影响路径,提供前后截图、录屏或预览步骤,并标明实际覆盖的状态。按变更选择浅色/深色、窄窗口、空/加载/错误状态、大量数据及键盘操作;共用主题或基础组件变更按主题文档扩大验证范围。模拟数据与预览材料不包含真实凭据。
|
||||
|
||||
涉及浏览器能力或访问环境差异时,按项目支持的部署方式核对相关 HTTP/HTTPS、localhost/实际部署地址与浏览器版本,记录实际覆盖的条件;本地预览正常不能替代部署环境中的交互验收。
|
||||
|
||||
## 验证
|
||||
|
||||
选择与变更相关的检查,纯文档修改无需重建前后端。行为修复补充能覆盖触发条件和结果的测试,不增加仅复述实现的断言。
|
||||
|
||||
| 范围 | 现有检查 | 仍需提供的证据 |
|
||||
| --- | --- | --- |
|
||||
| Rust | Rustfmt、严格 Clippy、后端测试、架构边界检查 | 相关失败场景和集成行为 |
|
||||
| 前端 | ESLint、类型检查与构建 | 页面一致性、交互和实际数据状态 |
|
||||
| 数据库迁移 | 冻结清单校验与 PR 基线比对 | 升级行为及兼容说明 |
|
||||
| 依赖安全 | Cargo Audit、pnpm Audit | 本次变更的权限、数据与信任边界 |
|
||||
| 工作流与容器 | Actionlint、相关容器检查 | 改动涉及的运行行为 |
|
||||
|
||||
前端从仓库根目录执行:
|
||||
|
||||
```bash
|
||||
pnpm --dir frontend format:check
|
||||
pnpm --dir frontend build
|
||||
git diff --check
|
||||
```
|
||||
|
||||
`build` 已包含类型检查;按当前架构约定,不维护独立前端测试代码。后端的工具链、命令与线程栈设置统一见 [修改与验收](docs/architecture.md#12-修改与验收)。
|
||||
|
||||
依赖 PostgreSQL / Redis 的测试使用 [专用测试环境](backend/migrations/README.md#本地测试库),未配置导致跳过时明确记录。CI 中已有对应服务和检查,见 [质量工作流](.github/workflows/_quality.yml) 与 [安全扫描](.github/workflows/security-scan.yml)。
|
||||
|
||||
集成变更应验证实际业务路径及结果,不能只凭健康检查或探测成功认定全链路正常。协议和网络改动按需核对认证、证书、代理、流式交付及客户端结果;记录影响复现的环境条件。缺少运行条件时说明缺口,不用模拟结果代替真实集成结论。
|
||||
|
||||
## 自动检查与合并条件
|
||||
|
||||
[CI](.github/workflows/ci.yml) 已按变更范围调用现有质量检查。新增自动规则优先使用现有工具,要求依据明确、结果稳定、误报可控;无法可靠判定的设计与视觉问题保留审查判断,避免把具体源码写法锁成测试。
|
||||
|
||||
AI 自动审查需要在所用工具中单独启用。其支持的检查范围和优先级过滤以工具说明为准;即使能读取 `AGENTS.md`,也不能保证所有可维护性和视觉问题都会被报告。配置参考 [Codex GitHub 审查文档](https://learn.chatgpt.com/docs/third-party/github)。
|
||||
|
||||
要由 GitHub 强制执行合并条件,维护者需配置对应分支的必需状态检查、审阅批准和讨论解决要求,并按需要设置新提交后的重新审阅。配置前核对实际检查名称与触发范围,见 [GitHub 分支保护](https://docs.github.com/en/repositories/configuring-branches-and-merges-in-your-repository/managing-protected-branches/about-protected-branches)。规则文档本身不会启用 AI 服务或更改分支设置。
|
||||
@@ -12,23 +12,10 @@
|
||||
[](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 协议。
|
||||
|
||||
@@ -123,6 +110,7 @@ curl http://127.0.0.1:8080/v1/models \
|
||||
- [系统架构](docs/architecture.md)
|
||||
- [管理端主题](docs/theme.md)
|
||||
- [数据库迁移](backend/migrations/README.md)
|
||||
- [贡献与审查](CONTRIBUTING.md)
|
||||
|
||||
## 社区
|
||||
|
||||
|
||||
Generated
+4
@@ -744,6 +744,7 @@ dependencies = [
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1433,6 +1434,7 @@ dependencies = [
|
||||
"cron",
|
||||
"futures",
|
||||
"gateway-core",
|
||||
"heck",
|
||||
"hex",
|
||||
"rand_core 0.6.4",
|
||||
"secrecy",
|
||||
@@ -1456,6 +1458,7 @@ dependencies = [
|
||||
"base64 0.23.1",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"flate2",
|
||||
"futures",
|
||||
"gateway-admin",
|
||||
"gateway-core",
|
||||
@@ -1471,6 +1474,7 @@ dependencies = [
|
||||
"tracing",
|
||||
"url",
|
||||
"uuid",
|
||||
"zstd",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -28,6 +28,7 @@ config = { version = "0.15.23", default-features = false, features = ["yaml"] }
|
||||
serde_json = "1"
|
||||
syn = { version = "3.0.3", features = ["full", "visit"] }
|
||||
tempfile = "3.27.0"
|
||||
tokio-rustls = { version = "=0.26.4", default-features = false }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -78,7 +78,10 @@ pub async fn run() -> Result<(), BootstrapError> {
|
||||
store.admin_ports(),
|
||||
vec![openai.admin_provider(), xai.admin_provider()],
|
||||
core.snapshot_control(),
|
||||
core.account_probe(),
|
||||
(
|
||||
core.account_probe(),
|
||||
host.proxy_probe(provider_openai::build_reqwest_client_with_custom_ca),
|
||||
),
|
||||
host.client_distribution_resolver(),
|
||||
host.system_operations(),
|
||||
)
|
||||
|
||||
@@ -182,6 +182,7 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[
|
||||
&[
|
||||
"client_distribution",
|
||||
"config",
|
||||
"proxy_probe",
|
||||
"serve",
|
||||
"system_update",
|
||||
"workers",
|
||||
|
||||
@@ -9,6 +9,151 @@ const REDIS_PASSWORD: &str = "222222222222222222222222222222222222222222222222";
|
||||
const ADMIN_PASSWORD: &str = "test-admin-password";
|
||||
const TOPOLOGY_CHILD_ENV: &str = "CPR_TEST_TOPOLOGY_CHILD";
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_should_use_provider_custom_ca_for_https_proxies() {
|
||||
use gateway_admin::ports::proxy::ProxyProbe;
|
||||
use gateway_core::account::OutboundProxy;
|
||||
use gateway_host::proxy_probe::HttpProxyProbe;
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
use tokio_rustls::{
|
||||
TlsAcceptor,
|
||||
rustls::{
|
||||
ServerConfig,
|
||||
pki_types::{CertificateDer, PrivateKeyDer, pem::PemObject as _},
|
||||
},
|
||||
};
|
||||
|
||||
const CHILD_ENV: &str = "CPR_TEST_PROXY_TLS_DIRECTORY";
|
||||
let Ok(directory) = std::env::var(CHILD_ENV) else {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
generate_proxy_test_certificates(directory.path());
|
||||
// 使用子进程隔离环境变量,避免并行测试读取到临时 CA 配置。
|
||||
for ca_env in ["CODEX_CA_CERTIFICATE", "SSL_CERT_FILE"] {
|
||||
let mut child = Command::new(std::env::current_exe().unwrap());
|
||||
child
|
||||
.args([
|
||||
"--exact",
|
||||
"bootstrap::proxy_probe_should_use_provider_custom_ca_for_https_proxies",
|
||||
"--nocapture",
|
||||
])
|
||||
.env(CHILD_ENV, directory.path())
|
||||
.env_remove("CODEX_CA_CERTIFICATE")
|
||||
.env_remove("SSL_CERT_FILE")
|
||||
.env(ca_env, directory.path().join("ca.pem"));
|
||||
if ca_env == "CODEX_CA_CERTIFICATE" {
|
||||
child.env(
|
||||
"SSL_CERT_FILE",
|
||||
directory.path().join("missing-fallback.pem"),
|
||||
);
|
||||
}
|
||||
let output = child.output().unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{ca_env}: {}\n{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
}
|
||||
return;
|
||||
};
|
||||
provider_openai::ensure_rustls_provider();
|
||||
let directory = std::path::Path::new(&directory);
|
||||
let certificate = CertificateDer::from_pem_file(directory.join("server.pem")).unwrap();
|
||||
let key = PrivateKeyDer::from_pem_file(directory.join("server.key")).unwrap();
|
||||
let acceptor = TlsAcceptor::from(Arc::new(
|
||||
ServerConfig::builder()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(vec![certificate], key)
|
||||
.unwrap(),
|
||||
));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let proxy =
|
||||
OutboundProxy::parse(&format!("https://{}", listener.local_addr().unwrap())).unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut stream = acceptor.accept(socket).await.unwrap();
|
||||
let mut request = Vec::new();
|
||||
while !request.ends_with(b"\r\n\r\n") {
|
||||
assert!(request.len() < 8192);
|
||||
request.push(stream.read_u8().await.unwrap());
|
||||
}
|
||||
assert!(request.starts_with(b"GET http://unresolvable.invalid/ip "));
|
||||
let body = "{\"ip\":\"203.0.113.8\"}";
|
||||
stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
stream.shutdown().await.unwrap();
|
||||
});
|
||||
let probe = HttpProxyProbe::new("http://unresolvable.invalid/ip")
|
||||
.with_client_builder(provider_openai::build_reqwest_client_with_custom_ca);
|
||||
let result = probe.test(&proxy).await;
|
||||
assert!(result.success, "{}", result.message);
|
||||
assert_eq!(result.exit_ip.unwrap().to_string(), "203.0.113.8");
|
||||
tokio::time::timeout(Duration::from_secs(5), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn generate_proxy_test_certificates(directory: &std::path::Path) {
|
||||
let openssl = |args: &[&str]| {
|
||||
let output = Command::new("openssl")
|
||||
.args(args)
|
||||
.current_dir(directory)
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
};
|
||||
openssl(&[
|
||||
"req",
|
||||
"-x509",
|
||||
"-newkey",
|
||||
"rsa:2048",
|
||||
"-nodes",
|
||||
"-keyout",
|
||||
"ca.key",
|
||||
"-out",
|
||||
"ca.pem",
|
||||
"-days",
|
||||
"1",
|
||||
"-subj",
|
||||
"/CN=Proxy Test CA",
|
||||
]);
|
||||
openssl(&[
|
||||
"req",
|
||||
"-newkey",
|
||||
"rsa:2048",
|
||||
"-nodes",
|
||||
"-keyout",
|
||||
"server.key",
|
||||
"-out",
|
||||
"server.csr",
|
||||
"-subj",
|
||||
"/CN=localhost",
|
||||
]);
|
||||
fs::write(directory.join("extensions"), "subjectAltName=IP:127.0.0.1\nbasicConstraints=critical,CA:FALSE\nkeyUsage=digitalSignature,keyEncipherment\nextendedKeyUsage=serverAuth\n").unwrap();
|
||||
openssl(&[
|
||||
"x509",
|
||||
"-req",
|
||||
"-in",
|
||||
"server.csr",
|
||||
"-CA",
|
||||
"ca.pem",
|
||||
"-CAkey",
|
||||
"ca.key",
|
||||
"-CAcreateserial",
|
||||
"-out",
|
||||
"server.pem",
|
||||
"-days",
|
||||
"1",
|
||||
"-extfile",
|
||||
"extensions",
|
||||
]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_load_complete_terminal_example() {
|
||||
parse_config(&valid_config()).expect("terminal config example");
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -19,6 +19,7 @@ bytes = "1.11.1"
|
||||
chrono = "0.4.45"
|
||||
futures = "0.3.32"
|
||||
gateway-core = { path = "../gateway-core" }
|
||||
heck = "0.5.0"
|
||||
hex = "0.4.3"
|
||||
rand_core = { version = "0.6.4", features = ["getrandom"] }
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
|
||||
@@ -25,7 +25,7 @@ pub use use_case::{
|
||||
account_groups::AccountGroupService, accounts::AccountsService, auth::AuthService,
|
||||
backup::BackupService, client_distribution::ClientDistributionService,
|
||||
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
|
||||
settings::SettingsService, system::SystemService, xai::XaiService,
|
||||
proxies::ProxiesService, settings::SettingsService, system::SystemService, xai::XaiService,
|
||||
};
|
||||
|
||||
use model::{AdminError, AdminErrorKind};
|
||||
@@ -149,6 +149,7 @@ pub enum AdminConfigError {
|
||||
/// 字段全部私有;调用方经 accessor 直接调用能力,不需要命名内部 `use_case` 模块。
|
||||
#[derive(Clone)]
|
||||
pub struct AdminServices {
|
||||
proxies: Arc<dyn ProxiesService>,
|
||||
auth: Arc<dyn AuthService>,
|
||||
accounts: Arc<dyn AccountsService>,
|
||||
account_groups: Arc<dyn AccountGroupService>,
|
||||
@@ -163,6 +164,11 @@ pub struct AdminServices {
|
||||
}
|
||||
|
||||
impl AdminServices {
|
||||
#[must_use]
|
||||
pub fn proxies(&self) -> &dyn ProxiesService {
|
||||
self.proxies.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn auth(&self) -> &dyn AuthService {
|
||||
self.auth.as_ref()
|
||||
@@ -247,10 +253,11 @@ pub async fn initialize(
|
||||
store: AdminStorePorts,
|
||||
providers: Vec<Arc<dyn ProviderAdmin>>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
probe: Arc<dyn AccountProbe>,
|
||||
probes: (Arc<dyn AccountProbe>, Arc<dyn ports::proxy::ProxyProbe>),
|
||||
client_distribution: Arc<dyn ClientDistributionResolver>,
|
||||
system: Arc<dyn SystemOperations>,
|
||||
) -> Result<AdminBundle, AdminError> {
|
||||
let (probe, proxy_probe) = probes;
|
||||
config
|
||||
.resolve_and_validate(Path::new("."))
|
||||
.map_err(|error| AdminError::invalid(error.to_string()))?;
|
||||
@@ -290,6 +297,12 @@ pub async fn initialize(
|
||||
backup_ports.object_store(),
|
||||
);
|
||||
let services = AdminServices {
|
||||
proxies: Arc::new(use_case::proxies::DefaultProxiesService::new(
|
||||
store.proxies(),
|
||||
proxy_probe,
|
||||
snapshot.clone(),
|
||||
registry.clone(),
|
||||
)),
|
||||
auth,
|
||||
accounts,
|
||||
account_groups: Arc::new(DefaultAccountGroupService::new(
|
||||
@@ -316,11 +329,13 @@ pub async fn initialize(
|
||||
openai: Arc::new(DefaultOpenAiService::new(
|
||||
openai,
|
||||
store.accounts(),
|
||||
store.proxies(),
|
||||
snapshot.clone(),
|
||||
)),
|
||||
xai: Arc::new(DefaultXaiService::new(
|
||||
xai,
|
||||
store.accounts(),
|
||||
store.proxies(),
|
||||
snapshot.clone(),
|
||||
)),
|
||||
backups,
|
||||
|
||||
@@ -221,7 +221,7 @@ pub struct UpdateAccount {
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
|
||||
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
|
||||
pub outbound_proxy: Option<super::proxies::AccountProxySelection>,
|
||||
}
|
||||
|
||||
/// 账号更新结果。
|
||||
@@ -239,7 +239,7 @@ pub struct BatchUpdateAccounts {
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
|
||||
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
|
||||
pub outbound_proxy: Option<super::proxies::AccountProxySelection>,
|
||||
}
|
||||
|
||||
/// 批量账号更新结果。
|
||||
|
||||
@@ -10,6 +10,9 @@ pub mod client_distribution;
|
||||
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 {
|
||||
@@ -108,6 +58,7 @@ pub struct CredentialDetails {
|
||||
|
||||
/// Provider 正式文档批量导入命令。
|
||||
pub struct ImportCredentials {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub settings: Option<AccountImportSettings>,
|
||||
pub context: MutationContext,
|
||||
pub document: ProviderDocument,
|
||||
@@ -132,6 +83,7 @@ pub struct CredentialImportResult {
|
||||
|
||||
/// Provider 解析导入文档时只接收不透明文档,不接触 revision 或审计上下文。
|
||||
pub struct PrepareCredentialImport {
|
||||
pub default_outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub document: ProviderDocument,
|
||||
}
|
||||
|
||||
@@ -175,6 +127,7 @@ pub struct PreparedCredentialImport {
|
||||
/// Admin 交给 Store 的导入事务命令。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CredentialImportCommit {
|
||||
pub outbound_proxy: Option<super::proxies::ImportProxyBinding>,
|
||||
pub settings: Option<AccountImportSettings>,
|
||||
pub prepared: PreparedCredentialImport,
|
||||
}
|
||||
@@ -255,6 +208,7 @@ pub struct PendingAuthorizationMutation {
|
||||
target: AuthorizationMutationTarget,
|
||||
owner_binding: AuthorizationOwnerBinding,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
outbound_proxy_id: Option<String>,
|
||||
}
|
||||
|
||||
impl PendingAuthorizationMutation {
|
||||
@@ -269,6 +223,7 @@ impl PendingAuthorizationMutation {
|
||||
target,
|
||||
owner_binding,
|
||||
outbound_proxy: None,
|
||||
outbound_proxy_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -299,6 +254,16 @@ impl PendingAuthorizationMutation {
|
||||
pub fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
|
||||
self.outbound_proxy.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_outbound_proxy_id(mut self, id: Option<String>) -> Self {
|
||||
self.outbound_proxy_id = id;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn outbound_proxy_id(&self) -> Option<&str> {
|
||||
self.outbound_proxy_id.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
impl PendingAuthorizationMutation {
|
||||
@@ -338,6 +303,9 @@ impl PendingAuthorizationMutation {
|
||||
Value::String(proxy.expose_url().to_owned()),
|
||||
);
|
||||
}
|
||||
if let Some(id) = &self.outbound_proxy_id {
|
||||
document.insert("outbound_proxy_id".to_owned(), Value::String(id.clone()));
|
||||
}
|
||||
document
|
||||
}
|
||||
|
||||
@@ -382,7 +350,8 @@ impl PendingAuthorizationMutation {
|
||||
started_request_id: document.started_request_id,
|
||||
},
|
||||
)
|
||||
.with_outbound_proxy(proxy))
|
||||
.with_outbound_proxy(proxy)
|
||||
.with_outbound_proxy_id(document.outbound_proxy_id))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -395,6 +364,8 @@ struct StoredAuthorizationMutationV1 {
|
||||
started_request_id: String,
|
||||
#[serde(default)]
|
||||
outbound_proxy_url: Option<String>,
|
||||
#[serde(default)]
|
||||
outbound_proxy_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -418,7 +389,7 @@ pub struct StartAuthorization {
|
||||
pub context: MutationContext,
|
||||
pub name: String,
|
||||
pub reauthorization: Option<ProviderAccountId>,
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub outbound_proxy: Option<super::proxies::AccountProxySelection>,
|
||||
}
|
||||
|
||||
/// OAuth 流程启动结果。
|
||||
@@ -750,6 +721,13 @@ pub struct ProviderQuotaWindow {
|
||||
pub provider_data: Option<ProviderDocument>,
|
||||
}
|
||||
|
||||
/// 账号用量统计周期,按周优先、月次之选择;短期限流窗口不参与统计面板。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum AccountUsagePeriod {
|
||||
Weekly,
|
||||
Monthly,
|
||||
}
|
||||
|
||||
/// Provider 解释 quota 所需的公共请求事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderQuotaRequest {
|
||||
@@ -900,23 +878,37 @@ impl ProviderQuota {
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回账号展开面板使用的当前额度窗口用量。
|
||||
/// 选择账号用量统计的周/月窗口,同时返回展示所需的周期事实。
|
||||
///
|
||||
/// Provider 已提供账号级本地用量时,固定窗口和无固定重置点的滚动窗口
|
||||
/// 都可作为代表用量;优先级与 Dashboard 代表性额度窗口一致,
|
||||
/// 同一优先级保持 Provider 投影顺序。
|
||||
/// 只使用边界完整、可归属到整个账号的窗口;周优先于月,同周期保持
|
||||
/// Provider 投影顺序。没有符合条件的窗口时不回退到短期或历史累计。
|
||||
#[must_use]
|
||||
pub fn representative_window_usage(&self) -> Option<&AccountUsage> {
|
||||
self.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, window)| {
|
||||
window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
|
||||
&& window.window_seconds.is_some()
|
||||
&& window.local_usage.is_some()
|
||||
})
|
||||
.min_by_key(|(index, window)| (quota_usage_priority(window), *index))
|
||||
.and_then(|(_, window)| window.local_usage.as_ref())
|
||||
pub fn usage_window(&self) -> Option<(&ProviderQuotaWindow, AccountUsagePeriod)> {
|
||||
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,105 @@
|
||||
//! 可复用的账号出口配置,以及脱敏后的连通性测试结果。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::account::OutboundProxy;
|
||||
|
||||
use super::{PageSize, Revision, account_groups::AccountGroupRef};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AccountProxySelection {
|
||||
Direct,
|
||||
Url(OutboundProxy),
|
||||
Saved(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ImportProxyBinding {
|
||||
pub id: String,
|
||||
pub proxy: OutboundProxy,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyListQuery {
|
||||
pub page: u32,
|
||||
pub page_size: PageSize,
|
||||
pub search: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProxyAccountRef {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub email: Option<String>,
|
||||
pub provider_kind: String,
|
||||
pub authentication_kind: String,
|
||||
pub plan_type: Option<String>,
|
||||
pub plan_type_display: Option<String>,
|
||||
pub groups: Vec<AccountGroupRef>,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// 按代理查询关联账号,分页与搜索均在存储层执行。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyAccountListQuery {
|
||||
pub proxy_id: String,
|
||||
pub page: u32,
|
||||
pub page_size: PageSize,
|
||||
pub search: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyAccountPage {
|
||||
pub items: Vec<ProxyAccountRef>,
|
||||
pub total: u64,
|
||||
pub page: u32,
|
||||
pub page_size: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProxyTestResult {
|
||||
pub success: bool,
|
||||
pub latency_ms: u64,
|
||||
pub exit_ip: Option<std::net::IpAddr>,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyRecord {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub proxy: OutboundProxy,
|
||||
pub revision: Revision,
|
||||
pub account_count: u64,
|
||||
pub last_test_at: Option<DateTime<Utc>>,
|
||||
pub last_test: Option<ProxyTestResult>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyPage {
|
||||
pub items: Vec<ProxyRecord>,
|
||||
pub total: u64,
|
||||
pub page: u32,
|
||||
pub page_size: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NewProxy {
|
||||
pub name: String,
|
||||
pub proxy: OutboundProxy,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UpdateProxy {
|
||||
pub id: String,
|
||||
pub revision: Revision,
|
||||
pub name: String,
|
||||
pub proxy: Option<OutboundProxy>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyMutation {
|
||||
pub config_revision: Revision,
|
||||
pub record: ProxyRecord,
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
//! 账号额度的只读容量估算,不参与金额结算或账号调度。
|
||||
|
||||
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 method: QuotaForecastMethod,
|
||||
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 key: String,
|
||||
pub label: String,
|
||||
pub window_seconds: u64,
|
||||
pub used_percent: Option<f64>,
|
||||
pub observed_at: Option<DateTime<Utc>>,
|
||||
pub start_at: Option<DateTime<Utc>>,
|
||||
pub reset_at: DateTime<Utc>,
|
||||
pub request_count: u64,
|
||||
pub tokens: Option<u64>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub cached_tokens: Option<u64>,
|
||||
pub known_cost_count: u64,
|
||||
pub partial_cost_count: u64,
|
||||
pub unavailable_cost_count: u64,
|
||||
pub usd: Option<f64>,
|
||||
pub sample_start_at: Option<DateTime<Utc>>,
|
||||
pub baseline_percent: f64,
|
||||
pub sampled_percent: Option<f64>,
|
||||
pub block_count: usize,
|
||||
pub observation_count: usize,
|
||||
pub missing_token_count: u64,
|
||||
pub excluded_request_count: u64,
|
||||
pub pending_request_count: u64,
|
||||
}
|
||||
|
||||
/// 优先预测真实的对应窗口;缺少对应周期时只给出明确标识的 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,
|
||||
method: QuotaForecastMethod::Cumulative,
|
||||
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 {
|
||||
key: window.key.clone(),
|
||||
label: window.label.clone(),
|
||||
window_seconds: seconds,
|
||||
used_percent: percent,
|
||||
observed_at,
|
||||
start_at: start,
|
||||
reset_at,
|
||||
request_count: usage.map_or(0, |usage| usage.request_count),
|
||||
tokens: usage.map(|usage| usage.tokens),
|
||||
input_tokens: usage.map(|usage| usage.input_tokens),
|
||||
output_tokens: usage.map(|usage| usage.output_tokens),
|
||||
cached_tokens: usage.map(|usage| usage.cached_tokens),
|
||||
known_cost_count: usage.map_or(0, |usage| usage.known_cost_count),
|
||||
partial_cost_count: 0,
|
||||
unavailable_cost_count: usage.map_or(0, |usage| usage.unavailable_cost_count),
|
||||
usd,
|
||||
sample_start_at: sample.map(|sample| sample.start_at),
|
||||
baseline_percent: sample.map_or(0.0, |sample| sample.baseline_percent),
|
||||
sampled_percent: sample.map(|sample| sample.sampled_percent),
|
||||
block_count: sample.map_or(0, |sample| sample.block_count),
|
||||
observation_count: sample.map_or(0, |sample| sample.observation_count),
|
||||
missing_token_count: usage.map_or(0, |usage| usage.missing_token_count),
|
||||
excluded_request_count: usage.map_or(0, |usage| usage.excluded_request_count),
|
||||
pending_request_count: sample.map_or(0, |sample| sample.pending_request_count),
|
||||
});
|
||||
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);
|
||||
self.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 && self.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
|
||||
|| (self.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.incomplete_tokens);
|
||||
self.estimated_tokens = tokens.and_then(|value| estimate_tokens(value, factor));
|
||||
self.estimated_usd = usd
|
||||
.filter(|_| !self.incomplete_cost)
|
||||
.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
|
||||
.filter(|_| !self.incomplete_cost)
|
||||
.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,200 @@
|
||||
//! 预测专用的配对采样;不改变完整交付用量或费用账本的统计口径。
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
impl QuotaForecastMethod {
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Cumulative => "cumulative",
|
||||
Self::Incremental => "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
|
||||
}
|
||||
@@ -3,5 +3,6 @@
|
||||
pub mod backup;
|
||||
pub mod client_distribution;
|
||||
pub mod provider;
|
||||
pub mod proxy;
|
||||
pub mod store;
|
||||
pub mod system;
|
||||
|
||||
@@ -8,6 +8,7 @@ use gateway_core::{
|
||||
operation::Operation,
|
||||
routing::{ProviderKind, UpstreamModelId},
|
||||
};
|
||||
use heck::ToUpperCamelCase;
|
||||
|
||||
use crate::model::observability::{
|
||||
CalculatedBillingBreakdown, DashboardWireProfile, ProviderBillingInput,
|
||||
@@ -20,6 +21,10 @@ use crate::model::provider_credentials::{
|
||||
ProviderProfileAvatar, ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderResetCreditResult, ProviderResetCredits, explicit_plan_type,
|
||||
};
|
||||
use crate::model::{
|
||||
provider_credentials::{ProviderDocument, ProviderQuotaWindow},
|
||||
quota_forecast_sampling::QuotaForecastObservation,
|
||||
};
|
||||
|
||||
/// Provider 管理失败的稳定分类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -170,6 +175,15 @@ 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 profile_statistics(
|
||||
&self,
|
||||
@@ -251,7 +265,20 @@ impl ProviderAdminRegistry {
|
||||
.ok_or_else(|| ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
/// 账号页和 Dashboard 共用的套餐展示投影,不修改持久化的原始套餐值。
|
||||
/// 账号目录与关联列表共用套餐补全和展示规则,已知账号套餐优先于额度快照。
|
||||
pub(crate) fn resolve_account_plan(
|
||||
&self,
|
||||
provider_kind: &str,
|
||||
plan_type: &mut Option<String>,
|
||||
quota: Option<&ProviderQuota>,
|
||||
) -> Option<String> {
|
||||
if let Some(quota) = quota {
|
||||
quota.fill_missing_plan_type(plan_type);
|
||||
}
|
||||
self.plan_type_display(provider_kind, plan_type.as_deref())
|
||||
}
|
||||
|
||||
/// 账号页和 Dashboard 共用的大驼峰套餐展示名称,不修改原始套餐值。
|
||||
pub(crate) fn plan_type_display(
|
||||
&self,
|
||||
provider_kind: &str,
|
||||
@@ -261,10 +288,12 @@ impl ProviderAdminRegistry {
|
||||
let provider = ProviderKind::new(provider_kind.to_owned())
|
||||
.ok()
|
||||
.and_then(|kind| self.providers.get(&kind));
|
||||
Some(provider.map_or_else(
|
||||
let display = provider.map_or_else(
|
||||
|| plan_type.to_owned(),
|
||||
|provider| provider.plan_type_display(plan_type),
|
||||
))
|
||||
);
|
||||
// 所有 Provider 及未知套餐统一排版,保留 Plus 后缀的含义。
|
||||
Some(display.replace('+', " Plus ").to_upper_camel_case())
|
||||
}
|
||||
|
||||
/// 返回所有已注册 Provider 的 Dashboard 上游身份画像。
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
|
||||
use super::store::AdminStoreResult;
|
||||
use crate::model::{
|
||||
MutationContext, Revision,
|
||||
proxies::{
|
||||
ImportProxyBinding, NewProxy, ProxyAccountListQuery, ProxyAccountPage, ProxyListQuery,
|
||||
ProxyMutation, ProxyPage, ProxyRecord, ProxyTestResult, UpdateProxy,
|
||||
},
|
||||
};
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxyStore: Send + Sync {
|
||||
/// 在凭据交换到提交期间保护选定代理的连接配置和测试结果。
|
||||
async fn reserve_import(&self, id: &str) -> AdminStoreResult<ProxyImportReservation>;
|
||||
async fn list(&self, query: ProxyListQuery) -> AdminStoreResult<ProxyPage>;
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage>;
|
||||
async fn get(&self, id: &str) -> AdminStoreResult<ProxyRecord>;
|
||||
/// 仅在账号仍绑定指定代理时解除关联,并清除账号保存的连接地址。
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &ProviderAccountId,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Revision>;
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation>;
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation>;
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Revision>;
|
||||
async fn record_test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
result: ProxyTestResult,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord>;
|
||||
}
|
||||
|
||||
/// 离开作用域时释放保护,错误返回和请求取消也遵循相同规则。
|
||||
pub trait ProxyImportGuard: Send + Sync {}
|
||||
|
||||
pub struct ProxyImportReservation {
|
||||
pub binding: ImportProxyBinding,
|
||||
pub guard: Box<dyn ProxyImportGuard>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxyProbe: Send + Sync {
|
||||
async fn test(&self, proxy: &OutboundProxy) -> ProxyTestResult;
|
||||
}
|
||||
@@ -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},
|
||||
};
|
||||
|
||||
@@ -110,11 +110,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,
|
||||
@@ -382,6 +382,7 @@ pub struct AdminAccountStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
groups: Arc<dyn AccountGroupStore>,
|
||||
proxies: Arc<dyn super::proxy::ProxyStore>,
|
||||
}
|
||||
|
||||
impl AdminAccountStorePorts {
|
||||
@@ -390,11 +391,13 @@ impl AdminAccountStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
groups: Arc<dyn AccountGroupStore>,
|
||||
proxies: Arc<dyn super::proxy::ProxyStore>,
|
||||
) -> Self {
|
||||
Self {
|
||||
accounts,
|
||||
runtime,
|
||||
groups,
|
||||
proxies,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -447,6 +450,11 @@ impl AdminStorePorts {
|
||||
self.accounts.groups.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn proxies(&self) -> Arc<dyn super::proxy::ProxyStore> {
|
||||
self.accounts.proxies.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn auth(&self) -> Arc<dyn AuthStore> {
|
||||
self.auth.clone()
|
||||
|
||||
@@ -28,6 +28,8 @@ use crate::{
|
||||
ProviderQuotaWindow, ProviderResetCreditResult, ProviderResetCredits,
|
||||
QuotaLocalUsageAttribution,
|
||||
},
|
||||
quota_forecast::{AccountQuotaForecastReport, account_quota_forecasts},
|
||||
quota_forecast_sampling::{QuotaForecastPoint, select_forecast_sample},
|
||||
},
|
||||
ports::{
|
||||
provider::ProviderAdminRegistry,
|
||||
@@ -83,6 +85,11 @@ pub trait AccountsService: Send + Sync {
|
||||
refresh: bool,
|
||||
) -> Result<AccountDirectoryItem, AdminError>;
|
||||
|
||||
async fn quota_forecast(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<AccountQuotaForecastReport, AdminError>;
|
||||
|
||||
async fn profile_statistics(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
@@ -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<_>>();
|
||||
@@ -279,12 +287,14 @@ impl DefaultAccountsService {
|
||||
std::slice::from_mut(&mut quota),
|
||||
)
|
||||
.await?;
|
||||
let usage = quota.representative_window_usage().cloned();
|
||||
quota.fill_missing_plan_type(&mut stored.account.plan_type);
|
||||
let usage = quota
|
||||
.usage_window()
|
||||
.and_then(|(window, _)| window.local_usage.clone());
|
||||
Ok(AccountDirectoryItem {
|
||||
plan_type_display: self.providers.plan_type_display(
|
||||
plan_type_display: self.providers.resolve_account_plan(
|
||||
stored.account.provider_kind.as_str(),
|
||||
stored.account.plan_type.as_deref(),
|
||||
&mut stored.account.plan_type,
|
||||
Some("a),
|
||||
),
|
||||
projection: stored.projection,
|
||||
usage,
|
||||
@@ -372,12 +382,14 @@ impl AccountsService for DefaultAccountsService {
|
||||
.into_iter()
|
||||
.zip(quotas)
|
||||
.map(|(mut item, quota)| {
|
||||
let usage = quota.representative_window_usage().cloned();
|
||||
quota.fill_missing_plan_type(&mut item.account.plan_type);
|
||||
let usage = quota
|
||||
.usage_window()
|
||||
.and_then(|(window, _)| window.local_usage.clone());
|
||||
AccountDirectoryItem {
|
||||
plan_type_display: self.providers.plan_type_display(
|
||||
plan_type_display: self.providers.resolve_account_plan(
|
||||
item.account.provider_kind.as_str(),
|
||||
item.account.plan_type.as_deref(),
|
||||
&mut item.account.plan_type,
|
||||
Some("a),
|
||||
),
|
||||
usage,
|
||||
account: item.account,
|
||||
@@ -577,6 +589,97 @@ impl AccountsService for DefaultAccountsService {
|
||||
self.load_directory_item(account_id, refresh).await
|
||||
}
|
||||
|
||||
async fn quota_forecast(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> 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, "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 profile_statistics(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -758,14 +861,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 {
|
||||
|
||||
@@ -8,10 +8,13 @@ pub mod client_distribution;
|
||||
pub mod client_keys;
|
||||
pub mod observability;
|
||||
pub mod openai;
|
||||
pub mod proxies;
|
||||
pub mod settings;
|
||||
pub mod system;
|
||||
pub mod xai;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind, MutationContext,
|
||||
@@ -21,7 +24,7 @@ use crate::{
|
||||
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
|
||||
PendingAuthorizationMutation, PreparedAuthorizationCommit,
|
||||
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
|
||||
StartAuthorization,
|
||||
ProviderQuotaRequest, StartAuthorization,
|
||||
},
|
||||
},
|
||||
ports::{
|
||||
@@ -96,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,
|
||||
@@ -109,13 +149,27 @@ async fn required_credential(
|
||||
.ok_or_else(|| AdminError::not_found("Provider 凭据不存在"))
|
||||
}
|
||||
|
||||
async fn import_proxy_binding(
|
||||
proxies: &dyn crate::ports::proxy::ProxyStore,
|
||||
id: Option<&str>,
|
||||
) -> Result<Option<crate::ports::proxy::ProxyImportReservation>, AdminError> {
|
||||
let Some(id) = id else { return Ok(None) };
|
||||
let reservation = proxies
|
||||
.reserve_import(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
Ok(Some(reservation))
|
||||
}
|
||||
|
||||
async fn pending_authorization(
|
||||
accounts: &dyn AccountStore,
|
||||
proxies: &dyn crate::ports::proxy::ProxyStore,
|
||||
provider_kind: &ProviderKind,
|
||||
command: &StartAuthorization,
|
||||
resource: &'static str,
|
||||
) -> Result<PendingAuthorizationMutation, AdminError> {
|
||||
let mut proxy = command.outbound_proxy.clone();
|
||||
let proxy;
|
||||
let mut proxy_id = None;
|
||||
let target = match &command.reauthorization {
|
||||
Some(account_id) => {
|
||||
proxy = required_credential(accounts, provider_kind, account_id, resource)
|
||||
@@ -126,9 +180,27 @@ async fn pending_authorization(
|
||||
account_id: account_id.clone(),
|
||||
}
|
||||
}
|
||||
None => AuthorizationMutationTarget::Create {
|
||||
name: command.name.clone(),
|
||||
},
|
||||
None => {
|
||||
use crate::model::proxies::AccountProxySelection;
|
||||
proxy = match &command.outbound_proxy {
|
||||
Some(AccountProxySelection::Url(value)) => Some(value.clone()),
|
||||
Some(AccountProxySelection::Saved(id)) => {
|
||||
let record = proxies
|
||||
.get(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
if !record.last_test.is_some_and(|test| test.success) {
|
||||
return Err(AdminError::conflict("请先测试代理连接"));
|
||||
}
|
||||
proxy_id = Some(record.id);
|
||||
Some(record.proxy)
|
||||
}
|
||||
Some(AccountProxySelection::Direct) | None => None,
|
||||
};
|
||||
AuthorizationMutationTarget::Create {
|
||||
name: command.name.clone(),
|
||||
}
|
||||
}
|
||||
};
|
||||
Ok(PendingAuthorizationMutation::new(
|
||||
provider_kind.clone(),
|
||||
@@ -137,7 +209,8 @@ async fn pending_authorization(
|
||||
&command.context,
|
||||
),
|
||||
)
|
||||
.with_outbound_proxy(proxy))
|
||||
.with_outbound_proxy(proxy)
|
||||
.with_outbound_proxy_id(proxy_id))
|
||||
}
|
||||
|
||||
fn validate_prepared_import(
|
||||
|
||||
@@ -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,
|
||||
@@ -58,6 +53,7 @@ pub trait OpenAiService: Send + Sync {
|
||||
pub(crate) struct DefaultOpenAiService {
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
}
|
||||
|
||||
@@ -66,71 +62,39 @@ impl DefaultOpenAiService {
|
||||
pub(crate) fn new(
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
accounts,
|
||||
proxies,
|
||||
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,
|
||||
) -> Result<CredentialImportResult, AdminError> {
|
||||
let context = command.context;
|
||||
let proxy_reservation = super::import_proxy_binding(
|
||||
self.proxies.as_ref(),
|
||||
command.outbound_proxy_id.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
let outbound_proxy = proxy_reservation
|
||||
.as_ref()
|
||||
.map(|reservation| reservation.binding.clone());
|
||||
let prepared = self
|
||||
.provider
|
||||
.prepare_import(PrepareCredentialImport {
|
||||
default_outbound_proxy: outbound_proxy
|
||||
.as_ref()
|
||||
.map(|binding| binding.proxy.clone()),
|
||||
document: command.document,
|
||||
})
|
||||
.await
|
||||
@@ -144,6 +108,7 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
.accounts
|
||||
.commit_credential_import(
|
||||
CredentialImportCommit {
|
||||
outbound_proxy,
|
||||
prepared,
|
||||
settings: command.settings,
|
||||
},
|
||||
@@ -151,12 +116,15 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "OpenAI credential import"))?;
|
||||
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);
|
||||
drop(proxy_reservation);
|
||||
publish_credentials_and_observe_quota(
|
||||
&self.provider,
|
||||
self.snapshot.as_ref(),
|
||||
result.config_revision,
|
||||
&result.credential_ids,
|
||||
&context.request_id,
|
||||
)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
@@ -166,6 +134,7 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
) -> Result<AuthorizationStarted, AdminError> {
|
||||
let pending = pending_authorization(
|
||||
self.accounts.as_ref(),
|
||||
self.proxies.as_ref(),
|
||||
self.provider.provider_kind(),
|
||||
&command,
|
||||
"OpenAI credential",
|
||||
@@ -203,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)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,261 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::{StreamExt, stream};
|
||||
use gateway_core::{
|
||||
account::{OutboundProxy, ProviderAccountId},
|
||||
routing::ProviderKind,
|
||||
runtime::SnapshotControl,
|
||||
};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use super::{map_store_error, publish_committed};
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind, MutationContext, Revision,
|
||||
provider_credentials::{ProviderQuotaRequest, explicit_plan_type},
|
||||
proxies::*,
|
||||
},
|
||||
ports::{
|
||||
provider::ProviderAdminRegistry,
|
||||
proxy::{ProxyProbe, ProxyStore},
|
||||
store::AdminStoreErrorKind,
|
||||
},
|
||||
};
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProxiesService: Send + Sync {
|
||||
async fn list(&self, query: ProxyListQuery) -> Result<ProxyPage, AdminError>;
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> Result<ProxyAccountPage, AdminError>;
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &str,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError>;
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError>;
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError>;
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError>;
|
||||
async fn test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyRecord, AdminError>;
|
||||
/// 探测未保存的连接地址,不写入代理记录或修改账号绑定。
|
||||
async fn probe(&self, proxy: &OutboundProxy) -> Result<ProxyTestResult, AdminError>;
|
||||
}
|
||||
|
||||
pub(crate) struct DefaultProxiesService {
|
||||
store: Arc<dyn ProxyStore>,
|
||||
probe: Arc<dyn ProxyProbe>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
providers: ProviderAdminRegistry,
|
||||
test_slots: Semaphore,
|
||||
}
|
||||
|
||||
impl DefaultProxiesService {
|
||||
pub(crate) fn new(
|
||||
store: Arc<dyn ProxyStore>,
|
||||
probe: Arc<dyn ProxyProbe>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
providers: ProviderAdminRegistry,
|
||||
) -> Self {
|
||||
Self {
|
||||
store,
|
||||
probe,
|
||||
snapshot,
|
||||
providers,
|
||||
test_slots: Semaphore::new(4),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxiesService for DefaultProxiesService {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &str,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError> {
|
||||
if proxy_id.is_empty() || proxy_id.len() > 128 || proxy_id.chars().any(char::is_control) {
|
||||
return Err(AdminError::invalid("代理 ID 不合法"));
|
||||
}
|
||||
let account_id = ProviderAccountId::new(account_id.to_owned())
|
||||
.map_err(|_| AdminError::invalid("账号 ID 不合法"))?;
|
||||
let revision = self
|
||||
.store
|
||||
.remove_account(proxy_id, &account_id, context)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
if error.kind() == AdminStoreErrorKind::Conflict {
|
||||
AdminError::conflict("账号的代理绑定已变化,请刷新后重试")
|
||||
} else {
|
||||
map_store_error(error, "proxy account")
|
||||
}
|
||||
})?;
|
||||
publish_committed(self.snapshot.as_ref(), revision).await?;
|
||||
Ok(revision)
|
||||
}
|
||||
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> Result<ProxyAccountPage, AdminError> {
|
||||
if query.proxy_id.is_empty()
|
||||
|| query.proxy_id.len() > 128
|
||||
|| query.proxy_id.chars().any(char::is_control)
|
||||
|| query.page == 0
|
||||
|| query.search.len() > 256
|
||||
|| query.search.chars().any(char::is_control)
|
||||
{
|
||||
return Err(AdminError::invalid("关联账号查询参数不合法"));
|
||||
}
|
||||
let mut page = self
|
||||
.store
|
||||
.list_accounts(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
page.items = stream::iter(page.items)
|
||||
.map(|mut account| async {
|
||||
// 与账号目录一致,缺失套餐时读取已有额度快照;限制并发且不触发上游刷新。
|
||||
let mut cached_quota = None;
|
||||
if explicit_plan_type(account.plan_type.as_deref()).is_none()
|
||||
&& let Ok(kind) = ProviderKind::new(account.provider_kind.clone())
|
||||
&& let Ok(provider) = self.providers.require(&kind)
|
||||
&& let Ok(account_id) = ProviderAccountId::new(account.id.clone())
|
||||
&& let Ok(quota) = provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id,
|
||||
refresh: false,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
cached_quota = Some(quota);
|
||||
}
|
||||
account.plan_type_display = self.providers.resolve_account_plan(
|
||||
&account.provider_kind,
|
||||
&mut account.plan_type,
|
||||
cached_quota.as_ref(),
|
||||
);
|
||||
account
|
||||
})
|
||||
.buffered(8)
|
||||
.collect()
|
||||
.await;
|
||||
Ok(page)
|
||||
}
|
||||
|
||||
async fn list(&self, query: ProxyListQuery) -> Result<ProxyPage, AdminError> {
|
||||
if query.page == 0 || query.search.len() > 256 || query.search.chars().any(char::is_control)
|
||||
{
|
||||
return Err(AdminError::invalid("代理查询参数不合法"));
|
||||
}
|
||||
self.store
|
||||
.list(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))
|
||||
}
|
||||
|
||||
async fn create(
|
||||
&self,
|
||||
mut command: NewProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError> {
|
||||
command.name = validate_name(&command.name)?;
|
||||
let result = self
|
||||
.store
|
||||
.create(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
mut command: UpdateProxy,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyMutation, AdminError> {
|
||||
command.name = validate_name(&command.name)?;
|
||||
let result = self
|
||||
.store
|
||||
.update(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<Revision, AdminError> {
|
||||
let result = self
|
||||
.store
|
||||
.delete(id, revision, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn probe(&self, proxy: &OutboundProxy) -> Result<ProxyTestResult, AdminError> {
|
||||
let _permit = self.test_slots.try_acquire().map_err(|_| {
|
||||
AdminError::new(AdminErrorKind::RateLimited, "代理测试繁忙,请稍后重试")
|
||||
})?;
|
||||
Ok(self.probe.test(proxy).await)
|
||||
}
|
||||
|
||||
async fn test(
|
||||
&self,
|
||||
id: &str,
|
||||
revision: Revision,
|
||||
context: &MutationContext,
|
||||
) -> Result<ProxyRecord, AdminError> {
|
||||
let _permit = self.test_slots.try_acquire().map_err(|_| {
|
||||
AdminError::new(AdminErrorKind::RateLimited, "代理测试繁忙,请稍后重试")
|
||||
})?;
|
||||
let record = self
|
||||
.store
|
||||
.get(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
if record.revision != revision {
|
||||
return Err(AdminError::conflict("代理已被修改,请刷新后重新测试"));
|
||||
}
|
||||
let result = self.probe.test(&record.proxy).await;
|
||||
self.store
|
||||
.record_test(id, revision, result, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_name(value: &str) -> Result<String, AdminError> {
|
||||
let value = value.trim();
|
||||
if value.is_empty() || value.chars().count() > 100 || value.chars().any(char::is_control) {
|
||||
return Err(AdminError::invalid("代理名称需要 1 至 100 个字符"));
|
||||
}
|
||||
Ok(value.to_owned())
|
||||
}
|
||||
@@ -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,
|
||||
@@ -49,6 +48,7 @@ pub trait XaiService: Send + Sync {
|
||||
pub(crate) struct DefaultXaiService {
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
}
|
||||
|
||||
@@ -57,11 +57,13 @@ impl DefaultXaiService {
|
||||
pub(crate) fn new(
|
||||
provider: Arc<dyn ProviderAdmin>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
accounts,
|
||||
proxies,
|
||||
snapshot,
|
||||
}
|
||||
}
|
||||
@@ -69,21 +71,25 @@ 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,
|
||||
) -> Result<CredentialImportResult, AdminError> {
|
||||
let context = command.context;
|
||||
let proxy_reservation = super::import_proxy_binding(
|
||||
self.proxies.as_ref(),
|
||||
command.outbound_proxy_id.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
let outbound_proxy = proxy_reservation
|
||||
.as_ref()
|
||||
.map(|reservation| reservation.binding.clone());
|
||||
let prepared = self
|
||||
.provider
|
||||
.prepare_import(PrepareCredentialImport {
|
||||
default_outbound_proxy: outbound_proxy
|
||||
.as_ref()
|
||||
.map(|binding| binding.proxy.clone()),
|
||||
document: command.document,
|
||||
})
|
||||
.await
|
||||
@@ -97,6 +103,7 @@ impl XaiService for DefaultXaiService {
|
||||
.accounts
|
||||
.commit_credential_import(
|
||||
CredentialImportCommit {
|
||||
outbound_proxy,
|
||||
prepared,
|
||||
settings: command.settings,
|
||||
},
|
||||
@@ -104,18 +111,15 @@ impl XaiService for DefaultXaiService {
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "xAI credential import"))?;
|
||||
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;
|
||||
}
|
||||
drop(proxy_reservation);
|
||||
publish_credentials_and_observe_quota(
|
||||
&self.provider,
|
||||
self.snapshot.as_ref(),
|
||||
result.config_revision,
|
||||
&result.credential_ids,
|
||||
&context.request_id,
|
||||
)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
@@ -125,6 +129,7 @@ impl XaiService for DefaultXaiService {
|
||||
) -> Result<AuthorizationStarted, AdminError> {
|
||||
let pending = pending_authorization(
|
||||
self.accounts.as_ref(),
|
||||
self.proxies.as_ref(),
|
||||
self.provider.provider_kind(),
|
||||
&command,
|
||||
"xAI credential",
|
||||
@@ -162,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)
|
||||
}
|
||||
|
||||
|
||||
@@ -157,7 +157,7 @@ fn representative_quota_should_prefer_account_wide_window_over_model_specific_wi
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn representative_window_usage_should_accept_account_wide_rolling_window_without_reset() {
|
||||
fn usage_window_should_not_use_daily_rolling_usage_for_weekly_statistics() {
|
||||
let usage = AccountUsage {
|
||||
account_id: "acct_xai".to_owned(),
|
||||
request_count: 5,
|
||||
@@ -194,14 +194,14 @@ fn representative_window_usage_should_accept_account_wide_rolling_window_without
|
||||
used_percent: None,
|
||||
reset_at: None,
|
||||
limit_reached: false,
|
||||
local_usage: Some(usage.clone()),
|
||||
local_usage: Some(usage),
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
};
|
||||
|
||||
assert_eq!(quota.representative_window_usage(), Some(&usage));
|
||||
assert_eq!(quota.usage_window(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2,3 +2,5 @@
|
||||
|
||||
mod backup;
|
||||
mod provider_credentials;
|
||||
mod quota_forecast;
|
||||
mod quota_forecast_sampling;
|
||||
|
||||
@@ -1,12 +1,135 @@
|
||||
use gateway_admin::model::provider_credentials::{
|
||||
AuthorizationMutationTarget, AuthorizationOwnerBinding, PendingAuthorizationMutation,
|
||||
ProviderQuota, ProviderQuotaWindow, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
|
||||
AccountUsagePeriod, AuthorizationMutationTarget, AuthorizationOwnerBinding,
|
||||
PendingAuthorizationMutation, ProviderQuota, ProviderQuotaWindow, ProviderQuotaWindowRole,
|
||||
QuotaLocalUsageAttribution,
|
||||
};
|
||||
use gateway_admin::model::{MutationActor, MutationContext};
|
||||
use gateway_core::account::ProviderAccountId;
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_prefer_weekly_without_changing_dashboard_selection() {
|
||||
let quota = usage_quota(vec![
|
||||
usage_window("short", "shortTerm", 18_000),
|
||||
usage_window("month", "monthly", 30 * 86_400),
|
||||
usage_window("week", "shortTerm", 7 * 86_400),
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
selected_usage_window("a),
|
||||
Some(("week", AccountUsagePeriod::Weekly))
|
||||
);
|
||||
assert_eq!(
|
||||
quota
|
||||
.representative_window()
|
||||
.map(|window| window.key.as_str()),
|
||||
Some("short")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_use_monthly_when_weekly_is_unavailable() {
|
||||
let mut model_week = usage_window("model-week", "shortTerm", 7 * 86_400);
|
||||
model_week.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
|
||||
let mut incomplete_week = usage_window("incomplete-week", "shortTerm", 7 * 86_400);
|
||||
incomplete_week.reset_at = None;
|
||||
let mut unknown_usage = usage_window("unknown-week", "shortTerm", 7 * 86_400);
|
||||
unknown_usage.local_usage = None;
|
||||
let quota = usage_quota(vec![
|
||||
model_week,
|
||||
incomplete_week,
|
||||
unknown_usage,
|
||||
usage_window("month", "monthly", 30 * 86_400),
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
selected_usage_window("a),
|
||||
Some(("month", AccountUsagePeriod::Monthly))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_not_fallback_to_short_daily_or_incomplete_windows() {
|
||||
let mut missing_duration = usage_window("unknown-month", "monthly", 30 * 86_400);
|
||||
missing_duration.window_seconds = None;
|
||||
let quota = usage_quota(vec![
|
||||
usage_window("short", "shortTerm", 18_000),
|
||||
usage_window("day", "shortTerm", 86_400),
|
||||
usage_window("invalid-month", "monthly", 86_400),
|
||||
missing_duration,
|
||||
]);
|
||||
|
||||
assert_eq!(quota.usage_window(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_preserve_provider_order_for_equal_periods() {
|
||||
let quota = usage_quota(vec![
|
||||
usage_window("first-week", "shortTerm", 7 * 86_400),
|
||||
usage_window("second-week", "shortTerm", 7 * 86_400),
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
selected_usage_window("a),
|
||||
Some(("first-week", AccountUsagePeriod::Weekly))
|
||||
);
|
||||
}
|
||||
|
||||
fn selected_usage_window(quota: &ProviderQuota) -> Option<(&str, AccountUsagePeriod)> {
|
||||
quota
|
||||
.usage_window()
|
||||
.map(|(window, period)| (window.key.as_str(), period))
|
||||
}
|
||||
|
||||
fn usage_quota(windows: Vec<ProviderQuotaWindow>) -> ProviderQuota {
|
||||
ProviderQuota {
|
||||
plan_type: None,
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows,
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_window(key: &str, group: &str, seconds: u64) -> ProviderQuotaWindow {
|
||||
ProviderQuotaWindow {
|
||||
key: key.to_owned(),
|
||||
group: group.to_owned(),
|
||||
label: key.to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(seconds),
|
||||
used_percent: Some(50.0),
|
||||
reset_at: Some(chrono::Utc::now()),
|
||||
limit_reached: false,
|
||||
local_usage: Some(gateway_admin::model::accounts::AccountUsage {
|
||||
account_id: "acct_test".to_owned(),
|
||||
request_count: 0,
|
||||
success_count: 0,
|
||||
input_tokens: None,
|
||||
output_tokens: None,
|
||||
cached_tokens: None,
|
||||
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: None,
|
||||
cost_coverage: Default::default(),
|
||||
costs: Vec::new(),
|
||||
last_used_at: None,
|
||||
request_buckets: Vec::new(),
|
||||
models: Vec::new(),
|
||||
}),
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dashboard_quota_should_preserve_unknown_and_actual_window_facts() {
|
||||
let mut quota = ProviderQuota {
|
||||
@@ -97,7 +220,21 @@ fn pending_mutation_v1_should_round_trip_all_targets_and_owners() {
|
||||
PendingAuthorizationMutation::from_storage_v1(Value::Object(
|
||||
expected.to_storage_v1()
|
||||
)),
|
||||
Ok(expected),
|
||||
Ok(expected.clone()),
|
||||
);
|
||||
let with_proxy = expected
|
||||
.with_outbound_proxy(Some(
|
||||
gateway_core::account::OutboundProxy::parse(
|
||||
"http://user:secret@proxy.example:8080",
|
||||
)
|
||||
.unwrap(),
|
||||
))
|
||||
.with_outbound_proxy_id(Some("proxy_saved".to_owned()));
|
||||
assert_eq!(
|
||||
PendingAuthorizationMutation::from_storage_v1(Value::Object(
|
||||
with_proxy.to_storage_v1()
|
||||
)),
|
||||
Ok(with_proxy)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,323 @@
|
||||
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().key, "week");
|
||||
assert_eq!(month.source.unwrap().key, "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_eq!(week.method, QuotaForecastMethod::Incremental);
|
||||
assert!(!week.low_sample);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_tokens_suppress_token_estimates_without_discarding_complete_costs() {
|
||||
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_eq!(week.estimated_tokens, None);
|
||||
assert_eq!(week.remaining_tokens, None);
|
||||
assert_eq!(week.estimated_usd, Some(10.0));
|
||||
assert_eq!(week.source.unwrap().missing_token_count, 1);
|
||||
}
|
||||
|
||||
#[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_or_unknown_costs_suppress_only_money_estimates() {
|
||||
for coverage in [
|
||||
CostCoverage {
|
||||
calculated_count: 9,
|
||||
partial_count: 1,
|
||||
..Default::default()
|
||||
},
|
||||
CostCoverage {
|
||||
calculated_count: 9,
|
||||
unavailable_count: 1,
|
||||
..Default::default()
|
||||
},
|
||||
CostCoverage::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!(week.estimated_usd.is_none());
|
||||
assert!(week.remaining_usd.is_none());
|
||||
assert_eq!(week.estimated_tokens, Some(5_000));
|
||||
}
|
||||
}
|
||||
|
||||
#[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,16 @@ 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,
|
||||
QuotaLocalUsageAttribution,
|
||||
},
|
||||
quota_forecast_sampling::{QuotaForecastHistory, QuotaForecastUsage},
|
||||
settings::{
|
||||
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
|
||||
RuntimeSettings,
|
||||
@@ -71,6 +71,8 @@ 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>,
|
||||
@@ -89,6 +91,8 @@ 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)),
|
||||
@@ -138,6 +142,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()
|
||||
@@ -377,6 +401,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 +433,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 +488,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 +506,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 +658,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(
|
||||
@@ -1391,7 +1431,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
|
||||
),
|
||||
(
|
||||
Some(" self_serve_business_prolite "),
|
||||
Some("openai display: self_serve_business_prolite"),
|
||||
Some("OpenaiDisplaySelfServeBusinessProlite"),
|
||||
)
|
||||
);
|
||||
|
||||
@@ -1409,16 +1449,27 @@ async fn accounts_list_should_return_complete_directory_semantics() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_should_fill_missing_plan_from_quota_without_overriding_known_subtypes() {
|
||||
for (stored_plan, quota_plan, expected) in [
|
||||
(None, Some("free"), Some("free")),
|
||||
(Some(" "), Some("free"), Some("free")),
|
||||
(Some("unknown"), Some("free"), Some("free")),
|
||||
for (stored_plan, quota_plan, expected, expected_display) in [
|
||||
(None, Some("free"), Some("free"), Some("OpenaiDisplayFree")),
|
||||
(
|
||||
Some(" "),
|
||||
Some("free"),
|
||||
Some("free"),
|
||||
Some("OpenaiDisplayFree"),
|
||||
),
|
||||
(
|
||||
Some("unknown"),
|
||||
Some("free"),
|
||||
Some("free"),
|
||||
Some("OpenaiDisplayFree"),
|
||||
),
|
||||
(
|
||||
Some("self_serve_business_prolite"),
|
||||
Some("team"),
|
||||
Some("self_serve_business_prolite"),
|
||||
Some("OpenaiDisplaySelfServeBusinessProlite"),
|
||||
),
|
||||
(None, None, None),
|
||||
(None, None, None, None),
|
||||
] {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
@@ -1444,10 +1495,7 @@ async fn accounts_should_fill_missing_plan_from_quota_without_overriding_known_s
|
||||
.expect("account list");
|
||||
let account = page.items.first().expect("account item");
|
||||
assert_eq!(account.account.plan_type.as_deref(), expected);
|
||||
assert_eq!(
|
||||
account.plan_type_display,
|
||||
expected.map(|plan| format!("openai display: {plan}"))
|
||||
);
|
||||
assert_eq!(account.plan_type_display.as_deref(), expected_display);
|
||||
for refresh in [false, true] {
|
||||
let detail = services
|
||||
.accounts()
|
||||
@@ -1692,6 +1740,174 @@ 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, QuotaForecastMethod,
|
||||
};
|
||||
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_eq!(result.forecasts[0].method, QuotaForecastMethod::Incremental);
|
||||
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());
|
||||
@@ -1746,10 +1962,8 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
.as_ref()
|
||||
.expect("quota window local usage");
|
||||
assert_eq!(usage.total_tokens, Some(4_330_000));
|
||||
assert_eq!(
|
||||
item.usage.as_ref().and_then(|usage| usage.total_tokens),
|
||||
Some(4_330_000),
|
||||
);
|
||||
// 短期额度条保留自己的本地用量,但不作为周/月统计面板的回退。
|
||||
assert!(item.usage.is_none());
|
||||
|
||||
let queries = store.quota_window_queries();
|
||||
assert_eq!(queries.len(), 1);
|
||||
@@ -1759,6 +1973,98 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
assert_eq!(queries[0].range.end, reset_at);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_and_quota_refresh_should_select_the_same_weekly_or_monthly_usage() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
let services = accounts_service(provider.clone(), store.clone()).await;
|
||||
let account_id = ProviderAccountId::new("acct_test").expect("account id");
|
||||
let reset_at = Utc::now() + TimeDelta::days(1);
|
||||
let short = ProviderQuotaWindow {
|
||||
key: "short".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(18_000),
|
||||
used_percent: Some(50.0),
|
||||
reset_at: Some(reset_at),
|
||||
limit_reached: false,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
};
|
||||
let week = ProviderQuotaWindow {
|
||||
key: "week".to_owned(),
|
||||
label: "周额度".to_owned(),
|
||||
window_seconds: Some(7 * 86_400),
|
||||
..short.clone()
|
||||
};
|
||||
let month = ProviderQuotaWindow {
|
||||
key: "month".to_owned(),
|
||||
group: "monthly".to_owned(),
|
||||
label: "月额度".to_owned(),
|
||||
window_seconds: Some(30 * 86_400),
|
||||
..short.clone()
|
||||
};
|
||||
for (windows, selected_key, tokens, duration) in [
|
||||
(vec![short.clone(), month.clone(), week], "week", 700, 7),
|
||||
(vec![short, month], "month", 3000, 30),
|
||||
] {
|
||||
provider.set_quota(ProviderQuota {
|
||||
windows,
|
||||
..empty_quota()
|
||||
});
|
||||
store.set_quota_window_usage(
|
||||
[("short", 50), ("week", 700), ("month", 3000)]
|
||||
.into_iter()
|
||||
.map(|(key, total)| AccountUsageWindowResult {
|
||||
account_id: account_id.to_string(),
|
||||
key: key.to_owned(),
|
||||
usage: quota_local_usage(account_id.as_str(), total),
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
let page = services
|
||||
.accounts()
|
||||
.list(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("list accounts");
|
||||
let refreshed = services
|
||||
.accounts()
|
||||
.quota(&account_id, true)
|
||||
.await
|
||||
.expect("refresh quota");
|
||||
let expected = quota_local_usage(account_id.as_str(), tokens);
|
||||
for item in [&page.items[0], &refreshed] {
|
||||
let usage = item.usage.as_ref().expect("selected window usage");
|
||||
assert_eq!(usage.total_tokens, expected.total_tokens);
|
||||
assert_eq!(
|
||||
item.quota
|
||||
.usage_window()
|
||||
.map(|(window, _)| window.key.as_str()),
|
||||
Some(selected_key)
|
||||
);
|
||||
}
|
||||
let queries = store.quota_window_queries();
|
||||
let selected = queries
|
||||
.iter()
|
||||
.find(|query| query.key == selected_key)
|
||||
.expect("selected query");
|
||||
assert_eq!(selected.range.start, reset_at - TimeDelta::days(duration));
|
||||
assert_eq!(selected.range.end, reset_at);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_windows() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
|
||||
@@ -5,6 +5,7 @@ mod backup;
|
||||
mod client_keys;
|
||||
mod observability;
|
||||
mod openai;
|
||||
mod proxies;
|
||||
mod settings;
|
||||
mod system;
|
||||
mod xai;
|
||||
@@ -44,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},
|
||||
@@ -82,6 +82,7 @@ pub(super) struct AdminHarness {
|
||||
default_password: String,
|
||||
session_ttl_minutes: u64,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
proxies: Arc<dyn gateway_admin::ports::proxy::ProxyStore>,
|
||||
account_runtime: Arc<dyn AccountRuntimeStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
@@ -101,6 +102,7 @@ impl AdminHarness {
|
||||
default_password: "strong-test-password".to_owned(),
|
||||
session_ttl_minutes: 60,
|
||||
accounts: unavailable.clone(),
|
||||
proxies: Arc::new(proxies::TestProxies::default()),
|
||||
account_runtime: unavailable.clone(),
|
||||
account_groups: Arc::new(UnavailableAccountGroupStore),
|
||||
auth: Arc::new(BootstrapAuthStore::default()),
|
||||
@@ -179,6 +181,14 @@ impl AdminHarness {
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn proxies(
|
||||
mut self,
|
||||
proxies: Arc<dyn gateway_admin::ports::proxy::ProxyStore>,
|
||||
) -> Self {
|
||||
self.proxies = proxies;
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn system(mut self, system: Arc<dyn SystemOperations>) -> Self {
|
||||
self.system = system;
|
||||
self
|
||||
@@ -196,6 +206,7 @@ impl AdminHarness {
|
||||
self.accounts,
|
||||
self.account_runtime,
|
||||
self.account_groups,
|
||||
self.proxies,
|
||||
),
|
||||
self.auth,
|
||||
self.client_keys,
|
||||
@@ -205,7 +216,7 @@ impl AdminHarness {
|
||||
),
|
||||
self.providers,
|
||||
Arc::new(NoopSnapshot),
|
||||
self.probe,
|
||||
(self.probe, Arc::new(proxies::TestProxies::default())),
|
||||
Arc::new(NoopClientDistribution),
|
||||
self.system,
|
||||
)
|
||||
@@ -360,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(
|
||||
|
||||
@@ -50,6 +50,7 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
|
||||
services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: Some(super::accounts::import_settings()),
|
||||
context: context("import-openai"),
|
||||
document: document(),
|
||||
@@ -57,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),
|
||||
@@ -107,6 +108,7 @@ async fn openai_import_should_expose_only_explicit_public_errors_without_committ
|
||||
let error = services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-pat-failure"),
|
||||
document: document(),
|
||||
@@ -132,6 +134,7 @@ async fn openai_import_should_refresh_quota_for_every_imported_account() {
|
||||
let result = services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-batch"),
|
||||
document: document(),
|
||||
@@ -139,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,
|
||||
@@ -176,6 +179,7 @@ async fn openai_import_should_remain_successful_when_quota_refresh_fails() {
|
||||
let result = services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-quota-failure"),
|
||||
document: document(),
|
||||
@@ -183,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,
|
||||
@@ -220,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),
|
||||
@@ -276,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);
|
||||
@@ -318,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),
|
||||
@@ -348,6 +352,7 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
|
||||
services
|
||||
.openai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-openai-error"),
|
||||
document: document(),
|
||||
@@ -401,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)));
|
||||
|
||||
@@ -429,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"),
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
use async_trait::async_trait;
|
||||
use gateway_admin::{
|
||||
model::{MutationContext, Revision, proxies::*},
|
||||
ports::{
|
||||
proxy::{ProxyProbe, ProxyStore},
|
||||
store::AdminStoreResult,
|
||||
},
|
||||
};
|
||||
use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct TestProxies {
|
||||
pub events: Option<super::accounts::EventLog>,
|
||||
pub accounts: Option<Vec<ProxyAccountRef>>,
|
||||
}
|
||||
|
||||
struct ImportGuard(super::accounts::EventLog);
|
||||
|
||||
impl gateway_admin::ports::proxy::ProxyImportGuard for ImportGuard {}
|
||||
|
||||
impl Drop for ImportGuard {
|
||||
fn drop(&mut self) {
|
||||
self.0.lock().unwrap().push("proxy.release");
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyStore for TestProxies {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
_: &str,
|
||||
_: &ProviderAccountId,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
|
||||
async fn reserve_import(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> AdminStoreResult<gateway_admin::ports::proxy::ProxyImportReservation> {
|
||||
let events = self
|
||||
.events
|
||||
.as_ref()
|
||||
.ok_or_else(|| super::unavailable("proxy"))?;
|
||||
events.lock().unwrap().push("proxy.reserve");
|
||||
Ok(gateway_admin::ports::proxy::ProxyImportReservation {
|
||||
binding: ImportProxyBinding {
|
||||
id: id.to_owned(),
|
||||
proxy: OutboundProxy::parse("http://127.0.0.1:8080").unwrap(),
|
||||
},
|
||||
guard: Box::new(ImportGuard(events.clone())),
|
||||
})
|
||||
}
|
||||
async fn list(&self, _: ProxyListQuery) -> AdminStoreResult<ProxyPage> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage> {
|
||||
let items = self
|
||||
.accounts
|
||||
.clone()
|
||||
.ok_or_else(|| super::unavailable("proxy"))?;
|
||||
Ok(ProxyAccountPage {
|
||||
total: items.len() as u64,
|
||||
items,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn get(&self, _: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn create(&self, _: NewProxy, _: &MutationContext) -> AdminStoreResult<ProxyMutation> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn update(&self, _: UpdateProxy, _: &MutationContext) -> AdminStoreResult<ProxyMutation> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn delete(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
async fn record_test(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
_: ProxyTestResult,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord> {
|
||||
Err(super::unavailable("proxy"))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyProbe for TestProxies {
|
||||
async fn test(&self, _: &OutboundProxy) -> ProxyTestResult {
|
||||
panic!("unexpected proxy probe")
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn linked_accounts_share_plan_resolution_and_only_read_cached_quota() {
|
||||
use super::accounts::{FakeProviderAdmin, events};
|
||||
use gateway_admin::model::{PageSize, provider_credentials::ProviderQuota};
|
||||
use std::sync::Arc;
|
||||
|
||||
for (stored, cached, expected) in [
|
||||
(None, Some("free"), Some("free")),
|
||||
(Some("unknown"), Some("free"), Some("free")),
|
||||
(Some("plus"), Some("free"), Some("plus")),
|
||||
(None, None, None),
|
||||
] {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
plan_type: cached.map(str::to_owned),
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![],
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
});
|
||||
let services = super::AdminHarness::new()
|
||||
.provider(provider.clone())
|
||||
.proxies(Arc::new(TestProxies {
|
||||
accounts: Some(vec![ProxyAccountRef {
|
||||
id: "acct_plan".to_owned(),
|
||||
name: "套餐测试".to_owned(),
|
||||
email: None,
|
||||
provider_kind: "openai".to_owned(),
|
||||
authentication_kind: "oauth".to_owned(),
|
||||
plan_type: stored.map(str::to_owned),
|
||||
plan_type_display: None,
|
||||
groups: vec![],
|
||||
enabled: true,
|
||||
}]),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let result = services
|
||||
.proxies()
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: "proxy_plan".to_owned(),
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
search: String::new(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.items[0].plan_type.as_deref(), expected);
|
||||
assert_eq!(
|
||||
result.items[0].plan_type_display.as_deref(),
|
||||
match expected {
|
||||
Some("free") => Some("OpenaiDisplayFree"),
|
||||
Some("plus") => Some("OpenaiDisplayPlus"),
|
||||
_ => None,
|
||||
}
|
||||
);
|
||||
let requests = provider.quota_requests();
|
||||
assert_eq!(requests.len(), usize::from(stored != Some("plus")));
|
||||
assert!(requests.iter().all(|request| !request.refresh));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_import_keeps_proxy_reserved_until_commit_and_releases_on_errors() {
|
||||
use super::accounts::{
|
||||
FakeAccountStore, FakeProviderAdmin, context, document, events, recorded,
|
||||
};
|
||||
use gateway_admin::model::provider_credentials::ImportCredentials;
|
||||
use gateway_admin::ports::provider::ProviderAdminErrorKind;
|
||||
use std::sync::Arc;
|
||||
|
||||
for kind in ["openai", "xai"] {
|
||||
for failure in [None, Some("prepare"), Some("commit")] {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new(kind, events.clone());
|
||||
let store = FakeAccountStore::new(kind, events.clone());
|
||||
if failure == Some("prepare") {
|
||||
provider.fail_next(ProviderAdminErrorKind::Unavailable);
|
||||
}
|
||||
if failure == Some("commit") {
|
||||
store.fail_next_commit();
|
||||
}
|
||||
let services = super::AdminHarness::new()
|
||||
.provider(provider)
|
||||
.accounts(store)
|
||||
.proxies(Arc::new(TestProxies {
|
||||
events: Some(events.clone()),
|
||||
..Default::default()
|
||||
}))
|
||||
.build()
|
||||
.await;
|
||||
let command = ImportCredentials {
|
||||
outbound_proxy_id: Some("proxy_import".to_owned()),
|
||||
settings: None,
|
||||
context: context("reserved-import"),
|
||||
document: document(),
|
||||
};
|
||||
let result = if kind == "openai" {
|
||||
services.openai().import_document(command).await
|
||||
} else {
|
||||
services.xai().import_document(command).await
|
||||
};
|
||||
assert_eq!(result.is_err(), failure.is_some());
|
||||
let events = recorded(&events);
|
||||
let expected = if failure == Some("prepare") {
|
||||
vec!["proxy.reserve", "provider.prepare_import", "proxy.release"]
|
||||
} else {
|
||||
vec![
|
||||
"proxy.reserve",
|
||||
"provider.prepare_import",
|
||||
"store.commit_import",
|
||||
"proxy.release",
|
||||
]
|
||||
};
|
||||
assert_eq!(&events[..expected.len()], expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -20,6 +22,7 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
|
||||
services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: Some(super::accounts::import_settings()),
|
||||
context: context("import-xai"),
|
||||
document: document(),
|
||||
@@ -27,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"
|
||||
]
|
||||
);
|
||||
@@ -61,6 +66,7 @@ async fn xai_import_should_refresh_quota_for_every_imported_account() {
|
||||
services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-xai-batch"),
|
||||
document: document(),
|
||||
@@ -68,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(),
|
||||
[
|
||||
@@ -96,6 +103,7 @@ async fn xai_import_should_remain_successful_when_quota_refresh_fails() {
|
||||
let result = services
|
||||
.xai()
|
||||
.import_document(ImportCredentials {
|
||||
outbound_proxy_id: None,
|
||||
settings: None,
|
||||
context: context("import-xai-quota-failure"),
|
||||
document: document(),
|
||||
@@ -103,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)
|
||||
|
||||
@@ -16,6 +16,7 @@ axum = { version = "0.8.9", default-features = false, features = ["json", "query
|
||||
base64 = "0.23.0"
|
||||
bytes = "1.11.1"
|
||||
chrono = { version = "0.4.45", features = ["serde"] }
|
||||
flate2 = "1.1.9"
|
||||
futures = "0.3.32"
|
||||
gateway-admin = { path = "../gateway-admin" }
|
||||
gateway-core = { path = "../gateway-core" }
|
||||
@@ -29,6 +30,7 @@ tower-http = { version = "0.7.0", features = ["cors", "fs", "request-id", "timeo
|
||||
tracing = "0.1.44"
|
||||
url = "2.5.8"
|
||||
uuid = { version = "1.23.3", features = ["v7"] }
|
||||
zstd = "=0.13.3"
|
||||
|
||||
[dev-dependencies]
|
||||
async-trait = "0.1.89"
|
||||
|
||||
@@ -53,6 +53,7 @@ impl AccountImportSettingsRequest {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct AccountImportRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub settings: Option<AccountImportSettingsRequest>,
|
||||
pub provider: String,
|
||||
pub data: Value,
|
||||
@@ -60,6 +61,9 @@ pub struct AccountImportRequest {
|
||||
|
||||
impl AccountImportRequest {
|
||||
pub fn validate(&self) -> Result<(), WireValidationError> {
|
||||
if let Some(id) = &self.outbound_proxy_id {
|
||||
require_wire_id(id, "outboundProxyId")?;
|
||||
}
|
||||
if let Some(settings) = &self.settings {
|
||||
settings.validate()?;
|
||||
}
|
||||
@@ -82,6 +86,7 @@ impl AccountImportRequest {
|
||||
Ok((
|
||||
provider,
|
||||
ImportCredentials {
|
||||
outbound_proxy_id: self.outbound_proxy_id,
|
||||
settings: self
|
||||
.settings
|
||||
.map(AccountImportSettingsRequest::into_settings)
|
||||
@@ -96,6 +101,7 @@ impl AccountImportRequest {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct StartAccountAuthorizationRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub account_id: Option<String>,
|
||||
@@ -129,7 +135,10 @@ impl StartAccountAuthorizationRequest {
|
||||
context,
|
||||
name: self.name,
|
||||
reauthorization,
|
||||
outbound_proxy: self.outbound_proxy_url.and_then(|value| value.0),
|
||||
outbound_proxy: super::wire::proxy_selection(
|
||||
self.outbound_proxy_id,
|
||||
self.outbound_proxy_url,
|
||||
)?,
|
||||
},
|
||||
))
|
||||
}
|
||||
@@ -188,6 +197,7 @@ impl CompleteAccountAuthorizationRequest {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct UpdateAccountRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub outbound_proxy_url: Option<super::wire::AccountProxyUpdate>,
|
||||
pub account_id: String,
|
||||
pub enabled: bool,
|
||||
@@ -209,7 +219,10 @@ impl UpdateAccountRequest {
|
||||
pub(super) fn into_command(self) -> Result<UpdateAccount, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(UpdateAccount {
|
||||
outbound_proxy: self.outbound_proxy_url.map(|value| value.0),
|
||||
outbound_proxy: super::wire::proxy_selection(
|
||||
self.outbound_proxy_id,
|
||||
self.outbound_proxy_url,
|
||||
)?,
|
||||
account_id: self.account_id,
|
||||
enabled: self.enabled,
|
||||
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
|
||||
|
||||
@@ -22,6 +22,10 @@ where
|
||||
post(batch_update_accounts::<S>),
|
||||
)
|
||||
.route("/api/admin/accounts/quota", get(account_quota::<S>))
|
||||
.route(
|
||||
"/api/admin/accounts/quota-forecast",
|
||||
get(account_quota_forecast::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/profile-statistics",
|
||||
get(account_profile_statistics::<S>),
|
||||
@@ -364,6 +368,27 @@ where
|
||||
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
|
||||
}
|
||||
|
||||
async fn account_quota_forecast<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()
|
||||
.quota_forecast(&account_id)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountQuotaForecastData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn account_profile_statistics<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
|
||||
@@ -26,13 +26,13 @@ use gateway_admin::model::{
|
||||
},
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
|
||||
CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation,
|
||||
CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels,
|
||||
ProviderProfileActivityInsights, ProviderProfileAvatar, ProviderProfileDailyUsage,
|
||||
ProviderProfileInvocation, ProviderProfileStatistics, ProviderProfileStatisticsSummary,
|
||||
ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult,
|
||||
ProviderResetCredits, RotateCredential, StartAuthorization,
|
||||
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,
|
||||
},
|
||||
};
|
||||
use gateway_core::{
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
//! Admin 领域结果到安全 HTTP wire 的展示投影。
|
||||
|
||||
use super::*;
|
||||
use gateway_admin::model::quota_forecast::{AccountQuotaForecast, AccountQuotaForecastReport};
|
||||
use gateway_admin::model::quota_forecast_sampling::QuotaForecastMethod;
|
||||
|
||||
pub(super) fn account_page_data(
|
||||
result: AccountDirectoryPage,
|
||||
@@ -69,6 +71,7 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> Ac
|
||||
let expires_at = account.access_token_expires_at.as_ref().map(china_rfc3339);
|
||||
let added_at = china_rfc3339(&account.created_at);
|
||||
let updated_at = china_rfc3339(&account.updated_at);
|
||||
let usage_period = quota.usage_window().map(|(_, period)| period);
|
||||
let (quota, refresh_token_expires_at) = account_quota_view(quota, rate_limited_until, now);
|
||||
AccountView {
|
||||
id: account.id.clone(),
|
||||
@@ -117,10 +120,105 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> Ac
|
||||
updated_at,
|
||||
updated_at_display: china_datetime(&account.updated_at),
|
||||
quota,
|
||||
usage: account_usage_view(usage, now),
|
||||
usage: account_usage_view(usage, usage_period, now),
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AccountQuotaForecastReport> for AccountQuotaForecastData {
|
||||
fn from(report: AccountQuotaForecastReport) -> Self {
|
||||
Self {
|
||||
account_id: report.account_id,
|
||||
generated_at: china_rfc3339(&report.generated_at),
|
||||
generated_at_display: china_datetime(&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 {
|
||||
key: source.key,
|
||||
label: source.label,
|
||||
window_days: source.window_seconds as f64 / 86_400.0,
|
||||
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)),
|
||||
start_at: source.start_at.map(|value| china_rfc3339(&value)),
|
||||
start_at_display: source
|
||||
.start_at
|
||||
.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)),
|
||||
reset_at: china_rfc3339(&source.reset_at),
|
||||
reset_at_display: china_datetime(&source.reset_at),
|
||||
request_count_display: format_number(source.request_count),
|
||||
tokens_display: display_optional_tokens(source.tokens),
|
||||
input_tokens_display: display_optional_tokens(source.input_tokens),
|
||||
output_tokens_display: display_optional_tokens(source.output_tokens),
|
||||
cached_tokens_display: display_optional_tokens(source.cached_tokens),
|
||||
known_cost_count: source.known_cost_count,
|
||||
known_cost_count_display: format_number(source.known_cost_count),
|
||||
partial_cost_count: source.partial_cost_count,
|
||||
partial_cost_count_display: format_number(source.partial_cost_count),
|
||||
unavailable_cost_count: source.unavailable_cost_count,
|
||||
unavailable_cost_count_display: format_number(source.unavailable_cost_count),
|
||||
usd_display: forecast_usd_display(source.usd),
|
||||
sample_start_at: source.sample_start_at.map(|value| china_rfc3339(&value)),
|
||||
sample_start_at_display: source
|
||||
.sample_start_at
|
||||
.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)),
|
||||
baseline_percent: source.baseline_percent,
|
||||
sampled_percent: source.sampled_percent,
|
||||
sampled_percent_display: source
|
||||
.sampled_percent
|
||||
.map_or_else(|| "—".to_owned(), |value| format!("{value:.1} 个百分点")),
|
||||
block_count: source.block_count,
|
||||
observation_count: source.observation_count,
|
||||
missing_token_count: source.missing_token_count,
|
||||
excluded_request_count: source.excluded_request_count,
|
||||
pending_request_count: source.pending_request_count,
|
||||
}),
|
||||
unavailable_reason: forecast.unavailable_reason,
|
||||
low_sample: forecast.low_sample,
|
||||
incomplete_cost: forecast.incomplete_cost,
|
||||
incomplete_tokens: forecast.incomplete_tokens,
|
||||
method: forecast.method.as_str(),
|
||||
method_display: match forecast.method {
|
||||
QuotaForecastMethod::Cumulative => "窗口累计估算",
|
||||
QuotaForecastMethod::Incremental => "近期分段估算",
|
||||
},
|
||||
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>>,
|
||||
@@ -211,6 +309,7 @@ pub(super) fn quota_local_usage(usage: &AccountUsage) -> Value {
|
||||
|
||||
pub(super) fn account_usage_view(
|
||||
usage: Option<AccountUsage>,
|
||||
period: Option<AccountUsagePeriod>,
|
||||
now: DateTime<Utc>,
|
||||
) -> AccountUsageView {
|
||||
let Some(usage) = usage else {
|
||||
@@ -228,6 +327,12 @@ pub(super) fn account_usage_view(
|
||||
"known"
|
||||
};
|
||||
AccountUsageView {
|
||||
window_label_display: match period {
|
||||
Some(AccountUsagePeriod::Weekly) => "周额度窗口",
|
||||
Some(AccountUsagePeriod::Monthly) => "月额度窗口",
|
||||
None => "周/月额度窗口",
|
||||
}
|
||||
.to_owned(),
|
||||
request_count: Some(usage.request_count),
|
||||
request_count_display: format_number(usage.request_count),
|
||||
input_tokens: usage.input_tokens,
|
||||
@@ -347,6 +452,7 @@ pub(super) fn display_optional_tokens(value: Option<u64>) -> String {
|
||||
|
||||
pub(super) fn empty_account_usage() -> AccountUsageView {
|
||||
AccountUsageView {
|
||||
window_label_display: "周/月额度窗口".to_owned(),
|
||||
request_count: None,
|
||||
request_count_display: "—".to_owned(),
|
||||
input_tokens: None,
|
||||
|
||||
@@ -17,6 +17,33 @@ impl<'de> Deserialize<'de> for AccountProxyUpdate {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn proxy_selection(
|
||||
id: Option<String>,
|
||||
url: Option<AccountProxyUpdate>,
|
||||
) -> Result<Option<gateway_admin::model::proxies::AccountProxySelection>, WireValidationError> {
|
||||
use gateway_admin::model::proxies::AccountProxySelection;
|
||||
match (id, url) {
|
||||
(Some(_), Some(_)) => Err(WireValidationError::new("outboundProxyId")),
|
||||
(Some(id), None) if id.is_empty() => Ok(Some(AccountProxySelection::Direct)),
|
||||
(Some(id), None) => {
|
||||
if id.len() > 128
|
||||
|| !id
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
|
||||
{
|
||||
return Err(WireValidationError::new("outboundProxyId"));
|
||||
}
|
||||
Ok(Some(AccountProxySelection::Saved(id)))
|
||||
}
|
||||
(None, Some(value)) => Ok(Some(
|
||||
value
|
||||
.0
|
||||
.map_or(AccountProxySelection::Direct, AccountProxySelection::Url),
|
||||
)),
|
||||
(None, None) => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号列表查询参数。
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
@@ -34,6 +61,7 @@ pub struct ListQuery {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct BatchUpdateAccountsRequest {
|
||||
pub outbound_proxy_id: Option<String>,
|
||||
pub outbound_proxy_url: Option<AccountProxyUpdate>,
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
@@ -84,7 +112,7 @@ impl BatchUpdateAccountsRequest {
|
||||
pub(super) fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(BatchUpdateAccounts {
|
||||
outbound_proxy: self.outbound_proxy_url.map(|value| value.0),
|
||||
outbound_proxy: proxy_selection(self.outbound_proxy_id, self.outbound_proxy_url)?,
|
||||
account_ids: self.account_ids,
|
||||
enabled: self.enabled,
|
||||
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
|
||||
@@ -244,6 +272,77 @@ 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 generated_at_display: 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 method: &'static str,
|
||||
pub method_display: &'static str,
|
||||
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 key: String,
|
||||
pub label: String,
|
||||
pub window_days: f64,
|
||||
pub used_percent: Option<f64>,
|
||||
pub used_percent_display: String,
|
||||
pub observed_at: Option<String>,
|
||||
pub observed_at_display: String,
|
||||
pub start_at: Option<String>,
|
||||
pub start_at_display: String,
|
||||
pub reset_at: String,
|
||||
pub reset_at_display: String,
|
||||
pub request_count_display: String,
|
||||
pub tokens_display: String,
|
||||
pub input_tokens_display: String,
|
||||
pub output_tokens_display: String,
|
||||
pub cached_tokens_display: String,
|
||||
pub known_cost_count: u64,
|
||||
pub known_cost_count_display: String,
|
||||
pub partial_cost_count: u64,
|
||||
pub partial_cost_count_display: String,
|
||||
pub unavailable_cost_count: u64,
|
||||
pub unavailable_cost_count_display: String,
|
||||
pub usd_display: String,
|
||||
pub sample_start_at: Option<String>,
|
||||
pub sample_start_at_display: String,
|
||||
pub baseline_percent: f64,
|
||||
pub sampled_percent: Option<f64>,
|
||||
pub sampled_percent_display: String,
|
||||
pub block_count: usize,
|
||||
pub observation_count: usize,
|
||||
pub missing_token_count: u64,
|
||||
pub excluded_request_count: u64,
|
||||
pub pending_request_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountGroupRefView {
|
||||
@@ -288,6 +387,7 @@ pub struct AccountQuotaWindowView {
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountUsageView {
|
||||
pub window_label_display: String,
|
||||
pub request_count: Option<u64>,
|
||||
pub request_count_display: String,
|
||||
pub input_tokens: Option<u64>,
|
||||
|
||||
@@ -16,6 +16,7 @@ pub mod client_keys;
|
||||
mod extract;
|
||||
pub mod observability;
|
||||
pub mod presenter;
|
||||
pub mod proxies;
|
||||
pub mod settings;
|
||||
pub mod system;
|
||||
pub mod wire;
|
||||
@@ -34,6 +35,7 @@ where
|
||||
{
|
||||
Router::new()
|
||||
.merge(account_groups::router::<S>())
|
||||
.merge(proxies::router::<S>())
|
||||
.merge(accounts::router::<S>())
|
||||
.merge(auth::router::<S>())
|
||||
.merge(backups::router::<S>())
|
||||
|
||||
@@ -0,0 +1,452 @@
|
||||
use axum::{
|
||||
Router,
|
||||
extract::State,
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
use gateway_admin::model::{
|
||||
PageSize, Revision,
|
||||
proxies::{
|
||||
NewProxy, ProxyAccountListQuery, ProxyListQuery, ProxyMutation, ProxyRecord,
|
||||
ProxyTestResult, UpdateProxy,
|
||||
},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{
|
||||
AdminAuth, AdminEnvelope, AdminError, AdminJson, AdminQuery, AdminResponse, AdminSessionState,
|
||||
PageMeta,
|
||||
accounts::{AccountGroupRefView, AccountProxyUpdate},
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ListQuery {
|
||||
page: Option<u32>,
|
||||
page_size: Option<u16>,
|
||||
search: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct AccountsQuery {
|
||||
proxy_id: String,
|
||||
page: Option<u32>,
|
||||
page_size: Option<u16>,
|
||||
search: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct RemoveAccountRequest {
|
||||
proxy_id: String,
|
||||
account_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct CreateRequest {
|
||||
name: String,
|
||||
proxy_url: AccountProxyUpdate,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct UpdateRequest {
|
||||
id: String,
|
||||
revision: u64,
|
||||
name: String,
|
||||
proxy_url: Option<AccountProxyUpdate>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct IdRequest {
|
||||
id: String,
|
||||
revision: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ProbeRequest {
|
||||
proxy_url: AccountProxyUpdate,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyTestView {
|
||||
success: bool,
|
||||
latency_ms: u64,
|
||||
exit_ip: Option<String>,
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl From<ProxyTestResult> for ProxyTestView {
|
||||
fn from(result: ProxyTestResult) -> Self {
|
||||
Self {
|
||||
success: result.success,
|
||||
latency_ms: result.latency_ms,
|
||||
exit_ip: result.exit_ip.map(|ip| ip.to_string()),
|
||||
message: result.message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyAccountView {
|
||||
id: String,
|
||||
name: String,
|
||||
email: Option<String>,
|
||||
provider: String,
|
||||
authentication_kind: String,
|
||||
plan_type: Option<String>,
|
||||
plan_type_display: Option<String>,
|
||||
groups: Vec<AccountGroupRefView>,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyView {
|
||||
id: String,
|
||||
name: String,
|
||||
endpoint: String,
|
||||
has_authentication: bool,
|
||||
revision: u64,
|
||||
account_count: u64,
|
||||
last_test_at: Option<String>,
|
||||
last_test: Option<ProxyTestView>,
|
||||
created_at: String,
|
||||
updated_at: String,
|
||||
}
|
||||
|
||||
impl From<ProxyRecord> for ProxyView {
|
||||
fn from(record: ProxyRecord) -> Self {
|
||||
let endpoint = record.proxy.endpoint();
|
||||
Self {
|
||||
id: record.id,
|
||||
name: record.name,
|
||||
has_authentication: record.proxy.expose_url() != endpoint,
|
||||
endpoint,
|
||||
revision: record.revision.get(),
|
||||
account_count: record.account_count,
|
||||
last_test_at: record.last_test_at.map(|at| at.to_rfc3339()),
|
||||
last_test: record.last_test.map(Into::into),
|
||||
created_at: record.created_at.to_rfc3339(),
|
||||
updated_at: record.updated_at.to_rfc3339(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ProxyPageView {
|
||||
items: Vec<ProxyView>,
|
||||
page: PageMeta,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ProxyAccountPageView {
|
||||
items: Vec<ProxyAccountView>,
|
||||
page: PageMeta,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct MutationView {
|
||||
record: ProxyView,
|
||||
config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<ProxyMutation> for MutationView {
|
||||
fn from(value: ProxyMutation) -> Self {
|
||||
Self {
|
||||
record: value.record.into(),
|
||||
config_revision: value.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn router<S>() -> Router<S>
|
||||
where
|
||||
S: AdminSessionState + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Router::new()
|
||||
.route("/api/admin/proxies", get(list::<S>))
|
||||
.route("/api/admin/proxies/accounts", get(list_accounts::<S>))
|
||||
.route(
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
post(remove_account::<S>),
|
||||
)
|
||||
.route("/api/admin/proxies/create", post(create::<S>))
|
||||
.route("/api/admin/proxies/update", post(update::<S>))
|
||||
.route("/api/admin/proxies/delete", post(delete::<S>))
|
||||
.route("/api/admin/proxies/test", post(test::<S>))
|
||||
.route("/api/admin/proxies/probe", post(probe::<S>))
|
||||
}
|
||||
|
||||
fn revision(value: u64) -> Result<Revision, AdminError> {
|
||||
if value > i64::MAX as u64 {
|
||||
return Err(AdminError::bad_request("代理版本不合法"));
|
||||
}
|
||||
Revision::new(value).map_err(|_| AdminError::bad_request("代理版本不合法"))
|
||||
}
|
||||
|
||||
fn map_error(error: gateway_admin::model::AdminError) -> AdminError {
|
||||
super::wire::map_admin_service_error(error)
|
||||
}
|
||||
|
||||
async fn list<S>(
|
||||
_: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<ListQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.list(ProxyListQuery {
|
||||
page: query.page.unwrap_or(1),
|
||||
page_size: PageSize::new(query.page_size.unwrap_or(20))
|
||||
.map_err(|_| AdminError::bad_request("分页大小不合法"))?,
|
||||
search: query.search.unwrap_or_default().trim().to_owned(),
|
||||
})
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
let total_pages = result.total.div_ceil(u64::from(result.page_size));
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyPageView {
|
||||
items: result.items.into_iter().map(Into::into).collect(),
|
||||
page: PageMeta::new(
|
||||
result.page,
|
||||
u32::from(result.page_size),
|
||||
result.total,
|
||||
u32::try_from(total_pages).unwrap_or(u32::MAX),
|
||||
),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn list_accounts<S>(
|
||||
_: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<AccountsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.list_accounts(ProxyAccountListQuery {
|
||||
proxy_id: query.proxy_id,
|
||||
page: query.page.unwrap_or(1),
|
||||
page_size: PageSize::new(query.page_size.unwrap_or(20))
|
||||
.map_err(|_| AdminError::bad_request("分页大小不合法"))?,
|
||||
search: query.search.unwrap_or_default().trim().to_owned(),
|
||||
})
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
let total_pages = result.total.div_ceil(u64::from(result.page_size));
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyAccountPageView {
|
||||
items: result
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|account| ProxyAccountView {
|
||||
id: account.id,
|
||||
name: account.name,
|
||||
email: account.email,
|
||||
provider: account.provider_kind,
|
||||
authentication_kind: account.authentication_kind,
|
||||
plan_type: account.plan_type,
|
||||
plan_type_display: account.plan_type_display,
|
||||
groups: account
|
||||
.groups
|
||||
.into_iter()
|
||||
.map(|group| AccountGroupRefView {
|
||||
id: group.id.to_string(),
|
||||
name: group.name,
|
||||
color: group.color.as_str().to_owned(),
|
||||
enabled: group.enabled,
|
||||
})
|
||||
.collect(),
|
||||
enabled: account.enabled,
|
||||
})
|
||||
.collect(),
|
||||
page: PageMeta::new(
|
||||
result.page,
|
||||
u32::from(result.page_size),
|
||||
result.total,
|
||||
u32::try_from(total_pages).unwrap_or(u32::MAX),
|
||||
),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn create<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<CreateRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let proxy = request
|
||||
.proxy_url
|
||||
.0
|
||||
.ok_or_else(|| AdminError::bad_request("代理 URL 不能为空"))?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.create(
|
||||
NewProxy {
|
||||
name: request.name,
|
||||
proxy,
|
||||
},
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::CREATED,
|
||||
AdminEnvelope::ok(MutationView::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn remove_account<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<RemoveAccountRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let revision = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.remove_account(
|
||||
&request.proxy_id,
|
||||
&request.account_id,
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"configRevision": revision.get()})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn update<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<UpdateRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let proxy = request
|
||||
.proxy_url
|
||||
.map(|value| {
|
||||
value
|
||||
.0
|
||||
.ok_or_else(|| AdminError::bad_request("代理 URL 不能为空"))
|
||||
})
|
||||
.transpose()?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.update(
|
||||
UpdateProxy {
|
||||
id: request.id,
|
||||
revision: revision(request.revision)?,
|
||||
name: request.name,
|
||||
proxy,
|
||||
},
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(MutationView::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn delete<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<IdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.delete(
|
||||
&request.id,
|
||||
revision(request.revision)?,
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"configRevision": result.get()})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn probe<S>(
|
||||
_: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<ProbeRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let proxy = request
|
||||
.proxy_url
|
||||
.0
|
||||
.ok_or_else(|| AdminError::bad_request("代理 URL 不能为空"))?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.probe(&proxy)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyTestView::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn test<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(request): AdminJson<IdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.proxies()
|
||||
.test(
|
||||
&request.id,
|
||||
revision(request.revision)?,
|
||||
&auth.context().mutation_context(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ProxyView::from(result)),
|
||||
))
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -124,6 +124,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 +174,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,
|
||||
|
||||
@@ -40,6 +40,15 @@ pub enum RequestDecodeError {
|
||||
/// 请求不是合法 JSON。
|
||||
#[error("request body must be valid JSON")]
|
||||
MalformedJson,
|
||||
/// 请求体解压后超过允许上限。
|
||||
#[error("decompressed request body is too large")]
|
||||
DecompressedBodyTooLarge,
|
||||
/// 请求使用了网关不支持的 `Content-Encoding`。
|
||||
#[error("request content-encoding `{encoding}` is not supported")]
|
||||
UnsupportedContentEncoding {
|
||||
/// Content-Encoding 标识。
|
||||
encoding: String,
|
||||
},
|
||||
/// 顶层不是 object。
|
||||
#[error("request body must be a JSON object")]
|
||||
ExpectedObject,
|
||||
@@ -99,6 +108,16 @@ impl RequestDecodeError {
|
||||
"Request body must be valid JSON.".to_owned(),
|
||||
None,
|
||||
),
|
||||
Self::DecompressedBodyTooLarge => (
|
||||
"request_too_large",
|
||||
"Decompressed request body exceeds the allowed size.".to_owned(),
|
||||
None,
|
||||
),
|
||||
Self::UnsupportedContentEncoding { encoding } => (
|
||||
"unsupported_content_encoding",
|
||||
format!("Content-Encoding `{encoding}` is not supported."),
|
||||
None,
|
||||
),
|
||||
Self::ExpectedObject => (
|
||||
"invalid_request_body",
|
||||
"Request body must be a JSON object.".to_owned(),
|
||||
|
||||
@@ -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 后提交下游,再持续驱动同一执行会话。
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
//! OpenAI Responses JSON 到 Core 路由事实与不透明 wire payload 的单一解码边界。
|
||||
|
||||
use std::{fmt, net::IpAddr};
|
||||
use std::{borrow::Cow, fmt, net::IpAddr};
|
||||
|
||||
use axum::http::HeaderMap;
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
|
||||
use gateway_protocol::openai::{
|
||||
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
|
||||
is_transport_managed_request_header,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
@@ -322,12 +323,87 @@ impl fmt::Debug for DecodedResponsesRequest {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// JSON 非法、顶层不是 object,或网关必须解释的路由字段无效时返回安全错误。
|
||||
/// 编码不支持、压缩正文损坏或超限、JSON 非法,或路由字段无效时返回安全错误。
|
||||
pub fn decode_request_with_headers(
|
||||
body: &[u8],
|
||||
headers: &HeaderMap,
|
||||
) -> Result<DecodedResponsesRequest, RequestDecodeError> {
|
||||
decode_request_inner(body, &OpenAiRequestHeaders::from_headers(headers))
|
||||
let body = decompress_request_body(body, headers)?;
|
||||
decode_request_inner(&body, &OpenAiRequestHeaders::from_headers(headers))
|
||||
}
|
||||
|
||||
/// 下游请求体解压后的最大字节数;防御解压炸弹。
|
||||
const MAX_DECOMPRESSED_REQUEST_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
/// 按 `Content-Encoding` 解压下游请求体。
|
||||
///
|
||||
/// 未压缩与 `identity` 借用原始正文;压缩正文在读取过程中限制展开大小,
|
||||
/// 避免先完整分配再检查。只接受单一编码,重复头和叠加编码不能只解释第一项。
|
||||
fn decompress_request_body<'a>(
|
||||
body: &'a [u8],
|
||||
headers: &HeaderMap,
|
||||
) -> Result<Cow<'a, [u8]>, RequestDecodeError> {
|
||||
let encoding = headers
|
||||
.get_all(axum::http::header::CONTENT_ENCODING)
|
||||
.iter()
|
||||
.map(|value| {
|
||||
value
|
||||
.to_str()
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
.join(",");
|
||||
match encoding.trim().to_ascii_lowercase().as_str() {
|
||||
"" | "identity" => Ok(Cow::Borrowed(body)),
|
||||
"gzip" => {
|
||||
let decoder = flate2::read::MultiGzDecoder::new(body);
|
||||
read_bounded(decoder).map(Cow::Owned)
|
||||
}
|
||||
"deflate" => {
|
||||
let decoder = flate2::read::ZlibDecoder::new(body);
|
||||
read_bounded(decoder).map(Cow::Owned)
|
||||
}
|
||||
"zstd" => {
|
||||
let mut decoder = zstd::stream::read::Decoder::with_buffer(body)
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
// 输出有界之外,也限制压缩帧声明的回溯窗口,防止解码器内部过量分配。
|
||||
decoder
|
||||
.window_log_max(MAX_DECOMPRESSED_REQUEST_BYTES.ilog2())
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
read_bounded(decoder).map(Cow::Owned)
|
||||
}
|
||||
other => Err(RequestDecodeError::UnsupportedContentEncoding {
|
||||
encoding: other.to_owned(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_bounded<R: std::io::Read>(mut reader: R) -> Result<Vec<u8>, RequestDecodeError> {
|
||||
let mut decoded = Vec::new();
|
||||
let mut chunk = [0; 8192];
|
||||
loop {
|
||||
let remaining = MAX_DECOMPRESSED_REQUEST_BYTES - decoded.len();
|
||||
let read_limit = chunk.len().min(remaining + 1);
|
||||
let read = reader
|
||||
.read(&mut chunk[..read_limit])
|
||||
.map_err(|_| RequestDecodeError::MalformedJson)?;
|
||||
if read == 0 {
|
||||
return Ok(decoded);
|
||||
}
|
||||
if read > remaining {
|
||||
return Err(RequestDecodeError::DecompressedBodyTooLarge);
|
||||
}
|
||||
// 默认 Vec 扩容和 read_to_end 的 EOF 探测可能突破输出上限;
|
||||
// 容量增长也受相同边界约束,越界探测只使用栈上分块缓冲区。
|
||||
let required = decoded.len() + read;
|
||||
if required > decoded.capacity() {
|
||||
let capacity = (decoded.capacity() * 2)
|
||||
.max(required)
|
||||
.min(MAX_DECOMPRESSED_REQUEST_BYTES);
|
||||
decoded.reserve_exact(capacity - decoded.len());
|
||||
}
|
||||
decoded.extend_from_slice(&chunk[..read]);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn decode_request_inner(
|
||||
@@ -458,7 +534,7 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
|
||||
if connection_headers
|
||||
.iter()
|
||||
.any(|connection_header| connection_header.eq_ignore_ascii_case(name))
|
||||
|| name.starts_with("sec-websocket-")
|
||||
|| is_transport_managed_request_header(name)
|
||||
|| name.starts_with("x-grok-")
|
||||
|| name.starts_with("x-xai-")
|
||||
{
|
||||
@@ -467,30 +543,8 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
|
||||
|
||||
!matches!(
|
||||
name,
|
||||
// 逐跳头、上游 authority 和重序列化后的实体长度由 transport 重建。
|
||||
"connection"
|
||||
| "keep-alive"
|
||||
| "proxy-connection"
|
||||
| "proxy-authenticate"
|
||||
| "proxy-authorization"
|
||||
| "te"
|
||||
| "trailer"
|
||||
| "transfer-encoding"
|
||||
| "upgrade"
|
||||
| "host"
|
||||
| "content-length"
|
||||
// 下游连接诊断事实只用于本地观测,不能伪装成 OpenAI 请求事实。
|
||||
| "forwarded"
|
||||
| "x-forwarded-for"
|
||||
| "x-forwarded-host"
|
||||
| "x-forwarded-proto"
|
||||
| "x-forwarded-port"
|
||||
| "x-real-ip"
|
||||
| "true-client-ip"
|
||||
| "cf-connecting-ip"
|
||||
| "x-request-id"
|
||||
// 下游鉴权和账号 cookie 绝不能成为上游账号身份。
|
||||
| "authorization"
|
||||
// 下游鉴权和账号 cookie 绝不能成为上游账号身份。
|
||||
"authorization"
|
||||
| "x-api-key"
|
||||
// Codex 的服务端托管认证标记只用于客户端能力判断,不代表上游身份。
|
||||
| "x-openai-actor-authorization"
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
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 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,6 @@
|
||||
mod handlers;
|
||||
mod presenter;
|
||||
|
||||
mod query {
|
||||
use gateway_admin::model::accounts::{AccountSortField, AccountStatus, SortDirection};
|
||||
use gateway_api::admin::accounts::ListQuery;
|
||||
@@ -340,6 +343,7 @@ mod response {
|
||||
#[test]
|
||||
fn account_usage_view_should_keep_unobserved_numbers_null() {
|
||||
let view = AccountUsageView {
|
||||
window_label_display: "周/月额度窗口".to_owned(),
|
||||
request_count: None,
|
||||
request_count_display: "-".to_owned(),
|
||||
input_tokens: None,
|
||||
@@ -374,6 +378,7 @@ mod response {
|
||||
models: Vec::new(),
|
||||
};
|
||||
let value = serde_json::to_value(view).expect("serialize account usage");
|
||||
assert_eq!(value["windowLabelDisplay"], "周/月额度窗口");
|
||||
assert!(value["inputTokens"].is_null());
|
||||
assert!(value["totalTokens"].is_null());
|
||||
assert!(value["reasoningTokens"].is_null());
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
use gateway_admin::model::{
|
||||
provider_credentials::AccountUsagePeriod,
|
||||
quota_forecast::{AccountQuotaForecast, AccountQuotaForecastReport, QuotaForecastSource},
|
||||
quota_forecast_sampling::QuotaForecastMethod,
|
||||
};
|
||||
use gateway_api::admin::accounts::AccountQuotaForecastData;
|
||||
|
||||
#[test]
|
||||
fn quota_forecast_projection_preserves_null_zero_and_sample_metadata() {
|
||||
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 {
|
||||
key: "week".to_owned(),
|
||||
label: "周额度".to_owned(),
|
||||
window_seconds: 7 * 86_400,
|
||||
used_percent: Some(100.0),
|
||||
observed_at: Some(now),
|
||||
start_at: Some(now),
|
||||
reset_at: now,
|
||||
request_count: 1_024,
|
||||
tokens: Some(1_000_000),
|
||||
input_tokens: Some(900_000),
|
||||
output_tokens: Some(100_000),
|
||||
cached_tokens: Some(800_000),
|
||||
known_cost_count: 1_000,
|
||||
partial_cost_count: 24,
|
||||
unavailable_cost_count: 0,
|
||||
usd: Some(0.1234),
|
||||
sample_start_at: Some(now),
|
||||
baseline_percent: 85.0,
|
||||
sampled_percent: Some(15.0),
|
||||
block_count: 3,
|
||||
observation_count: 96,
|
||||
missing_token_count: 0,
|
||||
excluded_request_count: 4,
|
||||
pending_request_count: 2,
|
||||
}),
|
||||
unavailable_reason: None,
|
||||
low_sample: false,
|
||||
incomplete_cost: true,
|
||||
incomplete_tokens: false,
|
||||
method: QuotaForecastMethod::Incremental,
|
||||
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"]["requestCountDisplay"], "1,024");
|
||||
assert_eq!(week["source"]["usdDisplay"], "$0.1234");
|
||||
assert_eq!(week["source"]["partialCostCount"], 24);
|
||||
assert_eq!(week["method"], "incremental");
|
||||
assert_eq!(week["source"]["sampledPercentDisplay"], "15.0 个百分点");
|
||||
assert_eq!(week["source"]["pendingRequestCount"], 2);
|
||||
assert!(week["source"].get("providerData").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());
|
||||
}
|
||||
@@ -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,
|
||||
@@ -82,6 +81,7 @@ mod auth;
|
||||
mod client_keys;
|
||||
mod errors;
|
||||
mod observability;
|
||||
mod proxies;
|
||||
mod settings;
|
||||
mod system;
|
||||
mod wire;
|
||||
@@ -96,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 {
|
||||
@@ -115,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),
|
||||
@@ -122,9 +126,15 @@ 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(unused.clone(), unused.clone(), account_groups.clone()),
|
||||
AdminAccountStorePorts::new(
|
||||
unused.clone(),
|
||||
unused.clone(),
|
||||
account_groups.clone(),
|
||||
Arc::new(proxies::MemoryProxies::default()),
|
||||
),
|
||||
auth.clone(),
|
||||
client_keys.clone(),
|
||||
unused,
|
||||
@@ -132,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 {
|
||||
@@ -144,7 +154,7 @@ impl AdminTestFixture {
|
||||
stores,
|
||||
providers,
|
||||
Arc::new(NoopSnapshot),
|
||||
Arc::new(NoopProbe),
|
||||
(Arc::new(NoopProbe), Arc::new(proxies::SuccessfulProbe)),
|
||||
Arc::new(StaticClientDistribution),
|
||||
system,
|
||||
)
|
||||
@@ -160,6 +170,8 @@ impl AdminTestFixture {
|
||||
ops_errors,
|
||||
dashboard_observation,
|
||||
dashboard_summary_range,
|
||||
provider_error,
|
||||
account,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -697,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]
|
||||
@@ -711,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"))
|
||||
}
|
||||
|
||||
@@ -732,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(
|
||||
@@ -930,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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -976,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(
|
||||
@@ -1004,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::{
|
||||
|
||||
@@ -0,0 +1,520 @@
|
||||
use std::sync::Mutex;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use gateway_admin::{
|
||||
model::{MutationContext, Revision, proxies::*},
|
||||
ports::{
|
||||
proxy::{ProxyProbe, ProxyStore},
|
||||
store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
|
||||
},
|
||||
};
|
||||
use gateway_core::account::{OutboundProxy, ProviderAccountId};
|
||||
use serde_json::{Value, json};
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState};
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct MemoryProxies(Mutex<Option<ProxyRecord>>);
|
||||
|
||||
fn missing() -> AdminStoreError {
|
||||
AdminStoreError::new(AdminStoreErrorKind::NotFound, "proxy", "missing proxy")
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyStore for MemoryProxies {
|
||||
async fn remove_account(
|
||||
&self,
|
||||
proxy_id: &str,
|
||||
account_id: &ProviderAccountId,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
if proxy_id != "proxy_test" || account_id.as_str() != "acct_linked" {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Conflict,
|
||||
"proxy account",
|
||||
"changed binding",
|
||||
));
|
||||
}
|
||||
Ok(Revision::new(4).unwrap())
|
||||
}
|
||||
|
||||
async fn reserve_import(
|
||||
&self,
|
||||
_: &str,
|
||||
) -> AdminStoreResult<gateway_admin::ports::proxy::ProxyImportReservation> {
|
||||
Err(missing())
|
||||
}
|
||||
async fn list(&self, query: ProxyListQuery) -> AdminStoreResult<ProxyPage> {
|
||||
let items: Vec<_> = self.0.lock().unwrap().iter().cloned().collect();
|
||||
Ok(ProxyPage {
|
||||
total: u64::try_from(items.len()).unwrap(),
|
||||
items,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn get(&self, _: &str) -> AdminStoreResult<ProxyRecord> {
|
||||
self.0.lock().unwrap().clone().ok_or_else(missing)
|
||||
}
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: ProxyAccountListQuery,
|
||||
) -> AdminStoreResult<ProxyAccountPage> {
|
||||
if !self
|
||||
.0
|
||||
.lock()
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.is_some_and(|record| record.id == query.proxy_id)
|
||||
{
|
||||
return Err(missing());
|
||||
}
|
||||
Ok(ProxyAccountPage {
|
||||
items: vec![ProxyAccountRef {
|
||||
id: "acct_linked".to_owned(),
|
||||
name: "工作账号".to_owned(),
|
||||
email: Some("work@example.invalid".to_owned()),
|
||||
provider_kind: "openai".to_owned(),
|
||||
authentication_kind: "oauth".to_owned(),
|
||||
plan_type: Some("plus".to_owned()),
|
||||
plan_type_display: None,
|
||||
groups: vec![],
|
||||
enabled: true,
|
||||
}],
|
||||
total: 21,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
async fn create(
|
||||
&self,
|
||||
command: NewProxy,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let record = ProxyRecord {
|
||||
id: "proxy_test".to_owned(),
|
||||
name: command.name,
|
||||
proxy: command.proxy,
|
||||
revision: Revision::new(1).unwrap(),
|
||||
account_count: 0,
|
||||
last_test: None,
|
||||
last_test_at: None,
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
};
|
||||
*self.0.lock().unwrap() = Some(record.clone());
|
||||
Ok(ProxyMutation {
|
||||
config_revision: record.revision,
|
||||
record,
|
||||
})
|
||||
}
|
||||
async fn update(
|
||||
&self,
|
||||
command: UpdateProxy,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let mut stored = self.0.lock().unwrap();
|
||||
let record = stored.as_mut().ok_or_else(missing)?;
|
||||
record.name = command.name;
|
||||
if let Some(proxy) = command.proxy {
|
||||
record.proxy = proxy;
|
||||
}
|
||||
record.revision = Revision::new(record.revision.get() + 1).unwrap();
|
||||
Ok(ProxyMutation {
|
||||
config_revision: record.revision,
|
||||
record: record.clone(),
|
||||
})
|
||||
}
|
||||
async fn delete(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
self.0.lock().unwrap().take().ok_or_else(missing)?;
|
||||
Ok(Revision::new(3).unwrap())
|
||||
}
|
||||
async fn record_test(
|
||||
&self,
|
||||
_: &str,
|
||||
_: Revision,
|
||||
result: ProxyTestResult,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyRecord> {
|
||||
let mut stored = self.0.lock().unwrap();
|
||||
let record = stored.as_mut().ok_or_else(missing)?;
|
||||
record.last_test = Some(result);
|
||||
record.last_test_at = Some(Utc::now());
|
||||
Ok(record.clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct SuccessfulProbe;
|
||||
|
||||
#[async_trait]
|
||||
impl ProxyProbe for SuccessfulProbe {
|
||||
async fn test(&self, proxy: &OutboundProxy) -> ProxyTestResult {
|
||||
assert_eq!(
|
||||
proxy.expose_url(),
|
||||
"http://test-user:private-password@proxy.example:8080/"
|
||||
);
|
||||
ProxyTestResult {
|
||||
success: true,
|
||||
latency_ms: 15,
|
||||
exit_ip: Some("203.0.113.2".parse().unwrap()),
|
||||
message: "Connected".to_owned(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn request(
|
||||
fixture: &AdminTestFixture,
|
||||
path: &str,
|
||||
body: Option<Value>,
|
||||
authenticated: bool,
|
||||
) -> (StatusCode, Value) {
|
||||
let mut builder = Request::builder()
|
||||
.header("x-request-id", "req_proxy_tests")
|
||||
.uri(path)
|
||||
.method(if body.is_some() { "POST" } else { "GET" });
|
||||
if authenticated {
|
||||
builder = builder.header(header::COOKIE, "cpr_admin_session=valid-session");
|
||||
}
|
||||
let body = body.map_or_else(Body::empty, |body| Body::from(body.to_string()));
|
||||
let response = gateway_api::admin::proxies::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(
|
||||
builder
|
||||
.header(header::CONTENT_TYPE, "application/json")
|
||||
.body(body)
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
|
||||
let text = std::str::from_utf8(&body).unwrap();
|
||||
assert!(!text.contains("private-password"));
|
||||
assert!(!text.contains("test-user"));
|
||||
(status, serde_json::from_slice(&body).unwrap())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_routes_save_reload_test_rename_and_delete_without_exposing_credentials() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let (status, created) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({
|
||||
"name": " Office ", "proxyUrl": "http://test-user:private-password@proxy.example:8080"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CREATED);
|
||||
assert_eq!(created["data"]["record"]["name"], "Office");
|
||||
assert_eq!(created["data"]["record"]["hasAuthentication"], true);
|
||||
assert_eq!(
|
||||
created["data"]["record"]["endpoint"],
|
||||
"http://proxy.example:8080/"
|
||||
);
|
||||
assert!(created["data"]["record"].get("proxyUrl").is_none());
|
||||
assert_eq!(created["data"]["record"]["accountCount"], 0);
|
||||
assert!(created["data"]["record"].get("accounts").is_none());
|
||||
|
||||
let (status, linked) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts?proxyId=proxy_test&page=2&pageSize=20&search=work",
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(
|
||||
linked["data"]["page"],
|
||||
json!({"page":2,"pageSize":20,"total":21,"totalPages":2})
|
||||
);
|
||||
assert_eq!(
|
||||
linked["data"]["items"],
|
||||
json!([{
|
||||
"id":"acct_linked", "name":"工作账号", "email":"work@example.invalid",
|
||||
"provider":"openai", "enabled":true, "authenticationKind":"oauth",
|
||||
"planType":"plus", "planTypeDisplay":"Plus", "groups":[]
|
||||
}])
|
||||
);
|
||||
|
||||
let (status, tested) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/test",
|
||||
Some(json!({"id": "proxy_test", "revision": 1})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(tested["data"]["lastTest"]["exitIp"], "203.0.113.2");
|
||||
let (_, listed) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies?page=1&pageSize=20",
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(listed["data"]["items"][0], tested["data"]);
|
||||
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/update",
|
||||
Some(json!({"id": "proxy_test", "revision": 1, "name": "Renamed"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/test",
|
||||
Some(json!({"id": "proxy_test", "revision": 1})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CONFLICT);
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/test",
|
||||
Some(json!({"id": "proxy_test", "revision": 2})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/delete",
|
||||
Some(json!({"id": "proxy_test", "revision": 2})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["page"]["total"], 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_checks_unsaved_address_without_creating_or_changing_records() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let draft = json!({
|
||||
"proxyUrl": "http://test-user:private-password@proxy.example:8080"
|
||||
});
|
||||
let (status, probed) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/probe",
|
||||
Some(draft.clone()),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(
|
||||
probed["data"],
|
||||
json!({"success": true, "latencyMs": 15, "exitIp": "203.0.113.2", "message": "Connected"})
|
||||
);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["page"]["total"], 0);
|
||||
|
||||
let (_, created) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({"name": "Saved", "proxyUrl": "http://saved.example:8080"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
let (status, _) = request(&fixture, "/api/admin/proxies/probe", Some(draft), true).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["items"][0], created["data"]["record"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_probe_rejects_empty_or_invalid_addresses() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (body, expected) in [
|
||||
(json!({"proxyUrl": ""}), StatusCode::BAD_REQUEST),
|
||||
(json!({"proxyUrl": null}), StatusCode::UNPROCESSABLE_ENTITY),
|
||||
(json!({}), StatusCode::UNPROCESSABLE_ENTITY),
|
||||
(
|
||||
json!({"proxyUrl": "ftp://test-user:private-password@proxy.example"}),
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
),
|
||||
(
|
||||
json!({"proxyUrl": "http://proxy.example:8080", "name": "Not saved"}),
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
request(&fixture, "/api/admin/proxies/probe", Some(body), true)
|
||||
.await
|
||||
.0,
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_routes_require_auth_and_reject_invalid_input() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
assert_eq!(
|
||||
request(&fixture, "/api/admin/proxies", None, false).await.0,
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts?proxyId=proxy_test",
|
||||
None,
|
||||
false
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
for query in [
|
||||
"",
|
||||
"proxyId=",
|
||||
"proxyId=proxy_test&page=0",
|
||||
"proxyId=proxy_test&pageSize=0",
|
||||
"proxyId=proxy_test&pageSize=201",
|
||||
"proxyId=proxy_test&search=%00",
|
||||
"proxyId=proxy_test&unknown=1",
|
||||
] {
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
&format!("/api/admin/proxies/accounts?{query}"),
|
||||
None,
|
||||
true
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::BAD_REQUEST
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts?proxyId=missing",
|
||||
None,
|
||||
true
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::NOT_FOUND
|
||||
);
|
||||
for action in [
|
||||
"create",
|
||||
"update",
|
||||
"delete",
|
||||
"test",
|
||||
"probe",
|
||||
"accounts/remove",
|
||||
] {
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
&format!("/api/admin/proxies/{action}"),
|
||||
Some(json!({})),
|
||||
false
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
}
|
||||
for query in ["page=0", "pageSize=0", "pageSize=201", "unknown=1"] {
|
||||
assert_eq!(
|
||||
request(&fixture, &format!("/api/admin/proxies?{query}"), None, true)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::BAD_REQUEST
|
||||
);
|
||||
}
|
||||
for body in [
|
||||
json!({"name":" ","proxyUrl":"http://proxy.example:8080"}),
|
||||
json!({"name":"Invalid","proxyUrl":""}),
|
||||
] {
|
||||
assert_eq!(
|
||||
request(&fixture, "/api/admin/proxies/create", Some(body), true)
|
||||
.await
|
||||
.0,
|
||||
StatusCode::BAD_REQUEST
|
||||
);
|
||||
}
|
||||
let (status, _) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({"name":"Invalid","proxyUrl":"ftp://test-user:private-password@proxy.example"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_account_removal_validates_binding_and_input() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
for (body, expected) in [
|
||||
(
|
||||
json!({"proxyId":"", "accountId":"acct_linked"}),
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
json!({"proxyId":"proxy_test", "accountId":""}),
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
json!({"proxyId":"proxy_test", "accountId":"invalid"}),
|
||||
StatusCode::BAD_REQUEST,
|
||||
),
|
||||
(
|
||||
json!({"proxyId":"proxy_test", "accountId":"acct_linked", "enabled":false}),
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
Some(body),
|
||||
true
|
||||
)
|
||||
.await
|
||||
.0,
|
||||
expected
|
||||
);
|
||||
}
|
||||
let (status, result) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
Some(json!({"proxyId":"proxy_test", "accountId":"acct_linked"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(result["data"], json!({"configRevision":4}));
|
||||
let (status, result) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/accounts/remove",
|
||||
Some(json!({"proxyId":"proxy_previous", "accountId":"acct_linked"})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CONFLICT);
|
||||
assert_eq!(result["message"], "账号的代理绑定已变化,请刷新后重试");
|
||||
}
|
||||
@@ -50,6 +50,7 @@ fn source_tree_should_match_frozen_machine_manifest() {
|
||||
"src/admin/observability/routes.rs",
|
||||
"src/admin/observability/wire.rs",
|
||||
"src/admin/presenter.rs",
|
||||
"src/admin/proxies.rs",
|
||||
"src/admin/settings.rs",
|
||||
"src/admin/system.rs",
|
||||
"src/admin/wire.rs",
|
||||
@@ -88,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",
|
||||
@@ -96,6 +99,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/admin/observability/mod.rs",
|
||||
"tests/admin/observability/query.rs",
|
||||
"tests/admin/observability/response.rs",
|
||||
"tests/admin/proxies.rs",
|
||||
"tests/admin/settings.rs",
|
||||
"tests/admin/system.rs",
|
||||
"tests/admin/wire.rs",
|
||||
@@ -110,7 +114,9 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/openai/endpoint.rs",
|
||||
"tests/openai/responses/http.rs",
|
||||
"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",
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -19,28 +19,55 @@ use super::{api_router, authenticated_client, authenticated_client_with_min_vers
|
||||
|
||||
pub(super) struct ModelsExecution {
|
||||
client: AuthenticatedClient,
|
||||
profiles: bool,
|
||||
presentation: Option<ModelPresentation>,
|
||||
}
|
||||
|
||||
impl ModelsExecution {
|
||||
pub(super) fn new() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
client: authenticated_client("sk_models_test"),
|
||||
profiles: false,
|
||||
presentation: None,
|
||||
})
|
||||
}
|
||||
|
||||
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),
|
||||
})
|
||||
}
|
||||
|
||||
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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -51,7 +78,7 @@ impl ModelsExecution {
|
||||
Some("26.825.51511"),
|
||||
None,
|
||||
),
|
||||
profiles: false,
|
||||
presentation: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -76,32 +103,12 @@ impl ExecutionService for ModelsExecution {
|
||||
}
|
||||
|
||||
fn public_model_profiles(&self, _: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
if !self.profiles {
|
||||
let Some(presentation) = &self.presentation else {
|
||||
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"),
|
||||
]),
|
||||
presentation.clone(),
|
||||
)]
|
||||
}
|
||||
|
||||
@@ -289,6 +296,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 +330,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())
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::io::Read;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::{
|
||||
Arc, Mutex,
|
||||
@@ -11,7 +12,7 @@ use axum::{
|
||||
extract::connect_info::ConnectInfo,
|
||||
http::{
|
||||
HeaderMap, HeaderValue, Request, StatusCode,
|
||||
header::{AUTHORIZATION, CONTENT_TYPE},
|
||||
header::{AUTHORIZATION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE},
|
||||
},
|
||||
};
|
||||
use bytes::Bytes;
|
||||
@@ -21,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,
|
||||
@@ -32,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};
|
||||
@@ -451,9 +452,26 @@ async fn captured_client_context(
|
||||
async fn captured_http_request(
|
||||
provider_name: &str,
|
||||
body: Value,
|
||||
mut headers: HeaderMap,
|
||||
headers: HeaderMap,
|
||||
peer_address: Option<SocketAddr>,
|
||||
) -> CapturedClientContext {
|
||||
let (response, observed) = http_request_with_body(
|
||||
provider_name,
|
||||
body.to_string().into(),
|
||||
headers,
|
||||
peer_address,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
||||
observed.expect("captured request context")
|
||||
}
|
||||
|
||||
async fn http_request_with_body(
|
||||
provider_name: &str,
|
||||
body: Bytes,
|
||||
mut headers: HeaderMap,
|
||||
peer_address: Option<SocketAddr>,
|
||||
) -> (axum::response::Response, Option<CapturedClientContext>) {
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
"Bearer sk_context_test".parse().expect("authorization"),
|
||||
@@ -464,7 +482,7 @@ async fn captured_http_request(
|
||||
client: authenticated_client_for_provider("sk_context_test", provider_name),
|
||||
});
|
||||
let mut request = Request::post("/v1/responses")
|
||||
.body(Body::from(body.to_string()))
|
||||
.body(Body::from(body))
|
||||
.expect("context request");
|
||||
*request.headers_mut() = headers;
|
||||
if let Some(peer_address) = peer_address {
|
||||
@@ -475,12 +493,75 @@ async fn captured_http_request(
|
||||
.oneshot(request)
|
||||
.await
|
||||
.expect("context response");
|
||||
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
||||
observed
|
||||
.lock()
|
||||
.expect("context capture lock")
|
||||
.clone()
|
||||
.expect("captured request context")
|
||||
let observed = observed.lock().expect("context capture lock").clone();
|
||||
(response, observed)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compressed_http_requests_should_preserve_execution_context_without_transport_headers() {
|
||||
let body = json!({
|
||||
"model": "model-a", "input": "hello", "previous_response_id": "resp_previous",
|
||||
"prompt_cache_key": "cache-key", "client_metadata": {"source": "test"}
|
||||
});
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
headers.insert("session_id", HeaderValue::from_static("session-test"));
|
||||
let expected = captured_http_request("openai", body.clone(), headers.clone(), None).await;
|
||||
for encoding in ["gzip", "deflate", "zstd"] {
|
||||
let compressed = super::request::encode_body(encoding, body.to_string().as_bytes());
|
||||
let mut headers = headers.clone();
|
||||
headers.insert(CONTENT_ENCODING, HeaderValue::from_static(encoding));
|
||||
headers.insert(
|
||||
CONTENT_LENGTH,
|
||||
HeaderValue::from_str(&compressed.len().to_string()).expect("compressed length"),
|
||||
);
|
||||
let (response, observed) =
|
||||
http_request_with_body("openai", compressed.into(), headers, None).await;
|
||||
// 捕获执行器在解码成功后主动返回 500;正文和业务头应与未压缩请求一致,
|
||||
// 压缩编码及长度不能泄漏到已解压正文的上游协议上下文。
|
||||
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
||||
assert_eq!(observed.as_ref(), Some(&expected), "{encoding}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_compressed_http_requests_should_fail_before_execution() {
|
||||
for (encoding, body, code, message) in [
|
||||
(
|
||||
"gzip",
|
||||
Bytes::from_static(b"broken gzip"),
|
||||
"invalid_json",
|
||||
"Request body must be valid JSON.",
|
||||
),
|
||||
(
|
||||
"br",
|
||||
Bytes::from_static(b"{}"),
|
||||
"unsupported_content_encoding",
|
||||
"Content-Encoding `br` is not supported.",
|
||||
),
|
||||
(
|
||||
"zstd",
|
||||
super::request::encode_body("zstd", std::io::repeat(0).take(64 * 1024 * 1024 + 1))
|
||||
.into(),
|
||||
"request_too_large",
|
||||
"Decompressed request body exceeds the allowed size.",
|
||||
),
|
||||
] {
|
||||
let headers = HeaderMap::from_iter([
|
||||
(CONTENT_ENCODING, HeaderValue::from_static(encoding)),
|
||||
(CONTENT_TYPE, HeaderValue::from_static("application/json")),
|
||||
]);
|
||||
let (response, observed) = http_request_with_body("openai", body, headers, None).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{encoding}");
|
||||
assert!(observed.is_none(), "invalid body must not start execution");
|
||||
let body = to_bytes(response.into_body(), 4096)
|
||||
.await
|
||||
.expect("error body");
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&body).expect("error JSON"),
|
||||
json!({"error": {"type": "invalid_request_error", "code": code, "message": message}})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn captured_http_compaction(provider_name: &str) -> CapturedClientContext {
|
||||
@@ -1738,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"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
mod http;
|
||||
mod request;
|
||||
mod websocket;
|
||||
|
||||
use axum::http::{HeaderMap, HeaderValue};
|
||||
@@ -369,6 +370,47 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_exclude_downstream_transport_headers_from_opaque_context() {
|
||||
let mut headers = HeaderMap::new();
|
||||
for name in [
|
||||
"cf-visitor",
|
||||
"cf-connecting-ip",
|
||||
"cf-connecting-ipv6",
|
||||
"cf-pseudo-ipv4",
|
||||
"cf-ray",
|
||||
"cf-ipcountry",
|
||||
"cf-warp-tag-id",
|
||||
"cf-worker",
|
||||
"cf-ew-via",
|
||||
"cdn-loop",
|
||||
"via",
|
||||
"forwarded",
|
||||
"x-forwarded-for",
|
||||
"x-forwarded-prefix",
|
||||
"accept-encoding",
|
||||
] {
|
||||
headers.insert(name, HeaderValue::from_static("downstream-only"));
|
||||
}
|
||||
headers.insert(
|
||||
"CF-Visitor",
|
||||
HeaderValue::from_static(r#"{"scheme":"https"}"#),
|
||||
);
|
||||
headers.insert("accept-encoding", HeaderValue::from_static("br, gzip"));
|
||||
// 正文未压缩;入口现在消费 Content-Encoding,但仍不得向上游透传。
|
||||
headers.insert("content-encoding", HeaderValue::from_static("identity"));
|
||||
headers.insert("x-openai-future-mode", HeaderValue::from_static("keep"));
|
||||
|
||||
let decoded =
|
||||
decode_request_with_headers(br#"{"model":"smart-code","input":"hello"}"#, &headers)
|
||||
.expect("decode request behind a reverse proxy");
|
||||
|
||||
assert_eq!(
|
||||
openai_protocol_context(&decoded).get("opaque_request_headers"),
|
||||
Some(&json!([["x-openai-future-mode", STANDARD.encode(b"keep")]])),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_preserve_unknown_nested_values_without_debug_disclosure() {
|
||||
let secret = "nested-private-value";
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
use std::io::{Read, Write, copy, repeat};
|
||||
|
||||
use axum::http::{HeaderMap, HeaderValue, header::CONTENT_ENCODING};
|
||||
use gateway_api::openai::responses::{RequestDecodeError, decode_request_with_headers};
|
||||
use serde_json::json;
|
||||
|
||||
use super::openai_wire_body;
|
||||
|
||||
const ENCODINGS: [&str; 3] = ["gzip", "deflate", "zstd"];
|
||||
const REQUEST: &[u8] = br#"{"model":"gpt-test","input":"hello","future_field":{"value":1}}"#;
|
||||
const DECOMPRESSED_LIMIT: u64 = 64 * 1024 * 1024;
|
||||
|
||||
pub(super) fn encode_body(encoding: &str, mut data: impl Read) -> Vec<u8> {
|
||||
match encoding {
|
||||
"gzip" => {
|
||||
let mut encoder =
|
||||
flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
|
||||
copy(&mut data, &mut encoder).expect("gzip encode");
|
||||
encoder.finish().expect("gzip finish")
|
||||
}
|
||||
"deflate" => {
|
||||
let mut encoder =
|
||||
flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::fast());
|
||||
copy(&mut data, &mut encoder).expect("deflate encode");
|
||||
encoder.finish().expect("deflate finish")
|
||||
}
|
||||
"zstd" => zstd::stream::encode_all(data, 3).expect("zstd encode"),
|
||||
_ => panic!("unsupported fixture encoding"),
|
||||
}
|
||||
}
|
||||
|
||||
fn headers(encoding: &str) -> HeaderMap {
|
||||
HeaderMap::from_iter([(
|
||||
CONTENT_ENCODING,
|
||||
HeaderValue::from_str(encoding).expect("encoding header"),
|
||||
)])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_preserve_compressed_request_fields() {
|
||||
let expected = serde_json::from_slice::<serde_json::Value>(REQUEST).expect("fixture JSON");
|
||||
for encoding in ENCODINGS {
|
||||
let compressed = encode_body(encoding, REQUEST);
|
||||
let decoded = decode_request_with_headers(&compressed, &headers(encoding)).expect(encoding);
|
||||
assert_eq!(
|
||||
openai_wire_body(&decoded),
|
||||
expected.as_object().expect("object")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_accept_plain_identity_and_case_insensitive_encodings() {
|
||||
for request_headers in [HeaderMap::new(), headers(""), headers(" identity ")] {
|
||||
assert!(decode_request_with_headers(REQUEST, &request_headers).is_ok());
|
||||
}
|
||||
let compressed = encode_body("gzip", REQUEST);
|
||||
assert!(decode_request_with_headers(&compressed, &headers(" GZip ")).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_reject_unknown_and_stacked_encodings() {
|
||||
let mut repeated = headers("identity");
|
||||
repeated.append(CONTENT_ENCODING, HeaderValue::from_static("gzip"));
|
||||
for request_headers in [headers("br"), headers("gzip, zstd"), repeated] {
|
||||
let error = decode_request_with_headers(REQUEST, &request_headers)
|
||||
.expect_err("unsupported encoding");
|
||||
assert!(matches!(
|
||||
error,
|
||||
RequestDecodeError::UnsupportedContentEncoding { .. }
|
||||
));
|
||||
assert_eq!(
|
||||
error.protocol_body().into_value()["error"]["code"],
|
||||
"unsupported_content_encoding"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_reject_corrupted_and_truncated_compressed_bodies() {
|
||||
for encoding in ENCODINGS {
|
||||
let mut truncated = encode_body(encoding, REQUEST);
|
||||
truncated.truncate(truncated.len() - 4);
|
||||
for body in [b"not-a-compressed-stream".as_slice(), truncated.as_slice()] {
|
||||
let error = decode_request_with_headers(body, &headers(encoding)).expect_err(encoding);
|
||||
assert_eq!(
|
||||
error.protocol_body().into_value(),
|
||||
json!({"error": {
|
||||
"type": "invalid_request_error", "code": "invalid_json",
|
||||
"message": "Request body must be valid JSON."
|
||||
}})
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_read_all_gzip_members_and_zstd_frames() {
|
||||
for encoding in ["gzip", "zstd"] {
|
||||
let mut compressed = encode_body(encoding, &REQUEST[..12]);
|
||||
compressed.extend(encode_body(encoding, &REQUEST[12..]));
|
||||
let decoded = decode_request_with_headers(&compressed, &headers(encoding)).expect(encoding);
|
||||
assert_eq!(openai_wire_body(&decoded)["input"], "hello");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_not_ignore_invalid_trailing_members() {
|
||||
for encoding in ["gzip", "zstd"] {
|
||||
let mut compressed = encode_body(encoding, REQUEST);
|
||||
compressed.extend(encode_body(encoding, b"invalid trailing JSON".as_slice()));
|
||||
assert_eq!(
|
||||
decode_request_with_headers(&compressed, &headers(encoding)).expect_err(encoding),
|
||||
RequestDecodeError::MalformedJson
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_reject_zstd_frames_with_an_excessive_window() {
|
||||
let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), 3).expect("zstd encoder");
|
||||
encoder.window_log(27).expect("128 MiB window");
|
||||
encoder.write_all(REQUEST).expect("write JSON");
|
||||
// 提前输出帧头,避免编码器根据最终小正文缩减窗口。
|
||||
encoder.flush().expect("flush frame header");
|
||||
let compressed = encoder.finish().expect("finish frame");
|
||||
assert_eq!(
|
||||
zstd::stream::decode_all(compressed.as_slice()).expect("valid zstd frame"),
|
||||
REQUEST
|
||||
);
|
||||
assert_eq!(
|
||||
decode_request_with_headers(&compressed, &headers("zstd")).expect_err("oversized window"),
|
||||
RequestDecodeError::MalformedJson
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_accept_the_exact_decompressed_limit() {
|
||||
let padding = repeat(b' ').take(DECOMPRESSED_LIMIT - REQUEST.len() as u64);
|
||||
let compressed = encode_body("zstd", REQUEST.chain(padding));
|
||||
let decoded = decode_request_with_headers(&compressed, &headers("zstd")).expect("exact limit");
|
||||
assert_eq!(decoded.metadata().requested_model(), "gpt-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_reject_bodies_above_the_decompressed_limit() {
|
||||
for encoding in ENCODINGS {
|
||||
let compressed = encode_body(encoding, repeat(0).take(DECOMPRESSED_LIMIT + 1));
|
||||
let error =
|
||||
decode_request_with_headers(&compressed, &headers(encoding)).expect_err(encoding);
|
||||
assert_eq!(
|
||||
error,
|
||||
RequestDecodeError::DecompressedBodyTooLarge,
|
||||
"{encoding}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_decode_should_stop_at_limit_before_reading_invalid_trailing_frames() {
|
||||
for encoding in ["zstd", "gzip"] {
|
||||
// 尾部故意损坏:若先完整解压再检查长度,就会错误地返回 invalid_json,
|
||||
// 这个断言验证越界时已经停止读取,而不只是最终返回了某个超限错误。
|
||||
let mut compressed = encode_body(encoding, repeat(0).take(DECOMPRESSED_LIMIT + 1));
|
||||
compressed.extend_from_slice(b"invalid trailing frame");
|
||||
let error =
|
||||
decode_request_with_headers(&compressed, &headers(encoding)).expect_err(encoding);
|
||||
assert_eq!(
|
||||
error,
|
||||
RequestDecodeError::DecompressedBodyTooLarge,
|
||||
"{encoding}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -1,8 +1,53 @@
|
||||
use axum::http::{HeaderMap, HeaderValue};
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use gateway_api::openai::responses::{OpenAiRequestHeaders, decode_response_create_with_context};
|
||||
use gateway_core::operation::Operation;
|
||||
use serde_json::json;
|
||||
|
||||
use super::decode_response_create;
|
||||
|
||||
#[test]
|
||||
fn response_create_should_exclude_cloudflare_headers_on_every_frame() {
|
||||
let mut headers = HeaderMap::new();
|
||||
for (name, value) in [
|
||||
("CF-Visitor", r#"{"scheme":"https"}"#),
|
||||
("cdn-loop", "cloudflare; loops=1"),
|
||||
("cf-warp-tag-id", "downstream-warp"),
|
||||
("cf-ipcountry", "US"),
|
||||
("cf-worker", "worker.example"),
|
||||
("cf-ew-via", "15"),
|
||||
("cf-future-proxy-field", "downstream-only"),
|
||||
] {
|
||||
headers.insert(name, HeaderValue::from_static(value));
|
||||
}
|
||||
headers.append("x-openai-future-mode", HeaderValue::from_static("first"));
|
||||
headers.append("x-openai-future-mode", HeaderValue::from_static("second"));
|
||||
let request_headers = OpenAiRequestHeaders::from_headers(&headers);
|
||||
|
||||
// 连接级头会用于同一 WebSocket 的每一帧,后续帧也不能恢复下游链路元数据。
|
||||
for input in ["hello", "continue"] {
|
||||
let decoded = decode_response_create_with_context(
|
||||
&json!({"type": "response.create", "model": "smart-code", "input": input}).to_string(),
|
||||
&request_headers,
|
||||
)
|
||||
.expect("decode response.create behind Cloudflare");
|
||||
let Operation::Generate(request) = decoded.operation() else {
|
||||
panic!("Responses must map to Generate");
|
||||
};
|
||||
assert_eq!(
|
||||
request
|
||||
.protocol_payload()
|
||||
.context()
|
||||
.get("opaque_request_headers"),
|
||||
Some(&json!([
|
||||
["x-openai-future-mode", STANDARD.encode(b"first")],
|
||||
["x-openai-future-mode", STANDARD.encode(b"second")],
|
||||
])),
|
||||
);
|
||||
assert_eq!(request.protocol_payload().body()["input"], input);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn latest_official_codex_response_create_fixture_decodes_unchanged() {
|
||||
// Audited against openai/codex main 94cbbddafc1776d5e377bca1b05932c697e82238.
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
|
||||
@@ -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,
|
||||
@@ -199,7 +199,7 @@ where
|
||||
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 +209,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)
|
||||
@@ -230,11 +230,22 @@ struct CurrentAttempt {
|
||||
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>,
|
||||
@@ -293,7 +304,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 +321,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 +339,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 +350,7 @@ where
|
||||
.await?;
|
||||
return Err(provider_engine_error(pending.error));
|
||||
}
|
||||
if self.finalized {
|
||||
if self.is_finalized() {
|
||||
return Ok(None);
|
||||
}
|
||||
loop {
|
||||
@@ -383,10 +400,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 +440,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 +497,7 @@ where
|
||||
|
||||
#[must_use]
|
||||
pub const fn is_finalized(&self) -> bool {
|
||||
self.finalized
|
||||
matches!(self.finalization, Some(RequestFinalization::Complete))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
@@ -573,14 +591,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 +606,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 +637,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))) => {
|
||||
@@ -679,11 +703,11 @@ where
|
||||
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 +722,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,
|
||||
@@ -795,14 +820,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 +886,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 +897,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 +920,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 +930,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 +943,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 +953,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 +969,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 +979,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 +1009,20 @@ 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,
|
||||
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 +1049,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)
|
||||
}
|
||||
|
||||
@@ -1220,15 +1254,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 +1406,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 +1429,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 +1487,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 +1503,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 +1550,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 +1563,7 @@ where
|
||||
&mut self,
|
||||
finalization: FailureFinalization,
|
||||
) -> Result<(), EngineError> {
|
||||
if self.finalized {
|
||||
if self.is_finalized() {
|
||||
return Ok(());
|
||||
}
|
||||
self.trace.record(
|
||||
@@ -1555,28 +1575,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 +1590,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 +1948,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()));
|
||||
|
||||
@@ -173,8 +173,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, ()>;
|
||||
}
|
||||
|
||||
@@ -978,8 +981,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 +1000,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 +1011,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 +1049,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 +1057,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 +1089,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 +1101,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 +1113,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 +1125,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 +1133,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]
|
||||
@@ -562,9 +505,14 @@ pub trait Provider: Send + Sync {
|
||||
|
||||
/// 选择一个未被排除的资源并返回 cold [`ProviderStream`]。
|
||||
///
|
||||
/// 返回成功、返回错误或准备 future 被取消前,均不得发送本次请求的上游握手或业务
|
||||
/// 载荷,也不得启动可独立完成这些发送的后台任务。发送只在返回的流被 poll 后发生,
|
||||
/// 保证 Core 能在真实出站前校验账号范围并登记 attempt。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 没有可用资源、请求无效或在产生 stream 前失败时返回 Provider 错误。
|
||||
/// 没有可用资源、请求无效或准备失败时返回 `NotSent` 错误;
|
||||
/// 可能已发送的失败必须通过 stream 返回,不得降级发送事实。
|
||||
async fn execute(
|
||||
&self,
|
||||
request: ProviderRequest,
|
||||
@@ -586,15 +534,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 +604,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 +618,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 +636,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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
//! 快照编译需要的 Provider 目录事实,不暴露执行或账号选择能力。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::identity::ProviderKind;
|
||||
|
||||
use super::{ModelCapabilities, ModelPresentation, UpstreamModelId};
|
||||
|
||||
/// 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,
|
||||
};
|
||||
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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,12 +2,16 @@
|
||||
|
||||
use std::collections::{BTreeSet, VecDeque};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use futures::StreamExt;
|
||||
use futures::channel::oneshot;
|
||||
use futures::executor::block_on;
|
||||
|
||||
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
@@ -16,8 +20,7 @@ use gateway_core::engine::continuation::{
|
||||
};
|
||||
use gateway_core::engine::execution::gateway_error_from_engine;
|
||||
use gateway_core::engine::provider::{
|
||||
Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities,
|
||||
ProviderRegistry, ProviderRequest, ProviderStream,
|
||||
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AttemptContext, AttemptCoordinator, AttemptRecord, AttemptTransport, CommitRequirement,
|
||||
@@ -41,14 +44,16 @@ use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, Provi
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{
|
||||
AccountRoutingSnapshot, ClientRoutingScope, ConfigRevision, FrozenAccountScope,
|
||||
ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
|
||||
RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
|
||||
ModelCapabilities, ProviderCatalogGeneration, ProviderKind, ProviderModel,
|
||||
ProviderModelCapabilities, PublicModelId, RoutingContext, RoutingPlan, RuntimeAccount,
|
||||
RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct FinalState {
|
||||
error_kind: Option<GatewayErrorKind>,
|
||||
diagnostic_trace_json: Option<String>,
|
||||
outcome: ExecutionOutcome,
|
||||
send_state: UpstreamSendState,
|
||||
@@ -101,6 +106,8 @@ struct StoreState {
|
||||
struct FakeStore {
|
||||
state: Mutex<StoreState>,
|
||||
failures: BTreeSet<StoreWriteFailure>,
|
||||
create_gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
finalize_gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
@@ -113,8 +120,8 @@ enum StoreWriteFailure {
|
||||
impl FakeStore {
|
||||
fn failing(failure: StoreWriteFailure) -> Self {
|
||||
Self {
|
||||
state: Mutex::new(StoreState::default()),
|
||||
failures: BTreeSet::from([failure]),
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -132,6 +139,10 @@ impl ExecutionStore for FakeStore {
|
||||
async fn create_model_request(&self, _request: NewModelRequest) -> Result<(), StoreError> {
|
||||
self.fails(StoreWriteFailure::Create)?;
|
||||
self.state.lock().expect("store lock").created += 1;
|
||||
let gate = self.create_gate.lock().expect("create gate lock").take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("create released");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -208,6 +219,7 @@ impl ExecutionStore for FakeStore {
|
||||
.expect("store lock")
|
||||
.finalizations
|
||||
.push(FinalState {
|
||||
error_kind: finalization.error.as_ref().map(|error| error.kind()),
|
||||
diagnostic_trace_json: finalization.diagnostic_trace_json,
|
||||
outcome: finalization.outcome,
|
||||
send_state: finalization.send_state,
|
||||
@@ -244,6 +256,14 @@ impl ExecutionStore for FakeStore {
|
||||
.total()
|
||||
.map(|total| total.amount().scaled()),
|
||||
});
|
||||
let gate = self
|
||||
.finalize_gate
|
||||
.lock()
|
||||
.expect("finalize gate lock")
|
||||
.take();
|
||||
if let Some(gate) = gate {
|
||||
gate.await.expect("finalize released");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -267,12 +287,22 @@ enum Script {
|
||||
items: Vec<Result<ProviderEvent, ProviderError>>,
|
||||
},
|
||||
Error(ProviderError),
|
||||
Pending,
|
||||
}
|
||||
|
||||
struct ScriptedProvider {
|
||||
scripts: Mutex<VecDeque<Script>>,
|
||||
contexts: Mutex<Vec<AttemptContext>>,
|
||||
operations: Mutex<Vec<Operation>>,
|
||||
released_leases: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
struct TrackedLease(Arc<AtomicUsize>);
|
||||
|
||||
impl Drop for TrackedLease {
|
||||
fn drop(&mut self) {
|
||||
self.0.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
impl ScriptedProvider {
|
||||
@@ -281,6 +311,7 @@ impl ScriptedProvider {
|
||||
scripts: Mutex::new(scripts.into()),
|
||||
contexts: Mutex::new(Vec::new()),
|
||||
operations: Mutex::new(Vec::new()),
|
||||
released_leases: Arc::new(AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -319,6 +350,7 @@ impl Provider for ScriptedProvider {
|
||||
.expect("one script per provider call");
|
||||
match script {
|
||||
Script::Error(error) => Err(error),
|
||||
Script::Pending => futures::future::pending().await,
|
||||
Script::Stream { account_id, items } => {
|
||||
let candidate = request.candidate();
|
||||
let metadata = ProviderCallMetadata::new(
|
||||
@@ -336,7 +368,7 @@ impl Provider for ScriptedProvider {
|
||||
Box::pin(futures::stream::iter(
|
||||
items.into_iter().map(canonical_provider_event),
|
||||
)),
|
||||
(),
|
||||
TrackedLease(Arc::clone(&self.released_leases)),
|
||||
))
|
||||
}
|
||||
Script::HangingStream { account_id, items } => {
|
||||
@@ -997,6 +1029,50 @@ fn failed_image_request_should_persist_failure() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_request_id_prefers_error_then_response_observation() {
|
||||
for error_request_id in [None, Some("upstream-error")] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let observation = ProviderResponseObservation::new(
|
||||
UpstreamTransport::new("http_sse").expect("transport"),
|
||||
)
|
||||
.with_status_code(200)
|
||||
.with_request_id(OpaqueUpstreamValue::new("upstream-opening"));
|
||||
let mut error =
|
||||
ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
|
||||
.with_status(400);
|
||||
if let Some(request_id) = error_request_id {
|
||||
error = error.with_upstream_request_id(OpaqueUpstreamValue::new(request_id));
|
||||
}
|
||||
let (coordinator, store, _) = coordinator(vec![Script::ObservedStream {
|
||||
account_id: "acct_observed",
|
||||
items: vec![Ok(ProviderEvent::observation(observation)), Err(error)],
|
||||
}]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
|
||||
assert!(matches!(
|
||||
block_on(session.collect_uncommitted()),
|
||||
Err(EngineError::Provider(error)) if error.kind() == ProviderErrorKind::InvalidRequest
|
||||
));
|
||||
let state = store.state.lock().expect("store lock");
|
||||
let finalization = &state.finalizations[0];
|
||||
assert_eq!(
|
||||
finalization.upstream_request_id.as_deref(),
|
||||
Some(error_request_id.unwrap_or("upstream-opening"))
|
||||
);
|
||||
assert_eq!(finalization.upstream_status_code, Some(400));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_observation_is_persisted_but_never_delivered() {
|
||||
let operation = generate_operation();
|
||||
@@ -1292,6 +1368,8 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
|
||||
Ok(ProviderEvent::observation(first_observation)),
|
||||
Err(
|
||||
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent)
|
||||
.with_status(504)
|
||||
.with_upstream_request_id(OpaqueUpstreamValue::new("discarded-error"))
|
||||
.with_replay_safe()
|
||||
.with_diagnostic(
|
||||
gateway_core::error::ProviderDiagnostic::new(
|
||||
@@ -1363,10 +1441,10 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
|
||||
.first_event_ms
|
||||
.is_some_and(|elapsed| elapsed < 987_654)
|
||||
);
|
||||
assert_eq!(state.intermediate_status_codes, vec![Some(503)]);
|
||||
assert_eq!(state.intermediate_status_codes, vec![Some(504)]);
|
||||
assert_eq!(
|
||||
state.intermediate_request_ids,
|
||||
vec![Some("discarded-request".to_owned())]
|
||||
vec![Some("discarded-error".to_owned())]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1933,7 +2011,17 @@ fn provider_metadata_for_another_account_fails_closed() {
|
||||
assert_eq!(provider.scripts.lock().expect("scripts lock").len(), 1);
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert!(state.attempts.is_empty());
|
||||
assert!(state.finalizations.is_empty());
|
||||
assert_eq!(state.created, 1);
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 0);
|
||||
assert_eq!(
|
||||
state.finalizations[0].send_state,
|
||||
UpstreamSendState::NotSent
|
||||
);
|
||||
assert_eq!(
|
||||
state.finalizations[0].error_kind,
|
||||
Some(GatewayErrorKind::Internal)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -3498,7 +3586,7 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_eligible_account_before_stream_does_not_create_request_detail() {
|
||||
fn no_eligible_account_before_stream_records_failure_without_fabricating_an_attempt() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, store, _) = coordinator(vec![Script::Error(ProviderError::new(
|
||||
@@ -3524,10 +3612,35 @@ fn no_eligible_account_before_stream_does_not_create_request_detail() {
|
||||
if error.kind() == ProviderErrorKind::NoEligibleAccount
|
||||
));
|
||||
assert!(session.provider_attempt_outcomes().is_empty());
|
||||
assert!(session.is_finalized());
|
||||
block_on(session.record_client_status(503)).expect("HTTP error status");
|
||||
block_on(session.cancel_and_finalize()).expect("repeated finalization");
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.created, 0);
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert!(state.finalizations.is_empty());
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
assert_eq!(state.recorded_statuses, [503]);
|
||||
let finalization = &state.finalizations[0];
|
||||
assert_eq!(finalization.outcome, ExecutionOutcome::Failed);
|
||||
assert_eq!(finalization.attempt_count, 0);
|
||||
assert_eq!(finalization.send_state, UpstreamSendState::NotSent);
|
||||
assert_eq!(
|
||||
finalization.error_kind,
|
||||
Some(GatewayErrorKind::NoAvailableProvider)
|
||||
);
|
||||
assert_eq!(finalization.upstream_transport, None);
|
||||
assert_eq!(finalization.total_tokens, None);
|
||||
assert_eq!(finalization.cost_ticks, None);
|
||||
assert_eq!(session.budget_charge().amount_usd.scaled(), 0);
|
||||
let trace: Value =
|
||||
serde_json::from_str(finalization.diagnostic_trace_json.as_deref().unwrap()).unwrap();
|
||||
let events = trace["events"].as_array().unwrap();
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|event| event["stage"] == "attempt.failed")
|
||||
);
|
||||
assert_eq!(events.last().unwrap()["stage"], "request.finished");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -3535,9 +3648,12 @@ fn expired_deadline_finalizes_without_calling_provider() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, store, provider) = coordinator(vec![]);
|
||||
let deadline = SystemTime::now() - Duration::from_secs(1);
|
||||
let mut request = model_request(&operation, deadline);
|
||||
request.started_at = deadline - Duration::from_secs(30);
|
||||
|
||||
let error = match block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::UNIX_EPOCH),
|
||||
request,
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
@@ -3551,11 +3667,485 @@ fn expired_deadline_finalizes_without_calling_provider() {
|
||||
assert!(matches!(error, gateway_core::engine::EngineError::Deadline));
|
||||
assert!(provider.contexts.lock().expect("contexts lock").is_empty());
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.created, 0);
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 0);
|
||||
assert_eq!(
|
||||
state.finalizations[0].send_state,
|
||||
UpstreamSendState::NotSent
|
||||
);
|
||||
assert_eq!(
|
||||
state.finalizations[0].error_kind,
|
||||
Some(GatewayErrorKind::Timeout)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preparation_failures_preserve_the_error_when_observation_writes_fail() {
|
||||
for failure in [
|
||||
None,
|
||||
Some(StoreWriteFailure::Create),
|
||||
Some(StoreWriteFailure::Finalize),
|
||||
] {
|
||||
let operation = generate_operation();
|
||||
let store = Arc::new(failure.map_or_else(FakeStore::default, FakeStore::failing));
|
||||
let (coordinator, store, provider) = coordinator_with_store(
|
||||
vec![Script::Error(ProviderError::new(
|
||||
ProviderErrorKind::ProviderInfrastructureUnavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
))],
|
||||
store,
|
||||
);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
block_on(session.next_event()),
|
||||
Err(EngineError::Provider(ref error))
|
||||
if error.kind() == ProviderErrorKind::ProviderInfrastructureUnavailable
|
||||
));
|
||||
assert!(session.is_finalized());
|
||||
let charge = session.budget_charge();
|
||||
block_on(session.cancel_and_finalize()).unwrap();
|
||||
assert!(block_on(session.next_event()).unwrap().is_none());
|
||||
assert_eq!(session.budget_charge().completed_at, charge.completed_at);
|
||||
assert_eq!(charge.amount_usd.scaled(), 0);
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
let state = store.state.lock().unwrap();
|
||||
assert!(state.attempts.is_empty());
|
||||
assert_eq!(
|
||||
state.created,
|
||||
usize::from(failure != Some(StoreWriteFailure::Create))
|
||||
);
|
||||
assert_eq!(state.finalizations.len(), usize::from(failure.is_none()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn interrupted_early_finalization_resumes_the_same_write_and_keeps_the_original_failure() {
|
||||
block_on(async {
|
||||
for create in [true, false] {
|
||||
for detach in [true, false] {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let store = Arc::new(if create {
|
||||
FakeStore {
|
||||
create_gate: Mutex::new(Some(gate)),
|
||||
..FakeStore::default()
|
||||
}
|
||||
} else {
|
||||
FakeStore {
|
||||
finalize_gate: Mutex::new(Some(gate)),
|
||||
..FakeStore::default()
|
||||
}
|
||||
});
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, provider) = coordinator_with_store(
|
||||
vec![Script::Error(ProviderError::new(
|
||||
ProviderErrorKind::NoEligibleAccount,
|
||||
UpstreamSendState::NotSent,
|
||||
))],
|
||||
store,
|
||||
);
|
||||
let mut session = coordinator
|
||||
.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
for _ in 0..2 {
|
||||
let mut next = Box::pin(session.next_event());
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
drop(next);
|
||||
assert!(!session.is_finalized());
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
}
|
||||
let completed_at = session.budget_charge().completed_at;
|
||||
release
|
||||
.send(())
|
||||
.expect("original write must still be alive");
|
||||
if detach {
|
||||
session.cancel_and_finalize().await.unwrap();
|
||||
} else {
|
||||
assert!(session.next_event().await.unwrap().is_none());
|
||||
}
|
||||
assert!(session.is_finalized());
|
||||
session.cancel_and_finalize().await.unwrap();
|
||||
assert_eq!(session.budget_charge().completed_at, completed_at);
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
|
||||
assert_eq!(
|
||||
state.finalizations[0].error_kind,
|
||||
Some(GatewayErrorKind::NoAvailableProvider)
|
||||
);
|
||||
let trace: Value = serde_json::from_str(
|
||||
state.finalizations[0]
|
||||
.diagnostic_trace_json
|
||||
.as_deref()
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
trace["events"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter(|event| event["stage"] == "request.finished")
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancellation_before_or_during_preparation_records_a_zero_attempt_terminal_request() {
|
||||
block_on(async {
|
||||
for before_start in [true, false] {
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, provider) = coordinator(vec![Script::Pending]);
|
||||
let cancellation = CancellationToken::new();
|
||||
if before_start {
|
||||
cancellation.cancel();
|
||||
}
|
||||
let result = coordinator
|
||||
.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
cancellation.clone(),
|
||||
)
|
||||
.await;
|
||||
if before_start {
|
||||
assert!(matches!(result, Err(EngineError::Cancelled)));
|
||||
assert!(provider.contexts.lock().unwrap().is_empty());
|
||||
} else {
|
||||
let mut session = result.unwrap();
|
||||
let mut next = Box::pin(session.next_event());
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
cancellation.cancel();
|
||||
assert!(matches!(next.await, Err(EngineError::Cancelled)));
|
||||
assert!(session.is_finalized());
|
||||
session.cancel_and_finalize().await.unwrap();
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
}
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled);
|
||||
assert_eq!(
|
||||
state.finalizations[0].send_state,
|
||||
UpstreamSendState::NotSent
|
||||
);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 0);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deadline_during_preparation_records_failure_without_an_upstream_attempt() {
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, provider) = coordinator(vec![Script::Pending]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_millis(300)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
block_on(session.next_event()),
|
||||
Err(EngineError::Deadline)
|
||||
));
|
||||
assert!(session.provider_attempt_outcomes().is_empty());
|
||||
assert!(session.is_finalized());
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
assert_eq!(
|
||||
state.finalizations[0].error_kind,
|
||||
Some(GatewayErrorKind::Timeout)
|
||||
);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 0);
|
||||
assert_eq!(
|
||||
state.finalizations[0].send_state,
|
||||
UpstreamSendState::NotSent
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_stream_send_facts_cannot_be_rewritten_as_a_zero_attempt_not_sent_request() {
|
||||
for send_state in [UpstreamSendState::Sent, UpstreamSendState::Ambiguous] {
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, _) = coordinator(vec![Script::Error(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
send_state,
|
||||
))]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
block_on(session.next_event()),
|
||||
Err(EngineError::Provider(ref error)) if error.send_state() == send_state
|
||||
));
|
||||
assert!(session.is_finalized());
|
||||
block_on(session.cancel_and_finalize()).unwrap();
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 0);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert!(state.finalizations.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_first_attempt_write_is_not_replaced_with_a_fabricated_zero_attempt_record() {
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, _) = coordinator_with_store(
|
||||
vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
UpstreamSendState::Sent,
|
||||
))],
|
||||
}],
|
||||
Arc::new(FakeStore::failing(StoreWriteFailure::Create)),
|
||||
);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
block_on(session.next_event()),
|
||||
Err(EngineError::Provider(_))
|
||||
));
|
||||
assert!(session.is_finalized());
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 0);
|
||||
assert!(state.finalizations.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn interrupted_first_attempt_write_never_restarts_selection_or_creates_a_zero_attempt_request() {
|
||||
block_on(async {
|
||||
for cancel in [true, false] {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, provider) = coordinator_with_store(
|
||||
vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
}],
|
||||
Arc::new(FakeStore {
|
||||
create_gate: Mutex::new(Some(gate)),
|
||||
..FakeStore::default()
|
||||
}),
|
||||
);
|
||||
let mut session = coordinator
|
||||
.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut next = Box::pin(session.next_event());
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
drop(next);
|
||||
assert_eq!(store.state.lock().unwrap().created, 1);
|
||||
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 0);
|
||||
// 首写属于可丢弃观测;丢失返回确认后不臆测入库成功,也不能退回零次重新创建。
|
||||
assert!(release.send(()).is_err());
|
||||
if cancel {
|
||||
session.cancel_and_finalize().await.unwrap();
|
||||
} else {
|
||||
session.collect_uncommitted().await.unwrap();
|
||||
session.commit_downstream(Some(200)).await.unwrap();
|
||||
}
|
||||
assert!(session.is_finalized());
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
let trace = session.trace().snapshot().unwrap();
|
||||
let events = trace["events"].as_array().unwrap();
|
||||
assert_eq!(events.last().unwrap()["data"]["attempts"], 1);
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.finalizations.is_empty());
|
||||
drop(state);
|
||||
drop(session);
|
||||
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejected_cold_stream_releases_its_lease_before_waiting_for_observation_writes() {
|
||||
block_on(async {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, provider) = coordinator_with_store(
|
||||
vec![Script::Stream {
|
||||
account_id: "acct_wrong",
|
||||
items: complete_stream(None),
|
||||
}],
|
||||
Arc::new(FakeStore {
|
||||
create_gate: Mutex::new(Some(gate)),
|
||||
..FakeStore::default()
|
||||
}),
|
||||
);
|
||||
let mut session = coordinator
|
||||
.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
Some(ProviderAccountId::new("acct_required").unwrap()),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut next = Box::pin(session.next_event());
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 1);
|
||||
release.send(()).unwrap();
|
||||
assert!(matches!(
|
||||
next.await,
|
||||
Err(EngineError::RequiredAccountMismatch)
|
||||
));
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 1);
|
||||
assert!(state.attempts.is_empty());
|
||||
assert_eq!(state.finalizations[0].attempt_count, 0);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn charged_terminal_write_resumes_without_replacing_success_failure_or_send_facts() {
|
||||
block_on(async {
|
||||
for failed in [true, false] {
|
||||
for cancel in [true, false] {
|
||||
let (release, gate) = oneshot::channel();
|
||||
let operation = generate_operation();
|
||||
let mut items = vec![
|
||||
Ok(GatewayEvent::Started(ResponseMeta::new(
|
||||
"charged-response",
|
||||
"gpt-5",
|
||||
))),
|
||||
Ok(GatewayEvent::ProviderCost(
|
||||
ProviderReportedCost::from_usd_ticks(25).unwrap(),
|
||||
)),
|
||||
];
|
||||
items.push(if failed {
|
||||
Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::Sent,
|
||||
))
|
||||
} else {
|
||||
Ok(GatewayEvent::Completed(ResponseMeta::new(
|
||||
"charged-response",
|
||||
"gpt-5",
|
||||
)))
|
||||
});
|
||||
let (coordinator, store, provider) = coordinator_with_store(
|
||||
vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items,
|
||||
}],
|
||||
Arc::new(FakeStore {
|
||||
finalize_gate: Mutex::new(Some(gate)),
|
||||
..FakeStore::default()
|
||||
}),
|
||||
);
|
||||
let mut session = coordinator
|
||||
.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
if failed {
|
||||
let mut collect = Box::pin(session.collect_uncommitted());
|
||||
assert!(futures::poll!(collect.as_mut()).is_pending());
|
||||
} else {
|
||||
session.collect_uncommitted().await.unwrap();
|
||||
let mut commit = Box::pin(session.commit_downstream(Some(200)));
|
||||
assert!(futures::poll!(commit.as_mut()).is_pending());
|
||||
}
|
||||
assert!(!session.is_finalized());
|
||||
let charge = session.budget_charge();
|
||||
assert_eq!(charge.amount_usd.scaled(), 25);
|
||||
release
|
||||
.send(())
|
||||
.expect("the original finalization still owns its gate");
|
||||
if cancel {
|
||||
session.cancel_and_finalize().await.unwrap();
|
||||
} else {
|
||||
assert!(session.next_event().await.unwrap().is_none());
|
||||
}
|
||||
assert!(session.is_finalized());
|
||||
assert_eq!(session.budget_charge().amount_usd, charge.amount_usd);
|
||||
assert_eq!(session.budget_charge().completed_at, charge.completed_at);
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.created, 1);
|
||||
assert_eq!(state.attempts.len(), 1);
|
||||
assert_eq!(state.finalizations.len(), 1);
|
||||
let finalization = &state.finalizations[0];
|
||||
assert_eq!(finalization.attempt_count, 1);
|
||||
assert_eq!(finalization.send_state, UpstreamSendState::Sent);
|
||||
assert_eq!(finalization.cost_ticks, Some(25));
|
||||
assert_eq!(
|
||||
finalization.outcome,
|
||||
if failed {
|
||||
ExecutionOutcome::Failed
|
||||
} else {
|
||||
ExecutionOutcome::Succeeded
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deadline_after_commit_is_incomplete_without_provider_circuit_failure() {
|
||||
let operation = generate_operation();
|
||||
|
||||
@@ -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};
|
||||
@@ -821,9 +1394,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 +1417,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 +1458,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 +1516,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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user