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 | ||
|
|
386b57bb77 | ||
|
|
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.164
|
||||
0.1.165
|
||||
|
||||
+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=
|
||||
|
||||
@@ -1640,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()
|
||||
@@ -1828,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)
|
||||
|
||||
@@ -3488,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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 {
|
||||
@@ -1519,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),
|
||||
)
|
||||
@@ -1771,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,
|
||||
@@ -1795,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 {
|
||||
@@ -1836,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
|
||||
}
|
||||
@@ -1858,11 +1947,22 @@ 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)
|
||||
@@ -1883,7 +1983,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
SessionID: sessionID,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingWS, reqModel, result.UpstreamModel),
|
||||
ChannelUsageFields: turnUsageFields,
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
reqLog.Error("openai.websocket_record_usage_failed",
|
||||
@@ -1896,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 的
|
||||
@@ -1944,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),
|
||||
@@ -2227,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
|
||||
}
|
||||
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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 时使用。
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
+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")
|
||||
|
||||
@@ -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)"},
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -374,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}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -13,7 +13,7 @@ import (
|
||||
)
|
||||
|
||||
type PromptAdminService interface {
|
||||
GetConfig() PublicConfig
|
||||
GetConfig() (PublicConfig, error)
|
||||
SaveConfig(context.Context, UpdateConfigRequest, int64) (PublicConfig, error)
|
||||
Probe(context.Context, ProbeRequest) ProbeResult
|
||||
Runtime(context.Context) RuntimeSnapshot
|
||||
@@ -31,7 +31,14 @@ func NewPromptAdminHandler(service PromptAdminService) *PromptAdminHandler {
|
||||
return &PromptAdminHandler{service: service}
|
||||
}
|
||||
|
||||
func (h *PromptAdminHandler) GetConfig(c *gin.Context) { response.Success(c, h.service.GetConfig()) }
|
||||
func (h *PromptAdminHandler) GetConfig(c *gin.Context) {
|
||||
config, err := h.service.GetConfig()
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, config)
|
||||
}
|
||||
|
||||
func (h *PromptAdminHandler) UpdateConfig(c *gin.Context) {
|
||||
var request UpdateConfigRequest
|
||||
|
||||
@@ -29,7 +29,9 @@ type fakePromptAdminService struct {
|
||||
deleteFilter func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error)
|
||||
}
|
||||
|
||||
func (s *fakePromptAdminService) GetConfig() PublicConfig { return s.config }
|
||||
func (s *fakePromptAdminService) GetConfig() (PublicConfig, error) {
|
||||
return s.config, nil
|
||||
}
|
||||
func (s *fakePromptAdminService) SaveConfig(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
|
||||
if s.save == nil {
|
||||
return PublicConfig{}, errors.New("unexpected SaveConfig call")
|
||||
@@ -156,6 +158,21 @@ func TestPromptAdminConfigRequiresVersionMapsConflictAndNeverEchoesToken(t *test
|
||||
})
|
||||
}
|
||||
|
||||
func TestPromptAdminGetConfigReturnsSecretFreeUnavailableError(t *testing.T) {
|
||||
const canary = "persisted-config-secret-canary"
|
||||
repository := &switchableSettingRepository{loadErr: errors.New("failed to load token " + canary)}
|
||||
manager := NewConfigManager(nil, repository, nil, prefixEncryptor{})
|
||||
require.Error(t, manager.Reload(context.Background()))
|
||||
service := &PromptService{config: manager}
|
||||
|
||||
response := promptAdminRequest(t, promptAdminRouter(service), http.MethodGet, "/admin/prompt-audit/config", nil)
|
||||
require.Equal(t, http.StatusServiceUnavailable, response.Code)
|
||||
require.Contains(t, response.Body.String(), ErrorCodeConfigUnavailable)
|
||||
require.NotContains(t, response.Body.String(), canary)
|
||||
require.NotContains(t, response.Body.String(), `"config_version"`)
|
||||
require.NotContains(t, response.Body.String(), `"token"`)
|
||||
}
|
||||
|
||||
func TestPromptAdminProbeSupportsTemporaryOrSavedTokenWithoutEcho(t *testing.T) {
|
||||
const canary = "probe-token-canary"
|
||||
for _, tc := range []struct {
|
||||
|
||||
@@ -166,7 +166,7 @@ func (s *PromptService) Evaluate(ctx context.Context, req Request) (*PromptDecis
|
||||
return s.evaluator.Evaluate(ctx, cfg, snapshot)
|
||||
}
|
||||
|
||||
func (s *PromptService) GetConfig() PublicConfig { return s.config.Public() }
|
||||
func (s *PromptService) GetConfig() (PublicConfig, error) { return s.config.Public() }
|
||||
|
||||
func (s *PromptService) SaveConfig(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
|
||||
return s.config.Save(ctx, req, actorID)
|
||||
|
||||
@@ -12,11 +12,12 @@ const (
|
||||
ConfigInvalidationChannel = "sub2api:prompt_guard:config:invalidate"
|
||||
PayloadKeyPrefix = "sub2api:prompt_audit:payload:"
|
||||
|
||||
ErrorCodeBlocked = "prompt_guard_blocked"
|
||||
ErrorCodeUnavailable = "prompt_guard_unavailable"
|
||||
ErrorCodeInvalidResponse = "prompt_guard_invalid_response"
|
||||
ErrorCodeConfigConflict = "prompt_audit_config_conflict"
|
||||
ErrorCodeRequiresEnabled = "prompt_guard_requires_audit_enabled"
|
||||
ErrorCodeBlocked = "prompt_guard_blocked"
|
||||
ErrorCodeUnavailable = "prompt_guard_unavailable"
|
||||
ErrorCodeInvalidResponse = "prompt_guard_invalid_response"
|
||||
ErrorCodeConfigConflict = "prompt_audit_config_conflict"
|
||||
ErrorCodeConfigUnavailable = "prompt_audit_config_unavailable"
|
||||
ErrorCodeRequiresEnabled = "prompt_guard_requires_audit_enabled"
|
||||
|
||||
DefaultGuardModel = "sileader/qwen3guard:0.6b"
|
||||
)
|
||||
|
||||
@@ -48,7 +48,7 @@ func (s *fakeConfigStore) EffectiveMode() Mode {
|
||||
return s.cfg.EffectiveMode()
|
||||
}
|
||||
func (s *fakeConfigStore) BlockingActivationDegraded() bool { return false }
|
||||
func (s *fakeConfigStore) Public() PublicConfig { return PublicConfig{} }
|
||||
func (s *fakeConfigStore) Public() (PublicConfig, error) { return PublicConfig{}, nil }
|
||||
func (s *fakeConfigStore) Save(context.Context, UpdateConfigRequest, int64) (PublicConfig, error) {
|
||||
return PublicConfig{}, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// panelRateLimitWindow 面板限流固定窗口时长(所有档位均按每分钟计数)。
|
||||
const panelRateLimitWindow = time.Minute
|
||||
|
||||
// panelRateLimitAllower 抽象底层限流原语,便于单测注入。
|
||||
type panelRateLimitAllower interface {
|
||||
Allow(ctx context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error)
|
||||
}
|
||||
|
||||
// PanelRateLimiter 面板(管理面 /api/v1)API 限流器。
|
||||
//
|
||||
// 设计要点:
|
||||
// - 认证接口按「用户 ID」维度计数:与客户端 IP 完全无关,反向代理/共享出口
|
||||
// (所有请求源 IP 坍缩为 127.0.0.1 等)不会互相误伤。
|
||||
// - 公开接口按安全客户端 IP 计数:仅统计全局单播地址,回环/内网/链路本地
|
||||
// 地址(反代内部转发地址)直接跳过,避免误拦整条反代链路的流量。
|
||||
// - 配置走进程内缓存(60s TTL),热路径零 DB 访问。
|
||||
// - Redis 异常一律 fail-open:限流是保护措施,不能反过来把面板打挂。
|
||||
type PanelRateLimiter struct {
|
||||
limiter panelRateLimitAllower
|
||||
settingService *service.SettingService
|
||||
}
|
||||
|
||||
// NewPanelRateLimiter 创建面板限流器。
|
||||
func NewPanelRateLimiter(redisClient *redis.Client, settingService *service.SettingService) *PanelRateLimiter {
|
||||
return &PanelRateLimiter{
|
||||
limiter: middleware.NewRateLimiter(redisClient),
|
||||
settingService: settingService,
|
||||
}
|
||||
}
|
||||
|
||||
// Global 认证面板接口的全局按用户限流(宽松档,覆盖所有登录后端点)。
|
||||
func (p *PanelRateLimiter) Global() gin.HandlerFunc {
|
||||
return p.userScoped("global", func(s service.PanelRateLimitSettings) int { return s.UserRPM })
|
||||
}
|
||||
|
||||
// Heavy 重查询接口的按用户限流(严格档,覆盖 usage/dashboard 等聚合统计端点)。
|
||||
// 与 Global 叠加计数:一次重查询同时消耗两档额度。
|
||||
func (p *PanelRateLimiter) Heavy() gin.HandlerFunc {
|
||||
return p.userScoped("heavy", func(s service.PanelRateLimitSettings) int { return s.HeavyRPM })
|
||||
}
|
||||
|
||||
func (p *PanelRateLimiter) userScoped(scope string, limitOf func(service.PanelRateLimitSettings) int) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if p == nil || p.limiter == nil || p.settingService == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
|
||||
if !settings.Enabled {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
limit := limitOf(settings)
|
||||
if limit <= 0 {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
subject, ok := GetAuthSubjectFromContext(c)
|
||||
if !ok || subject.UserID <= 0 {
|
||||
// 无认证主体(认证中间件缺位时的防御分支):放行,避免误伤
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if settings.ExemptAdmin {
|
||||
if role, hasRole := GetUserRoleFromContext(c); hasRole && role == service.RoleAdmin {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
key := "panel:" + scope + ":user:" + strconv.FormatInt(subject.UserID, 10)
|
||||
result, err := p.limiter.Allow(c.Request.Context(), key, limit, panelRateLimitWindow)
|
||||
if err != nil {
|
||||
// fail-open:Redis 异常不阻断面板访问
|
||||
slog.Warn("panel rate limit check failed, allowing request", "scope", scope, "error", err)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if !result.Allowed {
|
||||
abortPanelRateLimited(c, result.RetryAfter)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// PublicIP 无需认证的公开接口按客户端 IP 限流。
|
||||
// 使用与审计日志/会话绑定一致的安全客户端 IP 解析;解析结果为回环/内网/
|
||||
// 链路本地地址时跳过计数(这类地址通常是反代内部转发地址,按它计数会把
|
||||
// 整条反代链路的所有真实用户合并进同一个桶造成大面积误拦截)。
|
||||
func (p *PanelRateLimiter) PublicIP() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if p == nil || p.limiter == nil || p.settingService == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
|
||||
if !settings.Enabled || settings.PublicIPRPM <= 0 {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
clientIP := SecurityClientIP(c)
|
||||
if !isPubliclyRoutableClientIP(clientIP) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
result, err := p.limiter.Allow(c.Request.Context(), "panel:public:ip:"+clientIP, settings.PublicIPRPM, panelRateLimitWindow)
|
||||
if err != nil {
|
||||
slog.Warn("panel public rate limit check failed, allowing request", "error", err)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if !result.Allowed {
|
||||
abortPanelRateLimited(c, result.RetryAfter)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// isPubliclyRoutableClientIP 判断地址是否为可作为限流依据的全局单播地址。
|
||||
// 回环、RFC1918/ULA 内网、链路本地与未指定地址返回 false。
|
||||
func isPubliclyRoutableClientIP(clientIP string) bool {
|
||||
ip := net.ParseIP(clientIP)
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
if ip.IsLoopback() || ip.IsPrivate() || ip.IsUnspecified() ||
|
||||
ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
|
||||
return false
|
||||
}
|
||||
return ip.IsGlobalUnicast()
|
||||
}
|
||||
|
||||
func abortPanelRateLimited(c *gin.Context, retryAfter time.Duration) {
|
||||
if retryAfter <= 0 {
|
||||
retryAfter = panelRateLimitWindow
|
||||
}
|
||||
seconds := int64(retryAfter / time.Second)
|
||||
if retryAfter%time.Second > 0 {
|
||||
seconds++
|
||||
}
|
||||
c.Header("Retry-After", strconv.FormatInt(seconds, 10))
|
||||
AbortWithError(c, http.StatusTooManyRequests, "RATE_LIMITED", "Too many requests, please slow down and try again later")
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// panelRateLimitStubRepo 内存版 SettingRepository,仅覆盖本测试用到的方法。
|
||||
type panelRateLimitStubRepo struct {
|
||||
mu sync.Mutex
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) Get(_ context.Context, key string) (*service.Setting, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
value, ok := r.values[key]
|
||||
if !ok {
|
||||
return nil, service.ErrSettingNotFound
|
||||
}
|
||||
return &service.Setting{Key: key, Value: value}, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
value, ok := r.values[key]
|
||||
if !ok {
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) Set(_ context.Context, key, value string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string)
|
||||
}
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make(map[string]string, len(keys))
|
||||
for _, key := range keys {
|
||||
if value, ok := r.values[key]; ok {
|
||||
out[key] = value
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) SetMultiple(_ context.Context, settings map[string]string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string)
|
||||
}
|
||||
for key, value := range settings {
|
||||
r.values[key] = value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) GetAll(_ context.Context) (map[string]string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make(map[string]string, len(r.values))
|
||||
for key, value := range r.values {
|
||||
out[key] = value
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) Delete(_ context.Context, key string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
delete(r.values, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
// fakePanelAllower 内存计数版限流原语。
|
||||
type fakePanelAllower struct {
|
||||
mu sync.Mutex
|
||||
counts map[string]int64
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakePanelAllower) Allow(_ context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
if f.err != nil {
|
||||
return middleware.AllowResult{}, f.err
|
||||
}
|
||||
if f.counts == nil {
|
||||
f.counts = make(map[string]int64)
|
||||
}
|
||||
f.counts[key]++
|
||||
count := f.counts[key]
|
||||
result := middleware.AllowResult{Allowed: count <= int64(limit), Count: count}
|
||||
if !result.Allowed {
|
||||
result.RetryAfter = window
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func newPanelRateLimitTestService(t *testing.T, settingsJSON string) *service.SettingService {
|
||||
t.Helper()
|
||||
repo := &panelRateLimitStubRepo{}
|
||||
if settingsJSON != "" {
|
||||
repo.values = map[string]string{"panel_rate_limit_settings": settingsJSON}
|
||||
}
|
||||
return service.NewSettingService(repo, &config.Config{})
|
||||
}
|
||||
|
||||
type panelTestIdentity struct {
|
||||
userID int64
|
||||
role string
|
||||
}
|
||||
|
||||
func newPanelTestRouter(limiter gin.HandlerFunc, identity *panelTestIdentity) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
if identity != nil {
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: identity.userID})
|
||||
c.Set(string(ContextKeyUserRole), identity.role)
|
||||
c.Next()
|
||||
})
|
||||
}
|
||||
router.Use(limiter)
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
return router
|
||||
}
|
||||
|
||||
func performPanelRequest(router *gin.Engine, remoteAddr string) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = remoteAddr
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterGlobalPerUser(t *testing.T) {
|
||||
allower := &fakePanelAllower{}
|
||||
p := &PanelRateLimiter{
|
||||
limiter: allower,
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":2,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
|
||||
userA := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 1, role: service.RoleUser})
|
||||
userB := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 2, role: service.RoleUser})
|
||||
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
|
||||
// 用户 A 超限
|
||||
third := performPanelRequest(userA, "127.0.0.1:1000")
|
||||
require.Equal(t, http.StatusTooManyRequests, third.Code)
|
||||
require.NotEmpty(t, third.Header().Get("Retry-After"))
|
||||
require.Contains(t, third.Body.String(), "RATE_LIMITED")
|
||||
// 用户 B 不受影响(同一来源 IP 也互不干扰)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(userB, "127.0.0.1:1000").Code)
|
||||
|
||||
allower.mu.Lock()
|
||||
defer allower.mu.Unlock()
|
||||
require.Contains(t, allower.counts, "panel:global:user:1")
|
||||
require.Contains(t, allower.counts, "panel:global:user:2")
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterHeavyUsesHeavyRPM(t *testing.T) {
|
||||
allower := &fakePanelAllower{}
|
||||
p := &PanelRateLimiter{
|
||||
limiter: allower,
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":100,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
|
||||
router := newPanelTestRouter(p.Heavy(), &panelTestIdentity{userID: 7, role: service.RoleUser})
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
|
||||
allower.mu.Lock()
|
||||
defer allower.mu.Unlock()
|
||||
require.Contains(t, allower.counts, "panel:heavy:user:7")
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterAdminExemption(t *testing.T) {
|
||||
// 豁免开启:管理员不计数
|
||||
p := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
admin := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
|
||||
for i := 0; i < 5; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(admin, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
// 豁免关闭:管理员一样受限
|
||||
p2 := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":false,"public_ip_rpm":0}`),
|
||||
}
|
||||
admin2 := newPanelTestRouter(p2.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(admin2, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(admin2, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterDisabledOrMissingSubject(t *testing.T) {
|
||||
// 总开关关闭
|
||||
p := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":false,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
|
||||
}
|
||||
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
// 无认证主体:放行(防御分支)
|
||||
p2 := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
anonymous := newPanelTestRouter(p2.Global(), nil)
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(anonymous, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
// nil 限流器(测试环境注入 nil):直接放行
|
||||
var nilLimiter *PanelRateLimiter
|
||||
nilRouter := newPanelTestRouter(nilLimiter.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(nilRouter, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterFailOpenOnRedisError(t *testing.T) {
|
||||
p := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{err: errors.New("redis down")},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
|
||||
}
|
||||
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 5, role: service.RoleUser})
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
publicRouter := newPanelTestRouter(p.PublicIP(), nil)
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(publicRouter, "203.0.113.9:1000").Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterPublicIP(t *testing.T) {
|
||||
allower := &fakePanelAllower{}
|
||||
p := &PanelRateLimiter{
|
||||
limiter: allower,
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":0,"heavy_rpm":0,"exempt_admin":true,"public_ip_rpm":1}`),
|
||||
}
|
||||
router := newPanelTestRouter(p.PublicIP(), nil)
|
||||
|
||||
// 公网 IP:第二次被限
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "203.0.113.9:1000").Code)
|
||||
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "203.0.113.9:1000").Code)
|
||||
// 其他公网 IP 独立计数
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "198.51.100.7:1000").Code)
|
||||
|
||||
// 回环/内网地址(反代内部转发地址):跳过计数,绝不误拦
|
||||
for i := 0; i < 5; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "10.0.0.8:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "172.17.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "192.168.1.30:1000").Code)
|
||||
}
|
||||
|
||||
allower.mu.Lock()
|
||||
defer allower.mu.Unlock()
|
||||
require.Contains(t, allower.counts, "panel:public:ip:203.0.113.9")
|
||||
require.Contains(t, allower.counts, "panel:public:ip:198.51.100.7")
|
||||
for key := range allower.counts {
|
||||
require.NotContains(t, key, "127.0.0.1")
|
||||
require.NotContains(t, key, "10.0.0.8")
|
||||
require.NotContains(t, key, "172.17.0.1")
|
||||
require.NotContains(t, key, "192.168.1.30")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsPubliclyRoutableClientIP(t *testing.T) {
|
||||
require.True(t, isPubliclyRoutableClientIP("203.0.113.9"))
|
||||
require.True(t, isPubliclyRoutableClientIP("2001:db8::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("127.0.0.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("10.1.2.3"))
|
||||
require.False(t, isPubliclyRoutableClientIP("172.16.0.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("192.168.0.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("169.254.1.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("fe80::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("fc00::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("0.0.0.0"))
|
||||
require.False(t, isPubliclyRoutableClientIP(""))
|
||||
require.False(t, isPubliclyRoutableClientIP("not-an-ip"))
|
||||
}
|
||||
@@ -117,12 +117,16 @@ func registerRoutes(
|
||||
// API v1
|
||||
v1 := r.Group("/api/v1")
|
||||
|
||||
// 面板 API 限流器:认证接口按用户 ID、公开接口按安全客户端 IP,
|
||||
// 防止高频刷管理面接口打爆数据库(阈值可在系统设置中调整)。
|
||||
panelRateLimiter := middleware2.NewPanelRateLimiter(redisClient, settingService)
|
||||
|
||||
// 注册各模块路由
|
||||
routes.RegisterAuthRoutes(v1, h, jwtAuth, auditLog, redisClient, settingService)
|
||||
routes.RegisterUserRoutes(v1, h, jwtAuth, auditLog, settingService)
|
||||
routes.RegisterAdminRoutes(v1, h, adminAuth, auditLog, stepUpAuth, settingService)
|
||||
routes.RegisterAuthRoutes(v1, h, jwtAuth, auditLog, redisClient, settingService, panelRateLimiter)
|
||||
routes.RegisterUserRoutes(v1, h, jwtAuth, auditLog, settingService, panelRateLimiter)
|
||||
routes.RegisterAdminRoutes(v1, h, adminAuth, auditLog, stepUpAuth, settingService, panelRateLimiter)
|
||||
routes.RegisterGatewayRoutes(r, h, apiKeyAuth, apiKeyService, subscriptionService, opsService, settingService, compositeResolver, cfg)
|
||||
routes.RegisterPaymentRoutes(v1, h.Payment, h.PaymentWebhook, h.Admin.Payment, jwtAuth, adminAuth, auditLog, settingService)
|
||||
routes.RegisterPaymentRoutes(v1, h.Payment, h.PaymentWebhook, h.Admin.Payment, jwtAuth, adminAuth, auditLog, settingService, panelRateLimiter)
|
||||
|
||||
handler.RegisterPageRoutes(v1, cfg.Pricing.DataDir, gin.HandlerFunc(jwtAuth), gin.HandlerFunc(adminAuth), settingService)
|
||||
}
|
||||
|
||||
@@ -17,9 +17,12 @@ func RegisterAdminRoutes(
|
||||
auditLog middleware.AuditLogMiddleware,
|
||||
stepUpAuth middleware.StepUpAuthMiddleware,
|
||||
settingService *service.SettingService,
|
||||
panelRateLimiter *middleware.PanelRateLimiter,
|
||||
) {
|
||||
admin := v1.Group("/admin")
|
||||
admin.Use(gin.HandlerFunc(adminAuth))
|
||||
// 面板全局按用户限流(默认管理员豁免,可在系统设置中关闭豁免)
|
||||
admin.Use(panelRateLimiter.Global())
|
||||
// 审计中间件挂在认证之后:所有管理面变更类操作 + 敏感读取入审计日志
|
||||
admin.Use(gin.HandlerFunc(auditLog))
|
||||
admin.Use(middleware.AdminComplianceGuard(settingService))
|
||||
@@ -542,6 +545,9 @@ func registerSettingsRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
// 429默认回避配置
|
||||
adminSettings.GET("/rate-limit-429-cooldown", h.Admin.Setting.GetRateLimit429CooldownSettings)
|
||||
adminSettings.PUT("/rate-limit-429-cooldown", h.Admin.Setting.UpdateRateLimit429CooldownSettings)
|
||||
// 面板 API 限流配置
|
||||
adminSettings.GET("/panel-rate-limit", h.Admin.Setting.GetPanelRateLimitSettings)
|
||||
adminSettings.PUT("/panel-rate-limit", h.Admin.Setting.UpdatePanelRateLimitSettings)
|
||||
// 流超时处理配置
|
||||
adminSettings.GET("/stream-timeout", h.Admin.Setting.GetStreamTimeoutSettings)
|
||||
adminSettings.PUT("/stream-timeout", h.Admin.Setting.UpdateStreamTimeoutSettings)
|
||||
|
||||
@@ -20,6 +20,7 @@ func RegisterAuthRoutes(
|
||||
auditLog servermiddleware.AuditLogMiddleware,
|
||||
redisClient *redis.Client,
|
||||
settingService *service.SettingService,
|
||||
panelRateLimiter *servermiddleware.PanelRateLimiter,
|
||||
) {
|
||||
// 创建速率限制器
|
||||
rateLimiter := middleware.NewRateLimiter(redisClient)
|
||||
@@ -213,8 +214,10 @@ func RegisterAuthRoutes(
|
||||
)
|
||||
}
|
||||
|
||||
// 公开设置(无需认证)
|
||||
// 公开设置(无需认证):每次请求都会查询 DB,按客户端 IP 兜底限流,
|
||||
// 防止匿名高频刷接口打爆数据库(反代内部地址会被自动跳过,不会误伤)。
|
||||
settings := v1.Group("/settings")
|
||||
settings.Use(panelRateLimiter.PublicIP())
|
||||
{
|
||||
settings.GET("/public", h.Setting.GetPublicSettings)
|
||||
settings.GET("/email-unsubscribe", h.Setting.UnsubscribeNotificationEmail)
|
||||
@@ -224,6 +227,8 @@ func RegisterAuthRoutes(
|
||||
authenticated := v1.Group("")
|
||||
authenticated.Use(gin.HandlerFunc(jwtAuth))
|
||||
authenticated.Use(servermiddleware.BackendModeUserGuard(settingService))
|
||||
// 面板全局按用户限流
|
||||
authenticated.Use(panelRateLimiter.Global())
|
||||
{
|
||||
authenticated.GET("/auth/me", h.Auth.GetCurrentUser)
|
||||
// 撤销所有会话(需要认证)
|
||||
|
||||
@@ -33,6 +33,7 @@ func newAuthRoutesTestRouter(redisClient *redis.Client) *gin.Engine {
|
||||
}),
|
||||
redisClient,
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
|
||||
return router
|
||||
|
||||
@@ -25,7 +25,7 @@ func TestIngressRejectAdminRoutesRequireAdminAuthentication(t *testing.T) {
|
||||
})
|
||||
auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() })
|
||||
stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() })
|
||||
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil)
|
||||
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil, nil)
|
||||
|
||||
for _, path := range []string{
|
||||
"/api/v1/admin/ops/ingress-rejections",
|
||||
|
||||
@@ -20,11 +20,14 @@ func RegisterPaymentRoutes(
|
||||
adminAuth middleware.AdminAuthMiddleware,
|
||||
auditLog middleware.AuditLogMiddleware,
|
||||
settingService *service.SettingService,
|
||||
panelRateLimiter *middleware.PanelRateLimiter,
|
||||
) {
|
||||
// --- User-facing payment endpoints (authenticated) ---
|
||||
authenticated := v1.Group("/payment")
|
||||
authenticated.Use(gin.HandlerFunc(jwtAuth))
|
||||
authenticated.Use(middleware.BackendModeUserGuard(settingService))
|
||||
// 面板全局按用户限流
|
||||
authenticated.Use(panelRateLimiter.Global())
|
||||
{
|
||||
authenticated.GET("/config", paymentHandler.GetPaymentConfig)
|
||||
authenticated.GET("/checkout-info", paymentHandler.GetCheckoutInfo)
|
||||
|
||||
@@ -115,7 +115,7 @@ func TestPromptAuditAdminRoutesRejectUnauthenticatedAndNonAdminRequests(t *testi
|
||||
})
|
||||
auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() })
|
||||
stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() })
|
||||
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil)
|
||||
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil, nil)
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
|
||||
@@ -15,10 +15,13 @@ func RegisterUserRoutes(
|
||||
jwtAuth middleware.JWTAuthMiddleware,
|
||||
auditLog middleware.AuditLogMiddleware,
|
||||
settingService *service.SettingService,
|
||||
panelRateLimiter *middleware.PanelRateLimiter,
|
||||
) {
|
||||
authenticated := v1.Group("")
|
||||
authenticated.Use(gin.HandlerFunc(jwtAuth))
|
||||
authenticated.Use(middleware.BackendModeUserGuard(settingService))
|
||||
// 面板全局按用户限流:防止单个账号高频刷接口打爆数据库
|
||||
authenticated.Use(panelRateLimiter.Global())
|
||||
// 用户管理面变更类操作入审计(含 TOTP 启用/禁用、step-up 验证、密码修改等安全事件)
|
||||
authenticated.Use(gin.HandlerFunc(auditLog))
|
||||
{
|
||||
@@ -34,7 +37,7 @@ func RegisterUserRoutes(
|
||||
user.POST("/account-bindings/email", h.User.BindEmailIdentity)
|
||||
user.DELETE("/account-bindings/:provider", h.User.UnbindIdentity)
|
||||
user.POST("/auth-identities/bind/start", h.User.StartIdentityBinding)
|
||||
user.GET("/api-keys/:id/usage/daily", h.Usage.GetMyAPIKeyDailyUsage)
|
||||
user.GET("/api-keys/:id/usage/daily", panelRateLimiter.Heavy(), h.Usage.GetMyAPIKeyDailyUsage)
|
||||
user.GET("/platform-quotas", h.User.GetMyPlatformQuotas)
|
||||
|
||||
// 通知邮箱管理
|
||||
@@ -83,8 +86,9 @@ func RegisterUserRoutes(
|
||||
channels.GET("/available", h.AvailableChannel.List)
|
||||
}
|
||||
|
||||
// 使用记录
|
||||
// 使用记录(聚合统计属重查询,叠加更严格的按用户限流)
|
||||
usage := authenticated.Group("/usage")
|
||||
usage.Use(panelRateLimiter.Heavy())
|
||||
{
|
||||
usage.GET("", h.Usage.List)
|
||||
usage.GET("/errors", h.Usage.ListErrors)
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -698,6 +699,25 @@ func TestResolveAccountStatsCost_FallsBackToLiteLLM(t *testing.T) {
|
||||
require.InDelta(t, 0.2, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing.T) {
|
||||
channel := &Channel{
|
||||
ID: 1,
|
||||
Status: StatusActive,
|
||||
ApplyPricingToAccountStats: false,
|
||||
}
|
||||
cs := newTestChannelServiceForStats(t, channel, 10, "antigravity")
|
||||
bs := NewBillingService(&config.Config{}, nil)
|
||||
|
||||
result := resolveAccountStatsCost(
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "gemini-3.6-flash-low",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0,
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 9.15, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
|
||||
channel := &Channel{
|
||||
ID: 1,
|
||||
|
||||
@@ -794,6 +794,16 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode == http.StatusPaymentRequired && s.accountRepo != nil {
|
||||
stateCtx, cancel := openAIAccountStateContext(ctx)
|
||||
defer cancel()
|
||||
_ = s.accountRepo.SetTempUnschedulable(
|
||||
stateCtx,
|
||||
account.ID,
|
||||
time.Now().Add(30*time.Minute),
|
||||
"grok payment required",
|
||||
)
|
||||
}
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Grok Responses API returned %d: %s", resp.StatusCode, string(body)))
|
||||
}
|
||||
|
||||
|
||||
@@ -1588,7 +1588,8 @@ func TestAdminService_CreateCompositeRoute_NormalizesAndPersists(t *testing.T) {
|
||||
require.Equal(t, "router/gpt-", route.PublicModel)
|
||||
require.Equal(t, CompositeRouteMatchPrefix, route.MatchType)
|
||||
require.Equal(t, PlatformOpenAI, route.TargetPlatform)
|
||||
require.Equal(t, "router/gpt-", route.UpstreamModel)
|
||||
// prefix 路由留空 upstream_model 不再回填 public_model:留空表示透传原始请求模型。
|
||||
require.Equal(t, "", route.UpstreamModel)
|
||||
require.Equal(t, CompositeRouteEndpointResponses, route.Endpoint)
|
||||
require.Equal(t, 100, route.Priority)
|
||||
require.True(t, route.Enabled)
|
||||
@@ -1596,6 +1597,29 @@ func TestAdminService_CreateCompositeRoute_NormalizesAndPersists(t *testing.T) {
|
||||
require.Equal(t, route, routeRepo.created)
|
||||
}
|
||||
|
||||
// TestAdminService_CreateCompositeRoute_ExactEmptyUpstreamBackfillsPublicModel 锁定
|
||||
// 保守行为:exact 路由留空 upstream_model 仍回填 public_model(持久化/展示契约不变)。
|
||||
func TestAdminService_CreateCompositeRoute_ExactEmptyUpstreamBackfillsPublicModel(t *testing.T) {
|
||||
groupRepo := &groupRepoStubForAdmin{
|
||||
getByID: &Group{ID: 7, Platform: PlatformComposite},
|
||||
}
|
||||
routeRepo := &compositeRouteRepoStubForAdmin{nextID: 99}
|
||||
svc := &adminServiceImpl{groupRepo: groupRepo, compositeRouteRepo: routeRepo}
|
||||
|
||||
route, err := svc.CreateCompositeRoute(context.Background(), 7, CompositeRouteInput{
|
||||
PublicModel: "openrouter/gpt-5",
|
||||
MatchType: CompositeRouteMatchExact,
|
||||
TargetPlatform: PlatformOpenAI,
|
||||
Endpoint: CompositeRouteEndpointResponses,
|
||||
Enabled: true,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, route)
|
||||
require.Equal(t, CompositeRouteMatchExact, route.MatchType)
|
||||
require.Equal(t, "openrouter/gpt-5", route.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateAndDeleteCompositeRouteRequireRouteOwnership(t *testing.T) {
|
||||
groupRepo := &groupRepoStubForAdmin{
|
||||
getByID: &Group{ID: 7, Platform: PlatformComposite},
|
||||
|
||||
@@ -0,0 +1,526 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type antigravityCompatProtocol uint8
|
||||
|
||||
const (
|
||||
antigravityCompatChatCompletions antigravityCompatProtocol = iota
|
||||
antigravityCompatResponses
|
||||
)
|
||||
|
||||
const (
|
||||
// AntigravityCredentialRejectedClientMessage 是可安全返回给客户端的认证修复提示。
|
||||
AntigravityCredentialRejectedClientMessage = "Antigravity rejected the OAuth credential after refresh; reauthorize the account and verify project_id"
|
||||
// AntigravityCredentialRejectedReason 标识上游拒绝已刷新 OAuth 凭据。
|
||||
AntigravityCredentialRejectedReason GatewayFailureReason = "antigravity_oauth_credential_rejected"
|
||||
)
|
||||
|
||||
type antigravityCompatRequest struct {
|
||||
protocol antigravityCompatProtocol
|
||||
originalBody []byte
|
||||
claudeBody []byte
|
||||
originalModel string
|
||||
clientStream bool
|
||||
includeUsage bool
|
||||
startTime time.Time
|
||||
reasoningEffort *string
|
||||
}
|
||||
|
||||
type antigravityCompatUpstreamCall struct {
|
||||
request antigravityCompatRequest
|
||||
billingModel string
|
||||
prefix string
|
||||
proxyURL string
|
||||
accessToken string
|
||||
geminiBody []byte
|
||||
}
|
||||
|
||||
// ForwardAsChatCompletions 使用 Antigravity 原生 OAuth 账号转发 Chat Completions 请求。
|
||||
func (s *AntigravityGatewayService) ForwardAsChatCompletions(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
body []byte,
|
||||
_ *ParsedRequest,
|
||||
) (*ForwardResult, error) {
|
||||
if err := s.validateAntigravityCompatAccount(c, account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var request apicompat.ChatCompletionsRequest
|
||||
if json.Unmarshal(body, &request) != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
}
|
||||
if strings.TrimSpace(request.Model) == "" {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
}
|
||||
|
||||
responsesRequest, err := apicompat.ChatCompletionsToResponses(&request)
|
||||
if err != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
}
|
||||
claudeRequest, err := apicompat.ResponsesToAnthropicRequest(responsesRequest)
|
||||
if err != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
}
|
||||
preserveChatCompletionTokenLimit(&request, claudeRequest)
|
||||
claudeRequest.Stream = request.Stream
|
||||
claudeBody, err := json.Marshal(claudeRequest)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal anthropic request: %w", err)
|
||||
}
|
||||
|
||||
return s.forwardAntigravityCompat(ctx, c, account, antigravityCompatRequest{
|
||||
protocol: antigravityCompatChatCompletions,
|
||||
originalBody: body,
|
||||
claudeBody: claudeBody,
|
||||
originalModel: request.Model,
|
||||
clientStream: request.Stream,
|
||||
includeUsage: request.StreamOptions != nil && request.StreamOptions.IncludeUsage,
|
||||
startTime: time.Now(),
|
||||
reasoningEffort: extractCCReasoningEffortFromBody(body),
|
||||
})
|
||||
}
|
||||
|
||||
// ForwardAsResponses 使用 Antigravity 原生 OAuth 账号转发 Responses 请求。
|
||||
func (s *AntigravityGatewayService) ForwardAsResponses(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
body []byte,
|
||||
_ *ParsedRequest,
|
||||
) (*ForwardResult, error) {
|
||||
if err := s.validateAntigravityCompatAccount(c, account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var request apicompat.ResponsesRequest
|
||||
if json.Unmarshal(body, &request) != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
}
|
||||
if strings.TrimSpace(request.Model) == "" {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
}
|
||||
|
||||
claudeRequest, err := apicompat.ResponsesToAnthropicRequest(&request)
|
||||
if err != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
}
|
||||
claudeRequest.Stream = request.Stream
|
||||
claudeBody, err := json.Marshal(claudeRequest)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal anthropic request: %w", err)
|
||||
}
|
||||
|
||||
return s.forwardAntigravityCompat(ctx, c, account, antigravityCompatRequest{
|
||||
protocol: antigravityCompatResponses,
|
||||
originalBody: body,
|
||||
claudeBody: claudeBody,
|
||||
originalModel: request.Model,
|
||||
clientStream: request.Stream,
|
||||
startTime: time.Now(),
|
||||
reasoningEffort: ExtractResponsesReasoningEffortFromBody(body),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) validateAntigravityCompatAccount(c *gin.Context, account *Account) error {
|
||||
if account != nil && account.Platform == PlatformAntigravity && account.Type == AccountTypeOAuth {
|
||||
return nil
|
||||
}
|
||||
return s.writeAntigravityCompatError(
|
||||
c,
|
||||
http.StatusBadRequest,
|
||||
"invalid_request_error",
|
||||
"native OAuth account required for antigravity compatibility mode",
|
||||
)
|
||||
}
|
||||
|
||||
func preserveChatCompletionTokenLimit(request *apicompat.ChatCompletionsRequest, claudeRequest *apicompat.AnthropicRequest) {
|
||||
if request == nil || claudeRequest == nil {
|
||||
return
|
||||
}
|
||||
limit := request.MaxTokens
|
||||
if request.MaxCompletionTokens != nil {
|
||||
limit = request.MaxCompletionTokens
|
||||
}
|
||||
if limit != nil && *limit > 0 {
|
||||
claudeRequest.MaxTokens = *limit
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) forwardAntigravityCompat(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
request antigravityCompatRequest,
|
||||
) (*ForwardResult, error) {
|
||||
call, err := s.prepareAntigravityCompatCall(ctx, c, account, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{
|
||||
ctx: ctx,
|
||||
prefix: call.prefix,
|
||||
account: account,
|
||||
proxyURL: call.proxyURL,
|
||||
accessToken: call.accessToken,
|
||||
action: "streamGenerateContent",
|
||||
body: call.geminiBody,
|
||||
c: c,
|
||||
httpUpstream: s.httpUpstream,
|
||||
settingService: s.settingService,
|
||||
accountRepo: s.accountRepo,
|
||||
handleError: s.handleUpstreamError,
|
||||
requestedModel: request.originalModel,
|
||||
isStickySession: false,
|
||||
groupID: 0,
|
||||
sessionHash: "",
|
||||
})
|
||||
if err != nil {
|
||||
return nil, s.handleAntigravityCompatTransportError(c, err)
|
||||
}
|
||||
|
||||
return s.consumeAntigravityCompatResponse(ctx, c, account, call, result.resp)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) prepareAntigravityCompatCall(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
request antigravityCompatRequest,
|
||||
) (*antigravityCompatUpstreamCall, error) {
|
||||
var claudeRequest antigravity.ClaudeRequest
|
||||
if json.Unmarshal(request.claudeBody, &claudeRequest) != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body")
|
||||
}
|
||||
|
||||
mappedModel := s.getMappedModel(account, request.originalModel)
|
||||
if mappedModel == "" {
|
||||
MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
message := fmt.Sprintf("model %s not in whitelist", request.originalModel)
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusForbidden, "permission_error", message)
|
||||
}
|
||||
thinkingEnabled := claudeRequest.Thinking != nil &&
|
||||
(claudeRequest.Thinking.Type == "enabled" || claudeRequest.Thinking.Type == "adaptive")
|
||||
mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled)
|
||||
|
||||
if s.tokenProvider == nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured")
|
||||
}
|
||||
accessToken, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`),
|
||||
}
|
||||
}
|
||||
|
||||
projectID, err := resolveAntigravityProjectID(account)
|
||||
if err != nil {
|
||||
_ = s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return nil, err
|
||||
}
|
||||
geminiBody, err := s.buildAntigravityCompatGeminiBody(ctx, request.claudeBody, &claudeRequest, projectID, mappedModel)
|
||||
if err != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request")
|
||||
}
|
||||
|
||||
request.reasoningEffort = ApplyThinkingEnabledFallback(request.reasoningEffort, request.originalBody, mappedModel)
|
||||
return &antigravityCompatUpstreamCall{
|
||||
request: request,
|
||||
billingModel: mappedModel,
|
||||
prefix: logPrefix(getSessionID(c), account.Name),
|
||||
proxyURL: antigravityCompatProxyURL(account),
|
||||
accessToken: accessToken,
|
||||
geminiBody: geminiBody,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) buildAntigravityCompatGeminiBody(
|
||||
ctx context.Context,
|
||||
claudeBody []byte,
|
||||
claudeRequest *antigravity.ClaudeRequest,
|
||||
projectID string,
|
||||
mappedModel string,
|
||||
) ([]byte, error) {
|
||||
if strings.HasPrefix(strings.ToLower(mappedModel), "gemini-") {
|
||||
body, err := convertClaudeMessagesToGeminiGenerateContent(claudeBody)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body = ensureGeminiFunctionCallThoughtSignatures(body)
|
||||
body, err = injectIdentityPatchToGeminiRequest(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cleaned, cleanErr := cleanGeminiRequest(body); cleanErr == nil {
|
||||
body = cleaned
|
||||
}
|
||||
return s.wrapV1InternalRequest(projectID, mappedModel, body)
|
||||
}
|
||||
|
||||
options := s.getClaudeTransformOptions(ctx)
|
||||
options.EnableIdentityPatch = true
|
||||
return antigravity.TransformClaudeToGeminiWithOptions(claudeRequest, projectID, mappedModel, options)
|
||||
}
|
||||
|
||||
func antigravityCompatProxyURL(account *Account) string {
|
||||
if account.ProxyID == nil || account.Proxy == nil {
|
||||
return ""
|
||||
}
|
||||
return account.Proxy.URL()
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleAntigravityCompatTransportError(c *gin.Context, err error) error {
|
||||
if switchErr, ok := IsAntigravityAccountSwitchError(err); ok {
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: http.StatusServiceUnavailable,
|
||||
ForceCacheBilling: switchErr.IsStickySession,
|
||||
}
|
||||
}
|
||||
if c.Request.Context().Err() != nil {
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response")
|
||||
}
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) consumeAntigravityCompatResponse(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
call *antigravityCompatUpstreamCall,
|
||||
resp *http.Response,
|
||||
) (*ForwardResult, error) {
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode >= http.StatusBadRequest {
|
||||
return nil, s.handleAntigravityCompatHTTPError(ctx, c, account, call, resp)
|
||||
}
|
||||
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
if requestID != "" {
|
||||
c.Header("x-request-id", requestID)
|
||||
}
|
||||
streamResult, err := s.consumeAntigravityCompatSuccess(c, call, resp)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if streamResult.usage == nil {
|
||||
streamResult.usage = &ClaudeUsage{}
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *streamResult.usage,
|
||||
Model: call.request.originalModel,
|
||||
UpstreamModel: call.billingModel,
|
||||
Stream: call.request.clientStream,
|
||||
Duration: time.Since(call.request.startTime),
|
||||
FirstTokenMs: streamResult.firstTokenMs,
|
||||
ReasoningEffort: call.request.reasoningEffort,
|
||||
ClientDisconnect: streamResult.clientDisconnect,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) consumeAntigravityCompatSuccess(
|
||||
c *gin.Context,
|
||||
call *antigravityCompatUpstreamCall,
|
||||
resp *http.Response,
|
||||
) (*antigravityStreamResult, error) {
|
||||
if call.request.clientStream {
|
||||
if call.request.protocol == antigravityCompatChatCompletions {
|
||||
return s.handleChatCompletionsStreamingFromAntigravity(
|
||||
c,
|
||||
resp,
|
||||
call.request.startTime,
|
||||
call.request.originalModel,
|
||||
call.request.includeUsage,
|
||||
)
|
||||
}
|
||||
return s.handleResponsesStreamingFromAntigravity(c, resp, call.request.startTime, call.request.originalModel)
|
||||
}
|
||||
|
||||
if call.request.protocol == antigravityCompatChatCompletions {
|
||||
return s.handleChatCompletionsNonStreamingFromAntigravity(c, resp, call.request.startTime, call.request.originalModel)
|
||||
}
|
||||
return s.handleResponsesNonStreamingFromAntigravity(c, resp, call.request.startTime, call.request.originalModel)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleAntigravityCompatHTTPError(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
call *antigravityCompatUpstreamCall,
|
||||
resp *http.Response,
|
||||
) error {
|
||||
body := s.readUpstreamErrorBody(resp)
|
||||
s.handleUpstreamError(
|
||||
ctx,
|
||||
call.prefix,
|
||||
account,
|
||||
resp.StatusCode,
|
||||
resp.Header,
|
||||
body,
|
||||
call.request.originalModel,
|
||||
0,
|
||||
"",
|
||||
false,
|
||||
)
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
message := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(body)))
|
||||
event := OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
Kind: "failover",
|
||||
Message: message,
|
||||
Detail: s.getUpstreamErrorDetail(body),
|
||||
}
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
event.Stage = string(GatewayFailureStageAccountAuth)
|
||||
event.Scope = string(GatewayFailureScopeAccount)
|
||||
event.Reason = string(AntigravityCredentialRejectedReason)
|
||||
appendOpsUpstreamError(c, event)
|
||||
return antigravityCredentialRejectedError(resp, body)
|
||||
}
|
||||
appendOpsUpstreamError(c, event)
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: body,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
}
|
||||
}
|
||||
return s.writeMappedAntigravityCompatError(c, account, resp.StatusCode, resp.Header.Get("x-request-id"), body)
|
||||
}
|
||||
|
||||
func antigravityCredentialRejectedError(resp *http.Response, body []byte) *UpstreamFailoverError {
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: body,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
Stage: GatewayFailureStageAccountAuth,
|
||||
Scope: GatewayFailureScopeAccount,
|
||||
Reason: AntigravityCredentialRejectedReason,
|
||||
NextAccountAction: NextAccountRetry,
|
||||
ClientStatusCode: http.StatusBadGateway,
|
||||
ClientMessage: AntigravityCredentialRejectedClientMessage,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) writeAntigravityCompatError(
|
||||
c *gin.Context,
|
||||
status int,
|
||||
errType string,
|
||||
message string,
|
||||
) error {
|
||||
MarkResponseCommitted(c)
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": errType,
|
||||
"param": nil,
|
||||
"code": nil,
|
||||
},
|
||||
})
|
||||
return errors.New(message)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) writeMappedAntigravityCompatError(
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
upstreamStatus int,
|
||||
upstreamRequestID string,
|
||||
body []byte,
|
||||
) error {
|
||||
MarkResponseCommitted(c)
|
||||
message := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(body)))
|
||||
setOpsUpstreamError(c, upstreamStatus, message, s.getUpstreamErrorDetail(body))
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: upstreamStatus,
|
||||
UpstreamRequestID: upstreamRequestID,
|
||||
Kind: "http_error",
|
||||
Message: message,
|
||||
})
|
||||
c.JSON(mapUpstreamStatusCode(upstreamStatus), gin.H{
|
||||
"error": gin.H{
|
||||
"message": getPassthroughOrDefault(message, "Upstream request failed"),
|
||||
"type": "upstream_error",
|
||||
"param": nil,
|
||||
"code": nil,
|
||||
},
|
||||
})
|
||||
return fmt.Errorf("upstream error: %d %s", upstreamStatus, message)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleChatCompletionsNonStreamingFromAntigravity(
|
||||
c *gin.Context,
|
||||
resp *http.Response,
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
claudeResponse, result, err := s.collectClaudeStreamResponse(resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
return nil, s.mapAntigravityCompatCollectionError(c, err)
|
||||
}
|
||||
var anthropicResponse apicompat.AnthropicResponse
|
||||
if json.Unmarshal(claudeResponse, &anthropicResponse) != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
|
||||
}
|
||||
responsesResponse := apicompat.AnthropicToResponsesResponse(&anthropicResponse)
|
||||
c.JSON(http.StatusOK, apicompat.ResponsesToChatCompletions(responsesResponse, originalModel))
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleResponsesNonStreamingFromAntigravity(
|
||||
c *gin.Context,
|
||||
resp *http.Response,
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
claudeResponse, result, err := s.collectClaudeStreamResponse(resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
return nil, s.mapAntigravityCompatCollectionError(c, err)
|
||||
}
|
||||
var anthropicResponse apicompat.AnthropicResponse
|
||||
if json.Unmarshal(claudeResponse, &anthropicResponse) != nil {
|
||||
return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
|
||||
}
|
||||
c.JSON(http.StatusOK, apicompat.AnthropicToResponsesResponse(&anthropicResponse))
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) mapAntigravityCompatCollectionError(c *gin.Context, err error) error {
|
||||
var failoverError *UpstreamFailoverError
|
||||
if errors.As(err, &failoverError) {
|
||||
return err
|
||||
}
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return err
|
||||
}
|
||||
if strings.Contains(err.Error(), "stream data interval timeout") {
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_timeout", "Upstream stream data interval timeout")
|
||||
}
|
||||
if errors.Is(err, bufio.ErrTooLong) {
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "response_too_large", "Upstream response line too long")
|
||||
}
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
|
||||
}
|
||||
@@ -0,0 +1,473 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type antigravityCompatStreamAdapter interface {
|
||||
Emit(*apicompat.AnthropicStreamEvent, *antigravityClientWriter)
|
||||
Finalize(*antigravityClientWriter)
|
||||
WriteError(*antigravityClientWriter, string)
|
||||
}
|
||||
|
||||
type antigravityChatStreamAdapter struct {
|
||||
anthropicState *apicompat.AnthropicEventToResponsesState
|
||||
chatState *apicompat.ResponsesEventToChatState
|
||||
}
|
||||
|
||||
func newAntigravityChatStreamAdapter(model string, includeUsage bool) *antigravityChatStreamAdapter {
|
||||
anthropicState := apicompat.NewAnthropicEventToResponsesState()
|
||||
anthropicState.Model = model
|
||||
chatState := apicompat.NewResponsesEventToChatState()
|
||||
chatState.Model = model
|
||||
chatState.IncludeUsage = includeUsage
|
||||
return &antigravityChatStreamAdapter{
|
||||
anthropicState: anthropicState,
|
||||
chatState: chatState,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *antigravityChatStreamAdapter) Emit(event *apicompat.AnthropicStreamEvent, writer *antigravityClientWriter) {
|
||||
for _, responseEvent := range apicompat.AnthropicEventToResponsesEvents(event, a.anthropicState) {
|
||||
a.emitResponseEvent(&responseEvent, writer)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *antigravityChatStreamAdapter) Finalize(writer *antigravityClientWriter) {
|
||||
for _, responseEvent := range apicompat.FinalizeAnthropicResponsesStream(a.anthropicState) {
|
||||
a.emitResponseEvent(&responseEvent, writer)
|
||||
}
|
||||
for _, chunk := range apicompat.FinalizeResponsesChatStream(a.chatState) {
|
||||
if data, err := apicompat.ChatChunkToSSE(chunk); err == nil {
|
||||
writer.Write([]byte(data))
|
||||
}
|
||||
}
|
||||
writer.Write([]byte("data: [DONE]\n\n"))
|
||||
}
|
||||
|
||||
func (a *antigravityChatStreamAdapter) WriteError(writer *antigravityClientWriter, reason string) {
|
||||
writer.Fprintf("data: {\"error\":{\"message\":%q,\"type\":\"upstream_error\"}}\n\n", reason)
|
||||
}
|
||||
|
||||
func (a *antigravityChatStreamAdapter) emitResponseEvent(event *apicompat.ResponsesStreamEvent, writer *antigravityClientWriter) {
|
||||
for _, chunk := range apicompat.ResponsesEventToChatChunks(event, a.chatState) {
|
||||
if data, err := apicompat.ChatChunkToSSE(chunk); err == nil {
|
||||
writer.Write([]byte(data))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type antigravityResponsesStreamAdapter struct {
|
||||
anthropicState *apicompat.AnthropicEventToResponsesState
|
||||
}
|
||||
|
||||
func newAntigravityResponsesStreamAdapter(model string) *antigravityResponsesStreamAdapter {
|
||||
state := apicompat.NewAnthropicEventToResponsesState()
|
||||
state.Model = model
|
||||
return &antigravityResponsesStreamAdapter{anthropicState: state}
|
||||
}
|
||||
|
||||
func (a *antigravityResponsesStreamAdapter) Emit(event *apicompat.AnthropicStreamEvent, writer *antigravityClientWriter) {
|
||||
for _, responseEvent := range apicompat.AnthropicEventToResponsesEvents(event, a.anthropicState) {
|
||||
a.emitResponseEvent(responseEvent, writer)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *antigravityResponsesStreamAdapter) Finalize(writer *antigravityClientWriter) {
|
||||
for _, responseEvent := range apicompat.FinalizeAnthropicResponsesStream(a.anthropicState) {
|
||||
a.emitResponseEvent(responseEvent, writer)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *antigravityResponsesStreamAdapter) WriteError(writer *antigravityClientWriter, reason string) {
|
||||
writer.Fprintf("event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"upstream_error\",\"message\":%q}}\n\n", reason)
|
||||
}
|
||||
|
||||
func (a *antigravityResponsesStreamAdapter) emitResponseEvent(event apicompat.ResponsesStreamEvent, writer *antigravityClientWriter) {
|
||||
if data, err := apicompat.ResponsesEventToSSE(event); err == nil {
|
||||
writer.Write([]byte(data))
|
||||
}
|
||||
}
|
||||
|
||||
type antigravityCompatScanEvent struct {
|
||||
line string
|
||||
err error
|
||||
}
|
||||
|
||||
type antigravityCompatStreamSession struct {
|
||||
processor *antigravity.StreamingProcessor
|
||||
adapter antigravityCompatStreamAdapter
|
||||
writer *antigravityClientWriter
|
||||
usage *ClaudeUsage
|
||||
pendingEvents []apicompat.AnthropicStreamEvent
|
||||
firstTokenMs *int
|
||||
startTime time.Time
|
||||
meaningfulData bool
|
||||
}
|
||||
|
||||
func newAntigravityCompatStreamSession(
|
||||
model string,
|
||||
startTime time.Time,
|
||||
adapter antigravityCompatStreamAdapter,
|
||||
writer *antigravityClientWriter,
|
||||
) *antigravityCompatStreamSession {
|
||||
return &antigravityCompatStreamSession{
|
||||
processor: antigravity.NewStreamingProcessor(model),
|
||||
adapter: adapter,
|
||||
writer: writer,
|
||||
usage: &ClaudeUsage{},
|
||||
startTime: startTime,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) consume(line string) {
|
||||
claudeEvents := s.processor.ProcessLine(strings.TrimRight(line, "\r\n"))
|
||||
if len(claudeEvents) == 0 {
|
||||
return
|
||||
}
|
||||
s.consumeClaudeEvents(claudeEvents)
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) hasMeaningfulData() bool {
|
||||
return s.meaningfulData
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) finish() *antigravityStreamResult {
|
||||
finalEvents, usage := s.processor.Finish()
|
||||
mergeAntigravityCompatUsage(s.usage, usage)
|
||||
s.consumeClaudeEvents(finalEvents)
|
||||
s.adapter.Finalize(s.writer)
|
||||
return s.result(s.writer.Disconnected())
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) collectResult(clientDisconnect bool) *antigravityStreamResult {
|
||||
_, usage := s.processor.Finish()
|
||||
mergeAntigravityCompatUsage(s.usage, usage)
|
||||
return s.result(clientDisconnect)
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) result(clientDisconnect bool) *antigravityStreamResult {
|
||||
return &antigravityStreamResult{
|
||||
usage: s.usage,
|
||||
firstTokenMs: s.firstTokenMs,
|
||||
clientDisconnect: clientDisconnect,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) consumeClaudeEvents(data []byte) {
|
||||
var eventType string
|
||||
for _, line := range strings.Split(string(data), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
switch {
|
||||
case strings.HasPrefix(line, "event:"):
|
||||
eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
case strings.HasPrefix(line, "data:"):
|
||||
s.consumeClaudeData(eventType, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) consumeClaudeData(eventType, payload string) {
|
||||
var event apicompat.AnthropicStreamEvent
|
||||
if json.Unmarshal([]byte(payload), &event) != nil {
|
||||
return
|
||||
}
|
||||
if event.Type == "" {
|
||||
event.Type = eventType
|
||||
}
|
||||
if event.Usage != nil {
|
||||
mergeAnthropicUsage(s.usage, *event.Usage)
|
||||
}
|
||||
if event.Message != nil {
|
||||
mergeAnthropicUsage(s.usage, event.Message.Usage)
|
||||
}
|
||||
s.emitOrBuffer(event)
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) emitOrBuffer(event apicompat.AnthropicStreamEvent) {
|
||||
if s.meaningfulData {
|
||||
s.adapter.Emit(&event, s.writer)
|
||||
return
|
||||
}
|
||||
|
||||
s.pendingEvents = append(s.pendingEvents, event)
|
||||
if !isMeaningfulAntigravityCompatEvent(&event) {
|
||||
return
|
||||
}
|
||||
|
||||
s.meaningfulData = true
|
||||
ms := int(time.Since(s.startTime).Milliseconds())
|
||||
s.firstTokenMs = &ms
|
||||
for i := range s.pendingEvents {
|
||||
s.adapter.Emit(&s.pendingEvents[i], s.writer)
|
||||
}
|
||||
s.pendingEvents = nil
|
||||
}
|
||||
|
||||
func isMeaningfulAntigravityCompatEvent(event *apicompat.AnthropicStreamEvent) bool {
|
||||
if event == nil {
|
||||
return false
|
||||
}
|
||||
if event.Type == "message_stop" {
|
||||
return true
|
||||
}
|
||||
if event.ContentBlock != nil {
|
||||
block := event.ContentBlock
|
||||
return block.Type == "tool_use" ||
|
||||
block.Text != "" ||
|
||||
block.Thinking != "" ||
|
||||
block.Signature != "" ||
|
||||
block.Source != nil
|
||||
}
|
||||
if event.Delta != nil {
|
||||
delta := event.Delta
|
||||
return delta.Text != "" ||
|
||||
delta.PartialJSON != "" ||
|
||||
delta.Thinking != "" ||
|
||||
delta.Signature != "" ||
|
||||
delta.StopReason != ""
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func mergeAntigravityCompatUsage(dst *ClaudeUsage, src *antigravity.ClaudeUsage) {
|
||||
if dst == nil || src == nil {
|
||||
return
|
||||
}
|
||||
dst.InputTokens = src.InputTokens
|
||||
dst.OutputTokens = src.OutputTokens
|
||||
dst.CacheCreationInputTokens = src.CacheCreationInputTokens
|
||||
dst.CacheReadInputTokens = src.CacheReadInputTokens
|
||||
dst.ImageOutputTokens = src.ImageOutputTokens
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
c *gin.Context,
|
||||
resp *http.Response,
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
adapter antigravityCompatStreamAdapter,
|
||||
prefix string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
return nil, errors.New("streaming not supported")
|
||||
}
|
||||
|
||||
writer := newAntigravityClientWriter(c.Writer, flusher, prefix)
|
||||
writer.beforeFirstWrite = func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
c.Status(http.StatusOK)
|
||||
}
|
||||
session := newAntigravityCompatStreamSession(originalModel, startTime, adapter, writer)
|
||||
events, stopScanner, maxLineSize := s.startAntigravityCompatScanner(resp.Body)
|
||||
defer stopScanner()
|
||||
|
||||
timeout := s.antigravityCompatStreamTimeout()
|
||||
timeoutTimer, timeoutCh := newAntigravityCompatTimer(timeout)
|
||||
if timeoutTimer != nil {
|
||||
defer timeoutTimer.Stop()
|
||||
}
|
||||
keepaliveTicker, keepaliveCh := s.newAntigravityCompatKeepaliveTicker()
|
||||
if keepaliveTicker != nil {
|
||||
defer keepaliveTicker.Stop()
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case event, open := <-events:
|
||||
if !open {
|
||||
if !session.hasMeaningfulData() && !writer.Disconnected() {
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
return session.finish(), nil
|
||||
}
|
||||
if event.err != nil {
|
||||
return s.handleAntigravityCompatReadError(c, session, event.err, maxLineSize, prefix)
|
||||
}
|
||||
resetAntigravityCompatTimer(timeoutTimer, timeout)
|
||||
session.consume(event.line)
|
||||
|
||||
case <-timeoutCh:
|
||||
if writer.Disconnected() {
|
||||
return session.collectResult(true), nil
|
||||
}
|
||||
if !session.hasMeaningfulData() {
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (%s)", prefix)
|
||||
writeAntigravityCompatStreamError(c, adapter, writer, "stream_timeout")
|
||||
return session.collectResult(false), fmt.Errorf("stream data interval timeout")
|
||||
|
||||
case <-keepaliveCh:
|
||||
if session.hasMeaningfulData() && !writer.Disconnected() {
|
||||
writer.Write([]byte(": ping\n\n"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) startAntigravityCompatScanner(
|
||||
body io.Reader,
|
||||
) (<-chan antigravityCompatScanEvent, func(), int) {
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanner := bufio.NewScanner(body)
|
||||
scanBuf := getSSEScannerBuf64K()
|
||||
scanner.Buffer(scanBuf[:0], maxLineSize)
|
||||
|
||||
events := make(chan antigravityCompatScanEvent, 16)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer putSSEScannerBuf64K(scanBuf)
|
||||
defer close(events)
|
||||
send := func(event antigravityCompatScanEvent) bool {
|
||||
select {
|
||||
case events <- event:
|
||||
return true
|
||||
case <-done:
|
||||
return false
|
||||
}
|
||||
}
|
||||
for scanner.Scan() {
|
||||
if !send(antigravityCompatScanEvent{line: scanner.Text()}) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
send(antigravityCompatScanEvent{err: err})
|
||||
}
|
||||
}()
|
||||
return events, func() { close(done) }, maxLineSize
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) antigravityCompatStreamTimeout() time.Duration {
|
||||
if s.settingService == nil || s.settingService.cfg == nil {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) newAntigravityCompatKeepaliveTicker() (*time.Ticker, <-chan time.Time) {
|
||||
if s.settingService == nil || s.settingService.cfg == nil {
|
||||
return nil, nil
|
||||
}
|
||||
interval := time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
|
||||
if interval <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
return ticker, ticker.C
|
||||
}
|
||||
|
||||
func newAntigravityCompatTimer(timeout time.Duration) (*time.Timer, <-chan time.Time) {
|
||||
if timeout <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
timer := time.NewTimer(timeout)
|
||||
return timer, timer.C
|
||||
}
|
||||
|
||||
func resetAntigravityCompatTimer(timer *time.Timer, timeout time.Duration) {
|
||||
if timer == nil {
|
||||
return
|
||||
}
|
||||
if !timer.Stop() {
|
||||
select {
|
||||
case <-timer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
timer.Reset(timeout)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleAntigravityCompatReadError(
|
||||
c *gin.Context,
|
||||
session *antigravityCompatStreamSession,
|
||||
err error,
|
||||
maxLineSize int,
|
||||
prefix string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
if !session.hasMeaningfulData() && !session.writer.Disconnected() {
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
if disconnect, handled := handleStreamReadError(err, session.writer.Disconnected(), prefix); handled {
|
||||
return session.collectResult(disconnect), nil
|
||||
}
|
||||
if errors.Is(err, bufio.ErrTooLong) {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (%s): max_size=%d error=%v", prefix, maxLineSize, err)
|
||||
writeAntigravityCompatStreamError(c, session.adapter, session.writer, "response_too_large")
|
||||
return session.result(false), err
|
||||
}
|
||||
writeAntigravityCompatStreamError(c, session.adapter, session.writer, "stream_read_error")
|
||||
return nil, fmt.Errorf("stream read error: %w", err)
|
||||
}
|
||||
|
||||
func writeAntigravityCompatStreamError(
|
||||
c *gin.Context,
|
||||
adapter antigravityCompatStreamAdapter,
|
||||
writer *antigravityClientWriter,
|
||||
reason string,
|
||||
) {
|
||||
adapter.WriteError(writer, reason)
|
||||
MarkResponseCommitted(c)
|
||||
}
|
||||
|
||||
func antigravityCompatEmptyStreamError() error {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Empty Antigravity compatibility stream, triggering failover")
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
ResponseBody: []byte(`{"error":"empty stream response from upstream"}`),
|
||||
RetryableOnSameAccount: true,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleChatCompletionsStreamingFromAntigravity(
|
||||
c *gin.Context,
|
||||
resp *http.Response,
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
includeUsage bool,
|
||||
) (*antigravityStreamResult, error) {
|
||||
return s.handleAntigravityCompatStream(
|
||||
c,
|
||||
resp,
|
||||
startTime,
|
||||
originalModel,
|
||||
newAntigravityChatStreamAdapter(originalModel, includeUsage),
|
||||
"antigravity chat completions stream",
|
||||
)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleResponsesStreamingFromAntigravity(
|
||||
c *gin.Context,
|
||||
resp *http.Response,
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
return s.handleAntigravityCompatStream(
|
||||
c,
|
||||
resp,
|
||||
startTime,
|
||||
originalModel,
|
||||
newAntigravityResponsesStreamAdapter(originalModel),
|
||||
"antigravity responses stream",
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,617 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type antigravityCompatTokenCache struct {
|
||||
token string
|
||||
}
|
||||
|
||||
type antigravityCompatErrorReader struct {
|
||||
data []byte
|
||||
off int
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *antigravityCompatErrorReader) Read(p []byte) (int, error) {
|
||||
if r.off < len(r.data) {
|
||||
n := copy(p, r.data[r.off:])
|
||||
r.off += n
|
||||
return n, nil
|
||||
}
|
||||
return 0, r.err
|
||||
}
|
||||
|
||||
func (r *antigravityCompatErrorReader) Close() error { return nil }
|
||||
|
||||
func (c *antigravityCompatTokenCache) GetAccessToken(context.Context, string) (string, error) {
|
||||
return c.token, nil
|
||||
}
|
||||
|
||||
func (c *antigravityCompatTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *antigravityCompatTokenCache) DeleteAccessToken(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *antigravityCompatTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (c *antigravityCompatTokenCache) ReleaseRefreshLock(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func newAntigravityCompatService(cfg config.GatewayConfig, upstream HTTPUpstream) *AntigravityGatewayService {
|
||||
tokenProvider := NewAntigravityTokenProvider(
|
||||
nil,
|
||||
&antigravityCompatTokenCache{token: "fresh-oauth-token"},
|
||||
nil,
|
||||
)
|
||||
return NewAntigravityGatewayService(
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
tokenProvider,
|
||||
nil,
|
||||
upstream,
|
||||
NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: cfg}),
|
||||
nil,
|
||||
)
|
||||
}
|
||||
|
||||
func newAntigravityCompatAccount(accountType string) *Account {
|
||||
return &Account{
|
||||
ID: 3757,
|
||||
Name: "antigravity-compat",
|
||||
Platform: PlatformAntigravity,
|
||||
Type: accountType,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "stale-account-token",
|
||||
"project_id": "project-3757",
|
||||
"model_mapping": map[string]any{
|
||||
"gemini-3.1-pro-high": "gemini-3.1-pro-high",
|
||||
"claude-sonnet-4-5": "claude-sonnet-4-5",
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newAntigravityCompatContext(method, path string, body []byte) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(method, path, bytes.NewReader(body))
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func antigravityCompatSuccessResponse() *http.Response {
|
||||
body := `data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":3}}}` + "\n\n"
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"X-Request-Id": []string{"request-3757"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatOAuthUsesNativeTokenAndRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
body []byte
|
||||
call func(*AntigravityGatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error)
|
||||
}{
|
||||
{
|
||||
name: "chat completions",
|
||||
path: "/v1/chat/completions",
|
||||
body: []byte(`{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"Reply exactly: ok"}]}`),
|
||||
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
return svc.ForwardAsChatCompletions(ctx, c, account, body, nil)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "responses",
|
||||
path: "/v1/responses",
|
||||
body: []byte(`{"model":"gemini-3.1-pro-high","input":"Reply exactly: ok"}`),
|
||||
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
return svc.ForwardAsResponses(ctx, c, account, body, nil)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var authorization string
|
||||
var upstreamPath string
|
||||
var upstreamAlt string
|
||||
upstream := &queuedHTTPUpstreamStub{
|
||||
responses: []*http.Response{antigravityCompatSuccessResponse()},
|
||||
onCall: func(req *http.Request, _ *queuedHTTPUpstreamStub) {
|
||||
authorization = req.Header.Get("Authorization")
|
||||
upstreamPath = req.URL.Path
|
||||
upstreamAlt = req.URL.Query().Get("alt")
|
||||
},
|
||||
}
|
||||
svc := newAntigravityCompatService(
|
||||
config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
upstream,
|
||||
)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, tt.path, tt.body)
|
||||
|
||||
result, err := tt.call(svc, context.Background(), c, newAntigravityCompatAccount(AccountTypeOAuth), tt.body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "Bearer fresh-oauth-token", authorization)
|
||||
require.Equal(t, "/v1internal:streamGenerateContent", upstreamPath)
|
||||
require.Equal(t, "sse", upstreamAlt)
|
||||
require.Equal(t, "request-3757", result.RequestID)
|
||||
require.Equal(t, 8, result.Usage.InputTokens)
|
||||
require.Equal(t, 3, result.Usage.OutputTokens)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "ok")
|
||||
if tt.name == "chat completions" {
|
||||
require.Equal(t, "stop", gjson.Get(recorder.Body.String(), "choices.0.finish_reason").String())
|
||||
require.Equal(t, int64(8), gjson.Get(recorder.Body.String(), "usage.prompt_tokens").Int())
|
||||
require.Equal(t, int64(3), gjson.Get(recorder.Body.String(), "usage.completion_tokens").Int())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatRejectsUnsupportedAccountType(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
accountType string
|
||||
call func(*AntigravityGatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error)
|
||||
}{
|
||||
{
|
||||
name: "chat completions upstream",
|
||||
path: "/v1/chat/completions",
|
||||
accountType: AccountTypeUpstream,
|
||||
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
return svc.ForwardAsChatCompletions(ctx, c, account, body, nil)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "responses setup token",
|
||||
path: "/v1/responses",
|
||||
accountType: AccountTypeSetupToken,
|
||||
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
return svc.ForwardAsResponses(ctx, c, account, body, nil)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
body := []byte(`{"model":"gemini-3.1-pro-high"}`)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, tt.path, body)
|
||||
|
||||
result, err := tt.call(&AntigravityGatewayService{}, context.Background(), c, newAntigravityCompatAccount(tt.accountType), body)
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "native OAuth account required for antigravity compatibility mode")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatPreservesChatTokenLimit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
want int64
|
||||
}{
|
||||
{
|
||||
name: "legacy max_tokens below bridge floor",
|
||||
body: `{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}],"max_tokens":8}`,
|
||||
want: 8,
|
||||
},
|
||||
{
|
||||
name: "max_completion_tokens takes precedence",
|
||||
body: `{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}],"max_tokens":8,"max_completion_tokens":13}`,
|
||||
want: 13,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{antigravityCompatSuccessResponse()}}
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
||||
body := []byte(tt.body)
|
||||
c, _ := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", body)
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(
|
||||
context.Background(),
|
||||
c,
|
||||
newAntigravityCompatAccount(AccountTypeOAuth),
|
||||
body,
|
||||
nil,
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Len(t, upstream.requestBodies, 1)
|
||||
require.Equal(t, tt.want, gjson.GetBytes(upstream.requestBodies[0], "request.generationConfig.maxOutputTokens").Int())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatRoutesByMappedModelFamily(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
model string
|
||||
wantSessionID bool
|
||||
}{
|
||||
{model: "gemini-3.1-pro-high", wantSessionID: false},
|
||||
{model: "claude-sonnet-4-5", wantSessionID: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.model, func(t *testing.T) {
|
||||
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{antigravityCompatSuccessResponse()}}
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
||||
body := []byte(`{"model":"` + tt.model + `","messages":[{"role":"user","content":"ok"}],"max_tokens":8}`)
|
||||
c, _ := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", body)
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(
|
||||
context.Background(),
|
||||
c,
|
||||
newAntigravityCompatAccount(AccountTypeOAuth),
|
||||
body,
|
||||
nil,
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Len(t, upstream.requestBodies, 1)
|
||||
require.Equal(t, tt.model, gjson.GetBytes(upstream.requestBodies[0], "model").String())
|
||||
require.Equal(t, tt.wantSessionID, gjson.GetBytes(upstream.requestBodies[0], "request.sessionId").Exists())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatUnauthorizedIsCredentialFailure(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
Header: http.Header{"X-Request-Id": []string{"auth-3757"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"Invalid bearer token"}}`)),
|
||||
}}}
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
||||
body := []byte(`{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}]}`)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", body)
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(
|
||||
context.Background(),
|
||||
c,
|
||||
newAntigravityCompatAccount(AccountTypeOAuth),
|
||||
body,
|
||||
nil,
|
||||
)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage)
|
||||
require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
|
||||
require.Equal(t, AntigravityCredentialRejectedReason, failoverErr.Reason)
|
||||
require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
|
||||
require.Equal(t, http.StatusBadGateway, failoverErr.ClientStatusCode)
|
||||
require.Equal(t, AntigravityCredentialRejectedClientMessage, failoverErr.ClientMessage)
|
||||
require.Equal(t, "auth-3757", failoverErr.ResponseHeaders.Get("X-Request-Id"))
|
||||
require.Empty(t, recorder.Body.String())
|
||||
}
|
||||
|
||||
func TestAntigravityCompatEmptyStreamTriggersFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
run func(*AntigravityGatewayService, *gin.Context, *http.Response) (*antigravityStreamResult, error)
|
||||
}{
|
||||
{
|
||||
name: "chat completions",
|
||||
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
||||
return svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", true)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "responses",
|
||||
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
||||
return svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader("data: malformed\n\ndata: [DONE]\n\n")),
|
||||
}
|
||||
|
||||
result, err := tt.run(svc, c, resp)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||||
require.Empty(t, recorder.Body.String())
|
||||
require.Empty(t, recorder.Header().Get("Content-Type"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatUsageOnlyStreamTriggersFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
run func(*AntigravityGatewayService, *gin.Context, *http.Response) (*antigravityStreamResult, error)
|
||||
}{
|
||||
{
|
||||
name: "chat completions",
|
||||
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
||||
return svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", true)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "responses",
|
||||
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
||||
return svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`data: {"response":{"responseId":"resp_3757","usageMetadata":{"promptTokenCount":8}}}` + "\n\n",
|
||||
)),
|
||||
}
|
||||
|
||||
result, err := tt.run(svc, c, resp)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||||
require.Empty(t, recorder.Body.String())
|
||||
require.Empty(t, recorder.Header().Get("Content-Type"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatUsageOnlyNonStreamingTriggersFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
body []byte
|
||||
call func(*AntigravityGatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error)
|
||||
}{
|
||||
{
|
||||
name: "chat completions",
|
||||
path: "/v1/chat/completions",
|
||||
body: []byte(`{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}]}`),
|
||||
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
return svc.ForwardAsChatCompletions(ctx, c, account, body, nil)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "responses",
|
||||
path: "/v1/responses",
|
||||
body: []byte(`{"model":"gemini-3.1-pro-high","input":"ok"}`),
|
||||
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
return svc.ForwardAsResponses(ctx, c, account, body, nil)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`data: {"response":{"responseId":"resp_3757","usageMetadata":{"promptTokenCount":8}}}` + "\n\n",
|
||||
)),
|
||||
}}}
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, tt.path, tt.body)
|
||||
|
||||
result, err := tt.call(
|
||||
svc,
|
||||
context.Background(),
|
||||
c,
|
||||
newAntigravityCompatAccount(AccountTypeOAuth),
|
||||
tt.body,
|
||||
)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||||
require.Empty(t, recorder.Body.String())
|
||||
require.Empty(t, recorder.Header().Get("Content-Type"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatChatStreamMapsToolCallAndUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", nil)
|
||||
body := `data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"functionCall":{"id":"call_3757","name":"get_weather","args":{"city":"Tokyo"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":3}}}` + "\n\n"
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
|
||||
result, err := svc.handleChatCompletionsStreamingFromAntigravity(
|
||||
c,
|
||||
resp,
|
||||
time.Now(),
|
||||
"gemini-3.1-pro-high",
|
||||
true,
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 8, result.usage.InputTokens)
|
||||
require.Equal(t, 3, result.usage.OutputTokens)
|
||||
require.Contains(t, recorder.Body.String(), `"tool_calls"`)
|
||||
require.Contains(t, recorder.Body.String(), `"get_weather"`)
|
||||
require.Contains(t, recorder.Body.String(), `"finish_reason":"tool_calls"`)
|
||||
require.Contains(t, recorder.Body.String(), `"prompt_tokens":8`)
|
||||
require.Equal(t, 1, strings.Count(recorder.Body.String(), "data: [DONE]"))
|
||||
}
|
||||
|
||||
func TestAntigravityCompatFirstEventTimeoutTriggersFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(
|
||||
config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 1},
|
||||
nil,
|
||||
)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", nil)
|
||||
reader, writer := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: reader}
|
||||
type outcome struct {
|
||||
result *antigravityStreamResult
|
||||
err error
|
||||
}
|
||||
done := make(chan outcome, 1)
|
||||
|
||||
go func() {
|
||||
result, err := svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", false)
|
||||
done <- outcome{result: result, err: err}
|
||||
}()
|
||||
|
||||
select {
|
||||
case got := <-done:
|
||||
require.Nil(t, got.result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, got.err, &failoverErr)
|
||||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||||
require.Empty(t, recorder.Header().Get("Content-Type"))
|
||||
case <-time.After(2 * time.Second):
|
||||
_ = writer.Close()
|
||||
_ = reader.Close()
|
||||
t.Fatal("compat stream ignored StreamDataIntervalTimeout")
|
||||
}
|
||||
_ = writer.Close()
|
||||
_ = reader.Close()
|
||||
}
|
||||
|
||||
func TestAntigravityCompatClientDisconnectDrainsUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, _ := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
body := strings.Join([]string{
|
||||
`data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1}}}`,
|
||||
"",
|
||||
`data: {"response":{"responseId":"resp_3757","candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":15}}}`,
|
||||
"",
|
||||
}, "\n")
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
|
||||
result, err := svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", false)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.Equal(t, 15, result.usage.OutputTokens)
|
||||
}
|
||||
|
||||
func TestAntigravityCompatStreamErrorCommitsSingleTerminalFrame(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
||||
body := []byte(`data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1}}}` + "\n\n")
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: &antigravityCompatErrorReader{
|
||||
data: body,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: error"))
|
||||
}
|
||||
|
||||
func TestAntigravityCompatKeepaliveAfterFirstEvent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(
|
||||
config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamKeepaliveInterval: 1},
|
||||
nil,
|
||||
)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
||||
reader, writer := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: reader}
|
||||
done := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
_, err := svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
||||
done <- err
|
||||
}()
|
||||
_, err := io.WriteString(
|
||||
writer,
|
||||
`data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1}}}`+"\n\n",
|
||||
)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(1200 * time.Millisecond)
|
||||
require.NoError(t, writer.Close())
|
||||
require.NoError(t, <-done)
|
||||
require.Contains(t, recorder.Body.String(), ": ping\n\n")
|
||||
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
|
||||
require.NoError(t, reader.Close())
|
||||
}
|
||||
@@ -25,10 +25,11 @@ type antigravityStreamResult struct {
|
||||
// antigravityClientWriter 封装流式响应的客户端写入,自动检测断开并标记。
|
||||
// 断开后所有写入操作变为 no-op,调用方通过 Disconnected() 判断是否继续 drain 上游。
|
||||
type antigravityClientWriter struct {
|
||||
w gin.ResponseWriter
|
||||
flusher http.Flusher
|
||||
disconnected bool
|
||||
prefix string // 日志前缀,标识来源方法
|
||||
w gin.ResponseWriter
|
||||
flusher http.Flusher
|
||||
disconnected bool
|
||||
prefix string // 日志前缀,标识来源方法
|
||||
beforeFirstWrite func()
|
||||
}
|
||||
|
||||
func newAntigravityClientWriter(w gin.ResponseWriter, flusher http.Flusher, prefix string) *antigravityClientWriter {
|
||||
@@ -40,6 +41,7 @@ func (cw *antigravityClientWriter) Write(p []byte) bool {
|
||||
if cw.disconnected {
|
||||
return false
|
||||
}
|
||||
cw.prepareFirstWrite()
|
||||
if _, err := cw.w.Write(p); err != nil {
|
||||
cw.markDisconnected()
|
||||
return false
|
||||
@@ -53,6 +55,7 @@ func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool {
|
||||
if cw.disconnected {
|
||||
return false
|
||||
}
|
||||
cw.prepareFirstWrite()
|
||||
if _, err := fmt.Fprintf(cw.w, format, args...); err != nil {
|
||||
cw.markDisconnected()
|
||||
return false
|
||||
@@ -63,6 +66,15 @@ func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool {
|
||||
|
||||
func (cw *antigravityClientWriter) Disconnected() bool { return cw.disconnected }
|
||||
|
||||
func (cw *antigravityClientWriter) prepareFirstWrite() {
|
||||
if cw.beforeFirstWrite == nil {
|
||||
return
|
||||
}
|
||||
prepare := cw.beforeFirstWrite
|
||||
cw.beforeFirstWrite = nil
|
||||
prepare()
|
||||
}
|
||||
|
||||
func (cw *antigravityClientWriter) markDisconnected() {
|
||||
cw.disconnected = true
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during streaming (%s), continuing to drain upstream for billing", cw.prefix)
|
||||
@@ -751,9 +763,9 @@ func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int,
|
||||
return fmt.Errorf("%s", message)
|
||||
}
|
||||
|
||||
// handleClaudeStreamToNonStreaming 收集上游流式响应,转换为 Claude 非流式格式返回
|
||||
// collectClaudeStreamResponse 收集上游流式响应,转换为 Claude 非流式格式返回
|
||||
// 用于处理客户端非流式请求但上游只支持流式的情况
|
||||
func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
|
||||
func (s *AntigravityGatewayService) collectClaudeStreamResponse(resp *http.Response, startTime time.Time, originalModel string) ([]byte, *antigravityStreamResult, error) {
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
|
||||
@@ -766,6 +778,7 @@ func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Cont
|
||||
var last map[string]any
|
||||
var lastWithParts map[string]any
|
||||
var collectedParts []map[string]any // 收集所有 parts(包括 text、thinking、functionCall、inlineData 等)
|
||||
var meaningfulResponse bool
|
||||
|
||||
type scanEvent struct {
|
||||
line string
|
||||
@@ -827,7 +840,7 @@ func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Cont
|
||||
if errors.Is(ev.err, bufio.ErrTooLong) {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err)
|
||||
}
|
||||
return nil, ev.err
|
||||
return nil, nil, ev.err
|
||||
}
|
||||
|
||||
line := ev.line
|
||||
@@ -853,21 +866,23 @@ func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Cont
|
||||
continue
|
||||
}
|
||||
|
||||
// 记录首 token 时间
|
||||
if firstTokenMs == nil {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
|
||||
last = parsed
|
||||
|
||||
// 保留最后一个有 parts 的响应,并收集所有 parts
|
||||
if parts := extractGeminiParts(parsed); len(parts) > 0 {
|
||||
parts := extractGeminiParts(parsed)
|
||||
if len(parts) > 0 {
|
||||
lastWithParts = parsed
|
||||
|
||||
// 收集所有 parts(text、thinking、functionCall、inlineData 等)
|
||||
collectedParts = append(collectedParts, parts...)
|
||||
}
|
||||
if len(parts) > 0 || strings.TrimSpace(extractGeminiFinishReason(parsed)) != "" {
|
||||
meaningfulResponse = true
|
||||
if firstTokenMs == nil {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
}
|
||||
|
||||
case <-intervalCh:
|
||||
lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
|
||||
@@ -875,24 +890,24 @@ func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Cont
|
||||
continue
|
||||
}
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity claude non-stream)")
|
||||
return nil, fmt.Errorf("stream data interval timeout")
|
||||
return nil, nil, fmt.Errorf("stream data interval timeout")
|
||||
}
|
||||
}
|
||||
|
||||
returnResponse:
|
||||
// 选择最后一个有效响应
|
||||
finalResponse := pickGeminiCollectResult(last, lastWithParts)
|
||||
|
||||
// 处理空响应情况 — 触发同账号重试 + failover 切换账号
|
||||
if last == nil && lastWithParts == nil {
|
||||
if !meaningfulResponse {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (claude non-stream), triggering failover")
|
||||
return nil, &UpstreamFailoverError{
|
||||
return nil, nil, &UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
ResponseBody: []byte(`{"error":"empty stream response from upstream"}`),
|
||||
RetryableOnSameAccount: true,
|
||||
}
|
||||
}
|
||||
|
||||
// 选择最后一个有效响应
|
||||
finalResponse := pickGeminiCollectResult(last, lastWithParts)
|
||||
|
||||
// 将收集的所有 parts 合并到最终响应中
|
||||
if len(collectedParts) > 0 {
|
||||
finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts)
|
||||
@@ -901,27 +916,56 @@ returnResponse:
|
||||
// 序列化为 JSON(Gemini 格式)
|
||||
geminiBody, err := json.Marshal(finalResponse)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal gemini response: %w", err)
|
||||
return nil, nil, fmt.Errorf("failed to marshal gemini response: %w", err)
|
||||
}
|
||||
|
||||
// 转换 Gemini 响应为 Claude 格式
|
||||
claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody))
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
|
||||
return nil, nil, fmt.Errorf("failed to parse upstream response: %w", err)
|
||||
}
|
||||
|
||||
c.Data(http.StatusOK, "application/json", claudeResp)
|
||||
|
||||
// 转换为 service.ClaudeUsage
|
||||
usage := &ClaudeUsage{
|
||||
InputTokens: agUsage.InputTokens,
|
||||
OutputTokens: agUsage.OutputTokens,
|
||||
CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
|
||||
CacheReadInputTokens: agUsage.CacheReadInputTokens,
|
||||
ImageOutputTokens: agUsage.ImageOutputTokens,
|
||||
}
|
||||
|
||||
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
|
||||
return claudeResp, &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
|
||||
}
|
||||
|
||||
// handleClaudeStreamToNonStreaming 收集上游流式响应,转换为 Claude 非流式格式返回
|
||||
// 用于处理客户端非流式请求但上游只支持流式的情况
|
||||
func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
|
||||
claudeResp, streamRes, err := s.collectClaudeStreamResponse(resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
var failoverErr *UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
errMsg := "Failed to parse upstream response"
|
||||
errType := "upstream_error"
|
||||
if strings.Contains(err.Error(), "stream data interval timeout") {
|
||||
errMsg = "Upstream stream data interval timeout"
|
||||
errType = "upstream_timeout"
|
||||
} else if errors.Is(err, bufio.ErrTooLong) {
|
||||
errMsg = "Upstream response line too long"
|
||||
errType = "response_too_large"
|
||||
}
|
||||
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, errType, errMsg)
|
||||
}
|
||||
c.Data(http.StatusOK, "application/json", claudeResp)
|
||||
return streamRes, nil
|
||||
}
|
||||
|
||||
// handleClaudeStreamingResponse 处理 Claude 流式响应(Gemini SSE → Claude SSE 转换)
|
||||
|
||||
@@ -272,6 +272,17 @@ func (s *BillingService) initFallbackPricing() {
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
// Gemini 3.6 Flash (Google AI pricing: $1.50 input / $7.50 output /
|
||||
// $0.15 cached input per MTok). Antigravity's -high/-low/-medium/-tiered
|
||||
// aliases are matched below so unavailable remote pricing never records
|
||||
// token-bearing requests at $0.
|
||||
s.fallbackPrices["gemini-3.6-flash"] = &ModelPricing{
|
||||
InputPricePerToken: 1.5e-6,
|
||||
OutputPricePerToken: 7.5e-6,
|
||||
CacheReadPricePerToken: 0.15e-6,
|
||||
SupportsCacheBreakdown: false,
|
||||
}
|
||||
|
||||
// OpenAI GPT-5.4(业务指定价格)
|
||||
s.fallbackPrices["gpt-5.4"] = &ModelPricing{
|
||||
InputPricePerToken: 2.5e-6, // $2.5 per MTok
|
||||
@@ -630,6 +641,9 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
|
||||
if strings.Contains(modelLower, "gemini-3.1-pro") || strings.Contains(modelLower, "gemini-3-1-pro") {
|
||||
return s.fallbackPrices["gemini-3.1-pro"]
|
||||
}
|
||||
if strings.Contains(modelLower, "gemini-3.6-flash") || strings.Contains(modelLower, "gemini-3-6-flash") {
|
||||
return s.fallbackPrices["gemini-3.6-flash"]
|
||||
}
|
||||
|
||||
// DeepSeek V4 系列:仅匹配已知 V4 Pro/Flash 与官方兼容别名
|
||||
// (deepseek-chat / deepseek-reasoner → V4 Flash),未知 deepseek-* 型号不回退,避免误计价。
|
||||
|
||||
@@ -118,7 +118,12 @@ func normalizeCompositeRouteInput(input CompositeRouteInput) CompositeRouteInput
|
||||
input.TargetPlatform = strings.TrimSpace(input.TargetPlatform)
|
||||
input.UpstreamModel = strings.TrimSpace(input.UpstreamModel)
|
||||
input.Endpoint = normalizeCompositeRouteEndpoint(input.Endpoint)
|
||||
if input.UpstreamModel == "" {
|
||||
// 仅对 exact 路由把空 upstream_model 回填成 public_model:exact 命中时请求模型
|
||||
// 恒等于 public_model,回填只影响持久化/后台展示,保留原有契约不变。
|
||||
// prefix 路由留空则不回填——Resolve 会回退到具体请求模型,从而透传原始模型
|
||||
// (否则 public=deepseek-v4 的前缀路由会把 deepseek-v4-flash / deepseek-v4-pro
|
||||
// 都塌缩成固定的 deepseek-v4)。显式填写 upstream_model 时任何模式都原样固定转发。
|
||||
if input.UpstreamModel == "" && input.MatchType == CompositeRouteMatchExact {
|
||||
input.UpstreamModel = input.PublicModel
|
||||
}
|
||||
input.Notes = strings.TrimSpace(input.Notes)
|
||||
|
||||
@@ -107,6 +107,63 @@ func TestCompositeRouteResolverPrefersEndpointSpecificLongestPrefix(t *testing.T
|
||||
require.Equal(t, int64(2), decision.Route.ID)
|
||||
}
|
||||
|
||||
// TestCompositeRouteResolverPrefixEmptyUpstreamPassesThroughRequestedModel 验证:
|
||||
// 前缀匹配路由留空 upstream_model 时,转发的是具体请求模型(各自原样),而不是
|
||||
// 塌缩成 public_model。这是「留空 = 透传原始模型」语义的核心场景。
|
||||
func TestCompositeRouteResolverPrefixEmptyUpstreamPassesThroughRequestedModel(t *testing.T) {
|
||||
resolver := NewCompositeRouteResolver(compositeRouteRepoStub{
|
||||
routes: []CompositeModelRoute{
|
||||
{
|
||||
ID: 1,
|
||||
GroupID: 7,
|
||||
PublicModel: "deepseek-v4",
|
||||
MatchType: CompositeRouteMatchPrefix,
|
||||
TargetPlatform: PlatformOpenAI,
|
||||
UpstreamModel: "", // 留空 = 透传
|
||||
Endpoint: CompositeRouteEndpointAny,
|
||||
Priority: 100,
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
for _, model := range []string{"deepseek-v4-flash", "deepseek-v4-pro", "deepseek-v4"} {
|
||||
decision, err := resolver.Resolve(context.Background(), 7, model, CompositeRouteEndpointChatCompletions)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Matched, "model %q should match prefix route", model)
|
||||
require.Equal(t, CompositeRouteSourceExplicit, decision.Source)
|
||||
require.Equal(t, PlatformOpenAI, decision.TargetPlatform)
|
||||
require.Equal(t, model, decision.UpstreamModel, "model %q should pass through verbatim", model)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCompositeRouteResolverPrefixExplicitUpstreamStillFixed 验证:前缀匹配路由显式
|
||||
// 填写 upstream_model 时,所有命中请求仍转发同一个固定上游模型(行为不变)。
|
||||
func TestCompositeRouteResolverPrefixExplicitUpstreamStillFixed(t *testing.T) {
|
||||
resolver := NewCompositeRouteResolver(compositeRouteRepoStub{
|
||||
routes: []CompositeModelRoute{
|
||||
{
|
||||
ID: 1,
|
||||
GroupID: 7,
|
||||
PublicModel: "deepseek-v4",
|
||||
MatchType: CompositeRouteMatchPrefix,
|
||||
TargetPlatform: PlatformOpenAI,
|
||||
UpstreamModel: "deepseek-chat",
|
||||
Endpoint: CompositeRouteEndpointAny,
|
||||
Priority: 100,
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
for _, model := range []string{"deepseek-v4-flash", "deepseek-v4-pro"} {
|
||||
decision, err := resolver.Resolve(context.Background(), 7, model, CompositeRouteEndpointChatCompletions)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Matched)
|
||||
require.Equal(t, "deepseek-chat", decision.UpstreamModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompositeRouteResolverIgnoresDisabledRoutesAndFallsBackToDetector(t *testing.T) {
|
||||
resolver := NewCompositeRouteResolver(compositeRouteRepoStub{
|
||||
routes: []CompositeModelRoute{
|
||||
|
||||
@@ -178,6 +178,9 @@ const (
|
||||
// 敏感操作 step-up 2FA 设置
|
||||
SettingKeyStepUpEnabled = "step_up_enabled" // 敏感操作(导出/备份/S3配置/提升管理员等)要求 step-up 2FA,默认关闭
|
||||
|
||||
// 面板 API 限流设置(JSON:PanelRateLimitSettings)
|
||||
SettingKeyPanelRateLimitSettings = "panel_rate_limit_settings"
|
||||
|
||||
// 操作审计日志设置
|
||||
SettingKeyAuditLogRetentionDays = "audit_log_retention_days" // 审计日志保留天数(<=0 永久保留),默认 180
|
||||
|
||||
|
||||
@@ -1208,3 +1208,28 @@ func (s *GatewayService) claudeOAuthSystemPromptInjectionSettings(ctx context.Co
|
||||
}
|
||||
return s.settingService.GetClaudeOAuthSystemPromptInjectionSettings(ctx)
|
||||
}
|
||||
|
||||
// systemHasBillingAttributionBlock 检查请求体的 system 字段中是否包含真实 Claude Code
|
||||
// 客户端注入的 billing attribution block。该 block 格式稳定(见 gateway_billing_block.go),
|
||||
// 仅由真实 Claude Code CLI 生成;第三方客户端(opencode 等)不会生成此 block。
|
||||
//
|
||||
// 用于识别被上游 API 网关代理的真实 Claude Code 流量:此类请求的 User-Agent 被网关替换
|
||||
// 为 Go-http-client,但 body 保留了完整的客户端特征。如果不识别这类请求而走 mimicry
|
||||
// 重写 system,会破坏 Anthropic prompt cache 的前缀一致性,导致 messages 级缓存永不命中。
|
||||
func systemHasBillingAttributionBlock(body []byte) bool {
|
||||
system := gjson.GetBytes(body, "system")
|
||||
if !system.IsArray() {
|
||||
return false
|
||||
}
|
||||
found := false
|
||||
system.ForEach(func(_, item gjson.Result) bool {
|
||||
text := item.Get("text").String()
|
||||
if strings.HasPrefix(text, claudeCodeBillingHeaderPrefix) &&
|
||||
strings.Contains(text, claudeCodeEntrypointMarker) {
|
||||
found = true
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
return found
|
||||
}
|
||||
|
||||
@@ -170,6 +170,17 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
clientUserAgent = c.GetHeader("User-Agent")
|
||||
}
|
||||
isClaudeCode := IsClaudeCodeClient(ctx) || isClaudeCodeClient(clientUserAgent, parsed.MetadataUserID)
|
||||
|
||||
// 补充判定:上游 API 网关(如 new-api)转发真实 Claude Code 流量时,
|
||||
// UA 会变成 Go-http-client 但 body 保留了完整的 Claude Code 特征
|
||||
// (billing attribution block + metadata.user_id)。此时如果仍走 mimicry
|
||||
// 重写 system prompt,会破坏 Anthropic prompt cache 的前缀匹配——
|
||||
// 导致 messages 级缓存永远 miss、cache_creation 每轮全量重写。
|
||||
// 通过检查 body 中的 billing attribution block 来识别被代理的真实 CC 流量。
|
||||
if !isClaudeCode && parsed.MetadataUserID != "" {
|
||||
isClaudeCode = systemHasBillingAttributionBlock(body)
|
||||
}
|
||||
|
||||
shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCode
|
||||
|
||||
if shouldMimicClaudeCode {
|
||||
|
||||
@@ -37,15 +37,21 @@ func (s *GatewayService) ForwardAsResponses(
|
||||
) (*ForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// 1. Parse Responses request
|
||||
// 1. Lower Codex client-side tools to function tools understood by Anthropic.
|
||||
adaptedBody, clientToolMapping, err := adaptResponsesClientToolsForAnthropic(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("adapt responses client tools: %w", err)
|
||||
}
|
||||
|
||||
// 2. Parse Responses request
|
||||
var responsesReq apicompat.ResponsesRequest
|
||||
if err := json.Unmarshal(body, &responsesReq); err != nil {
|
||||
if err := json.Unmarshal(adaptedBody, &responsesReq); err != nil {
|
||||
return nil, fmt.Errorf("parse responses request: %w", err)
|
||||
}
|
||||
originalModel := responsesReq.Model
|
||||
clientStream := responsesReq.Stream
|
||||
|
||||
// 2. Convert Responses → Anthropic
|
||||
// 3. Convert Responses → Anthropic
|
||||
anthropicReq, err := apicompat.ResponsesToAnthropicRequest(&responsesReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("convert responses to anthropic: %w", err)
|
||||
@@ -182,14 +188,71 @@ func (s *GatewayService) ForwardAsResponses(
|
||||
var result *ForwardResult
|
||||
var handleErr error
|
||||
if clientStream {
|
||||
result, handleErr = s.handleResponsesStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
|
||||
result, handleErr = s.handleResponsesStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime, clientToolMapping)
|
||||
} else {
|
||||
result, handleErr = s.handleResponsesBufferedStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
|
||||
result, handleErr = s.handleResponsesBufferedStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime, clientToolMapping)
|
||||
}
|
||||
|
||||
return result, handleErr
|
||||
}
|
||||
|
||||
func adaptResponsesClientToolsForAnthropic(body []byte) ([]byte, apicompat.ResponsesClientToolMapping, error) {
|
||||
decoder := json.NewDecoder(bytes.NewReader(body))
|
||||
decoder.UseNumber()
|
||||
var requestBody map[string]any
|
||||
if err := decoder.Decode(&requestBody); err != nil {
|
||||
return body, apicompat.ResponsesClientToolMapping{}, err
|
||||
}
|
||||
additionalToolsChanged, err := liftResponsesAdditionalTools(requestBody)
|
||||
if err != nil {
|
||||
return body, apicompat.ResponsesClientToolMapping{}, err
|
||||
}
|
||||
|
||||
mapping, changed, err := apicompat.AdaptResponsesClientTools(requestBody)
|
||||
if err != nil {
|
||||
return body, apicompat.ResponsesClientToolMapping{}, err
|
||||
}
|
||||
changed = changed || additionalToolsChanged
|
||||
if !changed {
|
||||
return body, mapping, nil
|
||||
}
|
||||
rebuilt, err := json.Marshal(requestBody)
|
||||
if err != nil {
|
||||
return body, apicompat.ResponsesClientToolMapping{}, err
|
||||
}
|
||||
return rebuilt, mapping, nil
|
||||
}
|
||||
|
||||
func liftResponsesAdditionalTools(requestBody map[string]any) (bool, error) {
|
||||
input, ok := requestBody["input"].([]any)
|
||||
if !ok {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
tools, _ := requestBody["tools"].([]any)
|
||||
kept := make([]any, 0, len(input))
|
||||
changed := false
|
||||
for _, raw := range input {
|
||||
item, ok := raw.(map[string]any)
|
||||
if !ok || strings.TrimSpace(fmt.Sprint(item["type"])) != "additional_tools" {
|
||||
kept = append(kept, raw)
|
||||
continue
|
||||
}
|
||||
additional, ok := item["tools"].([]any)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("additional_tools.tools must be an array")
|
||||
}
|
||||
tools = append(tools, additional...)
|
||||
changed = true
|
||||
}
|
||||
if !changed {
|
||||
return false, nil
|
||||
}
|
||||
requestBody["tools"] = tools
|
||||
requestBody["input"] = kept
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// ExtractResponsesReasoningEffortFromBody reads Responses API reasoning.effort
|
||||
// and normalizes it for usage logging.
|
||||
func ExtractResponsesReasoningEffortFromBody(body []byte) *string {
|
||||
@@ -243,6 +306,7 @@ func (s *GatewayService) handleResponsesBufferedStreamingResponse(
|
||||
mappedModel string,
|
||||
reasoningEffort *string,
|
||||
startTime time.Time,
|
||||
clientToolMapping apicompat.ResponsesClientToolMapping,
|
||||
) (*ForwardResult, error) {
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
|
||||
@@ -358,6 +422,10 @@ func (s *GatewayService) handleResponsesBufferedStreamingResponse(
|
||||
c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8")
|
||||
if respBytes, err := json.Marshal(responsesResp); err == nil {
|
||||
respBytes = reverseToolNamesIfPresent(c, respBytes)
|
||||
respBytes, _, err = apicompat.RestoreResponsesClientToolPayload(respBytes, clientToolMapping)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("restore responses client tools: %w", err)
|
||||
}
|
||||
c.Data(http.StatusOK, "application/json; charset=utf-8", respBytes)
|
||||
} else {
|
||||
c.JSON(http.StatusOK, responsesResp)
|
||||
@@ -383,6 +451,7 @@ func (s *GatewayService) handleResponsesStreamingResponse(
|
||||
mappedModel string,
|
||||
reasoningEffort *string,
|
||||
startTime time.Time,
|
||||
clientToolMapping apicompat.ResponsesClientToolMapping,
|
||||
) (*ForwardResult, error) {
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
|
||||
@@ -397,6 +466,7 @@ func (s *GatewayService) handleResponsesStreamingResponse(
|
||||
|
||||
state := apicompat.NewAnthropicEventToResponsesState()
|
||||
state.Model = originalModel
|
||||
clientToolRestorer := apicompat.NewResponsesClientToolStreamRestorer(clientToolMapping)
|
||||
var usage ClaudeUsage
|
||||
var firstTokenMs *int
|
||||
firstChunk := true
|
||||
@@ -441,7 +511,7 @@ func (s *GatewayService) handleResponsesStreamingResponse(
|
||||
// Convert to Responses events
|
||||
events := apicompat.AnthropicEventToResponsesEvents(event, state)
|
||||
for _, evt := range events {
|
||||
sse, err := apicompat.ResponsesEventToSSE(evt)
|
||||
payload, err := json.Marshal(evt)
|
||||
if err != nil {
|
||||
logger.L().Warn("forward_as_responses stream: failed to marshal event",
|
||||
zap.Error(err),
|
||||
@@ -449,12 +519,23 @@ func (s *GatewayService) handleResponsesStreamingResponse(
|
||||
)
|
||||
continue
|
||||
}
|
||||
out := string(reverseToolNamesIfPresent(c, []byte(sse)))
|
||||
if _, err := fmt.Fprint(c.Writer, out); err != nil {
|
||||
logger.L().Info("forward_as_responses stream: client disconnected",
|
||||
payload = reverseToolNamesIfPresent(c, payload)
|
||||
payloads, _, err := clientToolRestorer.RestoreEvent(payload)
|
||||
if err != nil {
|
||||
logger.L().Warn("forward_as_responses stream: failed to restore client tools",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
return true // client disconnected
|
||||
continue
|
||||
}
|
||||
for _, restored := range payloads {
|
||||
eventType := gjson.GetBytes(restored, "type").String()
|
||||
if _, err := fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", eventType, restored); err != nil {
|
||||
logger.L().Info("forward_as_responses stream: client disconnected",
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
return true // client disconnected
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(events) > 0 {
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -10,10 +11,129 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAdaptResponsesClientToolsForAnthropic_FlattensNamespace(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := []byte(`{
|
||||
"model":"claude-fable-5",
|
||||
"input":[{"type":"function_call","call_id":"call_1","namespace":"codex_app","name":"read_thread","arguments":"{}"}],
|
||||
"tools":[{"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"read_thread","description":"Read a task","parameters":{"type":"object","properties":{}}}]}]
|
||||
}`)
|
||||
|
||||
adapted, mapping, err := adaptResponsesClientToolsForAnthropic(body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, apicompat.ResponsesNamespaceName{Namespace: "codex_app", Name: "read_thread"}, mapping.NamespaceTools["codex_app__read_thread"])
|
||||
|
||||
var request map[string]any
|
||||
require.NoError(t, json.Unmarshal(adapted, &request))
|
||||
tools := request["tools"].([]any)
|
||||
require.Len(t, tools, 1)
|
||||
tool := tools[0].(map[string]any)
|
||||
require.Equal(t, "function", tool["type"])
|
||||
require.Equal(t, "codex_app__read_thread", tool["name"])
|
||||
|
||||
input := request["input"].([]any)
|
||||
call := input[0].(map[string]any)
|
||||
require.Equal(t, "codex_app__read_thread", call["name"])
|
||||
require.NotContains(t, call, "namespace")
|
||||
}
|
||||
|
||||
func TestAdaptResponsesClientToolsForAnthropic_LiftsAdditionalTools(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model":"claude-fable-5",
|
||||
"input":[
|
||||
{"type":"additional_tools","tools":[
|
||||
{"type":"custom","name":"exec","description":"Run a command"},
|
||||
{"type":"namespace","name":"codex_app","tools":[
|
||||
{"type":"function","name":"read_thread","parameters":{"type":"object"}}
|
||||
]}
|
||||
]},
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"inspect"}]}
|
||||
]
|
||||
}`)
|
||||
|
||||
adapted, mapping, err := adaptResponsesClientToolsForAnthropic(body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, mapping.CustomTools["exec"])
|
||||
require.Equal(t, apicompat.ResponsesNamespaceName{Namespace: "codex_app", Name: "read_thread"}, mapping.NamespaceTools["codex_app__read_thread"])
|
||||
|
||||
var request map[string]any
|
||||
require.NoError(t, json.Unmarshal(adapted, &request))
|
||||
tools := request["tools"].([]any)
|
||||
require.Len(t, tools, 2)
|
||||
require.Equal(t, "function", tools[0].(map[string]any)["type"])
|
||||
require.Equal(t, "codex_app__read_thread", tools[1].(map[string]any)["name"])
|
||||
|
||||
input := request["input"].([]any)
|
||||
require.Len(t, input, 1)
|
||||
require.Equal(t, "message", input[0].(map[string]any)["type"])
|
||||
}
|
||||
|
||||
func namespaceToolAnthropicStream() string {
|
||||
return strings.Join([]string{
|
||||
`event: message_start`,
|
||||
`data: {"type":"message_start","message":{"id":"msg_namespace","type":"message","role":"assistant","content":[],"model":"claude-fable-5","stop_reason":"","usage":{"input_tokens":10}}}`,
|
||||
``,
|
||||
`event: content_block_start`,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_namespace","name":"codex_app__read_thread","input":{"thread_id":"123"}}}`,
|
||||
``,
|
||||
`event: content_block_stop`,
|
||||
`data: {"type":"content_block_stop","index":0}`,
|
||||
``,
|
||||
`event: message_delta`,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":5}}`,
|
||||
``,
|
||||
`event: message_stop`,
|
||||
`data: {"type":"message_stop"}`,
|
||||
``,
|
||||
}, "\n")
|
||||
}
|
||||
|
||||
func namespaceToolMapping() apicompat.ResponsesClientToolMapping {
|
||||
return apicompat.ResponsesClientToolMapping{NamespaceTools: map[string]apicompat.ResponsesNamespaceName{
|
||||
"codex_app__read_thread": {Namespace: "codex_app", Name: "read_thread"},
|
||||
}}
|
||||
}
|
||||
|
||||
func TestHandleResponsesBufferedStreamingResponse_RestoresNamespaceTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
resp := &http.Response{Body: io.NopCloser(strings.NewReader(namespaceToolAnthropicStream()))}
|
||||
|
||||
svc := &GatewayService{}
|
||||
_, err := svc.handleResponsesBufferedStreamingResponse(resp, c, "claude-fable-5", "claude-fable-5", nil, time.Now(), namespaceToolMapping())
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, rec.Body.String(), `"type":"function_call"`)
|
||||
require.Contains(t, rec.Body.String(), `"name":"read_thread"`)
|
||||
require.Contains(t, rec.Body.String(), `"namespace":"codex_app"`)
|
||||
require.NotContains(t, rec.Body.String(), `"name":"codex_app__read_thread"`)
|
||||
}
|
||||
|
||||
func TestHandleResponsesStreamingResponse_RestoresNamespaceTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
resp := &http.Response{Body: io.NopCloser(strings.NewReader(namespaceToolAnthropicStream()))}
|
||||
|
||||
svc := &GatewayService{}
|
||||
_, err := svc.handleResponsesStreamingResponse(resp, c, "claude-fable-5", "claude-fable-5", nil, time.Now(), namespaceToolMapping())
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, rec.Body.String(), `response.output_item.added`)
|
||||
require.Contains(t, rec.Body.String(), `"name":"read_thread"`)
|
||||
require.Contains(t, rec.Body.String(), `"namespace":"codex_app"`)
|
||||
require.NotContains(t, rec.Body.String(), `"name":"codex_app__read_thread"`)
|
||||
}
|
||||
|
||||
func TestExtractResponsesReasoningEffortFromBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -51,7 +171,7 @@ func TestHandleResponsesBufferedStreamingResponse_PreservesMessageStartCacheUsag
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleResponsesBufferedStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now())
|
||||
result, err := svc.handleResponsesBufferedStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now(), apicompat.ResponsesClientToolMapping{})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 12, result.Usage.InputTokens)
|
||||
@@ -87,7 +207,7 @@ func TestHandleResponsesStreamingResponse_PreservesMessageStartCacheUsage(t *tes
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleResponsesStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now())
|
||||
result, err := svc.handleResponsesStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now(), apicompat.ResponsesClientToolMapping{})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 20, result.Usage.InputTokens)
|
||||
@@ -198,7 +318,7 @@ func TestHandleResponsesBufferedStreamingResponse_CompactSSEFormat(t *testing.T)
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleResponsesBufferedStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now())
|
||||
result, err := svc.handleResponsesBufferedStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now(), apicompat.ResponsesClientToolMapping{})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 10, result.Usage.InputTokens)
|
||||
@@ -232,7 +352,7 @@ func TestHandleResponsesStreamingResponse_CompactSSEFormat(t *testing.T) {
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleResponsesStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now())
|
||||
result, err := svc.handleResponsesStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now(), apicompat.ResponsesClientToolMapping{})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 15, result.Usage.InputTokens)
|
||||
|
||||
@@ -219,7 +219,15 @@ func TestGetUserGroupRateMultiplier_UsesCacheAndSingleflight(t *testing.T) {
|
||||
}
|
||||
|
||||
close(start)
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
// Wait for every caller to have recorded its cache miss before releasing the
|
||||
// loader. A fixed sleep raced here: a goroutine that reached the cache after
|
||||
// the singleflight load had already finished got a hit instead of a miss, and
|
||||
// the miss assertion below saw 11 of 12. The miss counter is the observable
|
||||
// that says "all callers are now inside the singleflight group".
|
||||
require.Eventually(t, func() bool {
|
||||
_, miss, _, _, _ := GatewayUserGroupRateCacheStats()
|
||||
return miss == int64(concurrent)
|
||||
}, 5*time.Second, time.Millisecond, "all callers must miss the cache before the loader is released")
|
||||
close(unblock)
|
||||
wg.Wait()
|
||||
|
||||
|
||||
@@ -194,6 +194,64 @@ func TestGatewayServiceRecordUsage_PreservesRequestedAndUpstreamModels(t *testin
|
||||
require.Equal(t, mappedModel, *usageRepo.lastLog.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestGatewayServiceRecordUsage_PreservesChannelMappedUpstreamModel(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
|
||||
|
||||
err := svc.RecordUsage(context.Background(), &RecordUsageInput{
|
||||
Result: &ForwardResult{
|
||||
RequestID: "gateway_channel_mapping_models",
|
||||
Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
|
||||
Model: "gpt-5.6-terra",
|
||||
UpstreamModel: "gpt-5.6-terra",
|
||||
Duration: time.Second,
|
||||
},
|
||||
APIKey: &APIKey{ID: 501, Quota: 100},
|
||||
User: &User{ID: 601},
|
||||
Account: &Account{ID: 701},
|
||||
ChannelUsageFields: ChannelUsageFields{
|
||||
OriginalModel: "gpt-5.6-sol",
|
||||
ChannelMappedModel: "gpt-5.6-terra",
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usageRepo.lastLog)
|
||||
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
|
||||
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-terra", *usageRepo.lastLog.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestGatewayServiceRecordUsage_PreservesLoopedChannelAndAccountUpstreamModel(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
|
||||
|
||||
err := svc.RecordUsage(context.Background(), &RecordUsageInput{
|
||||
Result: &ForwardResult{
|
||||
RequestID: "gateway_looped_mapping_models",
|
||||
Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
|
||||
Model: "gpt-5.6-terra",
|
||||
UpstreamModel: "gpt-5.6-sol",
|
||||
Duration: time.Second,
|
||||
},
|
||||
APIKey: &APIKey{ID: 501, Quota: 100},
|
||||
User: &User{ID: 601},
|
||||
Account: &Account{ID: 701},
|
||||
ChannelUsageFields: ChannelUsageFields{
|
||||
OriginalModel: "gpt-5.6-sol",
|
||||
ChannelMappedModel: "gpt-5.6-terra",
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usageRepo.lastLog)
|
||||
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
|
||||
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-sol", *usageRepo.lastLog.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestGatewayServiceRecordUsage_EmptyImageSizeDefaultsBeforeBillingAndPersistence(t *testing.T) {
|
||||
imagePrice2K := 0.19
|
||||
groupID := int64(901)
|
||||
|
||||
@@ -986,7 +986,7 @@ func (s *GatewayService) buildRecordUsageLog(
|
||||
RequestID: requestID,
|
||||
Model: result.Model,
|
||||
RequestedModel: requestedModel,
|
||||
UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model),
|
||||
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
|
||||
ReasoningEffort: result.ReasoningEffort,
|
||||
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
|
||||
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
|
||||
|
||||
@@ -231,7 +231,11 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: evBody}
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: evBody,
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
|
||||
return nil, s.writeGeminiChatCompletionsMappedError(c, account, resp.StatusCode, requestID, evBody)
|
||||
|
||||
@@ -352,6 +352,68 @@ func TestGeminiErrorPolicyIntegration(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestPoolModeSkippedFailoverError — pool-mode accounts hitting
|
||||
// ErrorPolicySkipped must failover (align with other platform forwards)
|
||||
// instead of passing the upstream error through to the client.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestPoolModeSkippedFailoverError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := &GeminiMessagesCompatService{}
|
||||
|
||||
poolAccount := func(extra map[string]any) *Account {
|
||||
creds := map[string]any{"pool_mode": true}
|
||||
for k, v := range extra {
|
||||
creds[k] = v
|
||||
}
|
||||
return &Account{ID: 300, Type: AccountTypeAPIKey, Platform: PlatformGemini, Credentials: creds}
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
statusCode int
|
||||
expectFailover bool
|
||||
expectSameAccount bool
|
||||
}{
|
||||
{"pool_500_failover_no_same_account_retry", poolAccount(nil), 500, true, false},
|
||||
{"pool_429_failover_with_same_account_retry", poolAccount(nil), 429, true, true},
|
||||
{"pool_custom_retry_codes_500", poolAccount(map[string]any{
|
||||
"pool_mode_retry_status_codes": []any{float64(500)},
|
||||
}), 500, true, true},
|
||||
{"pool_400_not_failover_worthy", poolAccount(nil), 400, false, false},
|
||||
{"non_pool_account_keeps_passthrough", &Account{
|
||||
ID: 301, Type: AccountTypeAPIKey, Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
}, 500, false, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
body := []byte(`{"error":{"code":"bad_response_status_code","message":"openai_error"}}`)
|
||||
failoverErr := svc.poolModeSkippedFailoverError(c, tt.account, tt.statusCode, body, "req-1")
|
||||
|
||||
if !tt.expectFailover {
|
||||
require.Nil(t, failoverErr)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, failoverErr)
|
||||
require.Equal(t, tt.statusCode, failoverErr.StatusCode)
|
||||
require.Equal(t, body, failoverErr.ResponseBody)
|
||||
require.Equal(t, tt.expectSameAccount, failoverErr.RetryableOnSameAccount)
|
||||
require.True(t, failoverErr.ShouldRetryNextAccount())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestGeminiErrorPolicy_NilRateLimitService — verifies nil safety
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -945,6 +945,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
if upstreamReqID == "" {
|
||||
upstreamReqID = resp.Header.Get("x-goog-request-id")
|
||||
}
|
||||
if failoverErr := s.poolModeSkippedFailoverError(c, account, resp.StatusCode, respBody, upstreamReqID); failoverErr != nil {
|
||||
return nil, failoverErr
|
||||
}
|
||||
return nil, s.writeGeminiMappedError(c, account, http.StatusInternalServerError, upstreamReqID, respBody)
|
||||
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
|
||||
if policy == ErrorPolicyMatched {
|
||||
@@ -1452,6 +1455,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
|
||||
switch policy {
|
||||
case ErrorPolicySkipped:
|
||||
if failoverErr := s.poolModeSkippedFailoverError(c, account, resp.StatusCode, respBody, requestID); failoverErr != nil {
|
||||
return nil, failoverErr
|
||||
}
|
||||
respBody = unwrapIfNeeded(isOAuth, respBody)
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
@@ -1683,6 +1689,39 @@ func (s *GeminiMessagesCompatService) shouldFailoverGeminiUpstreamError(statusCo
|
||||
}
|
||||
}
|
||||
|
||||
// poolModeSkippedFailoverError 池模式账号命中 ErrorPolicySkipped 时构造 failover 错误:
|
||||
// 可 failover 的状态码返回 UpstreamFailoverError,交给 handler 层按 pool_mode_retry_count
|
||||
// 同账号重试后换号;返回 nil 表示不适用(非池模式或状态码不可 failover),由调用方透传。
|
||||
func (s *GeminiMessagesCompatService) poolModeSkippedFailoverError(c *gin.Context, account *Account, statusCode int, respBody []byte, upstreamRequestID string) *UpstreamFailoverError {
|
||||
if !account.IsPoolMode() || !s.shouldFailoverGeminiUpstreamError(statusCode) {
|
||||
return nil
|
||||
}
|
||||
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
}
|
||||
upstreamDetail = truncateString(string(respBody), maxBytes)
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: statusCode,
|
||||
UpstreamRequestID: upstreamRequestID,
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: statusCode,
|
||||
ResponseBody: respBody,
|
||||
RetryableOnSameAccount: account.IsPoolModeRetryableStatus(statusCode),
|
||||
}
|
||||
}
|
||||
|
||||
func sleepGeminiBackoff(attempt int) {
|
||||
delay := geminiRetryBaseDelay * time.Duration(1<<uint(attempt-1))
|
||||
if delay > geminiRetryMaxDelay {
|
||||
@@ -3428,17 +3467,12 @@ func normalizeGeminiRequestForAIStudio(body []byte) []byte {
|
||||
|
||||
func isClaudeWebSearchToolMap(tool map[string]any) bool {
|
||||
toolType, _ := tool["type"].(string)
|
||||
if strings.HasPrefix(toolType, "web_search") || toolType == "google_search" {
|
||||
return true
|
||||
}
|
||||
|
||||
name, _ := tool["name"].(string)
|
||||
switch strings.TrimSpace(name) {
|
||||
case "web_search", "google_search", "web_search_20250305":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
// A function named web_search is still a client-side function. This is
|
||||
// especially important for Chat Completions clients such as Hermes, whose
|
||||
// built-in runtime tools are represented as ordinary function tools.
|
||||
// Promote only explicitly typed server-side search tools to Gemini's
|
||||
// built-in googleSearch tool.
|
||||
return strings.HasPrefix(toolType, "web_search") || toolType == "google_search"
|
||||
}
|
||||
|
||||
// cleanToolSchema 清理工具的 JSON Schema,移除 Gemini 不支持的字段
|
||||
|
||||
@@ -170,6 +170,74 @@ func TestGeminiForwardAsChatCompletions_StreamsOpenAIChunksFromGeminiSSE(t *test
|
||||
require.Contains(t, out, "data: [DONE]")
|
||||
}
|
||||
|
||||
func TestGeminiForwardAsChatCompletions_FunctionNamedWebSearchStaysClientSide(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
httpStub := &geminiCompatHTTPUpstreamStub{
|
||||
response: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":1}}`,
|
||||
)),
|
||||
},
|
||||
}
|
||||
svc := &GeminiMessagesCompatService{
|
||||
httpUpstream: httpStub,
|
||||
cfg: &config.Config{},
|
||||
}
|
||||
account := &Account{
|
||||
ID: 103,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "gemini-api-key",
|
||||
},
|
||||
Concurrency: 1,
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
body := []byte(`{
|
||||
"model":"gemini-3.6-flash-high",
|
||||
"messages":[{"role":"user","content":"search and read"}],
|
||||
"tools":[
|
||||
{"type":"function","function":{"name":"web_search","description":"Search through the Hermes client","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}},
|
||||
{"type":"function","function":{"name":"read_file","description":"Read a local file","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}}
|
||||
]
|
||||
}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.NotNil(t, httpStub.lastReq)
|
||||
|
||||
postedBody, err := io.ReadAll(httpStub.lastReq.Body)
|
||||
require.NoError(t, err)
|
||||
|
||||
var posted map[string]any
|
||||
require.NoError(t, json.Unmarshal(postedBody, &posted))
|
||||
tools, ok := posted["tools"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, tools, 1, "Chat Completions function tools must not be promoted to Gemini built-ins by name")
|
||||
|
||||
functionTool, ok := tools[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
functionDecls, ok := functionTool["functionDeclarations"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, functionDecls, 2)
|
||||
webSearchDecl, ok := functionDecls[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
readFileDecl, ok := functionDecls[1].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "web_search", webSearchDecl["name"])
|
||||
require.Equal(t, "read_file", readFileDecl["name"])
|
||||
require.NotContains(t, functionTool, "googleSearch")
|
||||
require.NotContains(t, functionTool, "google_search")
|
||||
}
|
||||
|
||||
// TestConvertClaudeToolsToGeminiTools_CustomType 测试custom类型工具转换
|
||||
func TestConvertClaudeToolsToGeminiTools_CustomType(t *testing.T) {
|
||||
tests := []struct {
|
||||
|
||||
@@ -36,6 +36,14 @@ type ollamaUsageTestRepo struct {
|
||||
disableAutoAttempts atomic.Int64
|
||||
disableAutoCalls atomic.Int64
|
||||
groupResolveCalls atomic.Int64
|
||||
getByIDCalls atomic.Int64
|
||||
}
|
||||
|
||||
// GetByID counts loads so a test can wait for a caller to reach the point just
|
||||
// before the singleflight group, instead of guessing with a sleep.
|
||||
func (r *ollamaUsageTestRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
|
||||
r.getByIDCalls.Add(1)
|
||||
return r.upstreamBillingProbeAccountRepo.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func (r *ollamaUsageTestRepo) ListOllamaCloudUsageGroupAccounts(_ context.Context, anchors []*Account) ([]Account, error) {
|
||||
@@ -805,7 +813,19 @@ func TestOllamaCloudUsageRefreshSingleflightAndRunnerDeduplicateSharedGroup(t *t
|
||||
errs := make(chan error, 2)
|
||||
go func() { _, err := svc.Refresh(context.Background(), first.ID); errs <- err }()
|
||||
<-started
|
||||
// The first caller is now parked in the stub, having loaded the account twice
|
||||
// (once to build the group key, once inside the singleflight function).
|
||||
loadsBeforeSecond := repo.getByIDCalls.Load()
|
||||
go func() { _, err := svc.Refresh(context.Background(), second.ID); errs <- err }()
|
||||
// Only release the first caller once the second one has loaded its own
|
||||
// account, which happens immediately before it joins the singleflight group.
|
||||
// Releasing right after starting the goroutine raced: if the first refresh
|
||||
// finished first, the second became a fresh singleflight execution, re-read
|
||||
// the account, saw the LastAttemptAt just written, and failed with the 30s
|
||||
// manual-refresh 429 instead of sharing the in-flight result.
|
||||
require.Eventually(t, func() bool {
|
||||
return repo.getByIDCalls.Load() > loadsBeforeSecond
|
||||
}, 5*time.Second, time.Millisecond, "the second caller must reach the singleflight group before the first is released")
|
||||
close(release)
|
||||
require.NoError(t, <-errs)
|
||||
require.NoError(t, <-errs)
|
||||
|
||||
@@ -2006,6 +2006,37 @@ func TestAccountTestServiceGrokAPIKeyAllowsConfiguredHTTPWhenGlobalPolicyDoes(t
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
func TestAccountTestServiceGrokOAuthPaymentRequiredTemporarilyUnschedulesAccount(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
account := healthyGrokOAuthGatewayTestAccount(56, "access-token")
|
||||
repo := &grokQuotaAccountRepo{}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusPaymentRequired,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"code":"personal-team-blocked:spending-limit"}`)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
accountRepo: repo,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/56/test", nil)
|
||||
before := time.Now()
|
||||
|
||||
err := svc.testGrokAccountConnection(c, account, "grok")
|
||||
|
||||
require.Error(t, err)
|
||||
require.Equal(t, 1, repo.tempUnschedCalls)
|
||||
require.Equal(t, account.ID, repo.lastTempUnschedID)
|
||||
require.Equal(t, "grok payment required", repo.lastTempUnschedReason)
|
||||
require.WithinDuration(t, before.Add(30*time.Minute), repo.lastTempUnschedUntil, time.Second)
|
||||
require.Contains(t, recorder.Body.String(), `"type":"error"`)
|
||||
require.Contains(t, recorder.Body.String(), "Grok Responses API returned 402")
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -1335,6 +1335,71 @@ func TestOpenAIGatewayServiceRecordUsage_UsesRequestedModelAndUpstreamModelMetad
|
||||
require.Equal(t, 1, userRepo.deductCalls)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceRecordUsage_PreservesChannelMappedUpstreamModel(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
userRepo := &openAIRecordUsageUserRepoStub{}
|
||||
subRepo := &openAIRecordUsageSubRepoStub{}
|
||||
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil)
|
||||
|
||||
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
||||
Result: &OpenAIForwardResult{
|
||||
RequestID: "openai_channel_mapping_models",
|
||||
Model: "gpt-5.6-terra",
|
||||
UpstreamModel: "gpt-5.6-terra",
|
||||
Usage: OpenAIUsage{
|
||||
InputTokens: 20,
|
||||
OutputTokens: 10,
|
||||
},
|
||||
Duration: time.Second,
|
||||
},
|
||||
APIKey: &APIKey{ID: 10},
|
||||
User: &User{ID: 20},
|
||||
Account: &Account{ID: 30},
|
||||
ChannelUsageFields: ChannelUsageFields{
|
||||
OriginalModel: "gpt-5.6-sol",
|
||||
ChannelMappedModel: "gpt-5.6-terra",
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usageRepo.lastLog)
|
||||
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
|
||||
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-terra", *usageRepo.lastLog.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceRecordUsage_PreservesLoopedChannelAndAccountUpstreamModel(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
userRepo := &openAIRecordUsageUserRepoStub{}
|
||||
subRepo := &openAIRecordUsageSubRepoStub{}
|
||||
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil)
|
||||
|
||||
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
||||
Result: &OpenAIForwardResult{
|
||||
RequestID: "openai_looped_mapping_models",
|
||||
Model: "gpt-5.6-terra",
|
||||
UpstreamModel: "gpt-5.6-sol",
|
||||
Usage: OpenAIUsage{InputTokens: 20, OutputTokens: 10},
|
||||
Duration: time.Second,
|
||||
},
|
||||
APIKey: &APIKey{ID: 10},
|
||||
User: &User{ID: 20},
|
||||
Account: &Account{ID: 30},
|
||||
ChannelUsageFields: ChannelUsageFields{
|
||||
OriginalModel: "gpt-5.6-sol",
|
||||
ChannelMappedModel: "gpt-5.6-terra",
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usageRepo.lastLog)
|
||||
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
|
||||
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
|
||||
require.Equal(t, "gpt-5.6-sol", *usageRepo.lastLog.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceRecordUsage_BillsMappedRequestsUsingRequestedModel(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
userRepo := &openAIRecordUsageUserRepoStub{}
|
||||
|
||||
@@ -156,6 +156,119 @@ func sanitizeEncryptedReasoningInputItem(item any) (next any, changed bool, keep
|
||||
return inputItem, true, true
|
||||
}
|
||||
|
||||
// SanitizeOpenAICrossModeFailoverReasoning derives a failover attempt body from
|
||||
// the canonical request body by dropping provider-specific encrypted reasoning
|
||||
// input items in full (encrypted_content plus the coupled id/summary shape).
|
||||
//
|
||||
// This is the proactive counterpart to the reactive same-account
|
||||
// invalid_encrypted_content recovery in Forward: when a failover switches from an
|
||||
// OpenAI passthrough account (which forwards upstream-native encrypted reasoning,
|
||||
// e.g. Kiro) to a non-passthrough account (e.g. Bedrock Mantle) that rejects the
|
||||
// provider-specific reasoning IDs/shape, the whole reasoning item must go before
|
||||
// the request reaches the new upstream. Unlike trimOpenAIEncryptedReasoningItems,
|
||||
// which only strips the encrypted_content / null-content fields while preserving
|
||||
// the reasoning item's id and summary, this drops the entire item.
|
||||
//
|
||||
// The input slice is treated as immutable and is never mutated; a distinct slice
|
||||
// is returned only when changed is true.
|
||||
func SanitizeOpenAICrossModeFailoverReasoning(body []byte) (sanitized []byte, changed bool, err error) {
|
||||
if len(body) == 0 {
|
||||
return body, false, nil
|
||||
}
|
||||
if !gjson.GetBytes(body, "input").Exists() {
|
||||
return body, false, nil
|
||||
}
|
||||
var decoded map[string]any
|
||||
decoder := json.NewDecoder(bytes.NewReader(body))
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(&decoded); err != nil {
|
||||
return body, false, fmt.Errorf("decode cross-mode failover body: %w", err)
|
||||
}
|
||||
if !dropOpenAIEncryptedReasoningInputItems(decoded) {
|
||||
return body, false, nil
|
||||
}
|
||||
out, marshalErr := marshalOpenAIUpstreamJSON(decoded)
|
||||
if marshalErr != nil {
|
||||
return body, false, fmt.Errorf("serialize cross-mode failover body: %w", marshalErr)
|
||||
}
|
||||
return out, true, nil
|
||||
}
|
||||
|
||||
// dropOpenAIEncryptedReasoningInputItems removes reasoning input items that carry
|
||||
// provider-specific encrypted_content in full — including their coupled id and
|
||||
// summary — and reports whether anything changed. Contrast with
|
||||
// trimOpenAIEncryptedReasoningItems, which only strips fields while keeping the
|
||||
// reasoning item skeleton.
|
||||
func dropOpenAIEncryptedReasoningInputItems(reqBody map[string]any) bool {
|
||||
if len(reqBody) == 0 {
|
||||
return false
|
||||
}
|
||||
inputValue, has := reqBody["input"]
|
||||
if !has {
|
||||
return false
|
||||
}
|
||||
switch input := inputValue.(type) {
|
||||
case []any:
|
||||
filtered := input[:0]
|
||||
changed := false
|
||||
for _, item := range input {
|
||||
if isOpenAIEncryptedReasoningInputItem(item) {
|
||||
changed = true
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
if !changed {
|
||||
return false
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
delete(reqBody, "input")
|
||||
return true
|
||||
}
|
||||
reqBody["input"] = filtered
|
||||
return true
|
||||
case []map[string]any:
|
||||
filtered := input[:0]
|
||||
changed := false
|
||||
for _, item := range input {
|
||||
if isOpenAIEncryptedReasoningInputItem(item) {
|
||||
changed = true
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
if !changed {
|
||||
return false
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
delete(reqBody, "input")
|
||||
return true
|
||||
}
|
||||
reqBody["input"] = filtered
|
||||
return true
|
||||
case map[string]any:
|
||||
if isOpenAIEncryptedReasoningInputItem(input) {
|
||||
delete(reqBody, "input")
|
||||
return true
|
||||
}
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func isOpenAIEncryptedReasoningInputItem(item any) bool {
|
||||
inputItem, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if itemType, _ := inputItem["type"].(string); strings.TrimSpace(itemType) != "reasoning" {
|
||||
return false
|
||||
}
|
||||
_, has := inputItem["encrypted_content"]
|
||||
return has
|
||||
}
|
||||
|
||||
func IsOpenAIResponsesCompactPathForTest(c *gin.Context) bool {
|
||||
return isOpenAIResponsesCompactPath(c)
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestTrimOpenAIEncryptedReasoningItems_ContentNull(t *testing.T) {
|
||||
@@ -98,6 +99,53 @@ func TestTrimOpenAIEncryptedReasoningItems_NoReasoningItems(t *testing.T) {
|
||||
assert.False(t, changed)
|
||||
}
|
||||
|
||||
func TestSanitizeOpenAICrossModeFailoverReasoning_DropsWholeEncryptedItem(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.1","input":[` +
|
||||
`{"type":"message","role":"user","content":"hi"},` +
|
||||
`{"type":"reasoning","id":"rs_kiro_1","encrypted_content":"ENC","summary":[{"type":"summary_text","text":"t"}]},` +
|
||||
`{"type":"message","role":"assistant","content":"yo"}` +
|
||||
`]}`)
|
||||
|
||||
sanitized, changed, err := SanitizeOpenAICrossModeFailoverReasoning(body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
// The whole reasoning item is gone — id and summary go with encrypted_content,
|
||||
// unlike trimOpenAIEncryptedReasoningItems which keeps the skeleton.
|
||||
require.NotContains(t, string(sanitized), "reasoning")
|
||||
require.NotContains(t, string(sanitized), "rs_kiro_1")
|
||||
require.NotContains(t, string(sanitized), "summary_text")
|
||||
require.Equal(t, int64(2), gjson.GetBytes(sanitized, "input.#").Int())
|
||||
}
|
||||
|
||||
func TestSanitizeOpenAICrossModeFailoverReasoning_NoEncryptedIsNoop(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.1","input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"t"}]}]}`)
|
||||
sanitized, changed, err := SanitizeOpenAICrossModeFailoverReasoning(body)
|
||||
require.NoError(t, err)
|
||||
require.False(t, changed, "reasoning without encrypted_content must be preserved")
|
||||
require.Equal(t, string(body), string(sanitized))
|
||||
}
|
||||
|
||||
func TestSanitizeOpenAICrossModeFailoverReasoning_NoInputIsNoop(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.1"}`)
|
||||
sanitized, changed, err := SanitizeOpenAICrossModeFailoverReasoning(body)
|
||||
require.NoError(t, err)
|
||||
require.False(t, changed)
|
||||
require.Equal(t, string(body), string(sanitized))
|
||||
}
|
||||
|
||||
func TestSanitizeOpenAICrossModeFailoverReasoning_PreservesLargeIntegers(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.1","input":[` +
|
||||
`{"type":"reasoning","id":"rs_kiro_1","encrypted_content":"ENC"},` +
|
||||
`{"type":"message","role":"user","content":"hi"}` +
|
||||
`],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object","properties":{"id":{"const":9007199254740993}}}}}]}`)
|
||||
|
||||
sanitized, changed, err := SanitizeOpenAICrossModeFailoverReasoning(body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Contains(t, string(sanitized), `"const":9007199254740993`,
|
||||
"sanitization must not round JSON integers through float64")
|
||||
}
|
||||
|
||||
func TestTrimOpenAIEncryptedReasoningItems_ContentNullDropsBareSkeleton(t *testing.T) {
|
||||
reqBody := map[string]any{
|
||||
"input": []any{
|
||||
|
||||
@@ -256,7 +256,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
RequestID: requestID,
|
||||
Model: result.Model,
|
||||
RequestedModel: requestedModel,
|
||||
UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model),
|
||||
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
|
||||
ServiceTier: result.ServiceTier,
|
||||
ReasoningEffort: result.ReasoningEffort,
|
||||
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
|
||||
|
||||
@@ -207,8 +207,9 @@ func (e *OpenAIWSClientCloseError) Reason() string {
|
||||
|
||||
// OpenAIWSIngressHooks 定义入站 WS 每个 turn 的生命周期回调。
|
||||
type OpenAIWSIngressHooks struct {
|
||||
// InitialRequestModel 是首帧渠道映射前的请求模型,只用于 usage metadata
|
||||
// 的 reasoning effort 后缀推导,禁止用于上游请求或计费模型。
|
||||
// InitialRequestModel is the client-facing model from the first frame,
|
||||
// before channel or account mapping. Ingress modes preserve it for usage
|
||||
// attribution while MapRequestModel determines the upstream model.
|
||||
InitialRequestModel string
|
||||
// MaxReasoningEffort limits explicit reasoning effort values for this WS session.
|
||||
MaxReasoningEffort string
|
||||
@@ -216,7 +217,10 @@ type OpenAIWSIngressHooks struct {
|
||||
ReasoningEffortMappings []ReasoningEffortMapping
|
||||
BeforeTurn func(turn int) error
|
||||
BeforeRequest func(turn int, payload []byte, originalModel string) error
|
||||
AfterTurn func(turn int, result *OpenAIForwardResult, turnErr error)
|
||||
// MapRequestModel resolves the current turn's client model to the model
|
||||
// that must be written into the upstream response.create frame.
|
||||
MapRequestModel func(turn int, originalModel string) (string, error)
|
||||
AfterTurn func(turn int, result *OpenAIForwardResult, turnErr error)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) getOpenAIWSConnPool() *openAIWSConnPool {
|
||||
|
||||
@@ -163,7 +163,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
return rebuilt, nil
|
||||
}
|
||||
|
||||
parseClientPayload := func(raw []byte) (openAIWSClientPayload, error) {
|
||||
parseClientPayload := func(turn int, raw []byte) (openAIWSClientPayload, error) {
|
||||
trimmed := bytes.TrimSpace(raw)
|
||||
if len(trimmed) == 0 {
|
||||
return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "empty websocket request payload", nil)
|
||||
@@ -288,7 +288,17 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
normalized = rebuilt
|
||||
}
|
||||
}
|
||||
upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel))
|
||||
requestModel := originalModel
|
||||
if hooks != nil && hooks.MapRequestModel != nil {
|
||||
mappedModel, mapErr := hooks.MapRequestModel(turn, originalModel)
|
||||
if mapErr != nil {
|
||||
return openAIWSClientPayload{}, mapErr
|
||||
}
|
||||
if mappedModel = strings.TrimSpace(mappedModel); mappedModel != "" {
|
||||
requestModel = mappedModel
|
||||
}
|
||||
}
|
||||
upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(requestModel))
|
||||
if modelMissing || upstreamModel != originalModel {
|
||||
next, setErr := applyPayloadMutation(normalized, "model", upstreamModel)
|
||||
if setErr != nil {
|
||||
@@ -412,7 +422,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
firstPayload, err := parseClientPayload(firstClientMessage)
|
||||
firstPayload, err := parseClientPayload(1, firstClientMessage)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -513,7 +523,13 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
grokCacheIdentity := ""
|
||||
if account.Platform == PlatformGrok {
|
||||
grokCacheIdentity, err = resolveGrokWSCacheIdentity(c, account, grokCacheSeedPayload, currentBridgePayload.originalModel)
|
||||
grokCacheIdentity, err = resolveGrokWSCacheIdentity(
|
||||
c,
|
||||
account,
|
||||
grokCacheSeedPayload,
|
||||
currentBridgePayload.payloadRaw,
|
||||
currentBridgePayload.originalModel,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolve Grok websocket cache identity: %w", err)
|
||||
}
|
||||
@@ -573,7 +589,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
return fmt.Errorf("read client websocket request: %w", readErr)
|
||||
}
|
||||
nextPayload, parseErr := parseClientPayload(nextClientMessage)
|
||||
nextPayload, parseErr := parseClientPayload(turn+1, nextClientMessage)
|
||||
if parseErr != nil {
|
||||
return parseErr
|
||||
}
|
||||
@@ -792,7 +808,10 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
mappedModel := ""
|
||||
var mappedModelBytes []byte
|
||||
if originalModel != "" {
|
||||
mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel))
|
||||
mappedModel = strings.TrimSpace(gjson.GetBytes(payload, "model").String())
|
||||
if mappedModel == "" {
|
||||
mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel))
|
||||
}
|
||||
needModelReplace = mappedModel != "" && mappedModel != originalModel
|
||||
if needModelReplace {
|
||||
mappedModelBytes = []byte(mappedModel)
|
||||
@@ -1589,7 +1608,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
return fmt.Errorf("read client websocket request: %w", readErr)
|
||||
}
|
||||
|
||||
nextPayload, parseErr := parseClientPayload(nextClientMessage)
|
||||
nextPayload, parseErr := parseClientPayload(turn+1, nextClientMessage)
|
||||
if parseErr != nil {
|
||||
return parseErr
|
||||
}
|
||||
|
||||
@@ -286,7 +286,10 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
needModelReplace := false
|
||||
var mappedModelBytes []byte
|
||||
if originalModel != "" {
|
||||
mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel))
|
||||
mappedModel = strings.TrimSpace(gjson.GetBytes(body, "model").String())
|
||||
if mappedModel == "" {
|
||||
mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel))
|
||||
}
|
||||
needModelReplace = mappedModel != "" && mappedModel != originalModel
|
||||
if needModelReplace {
|
||||
mappedModelBytes = []byte(mappedModel)
|
||||
@@ -487,12 +490,12 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
return resultWithUsage(), terminalErr
|
||||
}
|
||||
|
||||
func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte, originalModel string) (string, error) {
|
||||
body, err := prepareOpenAIWSHTTPBridgeBody(payload)
|
||||
func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, seedPayload, currentPayload []byte, originalModel string) (string, error) {
|
||||
body, err := prepareOpenAIWSHTTPBridgeBody(seedPayload)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel)
|
||||
upstreamModel := resolveGrokWSUpstreamModel(account, currentPayload, originalModel)
|
||||
body, err = patchGrokResponsesBody(body, upstreamModel)
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -502,7 +505,11 @@ func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte
|
||||
|
||||
func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel string) string {
|
||||
upstreamModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
|
||||
if account != nil && originalModel != "" {
|
||||
originalModel = strings.TrimSpace(originalModel)
|
||||
// Shared ingress has already applied channel and account mappings when the
|
||||
// body model differs from the client-facing model. Only resolve from the
|
||||
// original model when the body still carries that original value.
|
||||
if account != nil && originalModel != "" && (upstreamModel == "" || upstreamModel == originalModel) {
|
||||
if mappedModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)); mappedModel != "" {
|
||||
upstreamModel = mappedModel
|
||||
}
|
||||
|
||||
@@ -520,7 +520,7 @@ func TestProxyOpenAIWSHTTPBridgeTurnPromotesCodexAdditionalToolsForMixedCache(t
|
||||
require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader))
|
||||
}
|
||||
|
||||
func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) {
|
||||
func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridgeAndPreservesMappedModels(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
bridgeResponse := func(responseID, requestID string, cachedTokens int) *http.Response {
|
||||
@@ -594,7 +594,14 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T)
|
||||
ginCtx.Request = req
|
||||
ginCtx.Set("api_key", &APIKey{ID: 7101})
|
||||
|
||||
errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "access-token", firstMessage, nil)
|
||||
errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "access-token", firstMessage, &OpenAIWSIngressHooks{
|
||||
MapRequestModel: func(_ int, originalModel string) (string, error) {
|
||||
if originalModel == "channel-alias" {
|
||||
return "grok-4.3", nil
|
||||
}
|
||||
return originalModel, nil
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer wsServer.Close()
|
||||
|
||||
@@ -626,7 +633,7 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T)
|
||||
require.Equal(t, "resp_grok_ws_1", gjson.GetBytes(completed, "response.id").String())
|
||||
|
||||
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"grok","stream":true,"previous_response_id":"resp_grok_ws_1","input":"second turn"}`))
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"channel-alias","stream":true,"previous_response_id":"resp_grok_ws_1","input":"second turn"}`))
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -637,6 +644,7 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T)
|
||||
require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String())
|
||||
require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
|
||||
require.Equal(t, "resp_grok_ws_2", gjson.GetBytes(completed, "response.id").String())
|
||||
require.Equal(t, "channel-alias", gjson.GetBytes(completed, "response.model").String())
|
||||
|
||||
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"grok-4.3","stream":true,"previous_response_id":"resp_grok_ws_2","input":"third turn with a different model"}`))
|
||||
@@ -666,7 +674,7 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T)
|
||||
require.Equal(t, "sub2api-grok/1.0", upstream.lastReq.Header.Get("User-Agent"))
|
||||
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.bodies[0], "model").String())
|
||||
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.bodies[1], "model").String())
|
||||
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.bodies[1], "model").String())
|
||||
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.bodies[2], "model").String())
|
||||
require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
|
||||
require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader))
|
||||
@@ -677,9 +685,10 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T)
|
||||
secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
|
||||
thirdIdentity := gjson.GetBytes(upstream.bodies[2], "prompt_cache_key").String()
|
||||
require.NotEmpty(t, firstIdentity)
|
||||
require.Equal(t, firstIdentity, secondIdentity)
|
||||
require.NotEmpty(t, secondIdentity)
|
||||
require.NotEqual(t, firstIdentity, secondIdentity)
|
||||
require.NotEmpty(t, thirdIdentity)
|
||||
require.NotEqual(t, firstIdentity, thirdIdentity)
|
||||
require.Equal(t, secondIdentity, thirdIdentity)
|
||||
require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
|
||||
require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
|
||||
require.Equal(t, thirdIdentity, upstream.requests[2].Header.Get(grokConversationIDHeader))
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
@@ -86,6 +87,7 @@ type RelayTraceEvent struct {
|
||||
|
||||
type relayState struct {
|
||||
usage Usage
|
||||
requestModelMu sync.RWMutex
|
||||
requestModel string
|
||||
lastResponseID string
|
||||
terminalEventType string
|
||||
@@ -164,6 +166,12 @@ func Relay(
|
||||
defer cancel()
|
||||
return upstreamConn.WriteFrame(writeCtx, msgType, payload)
|
||||
}
|
||||
writeClientFrameUpstream := func(msgType coderws.MessageType, payload []byte) error {
|
||||
if msgType == coderws.MessageText && strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" {
|
||||
state.setRequestModel(strings.TrimSpace(gjson.GetBytes(payload, "model").String()))
|
||||
}
|
||||
return writeUpstream(msgType, payload)
|
||||
}
|
||||
writeClient := func(msgType coderws.MessageType, payload []byte) error {
|
||||
writeCtx, cancel := context.WithTimeout(relayCtx, writeTimeout)
|
||||
defer cancel()
|
||||
@@ -215,7 +223,7 @@ func Relay(
|
||||
if !clientReaderStarted.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
go runClientToUpstream(relayCtx, clientConn, options.ReadClientFrame, writeUpstream, markActivity, clientToUpstreamFrames, onTrace, exitCh)
|
||||
go runClientToUpstream(relayCtx, clientConn, options.ReadClientFrame, writeClientFrameUpstream, markActivity, clientToUpstreamFrames, onTrace, exitCh)
|
||||
}
|
||||
if !options.StartClientAfterFirstDownstream {
|
||||
startClientReader()
|
||||
@@ -712,7 +720,7 @@ func emitTurnComplete(
|
||||
}
|
||||
requestModel := ""
|
||||
if state != nil {
|
||||
requestModel = state.requestModel
|
||||
requestModel = state.currentRequestModel()
|
||||
}
|
||||
onTurnComplete(RelayTurnResult{
|
||||
RequestModel: requestModel,
|
||||
@@ -873,13 +881,31 @@ func enrichResult(result *RelayResult, state *relayState, duration time.Duration
|
||||
if state == nil {
|
||||
return
|
||||
}
|
||||
result.RequestModel = state.requestModel
|
||||
result.RequestModel = state.currentRequestModel()
|
||||
result.Usage = state.usage
|
||||
result.RequestID = state.lastResponseID
|
||||
result.TerminalEventType = state.terminalEventType
|
||||
result.FirstTokenMs = state.firstTokenMs
|
||||
}
|
||||
|
||||
func (s *relayState) setRequestModel(model string) {
|
||||
if s == nil || model == "" {
|
||||
return
|
||||
}
|
||||
s.requestModelMu.Lock()
|
||||
s.requestModel = model
|
||||
s.requestModelMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *relayState) currentRequestModel() string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
s.requestModelMu.RLock()
|
||||
defer s.requestModelMu.RUnlock()
|
||||
return s.requestModel
|
||||
}
|
||||
|
||||
func isDisconnectError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
|
||||
@@ -472,6 +472,56 @@ func TestRelay_OnTurnComplete_PerTerminalEvent(t *testing.T) {
|
||||
require.Equal(t, 5, result.Usage.OutputTokens)
|
||||
}
|
||||
|
||||
func TestRelay_OnTurnComplete_UsesCurrentResponseCreateModel(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
clientConn := newPassthroughTestFrameConn(nil, false)
|
||||
upstreamConn := newPassthroughTestFrameConn(nil, false)
|
||||
firstPayload := []byte(`{"type":"response.create","model":"gpt-5.6-sol","input":[]}`)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
turns := make(chan RelayTurnResult, 2)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
_, _ = Relay(ctx, clientConn, upstreamConn, firstPayload, RelayOptions{
|
||||
OnTurnComplete: func(turn RelayTurnResult) {
|
||||
turns <- turn
|
||||
},
|
||||
})
|
||||
}()
|
||||
|
||||
upstreamConn.readCh <- passthroughTestFrame{
|
||||
msgType: coderws.MessageText,
|
||||
payload: []byte(`{"type":"response.completed","response":{"id":"resp_sol","usage":{"input_tokens":2,"output_tokens":1}}}`),
|
||||
}
|
||||
firstTurn := <-turns
|
||||
require.Equal(t, "gpt-5.6-sol", firstTurn.RequestModel)
|
||||
|
||||
clientConn.readCh <- passthroughTestFrame{
|
||||
msgType: coderws.MessageText,
|
||||
payload: []byte(`{"type":"response.create","model":"gpt-5.6-terra","input":[]}`),
|
||||
}
|
||||
require.Eventually(t, func() bool {
|
||||
return len(upstreamConn.Writes()) == 2
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
upstreamConn.readCh <- passthroughTestFrame{
|
||||
msgType: coderws.MessageText,
|
||||
payload: []byte(`{"type":"response.completed","response":{"id":"resp_terra","usage":{"input_tokens":3,"output_tokens":1}}}`),
|
||||
}
|
||||
secondTurn := <-turns
|
||||
require.Equal(t, "gpt-5.6-terra", secondTurn.RequestModel)
|
||||
|
||||
_ = clientConn.Close()
|
||||
_ = upstreamConn.Close()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("relay did not stop after client close")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRelay_OnTurnComplete_ProvidesTurnMetrics(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user