mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-17 08:04:00 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
073e92d171 | ||
|
|
4d9fedee20 | ||
|
|
fce41e318f | ||
|
|
8ae6d8f67e | ||
|
|
8219dcfc87 | ||
|
|
1d3b9665c8 | ||
|
|
c204d33b09 | ||
|
|
45dcce0e49 | ||
|
|
89d826be29 | ||
|
|
cb7b03795d | ||
|
|
a8b9ea22b7 | ||
|
|
9662cff2e7 | ||
|
|
fbfdcef818 | ||
|
|
0e82efe489 | ||
|
|
e803e3851c | ||
|
|
5912ae960c | ||
|
|
e77b4b260e | ||
|
|
e215c98c2c | ||
|
|
e29b93a1fb | ||
|
|
fd82dfd52d | ||
|
|
4853662f99 | ||
|
|
0fa577a19c | ||
|
|
b61e4bcc45 | ||
|
|
0ae151a238 | ||
|
|
b830bc14d6 | ||
|
|
8c4c3c09ce | ||
|
|
678eb22a40 | ||
|
|
c4d883b8da | ||
|
|
0de6d7e9ba | ||
|
|
363cc4994b | ||
|
|
f3d9491071 | ||
|
|
69648476d4 | ||
|
|
a04ce49016 | ||
|
|
bb9e74285e | ||
|
|
bba6a55e0f | ||
|
|
5935e674a8 | ||
|
|
ef4f99f292 | ||
|
|
93c32fa1a2 | ||
|
|
04f8cdb194 | ||
|
|
c0ab3a00ea | ||
|
|
814ecfba7c | ||
|
|
bd404c16f2 | ||
|
|
4ec9ceec4a | ||
|
|
46cbb7187b | ||
|
|
80acae16fa | ||
|
|
19c6007a22 | ||
|
|
0ed1a9f22a | ||
|
|
a29fce4a61 | ||
|
|
1225437099 | ||
|
|
5192abb6b5 | ||
|
|
40aae11888 | ||
|
|
177bf30867 | ||
|
|
d76c1af759 | ||
|
|
da283854f6 | ||
|
|
670b03f7e7 | ||
|
|
a163742fc9 | ||
|
|
04dc540b23 | ||
|
|
943f09d357 | ||
|
|
900194fab2 | ||
|
|
662444774f | ||
|
|
e24cb99b79 | ||
|
|
1e618dbc29 | ||
|
|
a3bbf35cbd | ||
|
|
caa1abb13a | ||
|
|
574dfad2dc | ||
|
|
bc0a6a7039 | ||
|
|
6876477371 | ||
|
|
b918874f81 | ||
|
|
20a2d12dde | ||
|
|
ca9e2b48ee | ||
|
|
e499a54a9d | ||
|
|
6564d376e5 | ||
|
|
c8d9af6ce1 | ||
|
|
2d9920ba7d | ||
|
|
0d7b6ae64c | ||
|
|
fd9ce53281 | ||
|
|
12abb54700 | ||
|
|
9261dd7734 | ||
|
|
0f73203e35 | ||
|
|
e2d9034d3d | ||
|
|
ba4bd0c0b4 | ||
|
|
61acf29125 | ||
|
|
3e1674a060 | ||
|
|
0b3fe95afd | ||
|
|
e5b325e481 | ||
|
|
33351c7bc7 | ||
|
|
b689e5b401 | ||
|
|
9096492b55 | ||
|
|
3d3aee2e72 | ||
|
|
591d47fb9b | ||
|
|
5350b3d98e | ||
|
|
10a4c6e3ad | ||
|
|
f3c7a1a8c4 | ||
|
|
f19e3816d0 | ||
|
|
30d0405388 | ||
|
|
5deeb7ef15 | ||
|
|
895e8247af | ||
|
|
2d0976ac2c | ||
|
|
4667f9012f | ||
|
|
af6928a268 | ||
|
|
358e4a89a1 | ||
|
|
b5e83156d9 | ||
|
|
fddc806db2 | ||
|
|
99b31067f7 | ||
|
|
f5c108c836 | ||
|
|
bbc8b6e906 | ||
|
|
118c1e8e72 | ||
|
|
48eb3766d2 | ||
|
|
280c1c8623 | ||
|
|
9f31df3fa8 | ||
|
|
9763765ebf | ||
|
|
74fcdf3d42 | ||
|
|
e01c917a97 | ||
|
|
9b54b46b06 | ||
|
|
8ea68bd689 | ||
|
|
ab326c96eb | ||
|
|
fe2c265c91 | ||
|
|
df249ee5d5 | ||
|
|
8aa425d22f | ||
|
|
ba92d70422 | ||
|
|
2f109e74ca | ||
|
|
5c52fa93d5 | ||
|
|
c418fd522f | ||
|
|
7045f89de4 |
@@ -32,7 +32,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
- name: Unit tests
|
||||
working-directory: backend
|
||||
run: make test-unit
|
||||
@@ -72,7 +72,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
- name: golangci-lint
|
||||
uses: golangci/golangci-lint-action@v9
|
||||
with:
|
||||
|
||||
@@ -115,7 +115,7 @@ jobs:
|
||||
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
|
||||
# Docker setup for GoReleaser
|
||||
- name: Set up QEMU
|
||||
|
||||
@@ -23,7 +23,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
- name: Run govulncheck
|
||||
working-directory: backend
|
||||
run: |
|
||||
|
||||
+4
-4
@@ -34,8 +34,8 @@
|
||||
### 开发工具
|
||||
|
||||
```bash
|
||||
# golangci-lint v2.7
|
||||
go install github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.7
|
||||
# golangci-lint(CI 用 v2.9,本地建议装同一版以免版本差异带来的噪音)
|
||||
go install github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.9
|
||||
|
||||
# pnpm (前端包管理)
|
||||
npm install -g pnpm
|
||||
@@ -47,13 +47,13 @@ npm install -g pnpm
|
||||
|
||||
| Workflow | 触发条件 | 检查内容 |
|
||||
|----------|----------|----------|
|
||||
| **backend-ci.yml** | push, pull_request | 单元测试 + 集成测试 + golangci-lint v2.7 |
|
||||
| **backend-ci.yml** | push, pull_request | 单元测试 + 集成测试 + golangci-lint v2.9 |
|
||||
| **security-scan.yml** | push, pull_request, 每周一 | govulncheck + gosec + pnpm audit |
|
||||
| **release.yml** | tag `v*` | 构建发布(PR 不触发) |
|
||||
|
||||
### CI 要求
|
||||
|
||||
- Go 版本必须是 **1.25.7**
|
||||
- Go 版本必须是 **1.26.5**:三个 workflow 都用 `go-version-file: backend/go.mod` 取版本,随后硬断言 `go version | grep -q 'go1.26.5'`。升级 Go 时要同时改 `backend/go.mod` 和 `backend-ci.yml`(两处)、`release.yml`、`security-scan.yml` 里的这句断言,否则 CI 会在版本校验步骤直接失败。
|
||||
- 前端使用 `pnpm install --frozen-lockfile`,必须提交 `pnpm-lock.yaml`
|
||||
|
||||
### 本地测试命令
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
# Sub2API
|
||||
|
||||
[](https://golang.org/)
|
||||
[](https://golang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://www.postgresql.org/)
|
||||
[](https://redis.io/)
|
||||
@@ -164,6 +164,16 @@ Please read the following carefully before using this project:
|
||||
<td><a href="https://hao.ai">hao.ai</a> is a high-speed, stable unified large-model API gateway for developers and teams. With a single API Key and a unified interface, you can access mainstream models such as GPT, Claude, and xAI Grok, with compatibility for common protocols and SDKs including OpenAI and Anthropic. The platform provides model routing, failover, team management, and complete request logs, with model prices as low as 15% of official reference pricing, helping users build AI applications more simply, more reliably, and at lower cost.</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.swiftproxy.net/?ref=sub2api"><img src="assets/partners/logos/swiftprox.png" alt="Swiftproxy" width="150"></a></td>
|
||||
<td>Swiftproxy is a high-performance proxy solution built for developers, providing stable and reliable residential and static residential proxy services. With 90M+ clean residential IPs, global coverage, flexible rotation, and precise geo-targeting, it helps projects such as web scraping, AI automation, browser automation, SEO monitoring, and multi-account management overcome access restrictions and improve workflow efficiency. It supports HTTP(S) and SOCKS5 protocols, integrates with popular automation tools like Playwright, Selenium, and Puppeteer, with dynamic proxy traffic that never expires until used and free testing available — <a href="https://www.swiftproxy.net/?ref=sub2api">start your free test now</a>!</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.duckip.cn/?keyword=cu7oog6y"><img src="assets/partners/logos/duckip.png" alt="DuckIP" width="150"></a></td>
|
||||
<td><a href="https://www.duckip.cn/?keyword=cu7oog6y">DuckIP</a> - 90M+ global residential network resources across 195+ countries and regions, with rotation and sticky sessions for public data collection, RAG updates, model evaluation, and multi-region data workloads. 🟢Residential Proxy - 20% Off; 🟢Static Residential Proxy - Starting at ¥50.00/IP; 🟢Unlimited Residential Proxy - Starting at ¥19.8/Hour. ✅Get 500M Free Trial.</td>
|
||||
</tr>
|
||||
|
||||
</table>
|
||||
|
||||
## Overview
|
||||
@@ -196,7 +206,7 @@ Community projects that extend or integrate with Sub2API:
|
||||
|
||||
| Component | Technology |
|
||||
|-----------|------------|
|
||||
| Backend | Go 1.25.7, Gin, Ent |
|
||||
| Backend | Go 1.26.5, Gin, Ent |
|
||||
| Frontend | Vue 3.4+, Vite 5+, TailwindCSS |
|
||||
| Database | PostgreSQL 15+ |
|
||||
| Cache/Queue | Redis 7+ |
|
||||
|
||||
+12
-2
@@ -4,7 +4,7 @@
|
||||
|
||||
# Sub2API
|
||||
|
||||
[](https://golang.org/)
|
||||
[](https://golang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://www.postgresql.org/)
|
||||
[](https://redis.io/)
|
||||
@@ -167,6 +167,16 @@
|
||||
<td><a href="https://hao.ai">hao.ai</a> 是面向开发者与团队的高速、稳定大模型统一 API 网关。通过一个 API Key 和统一接口,即可接入 GPT、Claude、xAI Grok 等主流模型,兼容 OpenAI、Anthropic 等常用协议与 SDK。平台提供模型路由、故障回退、团队管理及完整调用日志,模型价格低至官方参考价的 1.5 折,帮助用户更简单、更稳定、更低成本地构建 AI 应用。</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.swiftproxy.net/?ref=sub2api"><img src="assets/partners/logos/swiftprox.png" alt="Swiftproxy" width="150"></a></td>
|
||||
<td>Swiftproxy 是面向开发者的高性能代理解决方案,提供稳定可靠的住宅代理和静态住宅代理服务。拥有 9000 万+ 纯净住宅 IP,覆盖全球,支持灵活轮换和精准地理定位,帮助网页抓取、AI 自动化、浏览器自动化、SEO 监控和多账号管理等项目突破访问限制,提升工作流效率。支持 HTTP(S) 和 SOCKS5 协议,兼容 Playwright、Selenium、Puppeteer 等主流自动化工具,动态代理流量用完为止永不过期,支持免费测试 — <a href="https://www.swiftproxy.net/?ref=sub2api">立即开始免费测试</a>!</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.duckip.cn/?keyword=cu7oog6y"><img src="assets/partners/logos/duckip.png" alt="DuckIP" width="150"></a></td>
|
||||
<td><a href="https://www.duckip.cn/?keyword=cu7oog6y">DuckIP</a> - 9000 万+ 全球住宅网络资源,覆盖 195+ 国家和地区,支持轮换和粘性会话,适用于公共数据采集、RAG 更新、模型评估和多区域数据工作负载。🟢住宅代理 - 8 折优惠;🟢静态住宅代理 - ¥50.00/IP 起;🟢无限住宅代理 - ¥19.8/小时 起。✅免费领取 500M 试用流量。</td>
|
||||
</tr>
|
||||
|
||||
</table>
|
||||
|
||||
## 项目概述
|
||||
@@ -198,7 +208,7 @@ Sub2API 是一个 AI API 网关平台,用于分发和管理 AI 产品订阅的
|
||||
|
||||
| 组件 | 技术 |
|
||||
|------|------|
|
||||
| 后端 | Go 1.25.7, Gin, Ent |
|
||||
| 后端 | Go 1.26.5, Gin, Ent |
|
||||
| 前端 | Vue 3.4+, Vite 5+, TailwindCSS |
|
||||
| 数据库 | PostgreSQL 15+ |
|
||||
| 缓存/队列 | Redis 7+ |
|
||||
|
||||
+12
-2
@@ -4,7 +4,7 @@
|
||||
|
||||
# Sub2API
|
||||
|
||||
[](https://golang.org/)
|
||||
[](https://golang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://www.postgresql.org/)
|
||||
[](https://redis.io/)
|
||||
@@ -166,6 +166,16 @@
|
||||
<td><a href="https://hao.ai">hao.ai</a> は、開発者とチーム向けの高速で安定した大規模モデル統合 API ゲートウェイです。1 つの API Key と統一されたインターフェースで、GPT、Claude、xAI Grok などの主要モデルに接続でき、OpenAI や Anthropic などの一般的なプロトコルと SDK に対応しています。プラットフォームはモデルルーティング、フェイルオーバー、チーム管理、完全な呼び出しログを提供し、モデル価格は公式参考価格の 1.5 割から。よりシンプルに、より安定して、より低コストに AI アプリケーションを構築できます。</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.swiftproxy.net/?ref=sub2api"><img src="assets/partners/logos/swiftprox.png" alt="Swiftproxy" width="150"></a></td>
|
||||
<td>Swiftproxy は開発者向けの高性能プロキシソリューションで、安定して信頼できるレジデンシャルおよび静的レジデンシャルプロキシサービスを提供します。9,000 万以上のクリーンな住宅 IP を保有し、グローバルカバレッジ、柔軟なローテーション、精密なジオターゲティングにより、Web スクレイピング、AI オートメーション、ブラウザ自動化、SEO モニタリング、マルチアカウント管理などのプロジェクトがアクセス制限を克服し、ワークフロー効率を向上させます。HTTP(S) および SOCKS5 プロトコルに対応し、Playwright、Selenium、Puppeteer などの主要な自動化ツールと統合可能。動的プロキシトラフィックは使い切るまで期限切れなし、無料テストも利用可能 — <a href="https://www.swiftproxy.net/?ref=sub2api">今すぐ無料テストを開始</a>!</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.duckip.cn/?keyword=cu7oog6y"><img src="assets/partners/logos/duckip.png" alt="DuckIP" width="150"></a></td>
|
||||
<td><a href="https://www.duckip.cn/?keyword=cu7oog6y">DuckIP</a> - 195 以上の国と地域にわたる 9,000 万以上のグローバルレジデンシャルネットワークリソース。ローテーションとスティッキーセッションに対応し、パブリックデータ収集、RAG 更新、モデル評価、マルチリージョンデータワークロードに最適。🟢レジデンシャルプロキシ - 20% オフ;🟢スタティックレジデンシャルプロキシ - ¥50.00/IP から;🟢無制限レジデンシャルプロキシ - ¥19.8/時間 から。✅500M 無料トライアルを取得。</td>
|
||||
</tr>
|
||||
|
||||
</table>
|
||||
|
||||
## 概要
|
||||
@@ -197,7 +207,7 @@ Sub2API を拡張・統合するコミュニティプロジェクト:
|
||||
|
||||
| コンポーネント | 技術 |
|
||||
|-----------|------------|
|
||||
| バックエンド | Go 1.25.7, Gin, Ent |
|
||||
| バックエンド | Go 1.26.5, Gin, Ent |
|
||||
| フロントエンド | Vue 3.4+, Vite 5+, TailwindCSS |
|
||||
| データベース | PostgreSQL 15+ |
|
||||
| キャッシュ/キュー | Redis 7+ |
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 12 KiB |
@@ -1,32 +0,0 @@
|
||||
<svg
|
||||
width="340"
|
||||
height="112"
|
||||
viewBox="0 0 340 112"
|
||||
fill="none"
|
||||
xmlns="http://www.w3.org/2000/svg"
|
||||
role="img"
|
||||
aria-label="Hao.ai"
|
||||
>
|
||||
<!-- Fill the compact landscape canvas while preserving the homepage lockup proportions. -->
|
||||
<g transform="translate(170 56) scale(1.45) translate(-170 -56)">
|
||||
<!-- Homepage HaoAI mark, rendered in the brand green. -->
|
||||
<g transform="translate(66.55 12) scale(0.55)">
|
||||
<path
|
||||
d="M49 30C42.3726 30 37 35.3726 37 42V118C37 124.627 42.3726 130 49 130C55.6274 130 61 124.627 61 118V42C61 35.3726 55.6274 30 49 30ZM111 30C104.373 30 99 35.3726 99 42V118C99 124.627 104.373 130 111 130C117.627 130 123 124.627 123 118V42C123 35.3726 117.627 30 111 30ZM80 92.5C86.9036 92.5 92.5 86.9036 92.5 80C92.5 73.0964 86.9036 67.5 80 67.5C73.0964 67.5 67.5 73.0964 67.5 80C67.5 86.9036 73.0964 92.5 80 92.5Z"
|
||||
fill="#1DB954"
|
||||
fill-rule="evenodd"
|
||||
clip-rule="evenodd"
|
||||
/>
|
||||
</g>
|
||||
|
||||
<!-- Homepage wordmark converted from Inter SemiBold to vector outlines. -->
|
||||
<g transform="translate(150.35 69) scale(0.0166015625 -0.0166015625)" fill="#111111">
|
||||
<path d="M150 0L150 1490L417 1490L417 874L1110 874L1110 1490L1377 1490L1377 0L1110 0L1110 649L417 649L417 0Z" />
|
||||
<path transform="translate(1527 0)" d="M450 -23Q344 -23 259 15.5Q174 54 125 129Q76 204 76 314Q76 409 111.5 471Q147 533 207.5 570Q268 607 344 626Q420 645 502 653Q599 664 660 672Q721 680 749.5 698Q778 716 778 754L778 759Q778 814 756 852.5Q734 891 690.5 911.5Q647 932 583 932Q518 932 470.5 912Q423 892 393 860Q363 828 349 791L108 840Q141 939 210 1003.5Q279 1068 374.5 1100Q470 1132 582 1132Q661 1132 741.5 1113.5Q822 1095 889 1051Q956 1007 997 933Q1038 859 1038 747L1038 0L789 0L789 154L779 154Q754 107 710.5 66.5Q667 26 602.5 1.5Q538 -23 450 -23ZM517 170Q598 170 656.5 202Q715 234 747.5 286.5Q780 339 780 401L780 533Q767 523 736.5 514Q706 505 668.5 498Q631 491 594.5 486Q558 481 533 478Q475 470 428 452Q381 434 354.5 401.5Q328 369 328 317Q328 269 352.5 236.5Q377 204 419.5 187Q462 170 517 170Z" />
|
||||
<path transform="translate(2703 0)" d="M623 -23Q460 -23 340 49.5Q220 122 154.5 251.5Q89 381 89 553Q89 727 154.5 857.5Q220 988 340 1060Q460 1132 623 1132Q787 1132 907 1060Q1027 988 1092.5 857.5Q1158 727 1158 553Q1158 381 1092.5 251.5Q1027 122 907 49.5Q787 -23 623 -23ZM623 188Q714 188 774.5 237Q835 286 864.5 369Q894 452 894 554Q894 657 864.5 740Q835 823 774.5 872Q714 921 623 921Q532 921 472.5 872Q413 823 383.5 740.5Q354 658 354 554Q354 452 383.5 369Q413 286 472.5 237Q532 188 623 188Z" />
|
||||
<path transform="translate(3950 0)" fill="#1DB954" d="M326 -17Q259 -17 212.5 29Q166 75 166 143Q166 210 212.5 256Q259 302 326 302Q394 302 440.5 256Q487 210 487 143Q487 75 440.5 29Q394 -17 326 -17Z" />
|
||||
<path transform="translate(4603 0)" d="M450 -23Q344 -23 259 15.5Q174 54 125 129Q76 204 76 314Q76 409 111.5 471Q147 533 207.5 570Q268 607 344 626Q420 645 502 653Q599 664 660 672Q721 680 749.5 698Q778 716 778 754L778 759Q778 814 756 852.5Q734 891 690.5 911.5Q647 932 583 932Q518 932 470.5 912Q423 892 393 860Q363 828 349 791L108 840Q141 939 210 1003.5Q279 1068 374.5 1100Q470 1132 582 1132Q661 1132 741.5 1113.5Q822 1095 889 1051Q956 1007 997 933Q1038 859 1038 747L1038 0L789 0L789 154L779 154Q754 107 710.5 66.5Q667 26 602.5 1.5Q538 -23 450 -23ZM517 170Q598 170 656.5 202Q715 234 747.5 286.5Q780 339 780 401L780 533Q767 523 736.5 514Q706 505 668.5 498Q631 491 594.5 486Q558 481 533 478Q475 470 428 452Q381 434 354.5 401.5Q328 369 328 317Q328 269 352.5 236.5Q377 204 419.5 187Q462 170 517 170Z" />
|
||||
<path transform="translate(5779 0)" d="M138 0L138 1118L398 1118L398 0ZM268 1276Q206 1276 161.5 1318Q117 1360 117 1418Q117 1478 161.5 1519.5Q206 1561 268 1561Q331 1561 375.5 1519.5Q420 1478 420 1419Q420 1360 375.5 1318Q331 1276 268 1276Z" />
|
||||
</g>
|
||||
</g>
|
||||
</svg>
|
||||
|
Before Width: | Height: | Size: 3.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 10 KiB |
@@ -1 +1 @@
|
||||
0.1.172
|
||||
0.1.176
|
||||
|
||||
@@ -189,7 +189,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
|
||||
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
|
||||
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
|
||||
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver)
|
||||
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver, channelService)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService, settingService)
|
||||
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
|
||||
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
|
||||
@@ -208,7 +208,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
dataManagementHandler := admin.NewDataManagementHandler(dataManagementService)
|
||||
backupObjectStoreFactory := repository.NewS3BackupStoreFactory()
|
||||
dbDumper := repository.NewPgDumper(configConfig)
|
||||
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper)
|
||||
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper, leaderLockCache, db)
|
||||
imageStorageFactory := repository.ProvideImageStorageFactory()
|
||||
imageStorageSettingService := service.ProvideImageStorageSettingService(settingRepository, secretEncryptor, backupService, imageStorageFactory, configConfig)
|
||||
backupHandler := admin.NewBackupHandler(backupService, userService, imageStorageSettingService)
|
||||
|
||||
+26
-2
@@ -97,6 +97,10 @@ type Group struct {
|
||||
AudioTtsPricePerMillionChars *float64 `json:"audio_tts_price_per_million_chars,omitempty"`
|
||||
// STT 每小时价格(USD)
|
||||
AudioSttPricePerHour *float64 `json:"audio_stt_price_per_hour,omitempty"`
|
||||
// 是否按上下文长度应用模型阶梯价格
|
||||
LongContextPricingEnabled bool `json:"long_context_pricing_enabled,omitempty"`
|
||||
// 分组逐模型定价;优先级高于渠道和内置定价
|
||||
ModelPricing json.RawMessage `json:"model_pricing,omitempty"`
|
||||
// 是否仅允许 Claude Code 客户端
|
||||
ClaudeCodeOnly bool `json:"claude_code_only,omitempty"`
|
||||
// 非 Claude Code 请求降级使用的分组 ID
|
||||
@@ -245,9 +249,9 @@ func (*Group) scanValues(columns []string) ([]any, error) {
|
||||
values := make([]any, len(columns))
|
||||
for i := range columns {
|
||||
switch columns[i] {
|
||||
case group.FieldVideoModelPrices, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldReasoningEffortMappings:
|
||||
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldReasoningEffortMappings:
|
||||
values[i] = new([]byte)
|
||||
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
|
||||
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldLongContextPricingEnabled, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
|
||||
values[i] = new(sql.NullBool)
|
||||
case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldBatchImageDiscountMultiplier, group.FieldBatchImageHoldMultiplier, group.FieldVideoRateMultiplier, group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p, group.FieldWebSearchPricePerCall, group.FieldSearchPricePer1k, group.FieldAudioRealtimePricePerMin, group.FieldAudioTtsPricePerMillionChars, group.FieldAudioSttPricePerHour, group.FieldProfitMinMargin, group.FieldProfitSafetyBuffer:
|
||||
values[i] = new(sql.NullFloat64)
|
||||
@@ -531,6 +535,20 @@ func (_m *Group) assignValues(columns []string, values []any) error {
|
||||
_m.AudioSttPricePerHour = new(float64)
|
||||
*_m.AudioSttPricePerHour = value.Float64
|
||||
}
|
||||
case group.FieldLongContextPricingEnabled:
|
||||
if value, ok := values[i].(*sql.NullBool); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field long_context_pricing_enabled", values[i])
|
||||
} else if value.Valid {
|
||||
_m.LongContextPricingEnabled = value.Bool
|
||||
}
|
||||
case group.FieldModelPricing:
|
||||
if value, ok := values[i].(*[]byte); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field model_pricing", values[i])
|
||||
} else if value != nil && len(*value) > 0 {
|
||||
if err := json.Unmarshal(*value, &_m.ModelPricing); err != nil {
|
||||
return fmt.Errorf("unmarshal field model_pricing: %w", err)
|
||||
}
|
||||
}
|
||||
case group.FieldClaudeCodeOnly:
|
||||
if value, ok := values[i].(*sql.NullBool); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field claude_code_only", values[i])
|
||||
@@ -896,6 +914,12 @@ func (_m *Group) String() string {
|
||||
builder.WriteString(fmt.Sprintf("%v", *v))
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("long_context_pricing_enabled=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.LongContextPricingEnabled))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("model_pricing=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.ModelPricing))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("claude_code_only=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.ClaudeCodeOnly))
|
||||
builder.WriteString(", ")
|
||||
|
||||
@@ -94,6 +94,10 @@ const (
|
||||
FieldAudioTtsPricePerMillionChars = "audio_tts_price_per_million_chars"
|
||||
// FieldAudioSttPricePerHour holds the string denoting the audio_stt_price_per_hour field in the database.
|
||||
FieldAudioSttPricePerHour = "audio_stt_price_per_hour"
|
||||
// FieldLongContextPricingEnabled holds the string denoting the long_context_pricing_enabled field in the database.
|
||||
FieldLongContextPricingEnabled = "long_context_pricing_enabled"
|
||||
// FieldModelPricing holds the string denoting the model_pricing field in the database.
|
||||
FieldModelPricing = "model_pricing"
|
||||
// FieldClaudeCodeOnly holds the string denoting the claude_code_only field in the database.
|
||||
FieldClaudeCodeOnly = "claude_code_only"
|
||||
// FieldFallbackGroupID holds the string denoting the fallback_group_id field in the database.
|
||||
@@ -250,6 +254,8 @@ var Columns = []string{
|
||||
FieldAudioRealtimePricePerMin,
|
||||
FieldAudioTtsPricePerMillionChars,
|
||||
FieldAudioSttPricePerHour,
|
||||
FieldLongContextPricingEnabled,
|
||||
FieldModelPricing,
|
||||
FieldClaudeCodeOnly,
|
||||
FieldFallbackGroupID,
|
||||
FieldFallbackGroupIDOnInvalidRequest,
|
||||
@@ -364,6 +370,8 @@ var (
|
||||
AudioTtsPricePerMillionCharsValidator func(float64) error
|
||||
// AudioSttPricePerHourValidator is a validator for the "audio_stt_price_per_hour" field. It is called by the builders before save.
|
||||
AudioSttPricePerHourValidator func(float64) error
|
||||
// DefaultLongContextPricingEnabled holds the default value on creation for the "long_context_pricing_enabled" field.
|
||||
DefaultLongContextPricingEnabled bool
|
||||
// DefaultClaudeCodeOnly holds the default value on creation for the "claude_code_only" field.
|
||||
DefaultClaudeCodeOnly bool
|
||||
// DefaultModelRoutingEnabled holds the default value on creation for the "model_routing_enabled" field.
|
||||
@@ -604,6 +612,11 @@ func ByAudioSttPricePerHour(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldAudioSttPricePerHour, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByLongContextPricingEnabled orders the results by the long_context_pricing_enabled field.
|
||||
func ByLongContextPricingEnabled(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldLongContextPricingEnabled, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByClaudeCodeOnly orders the results by the claude_code_only field.
|
||||
func ByClaudeCodeOnly(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldClaudeCodeOnly, opts...).ToFunc()
|
||||
|
||||
@@ -245,6 +245,11 @@ func AudioSttPricePerHour(v float64) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldAudioSttPricePerHour, v))
|
||||
}
|
||||
|
||||
// LongContextPricingEnabled applies equality check predicate on the "long_context_pricing_enabled" field. It's identical to LongContextPricingEnabledEQ.
|
||||
func LongContextPricingEnabled(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldLongContextPricingEnabled, v))
|
||||
}
|
||||
|
||||
// ClaudeCodeOnly applies equality check predicate on the "claude_code_only" field. It's identical to ClaudeCodeOnlyEQ.
|
||||
func ClaudeCodeOnly(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v))
|
||||
@@ -2045,6 +2050,26 @@ func AudioSttPricePerHourNotNil() predicate.Group {
|
||||
return predicate.Group(sql.FieldNotNull(FieldAudioSttPricePerHour))
|
||||
}
|
||||
|
||||
// LongContextPricingEnabledEQ applies the EQ predicate on the "long_context_pricing_enabled" field.
|
||||
func LongContextPricingEnabledEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldLongContextPricingEnabled, v))
|
||||
}
|
||||
|
||||
// LongContextPricingEnabledNEQ applies the NEQ predicate on the "long_context_pricing_enabled" field.
|
||||
func LongContextPricingEnabledNEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldNEQ(FieldLongContextPricingEnabled, v))
|
||||
}
|
||||
|
||||
// ModelPricingIsNil applies the IsNil predicate on the "model_pricing" field.
|
||||
func ModelPricingIsNil() predicate.Group {
|
||||
return predicate.Group(sql.FieldIsNull(FieldModelPricing))
|
||||
}
|
||||
|
||||
// ModelPricingNotNil applies the NotNil predicate on the "model_pricing" field.
|
||||
func ModelPricingNotNil() predicate.Group {
|
||||
return predicate.Group(sql.FieldNotNull(FieldModelPricing))
|
||||
}
|
||||
|
||||
// ClaudeCodeOnlyEQ applies the EQ predicate on the "claude_code_only" field.
|
||||
func ClaudeCodeOnlyEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v))
|
||||
|
||||
@@ -4,6 +4,7 @@ package ent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
@@ -559,6 +560,26 @@ func (_c *GroupCreate) SetNillableAudioSttPricePerHour(v *float64) *GroupCreate
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (_c *GroupCreate) SetLongContextPricingEnabled(v bool) *GroupCreate {
|
||||
_c.mutation.SetLongContextPricingEnabled(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableLongContextPricingEnabled sets the "long_context_pricing_enabled" field if the given value is not nil.
|
||||
func (_c *GroupCreate) SetNillableLongContextPricingEnabled(v *bool) *GroupCreate {
|
||||
if v != nil {
|
||||
_c.SetLongContextPricingEnabled(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (_c *GroupCreate) SetModelPricing(v json.RawMessage) *GroupCreate {
|
||||
_c.mutation.SetModelPricing(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (_c *GroupCreate) SetClaudeCodeOnly(v bool) *GroupCreate {
|
||||
_c.mutation.SetClaudeCodeOnly(v)
|
||||
@@ -1042,6 +1063,10 @@ func (_c *GroupCreate) defaults() error {
|
||||
v := group.DefaultVideoRateMultiplier
|
||||
_c.mutation.SetVideoRateMultiplier(v)
|
||||
}
|
||||
if _, ok := _c.mutation.LongContextPricingEnabled(); !ok {
|
||||
v := group.DefaultLongContextPricingEnabled
|
||||
_c.mutation.SetLongContextPricingEnabled(v)
|
||||
}
|
||||
if _, ok := _c.mutation.ClaudeCodeOnly(); !ok {
|
||||
v := group.DefaultClaudeCodeOnly
|
||||
_c.mutation.SetClaudeCodeOnly(v)
|
||||
@@ -1237,6 +1262,9 @@ func (_c *GroupCreate) check() error {
|
||||
return &ValidationError{Name: "audio_stt_price_per_hour", err: fmt.Errorf(`ent: validator failed for field "Group.audio_stt_price_per_hour": %w`, err)}
|
||||
}
|
||||
}
|
||||
if _, ok := _c.mutation.LongContextPricingEnabled(); !ok {
|
||||
return &ValidationError{Name: "long_context_pricing_enabled", err: errors.New(`ent: missing required field "Group.long_context_pricing_enabled"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.ClaudeCodeOnly(); !ok {
|
||||
return &ValidationError{Name: "claude_code_only", err: errors.New(`ent: missing required field "Group.claude_code_only"`)}
|
||||
}
|
||||
@@ -1484,6 +1512,14 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(group.FieldAudioSttPricePerHour, field.TypeFloat64, value)
|
||||
_node.AudioSttPricePerHour = &value
|
||||
}
|
||||
if value, ok := _c.mutation.LongContextPricingEnabled(); ok {
|
||||
_spec.SetField(group.FieldLongContextPricingEnabled, field.TypeBool, value)
|
||||
_node.LongContextPricingEnabled = value
|
||||
}
|
||||
if value, ok := _c.mutation.ModelPricing(); ok {
|
||||
_spec.SetField(group.FieldModelPricing, field.TypeJSON, value)
|
||||
_node.ModelPricing = value
|
||||
}
|
||||
if value, ok := _c.mutation.ClaudeCodeOnly(); ok {
|
||||
_spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value)
|
||||
_node.ClaudeCodeOnly = value
|
||||
@@ -2396,6 +2432,36 @@ func (u *GroupUpsert) ClearAudioSttPricePerHour() *GroupUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (u *GroupUpsert) SetLongContextPricingEnabled(v bool) *GroupUpsert {
|
||||
u.Set(group.FieldLongContextPricingEnabled, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateLongContextPricingEnabled sets the "long_context_pricing_enabled" field to the value that was provided on create.
|
||||
func (u *GroupUpsert) UpdateLongContextPricingEnabled() *GroupUpsert {
|
||||
u.SetExcluded(group.FieldLongContextPricingEnabled)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (u *GroupUpsert) SetModelPricing(v json.RawMessage) *GroupUpsert {
|
||||
u.Set(group.FieldModelPricing, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateModelPricing sets the "model_pricing" field to the value that was provided on create.
|
||||
func (u *GroupUpsert) UpdateModelPricing() *GroupUpsert {
|
||||
u.SetExcluded(group.FieldModelPricing)
|
||||
return u
|
||||
}
|
||||
|
||||
// ClearModelPricing clears the value of the "model_pricing" field.
|
||||
func (u *GroupUpsert) ClearModelPricing() *GroupUpsert {
|
||||
u.SetNull(group.FieldModelPricing)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (u *GroupUpsert) SetClaudeCodeOnly(v bool) *GroupUpsert {
|
||||
u.Set(group.FieldClaudeCodeOnly, v)
|
||||
@@ -3534,6 +3600,41 @@ func (u *GroupUpsertOne) ClearAudioSttPricePerHour() *GroupUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (u *GroupUpsertOne) SetLongContextPricingEnabled(v bool) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetLongContextPricingEnabled(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateLongContextPricingEnabled sets the "long_context_pricing_enabled" field to the value that was provided on create.
|
||||
func (u *GroupUpsertOne) UpdateLongContextPricingEnabled() *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateLongContextPricingEnabled()
|
||||
})
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (u *GroupUpsertOne) SetModelPricing(v json.RawMessage) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetModelPricing(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateModelPricing sets the "model_pricing" field to the value that was provided on create.
|
||||
func (u *GroupUpsertOne) UpdateModelPricing() *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateModelPricing()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearModelPricing clears the value of the "model_pricing" field.
|
||||
func (u *GroupUpsertOne) ClearModelPricing() *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.ClearModelPricing()
|
||||
})
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (u *GroupUpsertOne) SetClaudeCodeOnly(v bool) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
@@ -4889,6 +4990,41 @@ func (u *GroupUpsertBulk) ClearAudioSttPricePerHour() *GroupUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (u *GroupUpsertBulk) SetLongContextPricingEnabled(v bool) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetLongContextPricingEnabled(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateLongContextPricingEnabled sets the "long_context_pricing_enabled" field to the value that was provided on create.
|
||||
func (u *GroupUpsertBulk) UpdateLongContextPricingEnabled() *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateLongContextPricingEnabled()
|
||||
})
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (u *GroupUpsertBulk) SetModelPricing(v json.RawMessage) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetModelPricing(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateModelPricing sets the "model_pricing" field to the value that was provided on create.
|
||||
func (u *GroupUpsertBulk) UpdateModelPricing() *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateModelPricing()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearModelPricing clears the value of the "model_pricing" field.
|
||||
func (u *GroupUpsertBulk) ClearModelPricing() *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.ClearModelPricing()
|
||||
})
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (u *GroupUpsertBulk) SetClaudeCodeOnly(v bool) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
|
||||
@@ -4,6 +4,7 @@ package ent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
@@ -787,6 +788,38 @@ func (_u *GroupUpdate) ClearAudioSttPricePerHour() *GroupUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (_u *GroupUpdate) SetLongContextPricingEnabled(v bool) *GroupUpdate {
|
||||
_u.mutation.SetLongContextPricingEnabled(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableLongContextPricingEnabled sets the "long_context_pricing_enabled" field if the given value is not nil.
|
||||
func (_u *GroupUpdate) SetNillableLongContextPricingEnabled(v *bool) *GroupUpdate {
|
||||
if v != nil {
|
||||
_u.SetLongContextPricingEnabled(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (_u *GroupUpdate) SetModelPricing(v json.RawMessage) *GroupUpdate {
|
||||
_u.mutation.SetModelPricing(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// AppendModelPricing appends value to the "model_pricing" field.
|
||||
func (_u *GroupUpdate) AppendModelPricing(v json.RawMessage) *GroupUpdate {
|
||||
_u.mutation.AppendModelPricing(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearModelPricing clears the value of the "model_pricing" field.
|
||||
func (_u *GroupUpdate) ClearModelPricing() *GroupUpdate {
|
||||
_u.mutation.ClearModelPricing()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (_u *GroupUpdate) SetClaudeCodeOnly(v bool) *GroupUpdate {
|
||||
_u.mutation.SetClaudeCodeOnly(v)
|
||||
@@ -1697,6 +1730,20 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if _u.mutation.AudioSttPricePerHourCleared() {
|
||||
_spec.ClearField(group.FieldAudioSttPricePerHour, field.TypeFloat64)
|
||||
}
|
||||
if value, ok := _u.mutation.LongContextPricingEnabled(); ok {
|
||||
_spec.SetField(group.FieldLongContextPricingEnabled, field.TypeBool, value)
|
||||
}
|
||||
if value, ok := _u.mutation.ModelPricing(); ok {
|
||||
_spec.SetField(group.FieldModelPricing, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AppendedModelPricing(); ok {
|
||||
_spec.AddModifier(func(u *sql.UpdateBuilder) {
|
||||
sqljson.Append(u, group.FieldModelPricing, value)
|
||||
})
|
||||
}
|
||||
if _u.mutation.ModelPricingCleared() {
|
||||
_spec.ClearField(group.FieldModelPricing, field.TypeJSON)
|
||||
}
|
||||
if value, ok := _u.mutation.ClaudeCodeOnly(); ok {
|
||||
_spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value)
|
||||
}
|
||||
@@ -2862,6 +2909,38 @@ func (_u *GroupUpdateOne) ClearAudioSttPricePerHour() *GroupUpdateOne {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (_u *GroupUpdateOne) SetLongContextPricingEnabled(v bool) *GroupUpdateOne {
|
||||
_u.mutation.SetLongContextPricingEnabled(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableLongContextPricingEnabled sets the "long_context_pricing_enabled" field if the given value is not nil.
|
||||
func (_u *GroupUpdateOne) SetNillableLongContextPricingEnabled(v *bool) *GroupUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetLongContextPricingEnabled(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (_u *GroupUpdateOne) SetModelPricing(v json.RawMessage) *GroupUpdateOne {
|
||||
_u.mutation.SetModelPricing(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// AppendModelPricing appends value to the "model_pricing" field.
|
||||
func (_u *GroupUpdateOne) AppendModelPricing(v json.RawMessage) *GroupUpdateOne {
|
||||
_u.mutation.AppendModelPricing(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearModelPricing clears the value of the "model_pricing" field.
|
||||
func (_u *GroupUpdateOne) ClearModelPricing() *GroupUpdateOne {
|
||||
_u.mutation.ClearModelPricing()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (_u *GroupUpdateOne) SetClaudeCodeOnly(v bool) *GroupUpdateOne {
|
||||
_u.mutation.SetClaudeCodeOnly(v)
|
||||
@@ -3802,6 +3881,20 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
|
||||
if _u.mutation.AudioSttPricePerHourCleared() {
|
||||
_spec.ClearField(group.FieldAudioSttPricePerHour, field.TypeFloat64)
|
||||
}
|
||||
if value, ok := _u.mutation.LongContextPricingEnabled(); ok {
|
||||
_spec.SetField(group.FieldLongContextPricingEnabled, field.TypeBool, value)
|
||||
}
|
||||
if value, ok := _u.mutation.ModelPricing(); ok {
|
||||
_spec.SetField(group.FieldModelPricing, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AppendedModelPricing(); ok {
|
||||
_spec.AddModifier(func(u *sql.UpdateBuilder) {
|
||||
sqljson.Append(u, group.FieldModelPricing, value)
|
||||
})
|
||||
}
|
||||
if _u.mutation.ModelPricingCleared() {
|
||||
_spec.ClearField(group.FieldModelPricing, field.TypeJSON)
|
||||
}
|
||||
if value, ok := _u.mutation.ClaudeCodeOnly(); ok {
|
||||
_spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value)
|
||||
}
|
||||
|
||||
@@ -934,6 +934,8 @@ var (
|
||||
{Name: "audio_realtime_price_per_min", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
|
||||
{Name: "audio_tts_price_per_million_chars", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
|
||||
{Name: "audio_stt_price_per_hour", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
|
||||
{Name: "long_context_pricing_enabled", Type: field.TypeBool, Default: true},
|
||||
{Name: "model_pricing", Type: field.TypeJSON, Nullable: true, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "claude_code_only", Type: field.TypeBool, Default: false},
|
||||
{Name: "fallback_group_id", Type: field.TypeInt64, Nullable: true},
|
||||
{Name: "fallback_group_id_on_invalid_request", Type: field.TypeInt64, Nullable: true},
|
||||
@@ -990,7 +992,7 @@ var (
|
||||
{
|
||||
Name: "group_sort_order",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[47]},
|
||||
Columns: []*schema.Column{GroupsColumns[49]},
|
||||
},
|
||||
{
|
||||
Name: "idx_groups_duplicate_operation_id_active",
|
||||
|
||||
+145
-1
@@ -21907,6 +21907,9 @@ type GroupMutation struct {
|
||||
addaudio_tts_price_per_million_chars *float64
|
||||
audio_stt_price_per_hour *float64
|
||||
addaudio_stt_price_per_hour *float64
|
||||
long_context_pricing_enabled *bool
|
||||
model_pricing *json.RawMessage
|
||||
appendmodel_pricing json.RawMessage
|
||||
claude_code_only *bool
|
||||
fallback_group_id *int64
|
||||
addfallback_group_id *int64
|
||||
@@ -24130,6 +24133,107 @@ func (m *GroupMutation) ResetAudioSttPricePerHour() {
|
||||
delete(m.clearedFields, group.FieldAudioSttPricePerHour)
|
||||
}
|
||||
|
||||
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
|
||||
func (m *GroupMutation) SetLongContextPricingEnabled(b bool) {
|
||||
m.long_context_pricing_enabled = &b
|
||||
}
|
||||
|
||||
// LongContextPricingEnabled returns the value of the "long_context_pricing_enabled" field in the mutation.
|
||||
func (m *GroupMutation) LongContextPricingEnabled() (r bool, exists bool) {
|
||||
v := m.long_context_pricing_enabled
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldLongContextPricingEnabled returns the old "long_context_pricing_enabled" field's value of the Group entity.
|
||||
// If the Group object wasn't provided to the builder, the object is fetched from the database.
|
||||
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
|
||||
func (m *GroupMutation) OldLongContextPricingEnabled(ctx context.Context) (v bool, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldLongContextPricingEnabled is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldLongContextPricingEnabled requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldLongContextPricingEnabled: %w", err)
|
||||
}
|
||||
return oldValue.LongContextPricingEnabled, nil
|
||||
}
|
||||
|
||||
// ResetLongContextPricingEnabled resets all changes to the "long_context_pricing_enabled" field.
|
||||
func (m *GroupMutation) ResetLongContextPricingEnabled() {
|
||||
m.long_context_pricing_enabled = nil
|
||||
}
|
||||
|
||||
// SetModelPricing sets the "model_pricing" field.
|
||||
func (m *GroupMutation) SetModelPricing(jm json.RawMessage) {
|
||||
m.model_pricing = &jm
|
||||
m.appendmodel_pricing = nil
|
||||
}
|
||||
|
||||
// ModelPricing returns the value of the "model_pricing" field in the mutation.
|
||||
func (m *GroupMutation) ModelPricing() (r json.RawMessage, exists bool) {
|
||||
v := m.model_pricing
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldModelPricing returns the old "model_pricing" field's value of the Group entity.
|
||||
// If the Group object wasn't provided to the builder, the object is fetched from the database.
|
||||
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
|
||||
func (m *GroupMutation) OldModelPricing(ctx context.Context) (v json.RawMessage, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldModelPricing is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldModelPricing requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldModelPricing: %w", err)
|
||||
}
|
||||
return oldValue.ModelPricing, nil
|
||||
}
|
||||
|
||||
// AppendModelPricing adds jm to the "model_pricing" field.
|
||||
func (m *GroupMutation) AppendModelPricing(jm json.RawMessage) {
|
||||
m.appendmodel_pricing = append(m.appendmodel_pricing, jm...)
|
||||
}
|
||||
|
||||
// AppendedModelPricing returns the list of values that were appended to the "model_pricing" field in this mutation.
|
||||
func (m *GroupMutation) AppendedModelPricing() (json.RawMessage, bool) {
|
||||
if len(m.appendmodel_pricing) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return m.appendmodel_pricing, true
|
||||
}
|
||||
|
||||
// ClearModelPricing clears the value of the "model_pricing" field.
|
||||
func (m *GroupMutation) ClearModelPricing() {
|
||||
m.model_pricing = nil
|
||||
m.appendmodel_pricing = nil
|
||||
m.clearedFields[group.FieldModelPricing] = struct{}{}
|
||||
}
|
||||
|
||||
// ModelPricingCleared returns if the "model_pricing" field was cleared in this mutation.
|
||||
func (m *GroupMutation) ModelPricingCleared() bool {
|
||||
_, ok := m.clearedFields[group.FieldModelPricing]
|
||||
return ok
|
||||
}
|
||||
|
||||
// ResetModelPricing resets all changes to the "model_pricing" field.
|
||||
func (m *GroupMutation) ResetModelPricing() {
|
||||
m.model_pricing = nil
|
||||
m.appendmodel_pricing = nil
|
||||
delete(m.clearedFields, group.FieldModelPricing)
|
||||
}
|
||||
|
||||
// SetClaudeCodeOnly sets the "claude_code_only" field.
|
||||
func (m *GroupMutation) SetClaudeCodeOnly(b bool) {
|
||||
m.claude_code_only = &b
|
||||
@@ -25435,7 +25539,7 @@ func (m *GroupMutation) Type() string {
|
||||
// order to get all numeric fields that were incremented/decremented, call
|
||||
// AddedFields().
|
||||
func (m *GroupMutation) Fields() []string {
|
||||
fields := make([]string, 0, 60)
|
||||
fields := make([]string, 0, 62)
|
||||
if m.created_at != nil {
|
||||
fields = append(fields, group.FieldCreatedAt)
|
||||
}
|
||||
@@ -25553,6 +25657,12 @@ func (m *GroupMutation) Fields() []string {
|
||||
if m.audio_stt_price_per_hour != nil {
|
||||
fields = append(fields, group.FieldAudioSttPricePerHour)
|
||||
}
|
||||
if m.long_context_pricing_enabled != nil {
|
||||
fields = append(fields, group.FieldLongContextPricingEnabled)
|
||||
}
|
||||
if m.model_pricing != nil {
|
||||
fields = append(fields, group.FieldModelPricing)
|
||||
}
|
||||
if m.claude_code_only != nil {
|
||||
fields = append(fields, group.FieldClaudeCodeOnly)
|
||||
}
|
||||
@@ -25702,6 +25812,10 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.AudioTtsPricePerMillionChars()
|
||||
case group.FieldAudioSttPricePerHour:
|
||||
return m.AudioSttPricePerHour()
|
||||
case group.FieldLongContextPricingEnabled:
|
||||
return m.LongContextPricingEnabled()
|
||||
case group.FieldModelPricing:
|
||||
return m.ModelPricing()
|
||||
case group.FieldClaudeCodeOnly:
|
||||
return m.ClaudeCodeOnly()
|
||||
case group.FieldFallbackGroupID:
|
||||
@@ -25831,6 +25945,10 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
|
||||
return m.OldAudioTtsPricePerMillionChars(ctx)
|
||||
case group.FieldAudioSttPricePerHour:
|
||||
return m.OldAudioSttPricePerHour(ctx)
|
||||
case group.FieldLongContextPricingEnabled:
|
||||
return m.OldLongContextPricingEnabled(ctx)
|
||||
case group.FieldModelPricing:
|
||||
return m.OldModelPricing(ctx)
|
||||
case group.FieldClaudeCodeOnly:
|
||||
return m.OldClaudeCodeOnly(ctx)
|
||||
case group.FieldFallbackGroupID:
|
||||
@@ -26155,6 +26273,20 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetAudioSttPricePerHour(v)
|
||||
return nil
|
||||
case group.FieldLongContextPricingEnabled:
|
||||
v, ok := value.(bool)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetLongContextPricingEnabled(v)
|
||||
return nil
|
||||
case group.FieldModelPricing:
|
||||
v, ok := value.(json.RawMessage)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetModelPricing(v)
|
||||
return nil
|
||||
case group.FieldClaudeCodeOnly:
|
||||
v, ok := value.(bool)
|
||||
if !ok {
|
||||
@@ -26713,6 +26845,9 @@ func (m *GroupMutation) ClearedFields() []string {
|
||||
if m.FieldCleared(group.FieldAudioSttPricePerHour) {
|
||||
fields = append(fields, group.FieldAudioSttPricePerHour)
|
||||
}
|
||||
if m.FieldCleared(group.FieldModelPricing) {
|
||||
fields = append(fields, group.FieldModelPricing)
|
||||
}
|
||||
if m.FieldCleared(group.FieldFallbackGroupID) {
|
||||
fields = append(fields, group.FieldFallbackGroupID)
|
||||
}
|
||||
@@ -26790,6 +26925,9 @@ func (m *GroupMutation) ClearField(name string) error {
|
||||
case group.FieldAudioSttPricePerHour:
|
||||
m.ClearAudioSttPricePerHour()
|
||||
return nil
|
||||
case group.FieldModelPricing:
|
||||
m.ClearModelPricing()
|
||||
return nil
|
||||
case group.FieldFallbackGroupID:
|
||||
m.ClearFallbackGroupID()
|
||||
return nil
|
||||
@@ -26924,6 +27062,12 @@ func (m *GroupMutation) ResetField(name string) error {
|
||||
case group.FieldAudioSttPricePerHour:
|
||||
m.ResetAudioSttPricePerHour()
|
||||
return nil
|
||||
case group.FieldLongContextPricingEnabled:
|
||||
m.ResetLongContextPricingEnabled()
|
||||
return nil
|
||||
case group.FieldModelPricing:
|
||||
m.ResetModelPricing()
|
||||
return nil
|
||||
case group.FieldClaudeCodeOnly:
|
||||
m.ResetClaudeCodeOnly()
|
||||
return nil
|
||||
|
||||
@@ -1133,80 +1133,84 @@ func init() {
|
||||
groupDescAudioSttPricePerHour := groupFields[35].Descriptor()
|
||||
// group.AudioSttPricePerHourValidator is a validator for the "audio_stt_price_per_hour" field. It is called by the builders before save.
|
||||
group.AudioSttPricePerHourValidator = groupDescAudioSttPricePerHour.Validators[0].(func(float64) error)
|
||||
// groupDescLongContextPricingEnabled is the schema descriptor for long_context_pricing_enabled field.
|
||||
groupDescLongContextPricingEnabled := groupFields[36].Descriptor()
|
||||
// group.DefaultLongContextPricingEnabled holds the default value on creation for the long_context_pricing_enabled field.
|
||||
group.DefaultLongContextPricingEnabled = groupDescLongContextPricingEnabled.Default.(bool)
|
||||
// groupDescClaudeCodeOnly is the schema descriptor for claude_code_only field.
|
||||
groupDescClaudeCodeOnly := groupFields[36].Descriptor()
|
||||
groupDescClaudeCodeOnly := groupFields[38].Descriptor()
|
||||
// group.DefaultClaudeCodeOnly holds the default value on creation for the claude_code_only field.
|
||||
group.DefaultClaudeCodeOnly = groupDescClaudeCodeOnly.Default.(bool)
|
||||
// groupDescModelRoutingEnabled is the schema descriptor for model_routing_enabled field.
|
||||
groupDescModelRoutingEnabled := groupFields[40].Descriptor()
|
||||
groupDescModelRoutingEnabled := groupFields[42].Descriptor()
|
||||
// group.DefaultModelRoutingEnabled holds the default value on creation for the model_routing_enabled field.
|
||||
group.DefaultModelRoutingEnabled = groupDescModelRoutingEnabled.Default.(bool)
|
||||
// groupDescMcpXMLInject is the schema descriptor for mcp_xml_inject field.
|
||||
groupDescMcpXMLInject := groupFields[41].Descriptor()
|
||||
groupDescMcpXMLInject := groupFields[43].Descriptor()
|
||||
// group.DefaultMcpXMLInject holds the default value on creation for the mcp_xml_inject field.
|
||||
group.DefaultMcpXMLInject = groupDescMcpXMLInject.Default.(bool)
|
||||
// groupDescSupportedModelScopes is the schema descriptor for supported_model_scopes field.
|
||||
groupDescSupportedModelScopes := groupFields[42].Descriptor()
|
||||
groupDescSupportedModelScopes := groupFields[44].Descriptor()
|
||||
// group.DefaultSupportedModelScopes holds the default value on creation for the supported_model_scopes field.
|
||||
group.DefaultSupportedModelScopes = groupDescSupportedModelScopes.Default.([]string)
|
||||
// groupDescSortOrder is the schema descriptor for sort_order field.
|
||||
groupDescSortOrder := groupFields[43].Descriptor()
|
||||
groupDescSortOrder := groupFields[45].Descriptor()
|
||||
// group.DefaultSortOrder holds the default value on creation for the sort_order field.
|
||||
group.DefaultSortOrder = groupDescSortOrder.Default.(int)
|
||||
// groupDescAllowMessagesDispatch is the schema descriptor for allow_messages_dispatch field.
|
||||
groupDescAllowMessagesDispatch := groupFields[44].Descriptor()
|
||||
groupDescAllowMessagesDispatch := groupFields[46].Descriptor()
|
||||
// group.DefaultAllowMessagesDispatch holds the default value on creation for the allow_messages_dispatch field.
|
||||
group.DefaultAllowMessagesDispatch = groupDescAllowMessagesDispatch.Default.(bool)
|
||||
// groupDescAllowLive is the schema descriptor for allow_live field.
|
||||
groupDescAllowLive := groupFields[45].Descriptor()
|
||||
groupDescAllowLive := groupFields[47].Descriptor()
|
||||
// group.DefaultAllowLive holds the default value on creation for the allow_live field.
|
||||
group.DefaultAllowLive = groupDescAllowLive.Default.(bool)
|
||||
// groupDescRequireOauthOnly is the schema descriptor for require_oauth_only field.
|
||||
groupDescRequireOauthOnly := groupFields[46].Descriptor()
|
||||
groupDescRequireOauthOnly := groupFields[48].Descriptor()
|
||||
// group.DefaultRequireOauthOnly holds the default value on creation for the require_oauth_only field.
|
||||
group.DefaultRequireOauthOnly = groupDescRequireOauthOnly.Default.(bool)
|
||||
// groupDescRequirePrivacySet is the schema descriptor for require_privacy_set field.
|
||||
groupDescRequirePrivacySet := groupFields[47].Descriptor()
|
||||
groupDescRequirePrivacySet := groupFields[49].Descriptor()
|
||||
// group.DefaultRequirePrivacySet holds the default value on creation for the require_privacy_set field.
|
||||
group.DefaultRequirePrivacySet = groupDescRequirePrivacySet.Default.(bool)
|
||||
// groupDescDefaultMappedModel is the schema descriptor for default_mapped_model field.
|
||||
groupDescDefaultMappedModel := groupFields[48].Descriptor()
|
||||
groupDescDefaultMappedModel := groupFields[50].Descriptor()
|
||||
// group.DefaultDefaultMappedModel holds the default value on creation for the default_mapped_model field.
|
||||
group.DefaultDefaultMappedModel = groupDescDefaultMappedModel.Default.(string)
|
||||
// group.DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save.
|
||||
group.DefaultMappedModelValidator = groupDescDefaultMappedModel.Validators[0].(func(string) error)
|
||||
// groupDescMessagesDispatchModelConfig is the schema descriptor for messages_dispatch_model_config field.
|
||||
groupDescMessagesDispatchModelConfig := groupFields[49].Descriptor()
|
||||
groupDescMessagesDispatchModelConfig := groupFields[51].Descriptor()
|
||||
// group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field.
|
||||
group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig)
|
||||
// groupDescModelsListConfig is the schema descriptor for models_list_config field.
|
||||
groupDescModelsListConfig := groupFields[50].Descriptor()
|
||||
groupDescModelsListConfig := groupFields[52].Descriptor()
|
||||
// group.DefaultModelsListConfig holds the default value on creation for the models_list_config field.
|
||||
group.DefaultModelsListConfig = groupDescModelsListConfig.Default.(domain.GroupModelsListConfig)
|
||||
// groupDescRpmLimit is the schema descriptor for rpm_limit field.
|
||||
groupDescRpmLimit := groupFields[51].Descriptor()
|
||||
groupDescRpmLimit := groupFields[53].Descriptor()
|
||||
// group.DefaultRpmLimit holds the default value on creation for the rpm_limit field.
|
||||
group.DefaultRpmLimit = groupDescRpmLimit.Default.(int)
|
||||
// groupDescMaxReasoningEffort is the schema descriptor for max_reasoning_effort field.
|
||||
groupDescMaxReasoningEffort := groupFields[52].Descriptor()
|
||||
groupDescMaxReasoningEffort := groupFields[54].Descriptor()
|
||||
// group.DefaultMaxReasoningEffort holds the default value on creation for the max_reasoning_effort field.
|
||||
group.DefaultMaxReasoningEffort = groupDescMaxReasoningEffort.Default.(string)
|
||||
// group.MaxReasoningEffortValidator is a validator for the "max_reasoning_effort" field. It is called by the builders before save.
|
||||
group.MaxReasoningEffortValidator = groupDescMaxReasoningEffort.Validators[0].(func(string) error)
|
||||
// groupDescReasoningEffortMappings is the schema descriptor for reasoning_effort_mappings field.
|
||||
groupDescReasoningEffortMappings := groupFields[53].Descriptor()
|
||||
groupDescReasoningEffortMappings := groupFields[55].Descriptor()
|
||||
// group.DefaultReasoningEffortMappings holds the default value on creation for the reasoning_effort_mappings field.
|
||||
group.DefaultReasoningEffortMappings = groupDescReasoningEffortMappings.Default.([]domain.ReasoningEffortMapping)
|
||||
// groupDescProfitControlEnabled is the schema descriptor for profit_control_enabled field.
|
||||
groupDescProfitControlEnabled := groupFields[54].Descriptor()
|
||||
groupDescProfitControlEnabled := groupFields[56].Descriptor()
|
||||
// group.DefaultProfitControlEnabled holds the default value on creation for the profit_control_enabled field.
|
||||
group.DefaultProfitControlEnabled = groupDescProfitControlEnabled.Default.(bool)
|
||||
// groupDescProfitMinMargin is the schema descriptor for profit_min_margin field.
|
||||
groupDescProfitMinMargin := groupFields[55].Descriptor()
|
||||
groupDescProfitMinMargin := groupFields[57].Descriptor()
|
||||
// group.DefaultProfitMinMargin holds the default value on creation for the profit_min_margin field.
|
||||
group.DefaultProfitMinMargin = groupDescProfitMinMargin.Default.(float64)
|
||||
// groupDescProfitSafetyBuffer is the schema descriptor for profit_safety_buffer field.
|
||||
groupDescProfitSafetyBuffer := groupFields[56].Descriptor()
|
||||
groupDescProfitSafetyBuffer := groupFields[58].Descriptor()
|
||||
// group.DefaultProfitSafetyBuffer holds the default value on creation for the profit_safety_buffer field.
|
||||
group.DefaultProfitSafetyBuffer = groupDescProfitSafetyBuffer.Default.(float64)
|
||||
idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin()
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package schema
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
|
||||
@@ -185,6 +187,13 @@ func (Group) Fields() []ent.Field {
|
||||
Min(0).
|
||||
SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}).
|
||||
Comment("STT 每小时价格(USD)"),
|
||||
field.Bool("long_context_pricing_enabled").
|
||||
Default(true).
|
||||
Comment("是否按上下文长度应用模型阶梯价格;默认开启以保持官方/渠道长上下文价"),
|
||||
field.JSON("model_pricing", json.RawMessage{}).
|
||||
Optional().
|
||||
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
|
||||
Comment("分组逐模型定价;优先级高于渠道和内置定价"),
|
||||
|
||||
// Claude Code 客户端限制 (added by migration 029)
|
||||
field.Bool("claude_code_only").
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
module github.com/Wei-Shaw/sub2api
|
||||
|
||||
go 1.26.5
|
||||
go 1.26.6
|
||||
|
||||
require (
|
||||
entgo.io/ent v0.14.5
|
||||
|
||||
@@ -1703,6 +1703,10 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
// 环境变量支持
|
||||
viper.AutomaticEnv()
|
||||
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
|
||||
if tz, ok := os.LookupEnv("TZ"); ok && strings.TrimSpace(tz) != "" {
|
||||
// AutomaticEnv 会先把 timezone 映射到 TIMEZONE;显式 Set 保证标准 TZ 变量优先。
|
||||
viper.Set("timezone", strings.TrimSpace(tz))
|
||||
}
|
||||
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
|
||||
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
|
||||
}
|
||||
|
||||
@@ -23,6 +23,38 @@ func resetViperWithJWTSecret(t *testing.T) {
|
||||
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
|
||||
}
|
||||
|
||||
func TestLoadTimezonePrecedence(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
fileTimezone string
|
||||
timezoneEnv string
|
||||
tzEnv string
|
||||
want string
|
||||
}{
|
||||
{name: "default", want: "Asia/Shanghai"},
|
||||
{name: "config_file", fileTimezone: "Europe/London", want: "Europe/London"},
|
||||
{name: "timezone_env", fileTimezone: "Europe/London", timezoneEnv: "UTC", want: "UTC"},
|
||||
{name: "tz_env", fileTimezone: "Europe/London", timezoneEnv: "UTC", tzEnv: "America/New_York", want: "America/New_York"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("TIMEZONE", tt.timezoneEnv)
|
||||
t.Setenv("TZ", tt.tzEnv)
|
||||
if tt.fileTimezone != "" {
|
||||
configFile := filepath.Join(t.TempDir(), "config.yaml")
|
||||
require.NoError(t, os.WriteFile(configFile, []byte("timezone: "+tt.fileTimezone+"\n"), 0o600))
|
||||
t.Setenv("CONFIG_FILE", configFile)
|
||||
}
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.want, cfg.Timezone)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadServerTimingConfig(t *testing.T) {
|
||||
t.Run("disabled by default", func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
|
||||
@@ -153,12 +153,12 @@ func (h *BackupHandler) GetDownloadURL(c *gin.Context) {
|
||||
response.BadRequest(c, "backup ID is required")
|
||||
return
|
||||
}
|
||||
url, err := h.backupService.GetBackupDownloadURL(c.Request.Context(), backupID)
|
||||
download, err := h.backupService.GetBackupDownloadURL(c.Request.Context(), backupID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{"url": url})
|
||||
response.Success(c, download)
|
||||
}
|
||||
|
||||
// ─── 恢复操作(需要重新输入管理员密码) ───
|
||||
|
||||
@@ -33,7 +33,7 @@ type createChannelRequest struct {
|
||||
GroupIDs []int64 `json:"group_ids"`
|
||||
ModelPricing []channelModelPricingRequest `json:"model_pricing"`
|
||||
ModelMapping map[string]map[string]string `json:"model_mapping"`
|
||||
BillingModelSource string `json:"billing_model_source" binding:"omitempty,oneof=requested upstream channel_mapped"`
|
||||
BillingModelSource string `json:"billing_model_source" binding:"omitempty,oneof=requested upstream channel_mapped response_model"`
|
||||
RestrictModels bool `json:"restrict_models"`
|
||||
Features string `json:"features"`
|
||||
FeaturesConfig map[string]any `json:"features_config"`
|
||||
@@ -48,7 +48,7 @@ type updateChannelRequest struct {
|
||||
GroupIDs *[]int64 `json:"group_ids"`
|
||||
ModelPricing *[]channelModelPricingRequest `json:"model_pricing"`
|
||||
ModelMapping map[string]map[string]string `json:"model_mapping"`
|
||||
BillingModelSource string `json:"billing_model_source" binding:"omitempty,oneof=requested upstream channel_mapped"`
|
||||
BillingModelSource string `json:"billing_model_source" binding:"omitempty,oneof=requested upstream channel_mapped response_model"`
|
||||
RestrictModels *bool `json:"restrict_models"`
|
||||
Features *string `json:"features"`
|
||||
FeaturesConfig map[string]any `json:"features_config"`
|
||||
|
||||
@@ -8,11 +8,11 @@ import (
|
||||
"log/slog"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/platform/liveattestation"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -96,15 +96,17 @@ func NewGroupHandler(adminService service.AdminService, dashboardService *servic
|
||||
|
||||
// CreateGroupRequest represents create group request
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok composite"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
|
||||
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok composite"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
|
||||
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
|
||||
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
|
||||
ModelPricing []service.ChannelModelPricing `json:"model_pricing"`
|
||||
// 图片生成计费配置(antigravity 和 gemini 平台使用,负数表示清除配置)
|
||||
AllowImageGeneration bool `json:"allow_image_generation"`
|
||||
AllowBatchImageGeneration bool `json:"allow_batch_image_generation"`
|
||||
@@ -162,16 +164,18 @@ type CreateGroupRequest struct {
|
||||
|
||||
// UpdateGroupRequest represents update group request
|
||||
type UpdateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Description *string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok composite"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
|
||||
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
|
||||
Name string `json:"name"`
|
||||
Description *string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok composite"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
|
||||
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
|
||||
LongContextPricingEnabled *bool `json:"long_context_pricing_enabled"`
|
||||
ModelPricing *[]service.ChannelModelPricing `json:"model_pricing"`
|
||||
// 图片生成计费配置(antigravity 和 gemini 平台使用,负数表示清除配置)
|
||||
AllowImageGeneration *bool `json:"allow_image_generation"`
|
||||
AllowBatchImageGeneration *bool `json:"allow_batch_image_generation"`
|
||||
@@ -508,6 +512,8 @@ func (h *GroupHandler) Create(c *gin.Context) {
|
||||
DailyLimitUSD: req.DailyLimitUSD.ToServiceInput(),
|
||||
WeeklyLimitUSD: req.WeeklyLimitUSD.ToServiceInput(),
|
||||
MonthlyLimitUSD: req.MonthlyLimitUSD.ToServiceInput(),
|
||||
LongContextPricingEnabled: req.LongContextPricingEnabled,
|
||||
ModelPricing: req.ModelPricing,
|
||||
AllowImageGeneration: req.AllowImageGeneration,
|
||||
AllowBatchImageGeneration: req.AllowBatchImageGeneration,
|
||||
ImageRateIndependent: req.ImageRateIndependent,
|
||||
@@ -635,6 +641,8 @@ func (h *GroupHandler) Update(c *gin.Context) {
|
||||
DailyLimitUSD: req.DailyLimitUSD.ToServiceInput(),
|
||||
WeeklyLimitUSD: req.WeeklyLimitUSD.ToServiceInput(),
|
||||
MonthlyLimitUSD: req.MonthlyLimitUSD.ToServiceInput(),
|
||||
LongContextPricingEnabled: req.LongContextPricingEnabled,
|
||||
ModelPricing: req.ModelPricing,
|
||||
AllowImageGeneration: req.AllowImageGeneration,
|
||||
AllowBatchImageGeneration: req.AllowBatchImageGeneration,
|
||||
ImageRateIndependent: req.ImageRateIndependent,
|
||||
@@ -726,12 +734,10 @@ func (h *GroupHandler) GetStats(c *gin.Context) {
|
||||
_ = groupID // TODO: implement actual stats
|
||||
}
|
||||
|
||||
// GetUsageSummary returns today's and cumulative cost for all groups.
|
||||
// GET /api/v1/admin/groups/usage-summary?timezone=Asia/Shanghai
|
||||
// GetUsageSummary returns today's, yesterday's, and cumulative cost for all groups.
|
||||
// GET /api/v1/admin/groups/usage-summary
|
||||
func (h *GroupHandler) GetUsageSummary(c *gin.Context) {
|
||||
userTZ := c.Query("timezone")
|
||||
now := timezone.NowInUserLocation(userTZ)
|
||||
todayStart := timezone.StartOfDayInUserLocation(now, userTZ)
|
||||
todayStart := service.GroupUsageTodayStart(time.Now())
|
||||
|
||||
results, err := h.dashboardService.GetGroupUsageSummary(c.Request.Context(), todayStart)
|
||||
if err != nil {
|
||||
|
||||
@@ -3,6 +3,8 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -62,6 +64,43 @@ type UpdateAPIKeyRequest struct {
|
||||
ResetRateLimitUsage *bool `json:"reset_rate_limit_usage"` // 重置限速用量
|
||||
}
|
||||
|
||||
func validAPIKeyLimit(v float64) bool { return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 }
|
||||
|
||||
func validateAPIKeyCreateRequest(req CreateAPIKeyRequest) error {
|
||||
if req.Quota != nil && !validAPIKeyLimit(*req.Quota) {
|
||||
return errors.New("invalid quota")
|
||||
}
|
||||
if req.RateLimit5h != nil && !validAPIKeyLimit(*req.RateLimit5h) {
|
||||
return errors.New("invalid rate_limit_5h")
|
||||
}
|
||||
if req.RateLimit1d != nil && !validAPIKeyLimit(*req.RateLimit1d) {
|
||||
return errors.New("invalid rate_limit_1d")
|
||||
}
|
||||
if req.RateLimit7d != nil && !validAPIKeyLimit(*req.RateLimit7d) {
|
||||
return errors.New("invalid rate_limit_7d")
|
||||
}
|
||||
if req.ExpiresInDays != nil && *req.ExpiresInDays <= 0 {
|
||||
return errors.New("invalid expires_in_days")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateAPIKeyUpdateRequest(req UpdateAPIKeyRequest) error {
|
||||
if req.Quota != nil && !validAPIKeyLimit(*req.Quota) {
|
||||
return errors.New("invalid quota")
|
||||
}
|
||||
if req.RateLimit5h != nil && !validAPIKeyLimit(*req.RateLimit5h) {
|
||||
return errors.New("invalid rate_limit_5h")
|
||||
}
|
||||
if req.RateLimit1d != nil && !validAPIKeyLimit(*req.RateLimit1d) {
|
||||
return errors.New("invalid rate_limit_1d")
|
||||
}
|
||||
if req.RateLimit7d != nil && !validAPIKeyLimit(*req.RateLimit7d) {
|
||||
return errors.New("invalid rate_limit_7d")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// List handles listing user's API keys with pagination
|
||||
// GET /api/v1/api-keys
|
||||
func (h *APIKeyHandler) List(c *gin.Context) {
|
||||
@@ -152,6 +191,10 @@ func (h *APIKeyHandler) Create(c *gin.Context) {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := validateAPIKeyCreateRequest(req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: numeric limits must be finite and non-negative, and expires_in_days must be greater than zero")
|
||||
return
|
||||
}
|
||||
|
||||
svcReq := service.CreateAPIKeyRequest{
|
||||
Name: req.Name,
|
||||
@@ -203,6 +246,10 @@ func (h *APIKeyHandler) Update(c *gin.Context) {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := validateAPIKeyUpdateRequest(req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: numeric limits must be finite and non-negative")
|
||||
return
|
||||
}
|
||||
|
||||
svcReq := service.UpdateAPIKeyRequest{
|
||||
IPWhitelist: req.IPWhitelist,
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateAPIKeyCreateRequest(t *testing.T) {
|
||||
zero, large, negative, nan, inf := 0.0, 1e100, -1.0, math.NaN(), math.Inf(1)
|
||||
positiveDays, zeroDays, negativeDays := 1, 0, -1
|
||||
require.NoError(t, validateAPIKeyCreateRequest(CreateAPIKeyRequest{}))
|
||||
require.NoError(t, validateAPIKeyCreateRequest(CreateAPIKeyRequest{Quota: &zero, RateLimit5h: &large, ExpiresInDays: &positiveDays}))
|
||||
|
||||
for _, req := range []CreateAPIKeyRequest{
|
||||
{Quota: &negative},
|
||||
{Quota: &nan},
|
||||
{RateLimit5h: &inf},
|
||||
{RateLimit1d: &negative},
|
||||
{RateLimit7d: &negative},
|
||||
{ExpiresInDays: &zeroDays},
|
||||
{ExpiresInDays: &negativeDays},
|
||||
} {
|
||||
require.Error(t, validateAPIKeyCreateRequest(req))
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateAPIKeyUpdateRequest(t *testing.T) {
|
||||
zero, large, negative, nan, inf := 0.0, 1e100, -1.0, math.NaN(), math.Inf(-1)
|
||||
require.NoError(t, validateAPIKeyUpdateRequest(UpdateAPIKeyRequest{Quota: &zero, RateLimit7d: &large}))
|
||||
|
||||
for _, req := range []UpdateAPIKeyRequest{
|
||||
{Quota: &negative},
|
||||
{RateLimit5h: &nan},
|
||||
{RateLimit1d: &inf},
|
||||
{RateLimit7d: &negative},
|
||||
} {
|
||||
require.Error(t, validateAPIKeyUpdateRequest(req))
|
||||
}
|
||||
}
|
||||
@@ -149,6 +149,7 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup {
|
||||
ProfitControlEnabled: g.ProfitControlEnabled,
|
||||
ProfitMinMargin: g.ProfitMinMargin,
|
||||
ProfitSafetyBuffer: g.ProfitSafetyBuffer,
|
||||
ModelPricing: g.ModelPricing,
|
||||
ModelRouting: g.ModelRouting,
|
||||
ModelRoutingEnabled: g.ModelRoutingEnabled,
|
||||
MCPXMLInject: g.MCPXMLInject,
|
||||
@@ -184,6 +185,7 @@ func groupFromServiceBase(g *service.Group) Group {
|
||||
DailyLimitUSD: g.DailyLimitUSD,
|
||||
WeeklyLimitUSD: g.WeeklyLimitUSD,
|
||||
MonthlyLimitUSD: g.MonthlyLimitUSD,
|
||||
LongContextPricingEnabled: g.LongContextPricingEnabled,
|
||||
AllowImageGeneration: g.AllowImageGeneration,
|
||||
AllowBatchImageGeneration: g.AllowBatchImageGeneration,
|
||||
ImageRateIndependent: g.ImageRateIndependent,
|
||||
|
||||
@@ -96,10 +96,11 @@ type Group struct {
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
Status string `json:"status"`
|
||||
|
||||
SubscriptionType string `json:"subscription_type"`
|
||||
DailyLimitUSD *float64 `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
|
||||
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
|
||||
SubscriptionType string `json:"subscription_type"`
|
||||
DailyLimitUSD *float64 `json:"daily_limit_usd"`
|
||||
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
|
||||
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
|
||||
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
|
||||
|
||||
// 图片生成计费配置(仅 antigravity 平台使用)
|
||||
AllowImageGeneration bool `json:"allow_image_generation"`
|
||||
@@ -164,9 +165,10 @@ type AdminGroup struct {
|
||||
// 分组利润控制(五个 token 平台分组可启用;margin/buffer 为小数存储)。
|
||||
// 仅管理员可见:这三个字段与同响应中的 rate_multiplier 相乘即可反推出
|
||||
// 运营方的上游成本上限,属于内部经营信息,不得下放到 dto.Group。
|
||||
ProfitControlEnabled bool `json:"profit_control_enabled"`
|
||||
ProfitMinMargin float64 `json:"profit_min_margin"`
|
||||
ProfitSafetyBuffer float64 `json:"profit_safety_buffer"`
|
||||
ProfitControlEnabled bool `json:"profit_control_enabled"`
|
||||
ProfitMinMargin float64 `json:"profit_min_margin"`
|
||||
ProfitSafetyBuffer float64 `json:"profit_safety_buffer"`
|
||||
ModelPricing []service.ChannelModelPricing `json:"model_pricing"`
|
||||
|
||||
// 模型路由配置(仅 anthropic 平台使用)
|
||||
ModelRouting map[string][]int64 `json:"model_routing"`
|
||||
|
||||
@@ -36,6 +36,9 @@ const (
|
||||
maxSameAccountRetries = 3
|
||||
// sameAccountRetryDelay 同账号重试间隔
|
||||
sameAccountRetryDelay = 500 * time.Millisecond
|
||||
// maxRequestScopedRetryDelay 限制请求级瞬时错误的指数退避上限,避免高重试配置
|
||||
// 将单次请求拖入分钟级等待。
|
||||
maxRequestScopedRetryDelay = 8 * time.Second
|
||||
// singleAccountBackoffDelay 单账号分组 503 退避重试固定延时。
|
||||
// Service 层在 SingleAccountRetry 模式下已做充分原地重试(最多 3 次、总等待 30s),
|
||||
// Handler 层只需短暂间隔后重新进入 Service 层即可。
|
||||
@@ -52,6 +55,21 @@ const (
|
||||
// 语义上等同于「无可用账号」:候选账号都不满足分组的利润约束。
|
||||
const profitVetoExhaustedMessage = "No available accounts: all candidates rejected by group profit control"
|
||||
|
||||
func sameAccountRetryDelayFor(failoverErr *service.UpstreamFailoverError, retryCount int) time.Duration {
|
||||
if failoverErr == nil || !failoverErr.RequestScopedTransient || retryCount <= 1 {
|
||||
return sameAccountRetryDelay
|
||||
}
|
||||
|
||||
delay := sameAccountRetryDelay
|
||||
for i := 1; i < retryCount; i++ {
|
||||
if delay >= maxRequestScopedRetryDelay/2 {
|
||||
return maxRequestScopedRetryDelay
|
||||
}
|
||||
delay *= 2
|
||||
}
|
||||
return delay
|
||||
}
|
||||
|
||||
// FailoverState 跨循环迭代共享的 failover 状态
|
||||
type FailoverState struct {
|
||||
SwitchCount int
|
||||
@@ -149,13 +167,15 @@ func (s *FailoverState) HandleFailoverError(
|
||||
// 重试次数上限 retryLimit 由调用方传入(账号级 pool_mode_retry_count 配置)。
|
||||
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit {
|
||||
s.SameAccountRetryCount[accountID]++
|
||||
retryDelay := sameAccountRetryDelayFor(failoverErr, s.SameAccountRetryCount[accountID])
|
||||
logger.FromContext(ctx).Warn("gateway.failover_same_account_retry",
|
||||
zap.Int64("account_id", accountID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("same_account_retry_count", s.SameAccountRetryCount[accountID]),
|
||||
zap.Int("same_account_retry_max", retryLimit),
|
||||
zap.Duration("retry_delay", retryDelay),
|
||||
)
|
||||
if !sleepWithContext(ctx, sameAccountRetryDelay) {
|
||||
if !sleepWithContext(ctx, retryDelay) {
|
||||
return FailoverCanceled
|
||||
}
|
||||
return FailoverContinue
|
||||
|
||||
@@ -31,6 +31,35 @@ func (m *mockTempUnscheduler) TempUnscheduleRetryableError(_ context.Context, ac
|
||||
m.calls = append(m.calls, tempUnscheduleCall{accountID: accountID, failoverErr: failoverErr})
|
||||
}
|
||||
|
||||
func TestSameAccountRetryDelayFor(t *testing.T) {
|
||||
capacityErr := &service.UpstreamFailoverError{RequestScopedTransient: true}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
retryCount int
|
||||
want time.Duration
|
||||
}{
|
||||
{name: "first retry", retryCount: 1, want: 500 * time.Millisecond},
|
||||
{name: "second retry", retryCount: 2, want: time.Second},
|
||||
{name: "third retry", retryCount: 3, want: 2 * time.Second},
|
||||
{name: "fourth retry", retryCount: 4, want: 4 * time.Second},
|
||||
{name: "fifth retry", retryCount: 5, want: 8 * time.Second},
|
||||
{name: "capped retry", retryCount: 10, want: 8 * time.Second},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
require.Equal(t, tc.want, sameAccountRetryDelayFor(capacityErr, tc.retryCount))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("non request scoped errors keep fixed delay", func(t *testing.T) {
|
||||
require.Equal(t, 500*time.Millisecond, sameAccountRetryDelayFor(&service.UpstreamFailoverError{}, 10))
|
||||
})
|
||||
|
||||
t.Run("nil error keeps fixed delay", func(t *testing.T) {
|
||||
require.Equal(t, 500*time.Millisecond, sameAccountRetryDelayFor(nil, 10))
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -1249,7 +1249,7 @@ func writeGrokModelsList(c *gin.Context, modelIDs []string) {
|
||||
|
||||
func grokModelSupportsConfigurableReasoning(modelID string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(modelID)) {
|
||||
case "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
|
||||
case "grok-4.6", "grok-4.6-latest", "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
|
||||
@@ -28,12 +28,8 @@ const (
|
||||
)
|
||||
|
||||
func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
type webSearchReq struct {
|
||||
Query string `json:"query" binding:"required"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
var req webSearchReq
|
||||
isXSearch := c.GetBool("grok_x_search_endpoint")
|
||||
var req grokStandaloneSearchRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"type": "invalid_request_error",
|
||||
@@ -41,7 +37,28 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
}})
|
||||
return
|
||||
}
|
||||
req.MaxResults = normalizeGrokWebSearchMaxResults(req.MaxResults)
|
||||
query := strings.TrimSpace(req.Query)
|
||||
if query == "" {
|
||||
query = strings.TrimSpace(req.Input)
|
||||
}
|
||||
if query == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"type": "invalid_request_error",
|
||||
"message": "query is required",
|
||||
}})
|
||||
return
|
||||
}
|
||||
req.Query = query
|
||||
maxResults := 0
|
||||
if req.MaxResults != nil {
|
||||
maxResults = *req.MaxResults
|
||||
}
|
||||
maxResults = normalizeGrokWebSearchMaxResults(maxResults)
|
||||
searchModel := resolveGrokStandaloneSearchModel()
|
||||
searchLabel := "web_search"
|
||||
if isXSearch {
|
||||
searchLabel = "x_search"
|
||||
}
|
||||
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil {
|
||||
@@ -55,7 +72,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
if apiKey.Group == nil || apiKey.Group.Platform != "grok" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"type": "invalid_request_error",
|
||||
"message": "web search is only supported for grok groups",
|
||||
"message": searchLabel + " is only supported for grok groups",
|
||||
}})
|
||||
return
|
||||
}
|
||||
@@ -79,7 +96,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
"role": "user", "content": req.Query,
|
||||
}},
|
||||
})
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, xai.DefaultTextModel, auditBody); decision != nil && !decision.AllowNextStage {
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, searchModel, auditBody); decision != nil && !decision.AllowNextStage {
|
||||
status := decision.HTTPStatus
|
||||
if status == 0 {
|
||||
status = http.StatusForbidden
|
||||
@@ -123,7 +140,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
// First attempt + up to 3 failover accounts (max 4 total).
|
||||
for attempt := 0; attempt < 4; attempt++ {
|
||||
selected, selectErr := h.gatewayService.SelectAccountWithLoadAwareness(
|
||||
c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0,
|
||||
c.Request.Context(), groupID, "", searchModel, failedAccounts, "", 0,
|
||||
)
|
||||
if selectErr != nil {
|
||||
if attempt == 0 {
|
||||
@@ -159,7 +176,11 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
account = selected.Account
|
||||
accountReleaseFunc = release
|
||||
|
||||
nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults)
|
||||
if isXSearch {
|
||||
nativeResp, providerName, err = h.doGrokNativeXSearch(c.Request.Context(), c, account, req, searchModel, maxResults)
|
||||
} else {
|
||||
nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, maxResults, searchModel)
|
||||
}
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
@@ -198,7 +219,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
// Request IDs are billing idempotency keys, so they must be unique per invocation.
|
||||
// Query/IP/UA hashes would collapse repeated identical searches into one charge.
|
||||
searchRequestID := "web_search:" + uuid.NewString()
|
||||
searchRequestID := searchLabel + ":" + uuid.NewString()
|
||||
if apiKey.Group != nil {
|
||||
if p := apiKey.Group.GetSearchPricePer1k(); p != nil && *p == 0 {
|
||||
logger.L().With(
|
||||
@@ -211,7 +232,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{
|
||||
RequestID: searchRequestID,
|
||||
Model: "grok-web-search",
|
||||
Model: "grok-" + strings.ReplaceAll(searchLabel, "_", "-"),
|
||||
SearchCount: 1,
|
||||
Duration: 0,
|
||||
},
|
||||
@@ -240,7 +261,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
"query": req.Query,
|
||||
"results": nativeResp.Results,
|
||||
"provider": providerName,
|
||||
"max_results": req.MaxResults,
|
||||
"max_results": maxResults,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -299,13 +320,13 @@ func (h *GatewayHandler) acquireWebSearchAccountSlot(
|
||||
|
||||
// doGrokNativeWebSearch executes web search using the Grok account's native capability
|
||||
// by calling the responses endpoint with web_search tool, then normalizes sources to unified format.
|
||||
func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Context, account *service.Account, query string, maxResults int) (*websearch.SearchResponse, string, error) {
|
||||
func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Context, account *service.Account, query string, maxResults int, model string) (*websearch.SearchResponse, string, error) {
|
||||
maxResults = normalizeGrokWebSearchMaxResults(maxResults)
|
||||
|
||||
// Build a minimal responses request that triggers Grok web search tool.
|
||||
// Ask for structured metadata because xAI action.sources commonly contains URLs only.
|
||||
searchBody := map[string]any{
|
||||
"model": xai.DefaultTextModel,
|
||||
"model": xai.ResolveDefaultTextModel(model),
|
||||
"input": buildGrokWebSearchPrompt(query, maxResults),
|
||||
"tools": []map[string]any{{"type": "web_search"}},
|
||||
"include": []string{"web_search_call.action.sources"},
|
||||
@@ -329,6 +350,23 @@ func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Conte
|
||||
}, "grok-native", nil
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) doGrokNativeXSearch(ctx context.Context, c *gin.Context, account *service.Account, req grokStandaloneSearchRequest, model string, maxResults int) (*websearch.SearchResponse, string, error) {
|
||||
maxResults = normalizeGrokWebSearchMaxResults(maxResults)
|
||||
bodyBytes, err := buildGrokXSearchResponsesBody(req, model)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
respBytes, err := h.gatewayService.DoGrokNativeResponsesJSON(ctx, account, bodyBytes)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
results := extractGrokWebSearchSources(respBytes, maxResults)
|
||||
return &websearch.SearchResponse{
|
||||
Results: results,
|
||||
Query: req.Query,
|
||||
}, "grok-native", nil
|
||||
}
|
||||
|
||||
func normalizeGrokWebSearchMaxResults(maxResults int) int {
|
||||
if maxResults <= 0 {
|
||||
return defaultGrokWebSearchResults
|
||||
@@ -377,7 +415,8 @@ func extractGrokWebSearchSources(body []byte, maxResults int) []websearch.Search
|
||||
|
||||
output := gjson.GetBytes(body, "output")
|
||||
output.ForEach(func(_, item gjson.Result) bool {
|
||||
if item.Get("type").String() == "web_search_call" {
|
||||
callType := item.Get("type").String()
|
||||
if callType == "web_search_call" || callType == "x_search_call" {
|
||||
sources := item.Get("action.sources")
|
||||
if sources.IsArray() {
|
||||
sources.ForEach(func(_, src gjson.Result) bool {
|
||||
|
||||
@@ -89,7 +89,7 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
model = "grok-voice-latest"
|
||||
}
|
||||
started := time.Now()
|
||||
proxyErr := h.gatewayService.ProxyGrokRealtime(c.Request.Context(), c, conn, selection.Account, token, model)
|
||||
audioObserved, proxyErr := h.gatewayService.ProxyGrokRealtime(c.Request.Context(), c, conn, selection.Account, token, model)
|
||||
elapsed := time.Since(started)
|
||||
if proxyErr != nil {
|
||||
reqLog.Info("grok_realtime.proxy_failed", zap.Error(proxyErr))
|
||||
@@ -98,20 +98,23 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// A relay normally returns a close error when either side closes normally.
|
||||
// Those sessions still consumed upstream audio time and must be billed.
|
||||
if elapsed > 0 {
|
||||
result := &service.OpenAIForwardResult{
|
||||
// One durable id per WS session so retries cannot collapse or double under client ids.
|
||||
RequestID: service.StableGrokRealtimeBillingRequestID(""),
|
||||
Model: model,
|
||||
Duration: elapsed,
|
||||
AudioUsage: &service.AudioUsage{Mode: "realtime", DurationOrUnits: elapsed.Minutes()},
|
||||
}
|
||||
if result := grokRealtimeBillingResult(model, elapsed, audioObserved); result != nil {
|
||||
h.recordGrokVoiceUsage(c, apiKey, selection.Account, subscription, "realtime", nil, result)
|
||||
}
|
||||
}
|
||||
|
||||
func grokRealtimeBillingResult(model string, elapsed time.Duration, audioObserved bool) *service.OpenAIForwardResult {
|
||||
if !audioObserved || elapsed <= 0 {
|
||||
return nil
|
||||
}
|
||||
return &service.OpenAIForwardResult{
|
||||
RequestID: service.StableGrokRealtimeBillingRequestID(""),
|
||||
Model: model,
|
||||
Duration: elapsed,
|
||||
AudioUsage: &service.AudioUsage{Mode: "realtime", DurationOrUnits: elapsed.Minutes()},
|
||||
}
|
||||
}
|
||||
|
||||
func isExpectedGrokRealtimeClose(err error) bool {
|
||||
if err == nil {
|
||||
return true
|
||||
|
||||
@@ -4,6 +4,7 @@ package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
coderws "github.com/coder/websocket"
|
||||
)
|
||||
@@ -23,3 +24,29 @@ func TestIsExpectedGrokRealtimeClose(t *testing.T) {
|
||||
t.Fatal("policy violations must not be treated as billable normal closes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokRealtimeBillingResultRequiresObservedAudio(t *testing.T) {
|
||||
if grokRealtimeBillingResult("grok-voice-latest", time.Second, false) != nil {
|
||||
t.Fatal("a session without observed audio must not be billed")
|
||||
}
|
||||
if grokRealtimeBillingResult("grok-voice-latest", 0, true) != nil {
|
||||
t.Fatal("zero-duration sessions must not be billed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokRealtimeBillingResultUsesForcedUniqueID(t *testing.T) {
|
||||
first := grokRealtimeBillingResult("grok-voice-latest", 90*time.Second, true)
|
||||
second := grokRealtimeBillingResult("grok-voice-latest", 90*time.Second, true)
|
||||
if first == nil || second == nil {
|
||||
t.Fatal("observed audio sessions should be billable")
|
||||
}
|
||||
if first.RequestID == "" {
|
||||
t.Fatalf("unexpected billing request ID %q", first.RequestID)
|
||||
}
|
||||
if first.RequestID == second.RequestID {
|
||||
t.Fatal("independent realtime connections must not share a billing request ID")
|
||||
}
|
||||
if first.AudioUsage == nil || first.AudioUsage.Mode != "realtime" || first.AudioUsage.DurationOrUnits != 1.5 {
|
||||
t.Fatalf("unexpected audio usage: %#v", first.AudioUsage)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -288,16 +288,18 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
sameAccountRetryCount[account.ID]++
|
||||
retryDelay := sameAccountRetryDelayFor(failoverErr, sameAccountRetryCount[account.ID])
|
||||
reqLog.Warn("openai_chat_completions.pool_mode_same_account_retry",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("retry_limit", retryLimit),
|
||||
zap.Int("retry_count", sameAccountRetryCount[account.ID]),
|
||||
zap.Duration("retry_delay", retryDelay),
|
||||
)
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
return
|
||||
case <-time.After(sameAccountRetryDelay):
|
||||
case <-time.After(retryDelay):
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -25,54 +25,211 @@ func newCompactBodySignalTestContext(t *testing.T, path string, body []byte) *gi
|
||||
|
||||
func TestNormalizeOpenAIResponsesCompactRequest_RemoteV2StaysOnResponses(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.6-sol",
|
||||
"stream":true,
|
||||
"store":true,
|
||||
"prompt_cache_key":"pck-signal-1",
|
||||
"reasoning":{"effort":"max","context":"all_turns"},
|
||||
"input":[
|
||||
{"type":"message","role":"user","content":"hello"},
|
||||
{"type":"compaction_trigger"}
|
||||
]
|
||||
}`)
|
||||
c := newCompactBodySignalTestContext(t, "/v1/responses", body)
|
||||
c.Request.Header.Set("x-codex-beta-features", "responses_websockets_v2, remote_compaction_v2, another_feature")
|
||||
tests := []struct {
|
||||
name string
|
||||
betaHeader string
|
||||
userAgent string
|
||||
}{
|
||||
{name: "headerless"},
|
||||
{name: "unrelated_header", betaHeader: "responses_websockets_v2"},
|
||||
{name: "wrong_case_header", betaHeader: "REMOTE_COMPACTION_V2"},
|
||||
{name: "declared_header", betaHeader: "remote_compaction_v2"},
|
||||
{name: "codex_cli_user_agent", userAgent: "codex_cli_rs/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color"},
|
||||
{name: "codex_desktop_user_agent", userAgent: "Codex Desktop/0.139.0 (Mac OS X 14; arm64) unknown"},
|
||||
}
|
||||
|
||||
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
|
||||
require.True(t, ok)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.6-sol",
|
||||
"stream":true,
|
||||
"store":true,
|
||||
"prompt_cache_key":"pck-signal-1",
|
||||
"reasoning":{"effort":"max","context":"all_turns"},
|
||||
"input":[
|
||||
{"type":"message","role":"user","content":"hello"},
|
||||
{"type":"compaction_trigger"}
|
||||
]
|
||||
}`)
|
||||
c := newCompactBodySignalTestContext(t, "/v1/responses", body)
|
||||
if tt.betaHeader != "" {
|
||||
c.Request.Header.Set("x-codex-beta-features", tt.betaHeader)
|
||||
}
|
||||
if tt.userAgent != "" {
|
||||
c.Request.Header.Set("User-Agent", tt.userAgent)
|
||||
}
|
||||
|
||||
require.Equal(t, "/v1/responses", c.Request.URL.Path)
|
||||
require.False(t, isOpenAIRemoteCompactPath(c))
|
||||
require.Equal(t, body, normalized)
|
||||
require.True(t, gjson.GetBytes(normalized, "stream").Bool())
|
||||
require.True(t, gjson.GetBytes(normalized, "store").Bool())
|
||||
require.Equal(t, "pck-signal-1", gjson.GetBytes(normalized, "prompt_cache_key").String())
|
||||
require.Equal(t, "max", gjson.GetBytes(normalized, "reasoning.effort").String())
|
||||
require.Equal(t, "all_turns", gjson.GetBytes(normalized, "reasoning.context").String())
|
||||
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
|
||||
require.True(t, ok)
|
||||
|
||||
reqStream, streamOK := parseOpenAICompatibleStream(normalized)
|
||||
require.True(t, streamOK)
|
||||
require.True(t, reqStream)
|
||||
require.Equal(t, "/v1/responses", c.Request.URL.Path)
|
||||
require.False(t, isOpenAILegacyCompactPath(c))
|
||||
require.Equal(t, body, normalized)
|
||||
require.True(t, gjson.GetBytes(normalized, "stream").Bool())
|
||||
require.True(t, gjson.GetBytes(normalized, "store").Bool())
|
||||
require.Equal(t, "pck-signal-1", gjson.GetBytes(normalized, "prompt_cache_key").String())
|
||||
require.Equal(t, "max", gjson.GetBytes(normalized, "reasoning.effort").String())
|
||||
require.Equal(t, "all_turns", gjson.GetBytes(normalized, "reasoning.context").String())
|
||||
legacyCompact := service.IsOpenAIResponsesCompactPath(c)
|
||||
nativeV2 := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(normalized)
|
||||
require.False(t, legacyCompact)
|
||||
require.True(t, nativeV2)
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityResponses,
|
||||
openAIResponsesRequiredCapabilityForRequest(false, nativeV2 || legacyCompact, service.PlatformOpenAI))
|
||||
|
||||
_, seedExists := c.Get(service.OpenAICompactSessionSeedKeyForTest())
|
||||
require.False(t, seedExists)
|
||||
_, streamMarkerExists := c.Get(service.OpenAICompactClientStreamKeyForTest())
|
||||
require.False(t, streamMarkerExists)
|
||||
reqStream, streamOK := parseOpenAICompatibleStream(normalized)
|
||||
require.True(t, streamOK)
|
||||
require.True(t, reqStream)
|
||||
|
||||
_, seedExists := c.Get(service.OpenAICompactSessionSeedKeyForTest())
|
||||
require.False(t, seedExists)
|
||||
_, streamMarkerExists := c.Get(service.OpenAICompactClientStreamKeyForTest())
|
||||
require.False(t, streamMarkerExists)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeOpenAIResponsesCompactRequest_RemoteV2PathAliasesStayOnResponses(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
body := []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`)
|
||||
for _, path := range []string{"/v1/responses/", "/backend-api/codex/responses"} {
|
||||
for _, path := range []string{
|
||||
"/v1/responses/",
|
||||
"/openai/v1/responses",
|
||||
"/responses",
|
||||
"/backend-api/codex/responses",
|
||||
} {
|
||||
t.Run(path, func(t *testing.T) {
|
||||
c := newCompactBodySignalTestContext(t, path, body)
|
||||
c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2")
|
||||
|
||||
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, path, c.Request.URL.Path)
|
||||
require.Equal(t, body, normalized)
|
||||
legacyCompact := service.IsOpenAIResponsesCompactPath(c)
|
||||
nativeV2 := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(normalized)
|
||||
require.False(t, legacyCompact)
|
||||
require.True(t, nativeV2)
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityResponses,
|
||||
openAIResponsesRequiredCapabilityForRequest(false, nativeV2 || legacyCompact, service.PlatformOpenAI))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesCompactionRoutingFlags(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
path string
|
||||
wantLegacyBefore bool
|
||||
wantNativeBefore bool
|
||||
wantLegacyAfter bool
|
||||
wantNativeAfter bool
|
||||
wantCapabilityAfter service.OpenAIEndpointCapability
|
||||
wantPathAfter string
|
||||
wantBodyUnchanged bool
|
||||
}{
|
||||
{
|
||||
name: "native_v2_stream_trigger",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
path: "/v1/responses",
|
||||
wantLegacyBefore: false,
|
||||
wantNativeBefore: true,
|
||||
wantLegacyAfter: false,
|
||||
wantNativeAfter: true,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
|
||||
wantPathAfter: "/v1/responses",
|
||||
wantBodyUnchanged: true,
|
||||
},
|
||||
{
|
||||
name: "native_v2_without_trigger",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"message","role":"user","content":"hello"}]}`),
|
||||
path: "/v1/responses",
|
||||
wantLegacyBefore: false,
|
||||
wantNativeBefore: false,
|
||||
wantLegacyAfter: false,
|
||||
wantNativeAfter: false,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityChatCompletions,
|
||||
wantPathAfter: "/v1/responses",
|
||||
wantBodyUnchanged: true,
|
||||
},
|
||||
{
|
||||
name: "explicit_compact",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
path: "/v1/responses/compact",
|
||||
wantLegacyBefore: true,
|
||||
wantNativeBefore: false,
|
||||
wantLegacyAfter: true,
|
||||
wantNativeAfter: false,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
|
||||
wantPathAfter: "/v1/responses/compact",
|
||||
},
|
||||
{
|
||||
name: "nested_compact",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
path: "/v1/responses/compact/detail",
|
||||
wantLegacyBefore: true,
|
||||
wantNativeBefore: false,
|
||||
wantLegacyAfter: true,
|
||||
wantNativeAfter: false,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
|
||||
wantPathAfter: "/v1/responses/compact/detail",
|
||||
},
|
||||
{
|
||||
name: "responses_subpath_with_native_signal",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
path: "/v1/responses/resp_123/responses",
|
||||
wantLegacyBefore: false,
|
||||
wantNativeBefore: false,
|
||||
wantLegacyAfter: false,
|
||||
wantNativeAfter: false,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityChatCompletions,
|
||||
wantPathAfter: "/v1/responses/resp_123/responses",
|
||||
wantBodyUnchanged: true,
|
||||
},
|
||||
{
|
||||
name: "stream_false_promotes",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","stream":false,"input":[{"type":"compaction_trigger"}]}`),
|
||||
path: "/v1/responses",
|
||||
wantLegacyBefore: false,
|
||||
wantNativeBefore: false,
|
||||
wantLegacyAfter: true,
|
||||
wantNativeAfter: false,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
|
||||
wantPathAfter: "/v1/responses/compact",
|
||||
},
|
||||
{
|
||||
name: "stream_absent_promotes",
|
||||
body: []byte(`{"model":"gpt-5.6-sol","input":[{"type":"compaction_trigger"}]}`),
|
||||
path: "/v1/responses",
|
||||
wantLegacyBefore: false,
|
||||
wantNativeBefore: false,
|
||||
wantLegacyAfter: true,
|
||||
wantNativeAfter: false,
|
||||
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
|
||||
wantPathAfter: "/v1/responses/compact",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := newCompactBodySignalTestContext(t, tt.path, tt.body)
|
||||
legacyBefore := service.IsOpenAIResponsesCompactPath(c)
|
||||
nativeBefore := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(tt.body)
|
||||
require.Equal(t, tt.wantLegacyBefore, legacyBefore)
|
||||
require.Equal(t, tt.wantNativeBefore, nativeBefore)
|
||||
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), tt.body)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, tt.wantPathAfter, c.Request.URL.Path)
|
||||
legacyAfter := service.IsOpenAIResponsesCompactPath(c)
|
||||
nativeAfter := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(normalized)
|
||||
require.Equal(t, tt.wantLegacyAfter, legacyAfter)
|
||||
require.Equal(t, tt.wantNativeAfter, nativeAfter)
|
||||
require.Equal(t, tt.wantCapabilityAfter,
|
||||
openAIResponsesRequiredCapabilityForRequest(false, nativeAfter || legacyAfter, service.PlatformOpenAI))
|
||||
if tt.wantBodyUnchanged {
|
||||
require.Equal(t, tt.body, normalized)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -106,31 +263,22 @@ func TestNormalizeOpenAIResponsesCompactRequest_NonRemoteV2BodySignalPromoted(t
|
||||
wantMarked bool
|
||||
}{
|
||||
{
|
||||
name: "no_header",
|
||||
body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
wantMarked: true,
|
||||
name: "stream_false_headerless",
|
||||
body: []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`),
|
||||
},
|
||||
{
|
||||
name: "unrelated_header",
|
||||
body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
betaHeader: "responses_websockets_v2",
|
||||
wantMarked: true,
|
||||
name: "stream_absent_headerless",
|
||||
body: []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`),
|
||||
},
|
||||
{
|
||||
name: "wrong_case_header",
|
||||
body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`),
|
||||
betaHeader: "REMOTE_COMPACTION_V2",
|
||||
wantMarked: true,
|
||||
},
|
||||
{
|
||||
name: "stream_false",
|
||||
name: "stream_false_declared_header",
|
||||
body: []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`),
|
||||
betaHeader: "remote_compaction_v2",
|
||||
},
|
||||
{
|
||||
name: "stream_absent",
|
||||
name: "stream_absent_wrong_case_header",
|
||||
body: []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`),
|
||||
betaHeader: "remote_compaction_v2",
|
||||
betaHeader: "REMOTE_COMPACTION_V2",
|
||||
},
|
||||
}
|
||||
|
||||
@@ -163,7 +311,7 @@ func TestNormalizeOpenAIResponsesCompactRequest_NoTriggerUntouched(t *testing.T)
|
||||
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "/v1/responses", c.Request.URL.Path)
|
||||
require.False(t, isOpenAIRemoteCompactPath(c))
|
||||
require.False(t, isOpenAILegacyCompactPath(c))
|
||||
require.Equal(t, body, normalized)
|
||||
require.True(t, gjson.GetBytes(normalized, "stream").Bool())
|
||||
}
|
||||
@@ -172,7 +320,6 @@ func TestNormalizeOpenAIResponsesCompactRequest_PathBasedNoDoubleSuffix(t *testi
|
||||
h := &OpenAIGatewayHandler{}
|
||||
body := []byte(`{"model":"gpt-5.5","stream":true,"store":true,"input":[{"type":"message","role":"user","content":"hello"}]}`)
|
||||
c := newCompactBodySignalTestContext(t, "/v1/responses/compact", body)
|
||||
c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2")
|
||||
|
||||
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
|
||||
require.True(t, ok)
|
||||
|
||||
@@ -105,21 +105,29 @@ func captureHandlerStructuredLog(t *testing.T) (*handlerInMemoryLogSink, func())
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsOpenAIRemoteCompactPath(t *testing.T) {
|
||||
require.False(t, isOpenAIRemoteCompactPath(nil))
|
||||
func TestIsOpenAILegacyCompactPath(t *testing.T) {
|
||||
require.False(t, isOpenAILegacyCompactPath(nil))
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil)
|
||||
require.True(t, isOpenAIRemoteCompactPath(c))
|
||||
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/responses/compact/", nil)
|
||||
require.True(t, isOpenAIRemoteCompactPath(c))
|
||||
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
require.False(t, isOpenAIRemoteCompactPath(c))
|
||||
for _, test := range []struct {
|
||||
path string
|
||||
want bool
|
||||
}{
|
||||
{path: "/v1/responses/compact", want: true},
|
||||
{path: "/v1/responses/compact/detail", want: true},
|
||||
{path: "/responses/compact/", want: true},
|
||||
{path: "/v1/responses", want: false},
|
||||
{path: "/openai/v1/responses", want: false},
|
||||
{path: "/responses", want: false},
|
||||
{path: "/backend-api/codex/responses", want: false},
|
||||
{path: "/v1/responses/resp_123/cancel", want: false},
|
||||
} {
|
||||
c.Request = httptest.NewRequest(http.MethodPost, test.path, nil)
|
||||
require.Equal(t, test.want, isOpenAILegacyCompactPath(c), test.path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogOpenAIRemoteCompactOutcome_Succeeded(t *testing.T) {
|
||||
|
||||
@@ -182,6 +182,16 @@ func openAIResponsesRequiredCapability(imageIntent bool, platform string) servic
|
||||
return service.OpenAIEndpointCapabilityChatCompletions
|
||||
}
|
||||
|
||||
// openAIResponsesRequiredCapabilityForRequest returns the endpoint capability
|
||||
// required by an image or Responses request. needsResponses includes both the
|
||||
// legacy /responses/compact endpoint and native remote compaction v2.
|
||||
func openAIResponsesRequiredCapabilityForRequest(imageIntent bool, needsResponses bool, platform string) service.OpenAIEndpointCapability {
|
||||
if needsResponses && platform == service.PlatformOpenAI {
|
||||
return service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
return openAIResponsesRequiredCapability(imageIntent, platform)
|
||||
}
|
||||
|
||||
func allowOpenAICompatibleMessagesDispatch(apiKey *service.APIKey) bool {
|
||||
if apiKey == nil || apiKey.Group == nil {
|
||||
return true
|
||||
@@ -288,6 +298,13 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
legacyCompact := service.IsOpenAIResponsesCompactPath(c)
|
||||
nativeV2 := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(body)
|
||||
if nativeV2 {
|
||||
// 原生 v2 压缩出站前补注 x-codex-beta-features: remote_compaction_v2,
|
||||
// 与真实 Codex 线型一致(网关链剥头后本级负责恢复,#5586)。
|
||||
service.MarkOpenAINativeCompactionV2(c)
|
||||
}
|
||||
// body-signal compact:上游 unary 等待期间向下游发 SSE 注释行心跳,防止
|
||||
// 反向代理空闲超时掐断长压缩连接(#3887)。首拍延迟一个心跳间隔,快速
|
||||
// 失败仍走 JSON+状态码链路;未标记客户端流式或间隔为 0 时是 no-op。
|
||||
@@ -377,6 +394,15 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
forwardBody := openAIModelMappedBody(body, channelMapping.Mapped, channelMapping.MappedModel, h.gatewayService.ReplaceModelInBody)
|
||||
seedOpenAIForwardImageIntentHint(c, channelMapping.Mapped, imageIntent)
|
||||
forwardModel := reqModel
|
||||
if channelMapping.Mapped {
|
||||
forwardModel = channelMapping.MappedModel
|
||||
}
|
||||
c.Request = c.Request.WithContext(service.WithOpenAIForwardModel(
|
||||
c.Request.Context(),
|
||||
forwardModel,
|
||||
legacyCompact,
|
||||
))
|
||||
|
||||
// 提前校验 function_call_output 是否具备可关联上下文,避免上游 400。
|
||||
if !h.validateFunctionCallOutputRequest(c, body, reqLog) {
|
||||
@@ -420,7 +446,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
|
||||
return
|
||||
}
|
||||
requireCompact := isOpenAIRemoteCompactPath(c)
|
||||
requireCompact := legacyCompact
|
||||
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
switchCount := 0
|
||||
@@ -437,7 +463,8 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// 仅对 OpenAI 平台生效:Grok 生图走独立的 forwardGrokResponses 路径,不应被过滤。
|
||||
// 复用前置权限与并发阶段在未修改 body 上确认的显式生图意图,避免大 tools 请求重复扫描。
|
||||
// 该判断已排除 Codex 被动 image_gen namespace,避免 CC-only 账号被误过滤(#4476)。
|
||||
requiredCapability := openAIResponsesRequiredCapability(imageIntent, requestPlatform)
|
||||
needsResponses := nativeV2 || legacyCompact
|
||||
requiredCapability := openAIResponsesRequiredCapabilityForRequest(imageIntent, needsResponses, requestPlatform)
|
||||
|
||||
// 分组利润控制:请求级装配定价上下文——pricingAt 固定本请求的
|
||||
// D 与计费高峰因子,选号、槽位终检与全部 failover 重入共用同一门与阈值。
|
||||
@@ -479,7 +506,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
)
|
||||
if len(failedAccountIDs) == 0 {
|
||||
if errors.Is(err, service.ErrNoAvailableCompactAccounts) {
|
||||
if legacyCompact && errors.Is(err, service.ErrNoAvailableCompactAccounts) {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "compact_not_supported", "No available accounts support /responses/compact", streamStarted)
|
||||
return
|
||||
@@ -610,16 +637,18 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
sameAccountRetryCount[account.ID]++
|
||||
retryDelay := sameAccountRetryDelayFor(failoverErr, sameAccountRetryCount[account.ID])
|
||||
reqLog.Warn("openai.pool_mode_same_account_retry",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("retry_limit", retryLimit),
|
||||
zap.Int("retry_count", sameAccountRetryCount[account.ID]),
|
||||
zap.Duration("retry_delay", retryDelay),
|
||||
)
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
return
|
||||
case <-time.After(sameAccountRetryDelay):
|
||||
case <-time.After(retryDelay):
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -733,12 +762,8 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
func isOpenAIRemoteCompactPath(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
normalizedPath := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/")
|
||||
return strings.HasSuffix(normalizedPath, "/responses/compact")
|
||||
func isOpenAILegacyCompactPath(c *gin.Context) bool {
|
||||
return service.IsOpenAIResponsesCompactPath(c)
|
||||
}
|
||||
|
||||
// isBareOpenAIResponsesPath 仅匹配裸 /responses 端点(无 /compact 等子路径),
|
||||
@@ -748,32 +773,27 @@ func isBareOpenAIResponsesPath(c *gin.Context) bool {
|
||||
return false
|
||||
}
|
||||
normalizedPath := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/")
|
||||
return strings.HasSuffix(normalizedPath, "/responses")
|
||||
}
|
||||
|
||||
func isOpenAIRemoteCompactionV2Request(c *gin.Context, body []byte) bool {
|
||||
stream, valid := parseOpenAICompatibleStream(body)
|
||||
if !valid || !stream || c == nil || c.Request == nil {
|
||||
switch normalizedPath {
|
||||
case EndpointResponses, "/openai/v1/responses", "/responses", "/backend-api/codex/responses":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
for _, header := range c.Request.Header.Values("x-codex-beta-features") {
|
||||
for _, feature := range strings.Split(header, ",") {
|
||||
if strings.TrimSpace(feature) == "remote_compaction_v2" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isOpenAIRemoteCompactionV2Request(body []byte) bool {
|
||||
stream, valid := parseOpenAICompatibleStream(body)
|
||||
return valid && stream && service.HasCompactionTriggerInInput(body)
|
||||
}
|
||||
|
||||
// normalizeOpenAIResponsesCompactRequest keeps Codex remote compaction v2 on
|
||||
// its native streaming /responses wire and preserves the legacy body-signal
|
||||
// promotion for clients that do not explicitly advertise that protocol.
|
||||
// promotion for non-streaming requests.
|
||||
// 返回归一化后的 body;ok=false 表示错误响应已写出,调用方应直接 return。
|
||||
func (h *OpenAIGatewayHandler) normalizeOpenAIResponsesCompactRequest(c *gin.Context, reqLog *zap.Logger, body []byte) ([]byte, bool) {
|
||||
isCompactRequest := service.IsOpenAIResponsesCompactPathForTest(c)
|
||||
isCompactRequest := isOpenAILegacyCompactPath(c)
|
||||
if !isCompactRequest && isBareOpenAIResponsesPath(c) && service.HasCompactionTriggerInInput(body) {
|
||||
if isOpenAIRemoteCompactionV2Request(c, body) {
|
||||
if isOpenAIRemoteCompactionV2Request(body) {
|
||||
return body, true
|
||||
}
|
||||
c.Request.URL.Path = strings.TrimRight(c.Request.URL.Path, "/") + "/compact"
|
||||
@@ -802,7 +822,7 @@ func (h *OpenAIGatewayHandler) normalizeOpenAIResponsesCompactRequest(c *gin.Con
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) logOpenAIRemoteCompactOutcome(c *gin.Context, startedAt time.Time) {
|
||||
if !isOpenAIRemoteCompactPath(c) {
|
||||
if !isOpenAILegacyCompactPath(c) {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1158,16 +1178,18 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
sameAccountRetryCount[account.ID]++
|
||||
retryDelay := sameAccountRetryDelayFor(failoverErr, sameAccountRetryCount[account.ID])
|
||||
reqLog.Warn("openai_messages.pool_mode_same_account_retry",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("retry_limit", retryLimit),
|
||||
zap.Int("retry_count", sameAccountRetryCount[account.ID]),
|
||||
zap.Duration("retry_delay", retryDelay),
|
||||
)
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
return
|
||||
case <-time.After(sameAccountRetryDelay):
|
||||
case <-time.After(retryDelay):
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -2023,6 +2045,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
MaxReasoningEffort: maxReasoningEffort,
|
||||
ReasoningEffortMappings: reasoningEffortMappings,
|
||||
BeforeRequest: func(turn int, payload []byte, originalModel string) error {
|
||||
c.Set(securityAuditWSTurnContextKey, turn)
|
||||
if turn == 1 {
|
||||
return nil
|
||||
}
|
||||
@@ -2698,7 +2721,11 @@ func (h *OpenAIGatewayHandler) ensureForwardErrorResponse(c *gin.Context, stream
|
||||
imageKeepalivePaddingOnly = adjustedSize < 0
|
||||
imageKeepaliveResponseWritten = adjustedSize >= 0
|
||||
}
|
||||
if service.IsResponseCommitted(c) || (!compactKeepaliveCommitted && imageKeepaliveResponseWritten) {
|
||||
compactKeepaliveHasMeaningfulOutput := compactKeepaliveCommitted && service.OpenAICompactKeepaliveAdjustedWrittenSize(c) > 0
|
||||
// Compact keepalive may have committed 200 headers without writing a
|
||||
// semantic SSE event. In that case the Responses stream still needs its
|
||||
// protocol-correct terminal response.failed event.
|
||||
if (service.IsResponseCommitted(c) && (!compactKeepaliveCommitted || compactKeepaliveHasMeaningfulOutput)) || (!compactKeepaliveCommitted && imageKeepaliveResponseWritten) {
|
||||
return false
|
||||
}
|
||||
if c.Writer.Written() && !imageKeepalivePaddingOnly {
|
||||
|
||||
@@ -335,6 +335,26 @@ func TestOpenAIEnsureForwardErrorResponse_AfterDeltaAppendsSingleValidResponseFa
|
||||
require.Equal(t, 1, errorEvents)
|
||||
}
|
||||
|
||||
func TestOpenAIEnsureForwardErrorResponse_CompactKeepaliveOnlyWritesResponseFailed(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
|
||||
service.MarkOpenAICompactClientStream(c)
|
||||
|
||||
stop := service.StartOpenAICompactSSEKeepalive(c, 5*time.Millisecond)
|
||||
defer stop()
|
||||
before := service.OpenAICompactKeepaliveAdjustedWrittenSize(c)
|
||||
require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond)
|
||||
require.Equal(t, before, service.OpenAICompactKeepaliveAdjustedWrittenSize(c))
|
||||
|
||||
h := &OpenAIGatewayHandler{}
|
||||
require.True(t, h.ensureForwardErrorResponse(c, false))
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), "event: response.failed\n")
|
||||
require.NotContains(t, w.Body.String(), "event: error\n")
|
||||
}
|
||||
|
||||
func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepaliveWritesSingleJSONFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
@@ -1816,6 +1836,37 @@ func (u *openAIHTTPPassthroughFailoverUpstream) calls() []int64 {
|
||||
return append([]int64(nil), u.accountIDs...)
|
||||
}
|
||||
|
||||
type openAIHTTPPassthroughAuthFailoverUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
accountIDs []int64
|
||||
statusCode int
|
||||
}
|
||||
|
||||
func (u *openAIHTTPPassthroughAuthFailoverUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
|
||||
u.mu.Lock()
|
||||
u.accountIDs = append(u.accountIDs, accountID)
|
||||
u.mu.Unlock()
|
||||
if accountID == 9911 {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_healthy","object":"response","model":"gpt-5.2","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`)),
|
||||
}, nil
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: u.statusCode,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"upstream credential rejected"}}`)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *openAIHTTPPassthroughAuthFailoverUpstream) calls() []int64 {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
return append([]int64(nil), u.accountIDs...)
|
||||
}
|
||||
|
||||
type openAIHTTPPassthroughSSERateLimitUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
@@ -2050,6 +2101,108 @@ func TestOpenAIResponses_APIKeyPassthroughPool5xxRetriesThenExhaustsMaxSwitches(
|
||||
require.Equal(t, "Upstream service temporarily unavailable", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponses_APIKeyPassthroughPoolAuthFailureRetriesThenSwitchesToHealthyAccount(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
}{
|
||||
{name: "401", statusCode: http.StatusUnauthorized},
|
||||
{name: "403", statusCode: http.StatusForbidden},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
groupID := int64(4203)
|
||||
accounts := []service.Account{
|
||||
{
|
||||
ID: 9910, Name: "pool-api-key", Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-pool",
|
||||
"base_url": "https://api.example.test",
|
||||
"pool_mode": true,
|
||||
"pool_mode_retry_count": float64(1),
|
||||
"pool_mode_retry_status_codes": []any{float64(tt.statusCode)},
|
||||
},
|
||||
Extra: map[string]any{"openai_passthrough": true},
|
||||
},
|
||||
{
|
||||
ID: 9911, Name: "fallback-api-key", Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 2,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-fallback",
|
||||
"base_url": "https://api.example.test",
|
||||
},
|
||||
Extra: map[string]any{"openai_passthrough": true},
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
cfg.Default.RateMultiplier = 1
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Gateway.MaxAccountSwitches = 1
|
||||
|
||||
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
|
||||
upstream := &openAIHTTPPassthroughAuthFailoverUpstream{statusCode: tt.statusCode}
|
||||
rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
|
||||
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheSvc.Stop)
|
||||
gatewaySvc := service.NewOpenAIGatewayService(
|
||||
accountRepo,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
cfg,
|
||||
nil,
|
||||
nil,
|
||||
service.NewBillingService(cfg, nil),
|
||||
rateLimitSvc,
|
||||
billingCacheSvc,
|
||||
upstream,
|
||||
&service.DeferredService{},
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
h := NewOpenAIGatewayHandler(
|
||||
gatewaySvc,
|
||||
service.NewConcurrencyService(nil),
|
||||
billingCacheSvc,
|
||||
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
cfg,
|
||||
)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.2","input":"hello","stream":false}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 1803, GroupID: &groupID,
|
||||
User: &service.User{ID: 1703, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
|
||||
})
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1703, Concurrency: 0})
|
||||
|
||||
h.Responses(c)
|
||||
|
||||
require.Equal(t, []int64{9910, 9910, 9911}, upstream.calls())
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "resp_healthy", gjson.GetBytes(rec.Body.Bytes(), "id").String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponses_APIKeyPassthroughSSERateLimitUsesConfiguredPoolRetry(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
groupID := int64(4204)
|
||||
|
||||
@@ -5,8 +5,7 @@ package handler
|
||||
// 槽位终检与生图跳门回归(handler 半程):
|
||||
// - 槽位获取成功后的利润终检:越线账号释放槽位并要求调用方排除重选,
|
||||
// 不写响应、不绑定粘连;
|
||||
// - openAIResponsesRequiredCapability 的生图意图映射钉死(scheduler 的
|
||||
// 跳门条件依赖 CapabilityResponses ⇔ 显式生图意图这一耦合)。
|
||||
// - openAIResponsesRequiredCapability 的请求能力映射覆盖生图与原生远程压缩。
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -141,10 +140,10 @@ func TestAcquireResponsesAccountSlotProfitRecheck(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
// scheduler 跳门条件依赖"CapabilityResponses 仅在显式生图意图时被要求"这一
|
||||
// 映射;后续若扩展该 capability 的用途,本测试失败提示同步收窄跳门条件。
|
||||
func TestOpenAIResponsesRequiredCapabilityPinsImageIntentMapping(t *testing.T) {
|
||||
func TestOpenAIResponsesRequiredCapabilityForRequest(t *testing.T) {
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityResponses, openAIResponsesRequiredCapability(true, service.PlatformOpenAI))
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityChatCompletions, openAIResponsesRequiredCapability(false, service.PlatformOpenAI))
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityChatCompletions, openAIResponsesRequiredCapability(true, service.PlatformGrok))
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityResponses, openAIResponsesRequiredCapabilityForRequest(false, true, service.PlatformOpenAI))
|
||||
require.Equal(t, service.OpenAIEndpointCapabilityChatCompletions, openAIResponsesRequiredCapabilityForRequest(false, true, service.PlatformGrok))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type grokStandaloneSearchRequest struct {
|
||||
Query string `json:"query"`
|
||||
Input string `json:"input"`
|
||||
MaxResults *int `json:"max_results"`
|
||||
AllowedXHandles []string `json:"allowed_x_handles"`
|
||||
ExcludedXHandles []string `json:"excluded_x_handles"`
|
||||
FromDate string `json:"from_date"`
|
||||
ToDate string `json:"to_date"`
|
||||
EnableImageUnderstanding *bool `json:"enable_image_understanding"`
|
||||
EnableVideoUnderstanding *bool `json:"enable_video_understanding"`
|
||||
}
|
||||
|
||||
// XSearch marks the standalone endpoint so WebSearch can use native x_search
|
||||
// while retaining its dedicated per-call billing contract.
|
||||
func (h *GatewayHandler) XSearch(c *gin.Context) {
|
||||
c.Set("grok_x_search_endpoint", true)
|
||||
h.WebSearch(c)
|
||||
}
|
||||
|
||||
func resolveGrokStandaloneSearchModel() string {
|
||||
return xai.ResolveDefaultTextModel(xai.RuntimeModelMappingOptions().DefaultText)
|
||||
}
|
||||
|
||||
func buildGrokXSearchResponsesBody(req grokStandaloneSearchRequest, model string) ([]byte, error) {
|
||||
input := strings.TrimSpace(req.Query)
|
||||
if input == "" {
|
||||
input = strings.TrimSpace(req.Input)
|
||||
}
|
||||
tool := map[string]any{"type": "x_search"}
|
||||
if len(req.AllowedXHandles) > 0 {
|
||||
tool["allowed_x_handles"] = req.AllowedXHandles
|
||||
}
|
||||
if len(req.ExcludedXHandles) > 0 {
|
||||
tool["excluded_x_handles"] = req.ExcludedXHandles
|
||||
}
|
||||
if strings.TrimSpace(req.FromDate) != "" {
|
||||
tool["from_date"] = strings.TrimSpace(req.FromDate)
|
||||
}
|
||||
if strings.TrimSpace(req.ToDate) != "" {
|
||||
tool["to_date"] = strings.TrimSpace(req.ToDate)
|
||||
}
|
||||
if req.EnableImageUnderstanding != nil {
|
||||
tool["enable_image_understanding"] = *req.EnableImageUnderstanding
|
||||
}
|
||||
if req.EnableVideoUnderstanding != nil {
|
||||
tool["enable_video_understanding"] = *req.EnableVideoUnderstanding
|
||||
}
|
||||
maxResults := 0
|
||||
if req.MaxResults != nil {
|
||||
maxResults = *req.MaxResults
|
||||
}
|
||||
return json.Marshal(map[string]any{
|
||||
"model": xai.ResolveDefaultTextModel(model),
|
||||
"input": buildGrokXSearchPrompt(input, maxResults),
|
||||
"tools": []map[string]any{tool},
|
||||
"tool_choice": "required",
|
||||
"include": []string{"x_search_call.action.sources"},
|
||||
"store": false,
|
||||
"stream": false,
|
||||
})
|
||||
}
|
||||
|
||||
func buildGrokXSearchPrompt(query string, maxResults int) string {
|
||||
return fmt.Sprintf(`Search X for the user query below. Return ONLY valid JSON with this exact shape: {"results":[{"url":"https://...","title":"post or page title","snippet":"concise factual summary"}]}. Return at most %d unique results. Every URL must be an actual x_search source. Populate a non-empty title and snippet for every result. Do not wrap the JSON in markdown.
|
||||
|
||||
User query:
|
||||
%s`, normalizeGrokWebSearchMaxResults(maxResults), query)
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestBuildGrokXSearchResponsesBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
understandImages := true
|
||||
understandVideos := false
|
||||
body, err := buildGrokXSearchResponsesBody(grokStandaloneSearchRequest{
|
||||
Query: "latest posts from xAI",
|
||||
AllowedXHandles: []string{"xai"},
|
||||
ExcludedXHandles: []string{"spam"},
|
||||
FromDate: "2026-08-01",
|
||||
ToDate: "2026-08-10",
|
||||
EnableImageUnderstanding: &understandImages,
|
||||
EnableVideoUnderstanding: &understandVideos,
|
||||
}, xai.DefaultTextModel)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultTextModel, gjson.GetBytes(body, "model").String())
|
||||
require.Contains(t, gjson.GetBytes(body, "input").String(), "latest posts from xAI")
|
||||
require.Contains(t, gjson.GetBytes(body, "input").String(), "Return ONLY valid JSON")
|
||||
require.Equal(t, "x_search_call.action.sources", gjson.GetBytes(body, "include.0").String())
|
||||
require.Equal(t, "required", gjson.GetBytes(body, "tool_choice").String())
|
||||
require.Equal(t, "x_search", gjson.GetBytes(body, "tools.0.type").String())
|
||||
require.Equal(t, "xai", gjson.GetBytes(body, "tools.0.allowed_x_handles.0").String())
|
||||
require.Equal(t, "spam", gjson.GetBytes(body, "tools.0.excluded_x_handles.0").String())
|
||||
require.Equal(t, "2026-08-01", gjson.GetBytes(body, "tools.0.from_date").String())
|
||||
require.Equal(t, "2026-08-10", gjson.GetBytes(body, "tools.0.to_date").String())
|
||||
require.True(t, gjson.GetBytes(body, "tools.0.enable_image_understanding").Bool())
|
||||
require.False(t, gjson.GetBytes(body, "tools.0.enable_video_understanding").Bool())
|
||||
require.False(t, gjson.GetBytes(body, "store").Bool())
|
||||
require.False(t, gjson.GetBytes(body, "stream").Bool())
|
||||
}
|
||||
|
||||
func TestBuildGrokXSearchResponsesBodyAcceptsInputAlias(t *testing.T) {
|
||||
t.Parallel()
|
||||
body, err := buildGrokXSearchResponsesBody(grokStandaloneSearchRequest{Input: "latest posts from xAI"}, xai.DefaultTextModel)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, gjson.GetBytes(body, "input").String(), "latest posts from xAI")
|
||||
}
|
||||
|
||||
func TestResolveGrokStandaloneSearchModelUsesRuntimeDefault(t *testing.T) {
|
||||
original := xai.RuntimeModelMappingOptions()
|
||||
t.Cleanup(func() { xai.SetRuntimeModelMappingOptions(original) })
|
||||
xai.SetRuntimeModelMappingOptions(xai.ModelMappingOptions{DefaultText: "grok-4.6"})
|
||||
|
||||
model := resolveGrokStandaloneSearchModel()
|
||||
body, err := buildGrokXSearchResponsesBody(grokStandaloneSearchRequest{Query: "latest posts from xAI"}, model)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "grok-4.6", model)
|
||||
require.Equal(t, model, gjson.GetBytes(body, "model").String())
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
@@ -12,6 +13,15 @@ import (
|
||||
)
|
||||
|
||||
const securityAuditCompletedContextKey = "sub2api.security_audit.completed"
|
||||
const securityAuditWSTurnContextKey = "sub2api.security_audit.ws_turn"
|
||||
const securityAuditWSDedupeContextKey = "sub2api.security_audit.ws_dedupe"
|
||||
|
||||
type securityAuditWSDedupeEntry struct {
|
||||
stage string
|
||||
turn int
|
||||
bodyHash [sha256.Size]byte
|
||||
decision securityaudit.Decision
|
||||
}
|
||||
|
||||
// cachesSecurityAuditCompletion reports whether a successful audit may be
|
||||
// reused for the rest of the gin request. WebSocket turns share one Context
|
||||
@@ -25,6 +35,15 @@ func cachesSecurityAuditCompletion(stage string) bool {
|
||||
}
|
||||
}
|
||||
|
||||
func isSecurityAuditWebSocketStage(stage string) bool {
|
||||
switch strings.TrimSpace(stage) {
|
||||
case "first_turn", "subsequent_turn":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) checkSecurityAudit(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte) *securityaudit.Decision {
|
||||
if h == nil {
|
||||
return nil
|
||||
@@ -76,27 +95,68 @@ func runSecurityAudit(c *gin.Context, reqLog *zap.Logger, coordinator *securitya
|
||||
return &decision
|
||||
}
|
||||
request := buildSecurityAuditRequest(c, apiKey, subject, protocol, model, body, stage)
|
||||
if reqLog != nil {
|
||||
reqLog.Info("security_audit.gateway_check_start",
|
||||
zap.String("request_id", request.RequestID), zap.Int64("user_id", request.UserID),
|
||||
zap.Int64("api_key_id", request.APIKeyID), zap.Int64p("group_id", request.GroupID),
|
||||
zap.String("endpoint", request.Endpoint), zap.String("provider", request.Provider),
|
||||
zap.String("protocol", request.Protocol), zap.String("model", request.Model), zap.String("stage", request.Stage),
|
||||
zap.Int("body_bytes", len(body)))
|
||||
if isSecurityAuditWebSocketStage(request.Stage) {
|
||||
if turnNo, ok := securityAuditWSTurn(c); ok {
|
||||
bodyHash := sha256.Sum256(body)
|
||||
if cached, exists := c.Get(securityAuditWSDedupeContextKey); exists {
|
||||
if entry, ok := cached.(securityAuditWSDedupeEntry); ok &&
|
||||
entry.stage == request.Stage && entry.turn == turnNo && entry.bodyHash == bodyHash {
|
||||
decision := entry.decision
|
||||
logSecurityAuditDone(reqLog, request, decision, true)
|
||||
return &decision
|
||||
}
|
||||
}
|
||||
logSecurityAuditStart(reqLog, request, len(body), false)
|
||||
decision := coordinator.Check(c.Request.Context(), request)
|
||||
if decision.Kind == securityaudit.DecisionAllow {
|
||||
c.Set(securityAuditWSDedupeContextKey, securityAuditWSDedupeEntry{
|
||||
stage: request.Stage, turn: turnNo, bodyHash: bodyHash, decision: decision,
|
||||
})
|
||||
}
|
||||
logSecurityAuditDone(reqLog, request, decision, false)
|
||||
return &decision
|
||||
}
|
||||
}
|
||||
logSecurityAuditStart(reqLog, request, len(body), false)
|
||||
decision := coordinator.Check(c.Request.Context(), request)
|
||||
if decision.AllowNextStage && cacheCompletion {
|
||||
c.Set(securityAuditCompletedContextKey, true)
|
||||
}
|
||||
if reqLog != nil {
|
||||
reqLog.Info("security_audit.gateway_check_done",
|
||||
zap.String("request_id", request.RequestID), zap.String("decision", string(decision.Kind)),
|
||||
zap.String("error_code", decision.ErrorCode), zap.Bool("allow_next_stage", decision.AllowNextStage),
|
||||
zap.String("stage", request.Stage))
|
||||
}
|
||||
logSecurityAuditDone(reqLog, request, decision, false)
|
||||
return &decision
|
||||
}
|
||||
|
||||
func logSecurityAuditStart(reqLog *zap.Logger, request securityaudit.Request, bodyBytes int, cached bool) {
|
||||
if reqLog == nil {
|
||||
return
|
||||
}
|
||||
reqLog.Info("security_audit.gateway_check_start",
|
||||
zap.String("request_id", request.RequestID), zap.Int64("user_id", request.UserID),
|
||||
zap.Int64("api_key_id", request.APIKeyID), zap.Int64p("group_id", request.GroupID),
|
||||
zap.String("endpoint", request.Endpoint), zap.String("provider", request.Provider),
|
||||
zap.String("protocol", request.Protocol), zap.String("model", request.Model), zap.String("stage", request.Stage),
|
||||
zap.Int("body_bytes", bodyBytes), zap.Bool("cached", cached))
|
||||
}
|
||||
|
||||
func logSecurityAuditDone(reqLog *zap.Logger, request securityaudit.Request, decision securityaudit.Decision, cached bool) {
|
||||
if reqLog == nil {
|
||||
return
|
||||
}
|
||||
reqLog.Info("security_audit.gateway_check_done",
|
||||
zap.String("request_id", request.RequestID), zap.String("decision", string(decision.Kind)),
|
||||
zap.String("error_code", decision.ErrorCode), zap.Bool("allow_next_stage", decision.AllowNextStage),
|
||||
zap.String("stage", request.Stage), zap.Bool("cached", cached))
|
||||
}
|
||||
|
||||
func securityAuditWSTurn(c *gin.Context) (int, bool) {
|
||||
turn, exists := c.Get(securityAuditWSTurnContextKey)
|
||||
if !exists {
|
||||
return 0, false
|
||||
}
|
||||
turnNo, ok := turn.(int)
|
||||
return turnNo, ok
|
||||
}
|
||||
|
||||
func buildSecurityAuditRequest(c *gin.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte, stage string) securityaudit.Request {
|
||||
legacy := buildContentModerationInput(c, apiKey, subject, protocol, model, body)
|
||||
request := securityaudit.Request{
|
||||
|
||||
@@ -11,6 +11,8 @@ import (
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/zap"
|
||||
"go.uber.org/zap/zaptest/observer"
|
||||
)
|
||||
|
||||
func TestCachesSecurityAuditCompletionSkipsWebSocketStages(t *testing.T) {
|
||||
@@ -18,6 +20,9 @@ func TestCachesSecurityAuditCompletionSkipsWebSocketStages(t *testing.T) {
|
||||
require.True(t, cachesSecurityAuditCompletion(""))
|
||||
require.False(t, cachesSecurityAuditCompletion("first_turn"))
|
||||
require.False(t, cachesSecurityAuditCompletion("subsequent_turn"))
|
||||
require.True(t, isSecurityAuditWebSocketStage("first_turn"))
|
||||
require.True(t, isSecurityAuditWebSocketStage("subsequent_turn"))
|
||||
require.False(t, isSecurityAuditWebSocketStage("http"))
|
||||
}
|
||||
|
||||
func TestRunSecurityAuditDoesNotSkipSubsequentWebSocketTurns(t *testing.T) {
|
||||
@@ -48,9 +53,119 @@ func TestRunSecurityAuditDoesNotSkipSubsequentWebSocketTurns(t *testing.T) {
|
||||
require.Equal(t, int64(2), engine.enqueues.Load(), "subsequent WebSocket turns must be audited again")
|
||||
}
|
||||
|
||||
func TestRunSecurityAuditDeduplicatesRepeatedPayloadWithinWebSocketTurn(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := &turnCountingEngine{mode: securityaudit.ModeBlocking}
|
||||
coordinator := securityaudit.NewCoordinator(nil, engine)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","response":{"input":"same turn"}}`)
|
||||
c.Set(securityAuditWSTurnContextKey, 2)
|
||||
first := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
second := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
require.NotNil(t, first)
|
||||
require.NotNil(t, second)
|
||||
require.True(t, first.AllowNextStage)
|
||||
require.True(t, second.AllowNextStage)
|
||||
require.Equal(t, int64(1), engine.evaluates.Load())
|
||||
|
||||
// The cache holds only one successful same-turn result.
|
||||
entry, exists := c.Get(securityAuditWSDedupeContextKey)
|
||||
require.True(t, exists)
|
||||
require.IsType(t, securityAuditWSDedupeEntry{}, entry)
|
||||
|
||||
c.Set(securityAuditWSTurnContextKey, 3)
|
||||
runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
require.Equal(t, int64(2), engine.evaluates.Load())
|
||||
}
|
||||
|
||||
func TestRunSecurityAuditDoesNotCacheFailedWebSocketDecision(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := &turnCountingEngine{
|
||||
mode: securityaudit.ModeBlocking,
|
||||
decisions: []*securityaudit.PromptDecision{
|
||||
{Kind: securityaudit.DecisionUnavailable, AllowNextStage: false},
|
||||
{Kind: securityaudit.DecisionAllow, AllowNextStage: true},
|
||||
},
|
||||
}
|
||||
coordinator := securityaudit.NewCoordinator(nil, engine)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(securityAuditWSTurnContextKey, 2)
|
||||
payload := []byte(`{"type":"response.create","response":{"input":"retry me"}}`)
|
||||
|
||||
first := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
_, cachedAfterFailure := c.Get(securityAuditWSDedupeContextKey)
|
||||
second := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
|
||||
require.False(t, first.AllowNextStage)
|
||||
require.False(t, cachedAfterFailure)
|
||||
require.True(t, second.AllowNextStage)
|
||||
require.Equal(t, int64(2), engine.evaluates.Load())
|
||||
}
|
||||
|
||||
func TestRunSecurityAuditDoesNotCacheFlaggedWebSocketDecision(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := &turnCountingEngine{
|
||||
mode: securityaudit.ModeBlocking,
|
||||
decisions: []*securityaudit.PromptDecision{
|
||||
{Kind: securityaudit.DecisionFlag, AllowNextStage: true},
|
||||
{Kind: securityaudit.DecisionAllow, AllowNextStage: true},
|
||||
},
|
||||
}
|
||||
coordinator := securityaudit.NewCoordinator(nil, engine)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(securityAuditWSTurnContextKey, 2)
|
||||
payload := []byte(`{"type":"response.create","response":{"input":"retry flagged"}}`)
|
||||
|
||||
first := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
_, cachedAfterFlag := c.Get(securityAuditWSDedupeContextKey)
|
||||
second := runSecurityAudit(c, nil, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
|
||||
require.Equal(t, securityaudit.DecisionFlag, first.Kind)
|
||||
require.True(t, first.AllowNextStage)
|
||||
require.False(t, cachedAfterFlag)
|
||||
require.Equal(t, securityaudit.DecisionAllow, second.Kind)
|
||||
require.Equal(t, int64(2), engine.evaluates.Load())
|
||||
}
|
||||
|
||||
func TestRunSecurityAuditLogsWebSocketChecksAndCacheHits(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := &turnCountingEngine{mode: securityaudit.ModeBlocking}
|
||||
coordinator := securityaudit.NewCoordinator(nil, engine)
|
||||
core, logs := observer.New(zap.InfoLevel)
|
||||
reqLog := zap.New(core)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(securityAuditWSTurnContextKey, 2)
|
||||
payload := []byte(`{"type":"response.create","response":{"input":"same turn"}}`)
|
||||
|
||||
runSecurityAudit(c, reqLog, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
runSecurityAudit(c, reqLog, coordinator, nil, nil, middleware2.AuthSubject{UserID: 7}, "openai_responses", "gpt-test", payload, "subsequent_turn")
|
||||
|
||||
startLogs := logs.FilterMessage("security_audit.gateway_check_start").All()
|
||||
require.Len(t, startLogs, 1)
|
||||
require.Equal(t, false, startLogs[0].ContextMap()["cached"])
|
||||
|
||||
doneLogs := logs.FilterMessage("security_audit.gateway_check_done").All()
|
||||
require.Len(t, doneLogs, 2)
|
||||
require.Equal(t, false, doneLogs[0].ContextMap()["cached"])
|
||||
require.Equal(t, true, doneLogs[1].ContextMap()["cached"])
|
||||
require.Equal(t, "allow", doneLogs[1].ContextMap()["decision"])
|
||||
require.Equal(t, "subsequent_turn", doneLogs[1].ContextMap()["stage"])
|
||||
require.Equal(t, int64(1), engine.evaluates.Load())
|
||||
}
|
||||
|
||||
type turnCountingEngine struct {
|
||||
mode securityaudit.Mode
|
||||
enqueues atomic.Int64
|
||||
mode securityaudit.Mode
|
||||
enqueues atomic.Int64
|
||||
evaluates atomic.Int64
|
||||
decisions []*securityaudit.PromptDecision
|
||||
}
|
||||
|
||||
func (e *turnCountingEngine) EffectiveMode() securityaudit.Mode { return e.mode }
|
||||
@@ -59,5 +174,9 @@ func (e *turnCountingEngine) Enqueue(context.Context, securityaudit.Request) err
|
||||
return nil
|
||||
}
|
||||
func (e *turnCountingEngine) Evaluate(context.Context, securityaudit.Request) (*securityaudit.PromptDecision, error) {
|
||||
call := e.evaluates.Add(1)
|
||||
if int(call) <= len(e.decisions) {
|
||||
return e.decisions[call-1], nil
|
||||
}
|
||||
return &securityaudit.PromptDecision{Kind: securityaudit.DecisionAllow, AllowNextStage: true}, nil
|
||||
}
|
||||
|
||||
@@ -423,19 +423,20 @@ func ChatCompletionsResponseToAnthropic(resp *ChatCompletionsResponse, model str
|
||||
// + the reasoning→thinking mapping in ResponsesToAnthropic.
|
||||
func chatMessageToAnthropicBlocks(message ChatMessage) []AnthropicContentBlock {
|
||||
var blocks []AnthropicContentBlock
|
||||
reasoning := message.reasoningText()
|
||||
|
||||
if message.ReasoningContent != "" {
|
||||
if reasoning != "" {
|
||||
blocks = append(blocks, AnthropicContentBlock{
|
||||
Type: "thinking",
|
||||
Thinking: message.ReasoningContent,
|
||||
Thinking: reasoning,
|
||||
})
|
||||
}
|
||||
|
||||
text := chatMessageContentText(message.Content)
|
||||
// DeepSeek reasoning-only fallback: when there is no text and no tool calls,
|
||||
// surface the reasoning content as visible text so the turn isn't empty.
|
||||
if text == "" && strings.TrimSpace(message.ReasoningContent) != "" && len(message.ToolCalls) == 0 {
|
||||
text = message.ReasoningContent
|
||||
if text == "" && strings.TrimSpace(reasoning) != "" && len(message.ToolCalls) == 0 {
|
||||
text = reasoning
|
||||
}
|
||||
if text != "" || len(message.ToolCalls) == 0 {
|
||||
blocks = append(blocks, AnthropicContentBlock{Type: "text", Text: text})
|
||||
@@ -608,11 +609,12 @@ func ChatCompletionsChunkToAnthropicEvents(
|
||||
|
||||
for _, choice := range chunk.Choices {
|
||||
// Reasoning content → thinking block.
|
||||
if choice.Delta.ReasoningContent != nil && *choice.Delta.ReasoningContent != "" {
|
||||
reasoning := choice.Delta.reasoningText()
|
||||
if reasoning != nil && *reasoning != "" {
|
||||
events = append(events, ensureCCAnthropicThinkingBlock(state)...)
|
||||
events = append(events, ccAnthropicDelta(state, &AnthropicDelta{
|
||||
Type: "thinking_delta",
|
||||
Thinking: *choice.Delta.ReasoningContent,
|
||||
Thinking: *reasoning,
|
||||
})...)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func readIssue5302Fixture(t *testing.T, name string, dst any) {
|
||||
t.Helper()
|
||||
payload, err := os.ReadFile(filepath.Join("testdata", "issue5302", name))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, json.Unmarshal(payload, dst))
|
||||
}
|
||||
|
||||
func TestChatReasoningAlias_AnthropicNonStreaming(t *testing.T) {
|
||||
var response ChatCompletionsResponse
|
||||
readIssue5302Fixture(t, "nonstream_reasoning.json", &response)
|
||||
|
||||
out := ChatCompletionsResponseToAnthropic(&response, "claude-sonnet-4-20250514")
|
||||
require.Len(t, out.Content, 2)
|
||||
require.Equal(t, "thinking", out.Content[0].Type)
|
||||
require.Equal(t, "fallback reasoning", out.Content[0].Thinking)
|
||||
require.Equal(t, "final answer", out.Content[1].Text)
|
||||
}
|
||||
|
||||
func TestChatReasoningAlias_AnthropicStreaming(t *testing.T) {
|
||||
var chunk ChatCompletionsChunk
|
||||
readIssue5302Fixture(t, "stream_reasoning.json", &chunk)
|
||||
|
||||
events := ChatCompletionsChunkToAnthropicEvents(&chunk, NewChatCompletionsToAnthropicStreamState("reasoning-model"))
|
||||
var thinking string
|
||||
for _, event := range events {
|
||||
if event.Delta != nil && event.Delta.Type == "thinking_delta" {
|
||||
thinking += event.Delta.Thinking
|
||||
}
|
||||
}
|
||||
require.Equal(t, "streamed fallback", thinking)
|
||||
}
|
||||
|
||||
func TestChatReasoningAlias_ResponsesSharedPaths(t *testing.T) {
|
||||
var response ChatCompletionsResponse
|
||||
readIssue5302Fixture(t, "nonstream_reasoning.json", &response)
|
||||
nonStream := ChatCompletionsResponseToResponses(&response, "reasoning-model", nil, false, nil)
|
||||
require.Len(t, nonStream.Output, 2)
|
||||
require.Equal(t, "reasoning", nonStream.Output[0].Type)
|
||||
require.Equal(t, "fallback reasoning", nonStream.Output[0].Summary[0].Text)
|
||||
|
||||
var chunk ChatCompletionsChunk
|
||||
readIssue5302Fixture(t, "stream_reasoning.json", &chunk)
|
||||
events := ChatCompletionsChunkToResponsesEvents(&chunk, NewChatCompletionsToResponsesStreamState("reasoning-model"))
|
||||
var deltas []string
|
||||
for _, event := range events {
|
||||
if event.Type == "response.reasoning_summary_text.delta" {
|
||||
deltas = append(deltas, event.Delta)
|
||||
}
|
||||
}
|
||||
require.Equal(t, []string{"streamed fallback"}, deltas)
|
||||
}
|
||||
|
||||
func TestChatReasoningAlias_ReasoningContentTakesPrecedence(t *testing.T) {
|
||||
var chunk ChatCompletionsChunk
|
||||
readIssue5302Fixture(t, "reasoning_content_precedence.json", &chunk)
|
||||
|
||||
events := ChatCompletionsChunkToResponsesEvents(&chunk, NewChatCompletionsToResponsesStreamState("reasoning-model"))
|
||||
var deltas []string
|
||||
for _, event := range events {
|
||||
if event.Type == "response.reasoning_summary_text.delta" {
|
||||
deltas = append(deltas, event.Delta)
|
||||
}
|
||||
}
|
||||
require.Equal(t, []string{"preferred reasoning"}, deltas)
|
||||
}
|
||||
@@ -63,6 +63,9 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
|
||||
if tool.Function != nil {
|
||||
declared[tool.Function.Name] = true
|
||||
}
|
||||
if strings.EqualFold(strings.TrimSpace(tool.Type), "x_search") {
|
||||
declared["x_search"] = true
|
||||
}
|
||||
}
|
||||
if tc := responsesToolChoiceToChatToolChoice(req.ToolChoice, declared); len(tc) > 0 {
|
||||
out.ToolChoice = tc
|
||||
@@ -847,6 +850,16 @@ func responsesToolsToChatTools(tools []ResponsesTool) ([]ChatTool, error) {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, flattened...)
|
||||
case "x_search":
|
||||
out = append(out, ChatTool{
|
||||
Type: "x_search",
|
||||
AllowedXHandles: tool.AllowedXHandles,
|
||||
ExcludedXHandles: tool.ExcludedXHandles,
|
||||
FromDate: tool.FromDate,
|
||||
ToDate: tool.ToDate,
|
||||
EnableImageUnderstanding: tool.EnableImageUnderstanding,
|
||||
EnableVideoUnderstanding: tool.EnableVideoUnderstanding,
|
||||
})
|
||||
}
|
||||
// 其余类型(web_search、image_generation 等服务端工具)在 chat 上游没有
|
||||
// 对应能力,维持丢弃。
|
||||
@@ -948,6 +961,15 @@ func responsesToolChoiceToChatToolChoice(raw json.RawMessage, declared map[strin
|
||||
}
|
||||
var name string
|
||||
switch rawString(choice["type"]) {
|
||||
case "x_search":
|
||||
if !declared["x_search"] {
|
||||
return nil
|
||||
}
|
||||
out, err := json.Marshal(map[string]any{"type": "x_search"})
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
return out
|
||||
case "tool_search":
|
||||
// tool_search 未被丢弃而是降级为同名 function 代理(见
|
||||
// responsesToolsToChatTools),强制选择它同样降级为 function 选择,
|
||||
@@ -1053,20 +1075,21 @@ func ChatCompletionsResponseToResponses(resp *ChatCompletionsResponse, model str
|
||||
|
||||
func chatMessageToResponsesOutput(message ChatMessage, customTools map[string]bool, toolSearch bool, namespaceTools map[string]NamespacedToolName) []ResponsesOutput {
|
||||
var outputs []ResponsesOutput
|
||||
if message.ReasoningContent != "" {
|
||||
reasoning := message.reasoningText()
|
||||
if reasoning != "" {
|
||||
outputs = append(outputs, ResponsesOutput{
|
||||
Type: "reasoning",
|
||||
ID: generateItemID(),
|
||||
Summary: []ResponsesSummary{{
|
||||
Type: "summary_text",
|
||||
Text: message.ReasoningContent,
|
||||
Text: reasoning,
|
||||
}},
|
||||
})
|
||||
}
|
||||
|
||||
text := chatMessageContentText(message.Content)
|
||||
if text == "" && strings.TrimSpace(message.ReasoningContent) != "" && len(message.ToolCalls) == 0 {
|
||||
text = message.ReasoningContent
|
||||
if text == "" && strings.TrimSpace(reasoning) != "" && len(message.ToolCalls) == 0 {
|
||||
text = reasoning
|
||||
}
|
||||
if text != "" || len(message.ToolCalls) == 0 {
|
||||
outputs = append(outputs, ResponsesOutput{
|
||||
@@ -1328,13 +1351,14 @@ func ChatCompletionsChunkToResponsesEvents(
|
||||
// (output_item.added + reasoning_summary_part.added) before the first
|
||||
// delta, otherwise a strict client discards the delta. The leading
|
||||
// empty-string reasoning delta upstreams send is filtered out.
|
||||
if choice.Delta.ReasoningContent != nil && *choice.Delta.ReasoningContent != "" {
|
||||
reasoning := choice.Delta.reasoningText()
|
||||
if reasoning != nil && *reasoning != "" {
|
||||
events = append(events, ensureChatReasoningItem(state)...)
|
||||
_, _ = state.Reasoning.WriteString(*choice.Delta.ReasoningContent)
|
||||
_, _ = state.Reasoning.WriteString(*reasoning)
|
||||
events = append(events, chatToResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.ReasoningIndex,
|
||||
SummaryIndex: 0,
|
||||
Delta: *choice.Delta.ReasoningContent,
|
||||
Delta: *reasoning,
|
||||
ItemID: state.ReasoningItemID,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -661,6 +661,20 @@ func TestResponsesToChatCompletionsRequest_DropsToolChoiceForDroppedTool(t *test
|
||||
require.Len(t, out.Tools, 1)
|
||||
assert.Empty(t, out.ToolChoice, "指向被丢弃服务端工具的 tool_choice 必须丢弃")
|
||||
|
||||
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
|
||||
Model: "glm-5.2",
|
||||
Input: json.RawMessage(`"hi"`),
|
||||
Tools: []ResponsesTool{
|
||||
{Type: "function", Name: "wait", Parameters: json.RawMessage(`{"type":"object","properties":{}}`)},
|
||||
{Type: "web_search"},
|
||||
{Type: "x_search"},
|
||||
},
|
||||
ToolChoice: json.RawMessage(`{"type":"function","name":"web_search"}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, out.Tools, 2)
|
||||
assert.Empty(t, out.ToolChoice, "surviving x_search must not keep a function tool_choice named web_search")
|
||||
|
||||
// 具名选择指向不存在的工具名。
|
||||
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
|
||||
Model: "glm-5.2",
|
||||
|
||||
@@ -419,6 +419,18 @@ func convertChatToolsToResponses(tools []ChatTool, functions []ChatFunction) []R
|
||||
var out []ResponsesTool
|
||||
|
||||
for _, t := range tools {
|
||||
if strings.EqualFold(strings.TrimSpace(t.Type), "x_search") {
|
||||
out = append(out, ResponsesTool{
|
||||
Type: "x_search",
|
||||
AllowedXHandles: t.AllowedXHandles,
|
||||
ExcludedXHandles: t.ExcludedXHandles,
|
||||
FromDate: t.FromDate,
|
||||
ToDate: t.ToDate,
|
||||
EnableImageUnderstanding: t.EnableImageUnderstanding,
|
||||
EnableVideoUnderstanding: t.EnableVideoUnderstanding,
|
||||
})
|
||||
continue
|
||||
}
|
||||
if t.Type != "function" || t.Function == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestChatCompletionsToResponsesPreservesXSearchTool(t *testing.T) {
|
||||
enabled := true
|
||||
req := &ChatCompletionsRequest{
|
||||
Model: "grok-4.5",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"latest xAI post"`)},
|
||||
},
|
||||
Tools: []ChatTool{{
|
||||
Type: "x_search",
|
||||
AllowedXHandles: []string{"xai"},
|
||||
ExcludedXHandles: []string{"spam"},
|
||||
FromDate: "2026-08-01",
|
||||
ToDate: "2026-08-10",
|
||||
EnableImageUnderstanding: &enabled,
|
||||
EnableVideoUnderstanding: &enabled,
|
||||
}},
|
||||
ToolChoice: json.RawMessage(`{"type":"x_search"}`),
|
||||
}
|
||||
|
||||
resp, err := ChatCompletionsToResponses(req)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, resp.Tools, 1)
|
||||
require.Equal(t, "x_search", resp.Tools[0].Type)
|
||||
require.Equal(t, []string{"xai"}, resp.Tools[0].AllowedXHandles)
|
||||
require.Equal(t, []string{"spam"}, resp.Tools[0].ExcludedXHandles)
|
||||
require.Equal(t, "2026-08-01", resp.Tools[0].FromDate)
|
||||
require.Equal(t, "2026-08-10", resp.Tools[0].ToDate)
|
||||
require.NotNil(t, resp.Tools[0].EnableImageUnderstanding)
|
||||
require.True(t, *resp.Tools[0].EnableImageUnderstanding)
|
||||
require.NotNil(t, resp.Tools[0].EnableVideoUnderstanding)
|
||||
require.True(t, *resp.Tools[0].EnableVideoUnderstanding)
|
||||
require.JSONEq(t, `{"type":"x_search"}`, string(resp.ToolChoice))
|
||||
}
|
||||
|
||||
func TestResponsesToChatCompletionsPreservesXSearchTool(t *testing.T) {
|
||||
enabled := true
|
||||
req := &ResponsesRequest{
|
||||
Model: "grok-4.5",
|
||||
Input: json.RawMessage(`"latest xAI post"`),
|
||||
Tools: []ResponsesTool{{
|
||||
Type: "x_search",
|
||||
AllowedXHandles: []string{"xai"},
|
||||
ExcludedXHandles: []string{"spam"},
|
||||
FromDate: "2026-08-01",
|
||||
ToDate: "2026-08-10",
|
||||
EnableImageUnderstanding: &enabled,
|
||||
EnableVideoUnderstanding: &enabled,
|
||||
}},
|
||||
ToolChoice: json.RawMessage(`{"type":"x_search"}`),
|
||||
}
|
||||
|
||||
chat, err := ResponsesToChatCompletionsRequest(req)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, chat.Tools, 1)
|
||||
require.Equal(t, "x_search", chat.Tools[0].Type)
|
||||
require.Equal(t, []string{"xai"}, chat.Tools[0].AllowedXHandles)
|
||||
require.Equal(t, []string{"spam"}, chat.Tools[0].ExcludedXHandles)
|
||||
require.JSONEq(t, `{"type":"x_search"}`, string(chat.ToolChoice))
|
||||
}
|
||||
|
||||
func TestResponsesToChatCompletionsXSearchToolChoiceString(t *testing.T) {
|
||||
chat, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
|
||||
Model: "grok-4.5",
|
||||
Input: json.RawMessage(`"latest xAI post"`),
|
||||
Tools: []ResponsesTool{{Type: "x_search"}},
|
||||
ToolChoice: json.RawMessage(`"x_search"`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `"x_search"`, string(chat.ToolChoice))
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"id": "chatcmpl-sanitized",
|
||||
"object": "chat.completion",
|
||||
"model": "reasoning-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "final answer",
|
||||
"reasoning": "fallback reasoning",
|
||||
"reasoning_details": [
|
||||
{"type": "reasoning.text", "text": "fallback reasoning"}
|
||||
]
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
]
|
||||
}
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
{
|
||||
"id": "chatcmpl-sanitized",
|
||||
"object": "chat.completion.chunk",
|
||||
"model": "reasoning-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"reasoning_content": "preferred reasoning",
|
||||
"reasoning": "fallback reasoning"
|
||||
},
|
||||
"finish_reason": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
{
|
||||
"id": "chatcmpl-sanitized",
|
||||
"object": "chat.completion.chunk",
|
||||
"model": "reasoning-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"reasoning": "streamed fallback",
|
||||
"reasoning_details": [
|
||||
{"type": "reasoning.text", "text": "streamed fallback"}
|
||||
]
|
||||
},
|
||||
"finish_reason": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -302,7 +302,7 @@ type ResponsesContentPart struct {
|
||||
|
||||
// ResponsesTool describes a tool in the Responses API.
|
||||
type ResponsesTool struct {
|
||||
Type string `json:"type"` // "function" | "custom" | "web_search" | "local_shell" etc.
|
||||
Type string `json:"type"` // "function" | "custom" | "web_search" | "x_search" | "local_shell" etc.
|
||||
Name string `json:"name,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters json.RawMessage `json:"parameters,omitempty"`
|
||||
@@ -311,6 +311,14 @@ type ResponsesTool struct {
|
||||
// type=namespace 的子工具列表(tools 与 children 二选一,语义相同)。
|
||||
Tools []ResponsesTool `json:"tools,omitempty"`
|
||||
Children []ResponsesTool `json:"children,omitempty"`
|
||||
|
||||
// type=x_search
|
||||
AllowedXHandles []string `json:"allowed_x_handles,omitempty"`
|
||||
ExcludedXHandles []string `json:"excluded_x_handles,omitempty"`
|
||||
FromDate string `json:"from_date,omitempty"`
|
||||
ToDate string `json:"to_date,omitempty"`
|
||||
EnableImageUnderstanding *bool `json:"enable_image_understanding,omitempty"`
|
||||
EnableVideoUnderstanding *bool `json:"enable_video_understanding,omitempty"`
|
||||
}
|
||||
|
||||
// UnmarshalJSON 容忍字符串形式的工具声明:codex 会以 "name" 简写声明 custom 工具,
|
||||
@@ -651,6 +659,7 @@ type ChatMessage struct {
|
||||
Role string `json:"role"` // "system" | "user" | "assistant" | "tool" | "function"
|
||||
Content json.RawMessage `json:"content,omitempty"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
Reasoning string `json:"reasoning,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
ToolCalls []ChatToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
@@ -674,8 +683,16 @@ type ChatImageURL struct {
|
||||
|
||||
// ChatTool describes a tool available to the model.
|
||||
type ChatTool struct {
|
||||
Type string `json:"type"` // "function"
|
||||
Type string `json:"type"` // "function" | "x_search"
|
||||
Function *ChatFunction `json:"function,omitempty"`
|
||||
|
||||
// type=x_search
|
||||
AllowedXHandles []string `json:"allowed_x_handles,omitempty"`
|
||||
ExcludedXHandles []string `json:"excluded_x_handles,omitempty"`
|
||||
FromDate string `json:"from_date,omitempty"`
|
||||
ToDate string `json:"to_date,omitempty"`
|
||||
EnableImageUnderstanding *bool `json:"enable_image_understanding,omitempty"`
|
||||
EnableVideoUnderstanding *bool `json:"enable_video_understanding,omitempty"`
|
||||
}
|
||||
|
||||
// ChatFunction describes a function tool definition.
|
||||
@@ -771,9 +788,24 @@ type ChatDelta struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content *string `json:"content,omitempty"` // pointer: omit when not present, null vs "" matters
|
||||
ReasoningContent *string `json:"reasoning_content,omitempty"`
|
||||
Reasoning *string `json:"reasoning,omitempty"`
|
||||
ToolCalls []ChatToolCall `json:"tool_calls,omitempty"`
|
||||
}
|
||||
|
||||
func (m ChatMessage) reasoningText() string {
|
||||
if m.ReasoningContent != "" {
|
||||
return m.ReasoningContent
|
||||
}
|
||||
return m.Reasoning
|
||||
}
|
||||
|
||||
func (d ChatDelta) reasoningText() *string {
|
||||
if d.ReasoningContent != nil {
|
||||
return d.ReasoningContent
|
||||
}
|
||||
return d.Reasoning
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared constants
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -115,11 +115,12 @@ type EndpointStat struct {
|
||||
ActualCost float64 `json:"actual_cost"` // 实际扣除
|
||||
}
|
||||
|
||||
// GroupUsageSummary represents today's and cumulative cost for a single group.
|
||||
// GroupUsageSummary represents today's, yesterday's, and cumulative cost for a single group.
|
||||
type GroupUsageSummary struct {
|
||||
GroupID int64 `json:"group_id"`
|
||||
TodayCost float64 `json:"today_cost"`
|
||||
TotalCost float64 `json:"total_cost"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
TodayCost float64 `json:"today_cost"`
|
||||
YesterdayCost float64 `json:"yesterday_cost"`
|
||||
TotalCost float64 `json:"total_cost"`
|
||||
}
|
||||
|
||||
// GroupStat represents usage statistics for a single group
|
||||
|
||||
@@ -83,6 +83,7 @@ func (o ModelMappingOptions) defaultText() string {
|
||||
|
||||
var defaultModels = []Model{
|
||||
// Text
|
||||
{ID: "grok-4.6", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.6"},
|
||||
{ID: "grok-4.5", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.5"},
|
||||
{ID: "grok-4.3", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.3"},
|
||||
{ID: "grok-3-mini", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 3 Mini"},
|
||||
@@ -106,6 +107,8 @@ var defaultModels = []Model{
|
||||
var grokTextResponsesModelAliases = map[string]string{
|
||||
"grok": DefaultTextModel,
|
||||
"grok-latest": DefaultTextModel,
|
||||
"grok-4.6": "grok-4.6",
|
||||
"grok-4.6-latest": "grok-4.6",
|
||||
"grok-4.5": DefaultTextModel,
|
||||
"grok-4.5-latest": DefaultTextModel,
|
||||
"grok-4.3": "grok-4.3",
|
||||
|
||||
@@ -52,11 +52,20 @@ func TestCanonicalImagineVideoModel(t *testing.T) {
|
||||
func TestIsGrokModelID(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.True(t, IsGrokModelID("grok-4.5"))
|
||||
require.True(t, IsGrokModelID("grok-4.6"))
|
||||
require.True(t, IsGrokModelID("x-ai/grok-4.3"))
|
||||
require.False(t, IsGrokModelID("gpt-5"))
|
||||
require.False(t, IsGrokModelID("claude-sonnet-4"))
|
||||
}
|
||||
|
||||
func TestDefaultModelsIncludesGrok46(t *testing.T) {
|
||||
t.Parallel()
|
||||
ids := DefaultModelIDs()
|
||||
require.Contains(t, ids, "grok-4.6")
|
||||
require.Equal(t, "grok-4.6", ResolveGrokTextResponsesModelID("grok-4.6"))
|
||||
require.Equal(t, "grok-4.6", ResolveGrokTextResponsesModelID("grok-4.6-latest"))
|
||||
}
|
||||
|
||||
func TestResolveGrokTextResponsesModelID(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.Equal(t, "grok-4.5", ResolveGrokTextResponsesModelID(""))
|
||||
|
||||
@@ -352,6 +352,8 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) {
|
||||
mapping := DefaultModelMapping()
|
||||
require.Equal(t, "grok-4.5", mapping["grok"])
|
||||
require.Equal(t, "grok-4.5", mapping["grok-latest"])
|
||||
require.Equal(t, "grok-4.6", mapping["grok-4.6"])
|
||||
require.Equal(t, "grok-4.6", mapping["grok-4.6-latest"])
|
||||
require.Equal(t, "grok-4.5", mapping["grok-4.5"])
|
||||
require.Equal(t, "grok-4.5", mapping["grok-4.5-latest"])
|
||||
require.Equal(t, "grok-build-0.1", mapping["grok-build"])
|
||||
|
||||
@@ -43,6 +43,12 @@ type QuotaSnapshot struct {
|
||||
LastProbeAt string `json:"last_probe_at,omitempty"`
|
||||
LastHeadersSeenAt string `json:"last_headers_seen_at,omitempty"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
// Model is the upstream id that produced these rate-limit headers.
|
||||
Model string `json:"model,omitempty"`
|
||||
// PlanFrom45Responses is inferred from a grok-4.5 Responses window
|
||||
// (8300/53M = Heavy). Carried across later non-4.5 overwrites.
|
||||
PlanFrom45Responses string `json:"plan_from_45_responses,omitempty"`
|
||||
PlanFrom45ResponsesAt string `json:"plan_from_45_responses_at,omitempty"`
|
||||
}
|
||||
|
||||
func (s *QuotaSnapshot) HasObservedHeaders() bool {
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
package xai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// GrokQuotaSignalMaxAge bounds how long a grok-4.5 Responses window can
|
||||
// influence SuperGrok vs Heavy inference.
|
||||
const GrokQuotaSignalMaxAge = 24 * time.Hour
|
||||
|
||||
const (
|
||||
grok45ResponsesModel = "grok-4.5"
|
||||
grokHeavyQuotaRequestLimit int64 = 8_300
|
||||
grokHeavyQuotaTokenLimit int64 = 53_000_000
|
||||
)
|
||||
|
||||
// MapJWTSubscriptionTier maps prod_auth.SubscriptionTier numeric JWT claims
|
||||
// to stable snake_case keys used by Grok Build / Mixpanel.
|
||||
func MapJWTSubscriptionTier(tier uint64) string {
|
||||
switch tier {
|
||||
case 0:
|
||||
return "free"
|
||||
case 1:
|
||||
return "supergrok"
|
||||
case 2:
|
||||
return "x_basic"
|
||||
case 3:
|
||||
return "x_premium"
|
||||
case 4:
|
||||
return "x_premium_plus"
|
||||
case 5:
|
||||
return "supergrok_heavy"
|
||||
case 6:
|
||||
return "supergrok_lite"
|
||||
case 7:
|
||||
return "supergrok_plus"
|
||||
default:
|
||||
return strconv.FormatUint(tier, 10)
|
||||
}
|
||||
}
|
||||
|
||||
// NormalizeSubscriptionTier canonicalizes display names, /user strings, and
|
||||
// JWT-derived keys onto the same snake_case identifiers.
|
||||
func NormalizeSubscriptionTier(raw string) string {
|
||||
t := strings.ToLower(strings.TrimSpace(raw))
|
||||
t = strings.ReplaceAll(t, "-", "_")
|
||||
t = strings.Join(strings.Fields(t), "_")
|
||||
switch t {
|
||||
case "free", "grok_free", "grokfree", "free_tier", "freetier", "grok_basic", "grokbasic":
|
||||
return "free"
|
||||
case "supergrok", "grokpro":
|
||||
return "supergrok"
|
||||
case "supergrok_lite", "supergroklite":
|
||||
return "supergrok_lite"
|
||||
case "supergrok_heavy", "supergrokheavy":
|
||||
return "supergrok_heavy"
|
||||
case "supergrok_pro", "supergrokpro":
|
||||
return "supergrok_pro"
|
||||
case "supergrok_plus", "supergrokplus":
|
||||
return "supergrok_plus"
|
||||
case "x_basic", "xbasic", "basic":
|
||||
return "x_basic"
|
||||
case "x_premium", "xpremium":
|
||||
return "x_premium"
|
||||
case "x_premium_plus", "xpremiumplus", "x_premium+":
|
||||
return "x_premium_plus"
|
||||
default:
|
||||
return t
|
||||
}
|
||||
}
|
||||
|
||||
// SubscriptionTierFromJWT decodes an access token payload (no signature check)
|
||||
// and maps the numeric or string `tier` claim.
|
||||
func SubscriptionTierFromJWT(jwt string) string {
|
||||
claims := DecodeJWTClaims(jwt)
|
||||
if claims == nil {
|
||||
return ""
|
||||
}
|
||||
raw, ok := claims["tier"]
|
||||
if !ok || raw == nil {
|
||||
return ""
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case float64:
|
||||
if v < 0 {
|
||||
return ""
|
||||
}
|
||||
return MapJWTSubscriptionTier(uint64(v))
|
||||
case json.Number:
|
||||
n, err := v.Int64()
|
||||
if err != nil || n < 0 {
|
||||
return NormalizeSubscriptionTier(v.String())
|
||||
}
|
||||
return MapJWTSubscriptionTier(uint64(n))
|
||||
case string:
|
||||
trimmed := strings.TrimSpace(v)
|
||||
if trimmed == "" {
|
||||
return ""
|
||||
}
|
||||
if n, err := strconv.ParseUint(trimmed, 10, 64); err == nil {
|
||||
return MapJWTSubscriptionTier(n)
|
||||
}
|
||||
return NormalizeSubscriptionTier(trimmed)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// CanonicalGrokPlan resolves SuperGrok vs Heavy when the provider label is
|
||||
// ambiguous (SuperGrokPro). JWT numeric claims are applied by the caller first.
|
||||
// Monthly $150/$1500 limits still win when present.
|
||||
// Rate-limit windows are only used when they came from grok-4.5 Responses.
|
||||
func CanonicalGrokPlan(monthlyLimitCents *float64, subscriptionTier string, quota *QuotaSnapshot) string {
|
||||
if plan := resolvePlan(monthlyLimitCents); plan != "" {
|
||||
return NormalizeSubscriptionTier(plan)
|
||||
}
|
||||
|
||||
normalized := NormalizeSubscriptionTier(subscriptionTier)
|
||||
switch normalized {
|
||||
case "free", "x_basic":
|
||||
return "free"
|
||||
case "supergrok_heavy":
|
||||
return "supergrok_heavy"
|
||||
case "supergrok_lite":
|
||||
return "supergrok_lite"
|
||||
case "supergrok_plus":
|
||||
return "supergrok_plus"
|
||||
}
|
||||
|
||||
if isAmbiguousGrokPaidPlan(normalized) {
|
||||
if hint := Grok45ResponsesPlanHint(quota, time.Time{}); hint != "" {
|
||||
return hint
|
||||
}
|
||||
return "supergrok"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func isAmbiguousGrokPaidPlan(normalized string) bool {
|
||||
switch normalized {
|
||||
case "supergrok", "supergrok_pro", "paid", "pro":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// IsGrok45ResponsesQuotaModel reports whether model is the grok-4.5 Responses
|
||||
// id (or a dated grok-4.5-* variant). Empty and other families are false.
|
||||
func IsGrok45ResponsesQuotaModel(model string) bool {
|
||||
m := strings.ToLower(strings.TrimSpace(StripGrokProviderPrefix(model)))
|
||||
return m == grok45ResponsesModel || strings.HasPrefix(m, grok45ResponsesModel+"-")
|
||||
}
|
||||
|
||||
// Grok45ResponsesPlanHint returns SuperGrok / Heavy inferred from a grok-4.5
|
||||
// Responses window. Other models' limits are ignored.
|
||||
func Grok45ResponsesPlanHint(quota *QuotaSnapshot, now time.Time) string {
|
||||
if quota == nil {
|
||||
return ""
|
||||
}
|
||||
if plan := NormalizeSubscriptionTier(quota.PlanFrom45Responses); plan == "supergrok" || plan == "supergrok_heavy" {
|
||||
if isQuotaTimestampFresh(quota.PlanFrom45ResponsesAt, now) {
|
||||
return plan
|
||||
}
|
||||
}
|
||||
if !IsGrok45ResponsesQuotaModel(quota.Model) || !IsQuotaSnapshotFresh(quota, now) {
|
||||
return ""
|
||||
}
|
||||
if quotaLooksLikeGrokHeavy(quota) {
|
||||
return "supergrok_heavy"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// ApplyGrok45ResponsesPlanSignal records a grok-4.5 Heavy/SuperGrok hint, or
|
||||
// copies the previous 4.5 hint when this observation is a different model.
|
||||
func (s *QuotaSnapshot) ApplyGrok45ResponsesPlanSignal(prev *QuotaSnapshot) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
observedAt := firstNonEmptyQuotaTime(s.LastHeadersSeenAt, s.UpdatedAt)
|
||||
if IsGrok45ResponsesQuotaModel(s.Model) && quotaHasLimitWindow(s) {
|
||||
if quotaLooksLikeGrokHeavy(s) {
|
||||
s.PlanFrom45Responses = "supergrok_heavy"
|
||||
s.PlanFrom45ResponsesAt = observedAt
|
||||
return
|
||||
}
|
||||
s.PlanFrom45Responses = "supergrok"
|
||||
s.PlanFrom45ResponsesAt = observedAt
|
||||
return
|
||||
}
|
||||
if prev != nil && strings.TrimSpace(prev.PlanFrom45Responses) != "" {
|
||||
s.PlanFrom45Responses = prev.PlanFrom45Responses
|
||||
s.PlanFrom45ResponsesAt = prev.PlanFrom45ResponsesAt
|
||||
}
|
||||
}
|
||||
|
||||
// QuotaSnapshotObservedAt prefers LastHeadersSeenAt over UpdatedAt so a later
|
||||
// snapshot rewrite cannot refresh a stale Heavy window.
|
||||
func QuotaSnapshotObservedAt(snapshot *QuotaSnapshot) (time.Time, bool) {
|
||||
if snapshot == nil {
|
||||
return time.Time{}, false
|
||||
}
|
||||
return parseQuotaTimestamp(firstNonEmptyQuotaTime(snapshot.LastHeadersSeenAt, snapshot.UpdatedAt))
|
||||
}
|
||||
|
||||
// IsQuotaSnapshotFresh reports whether a quota signal is recent enough to
|
||||
// distinguish SuperGrok from Heavy.
|
||||
func IsQuotaSnapshotFresh(snapshot *QuotaSnapshot, now time.Time) bool {
|
||||
observedAt, ok := QuotaSnapshotObservedAt(snapshot)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return isTimeFresh(observedAt, now)
|
||||
}
|
||||
|
||||
func isQuotaTimestampFresh(raw string, now time.Time) bool {
|
||||
parsed, ok := parseQuotaTimestamp(raw)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return isTimeFresh(parsed, now)
|
||||
}
|
||||
|
||||
func parseQuotaTimestamp(raw string) (time.Time, bool) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return time.Time{}, false
|
||||
}
|
||||
parsed, err := time.Parse(time.RFC3339, raw)
|
||||
if err != nil {
|
||||
return time.Time{}, false
|
||||
}
|
||||
return parsed, true
|
||||
}
|
||||
|
||||
func isTimeFresh(observedAt, now time.Time) bool {
|
||||
if now.IsZero() {
|
||||
now = time.Now()
|
||||
}
|
||||
age := now.Sub(observedAt)
|
||||
return age <= GrokQuotaSignalMaxAge && age >= -5*time.Minute
|
||||
}
|
||||
|
||||
func quotaHasLimitWindow(quota *QuotaSnapshot) bool {
|
||||
if quota == nil {
|
||||
return false
|
||||
}
|
||||
if quota.Requests != nil && quota.Requests.Limit != nil {
|
||||
return true
|
||||
}
|
||||
return quota.Tokens != nil && quota.Tokens.Limit != nil
|
||||
}
|
||||
|
||||
func quotaLooksLikeGrokHeavy(quota *QuotaSnapshot) bool {
|
||||
if quota == nil {
|
||||
return false
|
||||
}
|
||||
var requestLimit, tokenLimit int64
|
||||
if quota.Requests != nil && quota.Requests.Limit != nil {
|
||||
requestLimit = *quota.Requests.Limit
|
||||
}
|
||||
if quota.Tokens != nil && quota.Tokens.Limit != nil {
|
||||
tokenLimit = *quota.Tokens.Limit
|
||||
}
|
||||
return requestLimit >= grokHeavyQuotaRequestLimit || tokenLimit >= grokHeavyQuotaTokenLimit
|
||||
}
|
||||
|
||||
func firstNonEmptyQuotaTime(values ...string) string {
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
//go:build unit
|
||||
|
||||
package xai
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMapJWTSubscriptionTierNumber(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.Equal(t, "free", MapJWTSubscriptionTier(0))
|
||||
require.Equal(t, "supergrok", MapJWTSubscriptionTier(1))
|
||||
require.Equal(t, "x_basic", MapJWTSubscriptionTier(2))
|
||||
require.Equal(t, "x_premium", MapJWTSubscriptionTier(3))
|
||||
require.Equal(t, "x_premium_plus", MapJWTSubscriptionTier(4))
|
||||
require.Equal(t, "supergrok_heavy", MapJWTSubscriptionTier(5))
|
||||
require.Equal(t, "supergrok_lite", MapJWTSubscriptionTier(6))
|
||||
require.Equal(t, "supergrok_plus", MapJWTSubscriptionTier(7))
|
||||
require.Equal(t, "9", MapJWTSubscriptionTier(9))
|
||||
}
|
||||
|
||||
func TestNormalizeSubscriptionTierAliases(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.Equal(t, "free", NormalizeSubscriptionTier("Free"))
|
||||
require.Equal(t, "free", NormalizeSubscriptionTier(" FREE "))
|
||||
require.Equal(t, "supergrok", NormalizeSubscriptionTier("SuperGrok"))
|
||||
require.Equal(t, "supergrok_heavy", NormalizeSubscriptionTier("SuperGrok Heavy"))
|
||||
require.Equal(t, "supergrok_pro", NormalizeSubscriptionTier("SuperGrokPro"))
|
||||
require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("SuperGrok Lite"))
|
||||
require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("SuperGrokLite"))
|
||||
require.Equal(t, "x_basic", NormalizeSubscriptionTier("X Basic"))
|
||||
require.Equal(t, "free", NormalizeSubscriptionTier("free-tier"))
|
||||
require.Equal(t, "free", NormalizeSubscriptionTier("free_tier"))
|
||||
require.Equal(t, "free", NormalizeSubscriptionTier("grok-basic"))
|
||||
require.Equal(t, "free", NormalizeSubscriptionTier("grok_basic"))
|
||||
require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("supergrok_lite"))
|
||||
}
|
||||
|
||||
func TestSubscriptionTierFromJWTUsesNumericClaim(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.Equal(t, "supergrok_heavy", SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"tier": 5})))
|
||||
require.Equal(t, "free", SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"tier": 0})))
|
||||
require.Equal(t, "supergrok_lite", SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"tier": 6})))
|
||||
require.Empty(t, SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"sub": "user"})))
|
||||
require.Empty(t, SubscriptionTierFromJWT("not-a-jwt"))
|
||||
}
|
||||
|
||||
func TestCanonicalGrokPlanUsesOnlyGrok45ResponsesWindow(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
zero := float64(0)
|
||||
heavyReq, heavyTok := int64(8300), int64(53_000_000)
|
||||
superReq, superTok := int64(900), int64(15_000_000)
|
||||
fresh := time.Now().UTC().Format(time.RFC3339)
|
||||
stale := time.Now().Add(-GrokQuotaSignalMaxAge - time.Hour).UTC().Format(time.RFC3339)
|
||||
|
||||
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", nil))
|
||||
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrok", nil))
|
||||
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrok Heavy", nil))
|
||||
require.Empty(t, CanonicalGrokPlan(&zero, "", nil))
|
||||
require.Equal(t, "free", CanonicalGrokPlan(&zero, "free", nil))
|
||||
|
||||
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
|
||||
Model: "grok-4.5",
|
||||
Requests: &QuotaWindow{Limit: &heavyReq},
|
||||
Tokens: &QuotaWindow{Limit: &heavyTok},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}))
|
||||
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
|
||||
Model: "grok-4.6",
|
||||
Requests: &QuotaWindow{Limit: &heavyReq},
|
||||
Tokens: &QuotaWindow{Limit: &heavyTok},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}))
|
||||
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
|
||||
Requests: &QuotaWindow{Limit: &heavyReq},
|
||||
Tokens: &QuotaWindow{Limit: &heavyTok},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}))
|
||||
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
|
||||
Model: "grok-4.5",
|
||||
Requests: &QuotaWindow{Limit: &superReq},
|
||||
Tokens: &QuotaWindow{Limit: &superTok},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}))
|
||||
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
|
||||
Model: "grok-4.5",
|
||||
Requests: &QuotaWindow{Limit: &heavyReq},
|
||||
LastHeadersSeenAt: stale,
|
||||
}))
|
||||
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
|
||||
Model: "grok-4.6",
|
||||
Requests: &QuotaWindow{Limit: &superReq},
|
||||
PlanFrom45Responses: "supergrok_heavy",
|
||||
PlanFrom45ResponsesAt: fresh,
|
||||
}))
|
||||
require.Equal(t, "free", CanonicalGrokPlan(&zero, "free", &QuotaSnapshot{
|
||||
Model: "grok-4.5",
|
||||
Requests: &QuotaWindow{Limit: &heavyReq},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}))
|
||||
heavyCents := float64(SuperGrokHeavyLimitCents)
|
||||
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&heavyCents, "SuperGrokPro", nil))
|
||||
}
|
||||
|
||||
func TestApplyGrok45ResponsesPlanSignalCarriesHint(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
heavyReq := int64(8300)
|
||||
fresh := time.Now().UTC().Format(time.RFC3339)
|
||||
prev := &QuotaSnapshot{
|
||||
Model: "grok-4.5",
|
||||
Requests: &QuotaWindow{Limit: &heavyReq},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}
|
||||
prev.ApplyGrok45ResponsesPlanSignal(nil)
|
||||
require.Equal(t, "supergrok_heavy", prev.PlanFrom45Responses)
|
||||
|
||||
next := &QuotaSnapshot{
|
||||
Model: "grok-4.6",
|
||||
Requests: &QuotaWindow{Limit: int64Ptr(100)},
|
||||
LastHeadersSeenAt: fresh,
|
||||
}
|
||||
next.ApplyGrok45ResponsesPlanSignal(prev)
|
||||
require.Equal(t, "supergrok_heavy", next.PlanFrom45Responses)
|
||||
require.Equal(t, prev.PlanFrom45ResponsesAt, next.PlanFrom45ResponsesAt)
|
||||
}
|
||||
|
||||
func int64Ptr(v int64) *int64 { return &v }
|
||||
|
||||
func jwtWithClaims(t *testing.T, claims map[string]any) string {
|
||||
t.Helper()
|
||||
payload, err := json.Marshal(claims)
|
||||
require.NoError(t, err)
|
||||
return "header." + base64.RawURLEncoding.EncodeToString(payload) + ".sig"
|
||||
}
|
||||
@@ -3,8 +3,10 @@ package repository
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -196,6 +198,8 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
|
||||
group.FieldAudioRealtimePricePerMin,
|
||||
group.FieldAudioTtsPricePerMillionChars,
|
||||
group.FieldAudioSttPricePerHour,
|
||||
group.FieldLongContextPricingEnabled,
|
||||
group.FieldModelPricing,
|
||||
group.FieldClaudeCodeOnly,
|
||||
group.FieldFallbackGroupID,
|
||||
group.FieldFallbackGroupIDOnInvalidRequest,
|
||||
@@ -946,6 +950,14 @@ func groupEntityToService(g *dbent.Group) *service.Group {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
var modelPricing []service.ChannelModelPricing
|
||||
if len(g.ModelPricing) > 0 {
|
||||
if err := json.Unmarshal(g.ModelPricing, &modelPricing); err != nil {
|
||||
slog.Warn("group model_pricing unmarshal failed; falling back to channel/builtin pricing",
|
||||
"group_id", g.ID, "error", err)
|
||||
modelPricing = nil
|
||||
}
|
||||
}
|
||||
return &service.Group{
|
||||
ID: g.ID,
|
||||
Name: g.Name,
|
||||
@@ -980,6 +992,8 @@ func groupEntityToService(g *dbent.Group) *service.Group {
|
||||
AudioRealtimePricePerMin: g.AudioRealtimePricePerMin,
|
||||
AudioTTSPricePerMillionChars: g.AudioTtsPricePerMillionChars,
|
||||
AudioSTTPricePerHour: g.AudioSttPricePerHour,
|
||||
LongContextPricingEnabled: g.LongContextPricingEnabled,
|
||||
ModelPricing: modelPricing,
|
||||
DefaultValidityDays: g.DefaultValidityDays,
|
||||
ClaudeCodeOnly: g.ClaudeCodeOnly,
|
||||
FallbackGroupID: g.FallbackGroupID,
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"time"
|
||||
|
||||
@@ -59,6 +60,34 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
|
||||
return int64(len(data)), nil
|
||||
}
|
||||
|
||||
func (s *S3BackupStore) UploadFile(ctx context.Context, key string, filePath string, contentType string) (int64, error) {
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("open upload file: %w", err)
|
||||
}
|
||||
defer func() { _ = file.Close() }()
|
||||
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("stat upload file: %w", err)
|
||||
}
|
||||
sizeBytes := info.Size()
|
||||
|
||||
finish := servertiming.ObserveDependency(ctx, "s3")
|
||||
_, err = s.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: &s.bucket,
|
||||
Key: &key,
|
||||
Body: file,
|
||||
ContentLength: &sizeBytes,
|
||||
ContentType: &contentType,
|
||||
})
|
||||
finish()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("S3 PutObject file: %w", err)
|
||||
}
|
||||
return sizeBytes, nil
|
||||
}
|
||||
|
||||
func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser, error) {
|
||||
finish := servertiming.ObserveDependency(ctx, "s3")
|
||||
result, err := s.client.GetObject(ctx, &s3.GetObjectInput{
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestS3BackupStore_UploadFile(t *testing.T) {
|
||||
var received []byte
|
||||
var receivedLength int64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, http.MethodPut, r.Method)
|
||||
receivedLength = r.ContentLength
|
||||
var err error
|
||||
received, err = io.ReadAll(r.Body)
|
||||
require.NoError(t, err)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client, err := newS3Client(context.Background(), s3ClientParams{
|
||||
Endpoint: server.URL,
|
||||
Region: "auto",
|
||||
AccessKeyID: "test-ak",
|
||||
SecretAccessKey: "test-sk",
|
||||
ForcePathStyle: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
content := []byte("streamed backup payload")
|
||||
filePath := t.TempDir() + "/part.gz"
|
||||
require.NoError(t, os.WriteFile(filePath, content, 0o600))
|
||||
|
||||
store := &S3BackupStore{client: client, bucket: "backup-bucket"}
|
||||
size, err := store.UploadFile(context.Background(), "backup/part-1", filePath, "application/octet-stream")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(len(content)), size)
|
||||
require.Equal(t, int64(len(content)), receivedLength)
|
||||
require.Equal(t, content, received)
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Context, todayStart time.Time) (results []usagestats.GroupUsageSummary, err error) {
|
||||
todayStart = service.GroupUsageTodayStart(todayStart)
|
||||
yesterdayStart := service.GroupUsageYesterdayStart(todayStart)
|
||||
timezoneName := service.GroupUsageTimezoneName()
|
||||
todayDate := service.GroupUsageDate(todayStart)
|
||||
yesterdayDate := service.GroupUsageDate(yesterdayStart)
|
||||
|
||||
const query = `
|
||||
WITH state_values AS (
|
||||
SELECT
|
||||
COUNT(*) = 1
|
||||
AND MAX(timezone_name) = $3
|
||||
AND MAX(closed_before) <= $4::date AS valid,
|
||||
MAX(closed_before) AS closed_before,
|
||||
MAX(retained_from) AS retained_from
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
),
|
||||
state AS (
|
||||
SELECT
|
||||
CASE WHEN valid THEN closed_before ELSE DATE '1970-01-01' END AS closed_before,
|
||||
CASE WHEN valid THEN retained_from ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00' END AS retained_from,
|
||||
CASE
|
||||
WHEN valid THEN closed_before::timestamp AT TIME ZONE $3::text
|
||||
ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00'
|
||||
END AS tail_start,
|
||||
valid
|
||||
FROM state_values
|
||||
),
|
||||
historical AS (
|
||||
SELECT
|
||||
rollup.group_id,
|
||||
COALESCE(SUM(rollup.actual_cost), 0) AS actual_cost,
|
||||
COALESCE(SUM(rollup.actual_cost) FILTER (
|
||||
WHERE rollup.bucket_date = $5::date
|
||||
), 0) AS yesterday_cost
|
||||
FROM usage_group_daily_rollups rollup
|
||||
CROSS JOIN state
|
||||
WHERE state.valid
|
||||
AND rollup.bucket_date >= (state.retained_from AT TIME ZONE $3::text)::date
|
||||
AND rollup.bucket_date < state.closed_before
|
||||
GROUP BY rollup.group_id
|
||||
),
|
||||
tail AS (
|
||||
SELECT
|
||||
ul.group_id,
|
||||
COALESCE(SUM(ul.actual_cost), 0) AS actual_cost,
|
||||
COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $1), 0) AS today_cost,
|
||||
COALESCE(SUM(ul.actual_cost) FILTER (
|
||||
WHERE ul.created_at >= $2
|
||||
AND ul.created_at < $1
|
||||
), 0) AS yesterday_cost
|
||||
FROM usage_logs ul
|
||||
CROSS JOIN state
|
||||
WHERE ul.created_at >= state.tail_start
|
||||
GROUP BY ul.group_id
|
||||
)
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
COALESCE(historical.actual_cost, 0) + COALESCE(tail.actual_cost, 0) AS total_cost,
|
||||
COALESCE(tail.today_cost, 0) AS today_cost,
|
||||
COALESCE(historical.yesterday_cost, 0) + COALESCE(tail.yesterday_cost, 0) AS yesterday_cost
|
||||
FROM groups g
|
||||
LEFT JOIN historical ON historical.group_id = g.id
|
||||
LEFT JOIN tail ON tail.group_id = g.id
|
||||
ORDER BY g.id
|
||||
`
|
||||
|
||||
rows, err := r.sql.QueryContext(
|
||||
ctx,
|
||||
query,
|
||||
todayStart,
|
||||
yesterdayStart,
|
||||
timezoneName,
|
||||
todayDate,
|
||||
yesterdayDate,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := rows.Close(); closeErr != nil && err == nil {
|
||||
err = closeErr
|
||||
results = nil
|
||||
}
|
||||
}()
|
||||
|
||||
results = make([]usagestats.GroupUsageSummary, 0)
|
||||
for rows.Next() {
|
||||
var row usagestats.GroupUsageSummary
|
||||
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost, &row.YesterdayCost); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// SyncGroupUsageRollups 将服务端配置时区今日以前的用量发布为分组日桶。
|
||||
func (r *dashboardAggregationRepository) SyncGroupUsageRollups(ctx context.Context, todayStart time.Time) error {
|
||||
if r == nil || r.sql == nil {
|
||||
return nil
|
||||
}
|
||||
todayStart = service.GroupUsageTodayStart(todayStart)
|
||||
if db, ok := r.sql.(*sql.DB); ok {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
txRepo := newDashboardAggregationRepositoryWithSQL(tx)
|
||||
if err := txRepo.syncGroupUsageRollupsInTx(ctx, todayStart); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
return r.syncGroupUsageRollupsInTx(ctx, todayStart)
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) syncGroupUsageRollupsInTx(ctx context.Context, todayStart time.Time) error {
|
||||
var closedBefore string
|
||||
var previousRetainedFrom time.Time
|
||||
var stateTimezoneName string
|
||||
if err := scanSingleRow(ctx, r.sql, `
|
||||
SELECT closed_before::text, retained_from, timezone_name
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`, nil, &closedBefore, &previousRetainedFrom, &stateTimezoneName); err != nil {
|
||||
return fmt.Errorf("读取分组用量汇总水位: %w", err)
|
||||
}
|
||||
|
||||
todayDate := service.GroupUsageDate(todayStart)
|
||||
timezoneName := service.GroupUsageTimezoneName()
|
||||
timezoneChanged := stateTimezoneName != timezoneName
|
||||
var closedTime time.Time
|
||||
if !timezoneChanged {
|
||||
var err error
|
||||
closedTime, err = service.ParseGroupUsageDate(closedBefore)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析分组用量汇总水位 %q: %w", closedBefore, err)
|
||||
}
|
||||
todayDateTime, err := service.ParseGroupUsageDate(todayDate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if closedTime.After(todayDateTime) {
|
||||
return fmt.Errorf("分组用量汇总水位位于未来: %s", closedBefore)
|
||||
}
|
||||
if closedBefore == todayDate {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
var earliest sql.NullTime
|
||||
if err := scanSingleRow(ctx, r.sql, "SELECT MIN(created_at) FROM usage_logs", nil, &earliest); err != nil {
|
||||
return fmt.Errorf("读取最早用量记录: %w", err)
|
||||
}
|
||||
retainedFrom := todayStart
|
||||
if earliest.Valid {
|
||||
retainedFrom = earliest.Time.UTC()
|
||||
}
|
||||
retainedDate := service.GroupUsageDate(retainedFrom)
|
||||
retainedDateTime, err := service.ParseGroupUsageDate(retainedDate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rebuildStartDate := retainedDate
|
||||
if !timezoneChanged && closedTime.After(retainedDateTime) {
|
||||
rebuildStartDate = closedBefore
|
||||
}
|
||||
rebuildStart, err := service.ParseGroupUsageDate(rebuildStartDate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := r.sql.ExecContext(ctx, `
|
||||
DELETE FROM usage_group_daily_rollups
|
||||
WHERE bucket_date < $1::date
|
||||
OR (bucket_date >= $2::date AND bucket_date < $3::date)
|
||||
OR bucket_date >= $3::date
|
||||
`, retainedDate, rebuildStartDate, todayDate); err != nil {
|
||||
return fmt.Errorf("清理分组用量日桶: %w", err)
|
||||
}
|
||||
|
||||
if _, err := r.sql.ExecContext(ctx, `
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
SELECT
|
||||
(created_at AT TIME ZONE $3::text)::date AS bucket_date,
|
||||
group_id,
|
||||
COALESCE(SUM(actual_cost), 0) AS actual_cost,
|
||||
NOW()
|
||||
FROM usage_logs
|
||||
WHERE group_id IS NOT NULL
|
||||
AND created_at >= $1
|
||||
AND created_at < $2
|
||||
GROUP BY 1, 2
|
||||
ON CONFLICT (bucket_date, group_id)
|
||||
DO UPDATE SET
|
||||
actual_cost = EXCLUDED.actual_cost,
|
||||
computed_at = EXCLUDED.computed_at
|
||||
`, rebuildStart.UTC(), todayStart.UTC(), timezoneName); err != nil {
|
||||
return fmt.Errorf("重建分组用量日桶: %w", err)
|
||||
}
|
||||
|
||||
if _, err := r.sql.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = $1::date,
|
||||
retained_from = $2,
|
||||
timezone_name = $3,
|
||||
updated_at = NOW()
|
||||
WHERE id = 1
|
||||
`, todayDate, retainedFrom, timezoneName); err != nil {
|
||||
return fmt.Errorf("更新分组用量汇总水位: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func lockGroupUsageRollupState(ctx context.Context, tx *sql.Tx) error {
|
||||
var id int16
|
||||
if err := tx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`).Scan(&id); err != nil {
|
||||
return fmt.Errorf("锁定分组用量汇总水位: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func invalidateGroupUsageRollupsAt(ctx context.Context, tx *sql.Tx, affectedAt time.Time) error {
|
||||
timezoneName := service.GroupUsageTimezoneName()
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = LEAST(
|
||||
closed_before,
|
||||
($1::timestamptz AT TIME ZONE $2::text)::date
|
||||
),
|
||||
updated_at = NOW()
|
||||
WHERE id = 1
|
||||
`, affectedAt.UTC(), timezoneName)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
appTimezone "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func useGroupUsageRepositoryTestTimezone(t *testing.T, name string) {
|
||||
t.Helper()
|
||||
|
||||
previousName := appTimezone.Name()
|
||||
require.NoError(t, appTimezone.Init(name))
|
||||
t.Cleanup(func() { require.NoError(t, appTimezone.Init(previousName)) })
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsNoopsAtCurrentDate(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-08-14", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRebuildsWhenTimezoneChanges(t *testing.T) {
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
retainedFrom := time.Date(2026, 3, 1, 5, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-03-09", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
|
||||
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
|
||||
WithArgs("2026-03-01", "2026-03-01", "2026-03-09").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
|
||||
WithArgs(retainedFrom, todayStart, "America/New_York").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs("2026-03-09", retainedFrom, "America/New_York").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsPublishesWatermarkLast(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
retainedFrom := time.Date(2026, 5, 1, 3, 0, 0, 0, time.UTC)
|
||||
rebuildStart := time.Date(2026, 8, 12, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-08-13", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
|
||||
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
|
||||
WithArgs("2026-05-01", "2026-08-13", "2026-08-14").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
|
||||
WithArgs(rebuildStart, todayStart, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 2))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs("2026-08-14", retainedFrom, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRejectsFutureWatermark(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-08-15", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.SyncGroupUsageRollups(context.Background(), todayStart)
|
||||
require.ErrorContains(t, err, "未来")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryRecomputeRangeInvalidatesGroupRollupsBeforeDashboardRebuild(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
end := start.Add(time.Hour)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`DELETE FROM usage_dashboard_hourly`).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.RecomputeRange(context.Background(), start, end)
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryRecomputeRangeRebuildsGroupRollupsBeforeCommit(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
end := start.Add(time.Hour)
|
||||
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
|
||||
repo.clock = func() time.Time { return fixedNow }
|
||||
todayStart := service.GroupUsageTodayStart(fixedNow)
|
||||
startDate := service.GroupUsageDate(start)
|
||||
rebuildStart, err := service.ParseGroupUsageDate(startDate)
|
||||
require.NoError(t, err)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
for _, query := range []string{
|
||||
`DELETE FROM usage_dashboard_hourly WHERE`,
|
||||
`DELETE FROM usage_dashboard_hourly_users WHERE`,
|
||||
`DELETE FROM usage_dashboard_daily WHERE`,
|
||||
`DELETE FROM usage_dashboard_daily_users WHERE`,
|
||||
`INSERT INTO usage_dashboard_hourly_users`,
|
||||
`INSERT INTO usage_dashboard_daily_users`,
|
||||
`INSERT INTO usage_dashboard_hourly`,
|
||||
`INSERT INTO usage_dashboard_daily`,
|
||||
} {
|
||||
mock.ExpectExec(query).WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
}
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow(startDate, time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(start))
|
||||
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
|
||||
WithArgs(startDate, startDate, service.GroupUsageDate(todayStart)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
|
||||
WithArgs(rebuildStart.UTC(), todayStart.UTC(), "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(service.GroupUsageDate(todayStart), start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.RecomputeRange(context.Background(), start, end))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedInvalidatesEachBatchAndSyncs(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
|
||||
earliestDeletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
|
||||
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
|
||||
repo.clock = func() time.Time { return fixedNow }
|
||||
todayStart := service.GroupUsageTodayStart(fixedNow)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(cutoff, usageLogsCleanupBatchSize).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
|
||||
AddRow(earliestDeletedAt.Add(time.Hour)).
|
||||
AddRow(earliestDeletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(earliestDeletedAt, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionedSortsAndInvalidatesEachDropBeforeSync(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
|
||||
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
|
||||
juneStart := time.Date(2026, 6, 1, 0, 0, 0, 0, time.UTC)
|
||||
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
|
||||
repo.clock = func() time.Time { return fixedNow }
|
||||
todayStart := service.GroupUsageTodayStart(fixedNow)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectQuery(`SELECT c.relname`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
|
||||
AddRow("usage_logs_202606").
|
||||
AddRow("usage_logs_invalid").
|
||||
AddRow("usage_logs_202604").
|
||||
AddRow("usage_logs_202607"))
|
||||
|
||||
for _, partition := range []struct {
|
||||
name string
|
||||
start time.Time
|
||||
}{
|
||||
{name: "usage_logs_202604", start: aprilStart},
|
||||
{name: "usage_logs_202606", start: juneStart},
|
||||
} {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(partition.start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`DROP TABLE IF EXISTS "` + partition.name + `"`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectCommit()
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedFailureRollsBackWithoutSync(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
|
||||
deletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)SELECT ctid.*ORDER BY created_at ASC, id ASC.*DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(cutoff, usageLogsCleanupBatchSize).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(deletedAt, "Asia/Shanghai").
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.CleanupUsageLogs(context.Background(), cutoff)
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionFailureRollsBackAndStops(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
|
||||
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
|
||||
dropErr := errors.New("drop partition failed")
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectQuery(`SELECT c.relname`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
|
||||
AddRow("usage_logs_202606").
|
||||
AddRow("usage_logs_202604"))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(aprilStart, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`DROP TABLE IF EXISTS "usage_logs_202604"`).
|
||||
WillReturnError(dropErr)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.CleanupUsageLogs(context.Background(), cutoff)
|
||||
require.ErrorIs(t, err, dropErr)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func setGroupUsageRollupTestTimezone(t *testing.T) {
|
||||
t.Helper()
|
||||
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -14,7 +15,8 @@ import (
|
||||
)
|
||||
|
||||
type dashboardAggregationRepository struct {
|
||||
sql sqlExecutor
|
||||
sql sqlExecutor
|
||||
clock func() time.Time
|
||||
}
|
||||
|
||||
const usageLogsCleanupBatchSize = 10000
|
||||
@@ -33,7 +35,14 @@ func NewDashboardAggregationRepository(sqlDB *sql.DB) service.DashboardAggregati
|
||||
}
|
||||
|
||||
func newDashboardAggregationRepositoryWithSQL(sqlq sqlExecutor) *dashboardAggregationRepository {
|
||||
return &dashboardAggregationRepository{sql: sqlq}
|
||||
return &dashboardAggregationRepository{sql: sqlq, clock: time.Now}
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) now() time.Time {
|
||||
if r.clock != nil {
|
||||
return r.clock()
|
||||
}
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
func isPostgresDriver(db *sql.DB) bool {
|
||||
@@ -128,11 +137,23 @@ func (r *dashboardAggregationRepository) RecomputeRange(ctx context.Context, sta
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, start); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
txRepo := newDashboardAggregationRepositoryWithSQL(tx)
|
||||
if err := txRepo.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
if err := txRepo.syncGroupUsageRollupsInTx(ctx, service.GroupUsageTodayStart(r.now())); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
return r.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd)
|
||||
@@ -215,14 +236,35 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
|
||||
return err
|
||||
}
|
||||
if isPartitioned {
|
||||
return r.dropUsageLogsPartitions(ctx, cutoff)
|
||||
if err := r.dropUsageLogsPartitions(ctx, cutoff); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if err := r.cleanupUsageLogsBatches(ctx, cutoff); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.SyncGroupUsageRollups(ctx, service.GroupUsageTodayStart(r.now()))
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) cleanupUsageLogsBatches(ctx context.Context, cutoff time.Time) error {
|
||||
db, transactional := r.sql.(*sql.DB)
|
||||
for {
|
||||
if transactional {
|
||||
affected, err := cleanupUsageLogsBatchWithRollupInvalidation(ctx, db, cutoff)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected < usageLogsCleanupBatchSize {
|
||||
return nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
res, err := r.sql.ExecContext(ctx, `
|
||||
WITH victims AS (
|
||||
SELECT ctid
|
||||
FROM usage_logs
|
||||
WHERE created_at < $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
@@ -241,6 +283,66 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
|
||||
}
|
||||
}
|
||||
|
||||
func cleanupUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, cutoff time.Time) (int64, error) {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
rollback := func(err error) (int64, error) {
|
||||
_ = tx.Rollback()
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
rows, err := tx.QueryContext(ctx, `
|
||||
WITH victims AS (
|
||||
SELECT ctid
|
||||
FROM usage_logs
|
||||
WHERE created_at < $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE ctid IN (SELECT ctid FROM victims)
|
||||
RETURNING created_at
|
||||
`, cutoff.UTC(), usageLogsCleanupBatchSize)
|
||||
if err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
|
||||
var affected int64
|
||||
var earliestDeletedAt time.Time
|
||||
for rows.Next() {
|
||||
var deletedAt time.Time
|
||||
if err := rows.Scan(&deletedAt); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
affected++
|
||||
if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
|
||||
earliestDeletedAt = deletedAt
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
if affected > 0 {
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return affected, nil
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) CleanupUsageBillingDedup(ctx context.Context, cutoff time.Time) error {
|
||||
for {
|
||||
res, err := r.sql.ExecContext(ctx, `
|
||||
@@ -489,14 +591,16 @@ func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Con
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
_ = rows.Close()
|
||||
}()
|
||||
|
||||
cutoffMonth := truncateToMonthUTC(cutoff)
|
||||
type usageLogsPartition struct {
|
||||
name string
|
||||
month time.Time
|
||||
}
|
||||
partitions := make([]usageLogsPartition, 0)
|
||||
for rows.Next() {
|
||||
var name string
|
||||
if err := rows.Scan(&name); err != nil {
|
||||
_ = rows.Close()
|
||||
return err
|
||||
}
|
||||
if !strings.HasPrefix(name, "usage_logs_") {
|
||||
@@ -509,12 +613,56 @@ func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Con
|
||||
}
|
||||
month = month.UTC()
|
||||
if month.Before(cutoffMonth) {
|
||||
if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
|
||||
partitions = append(partitions, usageLogsPartition{name: name, month: month})
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return err
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sort.Slice(partitions, func(i, j int) bool {
|
||||
return partitions[i].month.Before(partitions[j].month)
|
||||
})
|
||||
if db, ok := r.sql.(*sql.DB); ok {
|
||||
for _, partition := range partitions {
|
||||
if err := dropUsageLogsPartitionWithRollupInvalidation(ctx, db, partition.name, partition.month); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return rows.Err()
|
||||
for _, partition := range partitions {
|
||||
if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(partition.name))); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func dropUsageLogsPartitionWithRollupInvalidation(ctx context.Context, db *sql.DB, name string, monthStart time.Time) error {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rollback := func(err error) error {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, monthStart); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) createUsageLogsPartition(ctx context.Context, month time.Time) error {
|
||||
|
||||
@@ -3,6 +3,7 @@ package repository
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
@@ -56,6 +57,10 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
|
||||
if groupIn == nil {
|
||||
return errors.New("group is nil")
|
||||
}
|
||||
modelPricing, err := json.Marshal(groupIn.ModelPricing)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal group model pricing: %w", err)
|
||||
}
|
||||
builder := client.Group.Create().
|
||||
SetName(groupIn.Name).
|
||||
SetDescription(groupIn.Description).
|
||||
@@ -88,6 +93,8 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
|
||||
SetNillableAudioRealtimePricePerMin(groupIn.AudioRealtimePricePerMin).
|
||||
SetNillableAudioTtsPricePerMillionChars(groupIn.AudioTTSPricePerMillionChars).
|
||||
SetNillableAudioSttPricePerHour(groupIn.AudioSTTPricePerHour).
|
||||
SetLongContextPricingEnabled(groupIn.LongContextPricingEnabled).
|
||||
SetModelPricing(modelPricing).
|
||||
SetDefaultValidityDays(groupIn.DefaultValidityDays).
|
||||
SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
|
||||
SetNillableFallbackGroupID(groupIn.FallbackGroupID).
|
||||
@@ -234,6 +241,10 @@ func (r *groupRepository) GetByIDLite(ctx context.Context, id int64) (*service.G
|
||||
}
|
||||
|
||||
func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) error {
|
||||
modelPricing, err := json.Marshal(groupIn.ModelPricing)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal group model pricing: %w", err)
|
||||
}
|
||||
builder := r.client.Group.UpdateOneID(groupIn.ID).
|
||||
SetName(groupIn.Name).
|
||||
SetDescription(groupIn.Description).
|
||||
@@ -260,6 +271,8 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
|
||||
SetNillableVideoPrice720p(groupIn.VideoPrice720P).
|
||||
SetNillableVideoPrice1080p(groupIn.VideoPrice1080P).
|
||||
SetVideoModelPrices(service.NormalizeVideoModelPrices(groupIn.VideoModelPrices)).
|
||||
SetLongContextPricingEnabled(groupIn.LongContextPricingEnabled).
|
||||
SetModelPricing(modelPricing).
|
||||
SetDefaultValidityDays(groupIn.DefaultValidityDays).
|
||||
SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
|
||||
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
|
||||
|
||||
@@ -0,0 +1,511 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/migrations"
|
||||
"github.com/lib/pq"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGroupUsageRollupTriggerInvalidatesCascadedHistoricalDelete(t *testing.T) {
|
||||
for _, partitioned := range []bool{false, true} {
|
||||
name := "ordinary"
|
||||
if partitioned {
|
||||
name = "partitioned"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, partitioned)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2020-01-02 08:00:00+08');
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1;
|
||||
DELETE FROM users WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var closedBefore string
|
||||
err = tx.QueryRowContext(ctx, `
|
||||
SELECT closed_before::text
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "2020-01-02", closedBefore)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerSerializesLateHistoricalInsertWithPublish(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
seedTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
_, err := seedTx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '2020-01-02'
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, seedTx.Commit())
|
||||
|
||||
syncTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = syncTx.Rollback() }()
|
||||
var stateID int16
|
||||
require.NoError(t, syncTx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`).Scan(&stateID))
|
||||
|
||||
lateTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = lateTx.Rollback() }()
|
||||
var lateBackendPID int
|
||||
require.NoError(t, lateTx.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&lateBackendPID))
|
||||
|
||||
insertResult := make(chan error, 1)
|
||||
go func() {
|
||||
_, insertErr := lateTx.ExecContext(ctx, `
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2020-01-02 09:00:00+08')
|
||||
`)
|
||||
insertResult <- insertErr
|
||||
}()
|
||||
|
||||
blocked, err := waitForGroupUsageRollupStateLock(ctx, lateBackendPID, insertResult)
|
||||
if err != nil || !blocked {
|
||||
_ = syncTx.Rollback()
|
||||
_ = lateTx.Rollback()
|
||||
require.NoError(t, err)
|
||||
require.True(t, blocked, "迟到写入必须等待正在发布水位的事务")
|
||||
}
|
||||
|
||||
_, err = syncTx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, syncTx.Commit())
|
||||
|
||||
select {
|
||||
case err = <-insertResult:
|
||||
require.NoError(t, err)
|
||||
case <-ctx.Done():
|
||||
t.Fatal("等待迟到写入完成超时")
|
||||
}
|
||||
require.NoError(t, lateTx.Commit())
|
||||
|
||||
var closedBefore string
|
||||
err = integrationDB.QueryRowContext(ctx, fmt.Sprintf(
|
||||
"SELECT closed_before::text FROM %s.usage_group_rollup_state WHERE id = 1",
|
||||
pq.QuoteIdentifier(schema),
|
||||
)).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "2020-01-02", closedBefore)
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerSerializesInsertTransactionAcrossMidnight(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
seedTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
_, err := seedTx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, seedTx.Commit())
|
||||
|
||||
syncTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = syncTx.Rollback() }()
|
||||
var stateID int16
|
||||
require.NoError(t, syncTx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`).Scan(&stateID))
|
||||
|
||||
insertTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = insertTx.Rollback() }()
|
||||
var insertBackendPID int
|
||||
require.NoError(t, insertTx.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&insertBackendPID))
|
||||
|
||||
insertResult := make(chan error, 1)
|
||||
go func() {
|
||||
_, insertErr := insertTx.ExecContext(ctx, `
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, CURRENT_TIMESTAMP)
|
||||
`)
|
||||
insertResult <- insertErr
|
||||
}()
|
||||
|
||||
blocked, err := waitForGroupUsageRollupStateLock(ctx, insertBackendPID, insertResult)
|
||||
if err != nil || !blocked {
|
||||
_ = syncTx.Rollback()
|
||||
_ = insertTx.Rollback()
|
||||
require.NoError(t, err)
|
||||
require.True(t, blocked, "跨越零点的在途写入必须与水位发布串行化")
|
||||
}
|
||||
|
||||
_, err = syncTx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date + 1
|
||||
WHERE id = 1
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, syncTx.Commit())
|
||||
|
||||
select {
|
||||
case err = <-insertResult:
|
||||
require.NoError(t, err)
|
||||
case <-ctx.Done():
|
||||
t.Fatal("等待跨零点写入完成超时")
|
||||
}
|
||||
require.NoError(t, insertTx.Commit())
|
||||
|
||||
var currentDate string
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, `
|
||||
SELECT (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date::text
|
||||
`).Scan(¤tDate))
|
||||
var closedBefore string
|
||||
err = integrationDB.QueryRowContext(ctx, fmt.Sprintf(
|
||||
"SELECT closed_before::text FROM %s.usage_group_rollup_state WHERE id = 1",
|
||||
pq.QuoteIdentifier(schema),
|
||||
)).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, currentDate, closedBefore)
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerKeepsWatermarkForTodayInsert(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1;
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, CURRENT_TIMESTAMP);
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var unchanged bool
|
||||
err = tx.QueryRowContext(ctx, `
|
||||
SELECT closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&unchanged)
|
||||
require.NoError(t, err)
|
||||
require.True(t, unchanged)
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerUsesSessionTimezoneAcrossDST(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
SET LOCAL TIME ZONE 'America/New_York';
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '2026-03-09',
|
||||
timezone_name = 'America/New_York'
|
||||
WHERE id = 1;
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2026-03-08 04:30:00+00');
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var closedBefore string
|
||||
err = tx.QueryRowContext(ctx, `
|
||||
SELECT closed_before::text
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "2026-03-07", closedBefore)
|
||||
}
|
||||
|
||||
func TestGroupUsageSummaryIncludesYesterdayAcrossWatermark(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
closedBefore string
|
||||
includeYesterday bool
|
||||
}{
|
||||
{name: "closed_rollup", closedBefore: "2026-08-14", includeYesterday: true},
|
||||
{name: "raw_tail", closedBefore: "2026-08-13", includeYesterday: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
|
||||
(1, 1, 10, 2, TIMESTAMPTZ '2026-08-12 12:00:00+08'),
|
||||
(2, 1, 10, 3, TIMESTAMPTZ '2026-08-13 12:00:00+08'),
|
||||
(3, 1, 10, 4, TIMESTAMPTZ '2026-08-14 12:00:00+08');
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
VALUES (DATE '2026-08-12', 10, 2, NOW());
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
if tt.includeYesterday {
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
VALUES (DATE '2026-08-13', 10, 3, NOW())
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = $1::date,
|
||||
retained_from = TIMESTAMPTZ '2026-08-12 00:00:00+08'
|
||||
WHERE id = 1
|
||||
`, tt.closedBefore)
|
||||
require.NoError(t, err)
|
||||
|
||||
repo := newUsageLogRepositoryWithSQL(nil, tx)
|
||||
result, err := repo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result, 1)
|
||||
require.InDelta(t, 9, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 4, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 3, result[0].YesterdayCost, 0.0000001)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupSyncRebuildsAfterTimezoneChange(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
SET LOCAL TIME ZONE 'UTC';
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
|
||||
(1, 1, 10, 3, TIMESTAMPTZ '2026-03-08 05:30:00+00'),
|
||||
(2, 1, 10, 5, TIMESTAMPTZ '2026-03-09 04:30:00+00');
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
VALUES (DATE '2026-03-08', 10, 99, NOW());
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '2026-03-09',
|
||||
retained_from = TIMESTAMPTZ '2026-03-08 05:30:00+00',
|
||||
timezone_name = 'Asia/Shanghai'
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
repo := newDashboardAggregationRepositoryWithSQL(tx)
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(ctx, todayStart))
|
||||
|
||||
var stateTimezone string
|
||||
var closedBefore string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
SELECT timezone_name, closed_before::text
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&stateTimezone, &closedBefore))
|
||||
require.Equal(t, "America/New_York", stateTimezone)
|
||||
require.Equal(t, "2026-03-09", closedBefore)
|
||||
|
||||
var rollupCost float64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
SELECT actual_cost
|
||||
FROM usage_group_daily_rollups
|
||||
WHERE bucket_date = DATE '2026-03-08' AND group_id = 10
|
||||
`).Scan(&rollupCost))
|
||||
require.InDelta(t, 3, rollupCost, 0.0000001)
|
||||
|
||||
usageRepo := newUsageLogRepositoryWithSQL(nil, tx)
|
||||
result, err := usageRepo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result, 1)
|
||||
require.InDelta(t, 8, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 5, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 3, result[0].YesterdayCost, 0.0000001)
|
||||
}
|
||||
|
||||
func TestGroupUsageSummaryUsesConfiguredDSTBoundaries(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
SET LOCAL TIME ZONE 'UTC';
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
|
||||
(1, 1, 10, 100, TIMESTAMPTZ '2026-03-08 04:30:00+00'),
|
||||
(2, 1, 10, 3, TIMESTAMPTZ '2026-03-08 05:30:00+00'),
|
||||
(3, 1, 10, 4, TIMESTAMPTZ '2026-03-09 03:30:00+00'),
|
||||
(4, 1, 10, 5, TIMESTAMPTZ '2026-03-09 04:30:00+00');
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '1970-01-01',
|
||||
retained_from = TIMESTAMPTZ '1970-01-01 00:00:00+00',
|
||||
timezone_name = 'America/New_York'
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
repo := newUsageLogRepositoryWithSQL(nil, tx)
|
||||
result, err := repo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result, 1)
|
||||
require.InDelta(t, 112, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 5, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 7, result[0].YesterdayCost, 0.0000001)
|
||||
}
|
||||
|
||||
func createGroupUsageRollupTriggerTestSchema(t *testing.T, ctx context.Context, partitioned bool) string {
|
||||
t.Helper()
|
||||
|
||||
schema := fmt.Sprintf("group_usage_rollup_trigger_%d", time.Now().UnixNano())
|
||||
quotedSchema := pq.QuoteIdentifier(schema)
|
||||
_, err := integrationDB.ExecContext(ctx, "CREATE SCHEMA "+quotedSchema)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DROP SCHEMA IF EXISTS "+quotedSchema+" CASCADE")
|
||||
})
|
||||
|
||||
tx, err := integrationDB.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
require.NoError(t, setGroupUsageRollupTriggerSearchPath(ctx, tx, quotedSchema))
|
||||
|
||||
usageLogsDDL := `
|
||||
CREATE TABLE usage_logs (
|
||||
id BIGINT PRIMARY KEY,
|
||||
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
group_id BIGINT REFERENCES groups(id) ON DELETE SET NULL,
|
||||
actual_cost NUMERIC(20, 10) NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL
|
||||
);
|
||||
`
|
||||
if partitioned {
|
||||
usageLogsDDL = `
|
||||
CREATE TABLE usage_logs (
|
||||
id BIGINT NOT NULL,
|
||||
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
group_id BIGINT REFERENCES groups(id) ON DELETE SET NULL,
|
||||
actual_cost NUMERIC(20, 10) NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL
|
||||
) PARTITION BY RANGE (created_at);
|
||||
CREATE TABLE usage_logs_default PARTITION OF usage_logs DEFAULT;
|
||||
`
|
||||
}
|
||||
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
CREATE TABLE users (id BIGINT PRIMARY KEY);
|
||||
CREATE TABLE groups (id BIGINT PRIMARY KEY);
|
||||
`+usageLogsDDL)
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, migrationName := range []string{
|
||||
"222_group_usage_daily_rollups.sql",
|
||||
"223_group_usage_rollup_timezone.sql",
|
||||
} {
|
||||
migrationSQL, readErr := migrations.FS.ReadFile(migrationName)
|
||||
require.NoError(t, readErr)
|
||||
for range 2 {
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
require.NoError(t, tx.Commit())
|
||||
|
||||
return schema
|
||||
}
|
||||
|
||||
func beginGroupUsageRollupTriggerTestTx(t *testing.T, ctx context.Context, schema string) *sql.Tx {
|
||||
t.Helper()
|
||||
|
||||
tx, err := integrationDB.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, setGroupUsageRollupTriggerSearchPath(ctx, tx, pq.QuoteIdentifier(schema)))
|
||||
return tx
|
||||
}
|
||||
|
||||
func setGroupUsageRollupTriggerSearchPath(ctx context.Context, tx *sql.Tx, quotedSchema string) error {
|
||||
_, err := tx.ExecContext(ctx, "SET LOCAL search_path TO "+quotedSchema)
|
||||
return err
|
||||
}
|
||||
|
||||
func waitForGroupUsageRollupStateLock(
|
||||
ctx context.Context,
|
||||
backendPID int,
|
||||
insertResult <-chan error,
|
||||
) (bool, error) {
|
||||
ticker := time.NewTicker(10 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case err := <-insertResult:
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return false, nil
|
||||
case <-ticker.C:
|
||||
var waitEventType sql.NullString
|
||||
err := integrationDB.QueryRowContext(ctx, `
|
||||
SELECT wait_event_type
|
||||
FROM pg_stat_activity
|
||||
WHERE pid = $1
|
||||
`, backendPID).Scan(&waitEventType)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if waitEventType.Valid && waitEventType.String == "Lock" {
|
||||
return true, nil
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return false, ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -291,6 +291,9 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
|
||||
return 0, fmt.Errorf("cleanup filters missing time range")
|
||||
}
|
||||
args = append(args, limit)
|
||||
if db, ok := r.sql.(*sql.DB); ok {
|
||||
return r.deleteUsageLogsBatchWithRollupInvalidation(ctx, db, whereClause, args)
|
||||
}
|
||||
query := fmt.Sprintf(`
|
||||
WITH target AS (
|
||||
SELECT id
|
||||
@@ -301,7 +304,7 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE id IN (SELECT id FROM target)
|
||||
RETURNING id
|
||||
RETURNING created_at
|
||||
`, whereClause, len(args))
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
@@ -320,6 +323,68 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func (r *usageCleanupRepository) deleteUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, whereClause string, args []any) (int64, error) {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
rollback := func(err error) (int64, error) {
|
||||
_ = tx.Rollback()
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
query := fmt.Sprintf(`
|
||||
WITH target AS (
|
||||
SELECT id
|
||||
FROM usage_logs
|
||||
WHERE %s
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $%d
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE id IN (SELECT id FROM target)
|
||||
RETURNING created_at
|
||||
`, whereClause, len(args))
|
||||
rows, err := tx.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
|
||||
var deleted int64
|
||||
var earliestDeletedAt time.Time
|
||||
for rows.Next() {
|
||||
var deletedAt time.Time
|
||||
if err := rows.Scan(&deletedAt); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
deleted++
|
||||
if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
|
||||
earliestDeletedAt = deletedAt
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
|
||||
if deleted > 0 {
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func buildUsageCleanupWhere(filters service.UsageCleanupFilters) (string, []any) {
|
||||
conditions := make([]string, 0, 8)
|
||||
args := make([]any, 0, 8)
|
||||
|
||||
@@ -21,6 +21,11 @@ func newSQLMock(t *testing.T) (*sql.DB, sqlmock.Sqlmock) {
|
||||
return db, mock
|
||||
}
|
||||
|
||||
func setUsageCleanupRollupTestTimezone(t *testing.T) {
|
||||
t.Helper()
|
||||
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
|
||||
}
|
||||
|
||||
func TestNewUsageCleanupRepository(t *testing.T) {
|
||||
db, _ := newSQLMock(t)
|
||||
repo := NewUsageCleanupRepository(nil, db)
|
||||
@@ -398,6 +403,7 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatchMissingRange(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
|
||||
setUsageCleanupRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
|
||||
@@ -412,9 +418,16 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
|
||||
Model: &model,
|
||||
}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery("DELETE FROM usage_logs").
|
||||
WithArgs(start, end, userID, "gpt-4", 2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(1)).AddRow(int64(2)))
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(start.Add(time.Hour)).AddRow(start.Add(2 * time.Hour)))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(start.Add(time.Hour), "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
|
||||
require.NoError(t, err)
|
||||
@@ -422,6 +435,62 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatchAtomicallyInvalidatesGroupRollups(t *testing.T) {
|
||||
setUsageCleanupRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
|
||||
start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
||||
end := start.Add(24 * time.Hour)
|
||||
firstDeletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
secondDeletedAt := firstDeletedAt.Add(time.Hour)
|
||||
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(start, end, 2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
|
||||
AddRow(firstDeletedAt).
|
||||
AddRow(secondDeletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(firstDeletedAt, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), deleted)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatchRollsBackWhenInvalidationFails(t *testing.T) {
|
||||
setUsageCleanupRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
|
||||
start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
||||
end := start.Add(24 * time.Hour)
|
||||
deletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(start, end, 1).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(deletedAt, "Asia/Shanghai").
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
_, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 1)
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
@@ -430,9 +499,13 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
|
||||
end := start.Add(24 * time.Hour)
|
||||
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery("DELETE FROM usage_logs").
|
||||
WithArgs(start, end, 5).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
_, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 5)
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUsageLogRepositoryGetAllGroupUsageSummaryUsesRollupTail(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newUsageLogRepositoryWithSQL(nil, db)
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
yesterdayStart := time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectQuery(`(?s)usage_group_rollup_state.*usage_group_daily_rollups.*created_at >= state\.tail_start`).
|
||||
WithArgs(todayStart, yesterdayStart, "America/New_York", "2026-03-09", "2026-03-08").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"group_id", "total_cost", "today_cost", "yesterday_cost"}).
|
||||
AddRow(int64(7), 12.5, 1.25, 2.5))
|
||||
|
||||
result, err := repo.GetAllGroupUsageSummary(context.Background(), todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(7), result[0].GroupID)
|
||||
require.InDelta(t, 12.5, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 1.25, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 2.5, result[0].YesterdayCost, 0.0000001)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
@@ -715,39 +715,9 @@ func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTim
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// GetAllGroupUsageSummary returns today's and cumulative actual_cost for every group.
|
||||
// todayStart is the start-of-day in the caller's timezone (UTC-based).
|
||||
// TODO(perf): This query scans ALL usage_logs rows for total_cost aggregation.
|
||||
// When usage_logs exceeds ~1M rows, consider adding a short-lived cache (30s)
|
||||
// or a materialized view / pre-aggregation table for cumulative costs.
|
||||
// GetAllGroupUsageSummary 返回所有分组在服务端配置时区内的今日、昨日与当前保留记录累计金额。
|
||||
func (r *usageLogRepository) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
query := `
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
COALESCE(SUM(ul.actual_cost), 0) AS total_cost,
|
||||
COALESCE(SUM(CASE WHEN ul.created_at >= $1 THEN ul.actual_cost ELSE 0 END), 0) AS today_cost
|
||||
FROM groups g
|
||||
LEFT JOIN usage_logs ul ON ul.group_id = g.id
|
||||
GROUP BY g.id
|
||||
`
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, todayStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var results []usagestats.GroupUsageSummary
|
||||
for rows.Next() {
|
||||
var row usagestats.GroupUsageSummary
|
||||
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
return r.getAllGroupUsageSummaryFromRollups(ctx, todayStart)
|
||||
}
|
||||
|
||||
// resolveModelDimensionExpression maps model source type to a safe SQL expression.
|
||||
|
||||
@@ -364,6 +364,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"daily_limit_usd": null,
|
||||
"weekly_limit_usd": null,
|
||||
"monthly_limit_usd": null,
|
||||
"long_context_pricing_enabled": false,
|
||||
"image_price_1k": null,
|
||||
"image_price_2k": null,
|
||||
"image_price_4k": null,
|
||||
@@ -1467,7 +1468,7 @@ func newContractDeps(t *testing.T) *contractDeps {
|
||||
settingRepo := newStubSettingRepo()
|
||||
settingService := service.NewSettingService(settingRepo, cfg)
|
||||
|
||||
adminService := service.NewAdminService(userRepo, groupRepo, &accountRepo, proxyRepo, apiKeyRepo, redeemRepo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
adminService := service.NewAdminService(userRepo, groupRepo, &accountRepo, proxyRepo, apiKeyRepo, redeemRepo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
authHandler := handler.NewAuthHandler(cfg, nil, userService, settingService, nil, redeemService, nil, nil)
|
||||
apiKeyHandler := handler.NewAPIKeyHandler(apiKeyService)
|
||||
usageHandler := handler.NewUsageHandler(usageService, apiKeyService, nil, nil)
|
||||
|
||||
@@ -315,6 +315,14 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.Gateway.WebSearch(c)
|
||||
})
|
||||
gateway.POST("/x_search", func(c *gin.Context) {
|
||||
if getGroupPlatform(c) != service.PlatformGrok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
|
||||
return
|
||||
}
|
||||
h.Gateway.XSearch(c)
|
||||
})
|
||||
}
|
||||
|
||||
// Gemini 原生 API 兼容层(Gemini SDK/CLI 直连)
|
||||
@@ -443,6 +451,14 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.Gateway.WebSearch(c)
|
||||
})
|
||||
r.POST("/x_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
if getGroupPlatform(c) != service.PlatformGrok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
|
||||
return
|
||||
}
|
||||
h.Gateway.XSearch(c)
|
||||
})
|
||||
|
||||
// Antigravity 模型列表
|
||||
r.GET("/antigravity/models", gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.Gateway.AntigravityModels)
|
||||
|
||||
@@ -48,6 +48,7 @@ func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
|
||||
"/models/*modelAction": {"gemini_v1beta_handler.go"},
|
||||
"/tts": {"grok_audio.go"},
|
||||
"/web_search": {"gateway_web_search.go"},
|
||||
"/x_search": {"gateway_web_search.go"},
|
||||
}
|
||||
excluded := map[string]string{
|
||||
"/messages/count_tokens": "tokenization only; it does not execute a model request",
|
||||
|
||||
@@ -53,7 +53,7 @@ func EvaluateAccountSchedulingThreshold(account *Account, thresholds map[string]
|
||||
var winner *accountSchedulingThresholdCandidate
|
||||
switch decision.Platform {
|
||||
case PlatformOpenAI:
|
||||
winner = pickLatestResetSchedulingCandidate(openAIThresholdCandidates(account), threshold, now)
|
||||
winner = pickLatestResetSchedulingCandidate(openAIThresholdCandidates(account, now), threshold, now)
|
||||
case PlatformAnthropic:
|
||||
winner = pickLatestResetSchedulingCandidate(anthropicThresholdCandidates(account), threshold, now)
|
||||
case PlatformGrok:
|
||||
@@ -149,7 +149,7 @@ func lookupAccountSchedulingThreshold(thresholds map[string]int, platform string
|
||||
return value, ok
|
||||
}
|
||||
|
||||
func openAIThresholdCandidates(account *Account) []*accountSchedulingThresholdCandidate {
|
||||
func openAIThresholdCandidates(account *Account, now time.Time) []*accountSchedulingThresholdCandidate {
|
||||
if account == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -157,8 +157,8 @@ func openAIThresholdCandidates(account *Account) []*accountSchedulingThresholdCa
|
||||
return nil
|
||||
}
|
||||
return []*accountSchedulingThresholdCandidate{
|
||||
openAIThresholdCandidate(account.Extra, "5h"),
|
||||
openAIThresholdCandidate(account.Extra, "7d"),
|
||||
openAIThresholdCandidate(account.Extra, "5h", now),
|
||||
openAIThresholdCandidate(account.Extra, "7d", now),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -219,7 +219,7 @@ func firstStringValue(values map[string]any, keys ...string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func openAIThresholdCandidate(extra map[string]any, window string) *accountSchedulingThresholdCandidate {
|
||||
func openAIThresholdCandidate(extra map[string]any, window string, now time.Time) *accountSchedulingThresholdCandidate {
|
||||
if len(extra) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -243,9 +243,12 @@ func openAIThresholdCandidate(extra map[string]any, window string) *accountSched
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if openAIQuotaWindowReset(extra, window, now) || openAICodexSnapshotStaleForPause(extra, now) {
|
||||
return nil
|
||||
}
|
||||
return &accountSchedulingThresholdCandidate{
|
||||
window: window,
|
||||
usedPercent: utilizationAsPercent(usedPercent),
|
||||
usedPercent: schedulingPercentValue(usedPercent),
|
||||
until: parseSchedulingResetAt(extra[resetAtKey]),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -96,7 +96,7 @@ func TestEvaluateAccountSchedulingThreshold_AnthropicIgnoresExpiredFiveHourWindo
|
||||
require.True(t, wantUntil.Equal(*decision.Until))
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_FractionalPlatformsKeepFractionSemantics(t *testing.T) {
|
||||
func TestEvaluateAccountSchedulingThreshold_OpenAIPreservesPercentageSemantics(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
@@ -104,17 +104,114 @@ func TestEvaluateAccountSchedulingThreshold_FractionalPlatformsKeepFractionSeman
|
||||
openAIAccount := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{
|
||||
"codex_5h_used_percent": 0.91,
|
||||
"codex_5h_used_percent": 1.0,
|
||||
"codex_5h_reset_at": openAIUntil.Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
|
||||
candidate := openAIThresholdCandidate(openAIAccount.Extra, "5h", now)
|
||||
require.NotNil(t, candidate)
|
||||
require.Equal(t, 1.0, candidate.usedPercent)
|
||||
|
||||
openAIDecision := EvaluateAccountSchedulingThreshold(openAIAccount, map[string]int{
|
||||
PlatformOpenAI: 90,
|
||||
}, now)
|
||||
require.False(t, openAIDecision.ShouldPause)
|
||||
|
||||
openAIAccount.Extra["codex_5h_used_percent"] = 91.0
|
||||
openAIDecision = EvaluateAccountSchedulingThreshold(openAIAccount, map[string]int{
|
||||
PlatformOpenAI: 90,
|
||||
}, now)
|
||||
require.True(t, openAIDecision.ShouldPause)
|
||||
require.Equal(t, 91.0, openAIDecision.UsedPercent)
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_OpenAISkipsStaleSnapshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{
|
||||
"codex_usage_updated_at": now.Add(-2 * time.Hour).Format(time.RFC3339),
|
||||
"codex_5h_used_percent": 100.0,
|
||||
"codex_5h_reset_at": now.Add(3 * time.Hour).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
|
||||
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
|
||||
|
||||
require.False(t, decision.ShouldPause)
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_OpenAISkipsResetWindow(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{
|
||||
"codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339),
|
||||
"codex_5h_used_percent": 100.0,
|
||||
"codex_5h_reset_at": now.Add(-time.Second).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
|
||||
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
|
||||
|
||||
require.False(t, decision.ShouldPause)
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_OpenAIPausesFreshExhaustedSnapshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
resetAt := now.Add(3 * time.Hour)
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{
|
||||
"codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339),
|
||||
"codex_5h_used_percent": 100.0,
|
||||
"codex_5h_reset_at": resetAt.Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
|
||||
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
|
||||
|
||||
require.True(t, decision.ShouldPause)
|
||||
require.Equal(t, "5h", decision.Window)
|
||||
require.Equal(t, 100.0, decision.UsedPercent)
|
||||
require.NotNil(t, decision.Until)
|
||||
require.True(t, resetAt.Equal(*decision.Until))
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_OpenAIPausesFreshExhaustedSevenDayWindow(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
resetAt := now.Add(5 * 24 * time.Hour)
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{
|
||||
"codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339),
|
||||
"codex_7d_used_percent": 95.0,
|
||||
"codex_7d_reset_at": resetAt.Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
|
||||
decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformOpenAI: 90}, now)
|
||||
|
||||
require.True(t, decision.ShouldPause)
|
||||
require.Equal(t, "7d", decision.Window)
|
||||
require.Equal(t, 95.0, decision.UsedPercent)
|
||||
require.NotNil(t, decision.Until)
|
||||
require.True(t, resetAt.Equal(*decision.Until))
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_AnthropicPreservesFractionalUtilizationSemantics(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
anthropicUntil := now.Add(5 * time.Hour)
|
||||
anthropicAccount := &Account{
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
//
|
||||
// upstreamModel 是最终发往上游的模型 ID。
|
||||
// totalCost 是本次请求的客户计费(倍率前),用于优先级 2。
|
||||
// serviceTier 是最终参与用户计费的 OpenAI 服务层级,用于优先级 3。
|
||||
func resolveAccountStatsCost(
|
||||
ctx context.Context,
|
||||
channelService *ChannelService,
|
||||
@@ -26,6 +27,7 @@ func resolveAccountStatsCost(
|
||||
tokens UsageTokens,
|
||||
requestCount int,
|
||||
totalCost float64,
|
||||
serviceTier string,
|
||||
) *float64 {
|
||||
if channelService == nil || upstreamModel == "" {
|
||||
return nil
|
||||
@@ -53,20 +55,22 @@ func resolveAccountStatsCost(
|
||||
|
||||
// 优先级 3:模型定价文件(LiteLLM)默认价格
|
||||
if billingService != nil {
|
||||
return tryModelFilePricing(billingService, upstreamModel, tokens)
|
||||
return tryModelFilePricing(billingService, upstreamModel, tokens, serviceTier)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// tryModelFilePricing 使用模型定价文件(LiteLLM/fallback)中的标准价格计算费用。
|
||||
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens) *float64 {
|
||||
// tryModelFilePricing 使用模型定价文件(LiteLLM/fallback)中的价格计算费用。
|
||||
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens, serviceTier string) *float64 {
|
||||
pricing, err := billingService.GetModelPricing(model)
|
||||
if err != nil || pricing == nil {
|
||||
return nil
|
||||
}
|
||||
if billingService.shouldApplySessionLongContextPricing(tokens, pricing) {
|
||||
breakdown, err := billingService.CalculateCost(model, tokens, 1)
|
||||
normalizedTier := normalizeBillingServiceTier(serviceTier)
|
||||
if normalizedTier == "priority" || normalizedTier == "flex" ||
|
||||
billingService.shouldApplySessionLongContextPricing(tokens, pricing) {
|
||||
breakdown, err := billingService.CalculateCostWithServiceTier(model, tokens, 1, normalizedTier)
|
||||
if err != nil || breakdown == nil || breakdown.TotalCost <= 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -241,7 +245,11 @@ func applyAccountStatsCost(
|
||||
if usageLog != nil && usageLog.ImageCount > 0 {
|
||||
requestCount = usageLog.ImageCount
|
||||
}
|
||||
serviceTier := ""
|
||||
if usageLog != nil && usageLog.ServiceTier != nil {
|
||||
serviceTier = *usageLog.ServiceTier
|
||||
}
|
||||
usageLog.AccountStatsCost = resolveAccountStatsCost(
|
||||
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost,
|
||||
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost, serviceTier,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -454,7 +454,7 @@ func TestTryModelFilePricing_Success(t *testing.T) {
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens)
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
|
||||
require.InDelta(t, 0.2, *result, 1e-12)
|
||||
@@ -473,18 +473,87 @@ func TestTryModelFilePricing_AppliesLongContextPricing(t *testing.T) {
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens)
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "")
|
||||
|
||||
require.NotNil(t, result)
|
||||
// Input and cache-read use the 2x input tier; output uses the 1.5x tier.
|
||||
require.InDelta(t, 0.233, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestTryModelFilePricing_AppliesServiceTierPricing(t *testing.T) {
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
|
||||
"gpt-5.6-sol": {
|
||||
InputPricePerToken: 0.001,
|
||||
InputPricePerTokenPriority: 0.002,
|
||||
OutputPricePerToken: 0.002,
|
||||
OutputPricePerTokenPriority: 0.004,
|
||||
CacheCreationPricePerToken: 0.003,
|
||||
CacheCreationPricePerTokenPriority: 0.006,
|
||||
CacheReadPricePerToken: 0.0005,
|
||||
CacheReadPricePerTokenPriority: 0.001,
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
CacheCreationTokens: 20,
|
||||
CacheReadTokens: 10,
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
serviceTier string
|
||||
want float64
|
||||
}{
|
||||
{name: "standard", serviceTier: "", want: 0.265},
|
||||
{name: "priority", serviceTier: "priority", want: 0.53},
|
||||
{name: "flex", serviceTier: "flex", want: 0.1325},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, tt.want, *result, 1e-12)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryModelFilePricing_CombinesPriorityAndLongContextPricing(t *testing.T) {
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
|
||||
"gpt-5.6-sol": {
|
||||
InputPricePerToken: 0.001,
|
||||
InputPricePerTokenPriority: 0.002,
|
||||
OutputPricePerToken: 0.002,
|
||||
OutputPricePerTokenPriority: 0.004,
|
||||
CacheCreationPricePerToken: 0.003,
|
||||
CacheCreationPricePerTokenPriority: 0.006,
|
||||
CacheReadPricePerToken: 0.0005,
|
||||
CacheReadPricePerTokenPriority: 0.001,
|
||||
LongContextInputThreshold: 100,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 1.5,
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{
|
||||
InputTokens: 101,
|
||||
OutputTokens: 10,
|
||||
CacheCreationTokens: 5,
|
||||
CacheReadTokens: 5,
|
||||
}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority")
|
||||
|
||||
require.NotNil(t, result)
|
||||
// priority 单价先应用,再叠加长上下文输入 2x、输出 1.5x。
|
||||
require.InDelta(t, 0.534, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestTryModelFilePricing_PricingNotFound(t *testing.T) {
|
||||
// "nonexistent-model" does not match any fallback pattern
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
|
||||
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
|
||||
result := tryModelFilePricing(bs, "nonexistent-model", tokens)
|
||||
result := tryModelFilePricing(bs, "nonexistent-model", tokens, "")
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -494,7 +563,7 @@ func TestTryModelFilePricing_NilFallback(t *testing.T) {
|
||||
"claude-sonnet-4": nil,
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens)
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -506,7 +575,7 @@ func TestTryModelFilePricing_ZeroCost(t *testing.T) {
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{} // all zero tokens → cost = 0 → nil
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens)
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -523,7 +592,7 @@ func TestTryModelFilePricing_WithImageOutput(t *testing.T) {
|
||||
OutputTokens: 50,
|
||||
ImageOutputTokens: 10,
|
||||
}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens)
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 + 10*0.01 = 0.1 + 0.1 + 0.1 = 0.3
|
||||
require.InDelta(t, 0.3, *result, 1e-12)
|
||||
@@ -544,7 +613,7 @@ func TestTryModelFilePricing_WithCacheTokens(t *testing.T) {
|
||||
CacheCreationTokens: 200,
|
||||
CacheReadTokens: 300,
|
||||
}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens)
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
|
||||
// = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
|
||||
@@ -561,7 +630,7 @@ func TestResolveAccountStatsCost_NilChannelService(t *testing.T) {
|
||||
nil, // channelService is nil
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 1, "claude-sonnet-4",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5,
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "",
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -577,7 +646,7 @@ func TestResolveAccountStatsCost_EmptyUpstreamModel(t *testing.T) {
|
||||
cs,
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 1, "", // empty upstream model
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5,
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "",
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -594,7 +663,7 @@ func TestResolveAccountStatsCost_GetChannelForGroupReturnsNil(t *testing.T) {
|
||||
cs,
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 99, "claude-sonnet-4", // groupID 99 has no channel
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5,
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "",
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -625,7 +694,7 @@ func TestResolveAccountStatsCost_HitsCustomRule(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, nil, // billingService not needed when custom rule hits
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 999.0, // totalCost ignored because custom rule hits
|
||||
tokens, 1, 999.0, "priority", // 自定义账号价格不叠加服务层级倍率
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.01 + 50*0.02 = 1.0 + 1.0 = 2.0
|
||||
@@ -647,7 +716,7 @@ func TestResolveAccountStatsCost_ApplyPricingToAccountStats_UsesTotalCost(t *tes
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 0.75, // totalCost = 0.75
|
||||
tokens, 1, 0.75, "priority", // 已完成用户计费,不再重复应用服务层级倍率
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 0.75, *result, 1e-12)
|
||||
@@ -665,7 +734,7 @@ func TestResolveAccountStatsCost_ApplyPricingToAccountStats_ZeroTotalCost_Return
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
UsageTokens{}, 1, 0.0, // totalCost = 0
|
||||
UsageTokens{}, 1, 0.0, "", // totalCost = 0
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -692,7 +761,7 @@ func TestResolveAccountStatsCost_FallsBackToLiteLLM(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 999.0, // totalCost ignored
|
||||
tokens, 1, 999.0, "", // totalCost ignored
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
|
||||
@@ -712,7 +781,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "gemini-3.6-flash-low",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0,
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "",
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 9.15, *result, 1e-12)
|
||||
@@ -736,7 +805,7 @@ func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "totally-unknown-model",
|
||||
tokens, 1, 0.0,
|
||||
tokens, 1, 0.0, "",
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -753,7 +822,7 @@ func TestResolveAccountStatsCost_NilBillingService_SkipsLiteLLM(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, nil, // billingService is nil
|
||||
1, 10, "claude-sonnet-4",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.0,
|
||||
UsageTokens{InputTokens: 100}, 1, 0.0, "",
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -786,13 +855,41 @@ func TestResolveAccountStatsCost_CustomRulePriorityOverApplyPricing(t *testing.T
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 99.0, // totalCost = 99.0 (would be used if ApplyPricing wins)
|
||||
tokens, 1, 99.0, "", // totalCost = 99.0 (would be used if ApplyPricing wins)
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// Custom rule: 100*0.05 = 5.0 (NOT 99.0 from totalCost)
|
||||
require.InDelta(t, 5.0, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestApplyAccountStatsCost_UsesUsageLogServiceTier(t *testing.T) {
|
||||
channel := &Channel{
|
||||
ID: 1,
|
||||
Status: StatusActive,
|
||||
ApplyPricingToAccountStats: false,
|
||||
}
|
||||
cs := newTestChannelServiceForStats(t, channel, 10, "openai")
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
|
||||
"gpt-5.6-sol": {
|
||||
InputPricePerToken: 0.001,
|
||||
InputPricePerTokenPriority: 0.002,
|
||||
OutputPricePerToken: 0.002,
|
||||
OutputPricePerTokenPriority: 0.004,
|
||||
},
|
||||
})
|
||||
serviceTier := "priority"
|
||||
usageLog := &UsageLog{ServiceTier: &serviceTier}
|
||||
|
||||
applyAccountStatsCost(
|
||||
context.Background(), usageLog, cs, bs,
|
||||
1, 10, "gpt-5.6-sol", "gpt-5.6-sol",
|
||||
UsageTokens{InputTokens: 100, OutputTokens: 50}, 999,
|
||||
)
|
||||
|
||||
require.NotNil(t, usageLog.AccountStatsCost)
|
||||
require.InDelta(t, 0.4, *usageLog.AccountStatsCost, 1e-12)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// helpers for resolveAccountStatsCost tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -610,10 +610,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
|
||||
}
|
||||
|
||||
// Align test routing with gateway behavior: OpenAI accounts apply normal
|
||||
// account model mapping, and compact mode applies compact-only mapping on top.
|
||||
// account model mapping. Native remote compaction v2 rides the ordinary
|
||||
// /responses wire and does NOT apply the legacy compact-only mapping
|
||||
// (post-#5641 semantics: compact_model_mapping is /responses/compact-only).
|
||||
testModelID = account.GetMappedModel(testModelID)
|
||||
if mode == AccountTestModeCompact {
|
||||
testModelID = resolveOpenAICompactForwardModel(account, testModelID)
|
||||
return s.testOpenAICompactConnection(c, account, testModelID)
|
||||
}
|
||||
|
||||
@@ -969,6 +970,7 @@ func (s *AccountTestService) observeGrokTestResponse(ctx context.Context, accoun
|
||||
resp.Body = io.NopCloser(bytes.NewReader(responseBody))
|
||||
}
|
||||
snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now)
|
||||
stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
|
||||
if snapshot != nil && s.accountRepo != nil {
|
||||
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
if limited {
|
||||
@@ -1073,7 +1075,7 @@ func (s *AccountTestService) testGrokResponsesConnection(c *gin.Context, ctx con
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
s.observeGrokTestResponse(ctx, account, resp)
|
||||
s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, testModelID), account, resp)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
@@ -1421,7 +1423,7 @@ User query:
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("standalone web_search probe failed: %s", err.Error()))
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
s.observeGrokTestResponse(ctx, account, resp)
|
||||
s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, grokDefaultResponsesModel), account, resp)
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -1975,8 +1977,10 @@ func (s *AccountTestService) testOpenAIChatCompletionsConnection(
|
||||
return s.processOpenAIChatCompletionsStream(c, resp.Body)
|
||||
}
|
||||
|
||||
// testOpenAICompactConnection probes /responses/compact and persists the
|
||||
// resulting capability state on the account.
|
||||
// testOpenAICompactConnection probes native remote compaction v2 (streaming
|
||||
// /responses with a compaction_trigger input item) and persists the resulting
|
||||
// capability state on the account. The legacy unary /responses/compact
|
||||
// endpoint has been sunset upstream (404, #5598/#5624) and is no longer probed.
|
||||
func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account *Account, testModelID string) error {
|
||||
ctx := c.Request.Context()
|
||||
credentialAccount := account
|
||||
@@ -2001,7 +2005,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
return s.sendErrorAndEnd(c, "No access token available")
|
||||
}
|
||||
apiURL = chatgptCodexAPIURL + "/compact"
|
||||
apiURL = chatgptCodexAPIURL
|
||||
case account.Type == AccountTypeAPIKey:
|
||||
authToken = account.GetOpenAIApiKey()
|
||||
if authToken == "" {
|
||||
@@ -2015,7 +2019,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
|
||||
}
|
||||
apiURL = appendOpenAIResponsesRequestPathSuffix(buildOpenAIResponsesURL(normalizedBaseURL), "/compact")
|
||||
apiURL = buildOpenAIResponsesURL(normalizedBaseURL)
|
||||
default:
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported account type: %s", account.Type))
|
||||
}
|
||||
@@ -2026,7 +2030,11 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
||||
c.Writer.Flush()
|
||||
|
||||
payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID))
|
||||
// 原生 v2 走普通 /responses 线:OAuth 与真实转发一致做上游模型归一化。
|
||||
if isOAuth {
|
||||
testModelID = normalizeOpenAIModelForUpstream(credentialAccount, testModelID)
|
||||
}
|
||||
payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID, isOAuth))
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) {
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
}
|
||||
@@ -2038,7 +2046,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
// v2 探测是流式请求;同时补注协商头,与真实 codex 出站线型一致。
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
ensureOpenAIRemoteCompactionV2BetaFeature(req.Header)
|
||||
if credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
|
||||
if authErr != nil {
|
||||
@@ -2060,6 +2070,12 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
if isOAuth {
|
||||
req.Host = "chatgpt.com"
|
||||
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
|
||||
// 指纹收敛:探测与真实转发走同一个 /responses 端点,身份也必须同构,
|
||||
// 否则探测流量会以「缺 x-codex-installation-id + 非收敛 session」的
|
||||
// 形态暴露在上游眼里。账号关闭收敛(off)时返回 nil,探测保持原样。
|
||||
if fpIDs := resolveCodexFingerprintIDsFromRequest(account, req.Header); fpIDs != nil {
|
||||
applyCodexFingerprintHeaders(req.Header, fpIDs)
|
||||
}
|
||||
}
|
||||
|
||||
// 账号级请求头覆写:测试请求与真实转发保持一致的最终头
|
||||
@@ -2073,7 +2089,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
resp, err := s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
|
||||
if err != nil {
|
||||
if s.accountRepo != nil {
|
||||
updates := buildOpenAICompactProbeExtraUpdates(nil, nil, err, time.Now())
|
||||
updates := buildOpenAICompactProbeExtraUpdates(nil, nil, err, false, time.Now())
|
||||
_ = s.accountRepo.UpdateExtra(ctx, account.ID, updates)
|
||||
mergeAccountExtra(account, updates)
|
||||
}
|
||||
@@ -2092,8 +2108,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
return s.testOpenAICompactConnection(c, account, testModelID)
|
||||
}
|
||||
|
||||
compactionFound := openAICompactProbeFoundCompactionItem(body)
|
||||
if s.accountRepo != nil {
|
||||
updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now())
|
||||
updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, compactionFound, time.Now())
|
||||
if codexUpdates, err := extractOpenAICodexProbeUpdates(resp); err == nil && len(codexUpdates) > 0 {
|
||||
updates = mergeExtraUpdates(updates, codexUpdates)
|
||||
}
|
||||
@@ -2115,7 +2132,11 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
|
||||
}
|
||||
|
||||
s.sendEvent(c, TestEvent{Type: "content", Text: "Compact probe succeeded"})
|
||||
if !compactionFound {
|
||||
return s.sendErrorAndEnd(c, "Upstream returned 2xx without a compaction output item (native remote compaction v2 unsupported on this chain)")
|
||||
}
|
||||
|
||||
s.sendEvent(c, TestEvent{Type: "content", Text: "Compact probe succeeded (native remote compaction v2)"})
|
||||
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -10,10 +10,16 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// compactProbeSSESuccessBody 是原生 v2 压缩成功的最小 SSE 形态:
|
||||
// output_item.done 携带 compaction item + response.completed。
|
||||
const compactProbeSSESuccessBody = "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"id\":\"cmp_probe\",\"encrypted_content\":\"blob\"}}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_probe\",\"output\":[]}}\n\n"
|
||||
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersistsSupport(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -38,8 +44,8 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersi
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-probe"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_probe","status":"completed"}`)),
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid-probe"}},
|
||||
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
accountRepo: repo,
|
||||
@@ -53,16 +59,22 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersi
|
||||
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, chatgptCodexAPIURL+"/compact", upstream.lastReq.URL.String())
|
||||
// 原生 v2:探测普通 /responses 线,不再打已下线的 /responses/compact。
|
||||
require.Equal(t, chatgptCodexAPIURL, upstream.lastReq.URL.String())
|
||||
require.Equal(t, "chatgpt.com", upstream.lastReq.Host)
|
||||
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Accept"))
|
||||
require.Equal(t, codexCLIVersion, upstream.lastReq.Header.Get("Version"))
|
||||
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
|
||||
require.Contains(t, upstream.lastReq.Header.Get("x-codex-beta-features"), "remote_compaction_v2")
|
||||
require.NotEmpty(t, upstream.lastReq.Header.Get("Session_Id"))
|
||||
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
|
||||
require.Equal(t, codexCLIUserAgent, upstream.lastReq.Header.Get("User-Agent"))
|
||||
require.Equal(t, "chatgpt-acc", upstream.lastReq.Header.Get("chatgpt-account-id"))
|
||||
require.Equal(t, "true", upstream.lastReq.Header.Get("x-openai-fedramp"))
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Bool())
|
||||
inputItems := gjson.GetBytes(upstream.lastBody, "input").Array()
|
||||
require.NotEmpty(t, inputItems)
|
||||
require.Equal(t, "compaction_trigger", inputItems[len(inputItems)-1].Get("type").String())
|
||||
|
||||
updates := <-updateCalls
|
||||
require.Equal(t, true, updates["openai_compact_supported"])
|
||||
@@ -114,7 +126,7 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactOAuth404MarksUnsu
|
||||
require.Contains(t, rec.Body.String(), `"type":"error"`)
|
||||
}
|
||||
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompactPath(t *testing.T) {
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesNativeResponsesPath(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
updateCalls := make(chan map[string]any, 1)
|
||||
@@ -127,8 +139,10 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompact
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://example.com/v1",
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://example.com/v1",
|
||||
// post-#5641:compact_model_mapping 仅作用于 legacy /responses/compact,
|
||||
// 原生 v2 探测不应用它。
|
||||
"compact_model_mapping": map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"},
|
||||
},
|
||||
}
|
||||
@@ -138,8 +152,8 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompact
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_probe_apikey","status":"completed"}`)),
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
accountRepo: repo,
|
||||
@@ -154,14 +168,16 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompact
|
||||
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, "https://example.com/v1/responses/compact", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "https://example.com/v1/responses", upstream.lastReq.URL.String())
|
||||
requireOpenAICodexProbeHeaders(t, upstream.lastReq.Header)
|
||||
require.Equal(t, "gpt-5.4-openai-compact", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.Contains(t, upstream.lastReq.Header.Get("x-codex-beta-features"), "remote_compaction_v2")
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String(),
|
||||
"原生 v2 探测不应用 compact_model_mapping")
|
||||
updates := <-updateCalls
|
||||
require.Equal(t, true, updates["openai_compact_supported"])
|
||||
}
|
||||
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBaseURLUsesV1Path(t *testing.T) {
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBaseURLUsesResponsesPath(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
updateCalls := make(chan map[string]any, 1)
|
||||
@@ -183,8 +199,8 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBase
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_probe_apikey_default","status":"completed"}`)),
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
accountRepo: repo,
|
||||
@@ -198,6 +214,112 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBase
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://api.openai.com/v1/responses/compact", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "https://api.openai.com/v1/responses", upstream.lastReq.URL.String())
|
||||
<-updateCalls
|
||||
}
|
||||
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompact2xxWithoutItemMarksUnsupported(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
updateCalls := make(chan map[string]any, 1)
|
||||
account := Account{
|
||||
ID: 5,
|
||||
Name: "openai-oauth-no-item",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "oauth-token",
|
||||
"chatgpt_account_id": "chatgpt-acc",
|
||||
},
|
||||
}
|
||||
repo := &snapshotUpdateAccountRepo{
|
||||
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
|
||||
updateExtraCalls: updateCalls,
|
||||
}
|
||||
// 200 但流里没有 compaction item:链路吞掉了 compaction_trigger 的形态
|
||||
//(#5478 的 "got 0 items"),必须判定为不支持。
|
||||
noItemBody := "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"message\",\"id\":\"msg_1\"}}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_x\",\"output\":[]}}\n\n"
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(noItemBody)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
accountRepo: repo,
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/5/test", bytes.NewReader(nil))
|
||||
|
||||
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
|
||||
require.Error(t, err)
|
||||
|
||||
updates := <-updateCalls
|
||||
require.Equal(t, false, updates["openai_compact_supported"])
|
||||
require.Contains(t, rec.Body.String(), `"type":"error"`)
|
||||
}
|
||||
|
||||
// 探测与真实转发走同一 /responses 端点,出站身份必须与真实 Codex 同构:
|
||||
// session/thread 为 UUID、携带 x-codex-installation-id(收敛账号用收敛值)。
|
||||
func TestAccountTestService_TestAccountConnection_OpenAICompactProbeIdentityMatchesRealTraffic(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
updateCalls := make(chan map[string]any, 1)
|
||||
account := Account{
|
||||
ID: 6,
|
||||
Name: "openai-oauth-identity",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "oauth-token",
|
||||
"chatgpt_account_id": "chatgpt-acc",
|
||||
},
|
||||
// 收敛是显式 opt-in(#5610),这里显式开启以验证探测身份与真实流量同构。
|
||||
Extra: map[string]any{"codex_fingerprint_mode": "session"},
|
||||
}
|
||||
repo := &snapshotUpdateAccountRepo{
|
||||
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
|
||||
updateExtraCalls: updateCalls,
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
|
||||
}}
|
||||
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/6/test", bytes.NewReader(nil))
|
||||
|
||||
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
|
||||
|
||||
// 显式 session 收敛模式:出站身份 = 账号级收敛值
|
||||
converged := resolveConvergedSessionID(&account)
|
||||
require.Equal(t, converged, upstream.lastReq.Header.Get("session-id"))
|
||||
require.Equal(t, converged, upstream.lastReq.Header.Get("session_id"))
|
||||
require.Equal(t, resolveConvergedInstallationID(&account), upstream.lastReq.Header.Get("x-codex-installation-id"),
|
||||
"真实 Codex 每个请求必带 installation-id,探测不得缺失")
|
||||
require.NotContains(t, upstream.lastReq.Header.Get("session-id"), "probe_compact",
|
||||
"探测标识不得是可被上游一眼识别的字面量")
|
||||
<-updateCalls
|
||||
}
|
||||
|
||||
func TestCompactProbeSessionID_IsUUIDShaped(t *testing.T) {
|
||||
for _, id := range []int64{0, 1, 987654} {
|
||||
got := compactProbeSessionID(id)
|
||||
_, err := uuid.Parse(got)
|
||||
require.NoError(t, err, "探测会话标识必须是 UUID 形态: %s", got)
|
||||
}
|
||||
require.Equal(t, compactProbeSessionID(7), compactProbeSessionID(7), "同账号应稳定复用同一会话")
|
||||
require.NotEqual(t, compactProbeSessionID(7), compactProbeSessionID(8))
|
||||
}
|
||||
|
||||
@@ -301,6 +301,10 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
|
||||
}
|
||||
|
||||
platform := NormalizeGroupPlatform(input.Platform)
|
||||
modelPricing, err := normalizeGroupModelPricing(platform, input.ModelPricing)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
maxReasoningEffort, err := normalizeMaxReasoningEffortForPlatform(platform, input.MaxReasoningEffort)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "INVALID_MAX_REASONING_EFFORT", "%v", err)
|
||||
@@ -459,6 +463,8 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
|
||||
DailyLimitUSD: dailyLimit,
|
||||
WeeklyLimitUSD: weeklyLimit,
|
||||
MonthlyLimitUSD: monthlyLimit,
|
||||
LongContextPricingEnabled: input.LongContextPricingEnabled,
|
||||
ModelPricing: modelPricing,
|
||||
AllowImageGeneration: allowImageGeneration,
|
||||
AllowBatchImageGeneration: allowBatchImageGeneration,
|
||||
ImageRateIndependent: input.ImageRateIndependent,
|
||||
@@ -635,6 +641,9 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 渠道缓存里存了 groupID → platform 的映射,改了平台要让它失效(见函数末尾)
|
||||
previousPlatform := group.Platform
|
||||
|
||||
if input.Name != "" {
|
||||
group.Name = input.Name
|
||||
}
|
||||
@@ -656,6 +665,16 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
if input.Status != "" {
|
||||
group.Status = input.Status
|
||||
}
|
||||
if input.LongContextPricingEnabled != nil {
|
||||
group.LongContextPricingEnabled = *input.LongContextPricingEnabled
|
||||
}
|
||||
if input.ModelPricing != nil {
|
||||
modelPricing, normalizeErr := normalizeGroupModelPricing(group.Platform, *input.ModelPricing)
|
||||
if normalizeErr != nil {
|
||||
return nil, normalizeErr
|
||||
}
|
||||
group.ModelPricing = modelPricing
|
||||
}
|
||||
|
||||
// 订阅相关字段
|
||||
if input.SubscriptionType != "" {
|
||||
@@ -884,6 +903,13 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, id)
|
||||
}
|
||||
|
||||
// 平台变了就失效渠道缓存:该缓存持有 groupID → platform,而渠道定价 / 模型映射 /
|
||||
// 模型白名单都按平台严格隔离。不失效的话,缓存最长 10 分钟仍按旧平台匹配,
|
||||
// 期间定价查不到会静默回落到 LiteLLM 价格表、映射与白名单也不生效。
|
||||
if group.Platform != previousPlatform && s.channelCacheInvalidator != nil {
|
||||
s.channelCacheInvalidator.InvalidateCache()
|
||||
}
|
||||
|
||||
// 如果指定了复制账号的源分组,同步绑定(替换当前分组的账号)
|
||||
if len(input.CopyAccountsFromGroupIDs) > 0 {
|
||||
// 去重源分组 IDs
|
||||
@@ -955,6 +981,28 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
return group, nil
|
||||
}
|
||||
|
||||
func normalizeGroupModelPricing(platform string, pricing []ChannelModelPricing) ([]ChannelModelPricing, error) {
|
||||
out := make([]ChannelModelPricing, len(pricing))
|
||||
for i := range pricing {
|
||||
out[i] = pricing[i].Clone()
|
||||
out[i].ID = 0
|
||||
out[i].ChannelID = 0
|
||||
if strings.TrimSpace(out[i].Platform) == "" {
|
||||
out[i].Platform = platform
|
||||
}
|
||||
for j := range out[i].Models {
|
||||
out[i].Models[j] = strings.TrimSpace(out[i].Models[j])
|
||||
}
|
||||
if len(out[i].Models) == 0 {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "GROUP_MODEL_PRICING_MODELS_REQUIRED", "group model pricing entry requires at least one model")
|
||||
}
|
||||
}
|
||||
if err := validatePricingEntries(out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) DeleteGroup(ctx context.Context, id int64) error {
|
||||
var groupKeys []string
|
||||
if s.authCacheInvalidator != nil {
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// groupPlatformRepoStub 只实现 UpdateGroup 走到的两个方法,其余靠内嵌接口占位。
|
||||
type groupPlatformRepoStub struct {
|
||||
GroupRepository
|
||||
group *Group
|
||||
updated *Group
|
||||
}
|
||||
|
||||
func (r *groupPlatformRepoStub) GetByID(_ context.Context, _ int64) (*Group, error) {
|
||||
cloned := *r.group
|
||||
return &cloned, nil
|
||||
}
|
||||
|
||||
func (r *groupPlatformRepoStub) Update(_ context.Context, group *Group) error {
|
||||
r.updated = group
|
||||
return nil
|
||||
}
|
||||
|
||||
type channelCacheInvalidatorSpy struct {
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *channelCacheInvalidatorSpy) InvalidateCache() { s.calls++ }
|
||||
|
||||
// 渠道缓存持有 groupID → platform,而渠道定价/模型映射/模型白名单都按平台严格隔离。
|
||||
// 改了分组平台却不失效缓存,最长 10 分钟内这些查找仍按旧平台匹配(静默走错价)。
|
||||
func TestUpdateGroupInvalidatesChannelCacheOnPlatformChange(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
fromPlatform string
|
||||
inputPlatform string
|
||||
wantCalls int
|
||||
}{
|
||||
{
|
||||
name: "platform changed invalidates",
|
||||
fromPlatform: PlatformAnthropic,
|
||||
inputPlatform: PlatformOpenAI,
|
||||
wantCalls: 1,
|
||||
},
|
||||
{
|
||||
name: "same platform does not invalidate",
|
||||
fromPlatform: PlatformAnthropic,
|
||||
inputPlatform: PlatformAnthropic,
|
||||
wantCalls: 0,
|
||||
},
|
||||
{
|
||||
// 请求里不带 platform 字段时不应该动缓存
|
||||
name: "platform omitted does not invalidate",
|
||||
fromPlatform: PlatformAnthropic,
|
||||
inputPlatform: "",
|
||||
wantCalls: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &groupPlatformRepoStub{group: &Group{ID: 7, Name: "g", Platform: tt.fromPlatform}}
|
||||
spy := &channelCacheInvalidatorSpy{}
|
||||
svc := &adminServiceImpl{groupRepo: repo, channelCacheInvalidator: spy}
|
||||
|
||||
got, err := svc.UpdateGroup(context.Background(), 7, &UpdateGroupInput{Platform: tt.inputPlatform})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got)
|
||||
require.Equal(t, tt.wantCalls, spy.calls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 依赖可以不注入(例如测试或裁剪构建),此时不应 panic——缓存靠 TTL 自然重建。
|
||||
func TestUpdateGroupWithoutChannelCacheInvalidator(t *testing.T) {
|
||||
repo := &groupPlatformRepoStub{group: &Group{ID: 7, Name: "g", Platform: PlatformAnthropic}}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
got, err := svc.UpdateGroup(context.Background(), 7, &UpdateGroupInput{Platform: PlatformOpenAI})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, PlatformOpenAI, got.Platform)
|
||||
}
|
||||
@@ -209,15 +209,17 @@ type AdminBoundAuthIdentityChannel struct {
|
||||
}
|
||||
|
||||
type CreateGroupInput struct {
|
||||
Name string
|
||||
Description string
|
||||
Platform string
|
||||
RateMultiplier float64
|
||||
IsExclusive bool
|
||||
SubscriptionType string // standard/subscription
|
||||
DailyLimitUSD *float64 // 日限额 (USD)
|
||||
WeeklyLimitUSD *float64 // 周限额 (USD)
|
||||
MonthlyLimitUSD *float64 // 月限额 (USD)
|
||||
Name string
|
||||
Description string
|
||||
Platform string
|
||||
RateMultiplier float64
|
||||
IsExclusive bool
|
||||
SubscriptionType string // standard/subscription
|
||||
DailyLimitUSD *float64 // 日限额 (USD)
|
||||
WeeklyLimitUSD *float64 // 周限额 (USD)
|
||||
MonthlyLimitUSD *float64 // 月限额 (USD)
|
||||
LongContextPricingEnabled bool
|
||||
ModelPricing []ChannelModelPricing
|
||||
// 图片生成计费配置(仅 antigravity 平台使用)
|
||||
AllowImageGeneration bool
|
||||
AllowBatchImageGeneration bool
|
||||
@@ -281,16 +283,18 @@ type CreateGroupInput struct {
|
||||
}
|
||||
|
||||
type UpdateGroupInput struct {
|
||||
Name string
|
||||
Description *string
|
||||
Platform string
|
||||
RateMultiplier *float64 // 使用指针以支持设置为0
|
||||
IsExclusive *bool
|
||||
Status string
|
||||
SubscriptionType string // standard/subscription
|
||||
DailyLimitUSD *float64 // 日限额 (USD)
|
||||
WeeklyLimitUSD *float64 // 周限额 (USD)
|
||||
MonthlyLimitUSD *float64 // 月限额 (USD)
|
||||
Name string
|
||||
Description *string
|
||||
Platform string
|
||||
RateMultiplier *float64 // 使用指针以支持设置为0
|
||||
IsExclusive *bool
|
||||
Status string
|
||||
SubscriptionType string // standard/subscription
|
||||
DailyLimitUSD *float64 // 日限额 (USD)
|
||||
WeeklyLimitUSD *float64 // 周限额 (USD)
|
||||
MonthlyLimitUSD *float64 // 月限额 (USD)
|
||||
LongContextPricingEnabled *bool
|
||||
ModelPricing *[]ChannelModelPricing
|
||||
// 图片生成计费配置(仅 antigravity 平台使用)
|
||||
AllowImageGeneration *bool
|
||||
AllowBatchImageGeneration *bool
|
||||
@@ -662,6 +666,14 @@ type adminServiceImpl struct {
|
||||
affiliateService adminRechargeAffiliateAccruer
|
||||
compositeRouteRepo CompositeModelRouteRepository
|
||||
compositeResolver *CompositeRouteResolver
|
||||
// 分组平台变更后用来失效渠道缓存;可为 nil(缓存会在 TTL 到期后自然重建)
|
||||
channelCacheInvalidator ChannelCacheInvalidator
|
||||
}
|
||||
|
||||
// ChannelCacheInvalidator 失效渠道缓存。
|
||||
// 窄接口,避免 admin 服务依赖整个 ChannelService——与 APIKeyAuthCacheInvalidator 同一思路。
|
||||
type ChannelCacheInvalidator interface {
|
||||
InvalidateCache()
|
||||
}
|
||||
|
||||
type adminRechargeAffiliateAccruer interface {
|
||||
@@ -695,6 +707,7 @@ func NewAdminService(
|
||||
affiliateService *AffiliateService,
|
||||
compositeRouteRepo CompositeModelRouteRepository,
|
||||
compositeResolver *CompositeRouteResolver,
|
||||
channelCacheInvalidator ChannelCacheInvalidator,
|
||||
) AdminService {
|
||||
return &adminServiceImpl{
|
||||
userRepo: userRepo,
|
||||
@@ -721,5 +734,7 @@ func NewAdminService(
|
||||
affiliateService: affiliateService,
|
||||
compositeRouteRepo: compositeRouteRepo,
|
||||
compositeResolver: compositeResolver,
|
||||
|
||||
channelCacheInvalidator: channelCacheInvalidator,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"html"
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -245,6 +246,36 @@ type UpdateAPIKeyRequest struct {
|
||||
ResetRateLimitUsage *bool `json:"reset_rate_limit_usage"` // Reset all usage counters to 0
|
||||
}
|
||||
|
||||
func validateAPIKeyLimit(v float64) error {
|
||||
if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
|
||||
return infraerrors.BadRequest("API_KEY_LIMIT_INVALID", "API key limits must be finite and non-negative")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateCreateAPIKeyRequest(req CreateAPIKeyRequest) error {
|
||||
for _, v := range []float64{req.Quota, req.RateLimit5h, req.RateLimit1d, req.RateLimit7d} {
|
||||
if err := validateAPIKeyLimit(v); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if req.ExpiresInDays != nil && *req.ExpiresInDays <= 0 {
|
||||
return infraerrors.BadRequest("API_KEY_EXPIRY_INVALID", "expires_in_days must be greater than zero")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateUpdateAPIKeyRequest(req UpdateAPIKeyRequest) error {
|
||||
for _, v := range []*float64{req.Quota, req.RateLimit5h, req.RateLimit1d, req.RateLimit7d} {
|
||||
if v != nil {
|
||||
if err := validateAPIKeyLimit(*v); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// APIKeyService API Key服务
|
||||
// RateLimitCacheInvalidator invalidates rate limit cache entries on manual reset.
|
||||
type RateLimitCacheInvalidator interface {
|
||||
@@ -428,6 +459,9 @@ func (s *APIKeyService) canUserBindGroup(ctx context.Context, user *User, group
|
||||
|
||||
// Create 创建API Key
|
||||
func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIKeyRequest) (*APIKey, error) {
|
||||
if err := validateCreateAPIKeyRequest(req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 验证用户存在
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
@@ -723,6 +757,9 @@ func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, erro
|
||||
|
||||
// Update 更新API Key
|
||||
func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req UpdateAPIKeyRequest) (*APIKey, error) {
|
||||
if err := validateUpdateAPIKeyRequest(req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get api key: %w", err)
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateCreateAPIKeyRequestNumericLimits(t *testing.T) {
|
||||
positiveExpiry := 1
|
||||
require.NoError(t, validateCreateAPIKeyRequest(CreateAPIKeyRequest{
|
||||
Quota: 1e100, RateLimit5h: 1e100, ExpiresInDays: &positiveExpiry,
|
||||
}))
|
||||
require.NoError(t, validateCreateAPIKeyRequest(CreateAPIKeyRequest{}))
|
||||
|
||||
invalidExpiry := 0
|
||||
tests := []CreateAPIKeyRequest{
|
||||
{Quota: -1},
|
||||
{Quota: math.NaN()},
|
||||
{Quota: math.Inf(1)},
|
||||
{RateLimit5h: -1},
|
||||
{RateLimit1d: math.NaN()},
|
||||
{RateLimit7d: math.Inf(-1)},
|
||||
{ExpiresInDays: &invalidExpiry},
|
||||
}
|
||||
for _, req := range tests {
|
||||
require.Error(t, validateCreateAPIKeyRequest(req))
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateUpdateAPIKeyRequestNumericLimits(t *testing.T) {
|
||||
zero, large, negative, nan, inf := 0.0, 1e100, -1.0, math.NaN(), math.Inf(1)
|
||||
require.NoError(t, validateUpdateAPIKeyRequest(UpdateAPIKeyRequest{Quota: &zero, RateLimit7d: &large}))
|
||||
|
||||
for _, req := range []UpdateAPIKeyRequest{
|
||||
{Quota: &negative},
|
||||
{RateLimit5h: &nan},
|
||||
{RateLimit1d: &inf},
|
||||
{RateLimit7d: &negative},
|
||||
} {
|
||||
require.Error(t, validateUpdateAPIKeyRequest(req))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
)
|
||||
|
||||
const defaultBackupPartSizeBytes int64 = 4 * 1024 * 1024 * 1024
|
||||
|
||||
// BackupPart 描述一个 gzip 字节分卷。
|
||||
type BackupPart struct {
|
||||
Index int `json:"index"`
|
||||
S3Key string `json:"s3_key"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
SHA256 string `json:"sha256,omitempty"`
|
||||
}
|
||||
|
||||
type localBackupPart struct {
|
||||
Index int
|
||||
Path string
|
||||
SizeBytes int64
|
||||
SHA256 string
|
||||
}
|
||||
|
||||
func splitBackupFile(srcPath string, partSize int64) (parts []localBackupPart, err error) {
|
||||
if partSize <= 0 {
|
||||
return nil, fmt.Errorf("backup part size must be positive")
|
||||
}
|
||||
|
||||
src, err := os.Open(srcPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open backup archive: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := src.Close(); err == nil && closeErr != nil {
|
||||
err = fmt.Errorf("close backup archive: %w", closeErr)
|
||||
}
|
||||
if err != nil {
|
||||
paths := make([]string, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
paths = append(paths, part.Path)
|
||||
}
|
||||
_ = cleanupBackupFiles(paths...)
|
||||
}
|
||||
}()
|
||||
|
||||
info, err := src.Stat()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stat backup archive: %w", err)
|
||||
}
|
||||
if info.Size() <= 0 {
|
||||
return nil, errors.New("backup archive is empty")
|
||||
}
|
||||
|
||||
remaining := info.Size()
|
||||
for index := 1; remaining > 0; index++ {
|
||||
partFile, createErr := os.CreateTemp("", "sub2api-backup-part-*")
|
||||
if createErr != nil {
|
||||
return nil, fmt.Errorf("create backup part: %w", createErr)
|
||||
}
|
||||
partPath := partFile.Name()
|
||||
partBytes := partSize
|
||||
if remaining < partBytes {
|
||||
partBytes = remaining
|
||||
}
|
||||
|
||||
hash := sha256.New()
|
||||
written, copyErr := io.CopyN(io.MultiWriter(partFile, hash), src, partBytes)
|
||||
closeErr := partFile.Close()
|
||||
if copyErr != nil {
|
||||
_ = os.Remove(partPath)
|
||||
return nil, fmt.Errorf("write backup part %d: %w", index, copyErr)
|
||||
}
|
||||
if closeErr != nil {
|
||||
_ = os.Remove(partPath)
|
||||
return nil, fmt.Errorf("close backup part %d: %w", index, closeErr)
|
||||
}
|
||||
|
||||
parts = append(parts, localBackupPart{
|
||||
Index: index,
|
||||
Path: partPath,
|
||||
SizeBytes: written,
|
||||
SHA256: hex.EncodeToString(hash.Sum(nil)),
|
||||
})
|
||||
remaining -= written
|
||||
}
|
||||
|
||||
return parts, nil
|
||||
}
|
||||
|
||||
func cleanupBackupFiles(paths ...string) error {
|
||||
var errs []error
|
||||
for _, path := range paths {
|
||||
if path == "" {
|
||||
continue
|
||||
}
|
||||
if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
errs = append(errs, fmt.Errorf("remove %s: %w", path, err))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSplitBackupFile_ReassemblesExactBytes(t *testing.T) {
|
||||
src := writeBackupArchiveFixture(t, []byte("0123456789abcdefg"))
|
||||
parts, err := splitBackupFile(src, 5)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, parts, 4)
|
||||
|
||||
var got bytes.Buffer
|
||||
for i, part := range parts {
|
||||
require.Equal(t, i+1, part.Index)
|
||||
require.LessOrEqual(t, part.SizeBytes, int64(5))
|
||||
data, readErr := os.ReadFile(part.Path)
|
||||
require.NoError(t, readErr)
|
||||
require.Equal(t, fmt.Sprintf("%x", sha256.Sum256(data)), part.SHA256)
|
||||
got.Write(data)
|
||||
}
|
||||
require.Equal(t, []byte("0123456789abcdefg"), got.Bytes())
|
||||
}
|
||||
|
||||
func TestSplitBackupFile_RejectsInvalidInput(t *testing.T) {
|
||||
src := writeBackupArchiveFixture(t, []byte("data"))
|
||||
|
||||
_, err := splitBackupFile(src, 0)
|
||||
require.Error(t, err)
|
||||
|
||||
empty := writeBackupArchiveFixture(t, nil)
|
||||
_, err = splitBackupFile(empty, 5)
|
||||
require.Error(t, err)
|
||||
|
||||
_, err = splitBackupFile(filepathForMissingBackupArchive(t), 5)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func writeBackupArchiveFixture(t *testing.T, content []byte) string {
|
||||
t.Helper()
|
||||
path := filepathForBackupArchive(t)
|
||||
require.NoError(t, os.WriteFile(path, content, 0o600))
|
||||
return path
|
||||
}
|
||||
|
||||
func filepathForBackupArchive(t *testing.T) string {
|
||||
t.Helper()
|
||||
return t.TempDir() + "/archive.gz"
|
||||
}
|
||||
|
||||
func filepathForMissingBackupArchive(t *testing.T) string {
|
||||
t.Helper()
|
||||
return t.TempDir() + "/missing.gz"
|
||||
}
|
||||
@@ -3,10 +3,14 @@ package service
|
||||
import (
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -26,7 +30,23 @@ const (
|
||||
settingKeyBackupSchedule = "backup_schedule"
|
||||
settingKeyBackupRecords = "backup_records"
|
||||
|
||||
maxBackupRecords = 100
|
||||
maxBackupRecords = 100
|
||||
backupObjectCleanupTimeout = 2 * time.Minute
|
||||
|
||||
// backupScheduledLeaderLockKey gates the scheduled full-database backup so
|
||||
// that only one instance in a clustered deployment performs the
|
||||
// dump-and-upload each cycle. Without it every instance runs the cron
|
||||
// independently, producing N concurrent pg_dumps against the same database,
|
||||
// N× peak memory while the archive is uploaded, and N identical objects that
|
||||
// overwrite the same timestamped key. Every other periodic job in this
|
||||
// package is already gated the same way; the scheduled backup was the last
|
||||
// one that still fanned out across every instance.
|
||||
backupScheduledLeaderLockKey = "backup:scheduled:leader"
|
||||
// backupScheduledLeaderLockTTL bounds crash recovery only; the lock is
|
||||
// released as soon as the backup finishes. It must exceed the job's
|
||||
// worst-case runtime (the scheduled backup context is bounded at 30m) so the
|
||||
// lock cannot expire mid-dump and let a peer start a second backup.
|
||||
backupScheduledLeaderLockTTL = 35 * time.Minute
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -61,6 +81,7 @@ type DBDumper interface {
|
||||
// BackupObjectStore abstracts object storage for backup files
|
||||
type BackupObjectStore interface {
|
||||
Upload(ctx context.Context, key string, body io.Reader, contentType string) (sizeBytes int64, err error)
|
||||
UploadFile(ctx context.Context, key string, filePath string, contentType string) (sizeBytes int64, err error)
|
||||
Download(ctx context.Context, key string) (io.ReadCloser, error)
|
||||
Delete(ctx context.Context, key string) error
|
||||
PresignURL(ctx context.Context, key string, expiry time.Duration) (string, error)
|
||||
@@ -98,21 +119,35 @@ type BackupScheduleConfig struct {
|
||||
|
||||
// BackupRecord 备份记录
|
||||
type BackupRecord struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"` // pending, running, completed, failed
|
||||
BackupType string `json:"backup_type"` // postgres
|
||||
FileName string `json:"file_name"`
|
||||
S3Key string `json:"s3_key"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
TriggeredBy string `json:"triggered_by"` // manual, scheduled
|
||||
ErrorMsg string `json:"error_message,omitempty"`
|
||||
StartedAt string `json:"started_at"`
|
||||
FinishedAt string `json:"finished_at,omitempty"`
|
||||
ExpiresAt string `json:"expires_at,omitempty"` // 过期时间
|
||||
Progress string `json:"progress,omitempty"` // "dumping", "uploading", ""
|
||||
RestoreStatus string `json:"restore_status,omitempty"` // "", "running", "completed", "failed"
|
||||
RestoreError string `json:"restore_error,omitempty"`
|
||||
RestoredAt string `json:"restored_at,omitempty"`
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"` // pending, running, completed, failed
|
||||
BackupType string `json:"backup_type"` // postgres
|
||||
FileName string `json:"file_name"`
|
||||
S3Key string `json:"s3_key"`
|
||||
Parts []BackupPart `json:"parts,omitempty"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
TriggeredBy string `json:"triggered_by"` // manual, scheduled
|
||||
ErrorMsg string `json:"error_message,omitempty"`
|
||||
StartedAt string `json:"started_at"`
|
||||
FinishedAt string `json:"finished_at,omitempty"`
|
||||
ExpiresAt string `json:"expires_at,omitempty"` // 过期时间
|
||||
Progress string `json:"progress,omitempty"` // "dumping", "uploading", ""
|
||||
RestoreStatus string `json:"restore_status,omitempty"` // "", "running", "completed", "failed"
|
||||
RestoreError string `json:"restore_error,omitempty"`
|
||||
RestoredAt string `json:"restored_at,omitempty"`
|
||||
}
|
||||
|
||||
// BackupDownloadPart 描述一个可下载的备份分卷。
|
||||
type BackupDownloadPart struct {
|
||||
Index int `json:"index"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
// BackupDownloadResponse 是单文件和分卷下载响应的兼容表示。
|
||||
type BackupDownloadResponse struct {
|
||||
URL string `json:"url,omitempty"`
|
||||
Parts []BackupDownloadPart `json:"parts,omitempty"`
|
||||
}
|
||||
|
||||
// BackupService 数据库备份恢复服务
|
||||
@@ -142,10 +177,19 @@ type BackupService struct {
|
||||
cronSched *cron.Cron
|
||||
cronEntryID cron.EntryID
|
||||
|
||||
wg sync.WaitGroup // 追踪活跃的备份/恢复 goroutine
|
||||
shuttingDown atomic.Bool // 阻止新备份启动
|
||||
bgCtx context.Context // 所有后台操作的 parent context
|
||||
bgCancel context.CancelFunc // 取消所有活跃后台操作
|
||||
// lockCache/db elect a single leader for the scheduled backup across
|
||||
// instances; instanceID identifies this process as the lock owner. Injected
|
||||
// via SetLeaderLock — when both are nil the backup runs ungated
|
||||
// (single-instance / test behavior).
|
||||
lockCache LeaderLockCache
|
||||
db *sql.DB
|
||||
instanceID string
|
||||
|
||||
wg sync.WaitGroup // 追踪活跃的备份/恢复 goroutine
|
||||
shuttingDown atomic.Bool // 阻止新备份启动
|
||||
bgCtx context.Context // 所有后台操作的 parent context
|
||||
bgCancel context.CancelFunc // 取消所有活跃后台操作
|
||||
partSizeBytes int64 // 分卷阈值;生产使用 4 GiB,测试可注入更小值
|
||||
}
|
||||
|
||||
func NewBackupService(
|
||||
@@ -165,9 +209,22 @@ func NewBackupService(
|
||||
dumper: dumper,
|
||||
bgCtx: bgCtx,
|
||||
bgCancel: bgCancel,
|
||||
partSizeBytes: defaultBackupPartSizeBytes,
|
||||
instanceID: uuid.NewString(),
|
||||
}
|
||||
}
|
||||
|
||||
// SetLeaderLock injects the leader-lock cache and DB used to elect a single
|
||||
// instance for the scheduled backup. When both are nil the scheduled backup runs
|
||||
// ungated (single-instance / test behavior).
|
||||
func (s *BackupService) SetLeaderLock(lockCache LeaderLockCache, db *sql.DB) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.lockCache = lockCache
|
||||
s.db = db
|
||||
}
|
||||
|
||||
// Start 启动定时备份调度器并清理孤立记录
|
||||
func (s *BackupService) Start() {
|
||||
s.cronSched = cron.New()
|
||||
@@ -191,33 +248,57 @@ func (s *BackupService) Start() {
|
||||
}
|
||||
}
|
||||
|
||||
// recoverStaleRecords 启动时将孤立的 running 记录标记为 failed
|
||||
// recoverStaleRecords 启动时将孤立的 running 记录标记为 failed,并清理已上传对象。
|
||||
func (s *BackupService) recoverStaleRecords() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
loadCtx, loadCancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer loadCancel()
|
||||
|
||||
records, err := s.loadRecords(ctx)
|
||||
records, err := s.loadRecords(loadCtx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for i := range records {
|
||||
if records[i].Status == "running" {
|
||||
staleRecord := records[i]
|
||||
records[i].Status = "failed"
|
||||
records[i].ErrorMsg = "interrupted by server restart"
|
||||
records[i].Progress = ""
|
||||
records[i].FinishedAt = time.Now().Format(time.RFC3339)
|
||||
_ = s.saveRecord(ctx, &records[i])
|
||||
s.saveRecoveredRecord(&records[i])
|
||||
|
||||
if cleanupErr := s.cleanupStaleBackupObjects(&staleRecord); cleanupErr != nil {
|
||||
records[i].ErrorMsg = fmt.Sprintf("interrupted by server restart; cleanup failed, manual deletion may be required: %v", cleanupErr)
|
||||
s.saveRecoveredRecord(&records[i])
|
||||
logger.LegacyPrintf("service.backup", "[Backup] failed to clean stale backup objects for %s: %v", records[i].ID, cleanupErr)
|
||||
}
|
||||
logger.LegacyPrintf("service.backup", "[Backup] recovered stale running record: %s", records[i].ID)
|
||||
}
|
||||
if records[i].RestoreStatus == "running" {
|
||||
records[i].RestoreStatus = "failed"
|
||||
records[i].RestoreError = "interrupted by server restart"
|
||||
_ = s.saveRecord(ctx, &records[i])
|
||||
s.saveRecoveredRecord(&records[i])
|
||||
logger.LegacyPrintf("service.backup", "[Backup] recovered stale restoring record: %s", records[i].ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) saveRecoveredRecord(record *BackupRecord) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := s.saveRecord(ctx, record); err != nil {
|
||||
logger.LegacyPrintf("service.backup", "[Backup] 保存恢复后的备份记录失败 %s: %v", record.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) cleanupStaleBackupObjects(record *BackupRecord) error {
|
||||
if len(backupObjectKeys(record)) == 0 {
|
||||
return nil
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), backupObjectCleanupTimeout)
|
||||
defer cancel()
|
||||
return s.deleteBackupObjects(ctx, record)
|
||||
}
|
||||
|
||||
// Stop 停止定时备份并等待活跃操作完成
|
||||
func (s *BackupService) Stop() {
|
||||
s.shuttingDown.Store(true)
|
||||
@@ -422,6 +503,16 @@ func (s *BackupService) runScheduledBackup() {
|
||||
ctx, cancel := context.WithTimeout(s.bgCtx, 30*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
// 多实例保护: 集群部署时只让 leader 执行定时备份, 避免每个实例各自对同一个
|
||||
// 数据库跑一次全量 dump、上传时峰值内存翻倍、以及多份同名对象互相覆盖。
|
||||
// 手动触发的备份 (CreateBackup/StartBackup) 不受此限, 运维仍可随时在任一节点强制备份。
|
||||
release, ok := tryAcquireSingletonLeaderLock(ctx, s.lockCache, s.db, backupScheduledLeaderLockKey, s.instanceID, backupScheduledLeaderLockTTL)
|
||||
if !ok {
|
||||
logger.LegacyPrintf("service.backup", "[Backup] 定时备份跳过: 本实例非 leader")
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
|
||||
// 读取定时备份配置中的过期天数
|
||||
schedule, _ := s.GetSchedule(ctx)
|
||||
expireDays := 14 // 默认14天过期
|
||||
@@ -452,7 +543,7 @@ func (s *BackupService) runScheduledBackup() {
|
||||
|
||||
// ─── 备份/恢复核心 ───
|
||||
|
||||
// CreateBackup 创建全量数据库备份并上传到 S3(流式处理)
|
||||
// CreateBackup 创建全量数据库备份并上传到 S3。
|
||||
// expireDays: 备份过期天数,0=永不过期,默认14天
|
||||
func (s *BackupService) CreateBackup(ctx context.Context, triggeredBy string, expireDays int) (*BackupRecord, error) {
|
||||
if s.shuttingDown.Load() {
|
||||
@@ -506,61 +597,27 @@ func (s *BackupService) CreateBackup(ctx context.Context, triggeredBy string, ex
|
||||
ExpiresAt: expiresAt,
|
||||
}
|
||||
|
||||
// 流式执行: pg_dump -> gzip -> S3 upload
|
||||
dumpReader, err := s.dumper.Dump(ctx)
|
||||
archivePath, sizeBytes, err := s.createCompressedBackupFile(ctx)
|
||||
if err != nil {
|
||||
record.Status = "failed"
|
||||
record.ErrorMsg = fmt.Sprintf("pg_dump failed: %v", err)
|
||||
record.ErrorMsg = err.Error()
|
||||
record.FinishedAt = time.Now().Format(time.RFC3339)
|
||||
_ = s.saveRecord(ctx, record)
|
||||
return record, fmt.Errorf("pg_dump: %w", err)
|
||||
return record, err
|
||||
}
|
||||
|
||||
// 使用 io.Pipe 将 gzip 压缩数据流式传递给 S3 上传
|
||||
pr, pw := io.Pipe()
|
||||
gzipDone := make(chan error, 1)
|
||||
go func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
pw.CloseWithError(fmt.Errorf("gzip goroutine panic: %v", r)) //nolint:errcheck
|
||||
gzipDone <- fmt.Errorf("gzip goroutine panic: %v", r)
|
||||
}
|
||||
}()
|
||||
gzWriter := gzip.NewWriter(pw)
|
||||
var gzErr error
|
||||
_, gzErr = io.Copy(gzWriter, dumpReader)
|
||||
if closeErr := gzWriter.Close(); closeErr != nil && gzErr == nil {
|
||||
gzErr = closeErr
|
||||
}
|
||||
if closeErr := dumpReader.Close(); closeErr != nil && gzErr == nil {
|
||||
gzErr = closeErr
|
||||
}
|
||||
if gzErr != nil {
|
||||
_ = pw.CloseWithError(gzErr)
|
||||
} else {
|
||||
_ = pw.Close()
|
||||
}
|
||||
gzipDone <- gzErr
|
||||
}()
|
||||
|
||||
contentType := "application/gzip"
|
||||
sizeBytes, err := objectStore.Upload(ctx, s3Key, pr, contentType)
|
||||
if err != nil {
|
||||
_ = pr.CloseWithError(err) // 确保 gzip goroutine 不会悬挂
|
||||
gzErr := <-gzipDone // 安全等待 gzip goroutine 完成
|
||||
record.Status = "failed"
|
||||
errMsg := fmt.Sprintf("S3 upload failed: %v", err)
|
||||
if gzErr != nil {
|
||||
errMsg = fmt.Sprintf("gzip/dump failed: %v", gzErr)
|
||||
}
|
||||
record.ErrorMsg = errMsg
|
||||
record.FinishedAt = time.Now().Format(time.RFC3339)
|
||||
_ = s.saveRecord(ctx, record)
|
||||
return record, fmt.Errorf("backup upload: %w", err)
|
||||
}
|
||||
<-gzipDone // 确保 gzip goroutine 已退出
|
||||
|
||||
defer func() { _ = cleanupBackupFiles(archivePath) }()
|
||||
record.SizeBytes = sizeBytes
|
||||
if err := s.saveRecord(ctx, record); err != nil {
|
||||
return nil, fmt.Errorf("save initial record: %w", err)
|
||||
}
|
||||
if err := s.uploadBackupArchive(ctx, record, objectStore, s3Cfg, archivePath); err != nil {
|
||||
record.Status = "failed"
|
||||
record.ErrorMsg = err.Error()
|
||||
record.FinishedAt = time.Now().Format(time.RFC3339)
|
||||
_ = s.saveRecord(ctx, record)
|
||||
return record, err
|
||||
}
|
||||
|
||||
record.Status = "completed"
|
||||
record.FinishedAt = time.Now().Format(time.RFC3339)
|
||||
if err := s.saveRecord(ctx, record); err != nil {
|
||||
@@ -656,78 +713,43 @@ func (s *BackupService) StartBackup(ctx context.Context, triggeredBy string, exp
|
||||
_ = s.saveRecord(context.Background(), record)
|
||||
}
|
||||
}()
|
||||
s.executeBackup(record, objectStore)
|
||||
s.executeBackup(record, objectStore, s3Cfg)
|
||||
}()
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// executeBackup 后台执行备份(独立于 HTTP context)
|
||||
func (s *BackupService) executeBackup(record *BackupRecord, objectStore BackupObjectStore) {
|
||||
func (s *BackupService) executeBackup(record *BackupRecord, objectStore BackupObjectStore, s3Cfg *BackupS3Config) {
|
||||
ctx, cancel := context.WithTimeout(s.bgCtx, 30*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
// 阶段1: pg_dump
|
||||
// 阶段1: pg_dump -> gzip 临时文件
|
||||
record.Progress = "dumping"
|
||||
_ = s.saveRecord(ctx, record)
|
||||
|
||||
dumpReader, err := s.dumper.Dump(ctx)
|
||||
archivePath, sizeBytes, err := s.createCompressedBackupFile(ctx)
|
||||
if err != nil {
|
||||
record.Status = "failed"
|
||||
record.ErrorMsg = fmt.Sprintf("pg_dump failed: %v", err)
|
||||
record.ErrorMsg = err.Error()
|
||||
record.Progress = ""
|
||||
record.FinishedAt = time.Now().Format(time.RFC3339)
|
||||
_ = s.saveRecord(context.Background(), record)
|
||||
return
|
||||
}
|
||||
defer func() { _ = cleanupBackupFiles(archivePath) }()
|
||||
record.SizeBytes = sizeBytes
|
||||
|
||||
// 阶段2: gzip + upload
|
||||
// 阶段2: 单对象或分卷上传
|
||||
record.Progress = "uploading"
|
||||
_ = s.saveRecord(ctx, record)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
gzipDone := make(chan error, 1)
|
||||
go func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
pw.CloseWithError(fmt.Errorf("gzip goroutine panic: %v", r)) //nolint:errcheck
|
||||
gzipDone <- fmt.Errorf("gzip goroutine panic: %v", r)
|
||||
}
|
||||
}()
|
||||
gzWriter := gzip.NewWriter(pw)
|
||||
var gzErr error
|
||||
_, gzErr = io.Copy(gzWriter, dumpReader)
|
||||
if closeErr := gzWriter.Close(); closeErr != nil && gzErr == nil {
|
||||
gzErr = closeErr
|
||||
}
|
||||
if closeErr := dumpReader.Close(); closeErr != nil && gzErr == nil {
|
||||
gzErr = closeErr
|
||||
}
|
||||
if gzErr != nil {
|
||||
_ = pw.CloseWithError(gzErr)
|
||||
} else {
|
||||
_ = pw.Close()
|
||||
}
|
||||
gzipDone <- gzErr
|
||||
}()
|
||||
|
||||
contentType := "application/gzip"
|
||||
sizeBytes, err := objectStore.Upload(ctx, record.S3Key, pr, contentType)
|
||||
if err != nil {
|
||||
_ = pr.CloseWithError(err) // 确保 gzip goroutine 不会悬挂
|
||||
gzErr := <-gzipDone // 安全等待 gzip goroutine 完成
|
||||
if err := s.uploadBackupArchive(ctx, record, objectStore, s3Cfg, archivePath); err != nil {
|
||||
record.Status = "failed"
|
||||
errMsg := fmt.Sprintf("S3 upload failed: %v", err)
|
||||
if gzErr != nil {
|
||||
errMsg = fmt.Sprintf("gzip/dump failed: %v", gzErr)
|
||||
}
|
||||
record.ErrorMsg = errMsg
|
||||
record.ErrorMsg = err.Error()
|
||||
record.Progress = ""
|
||||
record.FinishedAt = time.Now().Format(time.RFC3339)
|
||||
_ = s.saveRecord(context.Background(), record)
|
||||
return
|
||||
}
|
||||
<-gzipDone // 确保 gzip goroutine 已退出
|
||||
|
||||
record.SizeBytes = sizeBytes
|
||||
record.Status = "completed"
|
||||
@@ -738,6 +760,108 @@ func (s *BackupService) executeBackup(record *BackupRecord, objectStore BackupOb
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) createCompressedBackupFile(ctx context.Context) (string, int64, error) {
|
||||
dumpReader, err := s.dumper.Dump(ctx)
|
||||
if err != nil {
|
||||
return "", 0, fmt.Errorf("pg_dump: %w", err)
|
||||
}
|
||||
archive, err := os.CreateTemp("", "sub2api-backup-*.sql.gz")
|
||||
if err != nil {
|
||||
_ = dumpReader.Close()
|
||||
return "", 0, fmt.Errorf("create backup archive: %w", err)
|
||||
}
|
||||
archivePath := archive.Name()
|
||||
|
||||
gzWriter := gzip.NewWriter(archive)
|
||||
_, copyErr := io.Copy(gzWriter, dumpReader)
|
||||
if closeErr := gzWriter.Close(); copyErr == nil && closeErr != nil {
|
||||
copyErr = closeErr
|
||||
}
|
||||
if closeErr := dumpReader.Close(); copyErr == nil && closeErr != nil {
|
||||
copyErr = closeErr
|
||||
}
|
||||
if closeErr := archive.Close(); copyErr == nil && closeErr != nil {
|
||||
copyErr = closeErr
|
||||
}
|
||||
if copyErr != nil {
|
||||
_ = cleanupBackupFiles(archivePath)
|
||||
return "", 0, fmt.Errorf("gzip/dump failed: %w", copyErr)
|
||||
}
|
||||
|
||||
info, err := os.Stat(archivePath)
|
||||
if err != nil {
|
||||
_ = cleanupBackupFiles(archivePath)
|
||||
return "", 0, fmt.Errorf("stat backup archive: %w", err)
|
||||
}
|
||||
return archivePath, info.Size(), nil
|
||||
}
|
||||
|
||||
func (s *BackupService) uploadBackupArchive(ctx context.Context, record *BackupRecord, objectStore BackupObjectStore, cfg *BackupS3Config, archivePath string) error {
|
||||
info, err := os.Stat(archivePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("stat backup archive: %w", err)
|
||||
}
|
||||
partSize := s.partSizeBytes
|
||||
if partSize <= 0 {
|
||||
partSize = defaultBackupPartSizeBytes
|
||||
}
|
||||
if info.Size() <= partSize {
|
||||
if _, err := objectStore.UploadFile(ctx, record.S3Key, archivePath, "application/gzip"); err != nil {
|
||||
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), backupObjectCleanupTimeout)
|
||||
cleanupErr := deleteBackupObjectKeys(cleanupCtx, objectStore, record)
|
||||
cleanupCancel()
|
||||
return errors.Join(fmt.Errorf("backup upload: %w", err), cleanupErr)
|
||||
}
|
||||
record.Parts = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
localParts, err := splitBackupFile(archivePath, partSize)
|
||||
if err != nil {
|
||||
return fmt.Errorf("split backup archive: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
paths := make([]string, 0, len(localParts))
|
||||
for _, part := range localParts {
|
||||
paths = append(paths, part.Path)
|
||||
}
|
||||
_ = cleanupBackupFiles(paths...)
|
||||
}()
|
||||
if cfg == nil {
|
||||
return errors.New("backup S3 config is unavailable for split upload")
|
||||
}
|
||||
|
||||
record.S3Key = ""
|
||||
record.Parts = make([]BackupPart, 0, len(localParts))
|
||||
partRoot := strings.TrimRight(s.buildS3Key(cfg, record.ID), "/")
|
||||
for _, part := range localParts {
|
||||
record.Parts = append(record.Parts, BackupPart{
|
||||
Index: part.Index,
|
||||
S3Key: s.buildBackupPartKey(partRoot, part.Index),
|
||||
SizeBytes: part.SizeBytes,
|
||||
SHA256: part.SHA256,
|
||||
})
|
||||
}
|
||||
if err := s.saveRecord(ctx, record); err != nil {
|
||||
return fmt.Errorf("save split backup plan: %w", err)
|
||||
}
|
||||
for i, part := range localParts {
|
||||
if _, err := objectStore.UploadFile(ctx, record.Parts[i].S3Key, part.Path, "application/octet-stream"); err != nil {
|
||||
// PUT 可能已经在对象存储端成功、但客户端因超时收到错误;
|
||||
// 因此失败时清理整份分卷计划,而不只清理此前返回成功的卷。
|
||||
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), backupObjectCleanupTimeout)
|
||||
cleanupErr := deleteBackupObjectKeys(cleanupCtx, objectStore, record)
|
||||
cleanupCancel()
|
||||
return errors.Join(fmt.Errorf("upload backup part %d: %w", part.Index, err), cleanupErr)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *BackupService) buildBackupPartKey(root string, index int) string {
|
||||
return fmt.Sprintf("%s/payload.part-%06d", strings.TrimRight(root, "/"), index)
|
||||
}
|
||||
|
||||
// RestoreBackup 从 S3 下载备份并流式恢复到数据库
|
||||
func (s *BackupService) RestoreBackup(ctx context.Context, backupID string) error {
|
||||
s.opMu.Lock()
|
||||
@@ -770,7 +894,16 @@ func (s *BackupService) RestoreBackup(ctx context.Context, backupID string) erro
|
||||
return fmt.Errorf("init object store: %w", err)
|
||||
}
|
||||
|
||||
// 从 S3 流式下载
|
||||
if len(record.Parts) > 0 {
|
||||
archivePath, err := s.downloadBackupParts(ctx, objectStore, record.Parts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = cleanupBackupFiles(archivePath) }()
|
||||
return s.restoreArchive(ctx, archivePath)
|
||||
}
|
||||
|
||||
// 旧记录从 S3 流式下载
|
||||
body, err := objectStore.Download(ctx, record.S3Key)
|
||||
if err != nil {
|
||||
return fmt.Errorf("S3 download failed: %w", err)
|
||||
@@ -866,6 +999,29 @@ func (s *BackupService) executeRestore(record *BackupRecord, objectStore BackupO
|
||||
ctx, cancel := context.WithTimeout(s.bgCtx, 30*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
if len(record.Parts) > 0 {
|
||||
archivePath, err := s.downloadBackupParts(ctx, objectStore, record.Parts)
|
||||
if err != nil {
|
||||
record.RestoreStatus = "failed"
|
||||
record.RestoreError = err.Error()
|
||||
_ = s.saveRecord(context.Background(), record)
|
||||
return
|
||||
}
|
||||
defer func() { _ = cleanupBackupFiles(archivePath) }()
|
||||
if err := s.restoreArchive(ctx, archivePath); err != nil {
|
||||
record.RestoreStatus = "failed"
|
||||
record.RestoreError = fmt.Sprintf("pg restore: %v", err)
|
||||
_ = s.saveRecord(context.Background(), record)
|
||||
return
|
||||
}
|
||||
record.RestoreStatus = "completed"
|
||||
record.RestoredAt = time.Now().Format(time.RFC3339)
|
||||
if err := s.saveRecord(context.Background(), record); err != nil {
|
||||
logger.LegacyPrintf("service.backup", "[Backup] 保存恢复记录失败: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
body, err := objectStore.Download(ctx, record.S3Key)
|
||||
if err != nil {
|
||||
record.RestoreStatus = "failed"
|
||||
@@ -898,6 +1054,79 @@ func (s *BackupService) executeRestore(record *BackupRecord, objectStore BackupO
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) downloadBackupParts(ctx context.Context, objectStore BackupObjectStore, parts []BackupPart) (path string, err error) {
|
||||
if len(parts) == 0 {
|
||||
return "", errors.New("backup parts are empty")
|
||||
}
|
||||
ordered := append([]BackupPart(nil), parts...)
|
||||
sort.Slice(ordered, func(i, j int) bool { return ordered[i].Index < ordered[j].Index })
|
||||
for i, part := range ordered {
|
||||
if part.Index != i+1 || part.S3Key == "" || part.SizeBytes <= 0 {
|
||||
return "", fmt.Errorf("invalid backup part metadata at index %d", i+1)
|
||||
}
|
||||
}
|
||||
|
||||
archive, err := os.CreateTemp("", "sub2api-restore-*.sql.gz")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create restore archive: %w", err)
|
||||
}
|
||||
path = archive.Name()
|
||||
cleanup := func() {
|
||||
_ = archive.Close()
|
||||
_ = cleanupBackupFiles(path)
|
||||
}
|
||||
|
||||
for _, part := range ordered {
|
||||
body, downloadErr := objectStore.Download(ctx, part.S3Key)
|
||||
if downloadErr != nil {
|
||||
cleanup()
|
||||
return "", fmt.Errorf("download backup part %d: %w", part.Index, downloadErr)
|
||||
}
|
||||
hash := sha256.New()
|
||||
written, copyErr := io.Copy(io.MultiWriter(archive, hash), body)
|
||||
closeErr := body.Close()
|
||||
if copyErr != nil {
|
||||
cleanup()
|
||||
return "", fmt.Errorf("read backup part %d: %w", part.Index, copyErr)
|
||||
}
|
||||
if closeErr != nil {
|
||||
cleanup()
|
||||
return "", fmt.Errorf("close backup part %d: %w", part.Index, closeErr)
|
||||
}
|
||||
if written != part.SizeBytes {
|
||||
cleanup()
|
||||
return "", fmt.Errorf("backup part %d size mismatch: got %d, want %d", part.Index, written, part.SizeBytes)
|
||||
}
|
||||
if part.SHA256 != "" && !strings.EqualFold(part.SHA256, hex.EncodeToString(hash.Sum(nil))) {
|
||||
cleanup()
|
||||
return "", fmt.Errorf("backup part %d checksum mismatch", part.Index)
|
||||
}
|
||||
}
|
||||
if err := archive.Close(); err != nil {
|
||||
_ = cleanupBackupFiles(path)
|
||||
return "", fmt.Errorf("close restore archive: %w", err)
|
||||
}
|
||||
return path, nil
|
||||
}
|
||||
|
||||
func (s *BackupService) restoreArchive(ctx context.Context, archivePath string) error {
|
||||
archive, err := os.Open(archivePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open restore archive: %w", err)
|
||||
}
|
||||
defer func() { _ = archive.Close() }()
|
||||
|
||||
gzReader, err := gzip.NewReader(archive)
|
||||
if err != nil {
|
||||
return fmt.Errorf("gzip reader: %w", err)
|
||||
}
|
||||
defer func() { _ = gzReader.Close() }()
|
||||
if err := s.dumper.Restore(ctx, gzReader); err != nil {
|
||||
return fmt.Errorf("pg restore: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ─── 备份记录管理 ───
|
||||
|
||||
func (s *BackupService) ListBackups(ctx context.Context) ([]BackupRecord, error) {
|
||||
@@ -946,45 +1175,67 @@ func (s *BackupService) DeleteBackup(ctx context.Context, backupID string) error
|
||||
if found == nil {
|
||||
return ErrBackupNotFound
|
||||
}
|
||||
if found.Status == "running" {
|
||||
// 后台上传仍可能依赖 Parts 计划;删除对象会让随后完成的记录引用失效卷。
|
||||
return ErrBackupInProgress
|
||||
}
|
||||
|
||||
// 从 S3 删除
|
||||
if found.S3Key != "" && found.Status == "completed" {
|
||||
s3Cfg, err := s.loadS3Config(ctx)
|
||||
if err == nil && s3Cfg != nil && s3Cfg.IsConfigured() {
|
||||
objectStore, err := s.getOrCreateStore(ctx, s3Cfg)
|
||||
if err == nil {
|
||||
_ = objectStore.Delete(ctx, found.S3Key)
|
||||
}
|
||||
}
|
||||
// 从对象存储删除所有单文件或分卷对象。删除不完整时保留记录,便于重试。
|
||||
if err := s.deleteBackupObjects(ctx, found); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return s.saveRecordsLocked(ctx, remaining)
|
||||
}
|
||||
|
||||
// GetBackupDownloadURL 获取备份文件预签名下载 URL
|
||||
func (s *BackupService) GetBackupDownloadURL(ctx context.Context, backupID string) (string, error) {
|
||||
func (s *BackupService) GetBackupDownloadURL(ctx context.Context, backupID string) (BackupDownloadResponse, error) {
|
||||
var download BackupDownloadResponse
|
||||
record, err := s.GetBackupRecord(ctx, backupID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return download, err
|
||||
}
|
||||
if record.Status != "completed" {
|
||||
return "", infraerrors.BadRequest("BACKUP_NOT_COMPLETED", "backup is not completed")
|
||||
return download, infraerrors.BadRequest("BACKUP_NOT_COMPLETED", "backup is not completed")
|
||||
}
|
||||
|
||||
s3Cfg, err := s.loadS3Config(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return download, err
|
||||
}
|
||||
objectStore, err := s.getOrCreateStore(ctx, s3Cfg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return download, err
|
||||
}
|
||||
|
||||
if len(record.Parts) > 0 {
|
||||
parts := append([]BackupPart(nil), record.Parts...)
|
||||
sort.Slice(parts, func(i, j int) bool { return parts[i].Index < parts[j].Index })
|
||||
for i, part := range parts {
|
||||
if part.Index != i+1 || part.S3Key == "" || part.SizeBytes <= 0 {
|
||||
return download, fmt.Errorf("invalid backup part metadata at index %d", i+1)
|
||||
}
|
||||
url, presignErr := objectStore.PresignURL(ctx, part.S3Key, 1*time.Hour)
|
||||
if presignErr != nil {
|
||||
return download, fmt.Errorf("presign backup part %d: %w", part.Index, presignErr)
|
||||
}
|
||||
download.Parts = append(download.Parts, BackupDownloadPart{
|
||||
Index: part.Index,
|
||||
SizeBytes: part.SizeBytes,
|
||||
URL: url,
|
||||
})
|
||||
}
|
||||
return download, nil
|
||||
}
|
||||
if record.S3Key == "" {
|
||||
return download, errors.New("backup object key is empty")
|
||||
}
|
||||
url, err := objectStore.PresignURL(ctx, record.S3Key, 1*time.Hour)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("presign url: %w", err)
|
||||
return download, fmt.Errorf("presign url: %w", err)
|
||||
}
|
||||
return url, nil
|
||||
download.URL = url
|
||||
return download, nil
|
||||
}
|
||||
|
||||
// ─── 内部方法 ───
|
||||
@@ -1141,28 +1392,87 @@ func (s *BackupService) cleanupOldBackups(ctx context.Context, schedule *BackupS
|
||||
}
|
||||
}
|
||||
|
||||
// 删除 S3 上的文件
|
||||
var cleanupErrs []error
|
||||
deletedCount := 0
|
||||
for _, r := range toDelete {
|
||||
if r.S3Key != "" {
|
||||
_ = s.deleteS3Object(ctx, r.S3Key)
|
||||
if err := s.deleteBackupObjects(ctx, &r); err != nil {
|
||||
// 对象删除失败时保留记录,避免丢失后续重试所需的 key。
|
||||
toKeep = append(toKeep, r)
|
||||
cleanupErrs = append(cleanupErrs, fmt.Errorf("cleanup backup %s: %w", r.ID, err))
|
||||
continue
|
||||
}
|
||||
deletedCount++
|
||||
}
|
||||
|
||||
if len(toDelete) > 0 {
|
||||
logger.LegacyPrintf("service.backup", "[Backup] 自动清理了 %d 个过期备份", len(toDelete))
|
||||
return s.saveRecordsLocked(ctx, toKeep)
|
||||
if err := s.saveRecordsLocked(ctx, toKeep); err != nil {
|
||||
cleanupErrs = append(cleanupErrs, fmt.Errorf("save backup records after cleanup: %w", err))
|
||||
}
|
||||
if deletedCount > 0 {
|
||||
logger.LegacyPrintf("service.backup", "[Backup] 自动清理了 %d 个过期备份", deletedCount)
|
||||
}
|
||||
return errors.Join(cleanupErrs...)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *BackupService) deleteS3Object(ctx context.Context, key string) error {
|
||||
// backupObjectKeys 返回一条备份记录关联的全部对象 key。
|
||||
// 新记录使用 Parts,旧记录使用 S3Key;两者同时存在时也全部返回,便于清理异常残留对象。
|
||||
func backupObjectKeys(record *BackupRecord) []string {
|
||||
if record == nil {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(record.Parts)+1)
|
||||
seen := make(map[string]struct{}, len(record.Parts)+1)
|
||||
appendKey := func(key string) {
|
||||
if key == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[key]; ok {
|
||||
return
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
keys = append(keys, key)
|
||||
}
|
||||
appendKey(record.S3Key)
|
||||
parts := append([]BackupPart(nil), record.Parts...)
|
||||
sort.Slice(parts, func(i, j int) bool { return parts[i].Index < parts[j].Index })
|
||||
for _, part := range parts {
|
||||
appendKey(part.S3Key)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
// deleteBackupObjects 尝试删除记录关联的所有对象,并聚合删除错误。
|
||||
func (s *BackupService) deleteBackupObjects(ctx context.Context, record *BackupRecord) error {
|
||||
if len(backupObjectKeys(record)) == 0 {
|
||||
return nil
|
||||
}
|
||||
s3Cfg, err := s.loadS3Config(ctx)
|
||||
if err != nil || s3Cfg == nil {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if s3Cfg == nil || !s3Cfg.IsConfigured() {
|
||||
// 兼容没有配置对象存储的旧记录:记录仍可被删除。
|
||||
return nil
|
||||
}
|
||||
objectStore, err := s.getOrCreateStore(ctx, s3Cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return objectStore.Delete(ctx, key)
|
||||
return deleteBackupObjectKeys(ctx, objectStore, record)
|
||||
}
|
||||
|
||||
func deleteBackupObjectKeys(ctx context.Context, objectStore BackupObjectStore, record *BackupRecord) error {
|
||||
keys := backupObjectKeys(record)
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
var errs []error
|
||||
for _, key := range keys {
|
||||
if deleteErr := objectStore.Delete(ctx, key); deleteErr != nil {
|
||||
errs = append(errs, fmt.Errorf("delete backup object %q: %w", key, deleteErr))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
@@ -4,10 +4,13 @@ package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
@@ -21,8 +24,10 @@ import (
|
||||
// ─── Mocks ───
|
||||
|
||||
type mockSettingRepo struct {
|
||||
mu sync.Mutex
|
||||
data map[string]string
|
||||
mu sync.Mutex
|
||||
data map[string]string
|
||||
getValueErr error
|
||||
getValueCalls int
|
||||
}
|
||||
|
||||
func newMockSettingRepo() *mockSettingRepo {
|
||||
@@ -42,6 +47,10 @@ func (m *mockSettingRepo) Get(_ context.Context, key string) (*Setting, error) {
|
||||
func (m *mockSettingRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.getValueCalls++
|
||||
if m.getValueErr != nil {
|
||||
return "", m.getValueErr
|
||||
}
|
||||
v, ok := m.data[key]
|
||||
if !ok {
|
||||
return "", nil
|
||||
@@ -159,12 +168,49 @@ func (d *blockingDumper) Restore(_ context.Context, data io.Reader) error {
|
||||
}
|
||||
|
||||
type mockObjectStore struct {
|
||||
objects map[string][]byte
|
||||
mu sync.Mutex
|
||||
objects map[string][]byte
|
||||
mu sync.Mutex
|
||||
failUploadFileAt int
|
||||
uploadFileCalls int
|
||||
deletedKeys []string
|
||||
failDeleteKeys map[string]error
|
||||
}
|
||||
|
||||
type cancelingUploadFailureStore struct {
|
||||
*mockObjectStore
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (m *cancelingUploadFailureStore) UploadFile(_ context.Context, key string, filePath string, _ string) (int64, error) {
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
data, readErr := io.ReadAll(file)
|
||||
closeErr := file.Close()
|
||||
if readErr != nil {
|
||||
return 0, readErr
|
||||
}
|
||||
if closeErr != nil {
|
||||
return 0, closeErr
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
m.objects[key] = data
|
||||
m.mu.Unlock()
|
||||
m.cancel()
|
||||
return 0, fmt.Errorf("injected upload failure after object landed")
|
||||
}
|
||||
|
||||
func (m *cancelingUploadFailureStore) Delete(ctx context.Context, key string) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return m.mockObjectStore.Delete(ctx, key)
|
||||
}
|
||||
|
||||
func newMockObjectStore() *mockObjectStore {
|
||||
return &mockObjectStore{objects: make(map[string][]byte)}
|
||||
return &mockObjectStore{objects: make(map[string][]byte), failDeleteKeys: make(map[string]error)}
|
||||
}
|
||||
|
||||
func (m *mockObjectStore) Upload(_ context.Context, key string, body io.Reader, _ string) (int64, error) {
|
||||
@@ -178,6 +224,23 @@ func (m *mockObjectStore) Upload(_ context.Context, key string, body io.Reader,
|
||||
return int64(len(data)), nil
|
||||
}
|
||||
|
||||
func (m *mockObjectStore) UploadFile(ctx context.Context, key string, filePath string, contentType string) (int64, error) {
|
||||
m.mu.Lock()
|
||||
m.uploadFileCalls++
|
||||
call := m.uploadFileCalls
|
||||
failAt := m.failUploadFileAt
|
||||
m.mu.Unlock()
|
||||
if failAt > 0 && call == failAt {
|
||||
return 0, fmt.Errorf("injected upload failure at call %d", call)
|
||||
}
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer func() { _ = file.Close() }()
|
||||
return m.Upload(ctx, key, file, contentType)
|
||||
}
|
||||
|
||||
func (m *mockObjectStore) Download(_ context.Context, key string) (io.ReadCloser, error) {
|
||||
m.mu.Lock()
|
||||
data, ok := m.objects[key]
|
||||
@@ -190,6 +253,11 @@ func (m *mockObjectStore) Download(_ context.Context, key string) (io.ReadCloser
|
||||
|
||||
func (m *mockObjectStore) Delete(_ context.Context, key string) error {
|
||||
m.mu.Lock()
|
||||
m.deletedKeys = append(m.deletedKeys, key)
|
||||
if err, ok := m.failDeleteKeys[key]; ok {
|
||||
m.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
delete(m.objects, key)
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
@@ -405,6 +473,122 @@ func TestBackupService_CreateBackup_Streaming(t *testing.T) {
|
||||
store.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestBackupService_CreateBackup_SplitsCompressedArchive(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
dumpContent := entropyBackupFixture(512)
|
||||
dumper := &mockDumper{dumpData: dumpContent}
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, dumper, store)
|
||||
svc.partSizeBytes = 32
|
||||
|
||||
record, err := svc.CreateBackup(context.Background(), "manual", 14)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "completed", record.Status)
|
||||
require.Greater(t, len(record.Parts), 1)
|
||||
require.Empty(t, record.S3Key)
|
||||
|
||||
var compressed bytes.Buffer
|
||||
store.mu.Lock()
|
||||
for _, part := range record.Parts {
|
||||
data, ok := store.objects[part.S3Key]
|
||||
require.True(t, ok)
|
||||
require.LessOrEqual(t, len(data), 32)
|
||||
compressed.Write(data)
|
||||
}
|
||||
store.mu.Unlock()
|
||||
|
||||
gzReader, err := gzip.NewReader(bytes.NewReader(compressed.Bytes()))
|
||||
require.NoError(t, err)
|
||||
decompressed, err := io.ReadAll(gzReader)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, gzReader.Close())
|
||||
require.Equal(t, dumpContent, decompressed)
|
||||
}
|
||||
|
||||
func TestBackupService_StartBackup_SplitsCompressedArchive(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{dumpData: entropyBackupFixture(512)}, store)
|
||||
svc.partSizeBytes = 32
|
||||
|
||||
record, err := svc.StartBackup(context.Background(), "manual", 14)
|
||||
require.NoError(t, err)
|
||||
svc.wg.Wait()
|
||||
|
||||
final, err := svc.GetBackupRecord(context.Background(), record.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "completed", final.Status)
|
||||
require.Greater(t, len(final.Parts), 1)
|
||||
require.Empty(t, final.S3Key)
|
||||
}
|
||||
|
||||
func TestBackupService_StartBackup_UploadFailureCleansUploadedParts(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
store.failUploadFileAt = 2
|
||||
svc := newTestBackupService(repo, &mockDumper{dumpData: entropyBackupFixture(512)}, store)
|
||||
svc.partSizeBytes = 32
|
||||
|
||||
record, err := svc.StartBackup(context.Background(), "manual", 14)
|
||||
require.NoError(t, err)
|
||||
svc.wg.Wait()
|
||||
|
||||
final, err := svc.GetBackupRecord(context.Background(), record.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "failed", final.Status)
|
||||
require.NotEmpty(t, final.Parts)
|
||||
store.mu.Lock()
|
||||
deletedKeys := append([]string(nil), store.deletedKeys...)
|
||||
store.mu.Unlock()
|
||||
for _, part := range final.Parts {
|
||||
require.Contains(t, deletedKeys, part.S3Key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupService_UploadFailureCleanupUsesDetachedContext(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore())
|
||||
svc.partSizeBytes = 4
|
||||
|
||||
archive, err := os.CreateTemp("", "backup-upload-context-*.gz")
|
||||
require.NoError(t, err)
|
||||
archivePath := archive.Name()
|
||||
defer func() { _ = os.Remove(archivePath) }()
|
||||
_, err = archive.Write([]byte("0123456789"))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, archive.Close())
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
store := &cancelingUploadFailureStore{
|
||||
mockObjectStore: newMockObjectStore(),
|
||||
cancel: cancel,
|
||||
}
|
||||
record := &BackupRecord{ID: "cancel-cleanup", S3Key: "backups/cancel-cleanup.sql.gz"}
|
||||
|
||||
err = svc.uploadBackupArchive(ctx, record, store, &BackupS3Config{Prefix: "backups"}, archivePath)
|
||||
require.Error(t, err)
|
||||
require.NotContains(t, err.Error(), "context canceled")
|
||||
|
||||
store.mu.Lock()
|
||||
defer store.mu.Unlock()
|
||||
for _, part := range record.Parts {
|
||||
require.Contains(t, store.deletedKeys, part.S3Key)
|
||||
require.NotContains(t, store.objects, part.S3Key)
|
||||
}
|
||||
}
|
||||
|
||||
func entropyBackupFixture(size int) []byte {
|
||||
data := make([]byte, size)
|
||||
for i := range data {
|
||||
data[i] = byte((i*31 + 17) % 251)
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func TestBackupService_CreateBackup_DumpFailure(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
@@ -445,6 +629,56 @@ func TestBackupService_CreateBackup_ConcurrentBlocked(t *testing.T) {
|
||||
require.ErrorIs(t, err, ErrBackupInProgress)
|
||||
}
|
||||
|
||||
// TestBackupService_RunScheduledBackup_LeaderElection verifies the scheduled
|
||||
// backup is gated by a cross-instance leader lock: a non-leader instance skips
|
||||
// the dump entirely so a clustered deployment does not run N identical backups
|
||||
// against the same database, while the leader runs it and releases the lock
|
||||
// afterward. Manual backups (CreateBackup/StartBackup) are intentionally left
|
||||
// ungated and are covered by the other tests.
|
||||
func TestBackupService_RunScheduledBackup_LeaderElection(t *testing.T) {
|
||||
t.Run("non-leader skips", func(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("data")}, store)
|
||||
|
||||
// A peer already owns the lock, so this instance is not the leader.
|
||||
cache := &fakeLeaderLockCache{}
|
||||
peerRelease, ok := tryAcquireSingletonLeaderLock(context.Background(), cache, nil, backupScheduledLeaderLockKey, "peer", time.Minute)
|
||||
require.True(t, ok)
|
||||
defer peerRelease()
|
||||
|
||||
svc.SetLeaderLock(cache, nil)
|
||||
svc.runScheduledBackup()
|
||||
|
||||
store.mu.Lock()
|
||||
require.Empty(t, store.objects, "non-leader must not upload a backup")
|
||||
store.mu.Unlock()
|
||||
|
||||
records, err := svc.ListBackups(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, records, "non-leader must not create a backup record")
|
||||
require.Equal(t, "peer", cache.heldBy(backupScheduledLeaderLockKey), "peer keeps the lock")
|
||||
})
|
||||
|
||||
t.Run("leader runs and releases", func(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("-- dump\n")}, store)
|
||||
|
||||
cache := &fakeLeaderLockCache{}
|
||||
svc.SetLeaderLock(cache, nil)
|
||||
svc.runScheduledBackup()
|
||||
|
||||
records, err := svc.ListBackups(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Len(t, records, 1, "leader creates exactly one backup record")
|
||||
require.Equal(t, "completed", records[0].Status)
|
||||
require.Empty(t, cache.heldBy(backupScheduledLeaderLockKey), "leader releases the lock when done")
|
||||
})
|
||||
}
|
||||
|
||||
func TestBackupService_RestoreBackup_Streaming(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
@@ -466,6 +700,123 @@ func TestBackupService_RestoreBackup_Streaming(t *testing.T) {
|
||||
require.Equal(t, dumpContent, string(dumper.restored))
|
||||
}
|
||||
|
||||
func TestBackupService_RestoreBackup_SplitParts(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
dumpContent := entropyBackupFixture(512)
|
||||
dumper := &mockDumper{}
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, dumper, store)
|
||||
|
||||
compressed := gzipBackupBytes(t, dumpContent)
|
||||
parts := splitBackupBytes(compressed, 11)
|
||||
recordParts := make([]BackupPart, 0, len(parts))
|
||||
for i, data := range parts {
|
||||
key := fmt.Sprintf("backups/split-1/payload.part-%06d", i+1)
|
||||
store.objects[key] = data
|
||||
recordParts = append(recordParts, BackupPart{
|
||||
Index: i + 1,
|
||||
S3Key: key,
|
||||
SizeBytes: int64(len(data)),
|
||||
SHA256: fmt.Sprintf("%x", sha256.Sum256(data)),
|
||||
})
|
||||
}
|
||||
record := &BackupRecord{
|
||||
ID: "split-1",
|
||||
Status: "completed",
|
||||
Parts: recordParts,
|
||||
SizeBytes: int64(len(compressed)),
|
||||
}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), record))
|
||||
|
||||
require.NoError(t, svc.RestoreBackup(context.Background(), record.ID))
|
||||
require.Equal(t, dumpContent, dumper.restored)
|
||||
}
|
||||
|
||||
func TestBackupService_RestoreBackup_SplitPartsMissingPartDoesNotRestore(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
dumpContent := entropyBackupFixture(256)
|
||||
dumper := &mockDumper{}
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, dumper, store)
|
||||
|
||||
compressed := gzipBackupBytes(t, dumpContent)
|
||||
parts := splitBackupBytes(compressed, 11)
|
||||
recordParts := make([]BackupPart, 0, len(parts))
|
||||
for i, data := range parts {
|
||||
key := fmt.Sprintf("backups/split-missing/payload.part-%06d", i+1)
|
||||
store.objects[key] = data
|
||||
recordParts = append(recordParts, BackupPart{
|
||||
Index: i + 1,
|
||||
S3Key: key,
|
||||
SizeBytes: int64(len(data)),
|
||||
SHA256: fmt.Sprintf("%x", sha256.Sum256(data)),
|
||||
})
|
||||
}
|
||||
delete(store.objects, recordParts[1].S3Key)
|
||||
record := &BackupRecord{ID: "split-missing", Status: "completed", Parts: recordParts}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), record))
|
||||
|
||||
require.Error(t, svc.RestoreBackup(context.Background(), record.ID))
|
||||
require.Empty(t, dumper.restored)
|
||||
}
|
||||
|
||||
func TestBackupService_DownloadBackupPartsRejectsMismatchedMetadata(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
part BackupPart
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "size",
|
||||
part: BackupPart{Index: 1, S3Key: "backups/mismatch/size", SizeBytes: 4},
|
||||
want: "size mismatch",
|
||||
},
|
||||
{
|
||||
name: "checksum",
|
||||
part: BackupPart{Index: 1, S3Key: "backups/mismatch/checksum", SizeBytes: 3, SHA256: "bad-checksum"},
|
||||
want: "checksum mismatch",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
store.objects[tt.part.S3Key] = []byte("abc")
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
|
||||
_, err := svc.downloadBackupParts(context.Background(), store, []BackupPart{tt.part})
|
||||
require.ErrorContains(t, err, tt.want)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func gzipBackupBytes(t *testing.T, content []byte) []byte {
|
||||
t.Helper()
|
||||
var out bytes.Buffer
|
||||
writer := gzip.NewWriter(&out)
|
||||
_, err := writer.Write(content)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, writer.Close())
|
||||
return out.Bytes()
|
||||
}
|
||||
|
||||
func splitBackupBytes(data []byte, partSize int) [][]byte {
|
||||
parts := make([][]byte, 0, (len(data)+partSize-1)/partSize)
|
||||
for len(data) > 0 {
|
||||
size := partSize
|
||||
if len(data) < size {
|
||||
size = len(data)
|
||||
}
|
||||
parts = append(parts, append([]byte(nil), data[:size]...))
|
||||
data = data[size:]
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
||||
func TestBackupService_RestoreBackup_NotCompleted(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
@@ -512,6 +863,34 @@ func TestBackupService_DeleteBackup(t *testing.T) {
|
||||
require.ErrorIs(t, err, ErrBackupNotFound)
|
||||
}
|
||||
|
||||
func TestBackupService_DeleteBackup_RunningKeepsUploadObjects(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
parts := []BackupPart{
|
||||
{Index: 1, S3Key: "backups/running/payload.part-000001", SizeBytes: 3},
|
||||
{Index: 2, S3Key: "backups/running/payload.part-000002", SizeBytes: 3},
|
||||
}
|
||||
for _, part := range parts {
|
||||
store.objects[part.S3Key] = []byte("abc")
|
||||
}
|
||||
record := &BackupRecord{ID: "running-delete", Status: "running", Parts: parts}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), record))
|
||||
|
||||
err := svc.DeleteBackup(context.Background(), record.ID)
|
||||
require.ErrorIs(t, err, ErrBackupInProgress)
|
||||
store.mu.Lock()
|
||||
require.Empty(t, store.deletedKeys)
|
||||
for _, part := range parts {
|
||||
require.Contains(t, store.objects, part.S3Key)
|
||||
}
|
||||
store.mu.Unlock()
|
||||
got, getErr := svc.GetBackupRecord(context.Background(), record.ID)
|
||||
require.NoError(t, getErr)
|
||||
require.Equal(t, "running", got.Status)
|
||||
}
|
||||
|
||||
func TestBackupService_GetDownloadURL(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
@@ -523,9 +902,103 @@ func TestBackupService_GetDownloadURL(t *testing.T) {
|
||||
record, err := svc.CreateBackup(context.Background(), "manual", 14)
|
||||
require.NoError(t, err)
|
||||
|
||||
url, err := svc.GetBackupDownloadURL(context.Background(), record.ID)
|
||||
download, err := svc.GetBackupDownloadURL(context.Background(), record.ID)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, url, "https://presigned.example.com/")
|
||||
require.Contains(t, download.URL, "https://presigned.example.com/")
|
||||
}
|
||||
|
||||
func TestBackupService_DeleteBackup_SplitPartsFailureKeepsRecord(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
|
||||
parts := []BackupPart{
|
||||
{Index: 1, S3Key: "backups/split/payload.part-000001", SizeBytes: 3},
|
||||
{Index: 2, S3Key: "backups/split/payload.part-000002", SizeBytes: 3},
|
||||
{Index: 3, S3Key: "backups/split/payload.part-000003", SizeBytes: 3},
|
||||
}
|
||||
for _, part := range parts {
|
||||
store.objects[part.S3Key] = []byte("abc")
|
||||
}
|
||||
store.failDeleteKeys[parts[1].S3Key] = fmt.Errorf("delete failed")
|
||||
record := &BackupRecord{ID: "split-delete", Status: "completed", Parts: parts}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), record))
|
||||
|
||||
err := svc.DeleteBackup(context.Background(), record.ID)
|
||||
require.Error(t, err)
|
||||
|
||||
store.mu.Lock()
|
||||
deleted := append([]string(nil), store.deletedKeys...)
|
||||
store.mu.Unlock()
|
||||
for _, part := range parts {
|
||||
require.Contains(t, deleted, part.S3Key)
|
||||
}
|
||||
got, getErr := svc.GetBackupRecord(context.Background(), record.ID)
|
||||
require.NoError(t, getErr)
|
||||
require.Equal(t, record.ID, got.ID)
|
||||
store.mu.Lock()
|
||||
require.Contains(t, store.objects, parts[1].S3Key)
|
||||
store.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestBackupService_GetDownloadURL_SplitParts(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
|
||||
parts := []BackupPart{
|
||||
{Index: 2, S3Key: "backups/split/payload.part-000002", SizeBytes: 7},
|
||||
{Index: 1, S3Key: "backups/split/payload.part-000001", SizeBytes: 5},
|
||||
}
|
||||
record := &BackupRecord{ID: "split-download", Status: "completed", Parts: parts}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), record))
|
||||
|
||||
download, err := svc.GetBackupDownloadURL(context.Background(), record.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, download.URL)
|
||||
require.Len(t, download.Parts, 2)
|
||||
require.Equal(t, 1, download.Parts[0].Index)
|
||||
require.Equal(t, int64(5), download.Parts[0].SizeBytes)
|
||||
require.Equal(t, "https://presigned.example.com/backups/split/payload.part-000001", download.Parts[0].URL)
|
||||
require.Equal(t, 2, download.Parts[1].Index)
|
||||
}
|
||||
|
||||
func TestBackupService_CleanupOldBackups_SplitParts(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
now := time.Now()
|
||||
parts := []BackupPart{
|
||||
{Index: 1, S3Key: "backups/old/payload.part-000001", SizeBytes: 3},
|
||||
{Index: 2, S3Key: "backups/old/payload.part-000002", SizeBytes: 3},
|
||||
}
|
||||
for _, part := range parts {
|
||||
store.objects[part.S3Key] = []byte("abc")
|
||||
}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), &BackupRecord{
|
||||
ID: "new",
|
||||
Status: "completed",
|
||||
StartedAt: now.Format(time.RFC3339),
|
||||
}))
|
||||
require.NoError(t, svc.saveRecord(context.Background(), &BackupRecord{
|
||||
ID: "old",
|
||||
Status: "completed",
|
||||
StartedAt: now.Add(-2 * time.Hour).Format(time.RFC3339),
|
||||
Parts: parts,
|
||||
}))
|
||||
|
||||
err := svc.cleanupOldBackups(context.Background(), &BackupScheduleConfig{RetainCount: 1})
|
||||
require.NoError(t, err)
|
||||
_, err = svc.GetBackupRecord(context.Background(), "old")
|
||||
require.ErrorIs(t, err, ErrBackupNotFound)
|
||||
store.mu.Lock()
|
||||
for _, part := range parts {
|
||||
require.NotContains(t, store.objects, part.S3Key)
|
||||
}
|
||||
store.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestBackupService_ListBackups_Sorted(t *testing.T) {
|
||||
@@ -690,6 +1163,65 @@ func TestRecoverStaleRecords(t *testing.T) {
|
||||
require.Contains(t, r2.RestoreError, "server restart")
|
||||
}
|
||||
|
||||
func TestBackupService_RecoverStaleRecords_CleansBackupObjects(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
parts := []BackupPart{
|
||||
{Index: 1, S3Key: "backups/stale/payload.part-000001", SizeBytes: 3},
|
||||
{Index: 2, S3Key: "backups/stale/payload.part-000002", SizeBytes: 3},
|
||||
}
|
||||
for _, part := range parts {
|
||||
store.objects[part.S3Key] = []byte("abc")
|
||||
}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), &BackupRecord{
|
||||
ID: "stale-parts",
|
||||
Status: "running",
|
||||
Parts: parts,
|
||||
StartedAt: time.Now().Add(-time.Hour).Format(time.RFC3339),
|
||||
}))
|
||||
|
||||
svc.recoverStaleRecords()
|
||||
|
||||
record, err := svc.GetBackupRecord(context.Background(), "stale-parts")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "failed", record.Status)
|
||||
store.mu.Lock()
|
||||
defer store.mu.Unlock()
|
||||
for _, part := range parts {
|
||||
require.Contains(t, store.deletedKeys, part.S3Key)
|
||||
require.NotContains(t, store.objects, part.S3Key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupService_RecoverStaleRecords_PreservesKeysWhenCleanupFails(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, &mockDumper{}, store)
|
||||
part := BackupPart{Index: 1, S3Key: "backups/stale-failed/payload.part-000001", SizeBytes: 3}
|
||||
store.objects[part.S3Key] = []byte("abc")
|
||||
store.failDeleteKeys[part.S3Key] = fmt.Errorf("delete failed")
|
||||
require.NoError(t, svc.saveRecord(context.Background(), &BackupRecord{
|
||||
ID: "stale-cleanup-failed",
|
||||
Status: "running",
|
||||
Parts: []BackupPart{part},
|
||||
StartedAt: time.Now().Add(-time.Hour).Format(time.RFC3339),
|
||||
}))
|
||||
|
||||
svc.recoverStaleRecords()
|
||||
|
||||
record, err := svc.GetBackupRecord(context.Background(), "stale-cleanup-failed")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "failed", record.Status)
|
||||
require.Contains(t, record.ErrorMsg, "cleanup failed")
|
||||
require.Equal(t, part.S3Key, record.Parts[0].S3Key)
|
||||
store.mu.Lock()
|
||||
defer store.mu.Unlock()
|
||||
require.Contains(t, store.objects, part.S3Key)
|
||||
}
|
||||
|
||||
func TestGracefulShutdown(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
@@ -753,3 +1285,38 @@ func TestStartRestore_Async(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "completed", final.RestoreStatus)
|
||||
}
|
||||
|
||||
func TestBackupService_StartRestore_SplitParts(t *testing.T) {
|
||||
repo := newMockSettingRepo()
|
||||
seedS3Config(t, repo)
|
||||
dumpContent := entropyBackupFixture(384)
|
||||
dumper := &mockDumper{}
|
||||
store := newMockObjectStore()
|
||||
svc := newTestBackupService(repo, dumper, store)
|
||||
|
||||
compressed := gzipBackupBytes(t, dumpContent)
|
||||
parts := splitBackupBytes(compressed, 13)
|
||||
recordParts := make([]BackupPart, 0, len(parts))
|
||||
for i, data := range parts {
|
||||
key := fmt.Sprintf("backups/split-async/payload.part-%06d", i+1)
|
||||
store.objects[key] = data
|
||||
recordParts = append(recordParts, BackupPart{
|
||||
Index: i + 1,
|
||||
S3Key: key,
|
||||
SizeBytes: int64(len(data)),
|
||||
SHA256: fmt.Sprintf("%x", sha256.Sum256(data)),
|
||||
})
|
||||
}
|
||||
record := &BackupRecord{ID: "split-async", Status: "completed", Parts: recordParts}
|
||||
require.NoError(t, svc.saveRecord(context.Background(), record))
|
||||
|
||||
started, err := svc.StartRestore(context.Background(), record.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "running", started.RestoreStatus)
|
||||
svc.wg.Wait()
|
||||
|
||||
final, err := svc.GetBackupRecord(context.Background(), record.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "completed", final.RestoreStatus)
|
||||
require.Equal(t, dumpContent, dumper.restored)
|
||||
}
|
||||
|
||||
@@ -10,8 +10,8 @@ func TestCalculateSearchCost(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := &BillingService{}
|
||||
require.Equal(t, 0.0, s.CalculateSearchCost(0, floatPtr(10), 1).ActualCost)
|
||||
// nil price → default $10/1k: 5 calls = 0.05
|
||||
require.InDelta(t, 0.05, s.CalculateSearchCost(5, nil, 1).ActualCost, 1e-9)
|
||||
// nil price -> official xAI default $5/1k: 5 calls = 0.025
|
||||
require.InDelta(t, 0.025, s.CalculateSearchCost(5, nil, 1).ActualCost, 1e-9)
|
||||
// explicit 0 → free
|
||||
require.Equal(t, 0.0, s.CalculateSearchCost(5, floatPtr(0), 1).ActualCost)
|
||||
price := 10.0
|
||||
@@ -30,10 +30,10 @@ func TestCalculateAudioCost(t *testing.T) {
|
||||
require.InDelta(t, 1.5, s.CalculateAudioCost("tts", 0.1, cfg, 1).ActualCost, 1e-9)
|
||||
require.InDelta(t, 0.25, s.CalculateAudioCost("stt", 0.5, cfg, 1).ActualCost, 1e-9)
|
||||
require.Equal(t, 0.0, s.CalculateAudioCost("unknown", 1, cfg, 1).ActualCost)
|
||||
// nil config → defaults (realtime $0.10/min, tts $15/M, stt $0.36/hr)
|
||||
require.InDelta(t, 0.10, s.CalculateAudioCost("realtime", 1, nil, 1).ActualCost, 1e-9)
|
||||
// nil config -> official defaults (think-fast-1 $0.05/min, TTS $15/M, REST STT $0.10/hr)
|
||||
require.InDelta(t, 0.05, s.CalculateAudioCost("realtime", 1, nil, 1).ActualCost, 1e-9)
|
||||
require.InDelta(t, 15.0, s.CalculateAudioCost("tts", 1, nil, 1).ActualCost, 1e-9)
|
||||
require.InDelta(t, 0.36, s.CalculateAudioCost("stt", 1, nil, 1).ActualCost, 1e-9)
|
||||
require.InDelta(t, 0.10, s.CalculateAudioCost("stt", 1, nil, 1).ActualCost, 1e-9)
|
||||
// explicit 0 → free
|
||||
zero := 0.0
|
||||
require.Equal(t, 0.0, s.CalculateAudioCost("realtime", 1, &audioPriceConfig{RealtimePerMin: &zero}, 1).ActualCost)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
)
|
||||
|
||||
// APIKeyRateLimitCacheData holds rate limit usage data cached in Redis.
|
||||
@@ -104,6 +105,7 @@ type ModelPricing struct {
|
||||
CacheCreation1hPrice float64 // 1小时缓存创建每token价格 (USD)
|
||||
SupportsCacheBreakdown bool // 是否支持详细的缓存分类
|
||||
LongContextInputThreshold int // 超过阈值后按整次会话提升输入价格
|
||||
LongContextThresholdInclusive bool // 达到阈值即应用(xAI);默认保持严格大于以兼容既有模型
|
||||
LongContextInputMultiplier float64 // 长上下文整次会话输入倍率
|
||||
LongContextOutputMultiplier float64 // 长上下文整次会话输出倍率
|
||||
ImageOutputPricePerToken float64 // 图片输出 token 价格 (USD)
|
||||
@@ -585,32 +587,56 @@ func (s *BillingService) initFallbackPricing() {
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
// xAI Grok 4.5 (official docs: $2 input / $0.50 cached input / $6 output per MTok)
|
||||
// xAI Grok 4.5: $2 input / $0.30 cached input / $6 output below 200k.
|
||||
s.fallbackPrices["grok-4.5"] = &ModelPricing{
|
||||
InputPricePerToken: 2e-6,
|
||||
OutputPricePerToken: 6e-6,
|
||||
CacheReadPricePerToken: 0.5e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
InputPricePerToken: 2e-6,
|
||||
OutputPricePerToken: 6e-6,
|
||||
CacheReadPricePerToken: 0.3e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
LongContextInputThreshold: 200000,
|
||||
LongContextThresholdInclusive: true,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 2,
|
||||
}
|
||||
|
||||
// xAI Grok 4.3 (official docs: $1.25 input / $2.50 output per MTok)
|
||||
// xAI Grok 4.6 (docs.x.ai/developers/models: $2 input / $0.50 cached input /
|
||||
// $6 output per MTok under 200k prompt tokens; ≥200k is 2× on input,
|
||||
// cached input, and output).
|
||||
s.fallbackPrices["grok-4.6"] = &ModelPricing{
|
||||
InputPricePerToken: 2e-6,
|
||||
OutputPricePerToken: 6e-6,
|
||||
CacheReadPricePerToken: 0.5e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
LongContextInputThreshold: 200000,
|
||||
LongContextThresholdInclusive: true,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 2,
|
||||
}
|
||||
|
||||
// xAI Grok 4.3: $1.25 input / $0.20 cached / $2.50 output below 200k.
|
||||
s.fallbackPrices["grok-4.3"] = &ModelPricing{
|
||||
InputPricePerToken: 1.25e-6,
|
||||
OutputPricePerToken: 2.5e-6,
|
||||
CacheReadPricePerToken: 0.2e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
LongContextInputThreshold: 1000000,
|
||||
LongContextInputMultiplier: 1,
|
||||
InputPricePerToken: 1.25e-6,
|
||||
OutputPricePerToken: 2.5e-6,
|
||||
CacheReadPricePerToken: 0.2e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
LongContextInputThreshold: 200000,
|
||||
LongContextThresholdInclusive: true,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 2,
|
||||
}
|
||||
// xAI Grok Build 0.1 (official docs: $1 input / $0.20 cached input /
|
||||
// $2 output per MTok). Composer is available only through Grok Build and
|
||||
// has no standalone public API rate card, so its aliases use this coding
|
||||
// model rate instead of silently billing at zero.
|
||||
s.fallbackPrices["grok-build-0.1"] = &ModelPricing{
|
||||
InputPricePerToken: 1e-6,
|
||||
OutputPricePerToken: 2e-6,
|
||||
CacheReadPricePerToken: 0.2e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
InputPricePerToken: 1e-6,
|
||||
OutputPricePerToken: 2e-6,
|
||||
CacheReadPricePerToken: 0.2e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
LongContextInputThreshold: 200000,
|
||||
LongContextThresholdInclusive: true,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 2,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -803,8 +829,10 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
|
||||
}
|
||||
|
||||
switch modelLower {
|
||||
case "grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest":
|
||||
case "grok", "grok-latest", "grok-4.5", "grok-4.5-latest":
|
||||
return s.fallbackPrices["grok-4.5"]
|
||||
case "grok-4.6", "grok-4.6-latest":
|
||||
return s.fallbackPrices["grok-4.6"]
|
||||
case "grok-4.3",
|
||||
"grok-4.20-0309-reasoning",
|
||||
"grok-4.20-0309-non-reasoning",
|
||||
@@ -812,13 +840,83 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
|
||||
"grok-4.20-reasoning",
|
||||
"grok-4.20-non-reasoning":
|
||||
return s.fallbackPrices["grok-4.3"]
|
||||
case "grok-build", "grok-build-0.1", "grok-composer", "grok-composer-2.5-fast", "composer-2.5":
|
||||
case "grok-build", "grok-build-latest", "grok-build-0.1", "grok-composer", "grok-composer-2.5-fast", "composer-2.5":
|
||||
return s.fallbackPrices["grok-build-0.1"]
|
||||
}
|
||||
|
||||
// Unknown Grok text IDs (grok-5, dated snapshots, provider-prefixed) inherit
|
||||
// the current default text card so a new model cannot ship unbilled.
|
||||
if pricing := s.grokUnknownTextFamilyFallback(modelLower); pricing != nil {
|
||||
return pricing
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *BillingService) grokUnknownTextFamilyFallback(model string) *ModelPricing {
|
||||
if s == nil || !isGrokUnknownTextFamilyModel(model) {
|
||||
return nil
|
||||
}
|
||||
return s.fallbackPrices["grok-4.5"]
|
||||
}
|
||||
|
||||
func isGrokUnknownTextFamilyModel(model string) bool {
|
||||
native := strings.ToLower(strings.TrimSpace(xai.StripGrokProviderPrefix(model)))
|
||||
if isGrokMediaFamilyModel(native) {
|
||||
return false
|
||||
}
|
||||
switch {
|
||||
case native == "grok", native == "grok-latest":
|
||||
return true
|
||||
case strings.HasPrefix(native, "grok-build"),
|
||||
strings.HasPrefix(native, "grok-composer"),
|
||||
strings.HasPrefix(native, "composer-"):
|
||||
return true
|
||||
case len(native) > 5 && strings.HasPrefix(native, "grok-"):
|
||||
rest := native[len("grok-"):]
|
||||
return rest[0] >= '0' && rest[0] <= '9'
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// isGrokMediaFamilyModel matches ids that are billed per image/video/audio unit
|
||||
// rather than per token, so version-numbered media ids (grok-2-image-1212,
|
||||
// grok-5-video) cannot slip into the unknown-text fallback and pick up a token
|
||||
// card. "vision" is deliberately absent: multimodal chat models are token billed.
|
||||
func isGrokMediaFamilyModel(native string) bool {
|
||||
for _, marker := range []string{"imagine", "image", "video", "audio", "speech", "tts", "transcribe", "realtime"} {
|
||||
if strings.Contains(native, marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// HasIdentifiedTokenPricing 判断模型能否在价格表中被"确定性识别"出 token 价格。
|
||||
//
|
||||
// 与 GetModelPricing 的关键区别:本函数拒绝按子串猜系列的兜底。GetModelPricing 会
|
||||
// 让任意含 "haiku"/"opus"/"claude" 的名字(哪怕是不存在的型号)落到 getFallbackPricing
|
||||
// 的系列兜底价上,因此凡是模型名来自外部、且"能查到价"会直接影响计费金额的场景
|
||||
// (如按上游响应自报模型计费),都必须用本函数而不是 GetModelPricing 做准入判断。
|
||||
func (s *BillingService) HasIdentifiedTokenPricing(model string) bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
model = strings.ToLower(strings.TrimSpace(model))
|
||||
if model == "" {
|
||||
return false
|
||||
}
|
||||
if s.pricingService != nil {
|
||||
// 仅有图片价的条目不能用于 token 计费,口径与 GetModelPricing 保持一致。
|
||||
if pricing := s.pricingService.GetIdentifiedModelPricing(model); pricing != nil && !pricing.TokenPricingAbsent {
|
||||
return true
|
||||
}
|
||||
}
|
||||
pricing, ok := s.fallbackPrices[model]
|
||||
return ok && pricing != nil
|
||||
}
|
||||
|
||||
// GetModelPricing 获取模型价格配置
|
||||
func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) {
|
||||
// 标准化模型名称(转小写)
|
||||
@@ -926,9 +1024,11 @@ type CostInput struct {
|
||||
Ctx context.Context
|
||||
Model string
|
||||
GroupID *int64 // 用于渠道定价查找
|
||||
Group *Group
|
||||
Tokens UsageTokens
|
||||
RequestCount int // 按次计费时使用
|
||||
SizeTier string // 按次/图片模式的层级标签("1K","2K","4K","HD" 等)
|
||||
RequestCount int // 按次计费时使用
|
||||
UsageUnits float64 // 音频等连续计量单位(分钟/小时/百万字符)
|
||||
SizeTier string // 按次/图片模式的层级标签("1K","2K","4K","HD" 等)
|
||||
RateMultiplier float64
|
||||
ServiceTier string // "priority","flex","" 等
|
||||
Resolver *ModelPricingResolver // 定价解析器
|
||||
@@ -961,6 +1061,7 @@ func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown,
|
||||
resolved = input.Resolver.Resolve(input.Ctx, PricingInput{
|
||||
Model: input.Model,
|
||||
GroupID: input.GroupID,
|
||||
Group: input.Group,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -972,7 +1073,7 @@ func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown,
|
||||
var breakdown *CostBreakdown
|
||||
var err error
|
||||
switch resolved.Mode {
|
||||
case BillingModePerRequest, BillingModeImage:
|
||||
case BillingModePerRequest, BillingModeImage, BillingModeVideo:
|
||||
breakdown, err = s.calculatePerRequestCost(resolved, input)
|
||||
default: // BillingModeToken
|
||||
breakdown, err = s.calculateTokenCost(resolved, input)
|
||||
@@ -998,7 +1099,7 @@ func (s *BillingService) calculateTokenCost(resolved *ResolvedPricing, input Cos
|
||||
pricing = s.applyModelSpecificPricingPolicy(input.Model, pricing)
|
||||
|
||||
// 长上下文定价仅在无区间定价时应用(区间定价已包含上下文分层)
|
||||
applyLongCtx := len(resolved.Intervals) == 0
|
||||
applyLongCtx := len(resolved.Intervals) == 0 && resolved.longContextPricingEnabled
|
||||
if input.LongContextBillingEnabled != nil {
|
||||
applyLongCtx = applyLongCtx && *input.LongContextBillingEnabled
|
||||
}
|
||||
@@ -1133,9 +1234,13 @@ func (s *BillingService) computeCacheCreationCost(pricing *ModelPricing, tokens
|
||||
|
||||
// calculatePerRequestCost 按次/图片计费
|
||||
func (s *BillingService) calculatePerRequestCost(resolved *ResolvedPricing, input CostInput) (*CostBreakdown, error) {
|
||||
count := input.RequestCount
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
units := input.UsageUnits
|
||||
if units <= 0 {
|
||||
count := input.RequestCount
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
}
|
||||
units = float64(count)
|
||||
}
|
||||
|
||||
var unitPrice float64
|
||||
@@ -1154,7 +1259,7 @@ func (s *BillingService) calculatePerRequestCost(resolved *ResolvedPricing, inpu
|
||||
unitPrice = resolved.DefaultPerRequestPrice
|
||||
}
|
||||
|
||||
totalCost := unitPrice * float64(count)
|
||||
totalCost := unitPrice * units
|
||||
actualCost := totalCost * input.RateMultiplier
|
||||
|
||||
return &CostBreakdown{
|
||||
@@ -1256,6 +1361,9 @@ func (s *BillingService) shouldApplySessionLongContextPricing(tokens UsageTokens
|
||||
return false
|
||||
}
|
||||
totalInputTokens := tokens.InputTokens + tokens.CacheCreationTokens + tokens.CacheReadTokens
|
||||
if pricing.LongContextThresholdInclusive {
|
||||
return totalInputTokens >= pricing.LongContextInputThreshold
|
||||
}
|
||||
return totalInputTokens > pricing.LongContextInputThreshold
|
||||
}
|
||||
|
||||
@@ -1427,6 +1535,8 @@ const (
|
||||
defaultGrokImagineImagePrice2K = 0.02
|
||||
defaultGrokImagineImageQualityPrice1K = 0.05
|
||||
defaultGrokImagineImageQualityPrice2K = 0.07
|
||||
defaultGrokImagineImage20Price1K = 0.06 // default quality is Medium
|
||||
defaultGrokImagineImage20Price2K = 0.08
|
||||
|
||||
// 视频默认价为 xAI 官方**每秒**输出价格(USD/s),总价 = 每秒价 × 时长(秒)。
|
||||
defaultGrokImagineVideoPrice480P = 0.05
|
||||
@@ -1438,14 +1548,14 @@ const (
|
||||
// Codex alpha/search 网页搜索单次默认价:OpenAI 官方 web search 定价 $10/1000 次。
|
||||
defaultWebSearchPricePerCall = 0.01
|
||||
|
||||
// Grok /v1/web_search 与 SearchCount 附加费:与 Codex 对齐 $10/1000 次(按 1k 计价字段存储)。
|
||||
defaultSearchPricePer1k = 10.0
|
||||
// xAI server-side web/X search and code execution are $5/1000 calls.
|
||||
defaultSearchPricePer1k = 5.0
|
||||
|
||||
// Grok Voice 默认价(分组列 NULL 时使用;显式配 0 表示免费)。
|
||||
// 保守运营占位,运维可通过 groups.audio_* 覆盖。
|
||||
defaultAudioRealtimePricePerMin = 0.10
|
||||
// Generic realtime defaults to think-fast-1.0; think-fast-2.0 can be
|
||||
// configured independently through per-model group/channel pricing.
|
||||
defaultAudioRealtimePricePerMin = 0.05
|
||||
defaultAudioTTSPricePerMillionChars = 15.0
|
||||
defaultAudioSTTPricePerHour = 0.36
|
||||
defaultAudioSTTPricePerHour = 0.10
|
||||
)
|
||||
|
||||
// CalculateWebSearchCost 计算 Codex alpha/search 网页搜索按次费用。
|
||||
@@ -1703,6 +1813,12 @@ func (s *BillingService) getDefaultVideoPrice(model string, resolution string) f
|
||||
func getDefaultGrokImagineImagePrice(model string, imageSize string) (float64, bool) {
|
||||
model = strings.ToLower(strings.TrimSpace(model))
|
||||
switch model {
|
||||
case "grok-imagine-image-2.0":
|
||||
return getGrokImagineImageTierPrice(
|
||||
imageSize,
|
||||
defaultGrokImagineImage20Price1K,
|
||||
defaultGrokImagineImage20Price2K,
|
||||
), true
|
||||
case "grok-imagine-image-quality":
|
||||
return getGrokImagineImageTierPrice(
|
||||
imageSize,
|
||||
|
||||
@@ -871,8 +871,8 @@ func TestComputeTokenBreakdown_GptImage2ImageEditIssue4386(t *testing.T) {
|
||||
|
||||
cost := svc.computeTokenBreakdown(pricing, tokens, 1.0, "", false)
|
||||
|
||||
wantTextInput := float64(19) * 5e-6 // 0.000095
|
||||
wantImageInput := float64(352) * 8e-6 // 0.002816
|
||||
wantTextInput := float64(19) * 5e-6 // 0.000095
|
||||
wantImageInput := float64(352) * 8e-6 // 0.002816
|
||||
wantImageOutput := float64(439) * 30e-6 // 0.013170
|
||||
require.InDelta(t, wantTextInput, cost.InputCost, 1e-15, "InputCost 仅含文本输入")
|
||||
require.InDelta(t, wantImageInput, cost.ImageInputCost, 1e-15, "图片输入按 $8/1M 独立计费")
|
||||
@@ -1153,7 +1153,23 @@ func TestCalculateCostWithLongContext_PropagatesError(t *testing.T) {
|
||||
func TestGetModelPricing_Grok45OfficialFallback(t *testing.T) {
|
||||
svc := newTestBillingService()
|
||||
|
||||
for _, model := range []string{"grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest"} {
|
||||
for _, model := range []string{"grok", "grok-latest", "grok-4.5", "grok-4.5-latest"} {
|
||||
model := model
|
||||
t.Run(model, func(t *testing.T) {
|
||||
pricing, err := svc.GetModelPricing(model)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 2e-6, pricing.InputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 6e-6, pricing.OutputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 0.3e-6, pricing.CacheReadPricePerToken, 1e-12)
|
||||
require.False(t, pricing.SupportsCacheBreakdown)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetModelPricing_Grok46OfficialFallback(t *testing.T) {
|
||||
svc := newTestBillingService()
|
||||
|
||||
for _, model := range []string{"grok-4.6", "grok-4.6-latest"} {
|
||||
model := model
|
||||
t.Run(model, func(t *testing.T) {
|
||||
pricing, err := svc.GetModelPricing(model)
|
||||
@@ -1161,11 +1177,77 @@ func TestGetModelPricing_Grok45OfficialFallback(t *testing.T) {
|
||||
require.InDelta(t, 2e-6, pricing.InputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 6e-6, pricing.OutputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 0.5e-6, pricing.CacheReadPricePerToken, 1e-12)
|
||||
require.Equal(t, 200000, pricing.LongContextInputThreshold)
|
||||
require.InDelta(t, 2.0, pricing.LongContextInputMultiplier, 1e-12)
|
||||
require.InDelta(t, 2.0, pricing.LongContextOutputMultiplier, 1e-12)
|
||||
require.False(t, pricing.SupportsCacheBreakdown)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculateCostUnified_GroupLongContextToggleUsesPresetLadder(t *testing.T) {
|
||||
svc := newTestBillingService()
|
||||
resolver := NewModelPricingResolver(nil, svc)
|
||||
tokens := UsageTokens{InputTokens: 250000, OutputTokens: 1000}
|
||||
|
||||
off := &Group{LongContextPricingEnabled: false}
|
||||
disabled, err := svc.CalculateCostUnified(CostInput{
|
||||
Model: "grok-4.5", Group: off, Tokens: tokens, RateMultiplier: 1, Resolver: resolver,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
on := &Group{LongContextPricingEnabled: true}
|
||||
enabled, err := svc.CalculateCostUnified(CostInput{
|
||||
Model: "grok-4.5", Group: on, Tokens: tokens, RateMultiplier: 1, Resolver: resolver,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.False(t, disabled.LongContextBillingApplied)
|
||||
require.True(t, enabled.LongContextBillingApplied)
|
||||
require.InDelta(t, disabled.InputCost*2, enabled.InputCost, 1e-12)
|
||||
require.InDelta(t, disabled.OutputCost*2, enabled.OutputCost, 1e-12)
|
||||
}
|
||||
|
||||
func TestGetModelPricing_UnknownGrokTextFallsBackToGrok45(t *testing.T) {
|
||||
svc := newTestBillingService()
|
||||
baseline, err := svc.GetModelPricing("grok-4.5")
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, model := range []string{"grok-5", "grok-5-latest", "x-ai/grok-7", "grok-4.7-beta"} {
|
||||
pricing, err := svc.GetModelPricing(model)
|
||||
require.NoError(t, err, "model %s", model)
|
||||
require.InDelta(t, baseline.InputPricePerToken, pricing.InputPricePerToken, 1e-12, model)
|
||||
require.InDelta(t, baseline.OutputPricePerToken, pricing.OutputPricePerToken, 1e-12, model)
|
||||
require.InDelta(t, baseline.CacheReadPricePerToken, pricing.CacheReadPricePerToken, 1e-12, model)
|
||||
}
|
||||
|
||||
// Per-unit media ids must not inherit the text card just because they carry
|
||||
// a version number; they are billed by the image/video/audio paths instead.
|
||||
for _, model := range []string{"grok-2-image-1212", "grok-2-audio", "grok-5-video", "x-ai/grok-6-image"} {
|
||||
require.False(t, isGrokUnknownTextFamilyModel(model), "model %s", model)
|
||||
}
|
||||
// Multimodal chat models stay token billed.
|
||||
require.True(t, isGrokUnknownTextFamilyModel("grok-2-vision-1212"))
|
||||
|
||||
for _, model := range []string{
|
||||
"grok-imagine-image-3.0",
|
||||
"grok-imagine-video-2",
|
||||
"grok-voice-latest",
|
||||
"grok-web-search",
|
||||
"grok-x-search",
|
||||
"grok-speech-1",
|
||||
} {
|
||||
_, err := svc.GetModelPricing(model)
|
||||
require.Error(t, err, "non-text grok family %s must not inherit grok-4.5 token rates", model)
|
||||
require.ErrorIs(t, err, ErrModelPricingUnavailable)
|
||||
}
|
||||
|
||||
// Known cards stay on their own rate, not the 4.5 family floor.
|
||||
build, err := svc.GetModelPricing("grok-build-0.1")
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 1e-6, build.InputPricePerToken, 1e-12)
|
||||
}
|
||||
|
||||
func TestGetModelPricing_GrokCatalogFallbacks(t *testing.T) {
|
||||
svc := newTestBillingService()
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ const (
|
||||
// IsValid 检查 BillingMode 是否为合法值
|
||||
func (m BillingMode) IsValid() bool {
|
||||
switch m {
|
||||
case BillingModeToken, BillingModePerRequest, BillingModeImage, "":
|
||||
case BillingModeToken, BillingModePerRequest, BillingModeImage, BillingModeVideo, "":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
@@ -39,6 +39,10 @@ const (
|
||||
BillingModelSourceRequested = "requested"
|
||||
BillingModelSourceUpstream = "upstream"
|
||||
BillingModelSourceChannelMapped = "channel_mapped"
|
||||
// BillingModelSourceResponse bills by a trusted model declaration observed
|
||||
// in the successful upstream response. It is deliberately distinct from
|
||||
// "upstream", which means the model sent to the provider.
|
||||
BillingModelSourceResponse = "response_model"
|
||||
)
|
||||
|
||||
// Channel 渠道实体
|
||||
@@ -47,7 +51,7 @@ type Channel struct {
|
||||
Name string
|
||||
Description string
|
||||
Status string
|
||||
BillingModelSource string // "requested", "upstream", or "channel_mapped"
|
||||
BillingModelSource string // "requested", "upstream", "channel_mapped", or "response_model"
|
||||
RestrictModels bool // 是否限制模型(仅允许定价列表中的模型)
|
||||
Features string // 渠道特性描述(JSON 数组),用于支付页面展示
|
||||
FeaturesConfig map[string]any // 渠道功能配置(如 web search emulation)
|
||||
@@ -83,38 +87,38 @@ type AccountStatsPricingRule struct {
|
||||
|
||||
// ChannelModelPricing 渠道模型定价条目
|
||||
type ChannelModelPricing struct {
|
||||
ID int64
|
||||
ChannelID int64
|
||||
Platform string // 所属平台(anthropic/openai/gemini/...)
|
||||
Models []string // 绑定的模型列表
|
||||
BillingMode BillingMode // 计费模式
|
||||
InputPrice *float64 // 每 token 输入价格(USD)— 向后兼容 flat 定价
|
||||
OutputPrice *float64 // 每 token 输出价格(USD)
|
||||
CacheWritePrice *float64 // 缓存写入价格
|
||||
CacheReadPrice *float64 // 缓存读取价格
|
||||
ImageInputPrice *float64 // 图片输入 token 价格(如 gpt-image-2 图片编辑);未配置时回退文本输入价
|
||||
ImageOutputPrice *float64 // 图片输出价格(向后兼容)
|
||||
PerRequestPrice *float64 // 默认按次计费价格(USD)
|
||||
Intervals []PricingInterval // 区间定价列表
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
ID int64 `json:"id,omitempty"`
|
||||
ChannelID int64 `json:"channel_id,omitempty"`
|
||||
Platform string `json:"platform"` // 所属平台(anthropic/openai/gemini/...)
|
||||
Models []string `json:"models"`
|
||||
BillingMode BillingMode `json:"billing_mode"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
ImageInputPrice *float64 `json:"image_input_price"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
Intervals []PricingInterval `json:"intervals"`
|
||||
CreatedAt time.Time `json:"created_at,omitempty"`
|
||||
UpdatedAt time.Time `json:"updated_at,omitempty"`
|
||||
}
|
||||
|
||||
// PricingInterval 定价区间(token 区间 / 按次分层 / 图片分辨率分层)
|
||||
type PricingInterval struct {
|
||||
ID int64
|
||||
PricingID int64
|
||||
MinTokens int // 区间下界(含)
|
||||
MaxTokens *int // 区间上界(不含),nil = 无上限
|
||||
TierLabel string // 层级标签(按次/图片模式:1K, 2K, 4K, HD 等)
|
||||
InputPrice *float64 // token 模式:每 token 输入价
|
||||
OutputPrice *float64 // token 模式:每 token 输出价
|
||||
CacheWritePrice *float64 // token 模式:缓存写入价
|
||||
CacheReadPrice *float64 // token 模式:缓存读取价
|
||||
PerRequestPrice *float64 // 按次/图片模式:每次请求价格
|
||||
SortOrder int
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
ID int64 `json:"id,omitempty"`
|
||||
PricingID int64 `json:"pricing_id,omitempty"`
|
||||
MinTokens int `json:"min_tokens"`
|
||||
MaxTokens *int `json:"max_tokens"`
|
||||
TierLabel string `json:"tier_label"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
CreatedAt time.Time `json:"created_at,omitempty"`
|
||||
UpdatedAt time.Time `json:"updated_at,omitempty"`
|
||||
}
|
||||
|
||||
// IsActive 判断渠道是否启用
|
||||
@@ -311,7 +315,7 @@ func ValidateIntervals(intervals []PricingInterval, mode BillingMode) error {
|
||||
}
|
||||
|
||||
// per_request / image 模式按 tier_label 匹配,不做 token 区间重叠校验
|
||||
if mode == BillingModePerRequest || mode == BillingModeImage {
|
||||
if mode == BillingModePerRequest || mode == BillingModeImage || mode == BillingModeVideo {
|
||||
return nil
|
||||
}
|
||||
return validateIntervalOverlap(sorted)
|
||||
|
||||
@@ -108,7 +108,7 @@ type ChannelMappingResult struct {
|
||||
MappedModel string // 映射后的模型名(无映射时等于原始模型名)
|
||||
ChannelID int64 // 渠道 ID(0 = 无渠道关联)
|
||||
Mapped bool // 是否发生了映射
|
||||
BillingModelSource string // 计费模型来源("requested" / "upstream" / "channel_mapped")
|
||||
BillingModelSource string // 计费模型来源("requested" / "upstream" / "channel_mapped" / "response_model")
|
||||
}
|
||||
|
||||
// BuildModelMappingChain 根据映射结果和上游实际模型构建映射链描述。
|
||||
@@ -375,6 +375,13 @@ func channelLookupPlatform(ctx context.Context, groupPlatform string) string {
|
||||
}
|
||||
return groupPlatform
|
||||
}
|
||||
|
||||
// InvalidateCache 失效并重建渠道缓存。
|
||||
// 供渠道以外、但会影响渠道缓存内容的变更调用(如分组平台变更)。
|
||||
func (s *ChannelService) InvalidateCache() {
|
||||
s.invalidateCache()
|
||||
}
|
||||
|
||||
func (s *ChannelService) invalidateCache() {
|
||||
s.cache.Store((*channelCache)(nil))
|
||||
s.cacheSF.Forget("channel_cache")
|
||||
@@ -659,7 +666,7 @@ func validatePricingBillingMode(pricing []ChannelModelPricing) error {
|
||||
}
|
||||
|
||||
func checkBillingModeRequirements(p ChannelModelPricing) error {
|
||||
if p.BillingMode == BillingModePerRequest || p.BillingMode == BillingModeImage {
|
||||
if p.BillingMode == BillingModePerRequest || p.BillingMode == BillingModeImage || p.BillingMode == BillingModeVideo {
|
||||
if p.PerRequestPrice == nil && len(p.Intervals) == 0 {
|
||||
return infraerrors.BadRequest(
|
||||
"BILLING_MODE_MISSING_PRICE",
|
||||
@@ -962,19 +969,35 @@ func conflictsBetween(a, b modelEntry) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// toModelEntry 将模型名转换为 modelEntry
|
||||
// toModelEntry 将模型名转换为 modelEntry(用于模型映射的冲突检测)。
|
||||
// 归一化必须与 expandMappingToCache 写缓存键的方式一致:映射缓存只做 strings.ToLower。
|
||||
func toModelEntry(pattern string) modelEntry {
|
||||
prefix, isWild := splitWildcardSuffix(strings.ToLower(pattern))
|
||||
return modelEntry{pattern: pattern, prefix: prefix, wildcard: isWild}
|
||||
}
|
||||
|
||||
// toPricingModelEntry 将模型名转换为 modelEntry(用于模型定价的冲突检测)。
|
||||
//
|
||||
// 与 toModelEntry 的区别:定价缓存的键走 normalizeChannelPricingModelName
|
||||
// (额外做 TrimSpace,并把 claude-* 的 "." 换成 "-"),冲突检测必须用同一套归一化,
|
||||
// 否则两个校验时看着不同、写进缓存后键相同的定价会互相静默覆盖。
|
||||
func toPricingModelEntry(pattern string) modelEntry {
|
||||
// 先剥通配符再归一化,与 expandPricingToCache 的处理顺序保持一致
|
||||
prefix, isWild := splitWildcardSuffix(pattern)
|
||||
return modelEntry{
|
||||
pattern: pattern,
|
||||
prefix: normalizeChannelPricingModelName(prefix),
|
||||
wildcard: isWild,
|
||||
}
|
||||
}
|
||||
|
||||
// validateNoConflictingModels 检查定价列表中是否有冲突模型模式(同一平台下)。
|
||||
// 冲突包括:精确重复、通配符之间的前缀包含、通配符与精确名的前缀匹配。
|
||||
func validateNoConflictingModels(pricingList []ChannelModelPricing) error {
|
||||
byPlatform := make(map[string][]modelEntry)
|
||||
for _, p := range pricingList {
|
||||
for _, model := range p.Models {
|
||||
byPlatform[p.Platform] = append(byPlatform[p.Platform], toModelEntry(model))
|
||||
byPlatform[p.Platform] = append(byPlatform[p.Platform], toPricingModelEntry(model))
|
||||
}
|
||||
}
|
||||
for platform, entries := range byPlatform {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user