mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-19 16:17:56 +00:00
Compare commits
132
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8f6f459835 | ||
|
|
58ccea4eaa | ||
|
|
c253bd2c72 | ||
|
|
26cb59df05 | ||
|
|
58ea46e894 | ||
|
|
685455985b | ||
|
|
1ed3b6aef6 | ||
|
|
f211a630c8 | ||
|
|
37732dcd34 | ||
|
|
a341239596 | ||
|
|
1ba92449c7 | ||
|
|
9617775f9a | ||
|
|
fd42d3722d | ||
|
|
1870b58c1d | ||
|
|
269fbcac0e | ||
|
|
7e579cb28d | ||
|
|
938f1868ae | ||
|
|
4d19836189 | ||
|
|
c792b44bb2 | ||
|
|
1ea4150bf0 | ||
|
|
ed2da82396 | ||
|
|
6259940ef2 | ||
|
|
baaf59d417 | ||
|
|
c0325d24f9 | ||
|
|
c76ff55997 | ||
|
|
5c12960285 | ||
|
|
f6d6613ec6 | ||
|
|
5dfac8bffa | ||
|
|
632b4f4faa | ||
|
|
3d9a0a71bd | ||
|
|
1a3ecd2b94 | ||
|
|
8532344745 | ||
|
|
cddb03c0f1 | ||
|
|
a7a321232f | ||
|
|
16e4f7ecc3 | ||
|
|
bb6c3b4f6a | ||
|
|
ab0fcd1a0e | ||
|
|
8869775ed3 | ||
|
|
f66a42c40e | ||
|
|
7d633f5fc3 | ||
|
|
b2d1c3859a | ||
|
|
5253bb72b0 | ||
|
|
6bf335965a | ||
|
|
c46d07ca07 | ||
|
|
ab28a2d10b | ||
|
|
6793d5ac85 | ||
|
|
9f24a55305 | ||
|
|
c38c5beef8 | ||
|
|
10c8b70203 | ||
|
|
7cdca9e495 | ||
|
|
971544570d | ||
|
|
3c3bb2fa19 | ||
|
|
e330c243a8 | ||
|
|
8d82bb069c | ||
|
|
b8642ef674 | ||
|
|
76b70b1685 | ||
|
|
8a0a76aeb5 | ||
|
|
cb5e03a720 | ||
|
|
9aac3b73f7 | ||
|
|
e728545382 | ||
|
|
4b667ccd45 | ||
|
|
1977810cf2 | ||
|
|
79c2eb5020 | ||
|
|
5f19433103 | ||
|
|
396a9d1130 | ||
|
|
51016bb03a | ||
|
|
22fc0cdbf5 | ||
|
|
cb7841d85c | ||
|
|
e8ff2017c0 | ||
|
|
11e1e22882 | ||
|
|
baeac1f3de | ||
|
|
901a0439f1 | ||
|
|
073e92d171 | ||
|
|
d677d67dda | ||
|
|
4d9fedee20 | ||
|
|
fce41e318f | ||
|
|
8ae6d8f67e | ||
|
|
8219dcfc87 | ||
|
|
1d3b9665c8 | ||
|
|
44ef88f659 | ||
|
|
c204d33b09 | ||
|
|
45dcce0e49 | ||
|
|
89d826be29 | ||
|
|
cb7b03795d | ||
|
|
a8b9ea22b7 | ||
|
|
9662cff2e7 | ||
|
|
674570ca17 | ||
|
|
a288bab73a | ||
|
|
fbfdcef818 | ||
|
|
0e82efe489 | ||
|
|
e803e3851c | ||
|
|
5912ae960c | ||
|
|
e77b4b260e | ||
|
|
e215c98c2c | ||
|
|
e29b93a1fb | ||
|
|
fd82dfd52d | ||
|
|
4853662f99 | ||
|
|
0fa577a19c | ||
|
|
b61e4bcc45 | ||
|
|
0ae151a238 | ||
|
|
b830bc14d6 | ||
|
|
8c4c3c09ce | ||
|
|
678eb22a40 | ||
|
|
c4d883b8da | ||
|
|
0de6d7e9ba | ||
|
|
363cc4994b | ||
|
|
f3d9491071 | ||
|
|
76a13a5a8d | ||
|
|
69648476d4 | ||
|
|
a04ce49016 | ||
|
|
bb9e74285e | ||
|
|
bba6a55e0f | ||
|
|
5935e674a8 | ||
|
|
ef4f99f292 | ||
|
|
814ecfba7c | ||
|
|
bd404c16f2 | ||
|
|
c91fbeb0c0 | ||
|
|
c3fc4331c1 | ||
|
|
e087f1b736 | ||
|
|
145b0ac35e | ||
|
|
5ea03c178d | ||
|
|
1afb8264e5 | ||
|
|
2321c0e771 | ||
|
|
0b35370a7a | ||
|
|
9c36b75a7d | ||
|
|
fd9ce53281 | ||
|
|
0d5e3ca9b0 | ||
|
|
35e8ba2a3e | ||
|
|
a6d868f27b | ||
|
|
7d796f1119 | ||
|
|
3bff4b64b1 | ||
|
|
5e72deb7d4 |
@@ -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: |
|
||||
|
||||
+1
-1
@@ -53,7 +53,7 @@ npm install -g pnpm
|
||||
|
||||
### CI 要求
|
||||
|
||||
- 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 会在版本校验步骤直接失败。
|
||||
- Go 版本必须是 **1.26.6**:三个 workflow 都用 `go-version-file: backend/go.mod` 取版本,随后硬断言 `go version | grep -q 'go1.26.6'`。升级 Go 时要同时改 `backend/go.mod`、`backend-ci.yml`(两处)、`release.yml`、`security-scan.yml` 里的这句断言,**以及三个 Dockerfile 里的 Go 构建镜像**(`Dockerfile` / `deploy/Dockerfile` 的 `ARG GOLANG_IMAGE`、`backend/Dockerfile` 的 `FROM golang:`)。前者漏了 CI 会在版本校验步骤直接失败;**后者漏了 CI 不会报,而是等到有人用这些 Dockerfile 构建时才失败**(`go.mod requires go >= X (running Y; GOTOOLCHAIN=local)`)。
|
||||
- 前端使用 `pnpm install --frozen-lockfile`,必须提交 `pnpm-lock.yaml`
|
||||
|
||||
### 本地测试命令
|
||||
|
||||
+1
-1
@@ -8,7 +8,7 @@
|
||||
# =============================================================================
|
||||
|
||||
ARG NODE_IMAGE=node:24-alpine
|
||||
ARG GOLANG_IMAGE=golang:1.26.5-alpine
|
||||
ARG GOLANG_IMAGE=golang:1.26.6-alpine
|
||||
ARG ALPINE_IMAGE=alpine:3.21
|
||||
ARG POSTGRES_IMAGE=postgres:18-alpine
|
||||
ARG GOPROXY=https://goproxy.cn,direct
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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>
|
||||
|
||||
## 项目概述
|
||||
|
||||
@@ -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>
|
||||
|
||||
## 概要
|
||||
|
||||
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
@@ -1,4 +1,4 @@
|
||||
FROM golang:1.26.5-alpine
|
||||
FROM golang:1.26.6-alpine
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.1.173
|
||||
0.1.177
|
||||
|
||||
@@ -88,6 +88,7 @@ func provideCleanup(
|
||||
schedulerSnapshot *service.SchedulerSnapshotService,
|
||||
tokenRefresh *service.TokenRefreshService,
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
cnProviderBalanceCheck *service.CNProviderBalanceCheckService,
|
||||
codexVersionSync *service.OpenAICodexVersionSyncService,
|
||||
proxyExpiry *service.ProxyExpiryService,
|
||||
subscriptionExpiry *service.SubscriptionExpiryService,
|
||||
@@ -238,6 +239,12 @@ func provideCleanup(
|
||||
accountExpiry.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"CNProviderBalanceCheckService", func() error {
|
||||
if cnProviderBalanceCheck != nil {
|
||||
cnProviderBalanceCheck.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpenAICodexVersionSyncService", func() error {
|
||||
codexVersionSync.Stop()
|
||||
return nil
|
||||
|
||||
@@ -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)
|
||||
@@ -218,6 +218,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
|
||||
tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, compositeTokenCacheInvalidator, schedulerCache, configConfig, tempUnschedCache, privacyClientFactory, proxyRepository, oAuthRefreshAPI, openAIGatewayService)
|
||||
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService, tokenRefreshService)
|
||||
cnProviderQuotaService := service.ProvideCNProviderQuotaService(accountRepository, proxyRepository, httpUpstream, configConfig)
|
||||
cnProviderBalanceService := service.ProvideCNProviderBalanceService(accountRepository, proxyRepository, httpUpstream, configConfig)
|
||||
cnProviderHandler := admin.NewCNProviderHandler(cnProviderQuotaService, cnProviderBalanceService)
|
||||
proxyHandler := admin.NewProxyHandler(adminService)
|
||||
adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)
|
||||
promoHandler := admin.NewPromoHandler(promoService)
|
||||
@@ -277,7 +280,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
ollamaCloudUsageService := service.ProvideOllamaCloudUsageService(accountRepository, httpUpstream, settingService, secretEncryptor, configConfig, leaderLockCache, db)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
@@ -327,6 +330,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig)
|
||||
opsIngressRejectAggregator := service.ProvideOpsIngressRejectAggregator(opsRepository, opsService)
|
||||
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
|
||||
cnProviderBalanceCheckService := service.ProvideCNProviderBalanceCheckService(accountRepository, cnProviderBalanceService, cnProviderQuotaService, configConfig)
|
||||
openAICodexVersionSyncService := service.ProvideOpenAICodexVersionSyncService(settingRepository, settingService, gitHubReleaseClient)
|
||||
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
|
||||
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
|
||||
@@ -336,7 +340,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService)
|
||||
channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService)
|
||||
userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, openAICodexVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, auditLogService, promptService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, auditLogService, promptService)
|
||||
application := &Application{
|
||||
Server: httpServer,
|
||||
PromptAudit: promptService,
|
||||
@@ -380,6 +384,7 @@ func provideCleanup(
|
||||
schedulerSnapshot *service.SchedulerSnapshotService,
|
||||
tokenRefresh *service.TokenRefreshService,
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
cnProviderBalanceCheck *service.CNProviderBalanceCheckService,
|
||||
codexVersionSync *service.OpenAICodexVersionSyncService,
|
||||
proxyExpiry *service.ProxyExpiryService,
|
||||
subscriptionExpiry *service.SubscriptionExpiryService,
|
||||
@@ -529,6 +534,12 @@ func provideCleanup(
|
||||
accountExpiry.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"CNProviderBalanceCheckService", func() error {
|
||||
if cnProviderBalanceCheck != nil {
|
||||
cnProviderBalanceCheck.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpenAICodexVersionSyncService", func() error {
|
||||
codexVersionSync.Stop()
|
||||
return nil
|
||||
|
||||
@@ -66,6 +66,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
schedulerSnapshotSvc,
|
||||
tokenRefreshSvc,
|
||||
accountExpirySvc,
|
||||
nil, // cnProviderBalanceCheck
|
||||
codexVersionSyncSvc,
|
||||
proxyExpirySvc,
|
||||
subscriptionExpirySvc,
|
||||
|
||||
+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").
|
||||
|
||||
@@ -41,7 +41,8 @@ func (UserPlatformQuota) Fields() []ent.Field {
|
||||
// 注意:平台列表的单一权威源为 service.AllowedQuotaPlatforms;
|
||||
// 此处为 ent 构建期约束,需与 service.AllowedQuotaPlatforms 保持同步。
|
||||
switch s {
|
||||
case "anthropic", "openai", "gemini", "antigravity", "grok":
|
||||
case "anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek":
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("platform %q is not allowed", s)
|
||||
|
||||
+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
|
||||
|
||||
@@ -1025,6 +1025,10 @@ type GatewayConfig struct {
|
||||
|
||||
// Grok: Grok/xAI gateway scheduling and free-tier soft-gate settings.
|
||||
Grok GatewayGrokConfig `mapstructure:"grok"`
|
||||
|
||||
// CNProviders: 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)的余额检测配置。
|
||||
// 仅作用于 payg(按量付费)账号:周期探测余额,低于阈值则临时停调。
|
||||
CNProviders GatewayCNProvidersConfig `mapstructure:"cn_providers"`
|
||||
}
|
||||
|
||||
// GatewayGrokConfig holds Grok-specific gateway scheduling knobs.
|
||||
@@ -1057,6 +1061,18 @@ type GatewayGrokConfig struct {
|
||||
FreeQuotaStatsCacheSeconds int `mapstructure:"free_quota_stats_cache_seconds"`
|
||||
}
|
||||
|
||||
// GatewayCNProvidersConfig 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)的余额检测配置。
|
||||
//
|
||||
// 仅作用于 payg(按量付费)账号(kimi/deepseek 有公开余额端点;zhipu 无,仅靠响应式 429/402)。
|
||||
// - balance_check_enabled: 是否启用周期余额检测(默认 true)
|
||||
// - balance_threshold: 余额低于此值(账户货币单位,默认 0.5)触发临时停调
|
||||
// - balance_check_interval_minutes: 余额检测周期(分钟,默认 10)
|
||||
type GatewayCNProvidersConfig struct {
|
||||
BalanceCheckEnabled bool `mapstructure:"balance_check_enabled"`
|
||||
BalanceThreshold float64 `mapstructure:"balance_threshold"`
|
||||
BalanceCheckIntervalMinutes int `mapstructure:"balance_check_interval_minutes"`
|
||||
}
|
||||
|
||||
type GatewayLiveConfig struct {
|
||||
// MaxSessionDurationSeconds 是 Live 会话的硬上限。
|
||||
MaxSessionDurationSeconds int `mapstructure:"max_session_duration_seconds"`
|
||||
@@ -1703,6 +1719,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)
|
||||
}
|
||||
@@ -2351,6 +2371,10 @@ func setDefaults() {
|
||||
viper.SetDefault("gateway.grok.free_quota_soft_gate_percent", 95)
|
||||
viper.SetDefault("gateway.grok.free_quota_window_hours", 24)
|
||||
viper.SetDefault("gateway.grok.free_quota_stats_cache_seconds", 60)
|
||||
// 国产供应商余额检测(kimi/deepseek payg;zhipu 无余额端点,仅靠响应式 429/402)。
|
||||
viper.SetDefault("gateway.cn_providers.balance_check_enabled", true)
|
||||
viper.SetDefault("gateway.cn_providers.balance_threshold", 0.5)
|
||||
viper.SetDefault("gateway.cn_providers.balance_check_interval_minutes", 10)
|
||||
viper.SetDefault("gateway.image_concurrency.enabled", false)
|
||||
viper.SetDefault("gateway.image_concurrency.max_concurrent_requests", 0)
|
||||
viper.SetDefault("gateway.image_concurrency.overflow_mode", ImageConcurrencyOverflowModeReject)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -23,7 +23,27 @@ const (
|
||||
PlatformGemini = "gemini"
|
||||
PlatformAntigravity = "antigravity"
|
||||
PlatformGrok = "grok"
|
||||
PlatformComposite = "composite"
|
||||
// 国产 OpenAI 兼容供应商(经 OpenAI 网关转发,按 Chat Completions 协议)。
|
||||
PlatformKimi = "kimi" // Kimi (月之暗面 / Moonshot)
|
||||
PlatformZhipu = "zhipu" // 智谱 GLM (bigmodel)
|
||||
PlatformDeepseek = "deepseek" // DeepSeek
|
||||
PlatformComposite = "composite"
|
||||
)
|
||||
|
||||
// Account mode constants 区分国产供应商的「按量付费(余额)」与「Coding Plan」两种接入方式。
|
||||
// 存储于 credentials["account_mode"],决定 base_url 预设与额度监控方式。
|
||||
const (
|
||||
AccountModePayG = "payg" // 按量付费:消耗余额,做余额检测冷却
|
||||
AccountModeCoding = "coding" // Coding Plan:滚动用量窗口冷却(5h / weekly)
|
||||
)
|
||||
|
||||
// API protocol constants 国产供应商的上游 API 协议维度。存储于
|
||||
// credentials["api_protocol"],与 account_mode 正交:协议决定转发端点与格式,
|
||||
// 模式决定额度监控方式。同协议请求零转换直通;跨协议组合才走转换链。
|
||||
const (
|
||||
APIProtocolChatCompletions = "chat_completions" // OpenAI Chat Completions(默认)
|
||||
APIProtocolAnthropic = "anthropic" // 原生 Anthropic /v1/messages(适配 Claude Code)
|
||||
APIProtocolResponses = "responses" // OpenAI Responses(仅 deepseek,适配 Codex)
|
||||
)
|
||||
|
||||
// Account type constants
|
||||
|
||||
@@ -1025,7 +1025,8 @@ func (h *AccountHandler) Update(c *gin.Context) {
|
||||
// 当前请求。探测错误仅记录日志,不向上下文传播:探测失败时标记保持缺失,
|
||||
// 网关会按"现状即证据"默认走 Responses。
|
||||
func (h *AccountHandler) scheduleOpenAIResponsesProbe(account *service.Account) {
|
||||
if account == nil || account.Platform != service.PlatformOpenAI || account.Type != service.AccountTypeAPIKey {
|
||||
if account == nil || account.Type != service.AccountTypeAPIKey ||
|
||||
(account.Platform != service.PlatformOpenAI && !service.IsCNProvider(account.Platform)) {
|
||||
return
|
||||
}
|
||||
if h.accountTestService == nil {
|
||||
|
||||
@@ -57,17 +57,29 @@ type updateChannelRequest struct {
|
||||
}
|
||||
|
||||
type channelModelPricingRequest struct {
|
||||
Platform string `json:"platform" binding:"omitempty,max=50"`
|
||||
Models []string `json:"models" binding:"required,min=1,max=100"`
|
||||
BillingMode string `json:"billing_mode" binding:"omitempty,oneof=token per_request image"`
|
||||
InputPrice *float64 `json:"input_price" binding:"omitempty,min=0"`
|
||||
OutputPrice *float64 `json:"output_price" binding:"omitempty,min=0"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price" binding:"omitempty,min=0"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price" binding:"omitempty,min=0"`
|
||||
ImageInputPrice *float64 `json:"image_input_price" binding:"omitempty,min=0"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price" binding:"omitempty,min=0"`
|
||||
PerRequestPrice *float64 `json:"per_request_price" binding:"omitempty,min=0"`
|
||||
Intervals []pricingIntervalRequest `json:"intervals"`
|
||||
Platform string `json:"platform" binding:"omitempty,max=50"`
|
||||
Models []string `json:"models" binding:"required,min=1,max=100"`
|
||||
BillingMode string `json:"billing_mode" binding:"omitempty,oneof=token per_request image"`
|
||||
InputPrice *float64 `json:"input_price" binding:"omitempty,min=0"`
|
||||
OutputPrice *float64 `json:"output_price" binding:"omitempty,min=0"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price" binding:"omitempty,min=0"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price" binding:"omitempty,min=0"`
|
||||
ImageInputPrice *float64 `json:"image_input_price" binding:"omitempty,min=0"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price" binding:"omitempty,min=0"`
|
||||
PerRequestPrice *float64 `json:"per_request_price" binding:"omitempty,min=0"`
|
||||
Intervals []pricingIntervalRequest `json:"intervals"`
|
||||
TimePricing *channelTimePricingRequest `json:"time_pricing"`
|
||||
}
|
||||
|
||||
type channelTimePricingRequest struct {
|
||||
Timezone string `json:"timezone"`
|
||||
Periods []channelTimePricingPeriodRequest `json:"periods"`
|
||||
}
|
||||
|
||||
type channelTimePricingPeriodRequest struct {
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
Multiplier float64 `json:"multiplier"`
|
||||
}
|
||||
|
||||
type pricingIntervalRequest struct {
|
||||
@@ -108,18 +120,30 @@ type channelResponse struct {
|
||||
}
|
||||
|
||||
type channelModelPricingResponse struct {
|
||||
ID int64 `json:"id"`
|
||||
Platform string `json:"platform"`
|
||||
Models []string `json:"models"`
|
||||
BillingMode string `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 []pricingIntervalResponse `json:"intervals"`
|
||||
ID int64 `json:"id"`
|
||||
Platform string `json:"platform"`
|
||||
Models []string `json:"models"`
|
||||
BillingMode string `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 []pricingIntervalResponse `json:"intervals"`
|
||||
TimePricing *channelTimePricingResponse `json:"time_pricing"`
|
||||
}
|
||||
|
||||
type channelTimePricingResponse struct {
|
||||
Timezone string `json:"timezone"`
|
||||
Periods []channelTimePricingPeriodResponse `json:"periods"`
|
||||
}
|
||||
|
||||
type channelTimePricingPeriodResponse struct {
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
Multiplier float64 `json:"multiplier"`
|
||||
}
|
||||
|
||||
type pricingIntervalResponse struct {
|
||||
@@ -228,9 +252,25 @@ func pricingToResponse(p *service.ChannelModelPricing) channelModelPricingRespon
|
||||
ImageOutputPrice: p.ImageOutputPrice,
|
||||
PerRequestPrice: p.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
TimePricing: timePricingToResponse(p.TimePricing),
|
||||
}
|
||||
}
|
||||
|
||||
func timePricingToResponse(value *service.ChannelTimePricing) *channelTimePricingResponse {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
periods := make([]channelTimePricingPeriodResponse, 0, len(value.Periods))
|
||||
for _, period := range value.Periods {
|
||||
periods = append(periods, channelTimePricingPeriodResponse{
|
||||
StartTime: period.StartTime,
|
||||
EndTime: period.EndTime,
|
||||
Multiplier: period.Multiplier,
|
||||
})
|
||||
}
|
||||
return &channelTimePricingResponse{Timezone: value.Timezone, Periods: periods}
|
||||
}
|
||||
|
||||
func intervalToResponse(iv service.PricingInterval) pricingIntervalResponse {
|
||||
return pricingIntervalResponse{
|
||||
ID: iv.ID,
|
||||
@@ -280,11 +320,27 @@ func pricingRequestToService(reqs []channelModelPricingRequest) []service.Channe
|
||||
ImageOutputPrice: r.ImageOutputPrice,
|
||||
PerRequestPrice: r.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
TimePricing: timePricingRequestToService(r.TimePricing),
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func timePricingRequestToService(value *channelTimePricingRequest) *service.ChannelTimePricing {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
periods := make([]service.ChannelTimePricingPeriod, 0, len(value.Periods))
|
||||
for _, period := range value.Periods {
|
||||
periods = append(periods, service.ChannelTimePricingPeriod{
|
||||
StartTime: period.StartTime,
|
||||
EndTime: period.EndTime,
|
||||
Multiplier: period.Multiplier,
|
||||
})
|
||||
}
|
||||
return &service.ChannelTimePricing{Timezone: value.Timezone, Periods: periods}
|
||||
}
|
||||
|
||||
func accountStatsPricingRuleRequestToService(r accountStatsPricingRuleRequest) service.AccountStatsPricingRule {
|
||||
return service.AccountStatsPricingRule{
|
||||
Name: r.Name,
|
||||
@@ -512,9 +568,12 @@ func (h *ChannelHandler) GetModelDefaultPricing(c *gin.Context) {
|
||||
var platformToLiteLLMProvider = map[string]string{
|
||||
service.PlatformAnthropic: "anthropic",
|
||||
service.PlatformOpenAI: "openai",
|
||||
service.PlatformGemini: "google",
|
||||
service.PlatformGemini: "gemini",
|
||||
service.PlatformAntigravity: "anthropic",
|
||||
service.PlatformGrok: "xai",
|
||||
service.PlatformKimi: "moonshot",
|
||||
service.PlatformZhipu: "zhipu",
|
||||
service.PlatformDeepseek: "deepseek",
|
||||
}
|
||||
|
||||
// SyncPricingModels 返回 LiteLLM 定价目录中指定平台的最新模型列表
|
||||
|
||||
@@ -421,6 +421,49 @@ func TestPricingRequestToService_NilPriceFields(t *testing.T) {
|
||||
require.Nil(t, r.PerRequestPrice)
|
||||
}
|
||||
|
||||
func TestPricingRequestToService_TimePricing(t *testing.T) {
|
||||
req := channelModelPricingRequest{
|
||||
Models: []string{"gpt-5"},
|
||||
BillingMode: "token",
|
||||
TimePricing: &channelTimePricingRequest{
|
||||
Timezone: "Asia/Shanghai",
|
||||
Periods: []channelTimePricingPeriodRequest{{
|
||||
StartTime: "09:00", EndTime: "12:00", Multiplier: 2,
|
||||
}},
|
||||
},
|
||||
}
|
||||
|
||||
got := pricingRequestToService([]channelModelPricingRequest{req})
|
||||
require.Equal(t, "Asia/Shanghai", got[0].TimePricing.Timezone)
|
||||
require.Equal(t, 2.0, got[0].TimePricing.Periods[0].Multiplier)
|
||||
}
|
||||
|
||||
func TestPricingRequestToService_TimePricingNil(t *testing.T) {
|
||||
got := pricingRequestToService([]channelModelPricingRequest{{Models: []string{"gpt-5"}}})
|
||||
require.Nil(t, got[0].TimePricing)
|
||||
}
|
||||
|
||||
func TestPricingToResponse_TimePricing(t *testing.T) {
|
||||
got := pricingToResponse(&service.ChannelModelPricing{
|
||||
BillingMode: service.BillingModeToken,
|
||||
TimePricing: &service.ChannelTimePricing{
|
||||
Timezone: "Asia/Shanghai",
|
||||
Periods: []service.ChannelTimePricingPeriod{{
|
||||
StartTime: "14:00", EndTime: "18:00", Multiplier: 1.25,
|
||||
}},
|
||||
},
|
||||
})
|
||||
|
||||
require.NotNil(t, got.TimePricing)
|
||||
require.Equal(t, "Asia/Shanghai", got.TimePricing.Timezone)
|
||||
require.Equal(t, 1.25, got.TimePricing.Periods[0].Multiplier)
|
||||
}
|
||||
|
||||
func TestPricingToResponse_TimePricingNil(t *testing.T) {
|
||||
got := pricingToResponse(&service.ChannelModelPricing{})
|
||||
require.Nil(t, got.TimePricing)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 3. SyncPricingModels handler
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -459,7 +502,7 @@ func TestSyncPricingModels_ValidPlatform_EmptyService(t *testing.T) {
|
||||
svc := service.NewPricingService(nil, nil)
|
||||
router := setupSyncPricingModelsRouter(svc)
|
||||
|
||||
for _, platform := range []string{"anthropic", "openai", "gemini", "antigravity"} {
|
||||
for _, platform := range []string{"anthropic", "openai", "gemini", "antigravity", "grok", "kimi", "zhipu", "deepseek"} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/channels/pricing/sync-models?platform="+platform, nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// CNProviderHandler 暴露国产供应商(kimi/zhipu/deepseek)的额度与余额查询端点。
|
||||
//
|
||||
// - GET /admin/cn-providers/accounts/:id/quota Coding Plan 滚动窗口用量(kimi/zhipu)
|
||||
// - GET /admin/cn-providers/accounts/:id/balance payg 账号余额(kimi/deepseek)
|
||||
//
|
||||
// 智谱(zhipu)无余额端点,故同一账号仅 quota 或 balance 其一可用:服务端按账号
|
||||
// platform + account_mode 校验并返回明确错误(见 CNProvider*Service 的 load*Account)。
|
||||
type CNProviderHandler struct {
|
||||
quotaService *service.CNProviderQuotaService
|
||||
balanceService *service.CNProviderBalanceService
|
||||
}
|
||||
|
||||
func NewCNProviderHandler(
|
||||
quotaService *service.CNProviderQuotaService,
|
||||
balanceService *service.CNProviderBalanceService,
|
||||
) *CNProviderHandler {
|
||||
return &CNProviderHandler{
|
||||
quotaService: quotaService,
|
||||
balanceService: balanceService,
|
||||
}
|
||||
}
|
||||
|
||||
// QueryQuota 查询 Coding Plan 滚动窗口用量(5h + weekly)。
|
||||
func (h *CNProviderHandler) QueryQuota(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h == nil || h.quotaService == nil {
|
||||
response.BadRequest(c, "cn provider quota service is not enabled")
|
||||
return
|
||||
}
|
||||
result, err := h.quotaService.QueryUsage(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
// QueryBalance 查询 payg 账号余额。
|
||||
func (h *CNProviderHandler) QueryBalance(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h == nil || h.balanceService == nil {
|
||||
response.BadRequest(c, "cn provider balance service is not enabled")
|
||||
return
|
||||
}
|
||||
result, err := h.balanceService.QueryBalance(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
@@ -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 kimi zhipu deepseek 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 kimi zhipu deepseek 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 {
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
//go:build unit
|
||||
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 回归分组平台枚举:kimi/zhipu/deepseek 必须能通过 Create/Update 的 binding 校验
|
||||
// (历史 bug:调度/路由链路已支持 CN 平台分组,但 oneof 白名单漏加三平台,导致
|
||||
// 平台分组无法创建、CN 账号"无可用分组");非法值仍须被拒。
|
||||
func bindGroupPlatformJSON(t *testing.T, target any, body string) error {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
return c.ShouldBindJSON(target)
|
||||
}
|
||||
|
||||
func TestGroupPlatformBinding_AllowedPlatforms(t *testing.T) {
|
||||
allowed := []string{
|
||||
"anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "composite",
|
||||
}
|
||||
for _, platform := range allowed {
|
||||
t.Run("create_"+platform, func(t *testing.T) {
|
||||
var req CreateGroupRequest
|
||||
body := fmt.Sprintf(`{"name":"g","platform":%q}`, platform)
|
||||
require.NoError(t, bindGroupPlatformJSON(t, &req, body),
|
||||
"platform %q 应通过 CreateGroupRequest 校验", platform)
|
||||
require.Equal(t, platform, req.Platform)
|
||||
})
|
||||
t.Run("update_"+platform, func(t *testing.T) {
|
||||
var req UpdateGroupRequest
|
||||
body := fmt.Sprintf(`{"platform":%q}`, platform)
|
||||
require.NoError(t, bindGroupPlatformJSON(t, &req, body),
|
||||
"platform %q 应通过 UpdateGroupRequest 校验", platform)
|
||||
require.Equal(t, platform, req.Platform)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupPlatformBinding_RejectsInvalidPlatforms(t *testing.T) {
|
||||
invalid := []string{
|
||||
"moonshot", // 厂商别名,不是平台标识
|
||||
"Kimi", // 大小写敏感
|
||||
"openai ", // 尾随空格
|
||||
"glm",
|
||||
"bogus",
|
||||
}
|
||||
for _, platform := range invalid {
|
||||
t.Run("create_"+platform, func(t *testing.T) {
|
||||
var req CreateGroupRequest
|
||||
body := fmt.Sprintf(`{"name":"g","platform":%q}`, platform)
|
||||
require.Error(t, bindGroupPlatformJSON(t, &req, body),
|
||||
"platform %q 应被 CreateGroupRequest 拒绝", platform)
|
||||
})
|
||||
t.Run("update_"+platform, func(t *testing.T) {
|
||||
var req UpdateGroupRequest
|
||||
body := fmt.Sprintf(`{"platform":%q}`, platform)
|
||||
require.Error(t, bindGroupPlatformJSON(t, &req, body),
|
||||
"platform %q 应被 UpdateGroupRequest 拒绝", platform)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 守住 composite 路由目标不放行 CN:CN 平台不可作为 composite 路由目标
|
||||
// (DetectModelPlatform/isConcreteRequestPlatform 均无 CN 分支,放行即打开半实现路径)。
|
||||
func TestCompositeRouteTargetPlatform_StillExcludesCNProviders(t *testing.T) {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek"} {
|
||||
var req CompositeRouteRequest
|
||||
body := fmt.Sprintf(`{"public_model":"m","target_platform":%q}`, platform)
|
||||
require.Error(t, bindGroupPlatformJSON(t, &req, body),
|
||||
"composite target_platform %q 应保持被拒", platform)
|
||||
}
|
||||
}
|
||||
@@ -112,12 +112,13 @@ func TestUpdateUserPlatformQuotas_Success(t *testing.T) {
|
||||
if len(repo.upsertCalls) != 1 {
|
||||
t.Fatalf("UpsertForUser should be called once, got %d", len(repo.upsertCalls))
|
||||
}
|
||||
if repo.upsertCalls[0].userID != 42 || len(repo.upsertCalls[0].records) != len(service.AllowedQuotaPlatforms) {
|
||||
// upsert 记录数 = 请求体中给出的平台数(未给出的平台不落库)。
|
||||
if repo.upsertCalls[0].userID != 42 || len(repo.upsertCalls[0].records) != 5 {
|
||||
t.Errorf("unexpected upsert call: %+v", repo.upsertCalls[0])
|
||||
}
|
||||
// 缓存失效:按全部允许平台统一失效。
|
||||
if len(cache.deleteCalls) != 5 {
|
||||
t.Errorf("expected 5 cache delete calls, got %d: %+v", len(cache.deleteCalls), cache.deleteCalls)
|
||||
// 缓存失效:按全部允许平台统一失效(含 kimi/zhipu/deepseek)。
|
||||
if len(cache.deleteCalls) != len(service.AllowedQuotaPlatforms) {
|
||||
t.Errorf("expected %d cache delete calls, got %d: %+v", len(service.AllowedQuotaPlatforms), len(cache.deleteCalls), cache.deleteCalls)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ type AdminHandlers struct {
|
||||
GeminiOAuth *admin.GeminiOAuthHandler
|
||||
AntigravityOAuth *admin.AntigravityOAuthHandler
|
||||
GrokOAuth *admin.GrokOAuthHandler
|
||||
CNProvider *admin.CNProviderHandler
|
||||
Proxy *admin.ProxyHandler
|
||||
Redeem *admin.RedeemHandler
|
||||
Promo *admin.PromoHandler
|
||||
|
||||
@@ -255,6 +255,48 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
if err == nil && result != nil && result.FirstTokenMs != nil {
|
||||
service.SetOpsLatencyMs(c, service.OpsTimeToFirstTokenMsKey, int64(*result.FirstTokenMs))
|
||||
}
|
||||
// #5148 对齐:错误返回携带的部分 result(流中断前上游已计量的 usage)照常
|
||||
// 入账;failover 错误恒定 result=nil,不会重复计费。
|
||||
submitChatUsage := func(res *service.OpenAIForwardResult) {
|
||||
if res == nil {
|
||||
return
|
||||
}
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, res)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), res, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: res,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, res.UpstreamModel),
|
||||
PricingAt: pricingAt,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.openai_gateway.chat_completions"),
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
zap.String("model", reqModel),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("openai_chat_completions.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
if result != nil && result.ImageCount > 0 {
|
||||
reqLog.Warn("openai_chat_completions.forward_partial_error_with_image_result",
|
||||
@@ -339,6 +381,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated),
|
||||
zap.Error(err),
|
||||
)
|
||||
submitChatUsage(result)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -348,42 +391,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil)
|
||||
}
|
||||
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
PricingAt: pricingAt,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.openai_gateway.chat_completions"),
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
zap.String("model", reqModel),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("openai_chat_completions.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
submitChatUsage(result)
|
||||
reqLog.Debug("openai_chat_completions.request_completed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("switch_count", switchCount),
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
package handler
|
||||
|
||||
// CN 分组 /v1/messages 调度闸门回归(修复:正常途径创建的 CN 分组曾恒 403):
|
||||
// sanitizeGroupMessagesDispatchFields 对非 openai 平台强制 AllowMessagesDispatch
|
||||
// =false,故 CN 分组必须与 grok 一样在闸门处豁免,否则原生 Anthropic 直通
|
||||
//(Claude Code 主用例)永远不可达。
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAllowOpenAICompatibleMessagesDispatch_CNProvidersExempt(t *testing.T) {
|
||||
require.True(t, allowOpenAICompatibleMessagesDispatch(nil), "无 key 保持放行")
|
||||
|
||||
for _, platform := range []string{service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformGrok} {
|
||||
apiKey := &service.APIKey{Group: &service.Group{Platform: platform, AllowMessagesDispatch: false}}
|
||||
require.True(t, allowOpenAICompatibleMessagesDispatch(apiKey),
|
||||
"%s 分组必须豁免 allow_messages_dispatch 闸门", platform)
|
||||
}
|
||||
|
||||
// 非回归:openai 分组仍受开关控制。
|
||||
openaiOff := &service.APIKey{Group: &service.Group{Platform: service.PlatformOpenAI, AllowMessagesDispatch: false}}
|
||||
require.False(t, allowOpenAICompatibleMessagesDispatch(openaiOff))
|
||||
openaiOn := &service.APIKey{Group: &service.Group{Platform: service.PlatformOpenAI, AllowMessagesDispatch: true}}
|
||||
require.True(t, allowOpenAICompatibleMessagesDispatch(openaiOn))
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -78,7 +78,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) {
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
|
||||
if apiKey.Group != nil && !apiKey.Group.AllowMessagesDispatch {
|
||||
if !allowOpenAICompatibleMessagesDispatch(apiKey) {
|
||||
h.anthropicErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group does not allow /v1/messages dispatch")
|
||||
return
|
||||
|
||||
@@ -164,13 +164,11 @@ func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecord
|
||||
|
||||
func openAICompatibleRequestPlatform(ctx context.Context, apiKey *service.APIKey) string {
|
||||
if platform, ok := service.ResolvedTargetPlatformFromContext(ctx); ok {
|
||||
if platform == service.PlatformGrok {
|
||||
return service.PlatformGrok
|
||||
}
|
||||
return service.PlatformOpenAI
|
||||
// 保留 grok 与国产供应商原值,其他归一为 openai(与调度器精确匹配语义一致)。
|
||||
return service.NormalizeOpenAICompatiblePlatform(platform)
|
||||
}
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.Platform == service.PlatformGrok {
|
||||
return service.PlatformGrok
|
||||
if apiKey != nil && apiKey.Group != nil {
|
||||
return service.NormalizeOpenAICompatiblePlatform(apiKey.Group.Platform)
|
||||
}
|
||||
return service.PlatformOpenAI
|
||||
}
|
||||
@@ -182,6 +180,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
|
||||
@@ -189,11 +197,20 @@ func allowOpenAICompatibleMessagesDispatch(apiKey *service.APIKey) bool {
|
||||
if apiKey.Group.Platform == service.PlatformGrok {
|
||||
return true
|
||||
}
|
||||
// 国产供应商分组与 grok 同语义:/v1/messages 就是其主要服务形态(anthropic
|
||||
// 协议账号原生直通 Claude Code),无需 allow_messages_dispatch 开关授权——
|
||||
// 该开关对非 openai 平台恒被 sanitizeGroupMessagesDispatchFields 置 false,
|
||||
// 若不豁免,CN 分组将永远 403。
|
||||
if service.IsCNProvider(apiKey.Group.Platform) {
|
||||
return true
|
||||
}
|
||||
return apiKey.Group.AllowMessagesDispatch
|
||||
}
|
||||
|
||||
func openAICompatibleTextTargetAllowed(c *gin.Context, apiKey *service.APIKey, model string) bool {
|
||||
return compositeTargetPlatformAllowed(c, apiKey, model, service.PlatformOpenAI, service.PlatformGrok)
|
||||
return compositeTargetPlatformAllowed(c, apiKey, model,
|
||||
service.PlatformOpenAI, service.PlatformGrok,
|
||||
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek)
|
||||
}
|
||||
|
||||
// NewOpenAIGatewayHandler creates a new OpenAIGatewayHandler
|
||||
@@ -288,6 +305,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 +401,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 +453,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 +470,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 +513,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
|
||||
@@ -570,6 +604,50 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
if err == nil && result != nil && result.FirstTokenMs != nil {
|
||||
service.SetOpsLatencyMs(c, service.OpsTimeToFirstTokenMsKey, int64(*result.FirstTokenMs))
|
||||
}
|
||||
// #5148 对齐:错误返回携带的部分 result(流中断前上游已计量的 usage)照常
|
||||
// 入账;failover 错误恒定 result=nil,不会重复计费。
|
||||
submitResponsesUsage := func(res *service.OpenAIForwardResult) {
|
||||
if res == nil {
|
||||
return
|
||||
}
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, res)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), res, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: res,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, res.UpstreamModel),
|
||||
PricingAt: pricingAt,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.openai_gateway.responses"),
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
zap.String("model", reqModel),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("openai.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
if result != nil && result.ImageCount > 0 {
|
||||
reqLog.Warn("openai.forward_partial_error_with_image_result",
|
||||
@@ -669,6 +747,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated),
|
||||
zap.Error(err),
|
||||
}
|
||||
submitResponsesUsage(result)
|
||||
if shouldLogOpenAIForwardFailureAsWarn(c, wroteFallback) {
|
||||
reqLog.Warn("openai.forward_failed", fields...)
|
||||
return
|
||||
@@ -687,46 +766,8 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), nil)
|
||||
}
|
||||
|
||||
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
|
||||
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
PricingAt: pricingAt,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.openai_gateway.responses"),
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
zap.String("model", reqModel),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("openai.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
submitResponsesUsage(result)
|
||||
reqLog.Debug("openai.request_completed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("switch_count", switchCount),
|
||||
@@ -735,12 +776,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 等子路径),
|
||||
@@ -750,32 +787,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"
|
||||
@@ -804,7 +836,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
|
||||
}
|
||||
|
||||
@@ -1127,6 +1159,51 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
if err == nil && result != nil && result.FirstTokenMs != nil {
|
||||
service.SetOpsLatencyMs(c, service.OpsTimeToFirstTokenMsKey, int64(*result.FirstTokenMs))
|
||||
}
|
||||
// Forward 与错误一起返回的部分结果:流中断/客户端断开排水前上游已计量的
|
||||
// usage 照常入账,避免上游已产生消耗的请求完全漏记(#5148,对齐 anthropic
|
||||
// 网关同名修复)。failover 错误恒定 result=nil,不会重复计费。
|
||||
submitMessagesUsage := func(res *service.OpenAIForwardResult) {
|
||||
if res == nil {
|
||||
return
|
||||
}
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, res)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), res, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: res,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, res.UpstreamModel),
|
||||
PricingAt: pricingAt,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.openai_gateway.messages"),
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
zap.String("model", reqModel),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("openai_messages.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
if result != nil && result.ImageCount > 0 {
|
||||
reqLog.Warn("openai_messages.forward_partial_error_with_image_result",
|
||||
@@ -1201,6 +1278,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Error(err),
|
||||
)
|
||||
// 断开排水期间上游已计量的 usage 必须入账(此前直接 return 丢弃,
|
||||
// payg 上游照常计费而平台漏记)。
|
||||
submitMessagesUsage(result)
|
||||
return
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), false, nil)
|
||||
@@ -1210,6 +1290,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
zap.Bool("fallback_error_response_written", wroteFallback),
|
||||
zap.Error(err),
|
||||
)
|
||||
submitMessagesUsage(result)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -1219,44 +1300,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, nil)
|
||||
}
|
||||
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, result.UpstreamModel),
|
||||
PricingAt: pricingAt,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.openai_gateway.messages"),
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
zap.String("model", reqModel),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("openai_messages.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
submitMessagesUsage(result)
|
||||
reqLog.Debug("openai_messages.request_completed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("switch_count", switchCount),
|
||||
@@ -1402,12 +1446,10 @@ const (
|
||||
// 由 BeforeTurn 在每个 turn 开始时冻结,AfterTurn 的用量提交读取它;turn 在
|
||||
// 连接内串行推进,互斥锁只为跨用量提交 goroutine 的读取安全。
|
||||
//
|
||||
// 零值语义(重要):ws_v2 passthrough ingress 只实现了 AfterTurn,没有任何
|
||||
// turn 起始回调,BeforeTurn 永远不会被调用。此时本值保持零,RecordUsage 经
|
||||
// openAIUsagePricingAt 回退到记录时刻——与引入分组利润控制前的基线一致。
|
||||
// 绝不能用建连时刻初始化:那会把透传连接的所有 turn 钉死在建连时的高峰因子,
|
||||
// 客户端只要峰前一分钟建连并保活,整条连接就能全程按谷价结算,正是利润控制
|
||||
// 想堵的漏洞。透传 ingress 目前不做 turn 级利润复核,只有建连时的准入门。
|
||||
// ws_v2 passthrough ingress 没有 BeforeTurn,因此本值会保持零;AfterTurn 必须
|
||||
// 以 TurnStarted 已记录的所属 turn 开始时刻为回退,而不是用建连或记录时刻。
|
||||
// 这样每个 passthrough turn 都按自己的开始时刻计价,但不改变其仅在建连时执行
|
||||
// 准入门、没有 turn 级利润复核的既有行为。
|
||||
type openAIWSTurnPricing struct {
|
||||
mu sync.Mutex
|
||||
at time.Time
|
||||
@@ -1419,10 +1461,13 @@ func (p *openAIWSTurnPricing) freeze(at time.Time) {
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
func (p *openAIWSTurnPricing) current() time.Time {
|
||||
func (p *openAIWSTurnPricing) currentOr(fallback time.Time) time.Time {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return p.at
|
||||
if !p.at.IsZero() {
|
||||
return p.at
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
// recordOpenAIProfitVeto 记录 OpenAI 侧选号循环的一次利润门终检否决:把账号
|
||||
@@ -1685,6 +1730,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "missing first response.create message")
|
||||
return
|
||||
}
|
||||
firstTurnStartedAt := time.Now()
|
||||
if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "unsupported websocket message type")
|
||||
return
|
||||
@@ -2013,19 +2059,37 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
|
||||
maxReasoningEffort, reasoningEffortMappings, _ := openAIReasoningEffortPolicyForRequest(c, apiKey)
|
||||
var requestPayloadHash string
|
||||
var turnStartsMu sync.Mutex
|
||||
turnStarts := make(map[int]time.Time, 4)
|
||||
recordTurnStart := func(turn int, startedAt time.Time) {
|
||||
if turn <= 0 || startedAt.IsZero() {
|
||||
return
|
||||
}
|
||||
turnStartsMu.Lock()
|
||||
turnStarts[turn] = startedAt
|
||||
turnStartsMu.Unlock()
|
||||
}
|
||||
getTurnStart := func(turn int) time.Time {
|
||||
turnStartsMu.Lock()
|
||||
startedAt := turnStarts[turn]
|
||||
delete(turnStarts, turn)
|
||||
turnStartsMu.Unlock()
|
||||
return startedAt
|
||||
}
|
||||
// Passthrough rejects overlapping response.create frames, so one immutable
|
||||
// turn-tagged slot preserves the exact mapping used for the in-flight request.
|
||||
var turnChannelMapping atomic.Pointer[openAIWSTurnChannelMappingSnapshot]
|
||||
turnChannelMapping.Store(&openAIWSTurnChannelMappingSnapshot{turn: 1, mapping: channelMappingWS})
|
||||
// turn 级定价:BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号,
|
||||
// AfterTurn 的计费读取所属 turn 的时刻。零值起步的语义见
|
||||
// openAIWSTurnPricing 的注释——绝不能用建连时刻初始化。
|
||||
// turn 级定价:BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号;
|
||||
// passthrough 没有 BeforeTurn 时,AfterTurn 回退到 TurnStarted 的所属 turn 时刻。
|
||||
var turnPricing openAIWSTurnPricing
|
||||
hooks := &service.OpenAIWSIngressHooks{
|
||||
ClientLifecycleContext: clientLifecycleCtx,
|
||||
InitialRequestModel: reqModel,
|
||||
InitialTurnStartedAt: firstTurnStartedAt,
|
||||
MaxReasoningEffort: maxReasoningEffort,
|
||||
ReasoningEffortMappings: reasoningEffortMappings,
|
||||
TurnStarted: recordTurnStart,
|
||||
BeforeRequest: func(turn int, payload []byte, originalModel string) error {
|
||||
c.Set(securityAuditWSTurnContextKey, turn)
|
||||
if turn == 1 {
|
||||
@@ -2111,6 +2175,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return nil
|
||||
},
|
||||
AfterTurn: func(turn int, result *service.OpenAIForwardResult, turnErr error) {
|
||||
turnStart := getTurnStart(turn)
|
||||
// F1: cyber 标记按 turn 生命周期清理——defer 保证任意早返回路径都执行;
|
||||
// CyberBlocked 必须在 submit 前同步预捕获(task 闭包由 worker 池异步执行,
|
||||
// 届时 defer 已清除标记)。
|
||||
@@ -2178,7 +2243,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
turnRecordPricingAt := turnPricing.current()
|
||||
turnRecordPricingAt := turnPricing.currentOr(turnStart)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -7,16 +7,20 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestOpenAIWSTurnPricingZeroValue 钉死 WS turn 定价的零值语义:
|
||||
// 没有 turn 起始回调的 ingress 模式(ws_v2 passthrough 只实现 AfterTurn)
|
||||
// 必须让 pricingAt 保持零,由 RecordUsage 回退到记录时刻。
|
||||
//
|
||||
// 反例(本 PR 引入的回归):用建连时刻初始化,会把透传连接的所有 turn 钉死在
|
||||
// 建连时的高峰因子——客户端峰前一分钟建连并保活,整条连接就按谷价结算。
|
||||
func TestOpenAIWSTurnPricingZeroValue(t *testing.T) {
|
||||
var p openAIWSTurnPricing
|
||||
require.True(t, p.current().IsZero(),
|
||||
"未经 turn 起始回调冻结时必须保持零值,交由 RecordUsage 回退记录时刻")
|
||||
func TestOpenAIWSTurnPricingCurrentOr(t *testing.T) {
|
||||
fallback := time.Date(2024, time.January, 2, 2, 0, 0, 0, time.UTC)
|
||||
|
||||
t.Run("frozen time takes precedence", func(t *testing.T) {
|
||||
frozen := fallback.Add(time.Minute)
|
||||
var p openAIWSTurnPricing
|
||||
p.freeze(frozen)
|
||||
require.Equal(t, frozen, p.currentOr(fallback))
|
||||
})
|
||||
|
||||
t.Run("zero value falls back to turn start", func(t *testing.T) {
|
||||
var p openAIWSTurnPricing
|
||||
require.Equal(t, fallback, p.currentOr(fallback))
|
||||
})
|
||||
}
|
||||
|
||||
// TestOpenAIWSTurnPricingFreezePerTurn 钉死每个 turn 的 BeforeTurn 都会覆盖
|
||||
@@ -27,8 +31,8 @@ func TestOpenAIWSTurnPricingFreezePerTurn(t *testing.T) {
|
||||
turn2 := time.Now()
|
||||
|
||||
p.freeze(turn1)
|
||||
require.Equal(t, turn1, p.current())
|
||||
require.Equal(t, turn1, p.currentOr(time.Time{}))
|
||||
|
||||
p.freeze(turn2)
|
||||
require.Equal(t, turn2, p.current(), "后续 turn 必须使用自己的定价时刻")
|
||||
require.Equal(t, turn2, p.currentOr(time.Time{}), "后续 turn 必须使用自己的定价时刻")
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
@@ -23,6 +23,7 @@ func ProvideAdminHandlers(
|
||||
geminiOAuthHandler *admin.GeminiOAuthHandler,
|
||||
antigravityOAuthHandler *admin.AntigravityOAuthHandler,
|
||||
grokOAuthHandler *admin.GrokOAuthHandler,
|
||||
cnProviderHandler *admin.CNProviderHandler,
|
||||
proxyHandler *admin.ProxyHandler,
|
||||
redeemHandler *admin.RedeemHandler,
|
||||
promoHandler *admin.PromoHandler,
|
||||
@@ -63,6 +64,7 @@ func ProvideAdminHandlers(
|
||||
GeminiOAuth: geminiOAuthHandler,
|
||||
AntigravityOAuth: antigravityOAuthHandler,
|
||||
GrokOAuth: grokOAuthHandler,
|
||||
CNProvider: cnProviderHandler,
|
||||
Proxy: proxyHandler,
|
||||
Redeem: redeemHandler,
|
||||
Promo: promoHandler,
|
||||
@@ -253,6 +255,7 @@ var ProviderSet = wire.NewSet(
|
||||
admin.NewGeminiOAuthHandler,
|
||||
admin.NewAntigravityOAuthHandler,
|
||||
admin.NewGrokOAuthHandler,
|
||||
admin.NewCNProviderHandler,
|
||||
admin.NewProxyHandler,
|
||||
admin.NewRedeemHandler,
|
||||
admin.NewPromoHandler,
|
||||
|
||||
@@ -112,7 +112,8 @@ type GeminiImageSearch struct {
|
||||
|
||||
// GeminiToolConfig Gemini 工具配置
|
||||
type GeminiToolConfig struct {
|
||||
FunctionCallingConfig *GeminiFunctionCallingConfig `json:"functionCallingConfig,omitempty"`
|
||||
FunctionCallingConfig *GeminiFunctionCallingConfig `json:"functionCallingConfig,omitempty"`
|
||||
IncludeServerSideToolInvocations *bool `json:"includeServerSideToolInvocations,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiFunctionCallingConfig 函数调用配置
|
||||
|
||||
@@ -156,6 +156,13 @@ func TransformClaudeToGeminiWithOptions(claudeReq *ClaudeRequest, projectID, map
|
||||
Mode: "VALIDATED",
|
||||
},
|
||||
}
|
||||
// 内置工具(googleSearch)与函数调用混用时,上游要求显式开启
|
||||
// includeServerSideToolInvocations,否则返回 400(issue #5709)。
|
||||
// 与 raw 透传路的 enableMixedGeminiToolInvocations 注入保持同一语义。
|
||||
if hasMixedToolInvocations(tools) {
|
||||
enabled := true
|
||||
innerRequest.ToolConfig.IncludeServerSideToolInvocations = &enabled
|
||||
}
|
||||
}
|
||||
|
||||
if systemInstruction != nil {
|
||||
@@ -703,6 +710,21 @@ func isWebSearchTool(tool ClaudeTool) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// hasMixedToolInvocations 判断构建后的工具声明是否同时包含函数声明与内置工具
|
||||
// (googleSearch)。仅在两者并存时需要开启 includeServerSideToolInvocations。
|
||||
func hasMixedToolInvocations(declarations []GeminiToolDeclaration) bool {
|
||||
hasFunc, hasBuiltin := false, false
|
||||
for _, d := range declarations {
|
||||
if len(d.FunctionDeclarations) > 0 {
|
||||
hasFunc = true
|
||||
}
|
||||
if d.GoogleSearch != nil {
|
||||
hasBuiltin = true
|
||||
}
|
||||
}
|
||||
return hasFunc && hasBuiltin
|
||||
}
|
||||
|
||||
// buildTools 构建 tools
|
||||
func buildTools(tools []ClaudeTool) []GeminiToolDeclaration {
|
||||
if len(tools) == 0 {
|
||||
|
||||
@@ -565,3 +565,59 @@ func TestTransformClaudeToGeminiWithOptions_PreservesWebSearchAlongsideFunctions
|
||||
require.Equal(t, "get_weather", req.Request.Tools[0].FunctionDeclarations[0].Name)
|
||||
require.NotNil(t, req.Request.Tools[1].GoogleSearch)
|
||||
}
|
||||
|
||||
func TestGeminiToolConfig_IncludeServerSideToolInvocations(t *testing.T) {
|
||||
functionTool := ClaudeTool{
|
||||
Name: "get_weather",
|
||||
Description: "Get weather information",
|
||||
InputSchema: map[string]any{"type": "object"},
|
||||
}
|
||||
webSearchTool := ClaudeTool{
|
||||
Type: "web_search_20250305",
|
||||
Name: "web_search",
|
||||
}
|
||||
|
||||
transform := func(t *testing.T, tools []ClaudeTool) (V1InternalRequest, string) {
|
||||
t.Helper()
|
||||
body, err := TransformClaudeToGeminiWithOptions(&ClaudeRequest{
|
||||
Model: "claude-3-5-sonnet-latest",
|
||||
Messages: []ClaudeMessage{
|
||||
{
|
||||
Role: "user",
|
||||
Content: json.RawMessage(`[{"type":"text","text":"hello"}]`),
|
||||
},
|
||||
},
|
||||
Tools: tools,
|
||||
}, "project-1", "gemini-2.5-flash", DefaultTransformOptions())
|
||||
require.NoError(t, err)
|
||||
|
||||
var req V1InternalRequest
|
||||
require.NoError(t, json.Unmarshal(body, &req))
|
||||
return req, string(body)
|
||||
}
|
||||
|
||||
t.Run("mixed builtin and function tools enable server-side tool invocations", func(t *testing.T) {
|
||||
req, raw := transform(t, []ClaudeTool{functionTool, webSearchTool})
|
||||
|
||||
require.NotNil(t, req.Request.ToolConfig)
|
||||
require.NotNil(t, req.Request.ToolConfig.IncludeServerSideToolInvocations)
|
||||
require.True(t, *req.Request.ToolConfig.IncludeServerSideToolInvocations)
|
||||
require.Contains(t, raw, `"includeServerSideToolInvocations":true`)
|
||||
})
|
||||
|
||||
t.Run("function tools only leave the flag unset", func(t *testing.T) {
|
||||
req, raw := transform(t, []ClaudeTool{functionTool})
|
||||
|
||||
require.NotNil(t, req.Request.ToolConfig)
|
||||
require.Nil(t, req.Request.ToolConfig.IncludeServerSideToolInvocations)
|
||||
require.NotContains(t, raw, "includeServerSideToolInvocations")
|
||||
})
|
||||
|
||||
t.Run("web search only leaves the flag unset", func(t *testing.T) {
|
||||
req, raw := transform(t, []ClaudeTool{webSearchTool})
|
||||
|
||||
require.NotNil(t, req.Request.ToolConfig)
|
||||
require.Nil(t, req.Request.ToolConfig.IncludeServerSideToolInvocations)
|
||||
require.NotContains(t, raw, "includeServerSideToolInvocations")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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 选择,
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -430,7 +430,7 @@ func (r *ResponsesClientToolStreamRestorer) RestoreEvent(payload []byte) ([][]by
|
||||
if err := json.Unmarshal(payload, &wire); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if wire.Type == "response.completed" || wire.Type == "response.incomplete" || wire.Type == "response.failed" {
|
||||
if isResponsesClientToolTerminalEvent(wire.Type) {
|
||||
restored, changed, err := RestoreResponsesClientToolPayload(payload, r.adapter)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
@@ -465,6 +465,15 @@ func (r *ResponsesClientToolStreamRestorer) RestoreEvent(payload []byte) ([][]by
|
||||
return result, true, nil
|
||||
}
|
||||
|
||||
func isResponsesClientToolTerminalEvent(typ string) bool {
|
||||
switch strings.TrimSpace(typ) {
|
||||
case "response.completed", "response.done", "response.incomplete", "response.failed", "response.cancelled", "response.canceled":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (r *ResponsesClientToolStreamRestorer) clientToolEventPayload(payload []byte) bool {
|
||||
var raw struct {
|
||||
ItemID string `json:"item_id"`
|
||||
|
||||
@@ -174,3 +174,30 @@ func TestResponsesClientToolStreamRestorer_RawEventsPreserveUnknownFieldsAndOutp
|
||||
require.Len(t, done, 2)
|
||||
require.Equal(t, "pwd", done[1].Input)
|
||||
}
|
||||
|
||||
func TestResponsesClientToolStreamRestorer_RestoresAllTerminalEvents(t *testing.T) {
|
||||
for _, eventType := range []string{
|
||||
"response.completed",
|
||||
"response.done",
|
||||
"response.incomplete",
|
||||
"response.failed",
|
||||
"response.cancelled",
|
||||
"response.canceled",
|
||||
} {
|
||||
t.Run(eventType, func(t *testing.T) {
|
||||
restorer := NewResponsesClientToolStreamRestorer(ResponsesClientToolMapping{CustomTools: map[string]bool{"exec": true}})
|
||||
payload := []byte(`{"type":"` + eventType + `","sequence_number":7,"response":{"id":"resp_tools","output":[{"type":"function_call","id":"item_exec","call_id":"call_exec","name":"exec","arguments":"{\"input\":\"pwd\"}"}]}}`)
|
||||
|
||||
restored, changed, err := restorer.RestoreEvent(payload)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Len(t, restored, 1)
|
||||
require.Equal(t, eventType, gjson.GetBytes(restored[0], "type").String())
|
||||
require.Equal(t, int64(7), gjson.GetBytes(restored[0], "sequence_number").Int())
|
||||
require.Equal(t, "custom_tool_call", gjson.GetBytes(restored[0], "response.output.0.type").String())
|
||||
require.Equal(t, "pwd", gjson.GetBytes(restored[0], "response.output.0.input").String())
|
||||
require.False(t, gjson.GetBytes(restored[0], "response.output.0.arguments").Exists())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 工具,
|
||||
@@ -675,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.
|
||||
|
||||
@@ -45,6 +45,9 @@ func DefaultModelIDs() []string {
|
||||
// DefaultTestModel default model for testing OpenAI accounts
|
||||
const DefaultTestModel = "gpt-5.4"
|
||||
|
||||
// CodexUsageProbeModel is the model used for OAuth Codex usage probes.
|
||||
const CodexUsageProbeModel = "codex-auto-review"
|
||||
|
||||
// DefaultInstructions default instructions for non-Codex CLI requests.
|
||||
// 内容为真实 Codex CLI 的 GPT-5-Codex base prompt(codex 系模型默认)。
|
||||
//
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
@@ -71,6 +71,39 @@ var schedulerNeutralExtraKeys = map[string]struct{}{
|
||||
|
||||
const postgresParameterBatchSize = 50000
|
||||
|
||||
const codexFingerprintSeedCanonicalPattern = "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
|
||||
const codexFingerprintNilSeed = "00000000-0000-0000-0000-000000000000"
|
||||
|
||||
func codexFingerprintSeedValidSQL(extraExpr string) string {
|
||||
value := "(" + extraExpr + " ->> 'codex_fingerprint_seed')"
|
||||
return "(" + value + " ~ '" + codexFingerprintSeedCanonicalPattern + "' AND " + value + " <> '" + codexFingerprintNilSeed + "')"
|
||||
}
|
||||
|
||||
func ensureCodexFingerprintSeedSQL(extraExpr string) string {
|
||||
return "CASE WHEN platform = 'openai' AND type = 'oauth' THEN " +
|
||||
"jsonb_set(" + extraExpr + ", '{codex_fingerprint_seed}', " +
|
||||
"CASE WHEN " + codexFingerprintSeedValidSQL("extra") +
|
||||
" THEN to_jsonb(extra ->> 'codex_fingerprint_seed') ELSE to_jsonb(gen_random_uuid()::text) END, true) " +
|
||||
"ELSE " + extraExpr + " END"
|
||||
}
|
||||
|
||||
func stripCodexFingerprintSeedFromExtraUpdate(extra map[string]any) map[string]any {
|
||||
if extra == nil {
|
||||
return nil
|
||||
}
|
||||
if _, exists := extra["codex_fingerprint_seed"]; !exists {
|
||||
return extra
|
||||
}
|
||||
stripped := make(map[string]any, len(extra)-1)
|
||||
for key, value := range extra {
|
||||
if key == "codex_fingerprint_seed" {
|
||||
continue
|
||||
}
|
||||
stripped[key] = value
|
||||
}
|
||||
return stripped
|
||||
}
|
||||
|
||||
// NewAccountRepository 创建账户仓储实例。
|
||||
// 这是对外暴露的构造函数,返回接口类型以便于依赖注入。
|
||||
func NewAccountRepository(client *dbent.Client, sqlDB *sql.DB, schedulerCache service.SchedulerCache) service.AccountRepository {
|
||||
@@ -2520,6 +2553,7 @@ func (r *accountRepository) AutoPauseExpiredAccounts(ctx context.Context, now ti
|
||||
}
|
||||
|
||||
func (r *accountRepository) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||||
updates = stripCodexFingerprintSeedFromExtraUpdate(updates)
|
||||
if len(updates) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -2552,6 +2586,9 @@ func (r *accountRepository) UpdateExtra(ctx context.Context, id int64, updates m
|
||||
if clearProbeSnapshot {
|
||||
extraExpression = "(" + extraExpression + ") - 'upstream_billing_probe'"
|
||||
}
|
||||
if service.ShouldEnsureCodexFingerprintSeedForExtraUpdates(updates) {
|
||||
extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
|
||||
}
|
||||
result, err := client.ExecContext(
|
||||
ctx,
|
||||
"UPDATE accounts SET extra = "+extraExpression+", updated_at = NOW() WHERE id = $2 AND deleted_at IS NULL",
|
||||
@@ -2793,6 +2830,7 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
if len(ids) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
updates.Extra = stripCodexFingerprintSeedFromExtraUpdate(updates.Extra)
|
||||
|
||||
setClauses := make([]string, 0, 8)
|
||||
args := make([]any, 0, 8)
|
||||
@@ -2880,7 +2918,7 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
" AND "+ollamaCloudBaseURLMatchesSQL(credentialPlaceholder+"::jsonb ->> 'base_url'")+")")
|
||||
}
|
||||
|
||||
if len(updates.Extra) > 0 || len(ollamaGroupIdentityChanges) > 0 || ollamaProxyIdentityChanged != "" {
|
||||
if len(updates.Extra) > 0 || len(ollamaGroupIdentityChanges) > 0 || ollamaProxyIdentityChanged != "" || updates.EnsureCodexFingerprintSeed {
|
||||
extraExpression := "COALESCE(extra, '{}'::jsonb)"
|
||||
if len(updates.Extra) > 0 {
|
||||
payload, err := json.Marshal(updates.Extra)
|
||||
@@ -2919,6 +2957,9 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
} else if snapshotIdentityChanged != "" {
|
||||
extraExpression = "CASE WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot' ELSE " + extraExpression + " END"
|
||||
}
|
||||
if updates.EnsureCodexFingerprintSeed {
|
||||
extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
|
||||
}
|
||||
setClauses = append(setClauses, "extra = "+extraExpression)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
)
|
||||
|
||||
func TestBulkUpdateEnsuresCodexFingerprintSeedWithPerRowSQL(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{
|
||||
Extra: map[string]any{
|
||||
"codex_fingerprint_mode": "session",
|
||||
"codex_fingerprint_seed": "22222222-2222-4222-8222-222222222222",
|
||||
},
|
||||
EnsureCodexFingerprintSeed: true,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
query := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
require.Contains(t, query, "jsonb_set")
|
||||
require.Contains(t, query, "gen_random_uuid()::text")
|
||||
require.Contains(t, query, "platform = 'openai' AND type = 'oauth'")
|
||||
require.Contains(t, query, "to_jsonb(extra ->> 'codex_fingerprint_seed')")
|
||||
require.Contains(t, query, codexFingerprintSeedCanonicalPattern)
|
||||
require.NotContains(t, query, "22222222-2222-4222-8222-222222222222")
|
||||
require.NotEmpty(t, exec.execArgs)
|
||||
payload, ok := exec.execArgs[0][0].([]byte)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, `{"codex_fingerprint_mode":"session"}`, string(payload))
|
||||
}
|
||||
|
||||
func TestUpdateExtraEnsuresCodexFingerprintSeedAtomicallyWhenEnabling(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .*jsonb_set.*gen_random_uuid\(\)::text.*WHERE id = \$2 AND deleted_at IS NULL`).
|
||||
WithArgs(`{"codex_fingerprint_mode":"device"}`, int64(27)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(27), nil, nil, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
|
||||
err = repo.UpdateExtra(context.Background(), 27, map[string]any{
|
||||
"codex_fingerprint_mode": "device",
|
||||
"codex_fingerprint_seed": "22222222-2222-4222-8222-222222222222",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestBulkUpdateCodexFingerprintSeedRollsBackWhenUpdateFails(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .*gen_random_uuid\(\)::text.*WHERE id = ANY\(\$2\)`).
|
||||
WithArgs(sqlmock.AnyArg(), `{27,28}`).
|
||||
WillReturnError(errors.New("update failed"))
|
||||
mock.ExpectRollback()
|
||||
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{
|
||||
Extra: map[string]any{
|
||||
"codex_fingerprint_mode": "session",
|
||||
},
|
||||
EnsureCodexFingerprintSeed: true,
|
||||
})
|
||||
|
||||
require.EqualError(t, err, "update failed")
|
||||
require.Zero(t, rows)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestBulkUpdateCodexFingerprintSeedRollsBackWhenOutboxFails(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .*gen_random_uuid\(\)::text.*WHERE id = ANY\(\$2\)`).
|
||||
WithArgs(sqlmock.AnyArg(), `{27,28}`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 2))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WillReturnError(errors.New("outbox failed"))
|
||||
mock.ExpectRollback()
|
||||
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{
|
||||
Extra: map[string]any{
|
||||
"codex_fingerprint_mode": "full",
|
||||
},
|
||||
EnsureCodexFingerprintSeed: true,
|
||||
})
|
||||
|
||||
require.EqualError(t, err, "outbox failed")
|
||||
require.Zero(t, rows)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -16,7 +16,7 @@ import (
|
||||
|
||||
func (r *channelRepository) ListModelPricing(ctx context.Context, channelID int64) ([]service.ChannelModelPricing, error) {
|
||||
rows, err := r.db.QueryContext(ctx,
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, created_at, updated_at
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
|
||||
FROM channel_model_pricing WHERE channel_id = $1 ORDER BY id`, channelID,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -51,16 +51,20 @@ func (r *channelRepository) UpdateModelPricing(ctx context.Context, pricing *ser
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal models: %w", err)
|
||||
}
|
||||
timePricingJSON, err := marshalChannelTimePricing(pricing.TimePricing)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
billingMode := pricing.BillingMode
|
||||
if billingMode == "" {
|
||||
billingMode = service.BillingModeToken
|
||||
}
|
||||
result, err := r.db.ExecContext(ctx,
|
||||
`UPDATE channel_model_pricing
|
||||
SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_read_price = $6, image_input_price = $7, image_output_price = $8, per_request_price = $9, platform = $10, updated_at = NOW()
|
||||
WHERE id = $11`,
|
||||
SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_read_price = $6, image_input_price = $7, image_output_price = $8, per_request_price = $9, time_pricing = $10, platform = $11, updated_at = NOW()
|
||||
WHERE id = $12`,
|
||||
modelsJSON, billingMode, pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheReadPrice,
|
||||
pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice, pricing.Platform, pricing.ID,
|
||||
pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice, timePricingJSON, pricing.Platform, pricing.ID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("update model pricing: %w", err)
|
||||
@@ -91,7 +95,7 @@ func (r *channelRepository) ReplaceModelPricing(ctx context.Context, channelID i
|
||||
// batchLoadModelPricing 批量加载多个渠道的模型定价(含区间)
|
||||
func (r *channelRepository) batchLoadModelPricing(ctx context.Context, channelIDs []int64) (map[int64][]service.ChannelModelPricing, error) {
|
||||
rows, err := r.db.QueryContext(ctx,
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, created_at, updated_at
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
|
||||
FROM channel_model_pricing WHERE channel_id = ANY($1) ORDER BY channel_id, id`,
|
||||
pq.Array(channelIDs),
|
||||
)
|
||||
@@ -169,16 +173,22 @@ func scanModelPricingRows(rows *sql.Rows) ([]service.ChannelModelPricing, []int6
|
||||
for rows.Next() {
|
||||
var p service.ChannelModelPricing
|
||||
var modelsJSON []byte
|
||||
var timePricingJSON []byte
|
||||
if err := rows.Scan(
|
||||
&p.ID, &p.ChannelID, &p.Platform, &modelsJSON, &p.BillingMode,
|
||||
&p.InputPrice, &p.OutputPrice, &p.CacheWritePrice, &p.CacheReadPrice,
|
||||
&p.ImageInputPrice, &p.ImageOutputPrice, &p.PerRequestPrice, &p.CreatedAt, &p.UpdatedAt,
|
||||
&p.ImageInputPrice, &p.ImageOutputPrice, &p.PerRequestPrice, &timePricingJSON, &p.CreatedAt, &p.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, nil, fmt.Errorf("scan model pricing: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(modelsJSON, &p.Models); err != nil {
|
||||
p.Models = []string{}
|
||||
}
|
||||
timePricing, err := unmarshalChannelTimePricing(timePricingJSON)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
p.TimePricing = timePricing
|
||||
pricingIDs = append(pricingIDs, p.ID)
|
||||
result = append(result, p)
|
||||
}
|
||||
@@ -220,6 +230,10 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal models: %w", err)
|
||||
}
|
||||
timePricingJSON, err := marshalChannelTimePricing(pricing.TimePricing)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
billingMode := pricing.BillingMode
|
||||
if billingMode == "" {
|
||||
billingMode = service.BillingModeToken
|
||||
@@ -229,11 +243,11 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
|
||||
platform = "anthropic"
|
||||
}
|
||||
err = exec.QueryRowContext(ctx,
|
||||
`INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) RETURNING id, created_at, updated_at`,
|
||||
`INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, time_pricing)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING id, created_at, updated_at`,
|
||||
pricing.ChannelID, platform, modelsJSON, billingMode,
|
||||
pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheReadPrice,
|
||||
pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice,
|
||||
pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice, timePricingJSON,
|
||||
).Scan(&pricing.ID, &pricing.CreatedAt, &pricing.UpdatedAt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("insert model pricing: %w", err)
|
||||
@@ -249,6 +263,28 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
|
||||
return nil
|
||||
}
|
||||
|
||||
func marshalChannelTimePricing(config *service.ChannelTimePricing) (any, error) {
|
||||
if config == nil || len(config.Periods) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
data, err := json.Marshal(config)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal time pricing: %w", err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func unmarshalChannelTimePricing(data []byte) (*service.ChannelTimePricing, error) {
|
||||
if len(data) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var config service.ChannelTimePricing
|
||||
if err := json.Unmarshal(data, &config); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal time pricing: %w", err)
|
||||
}
|
||||
return &config, nil
|
||||
}
|
||||
|
||||
func createIntervalExec(ctx context.Context, exec dbExec, iv *service.PricingInterval) error {
|
||||
return exec.QueryRowContext(ctx,
|
||||
`INSERT INTO channel_pricing_intervals
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
var channelModelPricingTimePricingColumns = []string{
|
||||
"id", "channel_id", "platform", "models", "billing_mode", "input_price", "output_price",
|
||||
"cache_write_price", "cache_read_price", "image_input_price", "image_output_price",
|
||||
"per_request_price", "time_pricing", "created_at", "updated_at",
|
||||
}
|
||||
|
||||
const channelModelPricingTimePricingJSON = `{"timezone":"Asia/Shanghai","periods":[{"start_time":"09:00","end_time":"12:00","multiplier":2}]}`
|
||||
|
||||
func newChannelModelPricingTimePricingRepo(t *testing.T) (*channelRepository, sqlmock.Sqlmock) {
|
||||
t.Helper()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return &channelRepository{db: db}, mock
|
||||
}
|
||||
|
||||
func modelPricingTimePricingRow(timePricing any) *sqlmock.Rows {
|
||||
return sqlmock.NewRows(channelModelPricingTimePricingColumns).AddRow(
|
||||
int64(11), int64(7), "openai", `["gpt-5"]`, service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, timePricing,
|
||||
time.Date(2026, 8, 17, 0, 0, 0, 0, time.UTC), time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
|
||||
)
|
||||
}
|
||||
|
||||
func expectEmptyModelPricingIntervals(mock sqlmock.Sqlmock) {
|
||||
mock.ExpectQuery(`SELECT id, pricing_id, min_tokens, max_tokens, tier_label`).
|
||||
WithArgs(sqlmock.AnyArg()).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
}
|
||||
|
||||
func TestChannelModelPricingTimePricingListRoundTrip(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(`(?s)SELECT .*per_request_price, time_pricing, created_at, updated_at.*FROM channel_model_pricing.*channel_id = \$1`).
|
||||
WithArgs(int64(7)).
|
||||
WillReturnRows(modelPricingTimePricingRow(channelModelPricingTimePricingJSON))
|
||||
expectEmptyModelPricingIntervals(mock)
|
||||
|
||||
pricing, err := repo.ListModelPricing(context.Background(), 7)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, pricing, 1)
|
||||
require.NotNil(t, pricing[0].TimePricing)
|
||||
require.Equal(t, "Asia/Shanghai", pricing[0].TimePricing.Timezone)
|
||||
require.Len(t, pricing[0].TimePricing.Periods, 1)
|
||||
require.Equal(t, 2.0, pricing[0].TimePricing.Periods[0].Multiplier)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestChannelModelPricingTimePricingListNullAndMalformed(t *testing.T) {
|
||||
t.Run("SQL NULL maps to nil", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(`(?s)SELECT .*per_request_price, time_pricing, created_at, updated_at.*FROM channel_model_pricing.*channel_id = \$1`).
|
||||
WithArgs(int64(7)).
|
||||
WillReturnRows(modelPricingTimePricingRow(nil))
|
||||
expectEmptyModelPricingIntervals(mock)
|
||||
|
||||
pricing, err := repo.ListModelPricing(context.Background(), 7)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, pricing, 1)
|
||||
require.Nil(t, pricing[0].TimePricing)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
|
||||
t.Run("malformed JSON returns repository error", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(`(?s)SELECT .*per_request_price, time_pricing, created_at, updated_at.*FROM channel_model_pricing.*channel_id = \$1`).
|
||||
WithArgs(int64(7)).
|
||||
WillReturnRows(modelPricingTimePricingRow(`{"timezone":`))
|
||||
|
||||
_, err := repo.ListModelPricing(context.Background(), 7)
|
||||
require.Error(t, err)
|
||||
require.True(t, strings.Contains(err.Error(), "unmarshal time pricing"), "unexpected error: %v", err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
|
||||
func TestChannelModelPricingTimePricingCreateAndUpdateRoundTrip(t *testing.T) {
|
||||
pricing := &service.ChannelModelPricing{
|
||||
ID: 11,
|
||||
ChannelID: 7,
|
||||
Platform: "openai",
|
||||
Models: []string{"gpt-5"},
|
||||
TimePricing: &service.ChannelTimePricing{
|
||||
Timezone: "Asia/Shanghai",
|
||||
Periods: []service.ChannelTimePricingPeriod{{
|
||||
StartTime: "09:00", EndTime: "12:00", Multiplier: 2,
|
||||
}},
|
||||
},
|
||||
}
|
||||
|
||||
t.Run("create writes JSON", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, time_pricing)")).
|
||||
WithArgs(
|
||||
int64(7), "openai", []byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON,
|
||||
).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
|
||||
|
||||
require.NoError(t, repo.CreateModelPricing(context.Background(), pricing))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
|
||||
t.Run("update writes JSON and entry ID", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$9, time_pricing = \$10, platform = \$11.*WHERE id = \$12`).
|
||||
WithArgs(
|
||||
[]byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON, "openai", int64(11),
|
||||
).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
require.NoError(t, repo.UpdateModelPricing(context.Background(), pricing))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
|
||||
func TestChannelModelPricingTimePricingCreateAndUpdateWriteNullWhenDisabled(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
timePricing *service.ChannelTimePricing
|
||||
}{
|
||||
{name: "nil", timePricing: nil},
|
||||
{name: "empty periods", timePricing: &service.ChannelTimePricing{Timezone: "Asia/Shanghai"}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
newPricing := func() *service.ChannelModelPricing {
|
||||
return &service.ChannelModelPricing{
|
||||
ID: 11,
|
||||
ChannelID: 7,
|
||||
Platform: "openai",
|
||||
Models: []string{"gpt-5"},
|
||||
TimePricing: tt.timePricing,
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("create writes SQL NULL", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_read_price, image_input_price, image_output_price, per_request_price, time_pricing)")).
|
||||
WithArgs(
|
||||
int64(7), "openai", []byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
|
||||
|
||||
require.NoError(t, repo.CreateModelPricing(context.Background(), newPricing()))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
|
||||
t.Run("update writes SQL NULL", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$9, time_pricing = \$10, platform = \$11.*WHERE id = \$12`).
|
||||
WithArgs(
|
||||
[]byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, "openai", int64(11),
|
||||
).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
require.NoError(t, repo.UpdateModelPricing(context.Background(), newPricing()))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
dbmigrations "github.com/Wei-Shaw/sub2api/migrations"
|
||||
"github.com/google/uuid"
|
||||
"github.com/lib/pq"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func requireCanonicalUUIDString(t *testing.T, value string) {
|
||||
t.Helper()
|
||||
parsed, err := uuid.Parse(value)
|
||||
require.NoError(t, err)
|
||||
require.NotEqual(t, uuid.Nil, parsed)
|
||||
require.Equal(t, parsed.String(), value)
|
||||
}
|
||||
|
||||
func TestMigration225BackfillsOnlyEnabledOpenAIOAuthMissingOrMalformedSeeds(t *testing.T) {
|
||||
tx := testTx(t)
|
||||
ctx := context.Background()
|
||||
migrationSQL, err := dbmigrations.FS.ReadFile("225_backfill_codex_fingerprint_seed.sql")
|
||||
require.NoError(t, err)
|
||||
|
||||
var missingID, blankID, malformedID, validID, offID, apiKeyID int64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ('migration-225-missing', 'openai', 'oauth', '{"codex_fingerprint_mode":"session"}'::jsonb)
|
||||
RETURNING id
|
||||
`).Scan(&missingID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ('migration-225-blank', 'openai', 'oauth', '{"codex_fingerprint_mode":"device","codex_fingerprint_seed":""}'::jsonb)
|
||||
RETURNING id
|
||||
`).Scan(&blankID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ('migration-225-malformed', 'openai', 'oauth', '{"codex_fingerprint_mode":"full","codex_fingerprint_seed":"BAD"}'::jsonb)
|
||||
RETURNING id
|
||||
`).Scan(&malformedID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ('migration-225-valid', 'openai', 'oauth', '{"codex_fingerprint_mode":"session","codex_fingerprint_seed":"11111111-1111-4111-8111-111111111111"}'::jsonb)
|
||||
RETURNING id
|
||||
`).Scan(&validID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ('migration-225-off', 'openai', 'oauth', '{"codex_fingerprint_mode":"off"}'::jsonb)
|
||||
RETURNING id
|
||||
`).Scan(&offID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ('migration-225-apikey', 'openai', 'apikey', '{"codex_fingerprint_mode":"session"}'::jsonb)
|
||||
RETURNING id
|
||||
`).Scan(&apiKeyID))
|
||||
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
|
||||
seedsAfterFirst := map[int64]string{}
|
||||
for _, id := range []int64{missingID, blankID, malformedID, validID} {
|
||||
var seed string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT extra->>'codex_fingerprint_seed' FROM accounts WHERE id = $1`, id).Scan(&seed))
|
||||
requireCanonicalUUIDString(t, seed)
|
||||
seedsAfterFirst[id] = seed
|
||||
}
|
||||
require.Equal(t, "11111111-1111-4111-8111-111111111111", seedsAfterFirst[validID])
|
||||
|
||||
for _, id := range []int64{offID, apiKeyID} {
|
||||
var hasSeed bool
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT extra ? 'codex_fingerprint_seed' FROM accounts WHERE id = $1`, id).Scan(&hasSeed))
|
||||
require.False(t, hasSeed)
|
||||
}
|
||||
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
|
||||
for id, want := range seedsAfterFirst {
|
||||
var got string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT extra->>'codex_fingerprint_seed' FROM accounts WHERE id = $1`, id).Scan(&got))
|
||||
require.Equal(t, want, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBulkUpdateGeneratesDistinctStableCodexFingerprintSeedsPerEligibleRow(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
testName := "bulk-codex-seed-" + uuid.NewString()
|
||||
type fixture struct {
|
||||
name string
|
||||
accountType string
|
||||
extra string
|
||||
}
|
||||
fixtures := []fixture{
|
||||
{name: testName + "-missing-a", accountType: service.AccountTypeOAuth, extra: `{}`},
|
||||
{name: testName + "-missing-b", accountType: service.AccountTypeOAuth, extra: `{"codex_fingerprint_seed":"BAD"}`},
|
||||
{name: testName + "-valid", accountType: service.AccountTypeOAuth, extra: `{"codex_fingerprint_seed":"11111111-1111-4111-8111-111111111111"}`},
|
||||
{name: testName + "-apikey", accountType: service.AccountTypeAPIKey, extra: `{}`},
|
||||
}
|
||||
|
||||
ids := make([]int64, 0, len(fixtures))
|
||||
for _, f := range fixtures {
|
||||
var id int64
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, `
|
||||
INSERT INTO accounts (name, platform, type, extra)
|
||||
VALUES ($1, 'openai', $2, $3::jsonb)
|
||||
RETURNING id
|
||||
`, f.name, f.accountType, f.extra).Scan(&id))
|
||||
ids = append(ids, id)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(context.Background(), `DELETE FROM scheduler_outbox WHERE account_id = ANY($1)`, pq.Array(ids))
|
||||
_, _ = integrationDB.ExecContext(context.Background(), `DELETE FROM accounts WHERE id = ANY($1)`, pq.Array(ids))
|
||||
})
|
||||
|
||||
repo := newAccountRepositoryWithSQL(testEntClient(t), integrationDB, nil)
|
||||
updates := service.AccountBulkUpdate{
|
||||
Extra: map[string]any{
|
||||
"codex_fingerprint_mode": "session",
|
||||
},
|
||||
EnsureCodexFingerprintSeed: true,
|
||||
}
|
||||
rows, err := repo.BulkUpdate(ctx, ids, updates)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(len(ids)), rows)
|
||||
|
||||
readSeed := func(id int64) string {
|
||||
t.Helper()
|
||||
var seed string
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, `SELECT COALESCE(extra->>'codex_fingerprint_seed', '') FROM accounts WHERE id = $1`, id).Scan(&seed))
|
||||
return seed
|
||||
}
|
||||
firstSeeds := []string{readSeed(ids[0]), readSeed(ids[1]), readSeed(ids[2]), readSeed(ids[3])}
|
||||
requireCanonicalUUIDString(t, firstSeeds[0])
|
||||
requireCanonicalUUIDString(t, firstSeeds[1])
|
||||
require.NotEqual(t, firstSeeds[0], firstSeeds[1], "gen_random_uuid must be evaluated per eligible row")
|
||||
require.Equal(t, "11111111-1111-4111-8111-111111111111", firstSeeds[2])
|
||||
require.Empty(t, firstSeeds[3], "API-key accounts must not receive a Codex fingerprint seed")
|
||||
|
||||
rows, err = repo.BulkUpdate(ctx, ids, updates)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(len(ids)), rows)
|
||||
for i, want := range firstSeeds {
|
||||
require.Equal(t, want, readSeed(ids[i]), "retry must not rotate an existing valid seed")
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -46,9 +46,11 @@ func (s *openaiOAuthService) ExchangeCode(ctx context.Context, code, codeVerifie
|
||||
|
||||
var tokenResp openai.TokenResponse
|
||||
|
||||
authUA, authOriginator := service.CodexCanonicalAuthIdentity()
|
||||
resp, err := client.R().
|
||||
SetContext(ctx).
|
||||
SetHeader("User-Agent", "codex-cli/0.91.0").
|
||||
SetHeader("User-Agent", authUA).
|
||||
SetHeader("originator", authOriginator).
|
||||
SetFormDataFromValues(formData).
|
||||
SetSuccessResult(&tokenResp).
|
||||
Post(s.tokenURL)
|
||||
@@ -94,9 +96,11 @@ func (s *openaiOAuthService) refreshTokenWithClientID(ctx context.Context, refre
|
||||
|
||||
var tokenResp openai.TokenResponse
|
||||
|
||||
authUA, authOriginator := service.CodexCanonicalAuthIdentity()
|
||||
resp, err := client.R().
|
||||
SetContext(ctx).
|
||||
SetHeader("User-Agent", "codex-cli/0.91.0").
|
||||
SetHeader("User-Agent", authUA).
|
||||
SetHeader("originator", authOriginator).
|
||||
SetFormDataFromValues(formData).
|
||||
SetSuccessResult(&tokenResp).
|
||||
Post(s.tokenURL)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
@@ -77,6 +78,17 @@ func (s *OpenAIOAuthServiceSuite) TestExchangeCode_DefaultRedirectURI() {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
wantUA, wantOriginator := service.CodexCanonicalAuthIdentity()
|
||||
if got := r.Header.Get("User-Agent"); got != wantUA {
|
||||
errCh <- "user-agent mismatch"
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if got := r.Header.Get("originator"); got != wantOriginator {
|
||||
errCh <- "originator mismatch"
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"access_token":"at","refresh_token":"rt","token_type":"bearer","expires_in":3600}`)
|
||||
@@ -121,6 +133,17 @@ func (s *OpenAIOAuthServiceSuite) TestRefreshToken_FormFields() {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
wantUA, wantOriginator := service.CodexCanonicalAuthIdentity()
|
||||
if got := r.Header.Get("User-Agent"); got != wantUA {
|
||||
errCh <- "user-agent mismatch"
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if got := r.Header.Get("originator"); got != wantOriginator {
|
||||
errCh <- "originator mismatch"
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"access_token":"at2","refresh_token":"rt2","token_type":"bearer","expires_in":3600}`)
|
||||
|
||||
@@ -1000,6 +1000,8 @@ func filterSchedulerExtra(extra map[string]any) map[string]any {
|
||||
"openai_ws_force_http",
|
||||
"openai_responses_mode",
|
||||
"openai_responses_supported",
|
||||
"codex_fingerprint_mode",
|
||||
"codex_fingerprint_seed",
|
||||
"codex_5h_used_percent",
|
||||
"codex_7d_used_percent",
|
||||
"codex_5h_reset_at",
|
||||
|
||||
@@ -309,6 +309,8 @@ func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
|
||||
"openai_ws_force_http": true,
|
||||
"openai_responses_mode": "force_chat_completions",
|
||||
"openai_responses_supported": false,
|
||||
"codex_fingerprint_mode": "session",
|
||||
"codex_fingerprint_seed": "11111111-1111-4111-8111-111111111111",
|
||||
"mixed_scheduling": true,
|
||||
"unused_large_field": "drop-me",
|
||||
},
|
||||
@@ -321,6 +323,8 @@ func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
|
||||
require.Equal(t, true, got.Extra["openai_ws_force_http"])
|
||||
require.Equal(t, "force_chat_completions", got.Extra["openai_responses_mode"])
|
||||
require.Equal(t, false, got.Extra["openai_responses_supported"])
|
||||
require.Equal(t, "session", got.Extra["codex_fingerprint_mode"])
|
||||
require.Equal(t, "11111111-1111-4111-8111-111111111111", got.Extra["codex_fingerprint_seed"])
|
||||
require.Equal(t, true, got.Extra["mixed_scheduling"])
|
||||
require.Nil(t, got.Extra["unused_large_field"])
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -98,6 +98,35 @@ func TestUserPlatformQuotaRepository_BulkInsertInitial_GrokAllowed(t *testing.T)
|
||||
require.InDelta(t, 9.0, *rec.DailyLimitUSD, 1e-9)
|
||||
}
|
||||
|
||||
// TestUserPlatformQuotaRepository_BulkInsertInitial_CNProvidersAllowed 回归迁移 224:
|
||||
// kimi/zhipu/deepseek 平台必须能写入 user_platform_quotas(CHECK 约束已含国产供应商)。
|
||||
// 历史 bug:三个平台不在约束内 → 注册预填充 8 平台默认配额时整条多行 INSERT 中止 →
|
||||
// fail-open 吞错 → 新用户拿到零条配额记录(缺失配额行 = 无限额)。
|
||||
func TestUserPlatformQuotaRepository_BulkInsertInitial_CNProvidersAllowed(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
client := tx.Client()
|
||||
|
||||
userID := mustCreateUserForQuota(t, client)
|
||||
repo := NewUserPlatformQuotaRepository(client)
|
||||
|
||||
daily := 12.0
|
||||
records := []UserPlatformQuotaRecord{
|
||||
{UserID: userID, Platform: "kimi", DailyLimitUSD: &daily},
|
||||
{UserID: userID, Platform: "zhipu"},
|
||||
{UserID: userID, Platform: "deepseek"},
|
||||
}
|
||||
require.NoError(t, repo.BulkInsertInitial(txCtx, records),
|
||||
"kimi/zhipu/deepseek 平台应可写入(迁移 224 后 CHECK 约束已含国产供应商)")
|
||||
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek"} {
|
||||
rec, err := repo.GetByUserPlatform(txCtx, userID, platform)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, rec, "%s 配额行应已写入", platform)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserPlatformQuotaRepository_GetByUserPlatform(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
|
||||
@@ -70,23 +70,30 @@ func (r *userRepository) create(ctx context.Context, userIn *service.User, guard
|
||||
|
||||
// 统一使用 ent 的事务:保证用户与允许分组的更新原子化,
|
||||
// 并避免基于 *sql.Tx 手动构造 ent client 导致的 ExecQuerier 断言错误。
|
||||
tx, err := r.client.Tx(ctx)
|
||||
if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
|
||||
return err
|
||||
}
|
||||
|
||||
//
|
||||
// 注意:ent 的 Client.Tx 不感知上下文中的事务(只检查 driver 类型),
|
||||
// 因此必须显式检查 TxFromContext:当调用方已开启外部事务(如注册时的
|
||||
// “建用户 + 占用邀请码”原子事务),直接复用其 client,由调用方统一提交/回滚,
|
||||
// 否则用户写入会落入独立事务并自行提交,导致外层事务无法回滚(孤儿用户)。
|
||||
var txClient *dbent.Client
|
||||
txCtx := ctx
|
||||
if err == nil {
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
txClient = tx.Client()
|
||||
txCtx = dbent.NewTxContext(ctx, tx)
|
||||
var ownedTx *dbent.Tx
|
||||
if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
|
||||
txClient = existingTx.Client()
|
||||
} else {
|
||||
// 已处于外部事务中(ErrTxStarted),复用当前事务 client 并由调用方负责提交/回滚。
|
||||
if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
|
||||
txClient = existingTx.Client()
|
||||
} else {
|
||||
tx, err := r.client.Tx(ctx)
|
||||
switch {
|
||||
case errors.Is(err, dbent.ErrTxStarted):
|
||||
// r.client 本身已是事务绑定 client(client 注入式事务,如集成测试
|
||||
// 夹具 tx.Client()):直接复用,提交/回滚由 client 的持有方负责。
|
||||
txClient = r.client
|
||||
case err != nil:
|
||||
return err
|
||||
default:
|
||||
ownedTx = tx
|
||||
defer func() { _ = ownedTx.Rollback() }()
|
||||
txClient = tx.Client()
|
||||
txCtx = dbent.NewTxContext(ctx, tx)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -158,8 +165,8 @@ func (r *userRepository) create(ctx context.Context, userIn *service.User, guard
|
||||
return err
|
||||
}
|
||||
|
||||
if tx != nil {
|
||||
if err := tx.Commit(); err != nil {
|
||||
if ownedTx != nil {
|
||||
if err := ownedTx.Commit(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/ent/redeemcode"
|
||||
"github.com/Wei-Shaw/sub2api/ent/user"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestCreateWithEmailAliasGuardJoinsOuterTransaction 验证用户创建会加入调用方开启的
|
||||
// 外部 ent 事务(注册流程“建用户 + 占用邀请码”原子性的基础):
|
||||
// - 外层事务回滚后,用户与邀请码占用必须一并撤销(不得残留孤儿账号);
|
||||
// - 外层事务提交后,用户与邀请码占用同时生效。
|
||||
//
|
||||
// 回归背景:此前 create() 通过 r.client.Tx(ctx) 自开事务且自行 Commit —— ent 的
|
||||
// Client.Tx 不感知上下文事务(只检查 driver 类型),ErrTxStarted 分支实际是死代码,
|
||||
// 导致外层事务包不住用户写入;并发注册同一邀请码时,败者用户的写入已自行提交,
|
||||
// 即使注册被拒绝也会留下可登录的孤儿账号(1 个邀请码仍可生成任意数量账号)。
|
||||
func TestCreateWithEmailAliasGuardJoinsOuterTransaction(t *testing.T) {
|
||||
client := testEntClient(t)
|
||||
userRepo := NewUserRepository(client, integrationDB)
|
||||
redeemRepo := NewRedeemCodeRepository(client)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// 清理:本测试会真实提交少量数据,确保不影响同包其它集成测试。
|
||||
var committedUserEmails []string
|
||||
var seededCodeIDs []int64
|
||||
t.Cleanup(func() {
|
||||
if len(committedUserEmails) > 0 {
|
||||
_, _ = client.User.Delete().Where(user.EmailIn(committedUserEmails...)).Exec(ctx)
|
||||
}
|
||||
if len(seededCodeIDs) > 0 {
|
||||
_, _ = client.RedeemCode.Delete().Where(redeemcode.IDIn(seededCodeIDs...)).Exec(ctx)
|
||||
}
|
||||
})
|
||||
|
||||
seedCode := func(code string) int64 {
|
||||
_, err := client.RedeemCode.Create().
|
||||
SetCode(code).
|
||||
SetType(service.RedeemTypeInvitation).
|
||||
SetStatus(service.StatusUnused).
|
||||
SetValue(0).
|
||||
Save(ctx)
|
||||
require.NoError(t, err, "seed redeem code")
|
||||
c, err := client.RedeemCode.Query().Where(redeemcode.CodeEQ(code)).Only(ctx)
|
||||
require.NoError(t, err)
|
||||
seededCodeIDs = append(seededCodeIDs, c.ID)
|
||||
return c.ID
|
||||
}
|
||||
|
||||
t.Run("rollback removes user and releases claim", func(t *testing.T) {
|
||||
codeID := seedCode("ITX-RACE-ROLLBACK-001")
|
||||
tx, err := client.Tx(ctx)
|
||||
require.NoError(t, err)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
|
||||
u := &service.User{
|
||||
Email: "itx-rollback@example.com",
|
||||
PasswordHash: "test-password-hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 0,
|
||||
Concurrency: 1,
|
||||
}
|
||||
require.NoError(t, userRepo.CreateWithEmailAliasGuard(txCtx, u))
|
||||
require.Greater(t, u.ID, int64(0), "create 应回填用户 ID")
|
||||
require.NoError(t, redeemRepo.Use(txCtx, codeID, u.ID))
|
||||
require.NoError(t, tx.Rollback())
|
||||
|
||||
exists, err := userRepo.ExistsByEmail(ctx, "itx-rollback@example.com")
|
||||
require.NoError(t, err)
|
||||
require.False(t, exists, "回滚后不得残留孤儿用户")
|
||||
|
||||
after, err := client.RedeemCode.Get(ctx, codeID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, service.StatusUnused, after.Status, "回滚后邀请码应保持 unused")
|
||||
})
|
||||
|
||||
t.Run("commit persists user and claim together", func(t *testing.T) {
|
||||
codeID := seedCode("ITX-RACE-COMMIT-001")
|
||||
tx, err := client.Tx(ctx)
|
||||
require.NoError(t, err)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
|
||||
u := &service.User{
|
||||
Email: "itx-commit@example.com",
|
||||
PasswordHash: "test-password-hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 0,
|
||||
Concurrency: 1,
|
||||
}
|
||||
require.NoError(t, userRepo.CreateWithEmailAliasGuard(txCtx, u))
|
||||
require.NoError(t, redeemRepo.Use(txCtx, codeID, u.ID))
|
||||
require.NoError(t, tx.Commit())
|
||||
committedUserEmails = append(committedUserEmails, u.Email)
|
||||
|
||||
exists, err := userRepo.ExistsByEmail(ctx, "itx-commit@example.com")
|
||||
require.NoError(t, err)
|
||||
require.True(t, exists, "提交后用户应存在")
|
||||
|
||||
after, err := client.RedeemCode.Get(ctx, codeID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, service.StatusUsed, after.Status, "提交后邀请码应为 used")
|
||||
})
|
||||
}
|
||||
@@ -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,
|
||||
@@ -860,7 +861,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"force_email_on_third_party_signup": false,
|
||||
"default_concurrency": 5,
|
||||
"default_balance": 1.25,
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"auth_source_default_email_platform_quotas": null,
|
||||
"auth_source_default_github_platform_quotas": null,
|
||||
"auth_source_default_google_platform_quotas": null,
|
||||
@@ -1174,7 +1175,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"purchase_subscription_url": "",
|
||||
"table_default_page_size": 20,
|
||||
"table_page_size_options": [10, 20, 50],
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"auth_source_default_email_platform_quotas": null,
|
||||
"auth_source_default_github_platform_quotas": null,
|
||||
"auth_source_default_google_platform_quotas": 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)
|
||||
|
||||
@@ -58,6 +58,9 @@ func RegisterAdminRoutes(
|
||||
// Grok OAuth
|
||||
registerGrokOAuthRoutes(admin, h)
|
||||
|
||||
// 国产供应商(kimi/zhipu/deepseek)额度与余额
|
||||
registerCNProviderRoutes(admin, h)
|
||||
|
||||
// 代理管理
|
||||
registerProxyRoutes(admin, h, stepUpAuth)
|
||||
|
||||
@@ -482,6 +485,17 @@ func registerGrokOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
}
|
||||
}
|
||||
|
||||
// registerCNProviderRoutes 注册国产供应商(kimi/zhipu/deepseek)的额度与余额查询端点。
|
||||
func registerCNProviderRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
cn := admin.Group("/cn-providers")
|
||||
{
|
||||
// Coding Plan 滚动窗口用量(kimi/zhipu coding 账号)。
|
||||
cn.GET("/accounts/:id/quota", h.Admin.CNProvider.QueryQuota)
|
||||
// payg 账号余额(kimi/deepseek;zhipu 无余额端点)。
|
||||
cn.GET("/accounts/:id/balance", h.Admin.CNProvider.QueryBalance)
|
||||
}
|
||||
}
|
||||
|
||||
func registerProxyRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAuth middleware.StepUpAuthMiddleware) {
|
||||
proxies := admin.Group("/proxies")
|
||||
{
|
||||
|
||||
@@ -47,7 +47,9 @@ func RegisterGatewayRoutes(
|
||||
|
||||
isOpenAIResponsesCompatibleGatewayPlatform := func(c *gin.Context) bool {
|
||||
switch getGroupPlatform(c) {
|
||||
case service.PlatformOpenAI, service.PlatformGrok:
|
||||
case service.PlatformOpenAI, service.PlatformGrok,
|
||||
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek:
|
||||
// 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)与 openai/grok 一样经 OpenAI 网关转发。
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
@@ -58,7 +60,7 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
countTokensHandler := func(c *gin.Context) {
|
||||
switch getGroupPlatform(c) {
|
||||
case service.PlatformOpenAI:
|
||||
case service.PlatformOpenAI, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek:
|
||||
h.OpenAIGateway.CountTokens(c)
|
||||
case service.PlatformGrok:
|
||||
h.OpenAIGateway.GrokCountTokens(c)
|
||||
@@ -315,6 +317,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 +453,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",
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user