mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc893dd0b8 | ||
|
|
f18f3143e2 | ||
|
|
fead4c7ec3 | ||
|
|
d96b6a31ff | ||
|
|
ab73bc0c77 | ||
|
|
3ce8efc125 | ||
|
|
cc84cd8b4c | ||
|
|
95590b5530 | ||
|
|
b765a7f9f6 | ||
|
|
b72d487b85 | ||
|
|
ad34f89152 | ||
|
|
de6b189a6b | ||
|
|
a74e11c26a | ||
|
|
131d42d25d | ||
|
|
031c83b7e0 | ||
|
|
7a3fda57c8 | ||
|
|
16365199aa | ||
|
|
91a2281c7a | ||
|
|
ece9517091 | ||
|
|
4cc88e27b6 | ||
|
|
b468e428e9 | ||
|
|
a40d6de12e | ||
|
|
bc9173be15 | ||
|
|
eb6e3d1f1d | ||
|
|
a93bfb6623 | ||
|
|
8f47bd5fa0 | ||
|
|
beeb4b84ed | ||
|
|
aac44473aa | ||
|
|
465362e1af | ||
|
|
6ee2304dcd | ||
|
|
e94383a4c4 | ||
|
|
7d3a896fcd | ||
|
|
7dde9370e4 | ||
|
|
5c471485ab | ||
|
|
7b3ed2a961 | ||
|
|
be65c713ff | ||
|
|
78f78947f1 | ||
|
|
005a5d2a37 | ||
|
|
1f45c99de7 | ||
|
|
3e08106116 | ||
|
|
2db0cbd292 | ||
|
|
71d7f86883 | ||
|
|
7af28ca843 | ||
|
|
16dd3d8ee6 | ||
|
|
56b5f0df68 | ||
|
|
7ce6e8d652 | ||
|
|
a36222748d | ||
|
|
2447c44f85 | ||
|
|
fd7e2039d3 | ||
|
|
ba5fa6a38f | ||
|
|
6d99e668d2 | ||
|
|
0875143d98 | ||
|
|
d11b838702 | ||
|
|
6c7625800a | ||
|
|
1850e00955 | ||
|
|
291a737422 | ||
|
|
1614ae9c99 | ||
|
|
2730c1c43b | ||
|
|
0b5903d458 | ||
|
|
e9a58c1cb8 | ||
|
|
ef0ca5bdf5 | ||
|
|
bc3acd6e28 | ||
|
|
b6f9277515 | ||
|
|
2e2638c01d | ||
|
|
1763db3a2d | ||
|
|
bb0c38306f | ||
|
|
b403f88f51 | ||
|
|
f1eaed31ed | ||
|
|
0f72b7dca7 | ||
|
|
43d4bae246 | ||
|
|
5374ce2a0f | ||
|
|
7acc13a29b | ||
|
|
b99ea955d9 | ||
|
|
e6a3030beb | ||
|
|
3176cdee17 | ||
|
|
6b267e4f40 | ||
|
|
f06c8c5698 | ||
|
|
8bed40c67b | ||
|
|
a29a62112b | ||
|
|
eb03be0dbf | ||
|
|
6d956bdc20 | ||
|
|
333acde7d1 | ||
|
|
9994eaa70f | ||
|
|
a9866f03db | ||
|
|
0e39e21fa2 | ||
|
|
95cc6ee2c0 | ||
|
|
db6fbdbf29 | ||
|
|
988d4b577e | ||
|
|
ec23716ee4 | ||
|
|
e6eb23eaac | ||
|
|
37ed639d1e | ||
|
|
a5aae5db9a | ||
|
|
6c9b84cc7a | ||
|
|
386b57bb77 | ||
|
|
521db6869e | ||
|
|
30d146e8b5 | ||
|
|
1891faa68d | ||
|
|
41456b69dd | ||
|
|
c5d9d57940 | ||
|
|
41f12e874d | ||
|
|
67bb446b5c | ||
|
|
972a4a471e | ||
|
|
7b387e479d | ||
|
|
b5979050f3 | ||
|
|
8276fdf605 | ||
|
|
d2d4fb9221 | ||
|
|
1bdf99109b | ||
|
|
e70fbf3204 | ||
|
|
4dc0aedbfe | ||
|
|
1c0cb24c7e | ||
|
|
4fd9182aff | ||
|
|
38cc916f7d | ||
|
|
51f354f5c2 | ||
|
|
cb24522dd5 | ||
|
|
e46d55bc57 | ||
|
|
c81191b46a |
@@ -16,6 +16,7 @@ jobs:
|
||||
run: |
|
||||
/bin/bash -n deploy/apple-container.sh
|
||||
/bin/bash deploy/tests/apple-container-test.sh
|
||||
/bin/sh deploy/test-caddyfile-cache.sh
|
||||
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -102,18 +102,6 @@ Please read the following carefully before using this project:
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://apikl.ai"><img src="assets/partners/logos/apikl.png" alt="apikl" width="150"></a></td>
|
||||
<td>Thanks to Apikl for sponsoring this project! Built on Sub2API, the platform provides developers with relay services for Codex / Claude series models, focusing on long-term stability, high-speed direct connections, and excellent cost-effectiveness. It offers pay-as-you-go balance billing, enterprise-grade official invoices, and one-on-one dedicated support. <a href="https://apikl.ai">Register now</a> for a 1:1 top-up bonus — double your balance!
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://tokeneum.ai"><img src="assets/partners/logos/tokeneum.png" alt="tokeneum" width="150"></a></td>
|
||||
<td>Thanks to TokenEum for sponsoring this project! <a href="https://tokeneum.ai">TokenEum</a> is a comprehensive AI model aggregation platform and intelligent agent development company. It brings together top-tier international models — including Claude, Gemini, and OpenAI — alongside leading open-source models such as GLM, Qwen, and Kimi, offering a wide range of options across different quality and price tiers to suit every need. TokenEum also provides access to cutting-edge video generation models like Seedance2.0 and Happy Horse. Committed to transparency and honest business practices, TokenEum ensures all model information is accurate and reliable. Visit <a href="https://tokeneum.ai">tokeneum.ai</a> to get started.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://666api.work/sub2api"><img src="assets/partners/logos/666api.jpg" alt="666api" width="150"></a></td>
|
||||
<td>Thanks to 666api for sponsoring this project! <a href="https://666api.work/sub2api">666api</a> is an all-in-one platform offering:<br>
|
||||
@@ -124,23 +112,12 @@ Please read the following carefully before using this project:
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html"><img src="assets/partners/logos/byteplus.png" alt="BytePlus" width="150"></a></td>
|
||||
<td>Thanks to Dola seed for sponsoring this project! Dola Seed 2.0 is a full‑modal general large model independently developed by ByteDance for the global market. Built on a unified multimodal architecture, it supports joint understanding and generation of text, images, audio, and video. It natively enables agent collaboration, with strong reasoning, long‑task execution, tool integration, and coding capabilities. It is widely applicable to smart cockpits, personal assistants, education, customer support, marketing, retail, and other scenarios. It excels in multimodal perception, end‑to‑end complex task delivery, stable interaction, and data security, and is readily accessible and deployable via the ModelArk platform.Register via <a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html">this link</a> to get 500,000 tokens of free inference quota per model.<a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html"> >>中国大陆地区的开发者请点击这里</a></td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://sui-xiang.com/"><img src="assets/partners/logos/sui-xiang.jpg" alt="sui-xiang" width="150"></a></td>
|
||||
<td>Thanks to Suixiang AI Gateway for sponsoring this project! <a href="https://sui-xiang.com/">Suixiang AI Gateway</a> is a reliable and efficient API relay service provider offering relay services for Claude, Codex, Gemini, and more. A privacy-focused relay — no data reselling, no model dilution; privacy, transparency, and lightning-fast after-sales support. New accounts get ¥0.5 in trial credit daily by signing in; top-ups are 1:1, no subscription required, pay-as-you-go. Multi-line redundancy, cross-region disaster recovery, automatic failover, and uninterrupted long-link SSE. 99.9% availability — critical calls never fall behind.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.miyaip.com/?invitecode=sub2api"><img src="assets/partners/logos/miyaip.png" alt="miyaip" width="150"></a></td>
|
||||
<td>Thanks to MiyaIP for sponsoring this project! <a href="https://www.miyaip.com/?invitecode=sub2api">MiyaIP</a> is a platform dedicated to global residential proxy network services, committed to providing high-quality, pure overseas residential IP resources for enterprise developers, cross-border business teams, and AI application users. It delivers stable, independent overseas network environments for AI platforms, overseas SaaS, and other online services, supporting multi-region access testing and project environment isolation. Ideal for development and testing scenarios that require access to overseas AI services, such as: AI model platform access, AI development testing, AI SaaS service usage, AI API debugging, and multi-region network environment validation.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://anpin.ai"><img src="assets/partners/logos/anpin.jpg" alt="anpin" width="150"></a></td>
|
||||
<td>Thanks to <a href="https://anpin.ai">anpin.ai</a> for sponsoring this project! anpin.ai is a premium AI relay service platform dedicated to advancing AI accessibility. With an advanced technical architecture and globally distributed deployment, it provides users with a direct high-speed channel to the world's top-tier large language models.<br>
|
||||
|
||||
@@ -105,18 +105,6 @@
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://apikl.ai"><img src="assets/partners/logos/apikl.png" alt="apikl" width="150"></a></td>
|
||||
<td>感谢 Apikl 赞助本项目!平台基于 Sub2API 搭建,为开发者提供 Codex / Claude 系列模型的中转服务,专注于长期稳定、高速直连与高性价比。支持按量计费的余额结算、企业级正规发票及一对一专属对接。<a href="https://apikl.ai">立即注册</a>即享充值 1:1 赠送 — 余额翻倍!
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://tokeneum.ai"><img src="assets/partners/logos/tokeneum.png" alt="tokeneum" width="150"></a></td>
|
||||
<td>感谢 TokenEum 赞助本项目!<a href="https://tokeneum.ai">TokenEum</a> 是一家综合性 AI 模型聚合平台与智能体开发公司,汇聚 Claude、Gemini、OpenAI 等国际顶级模型,以及 GLM、Qwen、Kimi 等主流开源模型,提供不同质量与价格梯度的丰富选择,满足多样化需求。平台还接入了 Seedance2.0、Happy Horse 等前沿视频生成模型。秉持透明诚信的经营理念,TokenEum 确保所有模型信息真实可靠。访问 <a href="https://tokeneum.ai">tokeneum.ai</a> 开始使用。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://666api.work/sub2api"><img src="assets/partners/logos/666api.jpg" alt="666api" width="150"></a></td>
|
||||
<td>感谢 666api 赞助本项目!<a href="https://666api.work/sub2api">666api</a> 是一站式综合服务平台,提供:<br>
|
||||
@@ -127,23 +115,12 @@
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html"><img src="assets/partners/logos/huoshan.png" alt="HuoShan" width="150"></a></td>
|
||||
<td>感谢火山方舟 Agent Plan 模型赞助了本项目!方舟 Agent Plan 模型订阅套餐集成了包含 Doubao-Seed、Doubao-Seedance、Doubao-Seedream 等在内的字节跳动自研 SOTA 级模型,覆盖文本、代码、图像、视频等多模态任务。最新支持 MiniMax-M3、DeepSeek-V4 系列、GLM-5.1、Doubao-Seed-2.0 系列、Kimi-K2.6 等模型,工具不限。超全模态模型与 Harness 升级一步到位,深度支持 Agent 框架与 AI 编程工具。一次订阅,可以为不同任务切换合适的 AI 引擎。方舟 Coding Plan 为 Sub2Api 的用户提供了专属福利:通过<a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html">此链接</a>订阅方舟 Coding Plan,新客户首两个月享 2.5 折优惠 <a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html">>>For developers outside Mainland China, please click here</a></td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://sui-xiang.com/"><img src="assets/partners/logos/sui-xiang.jpg" alt="sui-xiang" width="150"></a></td>
|
||||
<td>感谢 随想AI网关 赞助本项目!<a href="https://sui-xiang.com/">随想AI网关</a> 是一家可靠高效的 API 中继服务提供商,提供 Claude、Codex、Gemini 等的中继服务。注重隐私的中转站·无数据倒卖·无模型掺水,隐私,透明,极速售后。新账户注册每日签到就送 0.5 元测试额度,充值额度 1:1,无需订阅,按量付费。多线路冗余、跨区域容灾、自动故障切换,长链路 SSE 不中断。99.9% 可用性,关键调用从不掉队。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.miyaip.com/?invitecode=sub2api"><img src="assets/partners/logos/miyaip.png" alt="miyaip" width="150"></a></td>
|
||||
<td>感谢 MiyaIP 赞助本项目!<a href="https://www.miyaip.com/?invitecode=sub2api">MiyaIP</a> 是一家专注于全球住宅代理网络服务的平台,致力于为企业开发者、跨境业务团队及AI 应用用户提供高质量、纯净的海外住宅 IP 资源。为 AI 平台、海外 SaaS 及其他在线服务提供稳定、独立的海外网络环境,支持多地区访问测试和项目环境隔离。适用于需要访问海外 AI 服务的开发和测试场景,例如:AI 模型平台访问、AI 开发测试、AI SaaS 服务使用、AI API 调试、多地区网络环境验证
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://anpin.ai"><img src="assets/partners/logos/anpin.jpg" alt="anpin" width="150"></a></td>
|
||||
<td>感谢 <a href="https://anpin.ai">anpin.ai</a> 赞助本项目!anpin.ai 是一家致力于推动 AI 普惠的高端 AI 中转服务平台。我们以先进的技术架构和全球分布式部署,为用户提供直达国际顶尖大模型的高速通道。<br>
|
||||
|
||||
@@ -104,41 +104,18 @@
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://apikl.ai"><img src="assets/partners/logos/apikl.png" alt="apikl" width="150"></a></td>
|
||||
<td>Apikl のご支援に感謝します!Sub2API をベースに構築された本プラットフォームは、開発者向けに Codex / Claude シリーズモデルの中継サービスを提供しています。長期安定性、高速直結、高いコストパフォーマンスを重視し、従量課金の残高ベース課金、エンタープライズグレードの正規請求書、1対1の専属サポートを提供します。<a href="https://apikl.ai">今すぐ登録</a>でチャージ 1:1 ボーナス — 残高が倍に!
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://tokeneum.ai"><img src="assets/partners/logos/tokeneum.png" alt="tokeneum" width="150"></a></td>
|
||||
<td>TokenEum のご支援に感謝します!<a href="https://tokeneum.ai">TokenEum</a> は総合的な AI モデル集約プラットフォームおよびインテリジェントエージェント開発企業です。Claude、Gemini、OpenAI などの世界トップクラスのモデルに加え、GLM、Qwen、Kimi などの主要なオープンソースモデルも集約しており、品質と価格の異なる豊富な選択肢を提供してあらゆるニーズに対応します。また、Seedance2.0 や Happy Horse などの最先端の動画生成モデルも利用可能です。TokenEum は透明性と誠実なビジネスを重視し、すべてのモデル情報の正確性と信頼性を保証します。<a href="https://tokeneum.ai">tokeneum.ai</a> でぜひお試しください。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://666api.work/sub2api"><img src="assets/partners/logos/666api.jpg" alt="666api" width="150"></a></td>
|
||||
<td>666api のご支援に感謝します!<a href="https://666api.work/sub2api">666api</a> はオールインワンプラットフォームで、以下を提供しています:⚡ API 中継 — グローバルモデルへの従量課金アクセス、100% 公式ソースから直接供給、公式価格の最大 75% オフ。独占特典:Zhipu GLM 50% オフ・DeepSeek V4-pro 50% オフ・Seedance2.0 8% オフ(ホワイトリスト)・HappyHorse Overseas 30% オフ(ホワイトリスト)🔑 GPT サブスクリプションアカウント — 同源 IP 込み・グローバル住宅 IP 💰 請求書発行対応
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html"><img src="assets/partners/logos/byteplus.png" alt="BytePlus" width="150"></a></td>
|
||||
<td>Dola seed のご支援に感謝します!Dola Seed 2.0 は、ByteDance がグローバル市場向けに独自開発した全モーダル汎用大規模モデルです。統一されたマルチモーダルアーキテクチャに基づき、テキスト・画像・音声・動画の共同理解と生成をサポートします。エージェント協調をネイティブに有効化し、強力な推論・長期タスク実行・ツール統合・コーディング能力を備えています。スマートコックピット、パーソナルアシスタント、教育、カスタマーサポート、マーケティング、リテールなど多様なシナリオに広く適用可能です。マルチモーダル知覚、エンドツーエンドの複雑タスク配信、安定したインタラクション、データセキュリティに優れ、ModelArk プラットフォームを通じて容易にアクセス・デプロイできます。<a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html">こちらのリンク</a>から登録すると、各モデルにつき 50 万トークンの無料推論クォータを獲得できます。<a href="https://dis.chatdesks.cn/chatdesk/hsyqsub2api.html"> >>中国大陸地域の開発者はこちらをクリック</a></td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://sui-xiang.com/"><img src="assets/partners/logos/sui-xiang.jpg" alt="sui-xiang" width="150"></a></td>
|
||||
<td>随想AI ゲートウェイのご支援に感謝します!<a href="https://sui-xiang.com/">随想AI ゲートウェイ</a> は信頼性と効率に優れた API 中継サービスプロバイダーで、Claude、Codex、Gemini などの中継サービスを提供しています。プライバシー重視の中継ステーション・データ転売なし・モデル水増しなし、プライバシー・透明性・超高速アフターサービス。新規アカウント登録後、毎日サインインで 0.5 元のテストクレジットがもらえ、チャージは 1:1、サブスクリプション不要、従量課金。マルチライン冗長、クロスリージョン災害復旧、自動フェイルオーバー、長時間 SSE 接続が途切れません。99.9% の可用性、重要な呼び出しは決して遅れません。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.miyaip.com/?invitecode=sub2api"><img src="assets/partners/logos/miyaip.png" alt="miyaip" width="150"></a></td>
|
||||
<td>MiyaIP のご支援に感謝します!<a href="https://www.miyaip.com/?invitecode=sub2api">MiyaIP</a> はグローバル住宅プロキシネットワークサービスに特化したプラットフォームで、企業開発者、越境ビジネスチーム、AI アプリケーションユーザーに高品質で純粋な海外住宅 IP リソースを提供することに注力しています。AI プラットフォーム、海外 SaaS、その他のオンラインサービスに安定した独立した海外ネットワーク環境を提供し、複数地域でのアクセステストやプロジェクト環境の分離をサポートします。海外 AI サービスへのアクセスが必要な開発・テストシナリオに最適です。例:AI モデルプラットフォームアクセス、AI 開発テスト、AI SaaS サービス利用、AI API デバッグ、複数地域のネットワーク環境検証など。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://anpin.ai"><img src="assets/partners/logos/anpin.jpg" alt="anpin" width="150"></a></td>
|
||||
<td><a href="https://anpin.ai">anpin.ai</a> のご支援に感謝します!anpin.ai は AI の普及推進に取り組むハイエンド AI 中継サービスプラットフォームです。先進的な技術アーキテクチャとグローバル分散デプロイメントにより、国際トップクラスの大規模モデルへの高速直通ルートを提供します。<br>
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 9.5 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 84 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 105 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 10 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 23 KiB |
@@ -1 +1 @@
|
||||
0.1.163
|
||||
0.1.165
|
||||
|
||||
+12
-1
@@ -105,6 +105,8 @@ type Group struct {
|
||||
SortOrder int `json:"sort_order,omitempty"`
|
||||
// 是否允许 /v1/messages 调度到此 OpenAI 分组
|
||||
AllowMessagesDispatch bool `json:"allow_messages_dispatch,omitempty"`
|
||||
// 是否允许此 OpenAI 分组访问 Live 接口
|
||||
AllowLive bool `json:"allow_live,omitempty"`
|
||||
// 仅允许非 apikey 类型账号关联到此分组
|
||||
RequireOauthOnly bool `json:"require_oauth_only,omitempty"`
|
||||
// 调度时仅允许 privacy 已成功设置的账号
|
||||
@@ -229,7 +231,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
|
||||
switch columns[i] {
|
||||
case 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.FieldRequireOauthOnly, group.FieldRequirePrivacySet:
|
||||
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:
|
||||
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:
|
||||
values[i] = new(sql.NullFloat64)
|
||||
@@ -537,6 +539,12 @@ func (_m *Group) assignValues(columns []string, values []any) error {
|
||||
} else if value.Valid {
|
||||
_m.AllowMessagesDispatch = value.Bool
|
||||
}
|
||||
case group.FieldAllowLive:
|
||||
if value, ok := values[i].(*sql.NullBool); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field allow_live", values[i])
|
||||
} else if value.Valid {
|
||||
_m.AllowLive = value.Bool
|
||||
}
|
||||
case group.FieldRequireOauthOnly:
|
||||
if value, ok := values[i].(*sql.NullBool); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field require_oauth_only", values[i])
|
||||
@@ -826,6 +834,9 @@ func (_m *Group) String() string {
|
||||
builder.WriteString("allow_messages_dispatch=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.AllowMessagesDispatch))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("allow_live=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.AllowLive))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("require_oauth_only=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.RequireOauthOnly))
|
||||
builder.WriteString(", ")
|
||||
|
||||
@@ -102,6 +102,8 @@ const (
|
||||
FieldSortOrder = "sort_order"
|
||||
// FieldAllowMessagesDispatch holds the string denoting the allow_messages_dispatch field in the database.
|
||||
FieldAllowMessagesDispatch = "allow_messages_dispatch"
|
||||
// FieldAllowLive holds the string denoting the allow_live field in the database.
|
||||
FieldAllowLive = "allow_live"
|
||||
// FieldRequireOauthOnly holds the string denoting the require_oauth_only field in the database.
|
||||
FieldRequireOauthOnly = "require_oauth_only"
|
||||
// FieldRequirePrivacySet holds the string denoting the require_privacy_set field in the database.
|
||||
@@ -236,6 +238,7 @@ var Columns = []string{
|
||||
FieldSupportedModelScopes,
|
||||
FieldSortOrder,
|
||||
FieldAllowMessagesDispatch,
|
||||
FieldAllowLive,
|
||||
FieldRequireOauthOnly,
|
||||
FieldRequirePrivacySet,
|
||||
FieldDefaultMappedModel,
|
||||
@@ -341,6 +344,8 @@ var (
|
||||
DefaultSortOrder int
|
||||
// DefaultAllowMessagesDispatch holds the default value on creation for the "allow_messages_dispatch" field.
|
||||
DefaultAllowMessagesDispatch bool
|
||||
// DefaultAllowLive holds the default value on creation for the "allow_live" field.
|
||||
DefaultAllowLive bool
|
||||
// DefaultRequireOauthOnly holds the default value on creation for the "require_oauth_only" field.
|
||||
DefaultRequireOauthOnly bool
|
||||
// DefaultRequirePrivacySet holds the default value on creation for the "require_privacy_set" field.
|
||||
@@ -576,6 +581,11 @@ func ByAllowMessagesDispatch(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldAllowMessagesDispatch, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByAllowLive orders the results by the allow_live field.
|
||||
func ByAllowLive(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldAllowLive, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByRequireOauthOnly orders the results by the require_oauth_only field.
|
||||
func ByRequireOauthOnly(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldRequireOauthOnly, opts...).ToFunc()
|
||||
|
||||
@@ -260,6 +260,11 @@ func AllowMessagesDispatch(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldAllowMessagesDispatch, v))
|
||||
}
|
||||
|
||||
// AllowLive applies equality check predicate on the "allow_live" field. It's identical to AllowLiveEQ.
|
||||
func AllowLive(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldAllowLive, v))
|
||||
}
|
||||
|
||||
// RequireOauthOnly applies equality check predicate on the "require_oauth_only" field. It's identical to RequireOauthOnlyEQ.
|
||||
func RequireOauthOnly(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldRequireOauthOnly, v))
|
||||
@@ -1985,6 +1990,16 @@ func AllowMessagesDispatchNEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldNEQ(FieldAllowMessagesDispatch, v))
|
||||
}
|
||||
|
||||
// AllowLiveEQ applies the EQ predicate on the "allow_live" field.
|
||||
func AllowLiveEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldAllowLive, v))
|
||||
}
|
||||
|
||||
// AllowLiveNEQ applies the NEQ predicate on the "allow_live" field.
|
||||
func AllowLiveNEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldNEQ(FieldAllowLive, v))
|
||||
}
|
||||
|
||||
// RequireOauthOnlyEQ applies the EQ predicate on the "require_oauth_only" field.
|
||||
func RequireOauthOnlyEQ(v bool) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldRequireOauthOnly, v))
|
||||
|
||||
@@ -607,6 +607,20 @@ func (_c *GroupCreate) SetNillableAllowMessagesDispatch(v *bool) *GroupCreate {
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (_c *GroupCreate) SetAllowLive(v bool) *GroupCreate {
|
||||
_c.mutation.SetAllowLive(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableAllowLive sets the "allow_live" field if the given value is not nil.
|
||||
func (_c *GroupCreate) SetNillableAllowLive(v *bool) *GroupCreate {
|
||||
if v != nil {
|
||||
_c.SetAllowLive(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (_c *GroupCreate) SetRequireOauthOnly(v bool) *GroupCreate {
|
||||
_c.mutation.SetRequireOauthOnly(v)
|
||||
@@ -948,6 +962,10 @@ func (_c *GroupCreate) defaults() error {
|
||||
v := group.DefaultAllowMessagesDispatch
|
||||
_c.mutation.SetAllowMessagesDispatch(v)
|
||||
}
|
||||
if _, ok := _c.mutation.AllowLive(); !ok {
|
||||
v := group.DefaultAllowLive
|
||||
_c.mutation.SetAllowLive(v)
|
||||
}
|
||||
if _, ok := _c.mutation.RequireOauthOnly(); !ok {
|
||||
v := group.DefaultRequireOauthOnly
|
||||
_c.mutation.SetRequireOauthOnly(v)
|
||||
@@ -1101,6 +1119,9 @@ func (_c *GroupCreate) check() error {
|
||||
if _, ok := _c.mutation.AllowMessagesDispatch(); !ok {
|
||||
return &ValidationError{Name: "allow_messages_dispatch", err: errors.New(`ent: missing required field "Group.allow_messages_dispatch"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.AllowLive(); !ok {
|
||||
return &ValidationError{Name: "allow_live", err: errors.New(`ent: missing required field "Group.allow_live"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.RequireOauthOnly(); !ok {
|
||||
return &ValidationError{Name: "require_oauth_only", err: errors.New(`ent: missing required field "Group.require_oauth_only"`)}
|
||||
}
|
||||
@@ -1334,6 +1355,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(group.FieldAllowMessagesDispatch, field.TypeBool, value)
|
||||
_node.AllowMessagesDispatch = value
|
||||
}
|
||||
if value, ok := _c.mutation.AllowLive(); ok {
|
||||
_spec.SetField(group.FieldAllowLive, field.TypeBool, value)
|
||||
_node.AllowLive = value
|
||||
}
|
||||
if value, ok := _c.mutation.RequireOauthOnly(); ok {
|
||||
_spec.SetField(group.FieldRequireOauthOnly, field.TypeBool, value)
|
||||
_node.RequireOauthOnly = value
|
||||
@@ -2224,6 +2249,18 @@ func (u *GroupUpsert) UpdateAllowMessagesDispatch() *GroupUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (u *GroupUpsert) SetAllowLive(v bool) *GroupUpsert {
|
||||
u.Set(group.FieldAllowLive, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateAllowLive sets the "allow_live" field to the value that was provided on create.
|
||||
func (u *GroupUpsert) UpdateAllowLive() *GroupUpsert {
|
||||
u.SetExcluded(group.FieldAllowLive)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (u *GroupUpsert) SetRequireOauthOnly(v bool) *GroupUpsert {
|
||||
u.Set(group.FieldRequireOauthOnly, v)
|
||||
@@ -3193,6 +3230,20 @@ func (u *GroupUpsertOne) UpdateAllowMessagesDispatch() *GroupUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (u *GroupUpsertOne) SetAllowLive(v bool) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetAllowLive(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateAllowLive sets the "allow_live" field to the value that was provided on create.
|
||||
func (u *GroupUpsertOne) UpdateAllowLive() *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateAllowLive()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (u *GroupUpsertOne) SetRequireOauthOnly(v bool) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
@@ -4345,6 +4396,20 @@ func (u *GroupUpsertBulk) UpdateAllowMessagesDispatch() *GroupUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (u *GroupUpsertBulk) SetAllowLive(v bool) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetAllowLive(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateAllowLive sets the "allow_live" field to the value that was provided on create.
|
||||
func (u *GroupUpsertBulk) UpdateAllowLive() *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateAllowLive()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (u *GroupUpsertBulk) SetRequireOauthOnly(v bool) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
|
||||
@@ -822,6 +822,20 @@ func (_u *GroupUpdate) SetNillableAllowMessagesDispatch(v *bool) *GroupUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (_u *GroupUpdate) SetAllowLive(v bool) *GroupUpdate {
|
||||
_u.mutation.SetAllowLive(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableAllowLive sets the "allow_live" field if the given value is not nil.
|
||||
func (_u *GroupUpdate) SetNillableAllowLive(v *bool) *GroupUpdate {
|
||||
if v != nil {
|
||||
_u.SetAllowLive(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (_u *GroupUpdate) SetRequireOauthOnly(v bool) *GroupUpdate {
|
||||
_u.mutation.SetRequireOauthOnly(v)
|
||||
@@ -1495,6 +1509,9 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if value, ok := _u.mutation.AllowMessagesDispatch(); ok {
|
||||
_spec.SetField(group.FieldAllowMessagesDispatch, field.TypeBool, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AllowLive(); ok {
|
||||
_spec.SetField(group.FieldAllowLive, field.TypeBool, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RequireOauthOnly(); ok {
|
||||
_spec.SetField(group.FieldRequireOauthOnly, field.TypeBool, value)
|
||||
}
|
||||
@@ -2627,6 +2644,20 @@ func (_u *GroupUpdateOne) SetNillableAllowMessagesDispatch(v *bool) *GroupUpdate
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (_u *GroupUpdateOne) SetAllowLive(v bool) *GroupUpdateOne {
|
||||
_u.mutation.SetAllowLive(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableAllowLive sets the "allow_live" field if the given value is not nil.
|
||||
func (_u *GroupUpdateOne) SetNillableAllowLive(v *bool) *GroupUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetAllowLive(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (_u *GroupUpdateOne) SetRequireOauthOnly(v bool) *GroupUpdateOne {
|
||||
_u.mutation.SetRequireOauthOnly(v)
|
||||
@@ -3330,6 +3361,9 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
|
||||
if value, ok := _u.mutation.AllowMessagesDispatch(); ok {
|
||||
_spec.SetField(group.FieldAllowMessagesDispatch, field.TypeBool, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AllowLive(); ok {
|
||||
_spec.SetField(group.FieldAllowLive, field.TypeBool, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RequireOauthOnly(); ok {
|
||||
_spec.SetField(group.FieldRequireOauthOnly, field.TypeBool, value)
|
||||
}
|
||||
|
||||
@@ -938,6 +938,7 @@ var (
|
||||
{Name: "supported_model_scopes", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "sort_order", Type: field.TypeInt, Default: 0},
|
||||
{Name: "allow_messages_dispatch", Type: field.TypeBool, Default: false},
|
||||
{Name: "allow_live", Type: field.TypeBool, Default: false},
|
||||
{Name: "require_oauth_only", Type: field.TypeBool, Default: false},
|
||||
{Name: "require_privacy_set", Type: field.TypeBool, Default: false},
|
||||
{Name: "default_mapped_model", Type: field.TypeString, Size: 100, Default: ""},
|
||||
|
||||
+55
-1
@@ -21911,6 +21911,7 @@ type GroupMutation struct {
|
||||
sort_order *int
|
||||
addsort_order *int
|
||||
allow_messages_dispatch *bool
|
||||
allow_live *bool
|
||||
require_oauth_only *bool
|
||||
require_privacy_set *bool
|
||||
default_mapped_model *string
|
||||
@@ -24226,6 +24227,42 @@ func (m *GroupMutation) ResetAllowMessagesDispatch() {
|
||||
m.allow_messages_dispatch = nil
|
||||
}
|
||||
|
||||
// SetAllowLive sets the "allow_live" field.
|
||||
func (m *GroupMutation) SetAllowLive(b bool) {
|
||||
m.allow_live = &b
|
||||
}
|
||||
|
||||
// AllowLive returns the value of the "allow_live" field in the mutation.
|
||||
func (m *GroupMutation) AllowLive() (r bool, exists bool) {
|
||||
v := m.allow_live
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldAllowLive returns the old "allow_live" 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) OldAllowLive(ctx context.Context) (v bool, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldAllowLive is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldAllowLive requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldAllowLive: %w", err)
|
||||
}
|
||||
return oldValue.AllowLive, nil
|
||||
}
|
||||
|
||||
// ResetAllowLive resets all changes to the "allow_live" field.
|
||||
func (m *GroupMutation) ResetAllowLive() {
|
||||
m.allow_live = nil
|
||||
}
|
||||
|
||||
// SetRequireOauthOnly sets the "require_oauth_only" field.
|
||||
func (m *GroupMutation) SetRequireOauthOnly(b bool) {
|
||||
m.require_oauth_only = &b
|
||||
@@ -24907,7 +24944,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, 51)
|
||||
fields := make([]string, 0, 52)
|
||||
if m.created_at != nil {
|
||||
fields = append(fields, group.FieldCreatedAt)
|
||||
}
|
||||
@@ -25037,6 +25074,9 @@ func (m *GroupMutation) Fields() []string {
|
||||
if m.allow_messages_dispatch != nil {
|
||||
fields = append(fields, group.FieldAllowMessagesDispatch)
|
||||
}
|
||||
if m.allow_live != nil {
|
||||
fields = append(fields, group.FieldAllowLive)
|
||||
}
|
||||
if m.require_oauth_only != nil {
|
||||
fields = append(fields, group.FieldRequireOauthOnly)
|
||||
}
|
||||
@@ -25155,6 +25195,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.SortOrder()
|
||||
case group.FieldAllowMessagesDispatch:
|
||||
return m.AllowMessagesDispatch()
|
||||
case group.FieldAllowLive:
|
||||
return m.AllowLive()
|
||||
case group.FieldRequireOauthOnly:
|
||||
return m.RequireOauthOnly()
|
||||
case group.FieldRequirePrivacySet:
|
||||
@@ -25266,6 +25308,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
|
||||
return m.OldSortOrder(ctx)
|
||||
case group.FieldAllowMessagesDispatch:
|
||||
return m.OldAllowMessagesDispatch(ctx)
|
||||
case group.FieldAllowLive:
|
||||
return m.OldAllowLive(ctx)
|
||||
case group.FieldRequireOauthOnly:
|
||||
return m.OldRequireOauthOnly(ctx)
|
||||
case group.FieldRequirePrivacySet:
|
||||
@@ -25592,6 +25636,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetAllowMessagesDispatch(v)
|
||||
return nil
|
||||
case group.FieldAllowLive:
|
||||
v, ok := value.(bool)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetAllowLive(v)
|
||||
return nil
|
||||
case group.FieldRequireOauthOnly:
|
||||
v, ok := value.(bool)
|
||||
if !ok {
|
||||
@@ -26180,6 +26231,9 @@ func (m *GroupMutation) ResetField(name string) error {
|
||||
case group.FieldAllowMessagesDispatch:
|
||||
m.ResetAllowMessagesDispatch()
|
||||
return nil
|
||||
case group.FieldAllowLive:
|
||||
m.ResetAllowLive()
|
||||
return nil
|
||||
case group.FieldRequireOauthOnly:
|
||||
m.ResetRequireOauthOnly()
|
||||
return nil
|
||||
|
||||
@@ -1141,40 +1141,44 @@ func init() {
|
||||
groupDescAllowMessagesDispatch := groupFields[39].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[40].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[40].Descriptor()
|
||||
groupDescRequireOauthOnly := groupFields[41].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[41].Descriptor()
|
||||
groupDescRequirePrivacySet := groupFields[42].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[42].Descriptor()
|
||||
groupDescDefaultMappedModel := groupFields[43].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[43].Descriptor()
|
||||
groupDescMessagesDispatchModelConfig := groupFields[44].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[44].Descriptor()
|
||||
groupDescModelsListConfig := groupFields[45].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[45].Descriptor()
|
||||
groupDescRpmLimit := groupFields[46].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[46].Descriptor()
|
||||
groupDescMaxReasoningEffort := groupFields[47].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[47].Descriptor()
|
||||
groupDescReasoningEffortMappings := groupFields[48].Descriptor()
|
||||
// group.DefaultReasoningEffortMappings holds the default value on creation for the reasoning_effort_mappings field.
|
||||
group.DefaultReasoningEffortMappings = groupDescReasoningEffortMappings.Default.([]domain.ReasoningEffortMapping)
|
||||
idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin()
|
||||
|
||||
@@ -198,6 +198,9 @@ func (Group) Fields() []ent.Field {
|
||||
field.Bool("allow_messages_dispatch").
|
||||
Default(false).
|
||||
Comment("是否允许 /v1/messages 调度到此 OpenAI 分组"),
|
||||
field.Bool("allow_live").
|
||||
Default(false).
|
||||
Comment("是否允许此 OpenAI 分组访问 Live 接口"),
|
||||
field.Bool("require_oauth_only").
|
||||
Default(false).
|
||||
Comment("仅允许非 apikey 类型账号关联到此分组"),
|
||||
|
||||
+9
-8
@@ -20,7 +20,7 @@ require (
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/google/wire v0.7.0
|
||||
github.com/gorilla/websocket v1.5.3
|
||||
github.com/imroc/req/v3 v3.57.0
|
||||
github.com/imroc/req/v3 v3.59.0
|
||||
github.com/klauspost/compress v1.18.2
|
||||
github.com/lib/pq v1.10.9
|
||||
github.com/patrickmn/go-cache v2.1.0+incompatible
|
||||
@@ -43,7 +43,7 @@ require (
|
||||
github.com/zeromicro/go-zero v1.9.4
|
||||
go.uber.org/zap v1.24.0
|
||||
golang.org/x/crypto v0.53.0
|
||||
golang.org/x/image v0.39.0
|
||||
golang.org/x/image v0.41.0
|
||||
golang.org/x/mod v0.37.0
|
||||
golang.org/x/net v0.56.0
|
||||
golang.org/x/sync v0.21.0
|
||||
@@ -141,7 +141,7 @@ require (
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
|
||||
github.com/quic-go/qpack v0.6.0 // indirect
|
||||
github.com/quic-go/quic-go v0.57.1 // indirect
|
||||
github.com/quic-go/quic-go v0.60.0 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/sagikazarmark/locafero v0.4.0 // indirect
|
||||
github.com/sagikazarmark/slog-shim v0.1.0 // indirect
|
||||
@@ -166,12 +166,12 @@ require (
|
||||
github.com/yusufpapurcu/wmi v1.2.4 // indirect
|
||||
github.com/zclconf/go-cty v1.14.4 // indirect
|
||||
github.com/zclconf/go-cty-yaml v1.1.0 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
|
||||
go.opentelemetry.io/otel v1.37.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.37.0 // indirect
|
||||
go.opentelemetry.io/otel/sdk v1.37.0 // indirect
|
||||
go.opentelemetry.io/otel/trace v1.37.0 // indirect
|
||||
go.opentelemetry.io/otel v1.41.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.41.0 // indirect
|
||||
go.opentelemetry.io/otel/sdk v1.41.0 // indirect
|
||||
go.opentelemetry.io/otel/trace v1.41.0 // indirect
|
||||
go.uber.org/atomic v1.10.0 // indirect
|
||||
go.uber.org/automaxprocs v1.6.0 // indirect
|
||||
go.uber.org/multierr v1.9.0 // indirect
|
||||
@@ -179,6 +179,7 @@ require (
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect
|
||||
golang.org/x/sys v0.46.0 // indirect
|
||||
golang.org/x/text v0.39.0 // indirect
|
||||
golang.org/x/time v0.12.0 // indirect
|
||||
golang.org/x/tools v0.47.0 // indirect
|
||||
google.golang.org/grpc v1.75.1 // indirect
|
||||
google.golang.org/protobuf v1.36.10 // indirect
|
||||
|
||||
+20
-44
@@ -166,8 +166,6 @@ github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
|
||||
github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE=
|
||||
github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4=
|
||||
@@ -185,8 +183,8 @@ github.com/hashicorp/hcl/v2 v2.18.1 h1:6nxnOJFku1EuSawSD81fuviYUV8DxFr3fp2dUi3ZY
|
||||
github.com/hashicorp/hcl/v2 v2.18.1/go.mod h1:ThLC89FV4p9MPW804KVbe/cEXoQ8NZEh+JtMeeGErHE=
|
||||
github.com/icholy/digest v1.1.0 h1:HfGg9Irj7i+IX1o1QAmPfIBNu/Q5A5Tu3n/MED9k9H4=
|
||||
github.com/icholy/digest v1.1.0/go.mod h1:QNrsSGQ5v7v9cReDI0+eyjsXGUoRSUZQHeQ5C4XLa0Y=
|
||||
github.com/imroc/req/v3 v3.57.0 h1:LMTUjNRUybUkTPn8oJDq8Kg3JRBOBTcnDhKu7mzupKI=
|
||||
github.com/imroc/req/v3 v3.57.0/go.mod h1:JL62ey1nvSLq81HORNcosvlf7SxZStONNqOprg0Pz00=
|
||||
github.com/imroc/req/v3 v3.59.0 h1:PqKhJHyBmJYob47LVuTHwRZE00ZO6icbLHe5Zra13jo=
|
||||
github.com/imroc/req/v3 v3.59.0/go.mod h1:cwLwwaE90Pq6ZMmKuHA9FivJbg3avnCH8/RA5lUUs74=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
@@ -222,8 +220,6 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk
|
||||
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
|
||||
github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
|
||||
github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM=
|
||||
github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
|
||||
github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI=
|
||||
@@ -257,8 +253,6 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
|
||||
github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
|
||||
github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
|
||||
github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
|
||||
github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
|
||||
github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=
|
||||
@@ -278,22 +272,22 @@ github.com/pquerna/otp v1.5.0 h1:NMMR+WrmaqXU4EzdGJEE1aUUI0AMRzsp96fFFWNPwxs=
|
||||
github.com/pquerna/otp v1.5.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg=
|
||||
github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
|
||||
github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
|
||||
github.com/quic-go/go-ossfuzz-seeds v0.1.0 h1:APacT+iIaNF6fd8AGEiN3bT/Jtkd2jz4v4TzM7MFjy0=
|
||||
github.com/quic-go/go-ossfuzz-seeds v0.1.0/go.mod h1:3IOHRbJIc+L6YKMwfDtJAM9Vj9k0YY4muhuyUYk5tbk=
|
||||
github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8=
|
||||
github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII=
|
||||
github.com/quic-go/quic-go v0.57.1 h1:25KAAR9QR8KZrCZRThWMKVAwGoiHIrNbT72ULHTuI10=
|
||||
github.com/quic-go/quic-go v0.57.1/go.mod h1:ly4QBAjHA2VhdnxhojRsCUOeJwKYg+taDlos92xb1+s=
|
||||
github.com/quic-go/quic-go v0.60.0 h1:xcQioE8OM66UQLeUMHltK1CCcOu3JbVB4JAQdDQSB+0=
|
||||
github.com/quic-go/quic-go v0.60.0/go.mod h1:wpKpjmPpftl30sL6pFh7REVpjbcCVy4zt2vDyK1TuJk=
|
||||
github.com/redis/go-redis/v9 v9.17.2 h1:P2EGsA4qVIM3Pp+aPocCJ7DguDHhqrXNhVcEp4ViluI=
|
||||
github.com/redis/go-redis/v9 v9.17.2/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370=
|
||||
github.com/refraction-networking/utls v1.8.2 h1:j4Q1gJj0xngdeH+Ox/qND11aEfhpgoEvV+S9iJ2IdQo=
|
||||
github.com/refraction-networking/utls v1.8.2/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
|
||||
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
|
||||
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
|
||||
github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
|
||||
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
|
||||
github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6keLGt6kNQ=
|
||||
github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4=
|
||||
github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE=
|
||||
@@ -322,8 +316,6 @@ github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8=
|
||||
github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY=
|
||||
github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0=
|
||||
github.com/spf13/cast v1.6.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
|
||||
github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
|
||||
github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
|
||||
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
|
||||
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||
github.com/spf13/viper v1.18.2 h1:LUXCnvUvSM6FXAsj6nnfc8Q2tp1dIgUfY9Kc8GsSOiQ=
|
||||
@@ -386,22 +378,22 @@ github.com/zclconf/go-cty-yaml v1.1.0 h1:nP+jp0qPHv2IhUVqmQSzjvqAWcObN0KBkUl2rWB
|
||||
github.com/zclconf/go-cty-yaml v1.1.0/go.mod h1:9YLUH4g7lOhVWqUbctnVlZ5KLpg7JAprQNgxSZ1Gyxs=
|
||||
github.com/zeromicro/go-zero v1.9.4 h1:aRLFoISqAYijABtkbliQC5SsI5TbizJpQvoHc9xup8k=
|
||||
github.com/zeromicro/go-zero v1.9.4/go.mod h1:a17JOTch25SWxBcUgJZYps60hygK3pIYdw7nGwlcS38=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 h1:jq9TW8u3so/bN+JPT166wjOI6/vQPF6Xe7nMNIltagk=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0/go.mod h1:p8pYQP+m5XfbZm9fxtSKAbM6oIllS7s2AfxrChvc7iw=
|
||||
go.opentelemetry.io/otel v1.37.0 h1:9zhNfelUvx0KBfu/gb+ZgeAfAgtWrfHJZcAqFC228wQ=
|
||||
go.opentelemetry.io/otel v1.37.0/go.mod h1:ehE/umFRLnuLa/vSccNq9oS1ErUlkkK71gMcN34UG8I=
|
||||
go.opentelemetry.io/otel v1.41.0 h1:YlEwVsGAlCvczDILpUXpIpPSL/VPugt7zHThEMLce1c=
|
||||
go.opentelemetry.io/otel v1.41.0/go.mod h1:Yt4UwgEKeT05QbLwbyHXEwhnjxNO6D8L5PQP51/46dE=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0 h1:t6wl9SPayj+c7lEIFgm4ooDBZVb01IhLB4InpomhRw8=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0/go.mod h1:iSDOcsnSA5INXzZtwaBPrKp/lWu/V14Dd+llD0oI2EA=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.24.0 h1:Xw8U6u2f8DK2XAkGRFV7BBLENgnTGX9i4rQRxJf+/vs=
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.24.0/go.mod h1:6KW1Fm6R/s6Z3PGXwSJN2K4eT6wQB3vXX6CVnYX9NmM=
|
||||
go.opentelemetry.io/otel/metric v1.37.0 h1:mvwbQS5m0tbmqML4NqK+e3aDiO02vsf/WgbsdpcPoZE=
|
||||
go.opentelemetry.io/otel/metric v1.37.0/go.mod h1:04wGrZurHYKOc+RKeye86GwKiTb9FKm1WHtO+4EVr2E=
|
||||
go.opentelemetry.io/otel/sdk v1.37.0 h1:ItB0QUqnjesGRvNcmAcU0LyvkVyGJ2xftD29bWdDvKI=
|
||||
go.opentelemetry.io/otel/sdk v1.37.0/go.mod h1:VredYzxUvuo2q3WRcDnKDjbdvmO0sCzOvVAiY+yUkAg=
|
||||
go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4=
|
||||
go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0=
|
||||
go.opentelemetry.io/otel/metric v1.41.0 h1:rFnDcs4gRzBcsO9tS8LCpgR0dxg4aaxWlJxCno7JlTQ=
|
||||
go.opentelemetry.io/otel/metric v1.41.0/go.mod h1:xPvCwd9pU0VN8tPZYzDZV/BMj9CM9vs00GuBjeKhJps=
|
||||
go.opentelemetry.io/otel/sdk v1.41.0 h1:YPIEXKmiAwkGl3Gu1huk1aYWwtpRLeskpV+wPisxBp8=
|
||||
go.opentelemetry.io/otel/sdk v1.41.0/go.mod h1:ahFdU0G5y8IxglBf0QBJXgSe7agzjE4GiTJ6HT9ud90=
|
||||
go.opentelemetry.io/otel/trace v1.41.0 h1:Vbk2co6bhj8L59ZJ6/xFTskY+tGAbOnCtQGVVa9TIN0=
|
||||
go.opentelemetry.io/otel/trace v1.41.0/go.mod h1:U1NU4ULCoxeDKc09yCWdWe+3QoyweJcISEVa1RBzOis=
|
||||
go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0=
|
||||
go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8=
|
||||
go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
|
||||
@@ -419,24 +411,16 @@ go.uber.org/zap v1.24.0/go.mod h1:2kMP+WWQ8aoFoedH3T2sq6iJ2yDWpHbP0f6MQbS9Gkg=
|
||||
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
|
||||
golang.org/x/arch v0.3.0 h1:02VY4/ZcO/gBOH6PUaoiptASxtXU10jazRCP865E97k=
|
||||
golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
|
||||
golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI=
|
||||
golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8=
|
||||
golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto=
|
||||
golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio=
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY=
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70=
|
||||
golang.org/x/image v0.39.0 h1:skVYidAEVKgn8lZ602XO75asgXBgLj9G/FE3RbuPFww=
|
||||
golang.org/x/image v0.39.0/go.mod h1:sIbmppfU+xFLPIG0FoVUTvyBMmgng1/XAMhQ2ft0hpA=
|
||||
golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM=
|
||||
golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU=
|
||||
golang.org/x/image v0.41.0 h1:8wS72eGJMJaBxK6okTzd4WaXumUlTVlb753MlsSvTCo=
|
||||
golang.org/x/image v0.41.0/go.mod h1:uIc348UZMSvS5Z65CVZ7iDPaNobNFEPeJ4kbqTOszmA=
|
||||
golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ=
|
||||
golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0=
|
||||
golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8=
|
||||
golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww=
|
||||
golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o=
|
||||
golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
|
||||
golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
@@ -448,22 +432,14 @@ golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBc
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
|
||||
golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw=
|
||||
golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4=
|
||||
golang.org/x/term v0.43.0/go.mod h1:lrhlHNdQJHO+1qVYiHfFKVuVioJIheAc3fBSMFYEIsk=
|
||||
golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc=
|
||||
golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y=
|
||||
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
|
||||
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
|
||||
golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus=
|
||||
golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM=
|
||||
golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
|
||||
golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
|
||||
golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c=
|
||||
golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI=
|
||||
golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q=
|
||||
golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
|
||||
@@ -913,6 +913,8 @@ type GatewayConfig struct {
|
||||
OpenAICompactModel string `mapstructure:"openai_compact_model"`
|
||||
// OpenAIWS: OpenAI Responses WebSocket 配置(默认开启,可按需回滚到 HTTP)
|
||||
OpenAIWS GatewayOpenAIWSConfig `mapstructure:"openai_ws"`
|
||||
// Live: ChatGPT Frameless Live 会话配置。
|
||||
Live GatewayLiveConfig `mapstructure:"live"`
|
||||
// OpenAIScheduler: OpenAI 高级调度器粘性逃逸配置
|
||||
OpenAIScheduler GatewayOpenAISchedulerConfig `mapstructure:"openai_scheduler"`
|
||||
// OpenAIHTTP2: OpenAI HTTP 上游协议策略(默认启用 HTTP/2,可按代理能力回退 HTTP/1.1)
|
||||
@@ -999,6 +1001,11 @@ type GatewayConfig struct {
|
||||
UserMessageQueue UserMessageQueueConfig `mapstructure:"user_message_queue"`
|
||||
}
|
||||
|
||||
type GatewayLiveConfig struct {
|
||||
// MaxSessionDurationSeconds 是 Live 会话的硬上限。
|
||||
MaxSessionDurationSeconds int `mapstructure:"max_session_duration_seconds"`
|
||||
}
|
||||
|
||||
// GatewayOpenAIHTTP2Config OpenAI HTTP 上游协议配置。
|
||||
// 默认启用 HTTP/2;在部分代理不兼容时按策略回退 HTTP/1.1。
|
||||
type GatewayOpenAIHTTP2Config struct {
|
||||
@@ -1633,20 +1640,7 @@ func LoadForBootstrap() (*Config, error) {
|
||||
func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
viper.SetConfigName("config")
|
||||
viper.SetConfigType("yaml")
|
||||
|
||||
// Add config paths in priority order
|
||||
// 1. DATA_DIR environment variable (highest priority)
|
||||
if dataDir := os.Getenv("DATA_DIR"); dataDir != "" {
|
||||
viper.AddConfigPath(dataDir)
|
||||
}
|
||||
// 2. Docker data directory
|
||||
viper.AddConfigPath("/app/data")
|
||||
// 3. Current directory
|
||||
viper.AddConfigPath(".")
|
||||
// 4. Config subdirectory
|
||||
viper.AddConfigPath("./config")
|
||||
// 5. System config directory
|
||||
viper.AddConfigPath("/etc/sub2api")
|
||||
configureConfigSource(viper.SetConfigFile, viper.AddConfigPath)
|
||||
|
||||
// 环境变量支持
|
||||
viper.AutomaticEnv()
|
||||
@@ -1821,6 +1815,22 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
func configureConfigSource(setConfigFile, addConfigPath func(string)) {
|
||||
if configFile := strings.TrimSpace(os.Getenv("CONFIG_FILE")); configFile != "" {
|
||||
setConfigFile(configFile)
|
||||
return
|
||||
}
|
||||
|
||||
// Add config paths in priority order.
|
||||
if dataDir := strings.TrimSpace(os.Getenv("DATA_DIR")); dataDir != "" {
|
||||
addConfigPath(dataDir)
|
||||
}
|
||||
addConfigPath("/app/data")
|
||||
addConfigPath(".")
|
||||
addConfigPath("./config")
|
||||
addConfigPath("/etc/sub2api")
|
||||
}
|
||||
|
||||
func setDefaults() {
|
||||
viper.SetDefault("run_mode", RunModeStandard)
|
||||
|
||||
@@ -2188,6 +2198,7 @@ func setDefaults() {
|
||||
viper.SetDefault("gateway.codex_image_generation_bridge_enabled", false)
|
||||
viper.SetDefault("gateway.openai_passthrough_allow_timeout_headers", false)
|
||||
viper.SetDefault("gateway.openai_compact_model", "gpt-5.4")
|
||||
viper.SetDefault("gateway.live.max_session_duration_seconds", 3600)
|
||||
// OpenAI Responses WebSocket(默认开启;可通过 force_http 紧急回滚)
|
||||
viper.SetDefault("gateway.openai_ws.enabled", true)
|
||||
viper.SetDefault("gateway.openai_ws.mode_router_v2_enabled", false)
|
||||
@@ -3039,6 +3050,9 @@ func (c *Config) Validate() error {
|
||||
(c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds > 0 && c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds < 30) {
|
||||
return fmt.Errorf("gateway.openai_high_effort_first_output_timeout_seconds must be 0 or between 30-1800 seconds")
|
||||
}
|
||||
if c.Gateway.Live.MaxSessionDurationSeconds <= 0 {
|
||||
c.Gateway.Live.MaxSessionDurationSeconds = 3600
|
||||
}
|
||||
if strings.TrimSpace(c.Gateway.ConnectionPoolIsolation) != "" {
|
||||
switch c.Gateway.ConnectionPoolIsolation {
|
||||
case ConnectionPoolIsolationProxy, ConnectionPoolIsolationAccount, ConnectionPoolIsolationAccountProxy:
|
||||
@@ -3477,14 +3491,12 @@ func generateJWTSecret(byteLength int) (string, error) {
|
||||
// GetServerAddress returns the server address (host:port) from config file or environment variable.
|
||||
// This is a lightweight function that can be used before full config validation,
|
||||
// such as during setup wizard startup.
|
||||
// Priority: config.yaml > environment variables > defaults
|
||||
// Priority: environment variables > config file > defaults.
|
||||
func GetServerAddress() string {
|
||||
v := viper.New()
|
||||
v.SetConfigName("config")
|
||||
v.SetConfigType("yaml")
|
||||
v.AddConfigPath(".")
|
||||
v.AddConfigPath("./config")
|
||||
v.AddConfigPath("/etc/sub2api")
|
||||
configureConfigSource(v.SetConfigFile, v.AddConfigPath)
|
||||
|
||||
// Support SERVER_HOST and SERVER_PORT environment variables
|
||||
v.AutomaticEnv()
|
||||
|
||||
@@ -17,6 +17,9 @@ import (
|
||||
func resetViperWithJWTSecret(t *testing.T) {
|
||||
t.Helper()
|
||||
viper.Reset()
|
||||
t.Cleanup(viper.Reset)
|
||||
t.Setenv("CONFIG_FILE", "")
|
||||
t.Setenv("DATA_DIR", "")
|
||||
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
|
||||
}
|
||||
|
||||
@@ -253,8 +256,53 @@ func TestLoadTrustedProxiesPresenceFromYAML(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadTrustedProxiesFromConfigFile(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
configFile := filepath.Join(t.TempDir(), "reporter-config.yaml")
|
||||
require.NoError(t, os.WriteFile(configFile, []byte(`server:
|
||||
host: 192.0.2.10
|
||||
trusted_proxies:
|
||||
- 127.0.0.1/32
|
||||
- 10.0.0.0/8
|
||||
- 172.16.0.0/12
|
||||
- 192.168.0.0/16
|
||||
`), 0o600))
|
||||
t.Setenv("CONFIG_FILE", configFile)
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "192.0.2.10", cfg.Server.Host)
|
||||
require.Equal(t, []string{"127.0.0.1/32", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16"}, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestConfigFileTakesPrecedenceOverDataDir(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
dataDir := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(dataDir, "config.yaml"), []byte("server:\n host: 192.0.2.20\n"), 0o600))
|
||||
configFile := filepath.Join(t.TempDir(), "explicit.yaml")
|
||||
require.NoError(t, os.WriteFile(configFile, []byte("server:\n host: 192.0.2.30\n"), 0o600))
|
||||
t.Setenv("DATA_DIR", dataDir)
|
||||
t.Setenv("CONFIG_FILE", configFile)
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "192.0.2.30", cfg.Server.Host)
|
||||
}
|
||||
|
||||
func TestLoadReturnsErrorForMissingConfigFile(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("CONFIG_FILE", filepath.Join(t.TempDir(), "missing.yaml"))
|
||||
|
||||
_, err := Load()
|
||||
require.ErrorContains(t, err, "read config error")
|
||||
}
|
||||
|
||||
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
|
||||
viper.Reset()
|
||||
t.Cleanup(viper.Reset)
|
||||
t.Setenv("CONFIG_FILE", "")
|
||||
t.Setenv("DATA_DIR", "")
|
||||
t.Setenv("JWT_SECRET", "")
|
||||
|
||||
cfg, err := LoadForBootstrap()
|
||||
@@ -1155,6 +1203,8 @@ func TestNormalizeStringSlice(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestGetServerAddressFromEnv(t *testing.T) {
|
||||
t.Setenv("CONFIG_FILE", "")
|
||||
t.Setenv("DATA_DIR", "")
|
||||
t.Setenv("SERVER_HOST", "127.0.0.1")
|
||||
t.Setenv("SERVER_PORT", "9090")
|
||||
|
||||
@@ -1164,6 +1214,17 @@ func TestGetServerAddressFromEnv(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetServerAddressFromConfigFile(t *testing.T) {
|
||||
configFile := filepath.Join(t.TempDir(), "setup.yaml")
|
||||
require.NoError(t, os.WriteFile(configFile, []byte("server:\n host: 192.0.2.40\n port: 9191\n"), 0o600))
|
||||
t.Setenv("CONFIG_FILE", configFile)
|
||||
t.Setenv("DATA_DIR", "")
|
||||
t.Setenv("SERVER_HOST", "")
|
||||
t.Setenv("SERVER_PORT", "")
|
||||
|
||||
require.Equal(t, "192.0.2.40:9191", GetServerAddress())
|
||||
}
|
||||
|
||||
func TestValidateAbsoluteHTTPURL(t *testing.T) {
|
||||
if err := ValidateAbsoluteHTTPURL("https://example.com/path"); err != nil {
|
||||
t.Fatalf("ValidateAbsoluteHTTPURL valid url error: %v", err)
|
||||
|
||||
@@ -133,6 +133,7 @@ var DefaultBedrockModelMapping = map[string]string{
|
||||
// Claude Fable
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
// Claude Opus
|
||||
"claude-opus-5": "us.anthropic.claude-opus-5-v1",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-opus-4-7": "us.anthropic.claude-opus-4-7-v1",
|
||||
"claude-opus-4-6-thinking": "us.anthropic.claude-opus-4-6-v1",
|
||||
|
||||
@@ -61,7 +61,7 @@ func (r *ollamaCloudUsageHandlerTestRepo) UpdateOllamaCloudUsageSnapshot(context
|
||||
func (r *ollamaCloudUsageHandlerTestRepo) DisableOllamaCloudUsageAutoRefresh(context.Context, *service.Account) error {
|
||||
return nil
|
||||
}
|
||||
func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, int) ([]service.Account, error) {
|
||||
func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]service.Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -273,4 +273,5 @@ func TestGetOllamaCloudUsageSettingsHandlerSuccess(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), `"enabled":false`)
|
||||
require.Contains(t, recorder.Body.String(), `"interval_minutes":60`)
|
||||
require.Contains(t, recorder.Body.String(), `"debounce_minutes":1`)
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
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"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -25,6 +26,16 @@ type GroupHandler struct {
|
||||
groupCapacityService *service.GroupCapacityService
|
||||
}
|
||||
|
||||
// GetLiveCapability 返回当前服务端是否具备生成 Live attestation 的运行环境。
|
||||
func (h *GroupHandler) GetLiveCapability(c *gin.Context) {
|
||||
err := liveattestation.NewProvider().Check(c.Request.Context())
|
||||
result := gin.H{"supported": err == nil}
|
||||
if err != nil {
|
||||
result["reason"] = err.Error()
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
type optionalLimitField struct {
|
||||
set bool
|
||||
value *float64
|
||||
@@ -125,6 +136,7 @@ type CreateGroupRequest struct {
|
||||
SupportedModelScopes []string `json:"supported_model_scopes"`
|
||||
// OpenAI Messages 调度配置(仅 openai 平台使用)
|
||||
AllowMessagesDispatch bool `json:"allow_messages_dispatch"`
|
||||
AllowLive bool `json:"allow_live"`
|
||||
RequireOAuthOnly bool `json:"require_oauth_only"`
|
||||
RequirePrivacySet bool `json:"require_privacy_set"`
|
||||
DefaultMappedModel string `json:"default_mapped_model"`
|
||||
@@ -183,6 +195,7 @@ type UpdateGroupRequest struct {
|
||||
SupportedModelScopes *[]string `json:"supported_model_scopes"`
|
||||
// OpenAI Messages 调度配置(仅 openai 平台使用)
|
||||
AllowMessagesDispatch *bool `json:"allow_messages_dispatch"`
|
||||
AllowLive *bool `json:"allow_live"`
|
||||
RequireOAuthOnly *bool `json:"require_oauth_only"`
|
||||
RequirePrivacySet *bool `json:"require_privacy_set"`
|
||||
DefaultMappedModel *string `json:"default_mapped_model"`
|
||||
@@ -499,6 +512,7 @@ func (h *GroupHandler) Create(c *gin.Context) {
|
||||
MCPXMLInject: req.MCPXMLInject,
|
||||
SupportedModelScopes: req.SupportedModelScopes,
|
||||
AllowMessagesDispatch: req.AllowMessagesDispatch,
|
||||
AllowLive: req.AllowLive,
|
||||
RequireOAuthOnly: req.RequireOAuthOnly,
|
||||
RequirePrivacySet: req.RequirePrivacySet,
|
||||
DefaultMappedModel: req.DefaultMappedModel,
|
||||
@@ -617,6 +631,7 @@ func (h *GroupHandler) Update(c *gin.Context) {
|
||||
MCPXMLInject: req.MCPXMLInject,
|
||||
SupportedModelScopes: req.SupportedModelScopes,
|
||||
AllowMessagesDispatch: req.AllowMessagesDispatch,
|
||||
AllowLive: req.AllowLive,
|
||||
RequireOAuthOnly: req.RequireOAuthOnly,
|
||||
RequirePrivacySet: req.RequirePrivacySet,
|
||||
DefaultMappedModel: req.DefaultMappedModel,
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
//go:build unit
|
||||
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// Saving settings is a whole-document PUT. A client that sends only the field it
|
||||
// cares about must not reset everything else: a payload as small as
|
||||
// `{"risk_control_enabled":true}` used to clear site_name, after which
|
||||
// getStringOrDefault rendered the empty value as the built-in default and the
|
||||
// login page silently changed name.
|
||||
|
||||
func TestUpdateSettingsPartialPayloadKeepsUnsentKeys(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeySiteName: "Example Gateway",
|
||||
service.SettingKeySiteSubtitle: "Example Gateway Platform",
|
||||
service.SettingKeySMTPHost: "smtp.example.com",
|
||||
service.SettingKeySMTPFrom: "noreply@example.com",
|
||||
service.SettingKeyTurnstileEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"risk_control_enabled": true}, nil)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyRiskControlEnabled],
|
||||
"the field the caller actually sent must be written")
|
||||
|
||||
require.Equal(t, "Example Gateway", repo.values[service.SettingKeySiteName])
|
||||
require.Equal(t, "Example Gateway Platform", repo.values[service.SettingKeySiteSubtitle])
|
||||
require.Equal(t, "smtp.example.com", repo.values[service.SettingKeySMTPHost])
|
||||
require.Equal(t, "noreply@example.com", repo.values[service.SettingKeySMTPFrom])
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyTurnstileEnabled])
|
||||
}
|
||||
|
||||
// A full payload keeps whole-document semantics: fields explicitly set to their
|
||||
// zero value are still cleared.
|
||||
func TestUpdateSettingsFullPayloadStillClearsSentEmptyFields(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeySiteName: "Example Gateway",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"site_name": ""}, nil)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
require.Equal(t, "", repo.values[service.SettingKeySiteName],
|
||||
"an explicitly sent empty value is a deliberate clear, not an omission")
|
||||
}
|
||||
|
||||
// smtp_from_email is the one request field whose JSON name differs from its
|
||||
// setting key; the alias keeps it from being treated as always-omitted.
|
||||
func TestUpdateSettingsSMTPFromAliasIsWritable(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeySMTPFrom: "old@example.com",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"smtp_from_email": "new@example.com"}, nil)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
require.Equal(t, "new@example.com", repo.values[service.SettingKeySMTPFrom])
|
||||
}
|
||||
@@ -154,6 +154,70 @@ func (h *SettingHandler) UpdateRateLimit429CooldownSettings(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// GetPanelRateLimitSettings 获取面板 API 限流配置
|
||||
// GET /api/v1/admin/settings/panel-rate-limit
|
||||
func (h *SettingHandler) GetPanelRateLimitSettings(c *gin.Context) {
|
||||
settings, err := h.settingService.GetPanelRateLimitSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.PanelRateLimitSettings{
|
||||
Enabled: settings.Enabled,
|
||||
UserRPM: settings.UserRPM,
|
||||
HeavyRPM: settings.HeavyRPM,
|
||||
ExemptAdmin: settings.ExemptAdmin,
|
||||
PublicIPRPM: settings.PublicIPRPM,
|
||||
})
|
||||
}
|
||||
|
||||
// UpdatePanelRateLimitSettingsRequest 更新面板 API 限流配置请求
|
||||
type UpdatePanelRateLimitSettingsRequest struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
UserRPM int `json:"user_rpm"`
|
||||
HeavyRPM int `json:"heavy_rpm"`
|
||||
ExemptAdmin bool `json:"exempt_admin"`
|
||||
PublicIPRPM int `json:"public_ip_rpm"`
|
||||
}
|
||||
|
||||
// UpdatePanelRateLimitSettings 更新面板 API 限流配置
|
||||
// PUT /api/v1/admin/settings/panel-rate-limit
|
||||
func (h *SettingHandler) UpdatePanelRateLimitSettings(c *gin.Context) {
|
||||
var req UpdatePanelRateLimitSettingsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
settings := &service.PanelRateLimitSettings{
|
||||
Enabled: req.Enabled,
|
||||
UserRPM: req.UserRPM,
|
||||
HeavyRPM: req.HeavyRPM,
|
||||
ExemptAdmin: req.ExemptAdmin,
|
||||
PublicIPRPM: req.PublicIPRPM,
|
||||
}
|
||||
|
||||
if err := h.settingService.SetPanelRateLimitSettings(c.Request.Context(), settings); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
updatedSettings, err := h.settingService.GetPanelRateLimitSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, dto.PanelRateLimitSettings{
|
||||
Enabled: updatedSettings.Enabled,
|
||||
UserRPM: updatedSettings.UserRPM,
|
||||
HeavyRPM: updatedSettings.HeavyRPM,
|
||||
ExemptAdmin: updatedSettings.ExemptAdmin,
|
||||
PublicIPRPM: updatedSettings.PublicIPRPM,
|
||||
})
|
||||
}
|
||||
|
||||
// GetStreamTimeoutSettings 获取流超时处理配置
|
||||
// GET /api/v1/admin/settings/stream-timeout
|
||||
func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) {
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
@@ -15,6 +16,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gin-gonic/gin/binding"
|
||||
)
|
||||
|
||||
// UpdateSettingsRequest 更新设置请求
|
||||
@@ -375,12 +377,70 @@ func (h *SettingHandler) ensureActorTotpForStepUp(c *gin.Context) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// settingKeyJSONAliases covers the request fields whose JSON name differs from
|
||||
// the setting key they persist to. Every other field of UpdateSettingsRequest
|
||||
// is named after its setting key.
|
||||
var settingKeyJSONAliases = map[string]string{
|
||||
"smtp_from_email": service.SettingKeySMTPFrom,
|
||||
}
|
||||
|
||||
// settingKeyByJSONName maps the value-typed top-level JSON fields of
|
||||
// UpdateSettingsRequest to the setting key each one writes. Resolved once from
|
||||
// the struct tags so new fields are covered without touching this file.
|
||||
//
|
||||
// Pointer-typed fields are deliberately excluded: they already carry their own
|
||||
// "omitted = keep the stored value" merge in UpdateSettings, and some of them
|
||||
// rely on being rewritten on every save to re-normalize fail-closed security
|
||||
// state (see TestUpdateSettingsMalformedForwardedClientIPHeadersRemainFailClosedWhenOmitted).
|
||||
// Only the value-typed fields are indistinguishable from a deliberate clear.
|
||||
var settingKeyByJSONName = buildSettingKeyByJSONName()
|
||||
|
||||
func buildSettingKeyByJSONName() map[string]string {
|
||||
t := reflect.TypeOf(UpdateSettingsRequest{})
|
||||
out := make(map[string]string, t.NumField())
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
field := t.Field(i)
|
||||
if field.Type.Kind() == reflect.Ptr {
|
||||
continue
|
||||
}
|
||||
name, _, _ := strings.Cut(field.Tag.Get("json"), ",")
|
||||
if name == "" || name == "-" {
|
||||
continue
|
||||
}
|
||||
if alias, ok := settingKeyJSONAliases[name]; ok {
|
||||
out[name] = alias
|
||||
continue
|
||||
}
|
||||
out[name] = name
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// omittedSettingKeys reports the setting keys this payload never mentioned.
|
||||
// Saving settings is a whole-document PUT, so without this a client that sends
|
||||
// only the one field it cares about resets every other field to a zero value.
|
||||
func omittedSettingKeys(sentFields map[string]json.RawMessage) service.OmittedSettingKeys {
|
||||
omitted := make(service.OmittedSettingKeys, len(settingKeyByJSONName))
|
||||
for jsonName, settingKey := range settingKeyByJSONName {
|
||||
if _, sent := sentFields[jsonName]; !sent {
|
||||
omitted[settingKey] = struct{}{}
|
||||
}
|
||||
}
|
||||
return omitted
|
||||
}
|
||||
|
||||
func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
var req UpdateSettingsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
var sentFields map[string]json.RawMessage
|
||||
if err := c.ShouldBindBodyWith(&sentFields, binding.JSON); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
var req UpdateSettingsRequest
|
||||
if err := c.ShouldBindBodyWith(&req, binding.JSON); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
omitted := omittedSettingKeys(sentFields)
|
||||
|
||||
previousSettings, err := h.settingService.GetAllSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
@@ -1692,7 +1752,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
},
|
||||
ForceEmailOnThirdPartySignup: boolValueOrDefault(req.ForceEmailOnThirdPartySignup, previousAuthSourceDefaults.ForceEmailOnThirdPartySignup),
|
||||
}
|
||||
if err := h.settingService.UpdateSettingsWithAuthSourceDefaults(c.Request.Context(), settings, authSourceDefaults); err != nil {
|
||||
if err := h.settingService.UpdateSettingsWithAuthSourceDefaultsOmitting(c.Request.Context(), settings, authSourceDefaults, omitted); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -110,6 +110,7 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
}
|
||||
|
||||
model := c.Query("model")
|
||||
requestID := strings.TrimSpace(c.Query("request_id"))
|
||||
billingMode := strings.TrimSpace(c.Query("billing_mode"))
|
||||
|
||||
var requestType *int16
|
||||
@@ -176,6 +177,7 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
RequestID: requestID,
|
||||
Model: model,
|
||||
ModelFilterSource: usagestats.ModelSourceRequested,
|
||||
RequestType: requestType,
|
||||
|
||||
@@ -107,6 +107,18 @@ func TestAdminUsageListExactTotalTrue(t *testing.T) {
|
||||
require.True(t, repo.listFilters.ExactTotal)
|
||||
}
|
||||
|
||||
func TestAdminUsageListRequestIDFilter(t *testing.T) {
|
||||
repo := &adminUsageRepoCapture{}
|
||||
router := newAdminUsageRequestTypeTestRouter(repo)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/usage?request_id=req-0123", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "req-0123", repo.listFilters.RequestID)
|
||||
}
|
||||
|
||||
func TestAdminUsageListInvalidExactTotal(t *testing.T) {
|
||||
repo := &adminUsageRepoCapture{}
|
||||
router := newAdminUsageRequestTypeTestRouter(repo)
|
||||
|
||||
@@ -2978,6 +2978,17 @@ func (r *oauthPendingFlowUserRepo) Create(ctx context.Context, user *service.Use
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *oauthPendingFlowUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
|
||||
aliasExists, err := r.ExistsByEmailAlias(ctx, user.Email)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if aliasExists {
|
||||
return service.ErrEmailExists
|
||||
}
|
||||
return r.Create(ctx, user)
|
||||
}
|
||||
|
||||
func (r *oauthPendingFlowUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
|
||||
entity, err := r.client.User.Get(ctx, id)
|
||||
if err != nil {
|
||||
@@ -3185,6 +3196,20 @@ func (r *oauthPendingFlowUserRepo) ExistsByEmail(ctx context.Context, email stri
|
||||
return count > 0, err
|
||||
}
|
||||
|
||||
func (r *oauthPendingFlowUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
|
||||
identity := service.NormalizeEmailForAliasDedup(email)
|
||||
emails, err := r.client.User.Query().Select(dbuser.FieldEmail).Strings(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
for _, stored := range emails {
|
||||
if service.NormalizeEmailForAliasDedup(stored) == identity {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (r *oauthPendingFlowUserRepo) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
|
||||
panic("unexpected RemoveGroupFromAllowedGroups call")
|
||||
}
|
||||
|
||||
@@ -42,6 +42,9 @@ func (h *BatchImageHandler) Submit(c *gin.Context) {
|
||||
if !h.checkSecurityAuditBeforeSubmit(c, &req) {
|
||||
return
|
||||
}
|
||||
if sessionID := service.ExtractClientSessionID(c); sessionID != "" {
|
||||
req.SessionID = &sessionID
|
||||
}
|
||||
got, err := h.service.Submit(c.Request.Context(), owner, req, c.GetHeader("Idempotency-Key"))
|
||||
if err != nil {
|
||||
batchImageError(c, err)
|
||||
|
||||
@@ -204,6 +204,7 @@ func groupFromServiceBase(g *service.Group) Group {
|
||||
FallbackGroupID: g.FallbackGroupID,
|
||||
FallbackGroupIDOnInvalidRequest: g.FallbackGroupIDOnInvalidRequest,
|
||||
AllowMessagesDispatch: g.AllowMessagesDispatch,
|
||||
AllowLive: g.AllowLive,
|
||||
RequireOAuthOnly: g.RequireOAuthOnly,
|
||||
RequirePrivacySet: g.RequirePrivacySet,
|
||||
RPMLimit: g.RPMLimit,
|
||||
@@ -669,6 +670,7 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog {
|
||||
MediaType: l.MediaType,
|
||||
UserAgent: l.UserAgent,
|
||||
IPAddress: l.IPAddress,
|
||||
SessionID: l.SessionID,
|
||||
CacheTTLOverridden: l.CacheTTLOverridden,
|
||||
BillingMode: l.BillingMode,
|
||||
CreatedAt: l.CreatedAt,
|
||||
|
||||
@@ -394,6 +394,15 @@ type RateLimit429CooldownSettings struct {
|
||||
CooldownSeconds int `json:"cooldown_seconds"`
|
||||
}
|
||||
|
||||
// PanelRateLimitSettings 面板 API 限流配置 DTO
|
||||
type PanelRateLimitSettings struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
UserRPM int `json:"user_rpm"`
|
||||
HeavyRPM int `json:"heavy_rpm"`
|
||||
ExemptAdmin bool `json:"exempt_admin"`
|
||||
PublicIPRPM int `json:"public_ip_rpm"`
|
||||
}
|
||||
|
||||
// StreamTimeoutSettings 流超时处理配置 DTO
|
||||
type StreamTimeoutSettings struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
|
||||
@@ -132,6 +132,8 @@ type Group struct {
|
||||
|
||||
// OpenAI Messages 调度开关(用户侧需要此字段判断是否展示 Claude Code 教程)
|
||||
AllowMessagesDispatch bool `json:"allow_messages_dispatch"`
|
||||
// OpenAI Live 接口开关
|
||||
AllowLive bool `json:"allow_live"`
|
||||
|
||||
// 账号过滤控制(仅 OpenAI/Antigravity 平台有效)
|
||||
RequireOAuthOnly bool `json:"require_oauth_only"`
|
||||
@@ -521,6 +523,9 @@ type UsageLog struct {
|
||||
UserAgent *string `json:"user_agent"`
|
||||
// IPAddress is visible to the owner of the usage record.
|
||||
IPAddress *string `json:"ip_address,omitempty"`
|
||||
// SessionID is the explicit client-provided request correlation identifier
|
||||
// (e.g. the session_id / X-Session-Id headers). Omitted when absent.
|
||||
SessionID *string `json:"session_id,omitempty"`
|
||||
|
||||
// Cache TTL Override 标记
|
||||
CacheTTLOverridden bool `json:"cache_ttl_overridden"`
|
||||
|
||||
@@ -31,9 +31,12 @@ const (
|
||||
EndpointGeminiModels = "/v1beta/models"
|
||||
)
|
||||
|
||||
const EndpointAntigravityGenerateContent = "/v1internal:streamGenerateContent"
|
||||
|
||||
// gin.Context keys used by the middleware and helpers below.
|
||||
const (
|
||||
ctxKeyInboundEndpoint = "_gateway_inbound_endpoint"
|
||||
ctxKeyInboundEndpoint = "_gateway_inbound_endpoint"
|
||||
ctxKeyActualUpstreamEndpoint = "_gateway_actual_upstream_endpoint"
|
||||
)
|
||||
|
||||
// ──────────────────────────────────────────────────────────
|
||||
@@ -297,6 +300,13 @@ func GetInboundEndpoint(c *gin.Context) string {
|
||||
// and the account platform. Handlers call this after scheduling an
|
||||
// account, passing account.Platform.
|
||||
func GetUpstreamEndpoint(c *gin.Context, platform string) string {
|
||||
if c != nil {
|
||||
if value, ok := c.Get(ctxKeyActualUpstreamEndpoint); ok {
|
||||
if endpoint, ok := value.(string); ok && endpoint != "" {
|
||||
return endpoint
|
||||
}
|
||||
}
|
||||
}
|
||||
inbound := GetInboundEndpoint(c)
|
||||
rawPath := ""
|
||||
if c != nil && c.Request != nil && c.Request.URL != nil {
|
||||
@@ -304,3 +314,15 @@ func GetUpstreamEndpoint(c *gin.Context, platform string) string {
|
||||
}
|
||||
return DeriveUpstreamEndpoint(inbound, rawPath, platform)
|
||||
}
|
||||
|
||||
func setActualUpstreamEndpoint(c *gin.Context, endpoint string) {
|
||||
if c != nil {
|
||||
c.Set(ctxKeyActualUpstreamEndpoint, strings.TrimSpace(endpoint))
|
||||
}
|
||||
}
|
||||
|
||||
func shouldUseAntigravityCompat(account *service.Account) bool {
|
||||
return account != nil &&
|
||||
account.Platform == service.PlatformAntigravity &&
|
||||
account.Type == service.AccountTypeOAuth
|
||||
}
|
||||
|
||||
@@ -146,6 +146,40 @@ func TestDeriveUpstreamEndpoint(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldUseAntigravityCompat(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account *service.Account
|
||||
want bool
|
||||
}{
|
||||
{"oauth", &service.Account{Platform: service.PlatformAntigravity, Type: service.AccountTypeOAuth}, true},
|
||||
{"setup token", &service.Account{Platform: service.PlatformAntigravity, Type: service.AccountTypeSetupToken}, false},
|
||||
{"upstream", &service.Account{Platform: service.PlatformAntigravity, Type: service.AccountTypeUpstream}, false},
|
||||
{"api key", &service.Account{Platform: service.PlatformAntigravity, Type: service.AccountTypeAPIKey}, false},
|
||||
{"anthropic oauth", &service.Account{Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth}, false},
|
||||
{"nil", nil, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, shouldUseAntigravityCompat(tt.account))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetUpstreamEndpointPrefersRuntimeOverride(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointChatCompletions, nil)
|
||||
c.Set(ctxKeyInboundEndpoint, EndpointChatCompletions)
|
||||
|
||||
setActualUpstreamEndpoint(c, EndpointAntigravityGenerateContent)
|
||||
require.Equal(t, EndpointAntigravityGenerateContent, GetUpstreamEndpoint(c, service.PlatformAntigravity))
|
||||
|
||||
setActualUpstreamEndpoint(c, "")
|
||||
require.Equal(t, EndpointMessages, GetUpstreamEndpoint(c, service.PlatformAntigravity))
|
||||
}
|
||||
|
||||
func TestResolveOpenAIUpstreamEndpointPrefersForwardResult(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -535,6 +535,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
|
||||
forceCacheBilling := fs.ForceCacheBilling
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
@@ -547,6 +548,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
SessionID: sessionID,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
ForceCacheBilling: forceCacheBilling,
|
||||
APIKeyService: h.apiKeyService,
|
||||
@@ -970,6 +972,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
|
||||
forceCacheBilling := fs.ForceCacheBilling
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
@@ -982,6 +985,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
SessionID: sessionID,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
ForceCacheBilling: forceCacheBilling,
|
||||
APIKeyService: h.apiKeyService,
|
||||
|
||||
@@ -240,6 +240,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
forwardBody = h.gatewayService.ReplaceModelInBody(body, channelMapping.MappedModel)
|
||||
}
|
||||
var result *service.ForwardResult
|
||||
setActualUpstreamEndpoint(c, "")
|
||||
if account.Platform == service.PlatformGemini {
|
||||
if h.geminiCompatService == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusBadGateway, "upstream_error", "Gemini compatibility service is not configured")
|
||||
@@ -249,6 +250,16 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
result, err = h.geminiCompatService.ForwardAsChatCompletions(c.Request.Context(), c, account, forwardBody)
|
||||
} else if shouldUseAntigravityCompat(account) {
|
||||
if h.antigravityGatewayService == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusBadGateway, "upstream_error", "Antigravity compatibility service is not configured")
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
return
|
||||
}
|
||||
setActualUpstreamEndpoint(c, EndpointAntigravityGenerateContent)
|
||||
result, err = h.antigravityGatewayService.ForwardAsChatCompletions(c.Request.Context(), c, account, forwardBody, parsedReq)
|
||||
} else {
|
||||
result, err = h.gatewayService.ForwardAsChatCompletions(c.Request.Context(), c, account, forwardBody, parsedReq)
|
||||
}
|
||||
@@ -298,6 +309,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
@@ -312,6 +324,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
reqLog.Error("gateway.cc.record_usage_failed",
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -53,6 +54,22 @@ func TestGatewayEnsureForwardErrorResponse_AppendsSSEAfterWritten(t *testing.T)
|
||||
assert.Contains(t, w.Body.String(), `data: {"type":"error"`)
|
||||
}
|
||||
|
||||
func TestGatewayEnsureForwardErrorResponse_SkipsCommittedSSEError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
_, _ = c.Writer.WriteString("event: error\ndata: {\"type\":\"error\"}\n\n")
|
||||
service.MarkResponseCommitted(c)
|
||||
|
||||
h := &GatewayHandler{}
|
||||
wrote := h.ensureForwardErrorResponse(c, true)
|
||||
|
||||
require.False(t, wrote)
|
||||
require.Equal(t, 1, strings.Count(w.Body.String(), "event: error"))
|
||||
}
|
||||
|
||||
// case B 回归:Anthropic-backed /responses,Writer 已被写过时
|
||||
// ensureForwardErrorResponse 仍要发 response.failed。
|
||||
func TestGatewayEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsResponseFailed(t *testing.T) {
|
||||
|
||||
@@ -228,7 +228,21 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
if channelMapping.Mapped {
|
||||
forwardBody = h.gatewayService.ReplaceModelInBody(body, channelMapping.MappedModel)
|
||||
}
|
||||
result, err := h.gatewayService.ForwardAsResponses(requestCtx, c, account, forwardBody, parsedReq)
|
||||
var result *service.ForwardResult
|
||||
setActualUpstreamEndpoint(c, "")
|
||||
if shouldUseAntigravityCompat(account) {
|
||||
if h.antigravityGatewayService == nil {
|
||||
h.responsesErrorResponse(c, http.StatusBadGateway, "upstream_error", "Antigravity compatibility service is not configured")
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
return
|
||||
}
|
||||
setActualUpstreamEndpoint(c, EndpointAntigravityGenerateContent)
|
||||
result, err = h.antigravityGatewayService.ForwardAsResponses(requestCtx, c, account, forwardBody, parsedReq)
|
||||
} else {
|
||||
result, err = h.gatewayService.ForwardAsResponses(requestCtx, c, account, forwardBody, parsedReq)
|
||||
}
|
||||
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
@@ -276,6 +290,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: result,
|
||||
@@ -290,6 +305,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
reqLog.Error("gateway.responses.record_usage_failed",
|
||||
|
||||
@@ -532,6 +532,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
|
||||
forceCacheBilling := fs.ForceCacheBilling
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{
|
||||
Result: result,
|
||||
@@ -549,6 +550,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
LongContextMultiplier: 2.0, // 超出部分双倍计费
|
||||
ForceCacheBilling: forceCacheBilling,
|
||||
APIKeyService: h.apiKeyService,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -465,6 +465,7 @@ func recordGrokMediaUsage(
|
||||
) {
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
payloadForHash := body
|
||||
if len(payloadForHash) == 0 && strings.TrimSpace(requestID) != "" {
|
||||
payloadForHash = []byte(requestID)
|
||||
@@ -493,6 +494,7 @@ func recordGrokMediaUsage(
|
||||
RequestPayloadHash: service.HashUsageRequestPayload(payloadForHash),
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: channelUsageFields,
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -233,6 +233,7 @@ func (h *OpenAIGatewayHandler) recordAlphaSearchUsage(
|
||||
) {
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
@@ -252,6 +253,7 @@ func (h *OpenAIGatewayHandler) recordAlphaSearchUsage(
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: channelMapping.ToUsageFields(requestedModel, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -340,6 +340,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
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) {
|
||||
@@ -355,6 +356,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
|
||||
@@ -243,6 +243,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
@@ -257,6 +258,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -42,6 +42,28 @@ func TestGatewayChatCredentialStopDoesNotSelectAnotherAccountAndReturnsSafe503(t
|
||||
require.NotContains(t, recorder.Body.String(), "client_secret")
|
||||
}
|
||||
|
||||
func TestGatewayChatAntigravityCredentialFailureReturnsActionableMessage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
(&GatewayHandler{}).handleCCFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
Stage: service.GatewayFailureStageAccountAuth,
|
||||
Scope: service.GatewayFailureScopeAccount,
|
||||
Reason: service.AntigravityCredentialRejectedReason,
|
||||
NextAccountAction: service.NextAccountRetry,
|
||||
ClientStatusCode: http.StatusBadGateway,
|
||||
ClientMessage: service.AntigravityCredentialRejectedClientMessage,
|
||||
ResponseBody: []byte(`{"error":{"message":"Invalid bearer token","refresh_token":"must-not-leak"}}`),
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusBadGateway, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), service.AntigravityCredentialRejectedClientMessage)
|
||||
require.NotContains(t, strings.ToLower(recorder.Body.String()), "bearer")
|
||||
require.NotContains(t, strings.ToLower(recorder.Body.String()), "refresh_token")
|
||||
}
|
||||
|
||||
func TestGatewayChatInferenceExhaustionRestoresRetryAfter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"runtime/debug"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
@@ -43,6 +44,49 @@ type OpenAIGatewayHandler struct {
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
type openAIWSTurnChannelMappingSnapshot struct {
|
||||
turn int
|
||||
mapping service.ChannelMappingResult
|
||||
}
|
||||
|
||||
var errOpenAIWSUnsupportedModelSwitch = errors.New("selected account does not support websocket model switch")
|
||||
|
||||
func newOpenAIWSUnsupportedModelSwitchError(model string) error {
|
||||
cause := fmt.Errorf("%w: model %q", errOpenAIWSUnsupportedModelSwitch, strings.TrimSpace(model))
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "model switch requires reconnect", cause)
|
||||
}
|
||||
|
||||
func shouldReportOpenAIWSProxyAccountFailure(err error) bool {
|
||||
return err != nil && !errors.Is(err, errOpenAIWSUnsupportedModelSwitch)
|
||||
}
|
||||
|
||||
func openAIWSTurnBillingModel(result *service.OpenAIForwardResult, mapping service.ChannelMappingResult, requestedModel, upstreamModel string) string {
|
||||
billingModel := ""
|
||||
if result != nil {
|
||||
billingModel = strings.TrimSpace(result.BillingModel)
|
||||
}
|
||||
if billingModel == "" {
|
||||
billingModel = strings.TrimSpace(upstreamModel)
|
||||
}
|
||||
if billingModel == "" {
|
||||
billingModel = strings.TrimSpace(requestedModel)
|
||||
}
|
||||
|
||||
requestedModel = strings.TrimSpace(requestedModel)
|
||||
switch mapping.BillingModelSource {
|
||||
case service.BillingModelSourceRequested:
|
||||
if requestedModel != "" {
|
||||
billingModel = requestedModel
|
||||
}
|
||||
case service.BillingModelSourceChannelMapped:
|
||||
mappedModel := strings.TrimSpace(mapping.MappedModel)
|
||||
if mappedModel != "" && mappedModel != requestedModel {
|
||||
billingModel = mappedModel
|
||||
}
|
||||
}
|
||||
return billingModel
|
||||
}
|
||||
|
||||
type grokMediaEligibilityProber interface {
|
||||
ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error)
|
||||
}
|
||||
@@ -386,6 +430,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
var passthroughFailoverState openAIPassthroughFailoverState
|
||||
|
||||
// 生图意图的 /v1/responses 请求必须调度到确实支持 Responses API 的账号,否则
|
||||
// 会在 forward 阶段被静默降级为无法生图的 Chat Completions 直转(#4417)。
|
||||
@@ -482,13 +527,17 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// 用扣除 compact 心跳字节的口径快照:心跳注释不构成语义响应,
|
||||
// 不能因心跳字节变化而放弃 failover 换号(#3887)。
|
||||
writerSizeBeforeForward := service.OpenAICompactKeepaliveAdjustedWrittenSize(c)
|
||||
// 跨 passthrough 边界的 failover:从 Kiro 等透传账号切到 Bedrock 等非透传账号前,
|
||||
// 从不可变的 canonical forwardBody 派生本次尝试 body 并整块剔除上游私有的加密
|
||||
// reasoning item(含耦合的 id/summary),避免非透传上游 400 拒绝 Kiro reasoning 形态。
|
||||
attemptBody := h.deriveOpenAIForwardAttemptBody(reqLog, forwardBody, account, &passthroughFailoverState)
|
||||
result, err := func() (*service.OpenAIForwardResult, error) {
|
||||
defer func() {
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
}()
|
||||
return h.gatewayService.Forward(c.Request.Context(), c, account, forwardBody)
|
||||
return h.gatewayService.Forward(c.Request.Context(), c, account, attemptBody)
|
||||
}()
|
||||
cyberBlockKeyHTTP := ""
|
||||
if service.GetOpsCyberPolicy(c) != nil {
|
||||
@@ -627,6 +676,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
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
|
||||
@@ -644,6 +694,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
@@ -1138,6 +1189,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
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) {
|
||||
@@ -1154,6 +1206,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
@@ -1515,7 +1568,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
previousResponseCanMove := !firstMessageToolCoverage.HasFunctionCallOutput || firstMessageToolCoverage.ContextCoversAllCallIDs
|
||||
reqLog = reqLog.With(
|
||||
zap.Bool("ws_ingress", true),
|
||||
zap.String("model", reqModel),
|
||||
zap.String("session_initial_model", reqModel),
|
||||
zap.Bool("has_previous_response_id", previousResponseID != ""),
|
||||
zap.String("previous_response_id_kind", previousResponseIDKind),
|
||||
)
|
||||
@@ -1767,6 +1820,10 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
reasoningEffortMappings = apiKey.Group.ReasoningEffortMappings
|
||||
}
|
||||
var requestPayloadHash string
|
||||
// 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})
|
||||
hooks := &service.OpenAIWSIngressHooks{
|
||||
InitialRequestModel: reqModel,
|
||||
MaxReasoningEffort: maxReasoningEffort,
|
||||
@@ -1791,6 +1848,22 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
}
|
||||
return nil
|
||||
},
|
||||
MapRequestModel: func(turn int, originalModel string) (string, error) {
|
||||
model := strings.TrimSpace(originalModel)
|
||||
if model == "" {
|
||||
model = reqModel
|
||||
}
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, model)
|
||||
mappedModelUnchanged := false
|
||||
if previous := turnChannelMapping.Load(); previous != nil && previous.turn < turn {
|
||||
mappedModelUnchanged = strings.TrimSpace(previous.mapping.MappedModel) == strings.TrimSpace(mapping.MappedModel)
|
||||
}
|
||||
if turn > 1 && !mappedModelUnchanged && !account.IsModelSupported(model) && !account.IsModelSupported(mapping.MappedModel) {
|
||||
return "", newOpenAIWSUnsupportedModelSwitchError(mapping.MappedModel)
|
||||
}
|
||||
turnChannelMapping.Store(&openAIWSTurnChannelMappingSnapshot{turn: turn, mapping: mapping})
|
||||
return mapping.MappedModel, nil
|
||||
},
|
||||
BeforeTurn: func(turn int) error {
|
||||
// turn==1 的会话屏蔽已由握手层检查覆盖;连接内 flag 只拦截后续 turn。
|
||||
if cyberBlockedThisConn {
|
||||
@@ -1832,7 +1905,27 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// 届时 defer 已清除标记)。
|
||||
defer clearCyberPolicyTurnState(c)
|
||||
releaseTurnSlots()
|
||||
h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, turnErr != nil, cyberBlockKey, clientRequestedUsageFields(c, channelMappingWS, reqModel, ""), requestPayloadHash)
|
||||
turnRequestedModel := reqModel
|
||||
turnUpstreamModel := ""
|
||||
if result != nil && turn > 1 {
|
||||
if model := strings.TrimSpace(result.Model); model != "" {
|
||||
turnRequestedModel = model
|
||||
}
|
||||
}
|
||||
if result != nil {
|
||||
turnUpstreamModel = strings.TrimSpace(result.UpstreamModel)
|
||||
}
|
||||
var turnMapping service.ChannelMappingResult
|
||||
if snapshot := turnChannelMapping.Load(); snapshot != nil && snapshot.turn == turn {
|
||||
turnMapping = snapshot.mapping
|
||||
} else {
|
||||
turnMapping, _ = h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, turnRequestedModel)
|
||||
}
|
||||
if turnUpstreamModel == "" {
|
||||
turnUpstreamModel = turnRequestedModel
|
||||
}
|
||||
turnUsageFields := turnMapping.ToUsageFields(turnRequestedModel, turnUpstreamModel)
|
||||
h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, turnRequestedModel, turnErr != nil, cyberBlockKey, turnUsageFields, requestPayloadHash)
|
||||
if service.GetOpsCyberPolicy(c) != nil {
|
||||
cyberBlockedThisConn = true
|
||||
}
|
||||
@@ -1854,14 +1947,26 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
if result == nil {
|
||||
return
|
||||
}
|
||||
result.BillingModel = openAIWSTurnBillingModel(result, turnMapping, turnRequestedModel, turnUpstreamModel)
|
||||
reqLog.Debug("openai.websocket_turn_billing",
|
||||
zap.Int("turn", turn),
|
||||
zap.String("turn_requested_model", turnRequestedModel),
|
||||
zap.String("turn_upstream_model", turnUpstreamModel),
|
||||
zap.String("billing_model", result.BillingModel),
|
||||
)
|
||||
// 排除 spark 影子:其 codex_* 仅由 QueryUsage(/wham/usage bengalfox)更新(外审第7轮 P1)。
|
||||
if account.Type == service.AccountTypeOAuth && !account.IsShadow() {
|
||||
h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(ctx, account.ID, result.ResponseHeaders)
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs)
|
||||
scheduleModel := turnUpstreamModel
|
||||
if scheduleModel == "" {
|
||||
scheduleModel = turnRequestedModel
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, scheduleModel, openAIForwardSucceededForScheduling(result), result.FirstTokenMs)
|
||||
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(ctx, result, func(taskCtx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
|
||||
@@ -1877,7 +1982,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingWS, reqModel, result.UpstreamModel),
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: turnUsageFields,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
reqLog.Error("openai.websocket_record_usage_failed",
|
||||
@@ -1890,11 +1996,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
},
|
||||
}
|
||||
|
||||
// 应用渠道模型映射到 WebSocket 首条消息
|
||||
wsFirstMessage := firstMessage
|
||||
if channelMappingWS.Mapped {
|
||||
wsFirstMessage = h.gatewayService.ReplaceModelInBody(firstMessage, channelMappingWS.MappedModel)
|
||||
}
|
||||
// 切组/会话失配防护:previous_response_id 未在当前分组命中粘连账号(StickyPreviousHit=false),
|
||||
// 说明该会话链不属于本次调度到的账号,原样转发会触发上游会话链鉴权失败(“鉴权失败,请检查 API Key”)。
|
||||
// 故剥离首包里的 previous_response_id,改用首包内 input 重建上下文;带 function_call_output 的
|
||||
@@ -1938,7 +2040,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
|
||||
if shouldReportOpenAIWSProxyAccountFailure(err) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), false, nil)
|
||||
}
|
||||
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
|
||||
proxyFailedFields := []zap.Field{
|
||||
zap.Int64("account_id", account.ID),
|
||||
@@ -2221,7 +2325,9 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
}
|
||||
|
||||
func credentialFailoverClientResponse(failoverErr *service.UpstreamFailoverError) (int, string) {
|
||||
_ = failoverErr
|
||||
if failoverErr != nil && failoverErr.Reason == service.AntigravityCredentialRejectedReason {
|
||||
return http.StatusBadGateway, service.AntigravityCredentialRejectedClientMessage
|
||||
}
|
||||
return http.StatusServiceUnavailable, service.GrokCredentialUnavailableClientMessage
|
||||
}
|
||||
|
||||
@@ -2884,6 +2990,8 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
userAgent = c.GetHeader("User-Agent")
|
||||
clientIPStr = strings.TrimSpace(ip.GetClientIP(c))
|
||||
}
|
||||
// 提前拍成标量,避免在下方 goroutine 内访问 gin.Context。
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
apiKeyPrefix := ""
|
||||
if apiKey != nil {
|
||||
apiKeyPrefix = keyPrefix(apiKey.Key, 8)
|
||||
@@ -2940,6 +3048,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIPStr,
|
||||
SessionID: sessionID,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: apiKeySvc,
|
||||
ChannelUsageFields: channelFields,
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -1369,6 +1370,229 @@ func TestOpenAIResponsesWebSocket_PassthroughUsageLogLeavesUserAgentNilWhenMissi
|
||||
require.Equal(t, "medium", *got.log.ReasoningEffort)
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_PassthroughTracksModelPerTurn(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"sol","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"terra","stream":false}`,
|
||||
channelMapping: map[string]string{
|
||||
"sol": "sol-channel",
|
||||
"terra": "terra-channel",
|
||||
},
|
||||
accountModelMapping: map[string]any{
|
||||
"sol": "gpt-5.6-sol",
|
||||
"terra": "gpt-5.6-terra",
|
||||
"sol-channel": "gpt-5.6-sol",
|
||||
"terra-channel": "gpt-5.6-terra",
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
|
||||
require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
|
||||
require.Len(t, got.clientEvents, 2)
|
||||
require.Equal(t, "sol", gjson.GetBytes(got.clientEvents[0], "response.model").String())
|
||||
require.Equal(t, "terra", gjson.GetBytes(got.clientEvents[1], "response.model").String())
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.Equal(t, "sol", got.logs[0].Model)
|
||||
require.Equal(t, "sol", got.logs[0].RequestedModel)
|
||||
require.NotNil(t, got.logs[0].UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-sol", *got.logs[0].UpstreamModel)
|
||||
require.NotNil(t, got.logs[0].ModelMappingChain)
|
||||
require.Equal(t, "sol→sol-channel→gpt-5.6-sol", *got.logs[0].ModelMappingChain)
|
||||
|
||||
require.Equal(t, "terra", got.logs[1].Model)
|
||||
require.Equal(t, "terra", got.logs[1].RequestedModel)
|
||||
require.NotNil(t, got.logs[1].UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-terra", *got.logs[1].UpstreamModel)
|
||||
require.NotNil(t, got.logs[1].ModelMappingChain)
|
||||
require.Equal(t, "terra→terra-channel→gpt-5.6-terra", *got.logs[1].ModelMappingChain)
|
||||
require.InDelta(t, got.logs[1].TotalCost*2, got.logs[0].TotalCost, 1e-12,
|
||||
"each turn must be billed with its own channel-mapped model")
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_UnchangedChannelTargetOutsideAccountMappingKeysRemainsValid(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","stream":false}`,
|
||||
channelMapping: map[string]string{
|
||||
"public-alias": "gpt-5.6-sol",
|
||||
},
|
||||
accountModelMapping: map[string]any{
|
||||
"public-alias": "gpt-5.6-terra",
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
|
||||
require.Len(t, got.clientEvents, 2)
|
||||
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[0], "response.model").String())
|
||||
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[1], "response.model").String())
|
||||
require.Len(t, got.logs, 2)
|
||||
for _, usageLog := range got.logs {
|
||||
require.Equal(t, "public-alias", usageLog.RequestedModel)
|
||||
require.NotNil(t, usageLog.UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-sol", *usageLog.UpstreamModel)
|
||||
require.NotNil(t, usageLog.ModelMappingChain)
|
||||
require.Equal(t, "public-alias→gpt-5.6-sol", *usageLog.ModelMappingChain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_PassthroughKeepsTurnMappingSnapshot(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"sol","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"sol","stream":false}`,
|
||||
channelMapping: map[string]string{
|
||||
"sol": "gpt-5.6-sol",
|
||||
},
|
||||
afterFirstUpstreamRequest: func(channelSvc *service.ChannelService) error {
|
||||
if channelSvc == nil {
|
||||
return errors.New("channel service is nil")
|
||||
}
|
||||
_, err := channelSvc.Update(context.Background(), 7701, &service.UpdateChannelInput{
|
||||
ModelMapping: map[string]map[string]string{
|
||||
service.PlatformOpenAI: {"sol": "gpt-5.6-terra"},
|
||||
},
|
||||
})
|
||||
return err
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
|
||||
require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.Equal(t, "sol", got.logs[0].Model)
|
||||
require.NotNil(t, got.logs[0].UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-sol", *got.logs[0].UpstreamModel)
|
||||
require.NotNil(t, got.logs[0].ModelMappingChain)
|
||||
require.Equal(t, "sol→gpt-5.6-sol", *got.logs[0].ModelMappingChain)
|
||||
require.InDelta(t, 40e-6, got.logs[0].TotalCost, 1e-12,
|
||||
"the in-flight turn must retain the channel-mapped billing model used when it was sent")
|
||||
|
||||
require.Equal(t, "sol", got.logs[1].Model)
|
||||
require.NotNil(t, got.logs[1].UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-terra", *got.logs[1].UpstreamModel)
|
||||
require.NotNil(t, got.logs[1].ModelMappingChain)
|
||||
require.Equal(t, "sol→gpt-5.6-terra", *got.logs[1].ModelMappingChain)
|
||||
require.InDelta(t, got.logs[1].TotalCost*2, got.logs[0].TotalCost, 1e-12,
|
||||
"the next turn must use the updated channel mapping")
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CtxPoolAppliesPerTurnMappingAndPreservesRequestedModel(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.6-terra","stream":false}`,
|
||||
ingressMode: service.OpenAIWSIngressModeCtxPool,
|
||||
billingModelSource: service.BillingModelSourceRequested,
|
||||
channelMapping: map[string]string{
|
||||
"gpt-5.6-terra": "gpt-5.6-sol",
|
||||
},
|
||||
})
|
||||
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
|
||||
require.Len(t, got.clientEvents, 2)
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.clientEvents[0], "response.model").String())
|
||||
require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.clientEvents[1], "response.model").String())
|
||||
|
||||
require.Len(t, got.logs, 2)
|
||||
require.Equal(t, "gpt-5.6-sol", got.logs[0].RequestedModel)
|
||||
require.Nil(t, got.logs[0].ModelMappingChain)
|
||||
require.InDelta(t, 40e-6, got.logs[0].TotalCost, 1e-12)
|
||||
require.Equal(t, "gpt-5.6-terra", got.logs[1].RequestedModel)
|
||||
require.NotNil(t, got.logs[1].ModelMappingChain)
|
||||
require.Equal(t, "gpt-5.6-terra→gpt-5.6-sol", *got.logs[1].ModelMappingChain)
|
||||
require.InDelta(t, 20e-6, got.logs[1].TotalCost, 1e-12,
|
||||
"BillingModelSourceRequested must use the client model before channel mapping")
|
||||
}
|
||||
|
||||
func TestOpenAIWSTurnBillingModelPreservesImagePricingModel(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
resultModel string
|
||||
mapping service.ChannelMappingResult
|
||||
requestedModel string
|
||||
upstreamModel string
|
||||
wantBillingModel string
|
||||
}{
|
||||
{
|
||||
name: "upstream billing preserves image model",
|
||||
resultModel: "gpt-image-2",
|
||||
mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceUpstream},
|
||||
requestedModel: "gpt-5.6-sol",
|
||||
upstreamModel: "gpt-5.6-sol",
|
||||
wantBillingModel: "gpt-image-2",
|
||||
},
|
||||
{
|
||||
name: "unmapped channel preserves image model",
|
||||
resultModel: "gpt-image-2",
|
||||
mapping: service.ChannelMappingResult{MappedModel: "gpt-5.6-sol", BillingModelSource: service.BillingModelSourceChannelMapped},
|
||||
requestedModel: "gpt-5.6-sol",
|
||||
upstreamModel: "gpt-5.6-sol",
|
||||
wantBillingModel: "gpt-image-2",
|
||||
},
|
||||
{
|
||||
name: "requested source overrides image model",
|
||||
resultModel: "gpt-image-2",
|
||||
mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceRequested},
|
||||
requestedModel: "public-image-alias",
|
||||
upstreamModel: "gpt-5.6-sol",
|
||||
wantBillingModel: "public-image-alias",
|
||||
},
|
||||
{
|
||||
name: "mapped channel source overrides image model",
|
||||
resultModel: "gpt-image-2",
|
||||
mapping: service.ChannelMappingResult{MappedModel: "priced-channel-model", BillingModelSource: service.BillingModelSourceChannelMapped},
|
||||
requestedModel: "public-image-alias",
|
||||
upstreamModel: "gpt-5.6-sol",
|
||||
wantBillingModel: "priced-channel-model",
|
||||
},
|
||||
{
|
||||
name: "text turn falls back to upstream model",
|
||||
mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceUpstream},
|
||||
requestedModel: "public-alias",
|
||||
upstreamModel: "gpt-5.6-sol",
|
||||
wantBillingModel: "gpt-5.6-sol",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := &service.OpenAIForwardResult{BillingModel: tt.resultModel}
|
||||
require.Equal(t, tt.wantBillingModel, openAIWSTurnBillingModel(result, tt.mapping, tt.requestedModel, tt.upstreamModel))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldReportOpenAIWSProxyAccountFailure(t *testing.T) {
|
||||
t.Run("unsupported client model switch does not penalize account", func(t *testing.T) {
|
||||
err := fmt.Errorf("wrapped ingress turn: %w", newOpenAIWSUnsupportedModelSwitchError("gpt-unsupported"))
|
||||
require.False(t, shouldReportOpenAIWSProxyAccountFailure(err))
|
||||
|
||||
var closeErr *service.OpenAIWSClientCloseError
|
||||
require.ErrorAs(t, err, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
|
||||
require.Equal(t, "model switch requires reconnect", closeErr.Reason())
|
||||
})
|
||||
|
||||
t.Run("upstream policy violation still penalizes account", func(t *testing.T) {
|
||||
err := service.NewOpenAIWSClientCloseError(
|
||||
coderws.StatusPolicyViolation,
|
||||
"upstream websocket authentication failed",
|
||||
errors.New("upstream rejected credentials"),
|
||||
)
|
||||
require.True(t, shouldReportOpenAIWSProxyAccountFailure(err))
|
||||
})
|
||||
|
||||
t.Run("generic proxy failure still penalizes account", func(t *testing.T) {
|
||||
require.True(t, shouldReportOpenAIWSProxyAccountFailure(errors.New("upstream websocket read failed")))
|
||||
})
|
||||
}
|
||||
|
||||
func TestSetOpenAIClientTransportHTTP(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -1516,14 +1740,22 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
}
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
firstPayload string
|
||||
userAgent *string
|
||||
channelMapping map[string]string
|
||||
firstPayload string
|
||||
secondPayload string
|
||||
userAgent *string
|
||||
ingressMode string
|
||||
channelMapping map[string]string
|
||||
billingModelSource string
|
||||
accountModelMapping map[string]any
|
||||
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
|
||||
}
|
||||
|
||||
type openAIResponsesWSUsageLogResult struct {
|
||||
log *service.UsageLog
|
||||
logs []*service.UsageLog
|
||||
upstreamFirstPayload []byte
|
||||
upstreamPayloads [][]byte
|
||||
clientEvents [][]byte
|
||||
}
|
||||
|
||||
type openAIWSUsageHandlerAccountRepoStub struct {
|
||||
@@ -1633,12 +1865,50 @@ func (s *openAIWSUsageHandlerUsageLogRepoStub) Create(ctx context.Context, log *
|
||||
|
||||
type openAIWSUsageHandlerChannelRepoStub struct {
|
||||
service.ChannelRepository
|
||||
mu sync.Mutex
|
||||
channels []service.Channel
|
||||
groupPlatforms map[int64]string
|
||||
}
|
||||
|
||||
func (s *openAIWSUsageHandlerChannelRepoStub) ListAll(ctx context.Context) ([]service.Channel, error) {
|
||||
return s.channels, nil
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]service.Channel, 0, len(s.channels))
|
||||
for i := range s.channels {
|
||||
out = append(out, *s.channels[i].Clone())
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *openAIWSUsageHandlerChannelRepoStub) GetByID(ctx context.Context, id int64) (*service.Channel, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for i := range s.channels {
|
||||
if s.channels[i].ID == id {
|
||||
return s.channels[i].Clone(), nil
|
||||
}
|
||||
}
|
||||
return nil, service.ErrChannelNotFound
|
||||
}
|
||||
|
||||
func (s *openAIWSUsageHandlerChannelRepoStub) Update(ctx context.Context, channel *service.Channel) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for i := range s.channels {
|
||||
if s.channels[i].ID == channel.ID {
|
||||
s.channels[i] = *channel.Clone()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return service.ErrChannelNotFound
|
||||
}
|
||||
|
||||
func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupIDs(ctx context.Context, channelID int64) ([]int64, error) {
|
||||
channel, err := s.GetByID(ctx, channelID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return append([]int64(nil), channel.GroupIDs...), nil
|
||||
}
|
||||
|
||||
func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupPlatforms(ctx context.Context, groupIDs []int64) (map[int64]string, error) {
|
||||
@@ -2147,8 +2417,13 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstreamPayloadCh := make(chan []byte, 1)
|
||||
turnCount := 1
|
||||
if strings.TrimSpace(tc.secondPayload) != "" {
|
||||
turnCount = 2
|
||||
}
|
||||
upstreamPayloadCh := make(chan []byte, turnCount)
|
||||
upstreamErrCh := make(chan error, 1)
|
||||
var channelSvc *service.ChannelService
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
|
||||
CompressionMode: coderws.CompressionContextTakeover,
|
||||
@@ -2161,29 +2436,39 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
_ = conn.CloseNow()
|
||||
}()
|
||||
|
||||
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
msgType, payload, readErr := conn.Read(readCtx)
|
||||
cancelRead()
|
||||
if readErr != nil {
|
||||
upstreamErrCh <- readErr
|
||||
return
|
||||
}
|
||||
if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
|
||||
upstreamErrCh <- errors.New("unexpected upstream websocket message type")
|
||||
return
|
||||
}
|
||||
upstreamPayloadCh <- payload
|
||||
for turn := 1; turn <= turnCount; turn++ {
|
||||
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
msgType, payload, readErr := conn.Read(readCtx)
|
||||
cancelRead()
|
||||
if readErr != nil {
|
||||
upstreamErrCh <- readErr
|
||||
return
|
||||
}
|
||||
if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
|
||||
upstreamErrCh <- errors.New("unexpected upstream websocket message type")
|
||||
return
|
||||
}
|
||||
upstreamPayloadCh <- payload
|
||||
if turn == 1 && tc.afterFirstUpstreamRequest != nil {
|
||||
if callbackErr := tc.afterFirstUpstreamRequest(channelSvc); callbackErr != nil {
|
||||
upstreamErrCh <- callbackErr
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(
|
||||
`{"type":"response.completed","response":{"id":"resp_usage_e2e","model":"gpt-5.4","usage":{"input_tokens":2,"output_tokens":1}}}`,
|
||||
))
|
||||
cancelWrite()
|
||||
if writeErr != nil {
|
||||
upstreamErrCh <- writeErr
|
||||
return
|
||||
response := fmt.Sprintf(
|
||||
`{"type":"response.completed","response":{"id":"resp_usage_e2e_%d","model":%q,"usage":{"input_tokens":2,"output_tokens":1}}}`,
|
||||
turn,
|
||||
gjson.GetBytes(payload, "model").String(),
|
||||
)
|
||||
writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(response))
|
||||
cancelWrite()
|
||||
if writeErr != nil {
|
||||
upstreamErrCh <- writeErr
|
||||
return
|
||||
}
|
||||
}
|
||||
_ = conn.Close(coderws.StatusNormalClosure, "done")
|
||||
upstreamErrCh <- nil
|
||||
}))
|
||||
defer upstreamServer.Close()
|
||||
@@ -2198,14 +2483,18 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"base_url": upstreamServer.URL,
|
||||
"api_key": "sk-test",
|
||||
"base_url": upstreamServer.URL,
|
||||
"model_mapping": tc.accountModelMapping,
|
||||
},
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
}
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.RunMode = config.RunModeSimple
|
||||
@@ -2221,17 +2510,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
|
||||
|
||||
accountRepo := &openAIWSUsageHandlerAccountRepoStub{account: account}
|
||||
usageRepo := &openAIWSUsageHandlerUsageLogRepoStub{created: make(chan *service.UsageLog, 1)}
|
||||
usageRepo := &openAIWSUsageHandlerUsageLogRepoStub{created: make(chan *service.UsageLog, turnCount)}
|
||||
|
||||
var channelSvc *service.ChannelService
|
||||
if len(tc.channelMapping) > 0 {
|
||||
channelSvc = service.NewChannelService(&openAIWSUsageHandlerChannelRepoStub{
|
||||
channels: []service.Channel{{
|
||||
ID: 7701,
|
||||
Name: "openai-ws-e2e-channel",
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{groupID},
|
||||
ModelMapping: map[string]map[string]string{service.PlatformOpenAI: tc.channelMapping},
|
||||
ID: 7701,
|
||||
Name: "openai-ws-e2e-channel",
|
||||
Status: service.StatusActive,
|
||||
GroupIDs: []int64{groupID},
|
||||
ModelMapping: map[string]map[string]string{service.PlatformOpenAI: tc.channelMapping},
|
||||
BillingModelSource: tc.billingModelSource,
|
||||
}},
|
||||
groupPlatforms: map[int64]string{groupID: service.PlatformOpenAI},
|
||||
}, nil, nil, nil)
|
||||
@@ -2314,26 +2603,44 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, event, err := clientConn.Read(readCtx)
|
||||
cancelRead()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
|
||||
clientEvents := make([][]byte, 0, turnCount)
|
||||
readCompleted := func() {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, event, readErr := clientConn.Read(readCtx)
|
||||
cancelRead()
|
||||
require.NoError(t, readErr)
|
||||
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
|
||||
clientEvents = append(clientEvents, append([]byte(nil), event...))
|
||||
}
|
||||
readCompleted()
|
||||
if turnCount == 2 {
|
||||
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.secondPayload))
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
readCompleted()
|
||||
}
|
||||
_ = clientConn.Close(coderws.StatusNormalClosure, "done")
|
||||
|
||||
var usageLog *service.UsageLog
|
||||
select {
|
||||
case usageLog = <-usageRepo.created:
|
||||
require.NotNil(t, usageLog)
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("等待 WebSocket usage log 写入超时")
|
||||
usageLogs := make([]*service.UsageLog, 0, turnCount)
|
||||
for len(usageLogs) < turnCount {
|
||||
select {
|
||||
case usageLog := <-usageRepo.created:
|
||||
require.NotNil(t, usageLog)
|
||||
usageLogs = append(usageLogs, usageLog)
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("等待 WebSocket usage log 写入超时")
|
||||
}
|
||||
}
|
||||
|
||||
var upstreamFirstPayload []byte
|
||||
select {
|
||||
case upstreamFirstPayload = <-upstreamPayloadCh:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("等待上游 WebSocket 首帧超时")
|
||||
upstreamPayloads := make([][]byte, 0, turnCount)
|
||||
for len(upstreamPayloads) < turnCount {
|
||||
select {
|
||||
case payload := <-upstreamPayloadCh:
|
||||
upstreamPayloads = append(upstreamPayloads, payload)
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("等待上游 WebSocket 请求帧超时")
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
@@ -2344,8 +2651,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
|
||||
return openAIResponsesWSUsageLogResult{
|
||||
log: usageLog,
|
||||
upstreamFirstPayload: upstreamFirstPayload,
|
||||
log: usageLogs[0],
|
||||
logs: usageLogs,
|
||||
upstreamFirstPayload: upstreamPayloads[0],
|
||||
upstreamPayloads: upstreamPayloads,
|
||||
clientEvents: clientEvents,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// openAIPassthroughFailoverState tracks whether this forwarding loop has attempted
|
||||
// an OpenAI passthrough account. Once it has, every subsequent non-passthrough
|
||||
// attempt must use a sanitized body because the immutable canonical body can carry
|
||||
// provider-specific encrypted reasoning produced by that passthrough upstream.
|
||||
type openAIPassthroughFailoverState struct {
|
||||
passthroughSeen bool
|
||||
}
|
||||
|
||||
// deriveOpenAIForwardAttemptBody returns the request body for the upcoming forward
|
||||
// attempt against account. It always derives from the immutable canonical body and
|
||||
// only removes encrypted reasoning input item(s) when the loop has already attempted
|
||||
// a passthrough account and the upcoming account is non-passthrough. This remains
|
||||
// sticky across retries and additional non-passthrough accounts. Attempts before
|
||||
// any passthrough account, and all passthrough attempts, forward the canonical body
|
||||
// unchanged. The canonical slice is never mutated.
|
||||
//
|
||||
// This method is invoked exactly once per forward attempt, immediately before the
|
||||
// Forward call, and advances the failover state as a side effect.
|
||||
func (h *OpenAIGatewayHandler) deriveOpenAIForwardAttemptBody(
|
||||
reqLog *zap.Logger,
|
||||
canonicalBody []byte,
|
||||
account *service.Account,
|
||||
state *openAIPassthroughFailoverState,
|
||||
) []byte {
|
||||
currentPassthrough := account.IsOpenAIPassthroughEnabled()
|
||||
if currentPassthrough {
|
||||
state.passthroughSeen = true
|
||||
return canonicalBody
|
||||
}
|
||||
if !state.passthroughSeen {
|
||||
return canonicalBody
|
||||
}
|
||||
|
||||
sanitized, changed, err := service.SanitizeOpenAICrossModeFailoverReasoning(canonicalBody)
|
||||
if err != nil {
|
||||
if reqLog != nil {
|
||||
reqLog.Warn("openai.failover_cross_mode_reasoning_sanitize_failed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
return canonicalBody
|
||||
}
|
||||
if !changed {
|
||||
return canonicalBody
|
||||
}
|
||||
if reqLog != nil {
|
||||
reqLog.Info("openai.failover_cross_mode_reasoning_stripped",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Bool("account_passthrough", currentPassthrough),
|
||||
zap.Bool("passthrough_seen", true),
|
||||
)
|
||||
}
|
||||
return sanitized
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// kiroReasoningCanonicalBody mirrors a native OpenAI Responses request whose input
|
||||
// carries a provider-specific encrypted reasoning item (Kiro shape: encrypted_content
|
||||
// coupled with an id and summary), plus surrounding message items.
|
||||
const kiroReasoningCanonicalBody = `{"model":"gpt-5.1","stream":false,"input":[` +
|
||||
`{"type":"message","role":"user","content":"hello"},` +
|
||||
`{"type":"reasoning","id":"rs_kiro_abc123","encrypted_content":"ENC_BLOB","summary":[{"type":"summary_text","text":"thinking"}]},` +
|
||||
`{"type":"message","role":"assistant","content":"hi"}` +
|
||||
`]}`
|
||||
|
||||
func newOpenAIPassthroughAccount(id int64, passthrough bool) *service.Account {
|
||||
extra := map[string]any{"openai_passthrough": passthrough}
|
||||
return &service.Account{
|
||||
ID: id,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Extra: extra,
|
||||
}
|
||||
}
|
||||
|
||||
func reasoningItemCount(t *testing.T, body []byte) int {
|
||||
t.Helper()
|
||||
count := 0
|
||||
gjson.GetBytes(body, "input").ForEach(func(_, item gjson.Result) bool {
|
||||
if item.Get("type").String() == "reasoning" {
|
||||
count++
|
||||
}
|
||||
return true
|
||||
})
|
||||
return count
|
||||
}
|
||||
|
||||
// TestDeriveOpenAIForwardAttemptBody_CrossModeStripsKiroReasoning drives the real
|
||||
// account-switch/request-body seam used inside OpenAIGatewayHandler.Responses: the
|
||||
// same deriveOpenAIForwardAttemptBody call the failover loop makes, fed real
|
||||
// *service.Account values in sequence. It proves the first passthrough attempt gets
|
||||
// the untouched reasoning item, and after a failover to a non-passthrough account
|
||||
// the second attempt body has the entire reasoning item removed — while the
|
||||
// immutable canonical body is never mutated.
|
||||
func TestDeriveOpenAIForwardAttemptBody_CrossModeStripsKiroReasoning(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
canonical := []byte(kiroReasoningCanonicalBody)
|
||||
|
||||
kiro := newOpenAIPassthroughAccount(1, true) // Kiro: openai_passthrough=true
|
||||
bedrock := newOpenAIPassthroughAccount(2, false) // Bedrock Mantle: passthrough disabled
|
||||
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
|
||||
// Attempt 1 — passthrough (Kiro): original reasoning item must survive intact.
|
||||
firstBody := h.deriveOpenAIForwardAttemptBody(nil, canonical, kiro, state)
|
||||
require.Equal(t, 1, reasoningItemCount(t, firstBody), "first passthrough attempt must keep the reasoning item")
|
||||
require.Equal(t, "ENC_BLOB", gjson.GetBytes(firstBody, "input.1.encrypted_content").String())
|
||||
require.Equal(t, "rs_kiro_abc123", gjson.GetBytes(firstBody, "input.1.id").String())
|
||||
require.JSONEq(t, kiroReasoningCanonicalBody, string(firstBody), "first attempt body must equal the canonical body")
|
||||
|
||||
// Attempt 2 — failover switches to non-passthrough (Bedrock): the whole
|
||||
// provider-specific encrypted reasoning item (id/summary/encrypted_content) is dropped.
|
||||
secondBody := h.deriveOpenAIForwardAttemptBody(nil, canonical, bedrock, state)
|
||||
require.Equal(t, 0, reasoningItemCount(t, secondBody), "cross-mode attempt must drop the reasoning item entirely")
|
||||
require.False(t, gjson.GetBytes(secondBody, "input.#(encrypted_content)").Exists(), "no encrypted_content may remain")
|
||||
require.NotContains(t, string(secondBody), "rs_kiro_abc123", "coupled reasoning id must be gone")
|
||||
require.NotContains(t, string(secondBody), "ENC_BLOB")
|
||||
|
||||
// Surviving message items are preserved.
|
||||
require.Equal(t, 2, int(gjson.GetBytes(secondBody, "input.#").Int()))
|
||||
require.Equal(t, "hello", gjson.GetBytes(secondBody, "input.0.content").String())
|
||||
require.Equal(t, "hi", gjson.GetBytes(secondBody, "input.1.content").String())
|
||||
|
||||
// The canonical body is immutable across attempts.
|
||||
require.JSONEq(t, kiroReasoningCanonicalBody, string(canonical), "canonical forwardBody must never be mutated")
|
||||
}
|
||||
|
||||
// TestDeriveOpenAIForwardAttemptBody_SameModePreservesReasoning proves that
|
||||
// non-passthrough attempts before any passthrough attempt, passthrough-family
|
||||
// failovers, and switching into passthrough forward the canonical reasoning item
|
||||
// unchanged.
|
||||
func TestDeriveOpenAIForwardAttemptBody_SameModePreservesReasoning(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
canonical := []byte(kiroReasoningCanonicalBody)
|
||||
|
||||
t.Run("non_passthrough_to_non_passthrough", func(t *testing.T) {
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
first := h.deriveOpenAIForwardAttemptBody(nil, canonical, newOpenAIPassthroughAccount(10, false), state)
|
||||
second := h.deriveOpenAIForwardAttemptBody(nil, canonical, newOpenAIPassthroughAccount(11, false), state)
|
||||
require.Equal(t, 1, reasoningItemCount(t, first))
|
||||
require.Equal(t, 1, reasoningItemCount(t, second), "Bedrock-family failover must preserve reasoning")
|
||||
require.JSONEq(t, kiroReasoningCanonicalBody, string(second))
|
||||
})
|
||||
|
||||
t.Run("passthrough_to_passthrough", func(t *testing.T) {
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
first := h.deriveOpenAIForwardAttemptBody(nil, canonical, newOpenAIPassthroughAccount(20, true), state)
|
||||
second := h.deriveOpenAIForwardAttemptBody(nil, canonical, newOpenAIPassthroughAccount(21, true), state)
|
||||
require.Equal(t, 1, reasoningItemCount(t, first))
|
||||
require.Equal(t, 1, reasoningItemCount(t, second), "Kiro-family failover must preserve reasoning")
|
||||
})
|
||||
|
||||
t.Run("non_passthrough_to_passthrough", func(t *testing.T) {
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
_ = h.deriveOpenAIForwardAttemptBody(nil, canonical, newOpenAIPassthroughAccount(30, false), state)
|
||||
second := h.deriveOpenAIForwardAttemptBody(nil, canonical, newOpenAIPassthroughAccount(31, true), state)
|
||||
require.Equal(t, 1, reasoningItemCount(t, second), "switching into passthrough must preserve reasoning")
|
||||
})
|
||||
|
||||
t.Run("same_account_pool_retry", func(t *testing.T) {
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
kiro := newOpenAIPassthroughAccount(40, true)
|
||||
_ = h.deriveOpenAIForwardAttemptBody(nil, canonical, kiro, state)
|
||||
retry := h.deriveOpenAIForwardAttemptBody(nil, canonical, kiro, state)
|
||||
require.Equal(t, 1, reasoningItemCount(t, retry), "same-account retry must preserve reasoning")
|
||||
})
|
||||
}
|
||||
|
||||
// TestDeriveOpenAIForwardAttemptBody_SanitizationSticksAcrossBedrockRetries proves
|
||||
// that once a passthrough account has been attempted, all later non-passthrough
|
||||
// retries and failovers remain sanitized instead of restoring the canonical Kiro
|
||||
// reasoning item after the first Bedrock failure.
|
||||
func TestDeriveOpenAIForwardAttemptBody_SanitizationSticksAcrossBedrockRetries(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
canonical := []byte(kiroReasoningCanonicalBody)
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
|
||||
kiro := newOpenAIPassthroughAccount(50, true)
|
||||
bedrockA := newOpenAIPassthroughAccount(51, false)
|
||||
bedrockB := newOpenAIPassthroughAccount(52, false)
|
||||
|
||||
first := h.deriveOpenAIForwardAttemptBody(nil, canonical, kiro, state)
|
||||
second := h.deriveOpenAIForwardAttemptBody(nil, canonical, bedrockA, state)
|
||||
retry := h.deriveOpenAIForwardAttemptBody(nil, canonical, bedrockA, state)
|
||||
nextAccount := h.deriveOpenAIForwardAttemptBody(nil, canonical, bedrockB, state)
|
||||
|
||||
require.Equal(t, 1, reasoningItemCount(t, first))
|
||||
require.Equal(t, 0, reasoningItemCount(t, second))
|
||||
require.Equal(t, 0, reasoningItemCount(t, retry), "same Bedrock account retry must remain sanitized")
|
||||
require.Equal(t, 0, reasoningItemCount(t, nextAccount), "later Bedrock account must remain sanitized")
|
||||
require.JSONEq(t, kiroReasoningCanonicalBody, string(canonical), "canonical forwardBody must never be mutated")
|
||||
}
|
||||
@@ -91,6 +91,8 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
zap.Bool("stream", parsed.Stream),
|
||||
zap.Bool("multipart", parsed.Multipart),
|
||||
zap.String("capability", string(parsed.RequiredCapability)),
|
||||
zap.String("img_quality", parsed.Quality),
|
||||
zap.String("img_size", parsed.Size),
|
||||
)
|
||||
|
||||
if !service.GroupAllowsImageGeneration(apiKey.Group) {
|
||||
@@ -374,6 +376,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
if result != nil {
|
||||
upstreamModel = result.UpstreamModel
|
||||
}
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
Result: result,
|
||||
@@ -388,6 +391,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, requestModel, upstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -12,11 +12,14 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
"go.uber.org/zap"
|
||||
"go.uber.org/zap/zaptest/observer"
|
||||
)
|
||||
|
||||
type openAIImagesFailoverAccountRepo struct {
|
||||
@@ -154,8 +157,10 @@ func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhen
|
||||
)
|
||||
handler.maxAccountSwitches = 10
|
||||
|
||||
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
||||
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","quality":"high","size":"1536x1024"}`)
|
||||
core, observedLogs := observer.New(zap.DebugLevel)
|
||||
requestCtx := logger.IntoContext(context.Background(), zap.New(core))
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)).WithContext(requestCtx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
@@ -173,6 +178,16 @@ func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhen
|
||||
|
||||
handler.Images(c)
|
||||
|
||||
accountSelectingLogs := observedLogs.FilterMessage("openai.images.account_selecting").All()
|
||||
require.NotEmpty(t, accountSelectingLogs)
|
||||
loggedFields := make(map[string]string)
|
||||
for _, field := range accountSelectingLogs[0].Context {
|
||||
loggedFields[field.Key] = field.String
|
||||
}
|
||||
require.Equal(t, "high", loggedFields["img_quality"])
|
||||
require.Equal(t, "1536x1024", loggedFields["img_size"])
|
||||
require.NotContains(t, loggedFields, "prompt")
|
||||
|
||||
require.Equal(t, []int64{1, 2}, upstream.calls())
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||||
require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func (h *OpenAIGatewayHandler) Live(c *gin.Context) {
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||||
return
|
||||
}
|
||||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||||
return
|
||||
}
|
||||
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformOpenAI {
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Live is not supported for this platform")
|
||||
return
|
||||
}
|
||||
if !liveEnabledForAPIKey(apiKey) {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", "Live is not enabled for this group")
|
||||
return
|
||||
}
|
||||
request, err := parseLiveCallRequest(c)
|
||||
if err != nil {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
model := strings.TrimSpace(gjson.GetBytes(request.Session, "model").String())
|
||||
reqLog := requestLogger(
|
||||
c,
|
||||
"handler.openai_gateway.live",
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
if decision := h.checkSecurityAudit(
|
||||
c,
|
||||
reqLog,
|
||||
apiKey,
|
||||
subject,
|
||||
service.ContentModerationProtocolOpenAIResponses,
|
||||
model,
|
||||
request.Session,
|
||||
); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||||
if h.billingCacheService == nil {
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Billing service unavailable")
|
||||
return
|
||||
}
|
||||
if err := h.billingCacheService.CheckBillingEligibility(
|
||||
c.Request.Context(),
|
||||
apiKey.User,
|
||||
apiKey,
|
||||
apiKey.Group,
|
||||
subscription,
|
||||
service.QuotaPlatform(c.Request.Context(), apiKey),
|
||||
); err != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
|
||||
userRelease, acquired, err := h.concurrencyHelper.TryAcquireUserSlot(
|
||||
c.Request.Context(),
|
||||
subject.UserID,
|
||||
subject.Concurrency,
|
||||
)
|
||||
if err != nil {
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Live concurrency unavailable")
|
||||
return
|
||||
}
|
||||
if !acquired {
|
||||
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Live concurrency limit reached")
|
||||
return
|
||||
}
|
||||
defer userRelease()
|
||||
|
||||
identity := liveCallIdentity(c, apiKey, subject.UserID, subscription)
|
||||
created, err := h.gatewayService.CreateLiveCall(c.Request.Context(), request, identity, subject.Concurrency)
|
||||
if err != nil {
|
||||
h.writeLiveCreateError(c, err)
|
||||
return
|
||||
}
|
||||
c.Header("Location", liveSidebandLocation(c.FullPath(), created.CallID))
|
||||
c.Data(http.StatusOK, "application/sdp", created.SDP)
|
||||
}
|
||||
|
||||
func parseLiveCallRequest(c *gin.Context) (*service.LiveCallRequest, error) {
|
||||
contentType := strings.ToLower(c.GetHeader("Content-Type"))
|
||||
if strings.HasPrefix(contentType, "multipart/form-data") {
|
||||
sdp := c.PostForm("sdp")
|
||||
session := json.RawMessage(c.PostForm("session"))
|
||||
request := &service.LiveCallRequest{SDP: sdp, Session: session}
|
||||
if err := service.ValidateLiveCallRequest(request); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
var request service.LiveCallRequest
|
||||
decoder := json.NewDecoder(c.Request.Body)
|
||||
if err := decoder.Decode(&request); err != nil {
|
||||
return nil, errors.New("request body must be valid JSON")
|
||||
}
|
||||
if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
|
||||
return nil, errors.New("request body must contain one JSON object")
|
||||
}
|
||||
if err := service.ValidateLiveCallRequest(&request); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &request, nil
|
||||
}
|
||||
|
||||
func liveSidebandLocation(fullPath, callID string) string {
|
||||
prefix := "/v1/live/"
|
||||
if strings.HasPrefix(fullPath, "/backend-api/codex/") {
|
||||
prefix = "/backend-api/codex/"
|
||||
}
|
||||
return prefix + url.PathEscape(callID)
|
||||
}
|
||||
|
||||
func liveCallIdentity(
|
||||
c *gin.Context,
|
||||
apiKey *service.APIKey,
|
||||
userID int64,
|
||||
subscription *service.UserSubscription,
|
||||
) service.LiveCallIdentity {
|
||||
var subscriptionID *int64
|
||||
if subscription != nil {
|
||||
value := subscription.ID
|
||||
subscriptionID = &value
|
||||
}
|
||||
return service.LiveCallIdentity{
|
||||
APIKeyID: apiKey.ID,
|
||||
UserID: userID,
|
||||
GroupID: apiKey.GroupID,
|
||||
SubscriptionID: subscriptionID,
|
||||
UserAgent: c.GetHeader("User-Agent"),
|
||||
IPAddress: ip.GetClientIP(c),
|
||||
InboundEndpoint: GetInboundEndpoint(c),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) writeLiveCreateError(c *gin.Context, err error) {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrLiveConcurrencyFull):
|
||||
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Live concurrency limit reached")
|
||||
case errors.Is(err, service.ErrLiveUnavailable):
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Live is unavailable")
|
||||
default:
|
||||
var attestationErr *service.LiveAttestationUnavailableError
|
||||
if errors.As(err, &attestationErr) {
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", attestationErr.Error())
|
||||
return
|
||||
}
|
||||
var upstreamErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &upstreamErr) && upstreamErr.StatusCode >= 400 && upstreamErr.StatusCode < 500 {
|
||||
h.errorResponse(c, upstreamErr.StatusCode, "invalid_request_error", "Live upstream rejected the request")
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusBadGateway, "api_error", "Live upstream request failed")
|
||||
}
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) LiveSideband(c *gin.Context) {
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||||
return
|
||||
}
|
||||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||||
return
|
||||
}
|
||||
if !liveEnabledForAPIKey(apiKey) {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", "Live is not enabled for this group")
|
||||
return
|
||||
}
|
||||
identity := service.LiveCallIdentity{
|
||||
APIKeyID: apiKey.ID,
|
||||
UserID: subject.UserID,
|
||||
GroupID: apiKey.GroupID,
|
||||
}
|
||||
record, err := h.gatewayService.GetLiveCallForIdentity(c.Request.Context(), c.Param("call_id"), identity)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrLiveIdentityMismatch) {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", "Live call belongs to another identity")
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Live call not found")
|
||||
return
|
||||
}
|
||||
downstream, err := coderws.Accept(c.Writer, c.Request, &coderws.AcceptOptions{
|
||||
InsecureSkipVerify: true,
|
||||
})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = downstream.CloseNow() }()
|
||||
if err := h.gatewayService.ProxyLiveSideband(c.Request.Context(), record, downstream); err != nil {
|
||||
_ = downstream.Close(coderws.StatusInternalError, "live sideband closed")
|
||||
return
|
||||
}
|
||||
_ = downstream.Close(coderws.StatusNormalClosure, "")
|
||||
}
|
||||
|
||||
func liveEnabledForAPIKey(apiKey *service.APIKey) bool {
|
||||
return apiKey != nil &&
|
||||
apiKey.Group != nil &&
|
||||
apiKey.Group.Platform == service.PlatformOpenAI &&
|
||||
apiKey.Group.AllowLive
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestParseLiveCallRequestMultipartPreservesSession(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
session := `{"model":"gpt-live-test","delegation":{"type":"client"},"instructions":"你好"}`
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
require.NoError(t, writer.WriteField("sdp", "v=0\r\n"))
|
||||
require.NoError(t, writer.WriteField("session", session))
|
||||
require.NoError(t, writer.Close())
|
||||
|
||||
request := httptest.NewRequest("POST", "/v1/live", &body)
|
||||
request.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = request
|
||||
|
||||
parsed, err := parseLiveCallRequest(context)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "v=0\r\n", parsed.SDP)
|
||||
require.JSONEq(t, session, string(parsed.Session))
|
||||
require.Equal(t, "client", jsonPathString(t, parsed.Session, "delegation", "type"))
|
||||
}
|
||||
|
||||
func TestParseLiveCallRequestJSONPreservesSessionWithoutDelegation(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := `{"sdp":"v=0\\r\\n","session":{"model":"gpt-live-test","instructions":"standalone"}}`
|
||||
request := httptest.NewRequest("POST", "/backend-api/codex/realtime/calls", bytes.NewBufferString(body))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = request
|
||||
|
||||
parsed, err := parseLiveCallRequest(context)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(parsed.Session), "delegation")
|
||||
require.Equal(t, "standalone", jsonPathString(t, parsed.Session, "instructions"))
|
||||
}
|
||||
|
||||
func TestParseLiveCallRequestRejectsInvalidJSONShape(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
testCases := []string{
|
||||
`{"session":{"type":"quicksilver"}}`,
|
||||
`{"sdp":"v=0\\r\\n","session":[]}`,
|
||||
`{"sdp":"v=0\\r\\n","session":null}`,
|
||||
`{"sdp":"v=0\\r\\n","session":{"type":"quicksilver"}} {}`,
|
||||
}
|
||||
for _, body := range testCases {
|
||||
request := httptest.NewRequest("POST", "/backend-api/codex/realtime/calls", bytes.NewBufferString(body))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = request
|
||||
_, err := parseLiveCallRequest(context)
|
||||
require.Error(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLiveSidebandLocationMatchesCreateRoute(t *testing.T) {
|
||||
require.Equal(t, "/v1/live/call_123", liveSidebandLocation("/v1/live", "call_123"))
|
||||
require.Equal(
|
||||
t,
|
||||
"/backend-api/codex/call_123",
|
||||
liveSidebandLocation("/backend-api/codex/realtime/calls", "call_123"),
|
||||
)
|
||||
}
|
||||
|
||||
func TestLiveEnabledForAPIKey(t *testing.T) {
|
||||
require.False(t, liveEnabledForAPIKey(nil))
|
||||
require.False(t, liveEnabledForAPIKey(&service.APIKey{}))
|
||||
require.False(t, liveEnabledForAPIKey(&service.APIKey{
|
||||
Group: &service.Group{Platform: service.PlatformOpenAI},
|
||||
}))
|
||||
require.False(t, liveEnabledForAPIKey(&service.APIKey{
|
||||
Group: &service.Group{Platform: service.PlatformAnthropic, AllowLive: true},
|
||||
}))
|
||||
require.True(t, liveEnabledForAPIKey(&service.APIKey{
|
||||
Group: &service.Group{Platform: service.PlatformOpenAI, AllowLive: true},
|
||||
}))
|
||||
}
|
||||
|
||||
func TestLiveAttestationErrorIsExplicit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
(&OpenAIGatewayHandler{}).writeLiveCreateError(context, &service.LiveAttestationUnavailableError{
|
||||
Reason: "Live attestation is only supported when Sub2API runs on macOS",
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "Sub2API runs on macOS")
|
||||
}
|
||||
|
||||
func jsonPathString(t *testing.T, raw json.RawMessage, keys ...string) string {
|
||||
t.Helper()
|
||||
var value any
|
||||
require.NoError(t, json.Unmarshal(raw, &value))
|
||||
current := value
|
||||
for _, key := range keys {
|
||||
object, ok := current.(map[string]any)
|
||||
require.True(t, ok)
|
||||
current = object[key]
|
||||
}
|
||||
result, ok := current.(string)
|
||||
require.True(t, ok)
|
||||
return result
|
||||
}
|
||||
@@ -26,6 +26,9 @@ type userHandlerRepoStub struct {
|
||||
}
|
||||
|
||||
func (s *userHandlerRepoStub) Create(context.Context, *service.User) error { return nil }
|
||||
func (s *userHandlerRepoStub) CreateWithEmailAliasGuard(context.Context, *service.User) error {
|
||||
return nil
|
||||
}
|
||||
func (s *userHandlerRepoStub) GetByID(context.Context, int64) (*service.User, error) {
|
||||
cloned := *s.user
|
||||
return &cloned, nil
|
||||
@@ -97,6 +100,9 @@ func (s *userHandlerRepoStub) BatchUpdateLimits(context.Context, []int64, *int,
|
||||
return 0, nil
|
||||
}
|
||||
func (s *userHandlerRepoStub) ExistsByEmail(context.Context, string) (bool, error) { return false, nil }
|
||||
func (s *userHandlerRepoStub) ExistsByEmailAlias(context.Context, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
func (s *userHandlerRepoStub) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
@@ -8,6 +8,8 @@ import (
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
ippkg "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
@@ -72,6 +74,52 @@ func NewRateLimiter(redisClient *redis.Client) *RateLimiter {
|
||||
}
|
||||
}
|
||||
|
||||
// AllowResult 单次固定窗口限流判定结果。
|
||||
type AllowResult struct {
|
||||
// Allowed 是否放行
|
||||
Allowed bool
|
||||
// Count 当前窗口内累计请求数(含本次)
|
||||
Count int64
|
||||
// RetryAfter 超限时距窗口重置的剩余时间(尽力而为;PTTL 不可用时回退为完整窗口)
|
||||
RetryAfter time.Duration
|
||||
}
|
||||
|
||||
// Allow 对给定 key(不含 "rate_limit:" 前缀)执行一次固定窗口计数判定。
|
||||
// 供需要自定义限流维度(如按用户 ID)的调用方使用;Redis 错误由调用方决定 fail-open/close。
|
||||
func (r *RateLimiter) Allow(ctx context.Context, key string, limit int, window time.Duration) (AllowResult, error) {
|
||||
redisKey := r.prefix + key
|
||||
windowMillis := windowTTLMillis(window)
|
||||
|
||||
count, repaired, err := rateLimitRun(ctx, r.redis, redisKey, windowMillis)
|
||||
if err != nil {
|
||||
return AllowResult{}, err
|
||||
}
|
||||
if repaired {
|
||||
log.Printf("[RateLimit] ttl repaired: key=%s window_ms=%d", redisKey, windowMillis)
|
||||
}
|
||||
|
||||
result := AllowResult{Allowed: count <= int64(limit), Count: count}
|
||||
if !result.Allowed {
|
||||
result.RetryAfter = window
|
||||
if ttl, ttlErr := r.redis.PTTL(ctx, redisKey).Result(); ttlErr == nil && ttl > 0 {
|
||||
result.RetryAfter = ttl
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// clientIPForRateLimit 返回 IP 维度限流使用的客户端地址。
|
||||
// 与审计日志/会话绑定/API Key IP ACL 共用同一套安全客户端 IP 解析
|
||||
// (SessionBindingContext 快照:兼容开关开启时信任反代转发头,关闭时走
|
||||
// server.trusted_proxies 可信链)。避免默认反代部署下 Gin ClientIP 恒等于
|
||||
// 代理地址、所有用户坍缩进同一个限流桶造成整体误拦截。
|
||||
func clientIPForRateLimit(c *gin.Context) string {
|
||||
if resolved := ippkg.GetSecurityClientIP(c, false); resolved != "" {
|
||||
return resolved
|
||||
}
|
||||
return c.ClientIP()
|
||||
}
|
||||
|
||||
// Limit 返回速率限制中间件
|
||||
// key: 限制类型标识
|
||||
// limit: 时间窗口内最大请求数
|
||||
@@ -88,32 +136,21 @@ func (r *RateLimiter) LimitWithOptions(key string, limit int, window time.Durati
|
||||
}
|
||||
|
||||
return func(c *gin.Context) {
|
||||
ip := c.ClientIP()
|
||||
redisKey := r.prefix + key + ":" + ip
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
windowMillis := windowTTLMillis(window)
|
||||
|
||||
// 使用 Lua 脚本原子操作增加计数并设置过期
|
||||
count, repaired, err := rateLimitRun(ctx, r.redis, redisKey, windowMillis)
|
||||
result, err := r.Allow(c.Request.Context(), key+":"+clientIPForRateLimit(c), limit, window)
|
||||
if err != nil {
|
||||
log.Printf("[RateLimit] redis error: key=%s mode=%s err=%v", redisKey, failureModeLabel(failureMode), err)
|
||||
log.Printf("[RateLimit] redis error: key=%s mode=%s err=%v", r.prefix+key, failureModeLabel(failureMode), err)
|
||||
if failureMode == RateLimitFailClose {
|
||||
abortRateLimit(c)
|
||||
abortRateLimit(c, window)
|
||||
return
|
||||
}
|
||||
// Redis 错误时放行,避免影响正常服务
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if repaired {
|
||||
log.Printf("[RateLimit] ttl repaired: key=%s window_ms=%d", redisKey, windowMillis)
|
||||
}
|
||||
|
||||
// 超过限制
|
||||
if count > int64(limit) {
|
||||
abortRateLimit(c)
|
||||
if !result.Allowed {
|
||||
abortRateLimit(c, result.RetryAfter)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -129,7 +166,14 @@ func windowTTLMillis(window time.Duration) int64 {
|
||||
return ttl
|
||||
}
|
||||
|
||||
func abortRateLimit(c *gin.Context) {
|
||||
func abortRateLimit(c *gin.Context, retryAfter time.Duration) {
|
||||
if retryAfter > 0 {
|
||||
seconds := int64(retryAfter / time.Second)
|
||||
if retryAfter%time.Second > 0 {
|
||||
seconds++
|
||||
}
|
||||
c.Header("Retry-After", strconv.FormatInt(seconds, 10))
|
||||
}
|
||||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
||||
"error": "rate limit exceeded",
|
||||
"message": "Too many requests, please try again later",
|
||||
|
||||
@@ -7,6 +7,8 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
ippkg "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -103,6 +105,87 @@ func TestRateLimiterDifferentIPsIndependent(t *testing.T) {
|
||||
require.Equal(t, http.StatusTooManyRequests, rec3.Code, "第一个 IP 的第二次请求应被限流")
|
||||
}
|
||||
|
||||
func TestRateLimiterAllow(t *testing.T) {
|
||||
originalRun := rateLimitRun
|
||||
var gotKey string
|
||||
count := int64(0)
|
||||
rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
|
||||
gotKey = key
|
||||
count++
|
||||
return count, false, nil
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
rateLimitRun = originalRun
|
||||
})
|
||||
|
||||
// PTTL 走真实客户端(不可达地址)→ 失败后 RetryAfter 应回退为完整窗口
|
||||
limiter := NewRateLimiter(redis.NewClient(&redis.Options{
|
||||
Addr: "127.0.0.1:1",
|
||||
DialTimeout: 50 * time.Millisecond,
|
||||
ReadTimeout: 50 * time.Millisecond,
|
||||
WriteTimeout: 50 * time.Millisecond,
|
||||
}))
|
||||
|
||||
res, err := limiter.Allow(context.Background(), "panel:global:user:42", 1, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.True(t, res.Allowed)
|
||||
require.Equal(t, int64(1), res.Count)
|
||||
require.Zero(t, res.RetryAfter)
|
||||
require.Equal(t, "rate_limit:panel:global:user:42", gotKey)
|
||||
|
||||
res, err = limiter.Allow(context.Background(), "panel:global:user:42", 1, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.False(t, res.Allowed)
|
||||
require.Equal(t, int64(2), res.Count)
|
||||
require.Equal(t, time.Minute, res.RetryAfter)
|
||||
}
|
||||
|
||||
func TestRateLimiterHonorsForwardedIPSnapshot(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
callCounts := make(map[string]int64)
|
||||
originalRun := rateLimitRun
|
||||
rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
|
||||
callCounts[key]++
|
||||
return callCounts[key], false, nil
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
rateLimitRun = originalRun
|
||||
})
|
||||
|
||||
limiter := NewRateLimiter(redis.NewClient(&redis.Options{Addr: "127.0.0.1:1"}))
|
||||
|
||||
router := gin.New()
|
||||
// 模拟 SessionBindingContext:开启转发 IP 兼容模式快照
|
||||
router.Use(func(c *gin.Context) {
|
||||
ippkg.SetForwardedIPSettings(c, true, nil)
|
||||
c.Next()
|
||||
})
|
||||
router.Use(limiter.Limit("fwd", 1, time.Second))
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
send := func(xff string) int {
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
// 所有请求都来自同一个反代地址
|
||||
req.RemoteAddr = "127.0.0.1:5678"
|
||||
req.Header.Set("X-Forwarded-For", xff)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec.Code
|
||||
}
|
||||
|
||||
// 反代后两个不同的真实客户端应各自独立计数,不因共享代理地址被合并限流
|
||||
require.Equal(t, http.StatusOK, send("198.51.100.1"))
|
||||
require.Equal(t, http.StatusOK, send("198.51.100.2"))
|
||||
// 同一真实客户端第二次请求应被限流
|
||||
require.Equal(t, http.StatusTooManyRequests, send("198.51.100.1"))
|
||||
|
||||
require.Contains(t, callCounts, "rate_limit:fwd:198.51.100.1")
|
||||
require.Contains(t, callCounts, "rate_limit:fwd:198.51.100.2")
|
||||
}
|
||||
|
||||
func TestRateLimiterSuccessAndLimit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -586,7 +586,7 @@ func (r *ResponsesClientToolStreamRestorer) restoreNamespaceEvent(event Response
|
||||
event.Item.Name, event.Item.Namespace = name.Name, name.Namespace
|
||||
}
|
||||
}
|
||||
if event.Type == "response.function_call_arguments.done" {
|
||||
if event.Type == "response.function_call_arguments.delta" || event.Type == "response.function_call_arguments.done" {
|
||||
if name, ok := r.adapter.NamespaceTools[event.Name]; ok {
|
||||
event.Name = name.Name
|
||||
}
|
||||
|
||||
@@ -146,7 +146,13 @@ func TestResponsesClientToolStreamRestorer_RestoresNamespaceLifecycle(t *testing
|
||||
require.Equal(t, "open", gjson.GetBytes(added[0], "item.name").String())
|
||||
require.Equal(t, "browser", gjson.GetBytes(added[0], "item.namespace").String())
|
||||
|
||||
done, changed, err := restorer.RestoreEvent([]byte(`{"type":"response.function_call_arguments.done","sequence_number":5,"output_index":0,"item_id":"i1","name":"browser__open","arguments":"{}"}`))
|
||||
delta, changed, err := restorer.RestoreEvent([]byte(`{"type":"response.function_call_arguments.delta","sequence_number":5,"output_index":0,"item_id":"i1","name":"browser__open","delta":"{\"url\":"}`))
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Len(t, delta, 1)
|
||||
require.Equal(t, "open", gjson.GetBytes(delta[0], "name").String())
|
||||
|
||||
done, changed, err := restorer.RestoreEvent([]byte(`{"type":"response.function_call_arguments.done","sequence_number":6,"output_index":0,"item_id":"i1","name":"browser__open","arguments":"{}"}`))
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Len(t, done, 1)
|
||||
|
||||
@@ -152,11 +152,7 @@ func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMess
|
||||
|
||||
case item.Type == "function_call_output":
|
||||
// function_call_output → user message with tool_result block
|
||||
outputContent := item.Output
|
||||
if outputContent == "" {
|
||||
outputContent = "(empty)"
|
||||
}
|
||||
contentJSON, _ := json.Marshal(outputContent)
|
||||
contentJSON := responsesFunctionOutputToAnthropicContent(item)
|
||||
block := AnthropicContentBlock{
|
||||
Type: "tool_result",
|
||||
ToolUseID: fromResponsesCallIDToAnthropic(item.CallID),
|
||||
@@ -216,6 +212,45 @@ func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMess
|
||||
return system, messages, nil
|
||||
}
|
||||
|
||||
func responsesFunctionOutputToAnthropicContent(item ResponsesInputItem) json.RawMessage {
|
||||
if len(item.outputRaw) == 0 {
|
||||
output := item.Output
|
||||
if output == "" {
|
||||
output = "(empty)"
|
||||
}
|
||||
content, _ := json.Marshal(output)
|
||||
return content
|
||||
}
|
||||
|
||||
var parts []ResponsesContentPart
|
||||
if err := json.Unmarshal(item.outputRaw, &parts); err == nil {
|
||||
blocks := make([]AnthropicContentBlock, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
switch part.Type {
|
||||
case "input_text", "output_text", "text":
|
||||
if part.Text != "" {
|
||||
blocks = append(blocks, AnthropicContentBlock{Type: "text", Text: part.Text})
|
||||
}
|
||||
case "input_image":
|
||||
if source := dataURIToAnthropicImageSource(part.ImageURL); source != nil {
|
||||
blocks = append(blocks, AnthropicContentBlock{Type: "image", Source: source})
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(blocks) > 0 {
|
||||
content, _ := json.Marshal(blocks)
|
||||
return content
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
content, _ := json.Marshal("(empty)")
|
||||
return content
|
||||
}
|
||||
}
|
||||
|
||||
content, _ := json.Marshal(item.Output)
|
||||
return content
|
||||
}
|
||||
|
||||
// normalizeAnthropicToolPairing rebuilds the message sequence so it satisfies
|
||||
// Anthropic's tool_use/tool_result invariants, which the naive item-by-item
|
||||
// conversion violates whenever the Responses history interleaves anything
|
||||
|
||||
@@ -163,3 +163,28 @@ func TestAnthropicPairing_SingleCall(t *testing.T) {
|
||||
require.True(t, hasToolUse(parseContentBlocks(msgs[1].Content), "call_A"))
|
||||
require.True(t, hasToolResult(parseContentBlocks(msgs[2].Content), "call_A"))
|
||||
}
|
||||
|
||||
func TestResponsesToAnthropic_FunctionOutputContentArray(t *testing.T) {
|
||||
msgs := convertAnthropic(t, `[
|
||||
{"type":"function_call","call_id":"call_A","name":"view_image","arguments":"{}"},
|
||||
{"type":"function_call_output","call_id":"call_A","output":[
|
||||
{"type":"input_text","text":"image loaded"},
|
||||
{"type":"input_image","image_url":"data:image/png;base64,YQ=="}
|
||||
]}
|
||||
]`)
|
||||
|
||||
require.Len(t, msgs, 2)
|
||||
resultBlocks := parseContentBlocks(msgs[1].Content)
|
||||
require.Len(t, resultBlocks, 1)
|
||||
require.Equal(t, "tool_result", resultBlocks[0].Type)
|
||||
|
||||
var content []AnthropicContentBlock
|
||||
require.NoError(t, json.Unmarshal(resultBlocks[0].Content, &content))
|
||||
require.Len(t, content, 2)
|
||||
require.Equal(t, "text", content[0].Type)
|
||||
require.Equal(t, "image loaded", content[0].Text)
|
||||
require.Equal(t, "image", content[1].Type)
|
||||
require.NotNil(t, content[1].Source)
|
||||
require.Equal(t, "image/png", content[1].Source.MediaType)
|
||||
require.Equal(t, "YQ==", content[1].Source.Data)
|
||||
}
|
||||
|
||||
@@ -122,6 +122,9 @@ func TestResponsesToAnthropic_MixedToolsProduceValidAnthropicTools(t *testing.T)
|
||||
assert.Equal(t, "web_search_20250305", tools[2].Type)
|
||||
assert.Equal(t, "web_search", tools[2].Name)
|
||||
assert.Empty(t, tools[2].InputSchema)
|
||||
serverToolWire, err := json.Marshal(tools[2])
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(serverToolWire), `"input_schema"`)
|
||||
}
|
||||
|
||||
func TestResponsesToAnthropic_DefaultToolNormalizesInputSchema(t *testing.T) {
|
||||
|
||||
@@ -4,7 +4,10 @@
|
||||
// formats can be served through a unified gateway.
|
||||
package apicompat
|
||||
|
||||
import "encoding/json"
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Anthropic Messages API types
|
||||
@@ -112,7 +115,7 @@ type AnthropicTool struct {
|
||||
Type string `json:"type,omitempty"` // e.g. "web_search_20250305" for server tools
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
InputSchema json.RawMessage `json:"input_schema"` // JSON Schema object
|
||||
InputSchema json.RawMessage `json:"input_schema,omitempty"` // JSON Schema object
|
||||
CacheControl *AnthropicCacheControl `json:"cache_control,omitempty"`
|
||||
}
|
||||
|
||||
@@ -260,7 +263,34 @@ type ResponsesInputItem struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
|
||||
// type=function_call_output
|
||||
Output string `json:"output,omitempty"`
|
||||
Output string `json:"output,omitempty"`
|
||||
outputRaw json.RawMessage
|
||||
}
|
||||
|
||||
func (i *ResponsesInputItem) UnmarshalJSON(data []byte) error {
|
||||
type alias ResponsesInputItem
|
||||
var wire struct {
|
||||
*alias
|
||||
Output json.RawMessage `json:"output"`
|
||||
}
|
||||
|
||||
*i = ResponsesInputItem{}
|
||||
wire.alias = (*alias)(i)
|
||||
if err := json.Unmarshal(data, &wire); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
output := bytes.TrimSpace(wire.Output)
|
||||
if len(output) == 0 || bytes.Equal(output, []byte("null")) {
|
||||
return nil
|
||||
}
|
||||
if err := json.Unmarshal(output, &i.Output); err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
i.outputRaw = append(i.outputRaw[:0], output...)
|
||||
i.Output = string(output)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ResponsesContentPart is a typed content part in a Responses message.
|
||||
|
||||
@@ -66,7 +66,7 @@ const DefaultCacheControlTTL = "5m"
|
||||
// CLICurrentVersion 是 sub2api 当前对外伪装的 Claude Code CLI 版本号(三段 semver)。
|
||||
// 用于 billing attribution block 中的 cc_version=X.Y.Z.{fp} 前缀以及 fingerprint 计算。
|
||||
// 必须与 DefaultHeaders["User-Agent"] 中的版本号严格一致;不一致会被 Anthropic 判第三方。
|
||||
const CLICurrentVersion = "2.1.161"
|
||||
const CLICurrentVersion = "2.1.220"
|
||||
|
||||
// FullClaudeCodeMimicryBetas 返回最"像"真实 Claude Code CLI 的完整 beta 列表,
|
||||
// 用于 OAuth 账号伪装成 Claude Code 时使用。
|
||||
@@ -147,6 +147,12 @@ var DefaultModels = []Model{
|
||||
DisplayName: "Claude Opus 4.8",
|
||||
CreatedAt: "2026-05-29T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-5",
|
||||
Type: "model",
|
||||
DisplayName: "Claude Opus 5",
|
||||
CreatedAt: "2026-07-25T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-5",
|
||||
Type: "model",
|
||||
|
||||
@@ -270,6 +270,7 @@ type UsageLogFilters struct {
|
||||
APIKeyID int64
|
||||
AccountID int64
|
||||
GroupID int64
|
||||
RequestID string
|
||||
Model string
|
||||
// ModelFilterSource controls how Model is matched. Empty preserves raw usage_logs.model semantics.
|
||||
ModelFilterSource string
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
package liveattestation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrUnsupportedPlatform = errors.New("live attestation is only supported when Sub2API runs on macOS; Windows support is not implemented yet")
|
||||
ErrChatGPTAppMissing = errors.New("live attestation requires the official ChatGPT app on the Sub2API server")
|
||||
)
|
||||
|
||||
// Provider 在发起 Live 请求前生成 ChatGPT DeviceCheck attestation。
|
||||
type Provider interface {
|
||||
Check(ctx context.Context) error
|
||||
Generate(ctx context.Context) (string, error)
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
//go:build darwin
|
||||
|
||||
package liveattestation
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
chatGPTApplicationPath = "/Applications/ChatGPT.app"
|
||||
attestationTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
type darwinProvider struct {
|
||||
appSessionID string
|
||||
appPaths []string
|
||||
}
|
||||
|
||||
type deviceSignals struct {
|
||||
SchemaVersion int `json:"schemaVersion"`
|
||||
PreferredLanguages []string `json:"preferredLanguages"`
|
||||
Locale string `json:"locale"`
|
||||
Timezone string `json:"timezone"`
|
||||
ScreenSizeSum int `json:"screenSizeSum"`
|
||||
ScreenScale float64 `json:"screenScale"`
|
||||
AppSessionID string `json:"appSessionId"`
|
||||
}
|
||||
|
||||
type macOSSignals struct {
|
||||
Locale string `json:"locale"`
|
||||
Languages []string `json:"languages"`
|
||||
Timezone string `json:"timezone"`
|
||||
Width float64 `json:"width"`
|
||||
Height float64 `json:"height"`
|
||||
Scale float64 `json:"scale"`
|
||||
}
|
||||
|
||||
func NewProvider() Provider {
|
||||
paths := []string{chatGPTApplicationPath}
|
||||
if home, err := os.UserHomeDir(); err == nil && strings.TrimSpace(home) != "" {
|
||||
paths = append(paths, filepath.Join(home, "Applications", "ChatGPT.app"))
|
||||
}
|
||||
return &darwinProvider{
|
||||
appSessionID: uuid.NewString(),
|
||||
appPaths: paths,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *darwinProvider) Check(ctx context.Context) error {
|
||||
checkCtx, cancel := context.WithTimeout(ctx, attestationTimeout)
|
||||
defer cancel()
|
||||
_, _, _, err := p.resolveRuntime(checkCtx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (p *darwinProvider) resolveRuntime(ctx context.Context) (string, string, string, error) {
|
||||
if runtime.GOARCH != "arm64" {
|
||||
return "", "", "", errors.New("live attestation currently requires Apple Silicon; Intel macOS is not supported")
|
||||
}
|
||||
appPath, err := p.findApplication()
|
||||
if err != nil {
|
||||
return "", "", "", err
|
||||
}
|
||||
resourcesPath := filepath.Join(appPath, "Contents", "Resources")
|
||||
nodePath := filepath.Join(resourcesPath, "cua_node", "bin", "node")
|
||||
modulePath := filepath.Join(resourcesPath, "native", "devicecheck.node")
|
||||
for filePath, label := range map[string]string{
|
||||
nodePath: "bundled Node.js runtime",
|
||||
modulePath: "DeviceCheck native module",
|
||||
} {
|
||||
if info, statErr := os.Stat(filePath); statErr != nil || info.IsDir() {
|
||||
return "", "", "", fmt.Errorf("%w: ChatGPT app is missing its %s", ErrChatGPTAppMissing, label)
|
||||
}
|
||||
}
|
||||
bundleID, err := readBundleIdentifier(ctx, appPath)
|
||||
if err != nil {
|
||||
return "", "", "", err
|
||||
}
|
||||
return nodePath, modulePath, bundleID, nil
|
||||
}
|
||||
|
||||
func (p *darwinProvider) Generate(ctx context.Context) (string, error) {
|
||||
runCtx, cancel := context.WithTimeout(ctx, attestationTimeout)
|
||||
defer cancel()
|
||||
nodePath, modulePath, bundleID, err := p.resolveRuntime(runCtx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
signals, err := p.readSignals(runCtx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
signalsJSON, err := json.Marshal(signals)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("encode Live attestation signals: %w", err)
|
||||
}
|
||||
|
||||
command := exec.CommandContext(runCtx, nodePath, "-e", deviceCheckScript)
|
||||
command.Env = []string{
|
||||
"PATH=/usr/bin:/bin",
|
||||
"SUB2API_DEVICECHECK_MODULE=" + modulePath,
|
||||
"SUB2API_ATTESTATION_BUNDLE_ID=" + bundleID,
|
||||
"SUB2API_ATTESTATION_SIGNALS=" + string(signalsJSON),
|
||||
}
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
command.Stdout = &stdout
|
||||
command.Stderr = &stderr
|
||||
if err := command.Run(); err != nil {
|
||||
if errors.Is(runCtx.Err(), context.DeadlineExceeded) {
|
||||
return "", errors.New("ChatGPT DeviceCheck token generation timed out")
|
||||
}
|
||||
reason := strings.TrimSpace(stderr.String())
|
||||
if len(reason) > 240 {
|
||||
reason = reason[:240]
|
||||
}
|
||||
if reason == "" {
|
||||
reason = err.Error()
|
||||
}
|
||||
return "", fmt.Errorf("ChatGPT DeviceCheck token generation failed: %s", reason)
|
||||
}
|
||||
header := strings.TrimSpace(stdout.String())
|
||||
if len(header) < 20 || len(header) > 16*1024 || !json.Valid([]byte(header)) {
|
||||
return "", errors.New("ChatGPT DeviceCheck returned a malformed attestation")
|
||||
}
|
||||
return header, nil
|
||||
}
|
||||
|
||||
func (p *darwinProvider) findApplication() (string, error) {
|
||||
for _, appPath := range p.appPaths {
|
||||
info, err := os.Stat(appPath)
|
||||
if err == nil && info.IsDir() {
|
||||
return appPath, nil
|
||||
}
|
||||
}
|
||||
return "", ErrChatGPTAppMissing
|
||||
}
|
||||
|
||||
func readBundleIdentifier(ctx context.Context, appPath string) (string, error) {
|
||||
infoPlist := filepath.Join(appPath, "Contents", "Info.plist")
|
||||
output, err := exec.CommandContext(
|
||||
ctx,
|
||||
"/usr/bin/plutil",
|
||||
"-extract",
|
||||
"CFBundleIdentifier",
|
||||
"raw",
|
||||
infoPlist,
|
||||
).Output()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: cannot read its bundle identifier", ErrChatGPTAppMissing)
|
||||
}
|
||||
bundleID := strings.TrimSpace(string(output))
|
||||
if !strings.HasPrefix(bundleID, "com.openai.") {
|
||||
return "", errors.New("the installed ChatGPT app has an unexpected bundle identifier")
|
||||
}
|
||||
return bundleID, nil
|
||||
}
|
||||
|
||||
func (p *darwinProvider) readSignals(ctx context.Context) (deviceSignals, error) {
|
||||
const script = `ObjC.import("Foundation"); ObjC.import("AppKit");
|
||||
const screen = $.NSScreen.mainScreen;
|
||||
const frame = screen.frame;
|
||||
JSON.stringify({
|
||||
locale: ObjC.unwrap($.NSLocale.currentLocale.localeIdentifier),
|
||||
languages: ObjC.deepUnwrap($.NSLocale.preferredLanguages),
|
||||
timezone: ObjC.unwrap($.NSTimeZone.localTimeZone.name),
|
||||
width: Number(frame.size.width),
|
||||
height: Number(frame.size.height),
|
||||
scale: Number(screen.backingScaleFactor)
|
||||
})`
|
||||
output, err := exec.CommandContext(ctx, "/usr/bin/osascript", "-l", "JavaScript", "-e", script).Output()
|
||||
if err != nil {
|
||||
return deviceSignals{}, fmt.Errorf("read macOS signals for Live attestation: %w", err)
|
||||
}
|
||||
var values macOSSignals
|
||||
if err := json.Unmarshal(output, &values); err != nil {
|
||||
return deviceSignals{}, fmt.Errorf("decode macOS signals for Live attestation: %w", err)
|
||||
}
|
||||
locale := truncateSignal(values.Locale, 64, "unknown")
|
||||
languages := values.Languages
|
||||
if len(languages) == 0 {
|
||||
languages = []string{locale}
|
||||
}
|
||||
if len(languages) > 16 {
|
||||
languages = languages[:16]
|
||||
}
|
||||
for index := range languages {
|
||||
languages[index] = truncateSignal(languages[index], 64, locale)
|
||||
}
|
||||
scale := values.Scale
|
||||
if scale <= 0 {
|
||||
scale = 1
|
||||
}
|
||||
return deviceSignals{
|
||||
SchemaVersion: 1,
|
||||
PreferredLanguages: languages,
|
||||
Locale: locale,
|
||||
Timezone: truncateSignal(values.Timezone, 64, "unknown"),
|
||||
ScreenSizeSum: max(0, int(values.Width+values.Height+0.5)),
|
||||
ScreenScale: scale,
|
||||
AppSessionID: truncateSignal(p.appSessionID, 128, uuid.NewString()),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func truncateSignal(value string, limit int, fallback string) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
value = fallback
|
||||
}
|
||||
if len(value) > limit {
|
||||
return value[:limit]
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
const deviceCheckScript = `
|
||||
const addon = require(process.env.SUB2API_DEVICECHECK_MODULE);
|
||||
const signals = JSON.parse(process.env.SUB2API_ATTESTATION_SIGNALS);
|
||||
const bundleID = process.env.SUB2API_ATTESTATION_BUNDLE_ID;
|
||||
|
||||
function head(major, value) {
|
||||
if (value < 24) return Buffer.from([major + value]);
|
||||
if (value <= 255) return Buffer.from([major + 24, value]);
|
||||
if (value <= 65535) {
|
||||
const out = Buffer.allocUnsafe(3);
|
||||
out[0] = major + 25;
|
||||
out.writeUInt16BE(value, 1);
|
||||
return out;
|
||||
}
|
||||
const out = Buffer.allocUnsafe(5);
|
||||
out[0] = major + 26;
|
||||
out.writeUInt32BE(value, 1);
|
||||
return out;
|
||||
}
|
||||
function uint(value) { return head(0, value); }
|
||||
function text(value) {
|
||||
const body = Buffer.from(value, "utf8");
|
||||
return Buffer.concat([head(96, body.length), body]);
|
||||
}
|
||||
function float(value) {
|
||||
if (Number.isSafeInteger(value) && value >= 0) return uint(value);
|
||||
const out = Buffer.allocUnsafe(9);
|
||||
out[0] = 251;
|
||||
out.writeDoubleBE(value, 1);
|
||||
return out;
|
||||
}
|
||||
function array(values) { return Buffer.concat([head(128, values.length), ...values]); }
|
||||
function map(entries) {
|
||||
return Buffer.concat([head(160, entries.length), ...entries.flatMap(([key, value]) => [uint(key), value])]);
|
||||
}
|
||||
function field(key, value) { return Buffer.concat([text(key), text(value)]); }
|
||||
function base64url(value) {
|
||||
return value.toString("base64").replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/u, "");
|
||||
}
|
||||
|
||||
(async () => {
|
||||
const result = await addon.generateToken();
|
||||
if (!result || !result.supported) throw new Error("DeviceCheck is not supported on this Mac");
|
||||
if (!result.tokenBase64) throw new Error("DeviceCheck returned no token");
|
||||
const fingerprint = map([
|
||||
[0, uint(signals.schemaVersion)],
|
||||
[1, array(signals.preferredLanguages.map(text))],
|
||||
[2, text(signals.locale)],
|
||||
[3, text(signals.timezone)],
|
||||
[4, uint(signals.screenSizeSum)],
|
||||
[5, float(signals.screenScale)],
|
||||
[6, text(signals.appSessionId)]
|
||||
]);
|
||||
const fields = [
|
||||
field("token", result.tokenBase64),
|
||||
field("bundle_id", bundleID),
|
||||
Buffer.concat([text("f"), head(64, fingerprint.length), fingerprint])
|
||||
];
|
||||
if (result.latencyMs != null) {
|
||||
fields.push(Buffer.concat([text("t"), float(result.latencyMs)]));
|
||||
}
|
||||
const token = "v1." + base64url(Buffer.concat([Buffer.from([160 + fields.length]), ...fields]));
|
||||
process.stdout.write(JSON.stringify({v: 1, s: 0, t: token}));
|
||||
})().catch((error) => {
|
||||
process.stderr.write(error instanceof Error ? error.message : String(error));
|
||||
process.exitCode = 1;
|
||||
});`
|
||||
@@ -0,0 +1,19 @@
|
||||
//go:build !darwin
|
||||
|
||||
package liveattestation
|
||||
|
||||
import "context"
|
||||
|
||||
type unsupportedProvider struct{}
|
||||
|
||||
func NewProvider() Provider {
|
||||
return unsupportedProvider{}
|
||||
}
|
||||
|
||||
func (unsupportedProvider) Check(context.Context) error {
|
||||
return ErrUnsupportedPlatform
|
||||
}
|
||||
|
||||
func (unsupportedProvider) Generate(context.Context) (string, error) {
|
||||
return "", ErrUnsupportedPlatform
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//go:build !darwin
|
||||
|
||||
package liveattestation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestUnsupportedProviderReturnsExplicitPlatformError(t *testing.T) {
|
||||
provider := NewProvider()
|
||||
if err := provider.Check(context.Background()); !errors.Is(err, ErrUnsupportedPlatform) {
|
||||
t.Fatalf("Check() error = %v, want ErrUnsupportedPlatform", err)
|
||||
}
|
||||
_, err := provider.Generate(context.Background())
|
||||
if !errors.Is(err, ErrUnsupportedPlatform) {
|
||||
t.Fatalf("Generate() error = %v, want ErrUnsupportedPlatform", err)
|
||||
}
|
||||
}
|
||||
@@ -3358,6 +3358,12 @@ func (r *accountRepository) FindByExtraField(ctx context.Context, key string, va
|
||||
// ListDueUpstreamBillingProbeAccounts bounds result hydration and network work
|
||||
// to limit. PostgreSQL must still filter and order all enabled candidates;
|
||||
// MATERIALIZED avoids repeating the defensive timestamp parse expression.
|
||||
// Go writes next_probe_at via RFC3339Nano (up to 9 fractional digits) while
|
||||
// jsonpath datetime() parses at most microseconds, so fractions beyond 6
|
||||
// digits are trimmed first — mirroring ListDueOllamaCloudUsageAccounts.
|
||||
// Without this, every nanosecond timestamp is treated as malformed and the
|
||||
// fail-open ordering pins the cycle to the lowest account IDs, starving the
|
||||
// rest of the pool.
|
||||
func (r *accountRepository) ListDueUpstreamBillingProbeAccounts(ctx context.Context, now time.Time, limit int) ([]service.Account, error) {
|
||||
if limit <= 0 {
|
||||
return []service.Account{}, nil
|
||||
@@ -3387,7 +3393,11 @@ func (r *accountRepository) ListDueUpstreamBillingProbeAccounts(ctx context.Cont
|
||||
jsonb_path_query_first_tz(
|
||||
jsonb_build_object(
|
||||
'value',
|
||||
replace(regexp_replace(next_probe_at, 'Z$', '+00:00'), 'T', ' ')
|
||||
replace(regexp_replace(regexp_replace(
|
||||
next_probe_at,
|
||||
'(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
|
||||
'\1\2'
|
||||
), 'Z$', '+00:00'), 'T', ' ')
|
||||
),
|
||||
'$.value.datetime()',
|
||||
'{}'::jsonb,
|
||||
|
||||
@@ -354,74 +354,182 @@ func canonicalJSON(raw string) string {
|
||||
return string(encoded)
|
||||
}
|
||||
|
||||
// ListDueOllamaCloudUsageAccounts returns at most one due representative per
|
||||
// exact API key before hydration, preventing one shared group from consuming a
|
||||
// whole runner cycle.
|
||||
func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context, now time.Time, limit int) ([]service.Account, error) {
|
||||
// ollamaCloudUsageParseRFC3339SQL reuses the verified RFC3339(/Nano) parse path
|
||||
// for a snapshot timestamp expression. Invalid or missing values fail open to NULL.
|
||||
//
|
||||
// The value is rewritten twice before it reaches jsonpath:
|
||||
// 1. Sub-second precision beyond 6 digits is truncated, because .datetime()
|
||||
// rejects more than microsecond resolution while Go emits 9 digits.
|
||||
// 2. A trailing "Z" is rewritten to "+00:00". jsonpath .datetime() only learned
|
||||
// to accept the ISO-8601 "Z" designator in PostgreSQL 17, and every timestamp
|
||||
// this service writes is UTC (hence "Z"). Without this rewrite the parse
|
||||
// silently yields NULL on PostgreSQL <= 16, which makes every due column NULL
|
||||
// and collapses ListDueOllamaCloudUsageAccounts into its fail-open branch.
|
||||
//
|
||||
// jsonpath (rather than a direct ::timestamptz cast) is required so that values
|
||||
// passing the shape regex but naming an impossible date (e.g. 2026-02-30) fail
|
||||
// open to NULL instead of aborting the whole query.
|
||||
func ollamaCloudUsageParseRFC3339SQL(expression string) string {
|
||||
return `CASE
|
||||
WHEN ` + expression + ` IS NULL THEN NULL
|
||||
WHEN ` + expression + ` ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:[0-9]{2})$'
|
||||
THEN jsonb_path_query_first_tz(
|
||||
to_jsonb(regexp_replace(
|
||||
regexp_replace(
|
||||
` + expression + `,
|
||||
'(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
|
||||
'\1\2'
|
||||
),
|
||||
'Z$',
|
||||
'+00:00'
|
||||
)),
|
||||
'$.datetime()', '{}'::jsonb, true
|
||||
) #>> '{}'
|
||||
ELSE NULL
|
||||
END`
|
||||
}
|
||||
|
||||
// ListDueOllamaCloudUsageAccounts returns at most one truly-due activity-driven
|
||||
// candidate per exact API key. Due timing (debounce, max-wait, failure backoff)
|
||||
// is evaluated in SQL before LIMIT so non-due active groups cannot starve due ones.
|
||||
// Account.LastUsedAt is stamped with the group MAX(last_used_at) for a service
|
||||
// pure-function recheck against races between list and refresh.
|
||||
//
|
||||
// Rules mirror service.ollamaCloudUsageAutoRefreshDueAt (keep both in sync):
|
||||
// - missing/invalid snapshot or times → fail-open first due
|
||||
// - success: activity after fetched_at;
|
||||
// due_at = GREATEST(LEAST(last_used+debounce, fetched+maxWait), fetched+minFetchInterval)
|
||||
// - failed/unauthorized: activity after last_attempt; activity_due = LEAST(...);
|
||||
// final due_at is not earlier than a valid next_refresh_at (invalid/missing fail-open)
|
||||
func (r *accountRepository) ListDueOllamaCloudUsageAccounts(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
debounce, maxWait time.Duration,
|
||||
limit int,
|
||||
) ([]service.Account, error) {
|
||||
if limit <= 0 {
|
||||
return []service.Account{}, nil
|
||||
}
|
||||
if r == nil || r.sql == nil {
|
||||
return nil, errors.New("account repository SQL executor not configured")
|
||||
}
|
||||
if debounce <= 0 {
|
||||
debounce = time.Minute
|
||||
}
|
||||
if maxWait <= 0 {
|
||||
maxWait = time.Hour
|
||||
}
|
||||
debounceSeconds := debounce.Seconds()
|
||||
maxWaitSeconds := maxWait.Seconds()
|
||||
minFetchIntervalSeconds := service.OllamaCloudUsageMinFetchInterval.Seconds()
|
||||
rows, err := r.sql.QueryContext(ctx, `
|
||||
WITH candidates AS (
|
||||
SELECT id, credentials ->> 'api_key' AS api_key,
|
||||
extra #>> '{ollama_cloud_usage_snapshot,next_refresh_at}' AS next_refresh_at
|
||||
WITH eligible AS (
|
||||
SELECT id,
|
||||
credentials ->> 'api_key' AS api_key,
|
||||
last_used_at,
|
||||
extra -> 'ollama_cloud_usage_snapshot' AS snapshot
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND status = 'active'
|
||||
AND `+ollamaCloudUsageEligibleSQL+`
|
||||
AND jsonb_typeof(extra -> 'ollama_cloud_usage_session') = 'string'
|
||||
AND extra @> '{"ollama_cloud_usage_auto_refresh": true}'::jsonb
|
||||
), group_activity AS (
|
||||
SELECT credentials ->> 'api_key' AS api_key,
|
||||
MAX(last_used_at) AS group_last_used_at
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND `+ollamaCloudUsageEligibleSQL+`
|
||||
AND jsonb_typeof(credentials -> 'api_key') = 'string'
|
||||
GROUP BY credentials ->> 'api_key'
|
||||
), joined AS (
|
||||
SELECT e.id, e.api_key, e.snapshot, g.group_last_used_at,
|
||||
e.snapshot #>> '{status}' AS status,
|
||||
e.snapshot #>> '{fetched_at}' AS fetched_at,
|
||||
e.snapshot #>> '{last_attempt_at}' AS last_attempt_at,
|
||||
e.snapshot #>> '{next_refresh_at}' AS next_refresh_at
|
||||
FROM eligible e
|
||||
JOIN group_activity g ON g.api_key = e.api_key
|
||||
), parsed AS MATERIALIZED (
|
||||
SELECT id, api_key, next_refresh_at,
|
||||
next_refresh_at ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:[0-9]{2})$' AS rfc3339_shape,
|
||||
jsonb_path_query_first_tz(
|
||||
to_jsonb(regexp_replace(
|
||||
next_refresh_at,
|
||||
'(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
|
||||
'\1\2'
|
||||
)),
|
||||
'$.datetime()', '{}'::jsonb, true
|
||||
) #>> '{}' AS parsed_next_refresh_at
|
||||
FROM candidates
|
||||
), due AS (
|
||||
SELECT id, api_key, snapshot, group_last_used_at, status,
|
||||
`+ollamaCloudUsageParseRFC3339SQL("fetched_at")+` AS parsed_fetched_at,
|
||||
`+ollamaCloudUsageParseRFC3339SQL("last_attempt_at")+` AS parsed_last_attempt_at,
|
||||
`+ollamaCloudUsageParseRFC3339SQL("next_refresh_at")+` AS parsed_next_refresh_at
|
||||
FROM joined
|
||||
), timed AS (
|
||||
SELECT *,
|
||||
CASE WHEN next_refresh_at IS NULL OR NOT rfc3339_shape OR parsed_next_refresh_at IS NULL THEN 0 ELSE 1 END AS due_class
|
||||
CASE
|
||||
WHEN status = 'ok'
|
||||
AND parsed_fetched_at IS NOT NULL
|
||||
AND group_last_used_at IS NOT NULL
|
||||
AND group_last_used_at > parsed_fetched_at::timestamptz
|
||||
THEN GREATEST(
|
||||
LEAST(
|
||||
group_last_used_at + make_interval(secs => $2::double precision),
|
||||
parsed_fetched_at::timestamptz + make_interval(secs => $3::double precision)
|
||||
),
|
||||
parsed_fetched_at::timestamptz + make_interval(secs => $5::double precision)
|
||||
)
|
||||
WHEN status IN ('failed', 'unauthorized')
|
||||
AND parsed_last_attempt_at IS NOT NULL
|
||||
AND group_last_used_at IS NOT NULL
|
||||
AND group_last_used_at > parsed_last_attempt_at::timestamptz
|
||||
THEN GREATEST(
|
||||
LEAST(
|
||||
group_last_used_at + make_interval(secs => $2::double precision),
|
||||
parsed_last_attempt_at::timestamptz + make_interval(secs => $3::double precision)
|
||||
),
|
||||
COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)
|
||||
)
|
||||
ELSE NULL
|
||||
END AS activity_due_at
|
||||
FROM parsed
|
||||
WHERE next_refresh_at IS NULL
|
||||
OR NOT rfc3339_shape
|
||||
OR parsed_next_refresh_at IS NULL
|
||||
OR parsed_next_refresh_at::timestamptz <= $1
|
||||
), candidates AS (
|
||||
SELECT *,
|
||||
CASE
|
||||
WHEN snapshot IS NULL OR snapshot = 'null'::jsonb OR status IS NULL
|
||||
OR status NOT IN ('ok', 'failed', 'unauthorized') THEN 0
|
||||
WHEN status = 'ok' AND parsed_fetched_at IS NULL THEN 0
|
||||
WHEN status IN ('failed', 'unauthorized') AND parsed_last_attempt_at IS NULL THEN 0
|
||||
WHEN activity_due_at IS NOT NULL AND $1 >= activity_due_at THEN 1
|
||||
ELSE NULL
|
||||
END AS due_class,
|
||||
activity_due_at AS due_at
|
||||
FROM timed
|
||||
), ranked AS (
|
||||
SELECT *, row_number() OVER (
|
||||
PARTITION BY api_key
|
||||
ORDER BY due_class,
|
||||
CASE WHEN rfc3339_shape AND parsed_next_refresh_at IS NOT NULL THEN parsed_next_refresh_at::timestamptz END NULLS FIRST,
|
||||
id
|
||||
) AS group_rank
|
||||
FROM due
|
||||
SELECT id, api_key, group_last_used_at, due_class, due_at,
|
||||
row_number() OVER (
|
||||
PARTITION BY api_key
|
||||
ORDER BY due_class,
|
||||
due_at NULLS FIRST,
|
||||
id
|
||||
) AS group_rank
|
||||
FROM candidates
|
||||
WHERE due_class IS NOT NULL
|
||||
)
|
||||
SELECT id
|
||||
SELECT id, group_last_used_at
|
||||
FROM ranked
|
||||
WHERE group_rank = 1
|
||||
ORDER BY due_class,
|
||||
CASE WHEN rfc3339_shape AND parsed_next_refresh_at IS NOT NULL THEN parsed_next_refresh_at::timestamptz END NULLS FIRST,
|
||||
id
|
||||
LIMIT $2
|
||||
`, now.UTC(), limit)
|
||||
ORDER BY due_class, due_at NULLS FIRST, id
|
||||
LIMIT $4
|
||||
`, now.UTC(), debounceSeconds, maxWaitSeconds, limit, minFetchIntervalSeconds)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
type dueRow struct {
|
||||
id int64
|
||||
groupLastUsed *time.Time
|
||||
}
|
||||
rowsOut := make([]dueRow, 0, limit)
|
||||
ids := make([]int64, 0, limit)
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
var row dueRow
|
||||
if err := rows.Scan(&row.id, &row.groupLastUsed); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
rowsOut = append(rowsOut, row)
|
||||
ids = append(ids, row.id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
@@ -430,11 +538,26 @@ func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context,
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]service.Account, 0, len(accounts))
|
||||
byID := make(map[int64]*service.Account, len(accounts))
|
||||
for _, account := range accounts {
|
||||
if account != nil {
|
||||
result = append(result, *account)
|
||||
byID[account.ID] = account
|
||||
}
|
||||
}
|
||||
result := make([]service.Account, 0, len(rowsOut))
|
||||
for _, row := range rowsOut {
|
||||
account := byID[row.id]
|
||||
if account == nil {
|
||||
continue
|
||||
}
|
||||
// Stamp group MAX(last_used_at) for service due evaluation.
|
||||
if row.groupLastUsed != nil {
|
||||
ts := row.groupLastUsed.UTC()
|
||||
account.LastUsedAt = &ts
|
||||
} else {
|
||||
account.LastUsedAt = nil
|
||||
}
|
||||
result = append(result, *account)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -23,81 +23,199 @@ func TestListDueOllamaCloudUsageAccountsOrderingLimitAndProxyHydration(t *testin
|
||||
Username: "user", Password: "pass", Status: service.StatusActive,
|
||||
})
|
||||
|
||||
createAccount := func(name, baseURL string, proxyID *int64, nextRefreshAt *time.Time) *service.Account {
|
||||
createAccount := func(name, baseURL string, proxyID *int64, snapshot map[string]any, lastUsed *time.Time) *service.Account {
|
||||
t.Helper()
|
||||
extra := map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
}
|
||||
if nextRefreshAt != nil {
|
||||
extra[service.OllamaCloudUsageSnapshotExtraKey] = map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK, "next_refresh_at": nextRefreshAt.UTC().Format(time.RFC3339Nano),
|
||||
}
|
||||
if snapshot != nil {
|
||||
extra[service.OllamaCloudUsageSnapshotExtraKey] = snapshot
|
||||
}
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": name, "base_url": baseURL},
|
||||
Extra: extra, ProxyID: proxyID,
|
||||
Extra: extra, ProxyID: proxyID, LastUsedAt: lastUsed,
|
||||
})
|
||||
}
|
||||
|
||||
uppercasePath := createAccount("ollama-uppercase-path", "https://ollama.com/V1", nil, nil)
|
||||
missingSnapshot := createAccount("ollama-due-missing", "HTTPS://WWW.OLLAMA.COM:443/v1", &proxy.ID, nil)
|
||||
oldest := now.Add(-2 * time.Hour)
|
||||
due := createAccount("ollama-due-oldest", "https://ollama.com", nil, &oldest)
|
||||
future := now.Add(time.Minute)
|
||||
_ = createAccount("ollama-not-due", "https://ollama.com", nil, &future)
|
||||
_ = createAccount("ollama-ineligible", "https://ollama.com.evil.test", nil, nil)
|
||||
uppercasePath := createAccount("ollama-uppercase-path", "https://ollama.com/V1", nil, nil, nil)
|
||||
missingSnapshot := createAccount("ollama-due-missing", "HTTPS://WWW.OLLAMA.COM:443/v1", &proxy.ID, nil, nil)
|
||||
fetched := now.Add(-2 * time.Hour)
|
||||
activity := now.Add(-5 * time.Minute)
|
||||
due := createAccount("ollama-due-activity", "https://ollama.com", nil, map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"next_refresh_at": fetched.Add(time.Hour).UTC().Format(time.RFC3339Nano),
|
||||
}, &activity)
|
||||
// Success snapshot without newer activity must not be listed.
|
||||
_ = createAccount("ollama-not-due-idle", "https://ollama.com", nil, map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": now.Add(-time.Hour).UTC().Format(time.RFC3339Nano),
|
||||
"last_attempt_at": now.Add(-time.Hour).UTC().Format(time.RFC3339Nano),
|
||||
"next_refresh_at": now.Add(-time.Minute).UTC().Format(time.RFC3339Nano),
|
||||
}, nil)
|
||||
_ = createAccount("ollama-ineligible", "https://ollama.com.evil.test", nil, nil, nil)
|
||||
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, 2)
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 2)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accounts, 2)
|
||||
require.Equal(t, missingSnapshot.ID, accounts[0].ID)
|
||||
require.Equal(t, due.ID, accounts[1].ID)
|
||||
require.NotNil(t, accounts[1].LastUsedAt)
|
||||
require.WithinDuration(t, activity.UTC(), accounts[1].LastUsedAt.UTC(), time.Second)
|
||||
require.NotContains(t, accountIDs(accounts), uppercasePath.ID)
|
||||
require.NotNil(t, accounts[0].Proxy)
|
||||
require.Equal(t, proxy.ID, accounts[0].Proxy.ID)
|
||||
require.Equal(t, proxy.URL(), accounts[0].Proxy.URL())
|
||||
}
|
||||
|
||||
func TestListDueOllamaCloudUsageAccountsParsesRFC3339NanoAndFailsOpen(t *testing.T) {
|
||||
// TestListDueOllamaCloudUsageAccountsParsesAllRFC3339Precisions pins the SQL
|
||||
// timestamp parse path across the sub-second precisions and zone spellings that
|
||||
// actually reach the database.
|
||||
//
|
||||
// Each fixture stores a fetched_at only two minutes old with activity 30s later,
|
||||
// so a correctly parsed row is NOT due (debounce and the min fetch interval both
|
||||
// place it in the future). A row whose timestamp fails to parse becomes NULL and
|
||||
// falls into the fail-open branch, which makes it due. Asserting on absence is
|
||||
// therefore what makes this test able to fail:
|
||||
//
|
||||
// - Go writes UTC times, i.e. the "Z" designator. jsonpath .datetime() only
|
||||
// accepts "Z" from PostgreSQL 17 on, so without the Z -> +00:00 rewrite in
|
||||
// ollamaCloudUsageParseRFC3339SQL every fixture here goes due on 14-16.
|
||||
// - 7/8/9 sub-second digits exceed the microsecond resolution .datetime()
|
||||
// allows and must be truncated first.
|
||||
//
|
||||
// Run against the oldest supported server to exercise the version-sensitive path:
|
||||
//
|
||||
// SUB2API_TEST_POSTGRES_IMAGE=postgres:15-alpine go test -tags integration ./internal/repository/
|
||||
func TestListDueOllamaCloudUsageAccountsParsesAllRFC3339Precisions(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Date(2026, time.July, 22, 14, 0, 0, 0, time.UTC)
|
||||
activity := now.Add(-30 * time.Second)
|
||||
|
||||
create := func(name, nextRefreshAt string) *service.Account {
|
||||
t.Helper()
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"},
|
||||
// All three spell the same instant, now-2m, with different precision/zone.
|
||||
notDue := map[string]string{
|
||||
"nano-z": "2026-07-22T13:58:00.123456789Z",
|
||||
"eight-positive": "2026-07-22T14:58:00.12345678+01:00",
|
||||
"seven-negative": "2026-07-22T11:58:00.1234567-02:00",
|
||||
}
|
||||
for name, fetchedAt := range notDue {
|
||||
_ = mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "ollama-precision-" + name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "precision-" + name, "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK, "next_refresh_at": nextRefreshAt,
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": fetchedAt,
|
||||
"last_attempt_at": fetchedAt,
|
||||
},
|
||||
},
|
||||
LastUsedAt: &activity,
|
||||
})
|
||||
}
|
||||
|
||||
sevenDigitsOffset := create("ollama-nano-seven", "2026-07-22T11:00:00.1234567-02:00")
|
||||
eightDigitsOffset := create("ollama-nano-eight", "2026-07-22T11:00:00.12345678+01:00")
|
||||
nineDigitsZ := create("ollama-nano-nine", "2026-07-22T09:00:00.123456789Z")
|
||||
invalidCalendar := create("ollama-nano-invalid", "2026-02-30T09:00:00.123456789Z")
|
||||
future := create("ollama-nano-future", "2026-07-22T15:00:00.123456789Z")
|
||||
// Guards against a vacuous pass: an genuinely due row must still come back.
|
||||
staleFetched := now.Add(-2 * time.Hour)
|
||||
due := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "ollama-precision-due", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "precision-due", "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": staleFetched.UTC().Format(time.RFC3339Nano),
|
||||
"last_attempt_at": staleFetched.UTC().Format(time.RFC3339Nano),
|
||||
},
|
||||
},
|
||||
LastUsedAt: &activity,
|
||||
})
|
||||
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, 10)
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 10)
|
||||
|
||||
require.NoError(t, err)
|
||||
ids := accountIDs(accounts)
|
||||
require.Contains(t, ids, due.ID, "a stale snapshot with fresh activity must be due")
|
||||
require.Len(t, ids, 1,
|
||||
"only the stale group may be due; extra rows mean a timestamp failed to parse and fell into the fail-open branch")
|
||||
}
|
||||
|
||||
func TestListDueOllamaCloudUsageAccountsUsesGroupMaxLastUsedAndFailsOpen(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Date(2026, time.July, 22, 14, 0, 0, 0, time.UTC)
|
||||
fetched := now.Add(-30 * time.Minute)
|
||||
older := now.Add(-10 * time.Minute)
|
||||
newer := now.Add(-2 * time.Minute)
|
||||
|
||||
leader := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "ollama-group-leader", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "shared-key", "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"next_refresh_at": fetched.Add(time.Hour).UTC().Format(time.RFC3339Nano),
|
||||
},
|
||||
},
|
||||
LastUsedAt: &older,
|
||||
})
|
||||
_ = mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "ollama-group-sibling", Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "shared-key", "base_url": "https://www.ollama.com/v1"},
|
||||
LastUsedAt: &newer,
|
||||
})
|
||||
invalid := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "ollama-invalid-snapshot", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "invalid-key", "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK, "fetched_at": "2026-02-30T09:00:00.123456789Z",
|
||||
},
|
||||
},
|
||||
})
|
||||
idle := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "ollama-idle-ok", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "idle-key", "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"next_refresh_at": fetched.Add(time.Hour).UTC().Format(time.RFC3339Nano),
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 10)
|
||||
|
||||
require.NoError(t, err, "invalid stored values must not abort the query")
|
||||
require.Equal(t, []int64{
|
||||
invalidCalendar.ID,
|
||||
nineDigitsZ.ID,
|
||||
eightDigitsOffset.ID,
|
||||
sevenDigitsOffset.ID,
|
||||
}, accountIDs(accounts))
|
||||
require.NotContains(t, accountIDs(accounts), future.ID)
|
||||
ids := accountIDs(accounts)
|
||||
require.Contains(t, ids, invalid.ID)
|
||||
require.Contains(t, ids, leader.ID)
|
||||
require.NotContains(t, ids, idle.ID)
|
||||
for _, account := range accounts {
|
||||
if account.ID == leader.ID {
|
||||
require.NotNil(t, account.LastUsedAt)
|
||||
require.WithinDuration(t, newer.UTC(), account.LastUsedAt.UTC(), time.Second,
|
||||
"group MAX(last_used_at) must come from the sibling")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLockAndMergeAccountProbeExtraCoalescesNullableOllamaGroupIdentity(t *testing.T) {
|
||||
@@ -413,3 +531,111 @@ func TestUpdateCredentialsUnchangedCredentialsPreserveManagedExtra(t *testing.T)
|
||||
require.NotContains(t, probeLoaded.Extra, service.UpstreamBillingProbeExtraKey,
|
||||
"changed credentials must keep clearing the probe snapshot")
|
||||
}
|
||||
|
||||
// TestListDueOllamaCloudUsageAccountsSQLDueRulesMatchService proves the SQL
|
||||
// candidate layer applies debounce / max-wait / failure-backoff before LIMIT,
|
||||
// matching service.ollamaCloudUsageIsAutoRefreshDue, and that >20 active-but-
|
||||
// not-yet-due groups cannot starve a truly due max-wait group.
|
||||
func TestListDueOllamaCloudUsageAccountsSQLDueRulesMatchService(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC)
|
||||
debounce := time.Minute
|
||||
maxWait := time.Hour
|
||||
|
||||
createOK := func(name string, fetched, lastUsed time.Time) *service.Account {
|
||||
t.Helper()
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK,
|
||||
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
|
||||
"next_refresh_at": fetched.Add(maxWait).UTC().Format(time.RFC3339Nano),
|
||||
},
|
||||
},
|
||||
LastUsedAt: &lastUsed,
|
||||
})
|
||||
}
|
||||
createFailed := func(name string, lastAttempt, lastUsed, nextRefresh time.Time, nextRefreshRaw string) *service.Account {
|
||||
t.Helper()
|
||||
snapshot := map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusFailed,
|
||||
"last_attempt_at": lastAttempt.UTC().Format(time.RFC3339Nano),
|
||||
"failure_count": 1,
|
||||
}
|
||||
if nextRefreshRaw != "" {
|
||||
snapshot["next_refresh_at"] = nextRefreshRaw
|
||||
} else {
|
||||
snapshot["next_refresh_at"] = nextRefresh.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"},
|
||||
Extra: map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: snapshot,
|
||||
},
|
||||
LastUsedAt: &lastUsed,
|
||||
})
|
||||
}
|
||||
|
||||
// 21 groups with activity after fetch but debounce not elapsed — previously
|
||||
// these alone could fill LIMIT 20 every minute and starve true due groups.
|
||||
notDueIDs := make(map[int64]struct{}, 21)
|
||||
for i := 0; i < 21; i++ {
|
||||
// fetched 10m ago, last used 10s ago → due_at = lastUsed+debounce = now+50s (not due)
|
||||
acc := createOK(fmt.Sprintf("ollama-not-due-debounce-%02d", i), now.Add(-10*time.Minute), now.Add(-10*time.Second))
|
||||
notDueIDs[acc.ID] = struct{}{}
|
||||
}
|
||||
|
||||
// Truly due via max-wait: fetched 2h ago, continuous activity 10s ago.
|
||||
// due_at = min(now-10s+1m, now-2h+1h) = now-1h → due.
|
||||
maxWaitDue := createOK("ollama-due-maxwait", now.Add(-2*time.Hour), now.Add(-10*time.Second))
|
||||
|
||||
// Success debounce elapsed: last used 2m ago with debounce 1m → due.
|
||||
debounceDue := createOK("ollama-due-debounce", now.Add(-30*time.Minute), now.Add(-2*time.Minute))
|
||||
|
||||
// Success still within debounce → not due.
|
||||
_ = createOK("ollama-not-due-fresh", now.Add(-30*time.Minute), now.Add(-20*time.Second))
|
||||
|
||||
// Failure blocked by next_refresh_at backoff even with new activity.
|
||||
_ = createFailed("ollama-fail-backoff", now.Add(-30*time.Minute), now.Add(-2*time.Minute), now.Add(10*time.Minute), "")
|
||||
|
||||
// Failure after backoff with new request → due.
|
||||
failDue := createFailed("ollama-fail-due", now.Add(-30*time.Minute), now.Add(-2*time.Minute), now.Add(-time.Minute), "")
|
||||
|
||||
// Invalid next_refresh_at must fail open (not abort query / not block activity due).
|
||||
failInvalidNext := createFailed("ollama-fail-invalid-next", now.Add(-30*time.Minute), now.Add(-2*time.Minute), time.Time{}, "not-a-timestamp")
|
||||
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, debounce, maxWait, 20)
|
||||
require.NoError(t, err)
|
||||
|
||||
ids := accountIDs(accounts)
|
||||
require.Contains(t, ids, maxWaitDue.ID, "max-wait due group must not be starved by not-yet-due activity groups")
|
||||
require.Contains(t, ids, debounceDue.ID, "success debounce elapsed must be due in SQL")
|
||||
require.Contains(t, ids, failDue.ID, "failure after backoff with new activity must be due in SQL")
|
||||
require.Contains(t, ids, failInvalidNext.ID, "invalid next_refresh_at must fail open to activity due")
|
||||
require.LessOrEqual(t, len(accounts), 20)
|
||||
|
||||
// Fixtures below match service.ollamaCloudUsageIsAutoRefreshDue semantics;
|
||||
// none of the not-yet-due groups may appear even when they outnumber the limit.
|
||||
for _, id := range ids {
|
||||
_, isNotDue := notDueIDs[id]
|
||||
require.False(t, isNotDue, "not-yet-due debounce group %d must not be returned by SQL LIMIT layer", id)
|
||||
}
|
||||
require.NotContains(t, ids, int64(0))
|
||||
|
||||
// Explicit not-due names must stay out: fresh success and failure still in backoff.
|
||||
for _, account := range accounts {
|
||||
require.NotContains(t, account.Name, "not-due")
|
||||
require.NotEqual(t, "ollama-fail-backoff", account.Name)
|
||||
require.NotEqual(t, "ollama-not-due-fresh", account.Name)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,13 +190,15 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC)
|
||||
debounce := time.Minute
|
||||
maxWait := time.Hour
|
||||
var capturedSQL string
|
||||
mock.ExpectQuery("WITH candidates AS").
|
||||
WithArgs(now, 20).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
mock.ExpectQuery("WITH eligible AS").
|
||||
WithArgs(now.UTC(), debounce.Seconds(), maxWait.Seconds(), 20, service.OllamaCloudUsageMinFetchInterval.Seconds()).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "group_last_used_at"}))
|
||||
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil)
|
||||
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, 20)
|
||||
accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, debounce, maxWait, 20)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, accounts)
|
||||
@@ -209,10 +211,24 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) {
|
||||
ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'"),
|
||||
"jsonb_typeof(extra -> 'ollama_cloud_usage_session') = 'string'",
|
||||
`extra @> '{"ollama_cloud_usage_auto_refresh": true}'::jsonb`,
|
||||
"parsed_next_refresh_at::timestamptz <= $1",
|
||||
"MAX(last_used_at) AS group_last_used_at",
|
||||
"PARTITION BY api_key",
|
||||
"WHERE group_rank = 1",
|
||||
"LIMIT $2",
|
||||
"LIMIT $4",
|
||||
"make_interval(secs => $2::double precision)",
|
||||
"make_interval(secs => $3::double precision)",
|
||||
// Minimum interval floor between successful fetches.
|
||||
"make_interval(secs => $5::double precision)",
|
||||
// jsonpath .datetime() only accepts the ISO-8601 "Z" designator from
|
||||
// PostgreSQL 17 on, and this service writes UTC timestamps. Without this
|
||||
// rewrite every parsed_* column is NULL on 14-16 and the due filter
|
||||
// collapses into its fail-open branch.
|
||||
`regexp_replace( regexp_replace( fetched_at, '(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$', '\1\2' ), 'Z$', '+00:00' )`,
|
||||
"group_last_used_at > parsed_fetched_at::timestamptz",
|
||||
"group_last_used_at > parsed_last_attempt_at::timestamptz",
|
||||
"$1 >= activity_due_at",
|
||||
"COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)",
|
||||
"ORDER BY due_class, due_at NULLS FIRST, id",
|
||||
} {
|
||||
require.Contains(t, normalized, clause)
|
||||
}
|
||||
|
||||
+104
@@ -49,3 +49,107 @@ func TestListDueUpstreamBillingProbeAccountsHandlesInvalidCalendarDate(t *testin
|
||||
require.Equal(t, invalidID, accounts[0].ID)
|
||||
require.Equal(t, dueID, accounts[1].ID)
|
||||
}
|
||||
|
||||
func insertUpstreamBillingProbeAccount(ctx context.Context, t *testing.T, tx sqlQueryer, name, nextProbeAt string) int64 {
|
||||
t.Helper()
|
||||
var id int64
|
||||
extra := fmt.Sprintf(`{
|
||||
"upstream_billing_probe_enabled": true,
|
||||
"upstream_billing_probe": {"status": "ok", "next_probe_at": %q}
|
||||
}`, nextProbeAt)
|
||||
err := scanSingleRow(ctx, tx, `
|
||||
INSERT INTO accounts (name, platform, type, status, extra)
|
||||
VALUES ($1, 'openai', $2, 'active', $3::jsonb)
|
||||
RETURNING id
|
||||
`, []any{name, service.AccountTypeAPIKey, extra}, &id)
|
||||
require.NoError(t, err)
|
||||
return id
|
||||
}
|
||||
|
||||
// rfc3339WithFraction renders t in UTC with an explicit fractional-second
|
||||
// suffix, e.g. "2026-07-25T17:29:00.123456789Z". The probes' Go writer uses
|
||||
// RFC3339Nano, so stored fractions carry up to 9 digits.
|
||||
func rfc3339WithFraction(t time.Time, fraction string) string {
|
||||
return fmt.Sprintf("%s.%sZ", t.UTC().Format("2006-01-02T15:04:05"), fraction)
|
||||
}
|
||||
|
||||
// Regression for the scheduled-probe starvation bug: Go persists next_probe_at
|
||||
// via RFC3339Nano (7-9 fractional digits), which jsonpath datetime() cannot
|
||||
// parse. Before the trim fix every such row was treated as malformed and
|
||||
// fail-open "due", so the cycle always returned the lowest account IDs and
|
||||
// higher IDs were never probed.
|
||||
func TestListDueUpstreamBillingProbeAccountsParsesNanosecondTimestamps(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Date(2026, time.July, 25, 17, 30, 0, 0, time.UTC)
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
UPDATE accounts
|
||||
SET extra = extra - 'upstream_billing_probe_enabled' - 'upstream_billing_probe'
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 22 low-ID accounts that are NOT due yet, stored with 9 fractional digits
|
||||
// exactly as the legacy writer produced them (no migration).
|
||||
notDue := now.Add(time.Hour)
|
||||
for i := 0; i < 22; i++ {
|
||||
insertUpstreamBillingProbeAccount(ctx, t, tx,
|
||||
fmt.Sprintf("probe-nano-not-due-%02d", i),
|
||||
rfc3339WithFraction(notDue.Add(time.Duration(i)*time.Second), "123456789"))
|
||||
}
|
||||
// 3 high-ID accounts that ARE due, with 9/8/7 fractional digits. Their
|
||||
// parsed order must follow due time, not insertion/ID order.
|
||||
dueThird := insertUpstreamBillingProbeAccount(ctx, t, tx,
|
||||
"probe-nano-due-9digits", rfc3339WithFraction(now.Add(-time.Minute), "123456789"))
|
||||
dueFirst := insertUpstreamBillingProbeAccount(ctx, t, tx,
|
||||
"probe-nano-due-8digits", rfc3339WithFraction(now.Add(-3*time.Minute), "12345678"))
|
||||
dueSecond := insertUpstreamBillingProbeAccount(ctx, t, tx,
|
||||
"probe-nano-due-7digits", rfc3339WithFraction(now.Add(-2*time.Minute), "1234567"))
|
||||
|
||||
accounts, err := repo.ListDueUpstreamBillingProbeAccounts(ctx, now, 20)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accounts, 3)
|
||||
require.Equal(t, dueFirst, accounts[0].ID)
|
||||
require.Equal(t, dueSecond, accounts[1].ID)
|
||||
require.Equal(t, dueThird, accounts[2].ID)
|
||||
}
|
||||
|
||||
// With more due accounts than the cycle limit, selection must follow the
|
||||
// parsed due time so accounts beyond the first <limit> IDs still get their
|
||||
// turn; before the fix the same lowest IDs monopolized every cycle.
|
||||
func TestListDueUpstreamBillingProbeAccountsSelectsEarliestDueAcrossIDs(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Date(2026, time.July, 25, 17, 30, 0, 0, time.UTC)
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
UPDATE accounts
|
||||
SET extra = extra - 'upstream_billing_probe_enabled' - 'upstream_billing_probe'
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 25 due accounts; the LOWEST IDs carry the LATEST due times, so a
|
||||
// correct query must pick the 20 highest-ID rows here.
|
||||
ids := make([]int64, 0, 25)
|
||||
for i := 0; i < 25; i++ {
|
||||
due := now.Add(-time.Duration(i+1) * time.Minute)
|
||||
ids = append(ids, insertUpstreamBillingProbeAccount(ctx, t, tx,
|
||||
fmt.Sprintf("probe-nano-order-%02d", i),
|
||||
rfc3339WithFraction(due, "123456789")))
|
||||
}
|
||||
|
||||
accounts, err := repo.ListDueUpstreamBillingProbeAccounts(ctx, now, 20)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, accounts, 20)
|
||||
got := make([]int64, 0, len(accounts))
|
||||
for _, account := range accounts {
|
||||
got = append(got, account.ID)
|
||||
}
|
||||
// Earliest due first == highest index first; the five most recently due
|
||||
// (lowest index / lowest IDs) fall outside the limit this cycle.
|
||||
want := make([]int64, 0, 20)
|
||||
for i := 24; i >= 5; i-- {
|
||||
want = append(want, ids[i])
|
||||
}
|
||||
require.Equal(t, want, got)
|
||||
}
|
||||
|
||||
@@ -32,6 +32,7 @@ func TestAccountRepositoryListDueUpstreamBillingProbeAccountsBoundsQuery(t *test
|
||||
require.Contains(t, normalized, "type = 'apikey'")
|
||||
require.Contains(t, normalized, `extra @> '{"upstream_billing_probe_enabled": true}'::jsonb`)
|
||||
require.Contains(t, normalized, "jsonb_path_query_first_tz")
|
||||
require.Contains(t, normalized, `'(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$'`)
|
||||
require.Contains(t, normalized, "parsed AS MATERIALIZED")
|
||||
require.Contains(t, normalized, "parsed_next_probe_at::timestamptz <= $1")
|
||||
require.Contains(t, normalized, "LIMIT $2")
|
||||
|
||||
@@ -199,6 +199,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
|
||||
group.FieldMcpXMLInject,
|
||||
group.FieldSupportedModelScopes,
|
||||
group.FieldAllowMessagesDispatch,
|
||||
group.FieldAllowLive,
|
||||
group.FieldDefaultMappedModel,
|
||||
group.FieldMessagesDispatchModelConfig,
|
||||
group.FieldModelsListConfig,
|
||||
@@ -948,6 +949,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
|
||||
SupportedModelScopes: g.SupportedModelScopes,
|
||||
SortOrder: g.SortOrder,
|
||||
AllowMessagesDispatch: g.AllowMessagesDispatch,
|
||||
AllowLive: g.AllowLive,
|
||||
RequireOAuthOnly: g.RequireOauthOnly,
|
||||
RequirePrivacySet: g.RequirePrivacySet,
|
||||
DefaultMappedModel: g.DefaultMappedModel,
|
||||
|
||||
@@ -750,7 +750,7 @@ INSERT INTO batch_image_jobs (
|
||||
batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price,
|
||||
pricing_snapshot_version,
|
||||
currency, hold_id,
|
||||
idempotency_key, request_hash, manifest_hash, retry_count, output_expires_at
|
||||
idempotency_key, request_hash, manifest_hash, retry_count, session_id, output_expires_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9,
|
||||
$10, $11, $12, $13, $14,
|
||||
@@ -760,7 +760,7 @@ INSERT INTO batch_image_jobs (
|
||||
$25, $26, $27, $28,
|
||||
$29,
|
||||
$30, $31,
|
||||
$32, $33, $34, $35, $36
|
||||
$32, $33, $34, $35, $36, $37
|
||||
)
|
||||
RETURNING `+batchImageJobColumns,
|
||||
params.BatchID, params.UserID, params.APIKeyID, params.AccountID, params.Provider, params.Model, params.TaskName, params.ParentBatchID, params.Status,
|
||||
@@ -771,7 +771,7 @@ RETURNING `+batchImageJobColumns,
|
||||
params.BatchDiscountMultiplier, params.HoldMultiplier, params.BillableUnitPrice, params.HoldUnitPrice,
|
||||
params.PricingSnapshotVersion,
|
||||
params.Currency, params.HoldID,
|
||||
params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.OutputExpiresAt,
|
||||
params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.SessionID, params.OutputExpiresAt,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -825,7 +825,7 @@ batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price
|
||||
pricing_snapshot_version,
|
||||
currency, hold_id,
|
||||
idempotency_key, request_hash, manifest_hash,
|
||||
retry_count, version, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at,
|
||||
retry_count, version, session_id, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at,
|
||||
last_error_code, last_error_message,
|
||||
created_at, updated_at, submitted_at, started_at, finished_at, settled_at`
|
||||
|
||||
@@ -838,6 +838,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) {
|
||||
var parentBatchID sql.NullString
|
||||
var holdAmount, actualCost sql.NullFloat64
|
||||
var holdID, idempotencyKey, requestHash, manifestHash sql.NullString
|
||||
var sessionID sql.NullString
|
||||
var outputExpiresAt, inputDeletedAt, outputDeletedAt, downloadedAt, userDeletedAt sql.NullTime
|
||||
var lastErrorCode, lastErrorMessage sql.NullString
|
||||
var submittedAt, startedAt, finishedAt, settledAt sql.NullTime
|
||||
@@ -852,7 +853,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) {
|
||||
&job.PricingSnapshotVersion,
|
||||
&job.Currency, &holdID,
|
||||
&idempotencyKey, &requestHash, &manifestHash,
|
||||
&job.RetryCount, &job.Version, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt,
|
||||
&job.RetryCount, &job.Version, &sessionID, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt,
|
||||
&lastErrorCode, &lastErrorMessage,
|
||||
&job.CreatedAt, &job.UpdatedAt, &submittedAt, &startedAt, &finishedAt, &settledAt,
|
||||
)
|
||||
@@ -874,6 +875,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) {
|
||||
job.IdempotencyKey = batchImageNullStringPtr(idempotencyKey)
|
||||
job.RequestHash = batchImageNullStringPtr(requestHash)
|
||||
job.ManifestHash = batchImageNullStringPtr(manifestHash)
|
||||
job.SessionID = batchImageNullStringPtr(sessionID)
|
||||
job.OutputExpiresAt = batchImageNullTimePtr(outputExpiresAt)
|
||||
job.InputDeletedAt = batchImageNullTimePtr(inputDeletedAt)
|
||||
job.OutputDeletedAt = batchImageNullTimePtr(outputDeletedAt)
|
||||
|
||||
@@ -29,11 +29,15 @@ const (
|
||||
// 格式: concurrency:user:{userID}
|
||||
userSlotKeyPrefix = "concurrency:user:"
|
||||
// 格式: concurrency:api_key:{apiKeyID}
|
||||
apiKeySlotKeyPrefix = "concurrency:api_key:"
|
||||
apiKeySlotKeyPrefix = "concurrency:api_key:"
|
||||
liveAccountSlotKeyPrefix = "concurrency:live:account:"
|
||||
liveUserSlotKeyPrefix = "concurrency:live:user:"
|
||||
liveAPIKeySlotKeyPrefix = "concurrency:live:api_key:"
|
||||
// API-key-scoped client WebSocket ingress leases use a shorter TTL than
|
||||
// ordinary request slots, because idle ingress sessions do not hold a turn slot.
|
||||
openAIWSIngressLeaseKeyPrefix = "concurrency:openai_ws_ingress:api_key:"
|
||||
openAIWSIngressLeaseTTLSeconds = 60
|
||||
liveLeaseTTLSeconds = 60
|
||||
// 等待队列计数器格式: concurrency:wait:{userID}
|
||||
waitQueueKeyPrefix = "concurrency:wait:"
|
||||
// 账号级等待队列计数器格式: wait:account:{accountID}
|
||||
@@ -59,7 +63,7 @@ const (
|
||||
var (
|
||||
// acquireScript 使用有序集合计数并在未达上限时添加槽位
|
||||
// 使用 Redis TIME 命令获取服务器时间,避免多实例时钟不同步问题
|
||||
// KEYS[1] = 有序集合键 (concurrency:account:{id} / concurrency:user:{id})
|
||||
// KEYS[1] = 普通槽位键,KEYS[2] = 对应 Live 槽位键
|
||||
// ARGV[1] = maxConcurrency
|
||||
// ARGV[2] = TTL(秒)
|
||||
// ARGV[3] = requestID
|
||||
@@ -69,6 +73,7 @@ var (
|
||||
-- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
|
||||
redis.replicate_commands()
|
||||
local key = KEYS[1]
|
||||
local liveKey = KEYS[2]
|
||||
local maxConcurrency = tonumber(ARGV[1])
|
||||
local ttl = tonumber(ARGV[2])
|
||||
local requestID = ARGV[3]
|
||||
@@ -80,6 +85,7 @@ var (
|
||||
|
||||
-- 清理过期槽位
|
||||
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
|
||||
redis.call('ZREMRANGEBYSCORE', liveKey, '-inf', now - 60)
|
||||
|
||||
-- 检查是否已存在(支持重试场景刷新时间戳)
|
||||
local exists = redis.call('ZSCORE', key, requestID)
|
||||
@@ -90,7 +96,7 @@ var (
|
||||
end
|
||||
|
||||
-- 检查是否达到并发上限
|
||||
local count = redis.call('ZCARD', key)
|
||||
local count = redis.call('ZCARD', key) + redis.call('ZCARD', liveKey)
|
||||
if count < maxConcurrency then
|
||||
redis.call('ZADD', key, now, requestID)
|
||||
redis.call('EXPIRE', key, ttl)
|
||||
@@ -102,13 +108,14 @@ var (
|
||||
|
||||
// getCountScript 统计有序集合中的槽位数量并清理过期条目
|
||||
// 使用 Redis TIME 命令获取服务器时间
|
||||
// KEYS[1] = 有序集合键
|
||||
// KEYS[1] = 普通槽位键,KEYS[2] = 对应 Live 槽位键
|
||||
// ARGV[1] = TTL(秒)
|
||||
getCountScript = redis.NewScript(`
|
||||
-- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
|
||||
-- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
|
||||
redis.replicate_commands()
|
||||
local key = KEYS[1]
|
||||
local liveKey = KEYS[2]
|
||||
local ttl = tonumber(ARGV[1])
|
||||
|
||||
-- 使用 Redis 服务器时间
|
||||
@@ -117,7 +124,60 @@ var (
|
||||
local expireBefore = now - ttl
|
||||
|
||||
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
|
||||
return redis.call('ZCARD', key)
|
||||
redis.call('ZREMRANGEBYSCORE', liveKey, '-inf', now - 60)
|
||||
return redis.call('ZCARD', key) + redis.call('ZCARD', liveKey)
|
||||
`)
|
||||
|
||||
acquireLiveLeaseScript = redis.NewScript(`
|
||||
redis.replicate_commands()
|
||||
local accountRegular = KEYS[1]
|
||||
local accountLive = KEYS[2]
|
||||
local userRegular = KEYS[3]
|
||||
local userLive = KEYS[4]
|
||||
local apiLive = KEYS[5]
|
||||
local accountMax = tonumber(ARGV[1])
|
||||
local userMax = tonumber(ARGV[2])
|
||||
local ttl = tonumber(ARGV[3])
|
||||
local leaseID = ARGV[4]
|
||||
local replacing = tonumber(ARGV[5])
|
||||
local now = tonumber(redis.call('TIME')[1])
|
||||
local liveExpireBefore = now - ttl
|
||||
redis.call('ZREMRANGEBYSCORE', accountLive, '-inf', liveExpireBefore)
|
||||
redis.call('ZREMRANGEBYSCORE', userLive, '-inf', liveExpireBefore)
|
||||
redis.call('ZREMRANGEBYSCORE', apiLive, '-inf', liveExpireBefore)
|
||||
if redis.call('ZSCORE', accountLive, leaseID) ~= false then
|
||||
return 1
|
||||
end
|
||||
local accountCount = redis.call('ZCARD', accountRegular) + redis.call('ZCARD', accountLive)
|
||||
local userCount = redis.call('ZCARD', userRegular) + redis.call('ZCARD', userLive)
|
||||
local allowance = 0
|
||||
if replacing == 1 then allowance = 1 end
|
||||
if accountMax > 0 and accountCount >= accountMax + allowance then return 0 end
|
||||
if userMax > 0 and userCount >= userMax + allowance then return 0 end
|
||||
redis.call('ZADD', accountLive, now, leaseID)
|
||||
redis.call('ZADD', userLive, now, leaseID)
|
||||
redis.call('ZADD', apiLive, now, leaseID)
|
||||
redis.call('EXPIRE', accountLive, ttl)
|
||||
redis.call('EXPIRE', userLive, ttl)
|
||||
redis.call('EXPIRE', apiLive, ttl)
|
||||
return 1
|
||||
`)
|
||||
|
||||
refreshLiveLeaseScript = redis.NewScript(`
|
||||
redis.replicate_commands()
|
||||
local ttl = tonumber(ARGV[1])
|
||||
local leaseID = ARGV[2]
|
||||
local now = tonumber(redis.call('TIME')[1])
|
||||
local expireBefore = now - ttl
|
||||
for _, key in ipairs(KEYS) do
|
||||
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
|
||||
if redis.call('ZSCORE', key, leaseID) == false then return 0 end
|
||||
end
|
||||
for _, key in ipairs(KEYS) do
|
||||
redis.call('ZADD', key, now, leaseID)
|
||||
redis.call('EXPIRE', key, ttl)
|
||||
end
|
||||
return 1
|
||||
`)
|
||||
|
||||
// trackSlotScript 记录 stats-only 槽位,不做并发上限判断。
|
||||
@@ -330,6 +390,18 @@ func apiKeySlotKey(apiKeyID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
|
||||
}
|
||||
|
||||
func liveAccountSlotKey(accountID int64) string {
|
||||
return fmt.Sprintf("%s%d", liveAccountSlotKeyPrefix, accountID)
|
||||
}
|
||||
|
||||
func liveUserSlotKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", liveUserSlotKeyPrefix, userID)
|
||||
}
|
||||
|
||||
func liveAPIKeySlotKey(apiKeyID int64) string {
|
||||
return fmt.Sprintf("%s%d", liveAPIKeySlotKeyPrefix, apiKeyID)
|
||||
}
|
||||
|
||||
func openAIWSIngressLeaseKey(apiKeyID int64) string {
|
||||
return fmt.Sprintf("%s%d", openAIWSIngressLeaseKeyPrefix, apiKeyID)
|
||||
}
|
||||
@@ -559,7 +631,7 @@ func runScriptInt64Pair(ctx context.Context, rdb *redis.Client, script *redis.Sc
|
||||
func (c *concurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
|
||||
key := accountSlotKey(accountID)
|
||||
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致
|
||||
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID)
|
||||
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key, liveAccountSlotKey(accountID)}, maxConcurrency, c.slotTTLSeconds, requestID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
@@ -583,7 +655,7 @@ func (c *concurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int
|
||||
func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID int64) (int, error) {
|
||||
key := accountSlotKey(accountID)
|
||||
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取
|
||||
result, err := getCountScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Int()
|
||||
result, err := getCountScript.Run(ctx, c.rdb, []string{key, liveAccountSlotKey(accountID)}, c.slotTTLSeconds).Int()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
@@ -605,14 +677,18 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou
|
||||
type accountCmd struct {
|
||||
accountID int64
|
||||
zcardCmd *redis.IntCmd
|
||||
liveCmd *redis.IntCmd
|
||||
}
|
||||
cmds := make([]accountCmd, 0, len(accountIDs))
|
||||
for _, accountID := range accountIDs {
|
||||
slotKey := accountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
|
||||
liveKey := liveAccountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
|
||||
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
|
||||
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
|
||||
cmds = append(cmds, accountCmd{
|
||||
accountID: accountID,
|
||||
zcardCmd: pipe.ZCard(ctx, slotKey),
|
||||
liveCmd: pipe.ZCard(ctx, liveKey),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -622,7 +698,7 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou
|
||||
|
||||
result := make(map[int64]int, len(accountIDs))
|
||||
for _, cmd := range cmds {
|
||||
result[cmd.accountID] = int(cmd.zcardCmd.Val())
|
||||
result[cmd.accountID] = int(cmd.zcardCmd.Val() + cmd.liveCmd.Val())
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -632,7 +708,7 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou
|
||||
func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
|
||||
key := userSlotKey(userID)
|
||||
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致
|
||||
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID)
|
||||
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key, liveUserSlotKey(userID)}, maxConcurrency, c.slotTTLSeconds, requestID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
@@ -656,7 +732,7 @@ func (c *concurrencyCache) ReleaseUserSlot(ctx context.Context, userID int64, re
|
||||
func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) (int, error) {
|
||||
key := userSlotKey(userID)
|
||||
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取
|
||||
result, err := getCountScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Int()
|
||||
result, err := getCountScript.Run(ctx, c.rdb, []string{key, liveUserSlotKey(userID)}, c.slotTTLSeconds).Int()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
@@ -716,6 +792,57 @@ func (c *concurrencyCache) ReleaseOpenAIWSIngressLease(ctx context.Context, apiK
|
||||
return c.rdb.ZRem(ctx, openAIWSIngressLeaseKey(apiKeyID), leaseID).Err()
|
||||
}
|
||||
|
||||
func (c *concurrencyCache) AcquireLiveLease(
|
||||
ctx context.Context,
|
||||
accountID int64,
|
||||
accountMax int,
|
||||
userID int64,
|
||||
userMax int,
|
||||
apiKeyID int64,
|
||||
leaseID string,
|
||||
replacingRegularSlots bool,
|
||||
) (bool, error) {
|
||||
if c == nil || c.rdb == nil || accountID <= 0 || userID <= 0 || apiKeyID <= 0 || leaseID == "" {
|
||||
return false, nil
|
||||
}
|
||||
replacing := 0
|
||||
if replacingRegularSlots {
|
||||
replacing = 1
|
||||
}
|
||||
result, err := acquireLiveLeaseScript.Run(ctx, c.rdb, []string{
|
||||
accountSlotKey(accountID),
|
||||
liveAccountSlotKey(accountID),
|
||||
userSlotKey(userID),
|
||||
liveUserSlotKey(userID),
|
||||
liveAPIKeySlotKey(apiKeyID),
|
||||
}, accountMax, userMax, liveLeaseTTLSeconds, leaseID, replacing).Int()
|
||||
return result == 1, err
|
||||
}
|
||||
|
||||
func (c *concurrencyCache) RefreshLiveLease(ctx context.Context, accountID, userID, apiKeyID int64, leaseID string) (bool, error) {
|
||||
if c == nil || c.rdb == nil || leaseID == "" {
|
||||
return false, nil
|
||||
}
|
||||
result, err := refreshLiveLeaseScript.Run(ctx, c.rdb, []string{
|
||||
liveAccountSlotKey(accountID),
|
||||
liveUserSlotKey(userID),
|
||||
liveAPIKeySlotKey(apiKeyID),
|
||||
}, liveLeaseTTLSeconds, leaseID).Int()
|
||||
return result == 1, err
|
||||
}
|
||||
|
||||
func (c *concurrencyCache) ReleaseLiveLease(ctx context.Context, accountID, userID, apiKeyID int64, leaseID string) error {
|
||||
if c == nil || c.rdb == nil || leaseID == "" {
|
||||
return nil
|
||||
}
|
||||
pipe := c.rdb.TxPipeline()
|
||||
pipe.ZRem(ctx, liveAccountSlotKey(accountID), leaseID)
|
||||
pipe.ZRem(ctx, liveUserSlotKey(userID), leaseID)
|
||||
pipe.ZRem(ctx, liveAPIKeySlotKey(apiKeyID), leaseID)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) {
|
||||
if len(apiKeyIDs) == 0 {
|
||||
return map[int64]int{}, nil
|
||||
@@ -731,14 +858,18 @@ func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKey
|
||||
type apiKeyCmd struct {
|
||||
apiKeyID int64
|
||||
zcardCmd *redis.IntCmd
|
||||
liveCmd *redis.IntCmd
|
||||
}
|
||||
cmds := make([]apiKeyCmd, 0, len(apiKeyIDs))
|
||||
for _, apiKeyID := range apiKeyIDs {
|
||||
slotKey := apiKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10)
|
||||
liveKey := liveAPIKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10)
|
||||
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
|
||||
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
|
||||
cmds = append(cmds, apiKeyCmd{
|
||||
apiKeyID: apiKeyID,
|
||||
zcardCmd: pipe.ZCard(ctx, slotKey),
|
||||
liveCmd: pipe.ZCard(ctx, liveKey),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -748,7 +879,7 @@ func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKey
|
||||
|
||||
result := make(map[int64]int, len(apiKeyIDs))
|
||||
for _, cmd := range cmds {
|
||||
result[cmd.apiKeyID] = int(cmd.zcardCmd.Val())
|
||||
result[cmd.apiKeyID] = int(cmd.zcardCmd.Val() + cmd.liveCmd.Val())
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -834,17 +965,21 @@ func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []
|
||||
id int64
|
||||
maxConcurrency int
|
||||
zcardCmd *redis.IntCmd
|
||||
liveCmd *redis.IntCmd
|
||||
getCmd *redis.StringCmd
|
||||
}
|
||||
cmds := make([]accountCmds, 0, len(accounts))
|
||||
for _, acc := range accounts {
|
||||
slotKey := accountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
|
||||
liveKey := liveAccountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
|
||||
waitKey := accountWaitKeyPrefix + strconv.FormatInt(acc.ID, 10)
|
||||
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
|
||||
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
|
||||
ac := accountCmds{
|
||||
id: acc.ID,
|
||||
maxConcurrency: acc.MaxConcurrency,
|
||||
zcardCmd: pipe.ZCard(ctx, slotKey),
|
||||
liveCmd: pipe.ZCard(ctx, liveKey),
|
||||
getCmd: pipe.Get(ctx, waitKey),
|
||||
}
|
||||
cmds = append(cmds, ac)
|
||||
@@ -856,7 +991,7 @@ func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []
|
||||
|
||||
loadMap := make(map[int64]*service.AccountLoadInfo, len(accounts))
|
||||
for _, ac := range cmds {
|
||||
currentConcurrency := int(ac.zcardCmd.Val())
|
||||
currentConcurrency := int(ac.zcardCmd.Val() + ac.liveCmd.Val())
|
||||
waitingCount := 0
|
||||
if v, err := ac.getCmd.Int(); err == nil {
|
||||
waitingCount = v
|
||||
@@ -894,17 +1029,21 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic
|
||||
id int64
|
||||
maxConcurrency int
|
||||
zcardCmd *redis.IntCmd
|
||||
liveCmd *redis.IntCmd
|
||||
getCmd *redis.StringCmd
|
||||
}
|
||||
cmds := make([]userCmds, 0, len(users))
|
||||
for _, u := range users {
|
||||
slotKey := userSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
|
||||
liveKey := liveUserSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
|
||||
waitKey := waitQueueKeyPrefix + strconv.FormatInt(u.ID, 10)
|
||||
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
|
||||
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
|
||||
uc := userCmds{
|
||||
id: u.ID,
|
||||
maxConcurrency: u.MaxConcurrency,
|
||||
zcardCmd: pipe.ZCard(ctx, slotKey),
|
||||
liveCmd: pipe.ZCard(ctx, liveKey),
|
||||
getCmd: pipe.Get(ctx, waitKey),
|
||||
}
|
||||
cmds = append(cmds, uc)
|
||||
@@ -916,7 +1055,7 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic
|
||||
|
||||
loadMap := make(map[int64]*service.UserLoadInfo, len(users))
|
||||
for _, uc := range cmds {
|
||||
currentConcurrency := int(uc.zcardCmd.Val())
|
||||
currentConcurrency := int(uc.zcardCmd.Val() + uc.liveCmd.Val())
|
||||
waitingCount := 0
|
||||
if v, err := uc.getCmd.Int(); err == nil {
|
||||
waitingCount = v
|
||||
|
||||
@@ -97,6 +97,42 @@ func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_ReapsCrashedLeaseW
|
||||
require.Equal(s.T(), int64(2), count)
|
||||
}
|
||||
|
||||
func (s *ConcurrencyCacheSuite) TestLiveLease_CountsTowardRegularAccountAndUserLimits() {
|
||||
liveCache, ok := s.cache.(service.LiveConcurrencyCache)
|
||||
require.True(s.T(), ok)
|
||||
|
||||
accountID := int64(9101)
|
||||
userID := int64(9102)
|
||||
apiKeyID := int64(9103)
|
||||
acquired, err := liveCache.AcquireLiveLease(
|
||||
s.ctx,
|
||||
accountID,
|
||||
1,
|
||||
userID,
|
||||
1,
|
||||
apiKeyID,
|
||||
"live-integration",
|
||||
false,
|
||||
)
|
||||
require.NoError(s.T(), err)
|
||||
require.True(s.T(), acquired)
|
||||
|
||||
regularAccount, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 1, "regular-account")
|
||||
require.NoError(s.T(), err)
|
||||
require.False(s.T(), regularAccount)
|
||||
regularUser, err := s.cache.AcquireUserSlot(s.ctx, userID, 1, "regular-user")
|
||||
require.NoError(s.T(), err)
|
||||
require.False(s.T(), regularUser)
|
||||
|
||||
refreshed, err := liveCache.RefreshLiveLease(s.ctx, accountID, userID, apiKeyID, "live-integration")
|
||||
require.NoError(s.T(), err)
|
||||
require.True(s.T(), refreshed)
|
||||
require.NoError(s.T(), liveCache.ReleaseLiveLease(s.ctx, accountID, userID, apiKeyID, "live-integration"))
|
||||
regularAccount, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 1, "regular-account")
|
||||
require.NoError(s.T(), err)
|
||||
require.True(s.T(), regularAccount)
|
||||
}
|
||||
|
||||
func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() {
|
||||
accountID := int64(10)
|
||||
reqID1, reqID2, reqID3 := "req1", "req2", "req3"
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestLiveLeaseReplacesRegularSlotsAndCountsTowardLimits(t *testing.T) {
|
||||
redisServer := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
|
||||
regular := NewConcurrencyCache(client, 15, 900)
|
||||
live, ok := regular.(service.LiveConcurrencyCache)
|
||||
require.True(t, ok)
|
||||
ctx := context.Background()
|
||||
|
||||
accountAcquired, err := regular.AcquireAccountSlot(ctx, 10, 1, "regular-account")
|
||||
require.NoError(t, err)
|
||||
require.True(t, accountAcquired)
|
||||
userAcquired, err := regular.AcquireUserSlot(ctx, 20, 1, "regular-user")
|
||||
require.NoError(t, err)
|
||||
require.True(t, userAcquired)
|
||||
|
||||
acquired, err := live.AcquireLiveLease(ctx, 10, 1, 20, 1, 30, "live-lease", true)
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
require.NoError(t, regular.ReleaseAccountSlot(ctx, 10, "regular-account"))
|
||||
require.NoError(t, regular.ReleaseUserSlot(ctx, 20, "regular-user"))
|
||||
|
||||
accountCount, err := regular.GetAccountConcurrency(ctx, 10)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, accountCount)
|
||||
userCount, err := regular.GetUserConcurrency(ctx, 20)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, userCount)
|
||||
accountAcquired, err = regular.AcquireAccountSlot(ctx, 10, 1, "ordinary-blocked")
|
||||
require.NoError(t, err)
|
||||
require.False(t, accountAcquired)
|
||||
|
||||
refreshed, err := live.RefreshLiveLease(ctx, 10, 20, 30, "live-lease")
|
||||
require.NoError(t, err)
|
||||
require.True(t, refreshed)
|
||||
require.NoError(t, live.ReleaseLiveLease(ctx, 10, 20, 30, "live-lease"))
|
||||
accountAcquired, err = regular.AcquireAccountSlot(ctx, 10, 1, "ordinary-allowed")
|
||||
require.NoError(t, err)
|
||||
require.True(t, accountAcquired)
|
||||
}
|
||||
|
||||
func TestLiveLeaseExpiresWithoutRefresh(t *testing.T) {
|
||||
redisServer := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
|
||||
regular := NewConcurrencyCache(client, 15, 900)
|
||||
live, ok := regular.(service.LiveConcurrencyCache)
|
||||
require.True(t, ok)
|
||||
ctx := context.Background()
|
||||
|
||||
acquired, err := live.AcquireLiveLease(ctx, 10, 1, 20, 1, 30, "expired-live", false)
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
|
||||
redisServer.FastForward(61 * time.Second)
|
||||
acquired, err = regular.AcquireAccountSlot(ctx, 10, 1, "ordinary-after-expiry")
|
||||
require.NoError(t, err)
|
||||
require.True(t, acquired)
|
||||
refreshed, err := live.RefreshLiveLease(ctx, 10, 20, 30, "expired-live")
|
||||
require.NoError(t, err)
|
||||
require.False(t, refreshed)
|
||||
}
|
||||
@@ -2,7 +2,10 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -10,6 +13,7 @@ import (
|
||||
)
|
||||
|
||||
const stickySessionPrefix = "sticky_session:"
|
||||
const liveCallPrefix = "live:call:"
|
||||
|
||||
type gatewayCache struct {
|
||||
rdb *redis.Client
|
||||
@@ -54,6 +58,7 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64
|
||||
|
||||
// Compile-time assertion: gatewayCache must implement CyberSessionBlockStore.
|
||||
var _ service.CyberSessionBlockStore = (*gatewayCache)(nil)
|
||||
var _ service.LiveCallStore = (*gatewayCache)(nil)
|
||||
|
||||
const cyberSessionBlockPrefix = "cyber_session_block:"
|
||||
|
||||
@@ -71,3 +76,142 @@ func (c *gatewayCache) IsCyberSessionBlocked(ctx context.Context, key string) (b
|
||||
}
|
||||
return n > 0, nil
|
||||
}
|
||||
|
||||
var claimLiveControllerScript = redis.NewScript(`
|
||||
local key = KEYS[1]
|
||||
local target = ARGV[1]
|
||||
local owner = ARGV[2]
|
||||
local current = redis.call('HGET', key, 'controller')
|
||||
if current == false or current == 'closed' then
|
||||
return 0
|
||||
end
|
||||
if target == 'observer' and current ~= 'pending' then
|
||||
return 0
|
||||
end
|
||||
if target == 'proxy' and current ~= 'pending' and current ~= 'observer' and
|
||||
(current ~= 'proxy' or redis.call('HGET', key, 'controller_owner') ~= owner) then
|
||||
return 0
|
||||
end
|
||||
redis.call('HSET', key, 'controller', target, 'controller_owner', owner)
|
||||
return 1
|
||||
`)
|
||||
|
||||
var markLiveCallClosedScript = redis.NewScript(`
|
||||
local key = KEYS[1]
|
||||
if redis.call('EXISTS', key) == 0 then
|
||||
return 0
|
||||
end
|
||||
if redis.call('HGET', key, 'controller') == 'closed' then
|
||||
return 0
|
||||
end
|
||||
redis.call('HSET', key, 'controller', 'closed', 'controller_owner', '')
|
||||
redis.call('EXPIRE', key, ARGV[1])
|
||||
return 1
|
||||
`)
|
||||
|
||||
var releaseLiveControllerScript = redis.NewScript(`
|
||||
local key = KEYS[1]
|
||||
if redis.call('HGET', key, 'controller') ~= 'proxy' or
|
||||
redis.call('HGET', key, 'controller_owner') ~= ARGV[1] then
|
||||
return 0
|
||||
end
|
||||
redis.call('HSET', key, 'controller', 'pending', 'controller_owner', '')
|
||||
return 1
|
||||
`)
|
||||
|
||||
func liveCallKey(callHash string) string {
|
||||
return liveCallPrefix + callHash
|
||||
}
|
||||
|
||||
func HashLiveCallID(callID string) string {
|
||||
sum := sha256.Sum256([]byte(callID))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func (c *gatewayCache) SaveLiveCall(ctx context.Context, record *service.LiveCallRecord, ttl time.Duration) error {
|
||||
if record == nil || record.CallHash == "" || record.CallID == "" {
|
||||
return fmt.Errorf("invalid live call record")
|
||||
}
|
||||
values := map[string]any{
|
||||
"call_id": record.CallID,
|
||||
"account_id": record.AccountID,
|
||||
"api_key_id": record.APIKeyID,
|
||||
"user_id": record.UserID,
|
||||
"group_id": record.GroupID,
|
||||
"subscription_id": record.SubscriptionID,
|
||||
"lease_id": record.LeaseID,
|
||||
"model": record.Model,
|
||||
"created_at": record.CreatedAt.UnixMilli(),
|
||||
"expires_at": record.ExpiresAt.UnixMilli(),
|
||||
"controller": record.Controller,
|
||||
"controller_owner": record.ControllerOwner,
|
||||
"user_agent": record.UserAgent,
|
||||
"ip_address": record.IPAddress,
|
||||
"inbound_endpoint": record.InboundEndpoint,
|
||||
"attestation": record.AttestationCiphertext,
|
||||
}
|
||||
key := liveCallKey(record.CallHash)
|
||||
pipe := c.rdb.TxPipeline()
|
||||
pipe.HSet(ctx, key, values)
|
||||
pipe.Expire(ctx, key, ttl)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *gatewayCache) GetLiveCall(ctx context.Context, callHash string) (*service.LiveCallRecord, error) {
|
||||
values, err := c.rdb.HGetAll(ctx, liveCallKey(callHash)).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(values) == 0 {
|
||||
return nil, service.ErrLiveCallNotFound
|
||||
}
|
||||
parseInt := func(field string) int64 {
|
||||
value, _ := strconv.ParseInt(values[field], 10, 64)
|
||||
return value
|
||||
}
|
||||
createdAt := time.UnixMilli(parseInt("created_at"))
|
||||
expiresAt := time.UnixMilli(parseInt("expires_at"))
|
||||
return &service.LiveCallRecord{
|
||||
CallID: values["call_id"],
|
||||
CallHash: callHash,
|
||||
AccountID: parseInt("account_id"),
|
||||
APIKeyID: parseInt("api_key_id"),
|
||||
UserID: parseInt("user_id"),
|
||||
GroupID: parseInt("group_id"),
|
||||
SubscriptionID: parseInt("subscription_id"),
|
||||
LeaseID: values["lease_id"],
|
||||
Model: values["model"],
|
||||
CreatedAt: createdAt,
|
||||
ExpiresAt: expiresAt,
|
||||
Controller: values["controller"],
|
||||
ControllerOwner: values["controller_owner"],
|
||||
UserAgent: values["user_agent"],
|
||||
IPAddress: values["ip_address"],
|
||||
InboundEndpoint: values["inbound_endpoint"],
|
||||
AttestationCiphertext: values["attestation"],
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *gatewayCache) ClaimLiveController(ctx context.Context, callHash, controller, owner string) (bool, error) {
|
||||
result, err := claimLiveControllerScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, controller, owner).Int()
|
||||
return result == 1, err
|
||||
}
|
||||
|
||||
func (c *gatewayCache) GetLiveController(ctx context.Context, callHash string) (string, error) {
|
||||
value, err := c.rdb.HGet(ctx, liveCallKey(callHash), "controller").Result()
|
||||
if err == redis.Nil {
|
||||
return "", service.ErrLiveCallNotFound
|
||||
}
|
||||
return value, err
|
||||
}
|
||||
|
||||
func (c *gatewayCache) ReleaseLiveController(ctx context.Context, callHash, owner string) (bool, error) {
|
||||
result, err := releaseLiveControllerScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, owner).Int()
|
||||
return result == 1, err
|
||||
}
|
||||
|
||||
func (c *gatewayCache) MarkLiveCallClosed(ctx context.Context, callHash string, ttl time.Duration) (bool, error) {
|
||||
result, err := markLiveCallClosedScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, int64(ttl.Seconds())).Int()
|
||||
return result == 1, err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGatewayCacheLiveCallIdentityAndController(t *testing.T) {
|
||||
redisServer := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
|
||||
cache, ok := NewGatewayCache(client).(service.LiveCallStore)
|
||||
require.True(t, ok)
|
||||
otherInstance, ok := NewGatewayCache(client).(service.LiveCallStore)
|
||||
require.True(t, ok)
|
||||
record := &service.LiveCallRecord{
|
||||
CallID: "call_secret",
|
||||
CallHash: HashLiveCallID("call_secret"),
|
||||
AccountID: 11,
|
||||
APIKeyID: 22,
|
||||
UserID: 33,
|
||||
GroupID: 44,
|
||||
LeaseID: "lease",
|
||||
Model: "gpt-live-test",
|
||||
AttestationCiphertext: "encrypted-attestation",
|
||||
CreatedAt: time.Now(),
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
Controller: service.LiveControllerPending,
|
||||
}
|
||||
require.NoError(t, cache.SaveLiveCall(context.Background(), record, time.Hour))
|
||||
|
||||
loaded, err := otherInstance.GetLiveCall(context.Background(), record.CallHash)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, record.CallID, loaded.CallID)
|
||||
require.Equal(t, record.AccountID, loaded.AccountID)
|
||||
require.Equal(t, record.AttestationCiphertext, loaded.AttestationCiphertext)
|
||||
|
||||
claimed, err := cache.ClaimLiveController(context.Background(), record.CallHash, service.LiveControllerObserver, "observer-1")
|
||||
require.NoError(t, err)
|
||||
require.True(t, claimed)
|
||||
claimed, err = cache.ClaimLiveController(context.Background(), record.CallHash, service.LiveControllerProxy, "proxy-1")
|
||||
require.NoError(t, err)
|
||||
require.True(t, claimed)
|
||||
controller, err := cache.GetLiveController(context.Background(), record.CallHash)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, service.LiveControllerProxy, controller)
|
||||
|
||||
released, err := cache.ReleaseLiveController(context.Background(), record.CallHash, "proxy-1")
|
||||
require.NoError(t, err)
|
||||
require.True(t, released)
|
||||
closed, err := cache.MarkLiveCallClosed(context.Background(), record.CallHash, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.True(t, closed)
|
||||
closed, err = cache.MarkLiveCallClosed(context.Background(), record.CallHash, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.False(t, closed)
|
||||
}
|
||||
@@ -90,6 +90,7 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
|
||||
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
|
||||
SetMcpXMLInject(groupIn.MCPXMLInject).
|
||||
SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch).
|
||||
SetAllowLive(groupIn.AllowLive).
|
||||
SetRequireOauthOnly(groupIn.RequireOAuthOnly).
|
||||
SetRequirePrivacySet(groupIn.RequirePrivacySet).
|
||||
SetDefaultMappedModel(groupIn.DefaultMappedModel).
|
||||
@@ -255,6 +256,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
|
||||
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
|
||||
SetMcpXMLInject(groupIn.MCPXMLInject).
|
||||
SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch).
|
||||
SetAllowLive(groupIn.AllowLive).
|
||||
SetRequireOauthOnly(groupIn.RequireOAuthOnly).
|
||||
SetRequirePrivacySet(groupIn.RequirePrivacySet).
|
||||
SetDefaultMappedModel(groupIn.DefaultMappedModel).
|
||||
|
||||
@@ -140,7 +140,20 @@ func dockerIsAvailable(ctx context.Context) bool {
|
||||
return cmd.Run() == nil
|
||||
}
|
||||
|
||||
// selectDockerImage resolves the container image for the harness.
|
||||
//
|
||||
// SUB2API_TEST_POSTGRES_IMAGE overrides the PostgreSQL image so the suite can be
|
||||
// run against the oldest documented-supported server, not just the newest. That
|
||||
// matters for SQL that behaves differently across major versions: jsonpath
|
||||
// .datetime() only accepts the ISO-8601 "Z" designator from PostgreSQL 17 on, so
|
||||
// a suite pinned to 18 cannot observe breakage on 14-16.
|
||||
//
|
||||
// SUB2API_TEST_POSTGRES_IMAGE=postgres:15-alpine go test -tags integration ./internal/repository/
|
||||
func selectDockerImage(ctx context.Context, preferred string) string {
|
||||
if override := strings.TrimSpace(os.Getenv("SUB2API_TEST_POSTGRES_IMAGE")); override != "" &&
|
||||
strings.HasPrefix(preferred, "postgres:") {
|
||||
return override
|
||||
}
|
||||
if dockerImageExists(ctx, preferred) {
|
||||
return preferred
|
||||
}
|
||||
|
||||
@@ -55,6 +55,9 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) {
|
||||
requireColumn(t, tx, "accounts", "session_window_status", "character varying", 20, true)
|
||||
requireIndex(t, tx, "accounts", "idx_accounts_autopause_expiry_due")
|
||||
|
||||
// groups: OpenAI Live 默认关闭,管理员显式开启后才可访问。
|
||||
requireColumn(t, tx, "groups", "allow_live", "boolean", 0, false)
|
||||
|
||||
// api_keys: key length should be 128
|
||||
requireColumn(t, tx, "api_keys", "key", "character varying", 128, false)
|
||||
|
||||
|
||||
@@ -40,8 +40,8 @@ func TestResolveModelDimensionExpression(t *testing.T) {
|
||||
want string
|
||||
}{
|
||||
{usagestats.ModelSourceRequested, "COALESCE(NULLIF(TRIM(requested_model), ''), model)"},
|
||||
{usagestats.ModelSourceUpstream, "COALESCE(NULLIF(TRIM(upstream_model), ''), COALESCE(NULLIF(TRIM(requested_model), ''), model))"},
|
||||
{usagestats.ModelSourceMapping, "(COALESCE(NULLIF(TRIM(requested_model), ''), model) || ' -> ' || COALESCE(NULLIF(TRIM(upstream_model), ''), COALESCE(NULLIF(TRIM(requested_model), ''), model)))"},
|
||||
{usagestats.ModelSourceUpstream, "COALESCE(NULLIF(TRIM(upstream_model), ''), model)"},
|
||||
{usagestats.ModelSourceMapping, "(COALESCE(NULLIF(TRIM(requested_model), ''), model) || ' -> ' || COALESCE(NULLIF(TRIM(upstream_model), ''), model))"},
|
||||
{"", "COALESCE(NULLIF(TRIM(requested_model), ''), model)"},
|
||||
{"invalid", "COALESCE(NULLIF(TRIM(requested_model), ''), model)"},
|
||||
}
|
||||
|
||||
@@ -79,6 +79,7 @@ var usageLogInsertArgTypes = [...]string{
|
||||
"text", // billing_tier
|
||||
"text", // billing_mode
|
||||
"numeric", // account_stats_cost
|
||||
"text", // session_id
|
||||
"timestamptz", // created_at
|
||||
}
|
||||
|
||||
@@ -274,6 +275,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7,
|
||||
@@ -281,7 +283,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
|
||||
$10, $11, $12, $13,
|
||||
$14, $15, $16, $17,
|
||||
$18, $19, $20, $21, $22, $23,
|
||||
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56
|
||||
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57
|
||||
)
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
RETURNING id, created_at
|
||||
@@ -728,10 +730,13 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
) AS (VALUES `)
|
||||
|
||||
args := make([]any, 0, len(keys)*56)
|
||||
// Each batch row prepends the synthetic input_index before the 57
|
||||
// usage-log column values.
|
||||
args := make([]any, 0, len(keys)*58)
|
||||
argPos := 1
|
||||
for idx, key := range keys {
|
||||
if idx > 0 {
|
||||
@@ -815,6 +820,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
)
|
||||
SELECT
|
||||
@@ -873,6 +879,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
FROM input
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
@@ -971,10 +978,11 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
) AS (VALUES `)
|
||||
|
||||
args := make([]any, 0, len(preparedList)*56)
|
||||
args := make([]any, 0, len(preparedList)*57)
|
||||
argPos := 1
|
||||
for idx, prepared := range preparedList {
|
||||
if idx > 0 {
|
||||
@@ -1055,6 +1063,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
)
|
||||
SELECT
|
||||
@@ -1113,6 +1122,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
FROM input
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
@@ -1179,6 +1189,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
|
||||
billing_tier,
|
||||
billing_mode,
|
||||
account_stats_cost,
|
||||
session_id,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7,
|
||||
@@ -1186,7 +1197,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
|
||||
$10, $11, $12, $13,
|
||||
$14, $15, $16, $17,
|
||||
$18, $19, $20, $21, $22, $23,
|
||||
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56
|
||||
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57
|
||||
)
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
`, prepared.args...)
|
||||
@@ -1227,6 +1238,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
||||
modelMappingChain := nullString(log.ModelMappingChain)
|
||||
billingTier := nullString(log.BillingTier)
|
||||
billingMode := nullString(log.BillingMode)
|
||||
sessionID := nullString(log.SessionID)
|
||||
requestedModel := strings.TrimSpace(log.RequestedModel)
|
||||
if requestedModel == "" {
|
||||
requestedModel = strings.TrimSpace(log.Model)
|
||||
@@ -1299,6 +1311,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
||||
billingTier,
|
||||
billingMode,
|
||||
log.AccountStatsCost, // account_stats_cost
|
||||
sessionID, // session_id
|
||||
createdAt,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -19,7 +19,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at"
|
||||
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, session_id, created_at"
|
||||
|
||||
func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) {
|
||||
query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1"
|
||||
@@ -116,6 +116,10 @@ func (r *usageLogRepository) ListWithFilters(ctx context.Context, params paginat
|
||||
conditions = append(conditions, fmt.Sprintf("group_id = $%d", len(args)+1))
|
||||
args = append(args, filters.GroupID)
|
||||
}
|
||||
if requestID := strings.TrimSpace(filters.RequestID); requestID != "" {
|
||||
conditions = append(conditions, fmt.Sprintf("request_id = $%d", len(args)+1))
|
||||
args = append(args, requestID)
|
||||
}
|
||||
conditions, args = appendUsageLogModelWhereCondition(conditions, args, filters.Model, filters.ModelFilterSource)
|
||||
conditions, args = appendRequestTypeOrStreamWhereCondition(conditions, args, filters.RequestType, filters.Stream)
|
||||
if filters.BillingType != nil {
|
||||
@@ -481,6 +485,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
billingTier sql.NullString
|
||||
billingMode sql.NullString
|
||||
accountStatsCost sql.NullFloat64
|
||||
sessionID sql.NullString
|
||||
createdAt time.Time
|
||||
)
|
||||
|
||||
@@ -541,6 +546,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
&billingTier,
|
||||
&billingMode,
|
||||
&accountStatsCost,
|
||||
&sessionID,
|
||||
&createdAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
@@ -661,6 +667,9 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
if accountStatsCost.Valid {
|
||||
log.AccountStatsCost = &accountStatsCost.Float64
|
||||
}
|
||||
if sessionID.Valid {
|
||||
log.SessionID = &sessionID.String
|
||||
}
|
||||
|
||||
return log, nil
|
||||
}
|
||||
|
||||
@@ -96,6 +96,7 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) {
|
||||
sqlmock.AnyArg(), // billing_tier
|
||||
sqlmock.AnyArg(), // billing_mode
|
||||
sqlmock.AnyArg(), // account_stats_cost
|
||||
sqlmock.AnyArg(), // session_id
|
||||
createdAt,
|
||||
).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(int64(99), createdAt))
|
||||
@@ -185,6 +186,7 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) {
|
||||
sqlmock.AnyArg(), // billing_tier
|
||||
sqlmock.AnyArg(), // billing_mode
|
||||
sqlmock.AnyArg(), // account_stats_cost
|
||||
sqlmock.AnyArg(), // session_id
|
||||
createdAt,
|
||||
).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(int64(100), createdAt))
|
||||
@@ -372,6 +374,23 @@ func TestUsageLogRepositoryListWithFiltersRequestTypePriority(t *testing.T) {
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageLogRepositoryListWithFiltersRequestID(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageLogRepository{sql: db}
|
||||
|
||||
filters := usagestats.UsageLogFilters{RequestID: " req-0123 "}
|
||||
|
||||
mock.ExpectQuery("SELECT .* FROM usage_logs WHERE request_id = \\$1 ORDER BY id DESC LIMIT \\$2 OFFSET \\$3").
|
||||
WithArgs("req-0123", 21, 0).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
|
||||
logs, page, err := repo.ListWithFilters(context.Background(), pagination.PaginationParams{Page: 1, PageSize: 20}, filters)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, logs)
|
||||
require.NotNil(t, page)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageLogRepositoryListWithFiltersRequestedModelSource(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageLogRepository{sql: db}
|
||||
@@ -826,6 +845,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
sql.NullString{},
|
||||
sql.NullString{},
|
||||
sql.NullFloat64{},
|
||||
sql.NullString{},
|
||||
now,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
@@ -900,6 +920,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
sql.NullString{}, // billing_tier
|
||||
sql.NullString{}, // billing_mode
|
||||
sql.NullFloat64{}, // account_stats_cost
|
||||
sql.NullString{}, // session_id
|
||||
now,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
@@ -957,6 +978,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
sql.NullString{}, // billing_tier
|
||||
sql.NullString{}, // billing_mode
|
||||
sql.NullFloat64{}, // account_stats_cost
|
||||
sql.NullString{}, // session_id
|
||||
now,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
@@ -1014,6 +1036,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
sql.NullString{}, // billing_tier
|
||||
sql.NullString{}, // billing_mode
|
||||
sql.NullFloat64{}, // account_stats_cost
|
||||
sql.NullString{}, // session_id
|
||||
now,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -750,11 +750,12 @@ func resolveModelDimensionExpressionWithAlias(modelType, alias string) string {
|
||||
return alias + "." + name
|
||||
}
|
||||
requestedExpr := fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("requested_model"), column("model"))
|
||||
upstreamExpr := fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("upstream_model"), column("model"))
|
||||
switch usagestats.NormalizeModelSource(modelType) {
|
||||
case usagestats.ModelSourceUpstream:
|
||||
return fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("upstream_model"), requestedExpr)
|
||||
return upstreamExpr
|
||||
case usagestats.ModelSourceMapping:
|
||||
return fmt.Sprintf("(%s || ' -> ' || COALESCE(NULLIF(TRIM(%s), ''), %s))", requestedExpr, column("upstream_model"), requestedExpr)
|
||||
return fmt.Sprintf("(%s || ' -> ' || %s)", requestedExpr, upstreamExpr)
|
||||
default:
|
||||
return requestedExpr
|
||||
}
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestUsageLog_SessionIDPersistence proves session_id round-trips from insert to
|
||||
// read and is omitted (NULL) when absent.
|
||||
func TestUsageLog_SessionIDPersistence(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := newUsageLogRepositoryWithSQL(client, integrationDB)
|
||||
|
||||
user := mustCreateUser(t, client, &service.User{Email: "session-id-" + uuid.NewString() + "@example.com"})
|
||||
apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-session-" + uuid.NewString(), Name: "k"})
|
||||
account := mustCreateAccount(t, client, &service.Account{Name: "acc-session-" + uuid.NewString()})
|
||||
|
||||
sessionID := "sess-" + uuid.NewString()
|
||||
|
||||
withSession := &service.UsageLog{
|
||||
UserID: user.ID,
|
||||
APIKeyID: apiKey.ID,
|
||||
AccountID: account.ID,
|
||||
RequestID: uuid.NewString(),
|
||||
Model: "claude-3",
|
||||
InputTokens: 10,
|
||||
OutputTokens: 5,
|
||||
TotalCost: 1.0,
|
||||
ActualCost: 1.0,
|
||||
SessionID: &sessionID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
_, err := repo.Create(ctx, withSession)
|
||||
require.NoError(t, err)
|
||||
require.NotZero(t, withSession.ID)
|
||||
|
||||
withoutSession := &service.UsageLog{
|
||||
UserID: user.ID,
|
||||
APIKeyID: apiKey.ID,
|
||||
AccountID: account.ID,
|
||||
RequestID: uuid.NewString(),
|
||||
Model: "claude-3",
|
||||
InputTokens: 7,
|
||||
OutputTokens: 3,
|
||||
TotalCost: 0.5,
|
||||
ActualCost: 0.5,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
_, err = repo.Create(ctx, withoutSession)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Round-trip: session id survives insert → read.
|
||||
got, err := repo.GetByID(ctx, withSession.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got.SessionID)
|
||||
require.Equal(t, sessionID, *got.SessionID)
|
||||
|
||||
// Omission: absent session id reads back as nil (NULL), not empty string.
|
||||
gotNone, err := repo.GetByID(ctx, withoutSession.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, gotNone.SessionID)
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newSessionIDUsageLog(sessionID *string) *service.UsageLog {
|
||||
return &service.UsageLog{
|
||||
UserID: 1,
|
||||
APIKeyID: 2,
|
||||
AccountID: 3,
|
||||
RequestID: "req-session-id",
|
||||
Model: "claude-3",
|
||||
InputTokens: 10,
|
||||
OutputTokens: 5,
|
||||
TotalCost: 1.0,
|
||||
ActualCost: 1.0,
|
||||
SessionID: sessionID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
}
|
||||
|
||||
// TestPrepareUsageLogInsert_SessionIDArgWiring pins the session_id column to the
|
||||
// arg slice / arg-type table so the five INSERT column lists stay in sync. session_id
|
||||
// is the penultimate arg (created_at is always last).
|
||||
func TestPrepareUsageLogInsert_SessionIDArgWiring(t *testing.T) {
|
||||
require.Len(t, usageLogInsertArgTypes, 57, "arg-type table must include session_id")
|
||||
|
||||
sessionID := "sess-persisted-123"
|
||||
prepared := prepareUsageLogInsert(newSessionIDUsageLog(&sessionID))
|
||||
|
||||
require.Len(t, prepared.args, len(usageLogInsertArgTypes),
|
||||
"prepared args must match the arg-type table length")
|
||||
|
||||
// created_at is last; session_id is the arg immediately before it.
|
||||
sessionArg := prepared.args[len(prepared.args)-2]
|
||||
ns, ok := sessionArg.(sql.NullString)
|
||||
require.True(t, ok, "session_id arg should be a sql.NullString, got %T", sessionArg)
|
||||
require.True(t, ns.Valid)
|
||||
require.Equal(t, sessionID, ns.String)
|
||||
|
||||
require.Equal(t, "text", usageLogInsertArgTypes[len(usageLogInsertArgTypes)-2],
|
||||
"session_id arg type must be text")
|
||||
}
|
||||
|
||||
// TestPrepareUsageLogInsert_SessionIDNullWhenAbsent proves an absent session id is
|
||||
// persisted as SQL NULL rather than an empty string.
|
||||
func TestPrepareUsageLogInsert_SessionIDNullWhenAbsent(t *testing.T) {
|
||||
prepared := prepareUsageLogInsert(newSessionIDUsageLog(nil))
|
||||
sessionArg := prepared.args[len(prepared.args)-2]
|
||||
ns, ok := sessionArg.(sql.NullString)
|
||||
require.True(t, ok, "session_id arg should be a sql.NullString, got %T", sessionArg)
|
||||
require.False(t, ns.Valid, "absent session id must be NULL, not empty string")
|
||||
|
||||
empty := ""
|
||||
preparedEmpty := prepareUsageLogInsert(newSessionIDUsageLog(&empty))
|
||||
nsEmpty := preparedEmpty.args[len(preparedEmpty.args)-2].(sql.NullString)
|
||||
require.False(t, nsEmpty.Valid, "empty session id must also be NULL")
|
||||
}
|
||||
|
||||
// TestUsageLogInsertQueries_IncludeSessionID guards that every generated INSERT path
|
||||
// and the SELECT column list reference session_id.
|
||||
func TestUsageLogInsertQueries_IncludeSessionID(t *testing.T) {
|
||||
require.Contains(t, usageLogSelectColumns, "session_id",
|
||||
"SELECT column list must include session_id")
|
||||
|
||||
sessionID := "sess-in-query"
|
||||
log := newSessionIDUsageLog(&sessionID)
|
||||
prepared := prepareUsageLogInsert(log)
|
||||
key := usageLogBatchKey(log.RequestID, log.APIKeyID)
|
||||
|
||||
batchQuery, batchArgs := buildUsageLogBatchInsertQuery([]string{key},
|
||||
map[string]usageLogInsertPrepared{key: prepared})
|
||||
require.Contains(t, batchQuery, "session_id")
|
||||
// Two column references (INSERT column list + SELECT ... FROM input) plus the CTE def.
|
||||
require.GreaterOrEqual(t, strings.Count(batchQuery, "session_id"), 3)
|
||||
require.Len(t, batchArgs, len(prepared.args)+1,
|
||||
"batch args include the synthetic input_index before usage-log values")
|
||||
|
||||
bestEffortQuery, bestEffortArgs := buildUsageLogBestEffortInsertQuery([]usageLogInsertPrepared{prepared})
|
||||
require.Contains(t, bestEffortQuery, "session_id")
|
||||
require.Len(t, bestEffortArgs, len(prepared.args))
|
||||
}
|
||||
@@ -43,6 +43,16 @@ func newUserRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *userRepos
|
||||
}
|
||||
|
||||
func (r *userRepository) Create(ctx context.Context, userIn *service.User) error {
|
||||
return r.create(ctx, userIn, false)
|
||||
}
|
||||
|
||||
// CreateWithEmailAliasGuard 见 service.UserRepository:在邮箱唯一性锁内复查收件箱身份,
|
||||
// 供注册路径使用。
|
||||
func (r *userRepository) CreateWithEmailAliasGuard(ctx context.Context, userIn *service.User) error {
|
||||
return r.create(ctx, userIn, true)
|
||||
}
|
||||
|
||||
func (r *userRepository) create(ctx context.Context, userIn *service.User, guardEmailAlias bool) error {
|
||||
if userIn == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -69,11 +79,16 @@ func (r *userRepository) Create(ctx context.Context, userIn *service.User) error
|
||||
}
|
||||
}
|
||||
|
||||
lockKeys := []string{normalizedEmailUniquenessLockKey(userIn.Email)}
|
||||
if guardEmailAlias {
|
||||
// 别名变体的字面量不同,唯一索引无法兜底;用收件箱身份锁把同一收件箱的并发注册串行化。
|
||||
lockKeys = append(lockKeys, emailAliasUniquenessLockKey(userIn.Email))
|
||||
}
|
||||
releaseEmailLock, err := lockRepositoryScopedKeys(
|
||||
txCtx,
|
||||
txClient,
|
||||
txAwareSQLExecutor(txCtx, r.sql, r.client),
|
||||
normalizedEmailUniquenessLockKey(userIn.Email),
|
||||
lockKeys...,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -84,6 +99,16 @@ func (r *userRepository) Create(ctx context.Context, userIn *service.User) error
|
||||
return err
|
||||
}
|
||||
|
||||
if guardEmailAlias {
|
||||
aliasExists, err := existsByEmailAliasWithClient(txCtx, txClient, userIn.Email)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if aliasExists {
|
||||
return service.ErrEmailExists
|
||||
}
|
||||
}
|
||||
|
||||
created, err := txClient.User.Create().
|
||||
SetEmail(userIn.Email).
|
||||
SetUsername(userIn.Username).
|
||||
@@ -904,6 +929,88 @@ func (r *userRepository) ExistsByEmail(ctx context.Context, email string) (bool,
|
||||
return r.client.User.Query().Where(userEmailLookupPredicate(email)).Exist(ctx)
|
||||
}
|
||||
|
||||
// emailAliasCandidateLimit 限制一次别名查重最多取回的候选行数。探针都以去点后的
|
||||
// 本地部分为前缀锚定(见 dotStrippedEmailExpr),正常收件箱的变体只有个位数;
|
||||
// 上限只是兜底,避免公开未鉴权的注册/发码端点把大表整张读进内存。
|
||||
const emailAliasCandidateLimit = 50
|
||||
|
||||
// ExistsByEmailAlias 见 service.UserRepository。软删除过滤沿用 ExistsByEmail 的默认行为。
|
||||
func (r *userRepository) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
|
||||
return existsByEmailAliasWithClient(ctx, clientFromContext(ctx, r.client), email)
|
||||
}
|
||||
|
||||
func existsByEmailAliasWithClient(ctx context.Context, client *dbent.Client, email string) (bool, error) {
|
||||
if client == nil {
|
||||
return false, nil
|
||||
}
|
||||
probes := service.EmailAliasDedupProbes(email)
|
||||
if len(probes) == 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
preds := make([]predicate.User, 0, 2*len(probes))
|
||||
for _, probe := range probes {
|
||||
preds = append(preds,
|
||||
dotStrippedEmailEQ(probe.Local+"@"+probe.Domain),
|
||||
// "+后缀"的内容未知,只能按前缀匹配。
|
||||
dotStrippedEmailLike(escapeLikeWildcards(probe.Local)+"+%@"+escapeLikeWildcards(probe.Domain)),
|
||||
)
|
||||
}
|
||||
candidates, err := client.User.Query().
|
||||
Where(dbuser.Or(preds...)).
|
||||
Limit(emailAliasCandidateLimit).
|
||||
Select(dbuser.FieldEmail).
|
||||
Strings(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
// 探针会有过度匹配(点号只在 Gmail 家族无意义),最终判定必须回到完整归一化规则。
|
||||
identity := service.NormalizeEmailForAliasDedup(email)
|
||||
for _, candidate := range candidates {
|
||||
if service.NormalizeEmailForAliasDedup(candidate) == identity {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// dotStrippedEmailExpr 渲染下面的表达式:去掉存量邮箱的大小写、首尾空白(与
|
||||
// userEmailLookupPredicate 的精确匹配口径一致,历史数据存在带空白的行)以及全部点号。
|
||||
//
|
||||
// REPLACE(LOWER(TRIM(email)), '.', '')
|
||||
//
|
||||
// 两侧都去点,因此一个域名探针即可同时覆盖 Gmail 点号变体与 FQDN 根点(user@gmail.com.)。
|
||||
// migrations/190 为同一表达式建了索引。
|
||||
func dotStrippedEmailExpr(b *entsql.Builder, s *entsql.Selector) *entsql.Builder {
|
||||
return b.WriteString("REPLACE(LOWER(TRIM(").
|
||||
Ident(s.C(dbuser.FieldEmail)).
|
||||
WriteString(")), '.', '')")
|
||||
}
|
||||
|
||||
func dotStrippedEmailEQ(value string) predicate.User {
|
||||
return predicate.User(func(s *entsql.Selector) {
|
||||
s.Where(entsql.P(func(b *entsql.Builder) {
|
||||
dotStrippedEmailExpr(b, s).WriteString(" = ").Arg(value)
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
func dotStrippedEmailLike(pattern string) predicate.User {
|
||||
return predicate.User(func(s *entsql.Selector) {
|
||||
s.Where(entsql.P(func(b *entsql.Builder) {
|
||||
dotStrippedEmailExpr(b, s).WriteString(" LIKE ").Arg(pattern).WriteString(` ESCAPE '\'`)
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
// escapeLikeWildcards 转义 LIKE 元字符:本地部分合法可含 % 与 _,不转义会扩大匹配面。
|
||||
var likeWildcardEscaper = strings.NewReplacer(`\`, `\\`, "%", `\%`, "_", `\_`)
|
||||
|
||||
func escapeLikeWildcards(value string) string {
|
||||
return likeWildcardEscaper.Replace(value)
|
||||
}
|
||||
|
||||
func ensureNormalizedEmailAvailableWithClient(ctx context.Context, client *dbent.Client, userID int64, email string) error {
|
||||
client = clientFromContext(ctx, client)
|
||||
if client == nil {
|
||||
@@ -951,6 +1058,16 @@ func normalizedEmailUniquenessLockKey(email string) string {
|
||||
return "users:normalized-email:" + normalized
|
||||
}
|
||||
|
||||
// emailAliasUniquenessLockKey 按收件箱身份(而非邮箱字面量)加锁,使同一收件箱的不同
|
||||
// 别名变体在注册时互斥。
|
||||
func emailAliasUniquenessLockKey(email string) string {
|
||||
identity := service.NormalizeEmailForAliasDedup(email)
|
||||
if identity == "" {
|
||||
return ""
|
||||
}
|
||||
return "users:email-alias-identity:" + identity
|
||||
}
|
||||
|
||||
func (r *userRepository) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
err := client.UserAllowedGroup.Create().
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func seedUserForAliasTest(t *testing.T, repo *userRepository, email string) {
|
||||
t.Helper()
|
||||
require.NoError(t, repo.Create(context.Background(), &service.User{
|
||||
Email: email,
|
||||
Username: email,
|
||||
PasswordHash: "hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
}))
|
||||
}
|
||||
|
||||
func TestUserRepositoryExistsByEmailAlias(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
stored string
|
||||
probe string
|
||||
want bool
|
||||
}{
|
||||
{"same address", "someone@gmail.com", "someone@gmail.com", true},
|
||||
{"gmail plus alias", "someone@gmail.com", "someone+bulk294@gmail.com", true},
|
||||
{"gmail dot trick", "d.axis.2026@gmail.com", "daxis2026@gmail.com", true},
|
||||
{"gmail dot trick both sides", "d.axis.2026@gmail.com", "da.xis.2026@gmail.com", true},
|
||||
{"stored plus alias found by canonical form", "someone+tag@gmail.com", "someone@gmail.com", true},
|
||||
{"googlemail is a gmail alias", "someone@googlemail.com", "some.one@gmail.com", true},
|
||||
{"fqdn root dot on probe", "d.axis.2026@gmail.com", "da.xis.2026@gmail.com.", true},
|
||||
{"fqdn root dot on stored row", "d.axis.2026@gmail.com.", "daxis2026@gmail.com", true},
|
||||
{"legacy row with spacing and case", " D.Axis.2026@Gmail.com ", "daxis2026@gmail.com", true},
|
||||
{"non-gmail plus alias", "first.last@qq.com", "first.last+tag@qq.com", true},
|
||||
{"different gmail inbox", "someone@gmail.com", "someoneelse@gmail.com", false},
|
||||
{"non-gmail dots are significant", "first.last@qq.com", "firstlast@qq.com", false},
|
||||
{"different domain", "someone@gmail.com", "someone@qq.com", false},
|
||||
{"distinct plus-prefixed locals", "+alice@gmail.com", "+bob@gmail.com", false},
|
||||
{"underscore is not a wildcard", "user_x@qq.com", "userax@qq.com", false},
|
||||
{"percent is not a wildcard", "a%b@qq.com", "axxb@qq.com", false},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
repo, _ := newUserEntRepo(t)
|
||||
seedUserForAliasTest(t, repo, tc.stored)
|
||||
|
||||
got, err := repo.ExistsByEmailAlias(context.Background(), tc.probe)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserRepositoryExistsByEmailAliasIgnoresMalformedInput(t *testing.T) {
|
||||
repo, _ := newUserEntRepo(t)
|
||||
seedUserForAliasTest(t, repo, "someone@gmail.com")
|
||||
|
||||
got, err := repo.ExistsByEmailAlias(context.Background(), "not-an-email")
|
||||
require.NoError(t, err)
|
||||
require.False(t, got)
|
||||
}
|
||||
|
||||
func TestUserRepositoryCreateWithEmailAliasGuard(t *testing.T) {
|
||||
repo, _ := newUserEntRepo(t)
|
||||
ctx := context.Background()
|
||||
seedUserForAliasTest(t, repo, "d.axis.2026@gmail.com")
|
||||
|
||||
// 注册路径:别名变体在唯一性锁内被拒绝。
|
||||
err := repo.CreateWithEmailAliasGuard(ctx, &service.User{
|
||||
Email: "da.xis.2026+free@googlemail.com",
|
||||
Username: "alias-variant",
|
||||
PasswordHash: "hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
})
|
||||
require.ErrorIs(t, err, service.ErrEmailExists)
|
||||
|
||||
// 不同收件箱仍可注册。
|
||||
require.NoError(t, repo.CreateWithEmailAliasGuard(ctx, &service.User{
|
||||
Email: "other.person@gmail.com",
|
||||
Username: "other-person",
|
||||
PasswordHash: "hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
}))
|
||||
|
||||
// 管理员建号(Create)不受别名限制。
|
||||
require.NoError(t, repo.Create(ctx, &service.User{
|
||||
Email: "daxis2026+support@gmail.com",
|
||||
Username: "admin-created",
|
||||
PasswordHash: "hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
}))
|
||||
}
|
||||
@@ -44,7 +44,7 @@ type ConfigStore interface {
|
||||
// It must stay false when blocking is not intended, even if config is
|
||||
// untrusted—otherwise default-off deployments fail closed for all traffic.
|
||||
BlockingActivationDegraded() bool
|
||||
Public() PublicConfig
|
||||
Public() (PublicConfig, error)
|
||||
Save(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error)
|
||||
RuntimeState() (expected int64, active int64, loadedAt *time.Time, loadError string)
|
||||
Encrypt(value string) (string, error)
|
||||
|
||||
@@ -194,15 +194,15 @@ func (m *ConfigManager) markUntrustedIfNoActiveSnapshot() {
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ConfigManager) Public() PublicConfig {
|
||||
func (m *ConfigManager) Public() (PublicConfig, error) {
|
||||
if m == nil {
|
||||
return PublicFromStorage(DefaultStorageConfig(), false)
|
||||
return PublicConfig{}, infraerrors.ServiceUnavailable(ErrorCodeConfigUnavailable, "提示词审计配置暂不可用")
|
||||
}
|
||||
snapshot := m.snapshot.Load()
|
||||
if snapshot == nil {
|
||||
return PublicFromStorage(DefaultStorageConfig(), false)
|
||||
return PublicConfig{}, infraerrors.ServiceUnavailable(ErrorCodeConfigUnavailable, "提示词审计配置暂不可用")
|
||||
}
|
||||
return PublicFromStorage(cloneStorageConfig(snapshot.storage), snapshot.active.RiskControlEnabled)
|
||||
return PublicFromStorage(cloneStorageConfig(snapshot.storage), snapshot.active.RiskControlEnabled), nil
|
||||
}
|
||||
|
||||
func (m *ConfigManager) Save(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
|
||||
|
||||
@@ -56,6 +56,57 @@ func TestConfigRuntimeLoadErrorIsStableBoundedAndSecretFree(t *testing.T) {
|
||||
require.LessOrEqual(t, len([]rune(message)), 160)
|
||||
}
|
||||
|
||||
func TestConfigManagerPublicRequiresSuccessfullyLoadedSnapshot(t *testing.T) {
|
||||
t.Run("absent persisted setting is legitimate default", func(t *testing.T) {
|
||||
manager := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
|
||||
SettingKeyPromptAuditConfig: "",
|
||||
SettingKeyRiskControl: "false",
|
||||
}}, nil, prefixEncryptor{})
|
||||
require.NoError(t, manager.Reload(context.Background()))
|
||||
|
||||
public, err := manager.Public()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), public.ConfigVersion)
|
||||
require.False(t, public.Enabled)
|
||||
})
|
||||
|
||||
t.Run("persisted config activation failure is unavailable", func(t *testing.T) {
|
||||
const canary = "persisted-token-canary"
|
||||
manager := NewConfigManager(nil, staticSettingRepository{values: map[string]string{
|
||||
SettingKeyPromptAuditConfig: `{"enabled":true,"config_version":9,"endpoints":[{"token_ciphertext":"` + canary + `"}]}`,
|
||||
SettingKeyRiskControl: "true",
|
||||
}}, nil, prefixEncryptor{})
|
||||
require.Error(t, manager.Reload(context.Background()))
|
||||
|
||||
public, err := manager.Public()
|
||||
require.Error(t, err)
|
||||
require.Empty(t, public)
|
||||
require.Equal(t, ErrorCodeConfigUnavailable, infraerrors.Reason(err))
|
||||
require.NotContains(t, err.Error(), canary)
|
||||
})
|
||||
|
||||
t.Run("reload failure preserves last successfully loaded snapshot", func(t *testing.T) {
|
||||
storage := DefaultStorageConfig()
|
||||
storage.ConfigVersion = 4
|
||||
storage.ChangeSummary = "trusted snapshot"
|
||||
raw, err := json.Marshal(storage)
|
||||
require.NoError(t, err)
|
||||
repository := &switchableSettingRepository{staticSettingRepository: staticSettingRepository{values: map[string]string{
|
||||
SettingKeyPromptAuditConfig: string(raw),
|
||||
SettingKeyRiskControl: "false",
|
||||
}}}
|
||||
manager := NewConfigManager(nil, repository, nil, prefixEncryptor{})
|
||||
require.NoError(t, manager.Reload(context.Background()))
|
||||
repository.loadErr = errors.New("settings unavailable")
|
||||
require.Error(t, manager.Reload(context.Background()))
|
||||
|
||||
public, err := manager.Public()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(4), public.ConfigVersion)
|
||||
require.Equal(t, "trusted snapshot", public.ChangeSummary)
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildNextStoragePreserveReplaceAndClearToken(t *testing.T) {
|
||||
manager := &ConfigManager{encryptor: prefixEncryptor{}}
|
||||
current := DefaultStorageConfig()
|
||||
@@ -137,6 +188,18 @@ func (errorSettingRepository) GetMultiple(context.Context, []string) (map[string
|
||||
return nil, errors.New("settings unavailable")
|
||||
}
|
||||
|
||||
type switchableSettingRepository struct {
|
||||
staticSettingRepository
|
||||
loadErr error
|
||||
}
|
||||
|
||||
func (r *switchableSettingRepository) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
|
||||
if r.loadErr != nil {
|
||||
return nil, r.loadErr
|
||||
}
|
||||
return r.staticSettingRepository.GetMultiple(ctx, keys)
|
||||
}
|
||||
|
||||
func TestConfigManagerStartupLoadFailureDoesNotBlockWhenBlockingNotIntended(t *testing.T) {
|
||||
// Settings unavailable and no prior blocking intent: stay ModeOff so the
|
||||
// gateway remains usable and admins can still disable/configure Prompt Audit.
|
||||
|
||||
@@ -222,7 +222,12 @@ func TestGuardEvaluatorFlagSharedDeadlineFailClosedAndContextCancel(t *testing.T
|
||||
elapsed := time.Since(started)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, 2, calls)
|
||||
require.Less(t, elapsed, 180*time.Millisecond)
|
||||
// The bound only has to prove the failover shared the first endpoint's
|
||||
// 70ms deadline instead of taking the second endpoint's own 500ms one.
|
||||
// An unshared deadline lands at ~535ms, so 350ms still fails loudly
|
||||
// while leaving room for scheduler delay on a busy CI machine. A
|
||||
// tighter bound made this test flaky, not stricter.
|
||||
require.Less(t, elapsed, 350*time.Millisecond)
|
||||
require.GreaterOrEqual(t, elapsed, 50*time.Millisecond)
|
||||
require.Equal(t, int64(1), metrics.Snapshot().Failovers)
|
||||
require.Equal(t, int64(1), metrics.Snapshot().Timeouts)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user